datacontract-cli 0.10.0__py3-none-any.whl → 0.10.37__py3-none-any.whl

datacontract/export/dbt_converter.py

@@ -1,42 +1,59 @@
-from typing import Dict
+from typing import Dict, Optional
 
 import yaml
 
+from datacontract.export.exporter import Exporter, _check_models_for_export
 from datacontract.export.sql_type_converter import convert_to_sql_type
-from datacontract.model.data_contract_specification import \
-    DataContractSpecification, Model, Field
+from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model
 
 
-def to_dbt_models_yaml(data_contract_spec: DataContractSpecification):
+class DbtExporter(Exporter):
+    def export(self, data_contract, model, server, sql_server_type, export_args) -> dict:
+        return to_dbt_models_yaml(data_contract, server)
+
+
+class DbtSourceExporter(Exporter):
+    def export(self, data_contract, model, server, sql_server_type, export_args) -> dict:
+        return to_dbt_sources_yaml(data_contract, server)
+
+
+class DbtStageExporter(Exporter):
+    def export(self, data_contract, model, server, sql_server_type, export_args) -> dict:
+        model_name, model_value = _check_models_for_export(data_contract, model, self.export_format)
+        return to_dbt_staging_sql(
+            data_contract,
+            model_name,
+            model_value,
+        )
+
+
+def to_dbt_models_yaml(data_contract_spec: DataContractSpecification, server: str = None) -> str:
     dbt = {
         "version": 2,
         "models": [],
     }
+
     for model_key, model_value in data_contract_spec.models.items():
-        dbt_model = _to_dbt_model(model_key, model_value, data_contract_spec)
+        dbt_model = _to_dbt_model(model_key, model_value, data_contract_spec, adapter_type=server)
         dbt["models"].append(dbt_model)
-    return yaml.dump(dbt, indent=2, sort_keys=False, allow_unicode=True)
+    return yaml.safe_dump(dbt, indent=2, sort_keys=False, allow_unicode=True)
 
 
 def to_dbt_staging_sql(data_contract_spec: DataContractSpecification, model_name: str, model_value: Model) -> str:
-    if data_contract_spec.models is None or len(data_contract_spec.models.items()) != 1:
-        print("Export to dbt-staging-sql currently only works with exactly one model in the data contract.")
-        return ""
-
     id = data_contract_spec.id
     columns = []
     for field_name, field in model_value.fields.items():
         # TODO escape SQL reserved key words, probably dependent on server type
         columns.append(field_name)
     return f"""
-    select 
+    select
         {", ".join(columns)}
     from {{{{ source('{id}', '{model_name}') }}}}
 """
 
 
 def to_dbt_sources_yaml(data_contract_spec: DataContractSpecification, server: str = None):
