datacontract-cli 0.10.33__py3-none-any.whl → 0.10.35__py3-none-any.whl

datacontract/engines/soda/connections/athena.py

@@ -0,0 +1,79 @@
+import os
+
+import yaml
+
+from datacontract.model.exceptions import DataContractException
+
+
+def to_athena_soda_configuration(server):
+    s3_region = os.getenv("DATACONTRACT_S3_REGION")
+    s3_access_key_id = os.getenv("DATACONTRACT_S3_ACCESS_KEY_ID")
+    s3_secret_access_key = os.getenv("DATACONTRACT_S3_SECRET_ACCESS_KEY")
+    s3_session_token = os.getenv("DATACONTRACT_S3_SESSION_TOKEN")
+
+    # Validate required parameters
+    if not s3_access_key_id:
+        raise DataContractException(
+            type="athena-connection",
+            name="missing_access_key_id",
+            reason="AWS access key ID is required. Set the DATACONTRACT_S3_ACCESS_KEY_ID environment variable.",
+            engine="datacontract",
+        )
+
+    if not s3_secret_access_key:
+        raise DataContractException(
+            type="athena-connection",
+            name="missing_secret_access_key",
+            reason="AWS secret access key is required. Set the DATACONTRACT_S3_SECRET_ACCESS_KEY environment variable.",
+            engine="datacontract",
+        )
+
+    if not hasattr(server, "schema_") or not server.schema_:
+        raise DataContractException(
+            type="athena-connection",
+            name="missing_schema",
+            reason="Schema is required for Athena connection. Specify the schema where your tables exist in the server configuration.",
+            engine="datacontract",
+        )
+
+    if not hasattr(server, "stagingDir") or not server.stagingDir:
+        raise DataContractException(
+            type="athena-connection",
+            name="missing_s3_staging_dir",
+            reason="S3 staging directory is required for Athena connection. This should be the Amazon S3 Query Result Location (e.g., 's3://my-bucket/athena-results/').",
+            engine="datacontract",
+        )
+
+    # Validate S3 staging directory format
+    if not server.stagingDir.startswith("s3://"):
+        raise DataContractException(
+            type="athena-connection",
+            name="invalid_s3_staging_dir",
+            reason=f"S3 staging directory must start with 's3://'. Got: {server.s3_staging_dir}. Example: 's3://my-bucket/athena-results/'",
+            engine="datacontract",
+        )
+
+    data_source = {
+        "type": "athena",
+        "access_key_id": s3_access_key_id,
+        "secret_access_key": s3_secret_access_key,
+        "schema": server.schema_,
+        "staging_dir": server.stagingDir,
+    }
+
+    if s3_region:
+        data_source["region_name"] = s3_region
+    elif server.region_name:
+        data_source["region_name"] = server.region_name
+
+    if server.catalog:
+        # Optional, Identify the name of the Data Source, also referred to as a Catalog. The default value is `awsdatacatalog`.
+        data_source["catalog"] = server.catalog
+
+    if s3_session_token:
+        data_source["aws_session_token"] = s3_session_token
+
+    soda_configuration = {f"data_source {server.type}": data_source}
+
+    soda_configuration_str = yaml.dump(soda_configuration)
+    return soda_configuration_str

datacontract/engines/soda/connections/duckdb_connection.py

