datacontract-cli 0.10.23__py3-none-any.whl → 0.10.40__py3-none-any.whl

datacontract/export/dbt_converter.py

@@ -27,7 +27,7 @@ class DbtStageExporter(Exporter):
         )
 
 
-def to_dbt_models_yaml(data_contract_spec: DataContractSpecification, server: str = None):
+def to_dbt_models_yaml(data_contract_spec: DataContractSpecification, server: str = None) -> str:
     dbt = {
         "version": 2,
         "models": [],
@@ -102,8 +102,11 @@ def _to_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
@@ -112,9 +115,28 @@ def _to_dbt_model(
         dbt_model["config"]["contract"] = {"enforced": True}
     if model_value.description is not None:
         dbt_model["description"] = model_value.description.strip().replace("\n", " ")
-    columns = _to_columns(data_contract_spec, model_value.fields, _supports_constraints(model_type), adapter_type)
+
+    # 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
 
 
@@ -123,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":
@@ -140,10 +162,13 @@ def _to_columns(
     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(data_contract_spec, field_name, field, supports_constraints, adapter_type)
+        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
 
@@ -161,6 +186,7 @@ def _to_column(
     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"
@@ -175,12 +201,15 @@ def _to_column(
         )
     if field.description is not None:
         column["description"] = field.description.strip().replace("\n", " ")
-    if field.required:
+    # Handle required/not_null constraint
+    if field.required or is_primary_key:
         if supports_constraints:
             column.setdefault("constraints", []).append({"type": "not_null"})
         else:
             column["data_tests"].append("not_null")
-    if field.unique:
+
+    # Handle unique constraint
+    if field.unique or is_primary_key:
         if supports_constraints:
             column.setdefault("constraints", []).append({"type": "unique"})
         else:

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)})"