-    source = {"name": data_contract_spec.id, "tables": []}
+    source = {"name": data_contract_spec.id}
     dbt = {
         "version": 2,
         "sources": [source],
@@ -44,38 +61,52 @@ def to_dbt_sources_yaml(data_contract_spec: DataContractSpecification, server: s
     if data_contract_spec.info.owner is not None:
         source["meta"] = {"owner": data_contract_spec.info.owner}
     if data_contract_spec.info.description is not None:
-        source["description"] = data_contract_spec.info.description
+        source["description"] = data_contract_spec.info.description.strip().replace("\n", " ")
     found_server = data_contract_spec.servers.get(server)
+    adapter_type = None
     if found_server is not None:
-        source["database"] = found_server.database
-        source["schema"] = found_server.schema_
+        adapter_type = found_server.type
+        if adapter_type == "bigquery":
+            source["database"] = found_server.project
+            source["schema"] = found_server.dataset
+        else:
+            source["database"] = found_server.database
+            source["schema"] = found_server.schema_
 
+    source["tables"] = []
     for model_key, model_value in data_contract_spec.models.items():
-        dbt_model = _to_dbt_source_table(model_key, model_value)
+        dbt_model = _to_dbt_source_table(data_contract_spec, model_key, model_value, adapter_type)
         source["tables"].append(dbt_model)
     return yaml.dump(dbt, indent=2, sort_keys=False, allow_unicode=True)
 
 
-def _to_dbt_source_table(model_key, model_value: Model) -> dict:
+def _to_dbt_source_table(
+    data_contract_spec: DataContractSpecification, model_key, model_value: Model, adapter_type: Optional[str]
+) -> dict:
     dbt_model = {
         "name": model_key,
     }
 
     if model_value.description is not None:
-        dbt_model["description"] = model_value.description
-    columns = _to_columns(model_value.fields, False, False)
+        dbt_model["description"] = model_value.description.strip().replace("\n", " ")
+    columns = _to_columns(data_contract_spec, model_value.fields, False, adapter_type)
     if columns:
         dbt_model["columns"] = columns
     return dbt_model
 
 
-def _to_dbt_model(model_key, model_value: Model, data_contract_spec: DataContractSpecification) -> dict:
+def _to_dbt_model(
+    model_key, model_value: Model, data_contract_spec: DataContractSpecification, adapter_type: Optional[str]
+) -> dict:
     dbt_model = {
         "name": model_key,
     }
     model_type = _to_dbt_model_type(model_value.type)
+
     dbt_model["config"] = {"meta": {"data_contract": data_contract_spec.id}}
-    dbt_model["config"]["materialized"] = model_type
+
+    if model_type:
+        dbt_model["config"]["materialized"] = model_type
 
     if data_contract_spec.info.owner is not None:
         dbt_model["config"]["meta"]["owner"] = data_contract_spec.info.owner
@@ -83,10 +114,29 @@ def _to_dbt_model(model_key, model_value: Model, data_contract_spec: DataContrac
     if _supports_constraints(model_type):
         dbt_model["config"]["contract"] = {"enforced": True}
     if model_value.description is not None:
-        dbt_model["description"] = model_value.description
-    columns = _to_columns(model_value.fields, _supports_constraints(model_type), True)
+        dbt_model["description"] = model_value.description.strip().replace("\n", " ")
+
+    # Handle model-level primaryKey (before columns for better YAML ordering)
+    primary_key_columns = []
+    if hasattr(model_value, "primaryKey") and model_value.primaryKey:
+        if isinstance(model_value.primaryKey, list) and len(model_value.primaryKey) > 1:
+            # Multiple columns: use dbt_utils.unique_combination_of_columns
+            dbt_model["data_tests"] = [
+                {"dbt_utils.unique_combination_of_columns": {"combination_of_columns": model_value.primaryKey}}
+            ]
+        elif isinstance(model_value.primaryKey, list) and len(model_value.primaryKey) == 1:
+            # Single column: handle at column level (pass to _to_columns)
+            primary_key_columns = model_value.primaryKey
+        elif isinstance(model_value.primaryKey, str):
+            # Single column as string: handle at column level
+            primary_key_columns = [model_value.primaryKey]
+
+    columns = _to_columns(
+        data_contract_spec, model_value.fields, _supports_constraints(model_type), adapter_type, primary_key_columns
+    )
     if columns:
         dbt_model["columns"] = columns
+
     return dbt_model
 
 
@@ -95,7 +145,7 @@ def _to_dbt_model_type(model_type):
     # Allowed values: table, view, incremental, ephemeral, materialized view
     # Custom values also possible
     if model_type is None:
-        return "table"
+        return None
     if model_type.lower() == "table":
         return "table"
     if model_type.lower() == "view":
@@ -107,48 +157,72 @@ def _supports_constraints(model_type):
     return model_type == "table" or model_type == "incremental"
 
 
-def _to_columns(fields: Dict[str, Field], supports_constraints: bool, supports_datatype: bool) -> list:
+def _to_columns(
+    data_contract_spec: DataContractSpecification,
+    fields: Dict[str, Field],
+    supports_constraints: bool,
+    adapter_type: Optional[str],
+    primary_key_columns: Optional[list] = None,
+) -> list:
     columns = []
+    primary_key_columns = primary_key_columns or []
     for field_name, field in fields.items():
-        column = _to_column(field, supports_constraints, supports_datatype)
-        column["name"] = field_name
+        is_primary_key = field_name in primary_key_columns
+        column = _to_column(data_contract_spec, field_name, field, supports_constraints, adapter_type, is_primary_key)
         columns.append(column)
     return columns
 
 
-def _to_column(field: Field, supports_constraints: bool, supports_datatype: bool) -> dict:
-    column = {}
-    dbt_type = convert_to_sql_type(field, "snowflake")
+def get_table_name_and_column_name(references: str) -> tuple[Optional[str], str]:
+    parts = references.split(".")
+    if len(parts) < 2:
+        return None, parts[0]
+    return parts[-2], parts[-1]
+
+
+def _to_column(
+    data_contract_spec: DataContractSpecification,
+    field_name: str,
+    field: Field,
+    supports_constraints: bool,
+    adapter_type: Optional[str],
+    is_primary_key: bool = False,
+) -> dict:
+    column = {"name": field_name}
+    adapter_type = adapter_type or "snowflake"
+    dbt_type = convert_to_sql_type(field, adapter_type)
+
+    column["data_tests"] = []
     if dbt_type is not None:
-        if supports_datatype:
-            column["data_type"] = dbt_type
-        else:
-            column.setdefault("tests", []).append(
-                {"dbt_expectations.dbt_expectations.expect_column_values_to_be_of_type": {"column_type": dbt_type}}
-            )
+        column["data_type"] = dbt_type
+    else:
+        column["data_tests"].append(
+            {"dbt_expectations.dbt_expectations.expect_column_values_to_be_of_type": {"column_type": dbt_type}}
+        )
     if field.description is not None:
-        column["description"] = field.description
-    if field.required:
+        column["description"] = field.description.strip().replace("\n", " ")
+    # Handle required/not_null constraint
+    if field.required or is_primary_key:
         if supports_constraints:
             column.setdefault("constraints", []).append({"type": "not_null"})
         else:
-            column.setdefault("tests", []).append("not_null")
-    if field.unique:
+            column["data_tests"].append("not_null")
+
+    # Handle unique constraint
+    if field.unique or is_primary_key:
         if supports_constraints:
             column.setdefault("constraints", []).append({"type": "unique"})
         else:
-            column.setdefault("tests", []).append("unique")
+            column["data_tests"].append("unique")
     if field.enum is not None and len(field.enum) > 0:
-        column.setdefault("tests", []).append({"accepted_values": {"values": field.enum}})
+        column["data_tests"].append({"accepted_values": {"values": field.enum}})
     if field.minLength is not None or field.maxLength is not None:
         length_test = {}
         if field.minLength is not None:
             length_test["min_value"] = field.minLength
         if field.maxLength is not None:
             length_test["max_value"] = field.maxLength
-        column.setdefault("tests", []).append(
-            {"dbt_expectations.expect_column_value_lengths_to_be_between": length_test}
-        )
+        column["data_tests"].append({"dbt_expectations.expect_column_value_lengths_to_be_between": length_test})
     if field.pii is not None:
         column.setdefault("meta", {})["pii"] = field.pii
     if field.classification is not None:
@@ -157,9 +231,7 @@ def _to_column(field: Field, supports_constraints: bool, supports_datatype: bool
         column.setdefault("tags", []).extend(field.tags)
     if field.pattern is not None:
         # Beware, the data contract pattern is a regex, not a like pattern
-        column.setdefault("tests", []).append(
-            {"dbt_expectations.expect_column_values_to_match_regex": {"regex": field.pattern}}
-        )
+        column["data_tests"].append({"dbt_expectations.expect_column_values_to_match_regex": {"regex": field.pattern}})
     if (
         field.minimum is not None
         or field.maximum is not None
@@ -171,7 +243,7 @@ def _to_column(field: Field, supports_constraints: bool, supports_datatype: bool
             range_test["min_value"] = field.minimum
         if field.maximum is not None:
             range_test["max_value"] = field.maximum
-        column.setdefault("tests", []).append({"dbt_expectations.expect_column_values_to_be_between": range_test})
+        column["data_tests"].append({"dbt_expectations.expect_column_values_to_be_between": range_test})
     elif (
         field.exclusiveMinimum is not None
         or field.exclusiveMaximum is not None
@@ -184,18 +256,18 @@ def _to_column(field: Field, supports_constraints: bool, supports_datatype: bool
         if field.exclusiveMaximum is not None:
             range_test["max_value"] = field.exclusiveMaximum
         range_test["strictly"] = True
-        column.setdefault("tests", []).append({"dbt_expectations.expect_column_values_to_be_between": range_test})
+        column["data_tests"].append({"dbt_expectations.expect_column_values_to_be_between": range_test})
     else:
         if field.minimum is not None:
-            column.setdefault("tests", []).append(
+            column["data_tests"].append(
                 {"dbt_expectations.expect_column_values_to_be_between": {"min_value": field.minimum}}
             )
         if field.maximum is not None:
-            column.setdefault("tests", []).append(
+            column["data_tests"].append(
                 {"dbt_expectations.expect_column_values_to_be_between": {"max_value": field.maximum}}
             )
         if field.exclusiveMinimum is not None:
-            column.setdefault("tests", []).append(
+            column["data_tests"].append(
                 {
                     "dbt_expectations.expect_column_values_to_be_between": {
                         "min_value": field.exclusiveMinimum,
@@ -204,7 +276,7 @@ def _to_column(field: Field, supports_constraints: bool, supports_datatype: bool
                 }
             )
         if field.exclusiveMaximum is not None:
-            column.setdefault("tests", []).append(
+            column["data_tests"].append(
                 {
                     "dbt_expectations.expect_column_values_to_be_between": {
                         "max_value": field.exclusiveMaximum,
@@ -212,6 +284,21 @@ def _to_column(field: Field, supports_constraints: bool, supports_datatype: bool
                     }
                 }
             )
+    if field.references is not None:
+        ref_source_name = data_contract_spec.id
+        table_name, column_name = get_table_name_and_column_name(field.references)
+        if table_name is not None and column_name is not None:
+            column["data_tests"].append(
+                {
+                    "relationships": {
+                        "to": f"""source("{ref_source_name}", "{table_name}")""",
+                        "field": f"{column_name}",
+                    }
+                }
+            )
+
+    if not column["data_tests"]:
+        column.pop("data_tests")
 
     # TODO: all constraints
     return column

datacontract/export/dcs_exporter.py

@@ -0,0 +1,6 @@
+from datacontract.export.exporter import Exporter
+
+
+class DcsExporter(Exporter):
+    def export(self, data_contract, model, server, sql_server_type, export_args) -> dict:
+        return data_contract.to_yaml()

datacontract/export/dqx_converter.py

@@ -0,0 +1,126 @@
+from typing import Any, Dict, List, Union
+
+import yaml
+
+from datacontract.export.exporter import Exporter, _check_models_for_export
+from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model, Quality
+
+
+class DqxKeys:
+    CHECK = "check"
+    ARGUMENTS = "arguments"
+    SPECIFICATION = "specification"
+    COL_NAME = "column"
+    COL_NAMES = "for_each_column"
+    COLUMNS = "columns"
+    FUNCTION = "function"
+
+
+class DqxExporter(Exporter):
+    """Exporter implementation for converting data contracts to DQX YAML file."""
+
+    def export(
+        self,
+        data_contract: DataContractSpecification,
+        model: Model,
+        server: str,
+        sql_server_type: str,
+        export_args: Dict[str, Any],
+    ) -> str:
+        """Exports a data contract to DQX format."""
+        model_name, model_value = _check_models_for_export(data_contract, model, self.export_format)
+        return to_dqx_yaml(model_value)
+
+
+def to_dqx_yaml(model_value: Model) -> str:
+    """
+    Converts the data contract's quality checks to DQX YAML format.
+
+    Args:
+        model_value (Model): The data contract to convert.
+
+    Returns:
+        str: YAML representation of the data contract's quality checks.
+    """
+    extracted_rules = extract_quality_rules(model_value)
+    return yaml.dump(extracted_rules, sort_keys=False, allow_unicode=True, default_flow_style=False)
+
+
+def process_quality_rule(rule: Quality, column_name: str) -> Dict[str, Any]:
+    """
+    Processes a single quality rule by injecting the column path into its arguments if absent.
+
+    Args:
+        rule (Quality): The quality rule to process.
+        column_name (str): The full path to the current column.
+
+    Returns:
+        dict: The processed quality rule specification.
+    """
+    rule_data = rule.model_extra
+    specification = rule_data[DqxKeys.SPECIFICATION]
+    check = specification[DqxKeys.CHECK]
+
+    if column_name:
+        arguments = check.setdefault(DqxKeys.ARGUMENTS, {})
+
+        if (
+            DqxKeys.COL_NAME not in arguments
+            and DqxKeys.COL_NAMES not in arguments
+            and DqxKeys.COLUMNS not in arguments
+        ):
+            if check[DqxKeys.FUNCTION] not in ("is_unique", "foreign_key"):
+                arguments[DqxKeys.COL_NAME] = column_name
+            else:
+                arguments[DqxKeys.COLUMNS] = [column_name]
+
+    return specification
+
+
+def extract_quality_rules(data: Union[Model, Field, Quality], column_path: str = "") -> List[Dict[str, Any]]:
+    """
+    Recursively extracts all quality rules from a data contract structure.
+
+    Args:
+        data (Union[Model, Field, Quality]): The data contract model, field, or quality rule.
+        column_path (str, optional): The current path in the schema hierarchy. Defaults to "".
+
+    Returns:
+        List[Dict[str, Any]]: A list of quality rule specifications.
+    """
+    quality_rules = []
+
+    if isinstance(data, Quality):
+        return [process_quality_rule(data, column_path)]
+
+    if isinstance(data, (Model, Field)):
+        for key, field in data.fields.items():
+            current_path = build_column_path(column_path, key)
+
+            if field.fields:
+                # Field is a struct-like object, recurse deeper
+                quality_rules.extend(extract_quality_rules(field, current_path))
+            else:
+                # Process quality rules at leaf fields
+                for rule in field.quality:
+                    quality_rules.append(process_quality_rule(rule, current_path))
+
+        # Process any quality rules attached directly to this level
+        for rule in data.quality:
+            quality_rules.append(process_quality_rule(rule, column_path))
+
+    return quality_rules
+
+
+def build_column_path(current_path: str, key: str) -> str:
+    """
+    Builds the full column path by concatenating parent path with current key.
+
+    Args:
+        current_path (str): The current path prefix.
+        key (str): The current field's key.
+
+    Returns:
+        str: The full path.
+    """
+    return f"{current_path}.{key}" if current_path else key

datacontract/export/duckdb_type_converter.py

@@ -0,0 +1,57 @@
+from typing import Dict
+
+from datacontract.model.data_contract_specification import Field
+
+
+# https://duckdb.org/docs/data/csv/overview.html
+# ['SQLNULL', 'BOOLEAN', 'BIGINT', 'DOUBLE', 'TIME', 'DATE', 'TIMESTAMP', 'VARCHAR']
+def convert_to_duckdb_csv_type(field) -> None | str:
+    datacontract_type = field.type
+    if datacontract_type is None:
+        return "VARCHAR"
+    if datacontract_type.lower() in ["string", "varchar", "text"]:
+        return "VARCHAR"
+    if datacontract_type.lower() in ["timestamp", "timestamp_tz"]:
+        return "TIMESTAMP"
+    if datacontract_type.lower() in ["timestamp_ntz"]:
+        return "TIMESTAMP"
+    if datacontract_type.lower() in ["date"]:
+        return "DATE"
+    if datacontract_type.lower() in ["time"]:
+        return "TIME"
+    if datacontract_type.lower() in ["number", "decimal", "numeric"]:
+        # precision and scale not supported by data contract
+        return "VARCHAR"
+    if datacontract_type.lower() in ["float", "double"]:
+        return "DOUBLE"
+    if datacontract_type.lower() in ["integer", "int", "long", "bigint"]:
+        return "BIGINT"
+    if datacontract_type.lower() in ["boolean"]:
+        return "BOOLEAN"
+    if datacontract_type.lower() in ["object", "record", "struct"]:
+        # not supported in CSV
+        return "VARCHAR"
+    if datacontract_type.lower() in ["bytes"]:
+        # not supported in CSV
+        return "VARCHAR"
+    if datacontract_type.lower() in ["array"]:
+        return "VARCHAR"
+    if datacontract_type.lower() in ["null"]:
+        return "SQLNULL"
+    return "VARCHAR"
+
+
+def convert_to_duckdb_json_type(field: Field) -> None | str:
+    datacontract_type = field.type
+    if datacontract_type is None:
+        return "VARCHAR"
+    if datacontract_type.lower() in ["array"]:
+        return convert_to_duckdb_json_type(field.items) + "[]"  # type: ignore
+    if datacontract_type.lower() in ["object", "record", "struct"]:
+        return convert_to_duckdb_object(field.fields)
+    return convert_to_duckdb_csv_type(field)
+
+
+def convert_to_duckdb_object(fields: Dict[str, Field]):
+    columns = [f'"{x[0]}" {convert_to_duckdb_json_type(x[1])}' for x in fields.items()]
+    return f"STRUCT({', '.join(columns)})"