@@ -71,6 +71,9 @@ def get_duckdb_connection(
         elif server.format == "delta":
             con.sql("update extensions;")  # Make sure we have the latest delta extension
             con.sql(f"""CREATE VIEW "{model_name}" AS SELECT * FROM delta_scan('{model_path}');""")
+        table_info = con.sql(f"PRAGMA table_info('{model_name}');").fetchdf()
+        if table_info is not None and not table_info.empty:
+            run.log_info(f"DuckDB Table Info: {table_info.to_string(index=False)}")
     return con
 
 

datacontract/export/avro_converter.py

@@ -44,12 +44,18 @@ def to_avro_field(field, field_name):
     avro_type = to_avro_type(field, field_name)
     avro_field["type"] = avro_type if is_required_avro else ["null", avro_type]
 
-    if avro_field["type"] == "enum":
-        avro_field["type"] = {
+    # Handle enum types - both required and optional
+    if avro_type == "enum" or (isinstance(avro_field["type"], list) and "enum" in avro_field["type"]):
+        enum_def = {
             "type": "enum",
             "name": field.title,
             "symbols": field.enum,
         }
+        if is_required_avro:
+            avro_field["type"] = enum_def
+        else:
+            # Replace "enum" with the full enum definition in the union
+            avro_field["type"] = ["null", enum_def]
 
     if field.config:
         if "avroDefault" in field.config:
@@ -77,6 +83,10 @@ def to_avro_type(field: Field, field_name: str) -> str | dict:
         if "avroType" in field.config:
             return field.config["avroType"]
 
+    # Check for enum fields based on presence of enum list and avroType config
+    if field.enum and field.config and field.config.get("avroType") == "enum":
+        return "enum"
+
     if field.type is None:
         return "null"
     if field.type in ["string", "varchar", "text"]:

datacontract/export/dqx_converter.py

@@ -0,0 +1,121 @@
+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]
+
+    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/exporter.py

@@ -46,6 +46,7 @@ class ExportFormat(str, Enum):
     iceberg = "iceberg"
     custom = "custom"
     excel = "excel"
+    dqx = "dqx"
 
     @classmethod
     def get_supported_formats(cls):

datacontract/export/exporter_factory.py

@@ -197,6 +197,12 @@ exporter_factory.register_lazy_exporter(
     class_name="MarkdownExporter",
 )
 
+exporter_factory.register_lazy_exporter(
+    name=ExportFormat.dqx,
+    module_path="datacontract.export.dqx_converter",
+    class_name="DqxExporter",
+)
+
 exporter_factory.register_lazy_exporter(
     name=ExportFormat.iceberg, module_path="datacontract.export.iceberg_converter", class_name="IcebergExporter"
 )

datacontract/export/markdown_converter.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Dict, List
 
 from pydantic import BaseModel
 
@@ -12,6 +12,9 @@ from datacontract.model.data_contract_specification import (
     ServiceLevel,
 )
 
+TAB = " "
+ARROW = "↳"
+
 
 class MarkdownExporter(Exporter):
     """Exporter implementation for converting data contracts to Markdown."""
@@ -70,7 +73,8 @@ def obj_attributes_to_markdown(obj: BaseModel, excluded_fields: set = set(), is_
     else:
         bullet_char = "-"
         newline_char = "\n"
-    obj_model = obj.model_dump(exclude_unset=True, exclude=excluded_fields)
+    model_attributes_to_include = set(obj.__class__.model_fields.keys())
+    obj_model = obj.model_dump(exclude_unset=True, include=model_attributes_to_include, exclude=excluded_fields)
     description_value = obj_model.pop("description", None)
     attributes = [
         (f"{bullet_char} `{attr}`" if value is True else f"{bullet_char} **{attr}:** {value}")
@@ -78,7 +82,8 @@ def obj_attributes_to_markdown(obj: BaseModel, excluded_fields: set = set(), is_
         if value
     ]
     description = f"*{description_to_markdown(description_value)}*"
-    return newline_char.join([description] + attributes)
+    extra = [extra_to_markdown(obj)] if obj.model_extra else []
+    return newline_char.join([description] + attributes + extra)
 
 
 def servers_to_markdown(servers: Dict[str, Server]) -> str:
@@ -153,8 +158,8 @@ def field_to_markdown(field_name: str, field: Field, level: int = 0) -> str:
     Returns:
         str: A Markdown table rows for the field.
     """
-    tabs = " " * level
-    arrow = "↳" if level > 0 else ""
+    tabs = TAB * level
+    arrow = ARROW if level > 0 else ""
     column_name = f"{tabs}{arrow} {field_name}"
 
     attributes = obj_attributes_to_markdown(field, {"type", "fields", "items", "keys", "values"}, True)
@@ -206,3 +211,108 @@ def service_level_to_markdown(service_level: ServiceLevel | None) -> str:
 
 def description_to_markdown(description: str | None) -> str:
     return (description or "No description.").replace("\n", "<br>")
+
+
+def array_of_dict_to_markdown(array: List[Dict[str, str]]) -> str:
+    """
+    Convert a list of dictionaries to a Markdown table.
+
+    Args:
+        array (List[Dict[str, str]]): A list of dictionaries where each dictionary represents a row in the table.
+
+    Returns:
+        str: A Markdown formatted table.
+    """
+    if not array:
+        return ""
+
+    headers = []
+
+    for item in array:
+        headers += item.keys()
+    headers = list(dict.fromkeys(headers))  # Preserve order and remove duplicates
+
+    markdown_parts = [
+        "| " + " | ".join(headers) + " |",
+        "| " + " | ".join(["---"] * len(headers)) + " |",
+    ]
+
+    for row in array:
+        element = row
+        markdown_parts.append(
+            "| "
+            + " | ".join(
+                f"{str(element.get(header, ''))}".replace("\n", "<br>").replace("\t", TAB) for header in headers
+            )
+            + " |"
+        )
+
+    return "\n".join(markdown_parts) + "\n"
+
+
+def array_to_markdown(array: List[str]) -> str:
+    """
+    Convert a list of strings to a Markdown formatted list.
+
+    Args:
+        array (List[str]): A list of strings to convert.
+
+    Returns:
+        str: A Markdown formatted list.
+    """
+    if not array:
+        return ""
+    return "\n".join(f"- {item}" for item in array) + "\n"
+
+
+def dict_to_markdown(dictionary: Dict[str, str]) -> str:
+    """
+    Convert a dictionary to a Markdown formatted list.
+
+    Args:
+        dictionary (Dict[str, str]): A dictionary where keys are item names and values are item descriptions.
+
+    Returns:
+        str: A Markdown formatted list of items.
+    """
+    if not dictionary:
+        return ""
+
+    markdown_parts = []
+    for key, value in dictionary.items():
+        if isinstance(value, dict):
+            markdown_parts.append(f"- {key}")
+            nested_markdown = dict_to_markdown(value)
+            if nested_markdown:
+                nested_lines = nested_markdown.split("\n")
+                for line in nested_lines:
+                    if line.strip():
+                        markdown_parts.append(f"  {line}")
+        else:
+            markdown_parts.append(f"- {key}: {value}")
+    return "\n".join(markdown_parts) + "\n"
+
+
+def extra_to_markdown(obj: BaseModel) -> str:
+    """
+    Convert the extra attributes of a data contract to Markdown format.
+    Args:
+        obj (BaseModel): The data contract object containing extra attributes.
+    Returns:
+        str: A Markdown formatted string representing the extra attributes of the data contract.
+    """
+    markdown_part = ""
+    extra = obj.model_extra
+    if extra:
+        for key_extra, value_extra in extra.items():
+            markdown_part += f"\n### {key_extra.capitalize()}\n"
+            if isinstance(value_extra, list) and len(value_extra):
+                if isinstance(value_extra[0], dict):
+                    markdown_part += array_of_dict_to_markdown(value_extra)
+                elif isinstance(value_extra[0], str):
+                    markdown_part += array_to_markdown(value_extra)
+            elif isinstance(value_extra, dict):
+                markdown_part += dict_to_markdown(value_extra)
+            else:
+                markdown_part += f"{str(value_extra)}\n"
+    return markdown_part

datacontract/export/mermaid_exporter.py

@@ -27,31 +27,33 @@ def dcs_to_mermaid(data_contract_spec: DataContractSpecification) -> str | None:
         mmd_references = []
 
         for model_name, model in data_contract_spec.models.items():
+            clean_model = _sanitize_name(model_name)
             entity_block = ""
 
             for field_name, field in model.fields.items():
                 clean_name = _sanitize_name(field_name)
-                indicators = ""
+                field_type = field.type or "unknown"
 
-                if field.primaryKey or (field.unique and field.required):
-                    indicators += "🔑"
-                if field.references:
-                    indicators += "⌘"
+                is_pk = bool(field.primaryKey or (field.unique and field.required))
+                is_fk = bool(field.references)
 
-                field_type = field.type or "unknown"
-                entity_block += f"\t{clean_name}{indicators} {field_type}\n"
+                entity_block += _field_line(clean_name, field_type, pk=is_pk, uk=bool(field.unique), fk=is_fk)
 
                 if field.references:
-                    referenced_model = field.references.split(".")[0] if "." in field.references else ""
+                    references = field.references.replace(".", "·")
+                    parts = references.split("·")
+                    referenced_model = _sanitize_name(parts[0]) if len(parts) > 0 else ""
+                    referenced_field = _sanitize_name(parts[1]) if len(parts) > 1 else ""
                     if referenced_model:
-                        mmd_references.append(f'"📑{referenced_model}"' + "}o--{ ||" + f'"📑{model_name}"')
+                        label = referenced_field or clean_name
+                        mmd_references.append(f'"**{referenced_model}**" ||--o{{ "**{clean_model}**" : {label}')
 
-            mmd_entity += f'\t"**{model_name}**"' + "{\n" + entity_block + "}\n"
+            mmd_entity += f'\t"**{clean_model}**" {{\n{entity_block}}}\n'
 
         if mmd_references:
             mmd_entity += "\n" + "\n".join(mmd_references)
 
-        return f"{mmd_entity}\n"
+        return mmd_entity + "\n"
 
     except Exception as e:
         print(f"Error generating DCS mermaid diagram: {e}")
@@ -95,3 +97,14 @@ def odcs_to_mermaid(data_contract_spec: OpenDataContractStandard) -> str | None:
 
 def _sanitize_name(name: str) -> str:
     return name.replace("#", "Nb").replace(" ", "_").replace("/", "by")
+
+
+def _field_line(name: str, field_type: str, pk: bool = False, uk: bool = False, fk: bool = False) -> str:
+    indicators = ""
+    if pk:
+        indicators += "🔑"
+    if uk:
+        indicators += "🔒"
+    if fk:
+        indicators += "⌘"
+    return f"\t{name}{indicators} {field_type}\n"

datacontract/export/spark_converter.py

@@ -1,3 +1,5 @@
+import json
+
 from pyspark.sql import types
 
 from datacontract.export.exporter import Exporter
@@ -104,7 +106,8 @@ def to_struct_field(field: Field, field_name: str) -> types.StructField:
         types.StructField: The corresponding Spark StructField.
     """
     data_type = to_spark_data_type(field)
-    return types.StructField(name=field_name, dataType=data_type, nullable=not field.required)
+    metadata = to_spark_metadata(field)
+    return types.StructField(name=field_name, dataType=data_type, nullable=not field.required, metadata=metadata)
 
 
 def to_spark_data_type(field: Field) -> types.DataType:
@@ -152,7 +155,25 @@ def to_spark_data_type(field: Field) -> types.DataType:
         return types.DateType()
     if field_type == "bytes":
         return types.BinaryType()
-    return types.StringType() # default if no condition is met
+    return types.StringType()  # default if no condition is met
+
+
+def to_spark_metadata(field: Field) -> dict[str, str]:
+    """
+    Convert a field to a Spark metadata dictonary.
+
+    Args:
+        field (Field): The field to convert.
+
+    Returns:
+        dict: dictionary that can be supplied to Spark as metadata for a StructField
+    """
+
+    metadata = {}
+    if field.description:
+        metadata["comment"] = field.description
+
+    return metadata
 
 
 def print_schema(dtype: types.DataType) -> str:
@@ -192,7 +213,11 @@ def print_schema(dtype: types.DataType) -> str:
         name = f'"{column.name}"'
         data_type = indent(print_schema(column.dataType), 1)
         nullable = indent(f"{column.nullable}", 1)
-        return f"StructField({name},\n{data_type},\n{nullable}\n)"
+        if column.metadata:
+            metadata = indent(f"{json.dumps(column.metadata)}", 1)
+            return f"StructField({name},\n{data_type},\n{nullable},\n{metadata}\n)"
+        else:
+            return f"StructField({name},\n{data_type},\n{nullable}\n)"
 
     def format_struct_type(struct_type: types.StructType) -> str:
         """

datacontract/export/sql_type_converter.py

@@ -3,6 +3,9 @@ from datacontract.model.data_contract_specification import Field
 
 
 def convert_to_sql_type(field: Field, server_type: str) -> str:
+    if field.config and "physicalType" in field.config:
+        return field.config["physicalType"]
+
     if server_type == "snowflake":
         return convert_to_snowflake(field)
     elif server_type == "postgres":
@@ -19,6 +22,7 @@ def convert_to_sql_type(field: Field, server_type: str) -> str:
         return convert_type_to_bigquery(field)
     elif server_type == "trino":
         return convert_type_to_trino(field)
+
     return field.type
 
 

datacontract/imports/avro_importer.py

@@ -130,13 +130,23 @@ def import_record_fields(record_fields: List[avro.schema.Field]) -> Dict[str, Fi
             imported_field.fields = import_record_fields(field.type.fields)
         elif field.type.type == "union":
             imported_field.required = False
-            type = import_type_of_optional_field(field)
-            imported_field.type = type
-            if type == "record":
-                imported_field.fields = import_record_fields(get_record_from_union_field(field).fields)
-            elif type == "array":
-                imported_field.type = "array"
-                imported_field.items = import_avro_array_items(get_array_from_union_field(field))
+            # Check for enum in union first, since it needs special handling
+            enum_schema = get_enum_from_union_field(field)
+            if enum_schema:
+                imported_field.type = "string"
+                imported_field.enum = enum_schema.symbols
+                imported_field.title = enum_schema.name
+                if not imported_field.config:
+                    imported_field.config = {}
+                imported_field.config["avroType"] = "enum"
+            else:
+                type = import_type_of_optional_field(field)
+                imported_field.type = type
+                if type == "record":
+                    imported_field.fields = import_record_fields(get_record_from_union_field(field).fields)
+                elif type == "array":
+                    imported_field.type = "array"
+                    imported_field.items = import_avro_array_items(get_array_from_union_field(field))
         elif field.type.type == "array":
             imported_field.type = "array"
             imported_field.items = import_avro_array_items(field.type)
@@ -277,6 +287,22 @@ def get_array_from_union_field(field: avro.schema.Field) -> avro.schema.ArraySch
     return None
 
 
+def get_enum_from_union_field(field: avro.schema.Field) -> avro.schema.EnumSchema | None:
+    """
+    Get the enum schema from a union field.
+
+    Args:
+        field: The Avro field with a union type.
+
+    Returns:
+        The enum schema if found, None otherwise.
+    """
+    for field_type in field.type.schemas:
+        if field_type.type == "enum":
+            return field_type
+    return None
+
+
 def map_type_from_avro(avro_type_str: str) -> str:
     """
     Map Avro type strings to data contract type strings.

datacontract/imports/odcs_v3_importer.py

@@ -131,14 +131,18 @@ def import_servers(odcs: OpenDataContractStandard) -> Dict[str, Server] | None:
         server.host = odcs_server.host
         server.port = odcs_server.port
         server.catalog = odcs_server.catalog
+        server.stagingDir = odcs_server.stagingDir
         server.topic = getattr(odcs_server, "topic", None)
         server.http_path = getattr(odcs_server, "http_path", None)
         server.token = getattr(odcs_server, "token", None)
         server.driver = getattr(odcs_server, "driver", None)
         server.roles = import_server_roles(odcs_server.roles)
         server.storageAccount = (
-            re.search(r"(?:@|://)([^.]+)\.", odcs_server.location, re.IGNORECASE) if server.type == "azure" else None
+            to_azure_storage_account(odcs_server.location)
+            if server.type == "azure" and "://" in server.location
+            else None
         )
+
         servers[server_name] = server
     return servers
 
@@ -413,3 +417,28 @@ def import_tags(odcs: OpenDataContractStandard) -> List[str] | None:
     if odcs.tags is None:
         return None
     return odcs.tags
+
+
+def to_azure_storage_account(location: str) -> str | None:
+    """
+    Converts a storage location string to extract the storage account name.
+    ODCS v3.0 has no explicit field for the storage account. It uses the location field, which is a URI.
+
+    This function parses a storage location string to identify and return the
+    storage account name. It handles two primary patterns:
+    1. Protocol://containerName@storageAccountName
+    2. Protocol://storageAccountName
+
+    :param location: The storage location string to parse, typically following
+                     the format protocol://containerName@storageAccountName. or
+                     protocol://storageAccountName.
+    :return: The extracted storage account name if found, otherwise None
+    """
+    # to catch protocol://containerName@storageAccountName. pattern from location
+    match = re.search(r"(?<=@)([^.]*)", location, re.IGNORECASE)
+    if match:
+        return match.group()
+    else:
+        # to catch protocol://storageAccountName. pattern from location
+        match = re.search(r"(?<=//)(?!@)([^.]*)", location, re.IGNORECASE)
+    return match.group() if match else None