datacontract-cli 0.10.14__py3-none-any.whl → 0.10.16__py3-none-any.whl

datacontract/export/great_expectations_converter.py

@@ -1,49 +1,118 @@
+"""
+This module provides functionalities to export data contracts to Great Expectations suites.
+It includes definitions for exporting different types of data (pandas, Spark, SQL) into
+Great Expectations expectations format.
+"""
+
 import json
-from typing import Dict, List, Any
+from enum import Enum
+from typing import Any, Dict, List
 
 import yaml
 
-from datacontract.model.data_contract_specification import DataContractSpecification, Field, Quality
-from datacontract.export.exporter import Exporter, _check_models_for_export
+from datacontract.export.exporter import (
+    Exporter,
+    _check_models_for_export,
+)
+from datacontract.export.pandas_type_converter import convert_to_pandas_type
+from datacontract.export.spark_converter import to_spark_data_type
+from datacontract.export.sql_type_converter import convert_to_sql_type
+from datacontract.model.data_contract_specification import (
+    DataContractSpecification,
+    Field,
+    Quality,
+)
+
+
+class GreatExpectationsEngine(Enum):
+    """Enum to represent the type of data engine for expectations.
+
+    Attributes:
+        pandas (str): Represents the Pandas engine type.
+        spark (str): Represents the Spark engine type.
+        sql (str): Represents the SQL engine type.
+    """
+
+    pandas = "pandas"
+    spark = "spark"
+    sql = "sql"
+
 
+class GreatExpectationsExporter(Exporter):
+    """Exporter class to convert data contracts to Great Expectations suites.
+
+    Methods:
+        export: Converts a data contract model to a Great Expectations suite.
+
+    """
 
-class GreateExpectationsExporter(Exporter):
     def export(self, data_contract, model, server, sql_server_type, export_args) -> dict:
+        """Exports a data contract model to a Great Expectations suite.
+
+        Args:
+            data_contract (DataContractSpecification): The data contract specification.
+            model (str): The model name to export.
+            server (str): The server information.
+            sql_server_type (str): Type of SQL server (e.g., "snowflake").
+            export_args (dict): Additional arguments for export, such as "suite_name" and "engine".
+
+        Returns:
+            dict: A dictionary representation of the Great Expectations suite.
+        """
+        expectation_suite_name = export_args.get("suite_name")
+        engine = export_args.get("engine")
         model_name, model_value = _check_models_for_export(data_contract, model, self.export_format)
-        return to_great_expectations(
-            data_contract,
-            model_name,
-        )
+        sql_server_type = "snowflake" if sql_server_type == "auto" else sql_server_type
+        return to_great_expectations(data_contract, model_name, expectation_suite_name, engine, sql_server_type)
 
 
-def to_great_expectations(data_contract_spec: DataContractSpecification, model_key: str) -> str:
-    """
-    Convert each model in the contract to a Great Expectation suite
-    @param data_contract_spec: data contract to export to great expectations
-    @param model_key: model to great expectations to
-    @return: a dictionary of great expectation suites
+def to_great_expectations(
+    data_contract_spec: DataContractSpecification,
+    model_key: str,
+    expectation_suite_name: str | None = None,
+    engine: str | None = None,
+    sql_server_type: str = "snowflake",
+) -> str:
+    """Converts a data contract model to a Great Expectations suite.
+
+    Args:
+        data_contract_spec (DataContractSpecification): The data contract specification.
+        model_key (str): The model key.
+        expectation_suite_name (str | None): Optional suite name for the expectations.
+        engine (str | None): Optional engine type (e.g., "pandas", "spark").
+        sql_server_type (str): The type of SQL server (default is "snowflake").
+
+    Returns:
+        str: JSON string of the Great Expectations suite.
     """
     expectations = []
+    if not expectation_suite_name:
+        expectation_suite_name = "{model_key}.{contract_version}".format(
+            model_key=model_key, contract_version=data_contract_spec.info.version
+        )
     model_value = data_contract_spec.models.get(model_key)
     quality_checks = get_quality_checks(data_contract_spec.quality)
-    expectations.extend(model_to_expectations(model_value.fields))
+    expectations.extend(model_to_expectations(model_value.fields, engine, sql_server_type))
     expectations.extend(checks_to_expectations(quality_checks, model_key))
-    model_expectation_suite = to_suite(model_key, data_contract_spec.info.version, expectations)
+    model_expectation_suite = to_suite(expectations, expectation_suite_name)
 
     return model_expectation_suite
 
 
-def to_suite(
-    model_key: str,
-    contract_version: str,
-    expectations: List[Dict[str, Any]],
-) -> str:
+def to_suite(expectations: List[Dict[str, Any]], expectation_suite_name: str) -> str:
+    """Converts a list of expectations to a JSON-formatted suite.
+
+    Args:
+        expectations (List[Dict[str, Any]]): List of expectations.
+        expectation_suite_name (str): Name of the expectation suite.
+
+    Returns:
+        str: JSON string of the expectation suite.
+    """
     return json.dumps(
         {
             "data_asset_type": "null",
-            "expectation_suite_name": "user-defined.{model_key}.{contract_version}".format(
-                model_key=model_key, contract_version=contract_version
-            ),
+            "expectation_suite_name": expectation_suite_name,
             "expectations": expectations,
             "meta": {},
         },
@@ -51,22 +120,53 @@ def to_suite(
     )
 
 
-def model_to_expectations(fields: Dict[str, Field]) -> List[Dict[str, Any]]:
-    """
-    Convert the model information to expectations
-    @param fields: model field
-    @return: list of expectations
+def model_to_expectations(fields: Dict[str, Field], engine: str | None, sql_server_type: str) -> List[Dict[str, Any]]:
+    """Converts model fields to a list of expectations.
+
+    Args:
+        fields (Dict[str, Field]): Dictionary of model fields.
+        engine (str | None): Engine type (e.g., "pandas", "spark").
+        sql_server_type (str): SQL server type.
+
+    Returns:
+        List[Dict[str, Any]]: List of expectations.
     """
     expectations = []
     add_column_order_exp(fields, expectations)
     for field_name, field in fields.items():
-        add_field_expectations(field_name, field, expectations)
+        add_field_expectations(field_name, field, expectations, engine, sql_server_type)
     return expectations
 
 
-def add_field_expectations(field_name, field: Field, expectations: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+def add_field_expectations(
+    field_name,
+    field: Field,
+    expectations: List[Dict[str, Any]],
+    engine: str | None,
+    sql_server_type: str,
+) -> List[Dict[str, Any]]:
+    """Adds expectations for a specific field based on its properties.
+
+    Args:
+        field_name (str): The name of the field.
+        field (Field): The field object.
+        expectations (List[Dict[str, Any]]): The expectations list to update.
+        engine (str | None): Engine type (e.g., "pandas", "spark").
+        sql_server_type (str): SQL server type.
+
+    Returns:
+        List[Dict[str, Any]]: Updated list of expectations.
+    """
     if field.type is not None:
-        expectations.append(to_column_types_exp(field_name, field.type))
+        if engine == GreatExpectationsEngine.spark.value:
+            field_type = to_spark_data_type(field).__class__.__name__
+        elif engine == GreatExpectationsEngine.pandas.value:
+            field_type = convert_to_pandas_type(field)
+        elif engine == GreatExpectationsEngine.sql.value:
+            field_type = convert_to_sql_type(field, sql_server_type)
+        else:
+            field_type = field.type
+        expectations.append(to_column_types_exp(field_name, field_type))
     if field.unique:
         expectations.append(to_column_unique_exp(field_name))
     if field.maxLength is not None or field.minLength is not None:
@@ -74,11 +174,16 @@ def add_field_expectations(field_name, field: Field, expectations: List[Dict[str
     if field.minimum is not None or field.maximum is not None:
         expectations.append(to_column_min_max_exp(field_name, field.minimum, field.maximum))
 
-    # TODO: all constraints
     return expectations
 
 
 def add_column_order_exp(fields: Dict[str, Field], expectations: List[Dict[str, Any]]):
+    """Adds expectation for column ordering.
+
+    Args:
+        fields (Dict[str, Field]): Dictionary of fields.
+        expectations (List[Dict[str, Any]]): The expectations list to update.
+    """
     expectations.append(
         {
             "expectation_type": "expect_table_columns_to_match_ordered_list",
@@ -89,6 +194,15 @@ def add_column_order_exp(fields: Dict[str, Field], expectations: List[Dict[str,
 
 
 def to_column_types_exp(field_name, field_type) -> Dict[str, Any]:
+    """Creates a column type expectation.
+
+    Args:
+        field_name (str): The name of the field.
+        field_type (str): The type of the field.
+
+    Returns:
+        Dict[str, Any]: Column type expectation.
+    """
     return {
         "expectation_type": "expect_column_values_to_be_of_type",
         "kwargs": {"column": field_name, "type_": field_type},
@@ -97,18 +211,54 @@ def to_column_types_exp(field_name, field_type) -> Dict[str, Any]:
 
 
 def to_column_unique_exp(field_name) -> Dict[str, Any]:
-    return {"expectation_type": "expect_column_values_to_be_unique", "kwargs": {"column": field_name}, "meta": {}}
+    """Creates a column uniqueness expectation.
+
+    Args:
+        field_name (str): The name of the field.
+
+    Returns:
+        Dict[str, Any]: Column uniqueness expectation.
+    """
+    return {
+        "expectation_type": "expect_column_values_to_be_unique",
+        "kwargs": {"column": field_name},
+        "meta": {},
+    }
 
 
 def to_column_length_exp(field_name, min_length, max_length) -> Dict[str, Any]:
+    """Creates a column length expectation.
+
+    Args:
+        field_name (str): The name of the field.
+        min_length (int | None): Minimum length.
+        max_length (int | None): Maximum length.
+
+    Returns:
+        Dict[str, Any]: Column length expectation.
+    """
     return {
         "expectation_type": "expect_column_value_lengths_to_be_between",
-        "kwargs": {"column": field_name, "min_value": min_length, "max_value": max_length},
+        "kwargs": {
+            "column": field_name,
+            "min_value": min_length,
+            "max_value": max_length,
+        },
         "meta": {},
     }
 
 
 def to_column_min_max_exp(field_name, minimum, maximum) -> Dict[str, Any]:
+    """Creates a column min-max value expectation.
+
+    Args:
+        field_name (str): The name of the field.
+        minimum (float | None): Minimum value.
+        maximum (float | None): Maximum value.
+
+    Returns:
+        Dict[str, Any]: Column min-max value expectation.
+    """
     return {
         "expectation_type": "expect_column_values_to_be_between",
         "kwargs": {"column": field_name, "min_value": minimum, "max_value": maximum},
@@ -117,6 +267,14 @@ def to_column_min_max_exp(field_name, minimum, maximum) -> Dict[str, Any]:
 
 
 def get_quality_checks(quality: Quality) -> Dict[str, Any]:
+    """Retrieves quality checks defined in a data contract.
+
+    Args:
+        quality (Quality): Quality object from the data contract.
+
+    Returns:
+        Dict[str, Any]: Dictionary of quality checks.
+    """
     if quality is None:
         return {}
     if quality.type is None:
@@ -131,11 +289,14 @@ def get_quality_checks(quality: Quality) -> Dict[str, Any]:
 
 
 def checks_to_expectations(quality_checks: Dict[str, Any], model_key: str) -> List[Dict[str, Any]]:
-    """
-    Get the quality definition for each model to the model expectation list
-    @param quality_checks: dictionary of quality checks by model
-    @param model_key: id of the model
-    @return: the list of expectations for that model
+    """Converts quality checks to a list of expectations.
+
+    Args:
+        quality_checks (Dict[str, Any]): Dictionary of quality checks by model.
+        model_key (str): The model key.
+
+    Returns:
+        List[Dict[str, Any]]: List of expectations for the model.
     """
     if quality_checks is None or model_key not in quality_checks:
         return []
@@ -148,3 +309,4 @@ def checks_to_expectations(quality_checks: Dict[str, Any], model_key: str) -> Li
     if isinstance(model_quality_checks, str):
         expectation_list = json.loads(model_quality_checks)
         return expectation_list
+    return []

datacontract/export/html_export.py

@@ -7,8 +7,8 @@ import pytz
 import yaml
 from jinja2 import Environment, PackageLoader, select_autoescape
 
-from datacontract.model.data_contract_specification import DataContractSpecification
 from datacontract.export.exporter import Exporter
+from datacontract.model.data_contract_specification import DataContractSpecification
 
 
 class HtmlExporter(Exporter):

datacontract/export/iceberg_converter.py

@@ -0,0 +1,188 @@
+from pyiceberg import types
+from pyiceberg.schema import Schema, assign_fresh_schema_ids
+
+from datacontract.export.exporter import Exporter
+from datacontract.model.data_contract_specification import (
+    DataContractSpecification,
+    Field,
+    Model,
+)
+
+
+class IcebergExporter(Exporter):
+    """
+    Exporter class for exporting data contracts to Iceberg schemas.
+    """
+
+    def export(
+        self,
+        data_contract: DataContractSpecification,
+        model,
+        server,
+        sql_server_type,
+        export_args,
+    ):
+        """
+        Export the given data contract model to an Iceberg schema.
+
+        Args:
+            data_contract (DataContractSpecification): The data contract specification.
+            model: The model to export, currently just supports one model.
+            server: Not used in this implementation.
+            sql_server_type: Not used in this implementation.
+            export_args: Additional arguments for export.
+
+        Returns:
+            str: A string representation of the Iceberg json schema.
+        """
+
+        return to_iceberg(data_contract, model)
+
+
+def to_iceberg(contract: DataContractSpecification, model: str) -> str:
+    """
+    Converts a DataContractSpecification into an Iceberg json schema string. JSON string follows https://iceberg.apache.org/spec/#appendix-c-json-serialization.
+
+    Args:
+        contract (DataContractSpecification): The data contract specification containing models.
+        model: The model to export, currently just supports one model.
+
+    Returns:
+        str: A string representation of the Iceberg json schema.
+    """
+    if model is None or model == "all":
+        if len(contract.models.items()) != 1:
+            # Iceberg doesn't have a way to combine multiple models into a single schema, an alternative would be to export json lines
+            raise Exception(f"Can only output one model at a time, found {len(contract.models.items())} models")
+        for model_name, model in contract.models.items():
+            schema = to_iceberg_schema(model)
+    else:
+        if model not in contract.models:
+            raise Exception(f"model {model} not found in contract")
+        schema = to_iceberg_schema(contract.models[model])
+
+    return schema.model_dump_json()
+
+
+def to_iceberg_schema(model: Model) -> types.StructType:
+    """
+    Convert a model to a Iceberg schema.
+
+    Args:
+        model (Model): The model to convert.
+
+    Returns:
+        types.StructType: The corresponding Iceberg schema.
+    """
+    iceberg_fields = []
+    primary_keys = []
+    for field_name, spec_field in model.fields.items():
+        iceberg_field = make_field(field_name, spec_field)
+        iceberg_fields.append(iceberg_field)
+
+        if spec_field.primaryKey:
+            primary_keys.append(iceberg_field.name)
+
+    schema = Schema(*iceberg_fields)
+
+    # apply non-0 field IDs so we can set the identifier fields for the schema
+    schema = assign_fresh_schema_ids(schema)
+    for field in schema.fields:
+        if field.name in primary_keys:
+            schema.identifier_field_ids.append(field.field_id)
+
+    return schema
+
+
+def make_field(field_name, field):
+    field_type = get_field_type(field)
+
+    # Note: might want to re-populate field_id from config['icebergFieldId'] if it exists, however, it gets
+    # complicated since field_ids impact the list and map element_ids, and the importer is not keeping track of those.
+    # Even if IDs are re-constituted, it seems like the SDK code would still reset them before any operation against a catalog,
+    # so it's likely not worth it.
+
+    # Note 2: field_id defaults to 0 to signify that the exporter is not attempting to populate meaningful values.
+    # also, the Iceberg sdk catalog code will re-set the fieldIDs prior to executing any table operations on the schema
+    # ref: https://github.com/apache/iceberg-python/pull/1072
+    return types.NestedField(field_id=0, name=field_name, field_type=field_type, required=field.required)
+
+
+def make_list(item):
+    field_type = get_field_type(item)
+
+    # element_id defaults to 0 to signify that the exporter is not attempting to populate meaningful values (see #make_field)
+    return types.ListType(element_id=0, element_type=field_type, element_required=item.required)
+
+
+def make_map(field):
+    key_type = get_field_type(field.keys)
+    value_type = get_field_type(field.values)
+
+    # key_id and value_id defaults to 0 to signify that the exporter is not attempting to populate meaningful values (see #make_field)
+    return types.MapType(
+        key_id=0, key_type=key_type, value_id=0, value_type=value_type, value_required=field.values.required
+    )
+
+
+def to_struct_type(fields: dict[str, Field]) -> types.StructType:
+    """
+    Convert a dictionary of fields to a Iceberg StructType.
+
+    Args:
+        fields (dict[str, Field]): The fields to convert.
+
+    Returns:
+        types.StructType: The corresponding Iceberg StructType.
+    """
+    struct_fields = []
+    for field_name, field in fields.items():
+        struct_field = make_field(field_name, field)
+        struct_fields.append(struct_field)
+    return types.StructType(*struct_fields)
+
+
+def get_field_type(field: Field) -> types.IcebergType:
+    """
+    Convert a field to a Iceberg IcebergType.
+
+    Args:
+        field (Field): The field to convert.
+
+    Returns:
+        types.IcebergType: The corresponding Iceberg IcebergType.
+    """
+    field_type = field.type
+    if field_type is None or field_type in ["null"]:
+        return types.NullType()
+    if field_type == "array":
+        return make_list(field.items)
+    if field_type == "map":
+        return make_map(field)
+    if field_type in ["object", "record", "struct"]:
+        return to_struct_type(field.fields)
+    if field_type in ["string", "varchar", "text"]:
+        return types.StringType()
+    if field_type in ["number", "decimal", "numeric"]:
+        precision = field.precision if field.precision is not None else 38
+        scale = field.scale if field.scale is not None else 0
+        return types.DecimalType(precision=precision, scale=scale)
+    if field_type in ["integer", "int"]:
+        return types.IntegerType()
+    if field_type in ["bigint", "long"]:
+        return types.LongType()
+    if field_type == "float":
+        return types.FloatType()
+    if field_type == "double":
+        return types.DoubleType()
+    if field_type == "boolean":
+        return types.BooleanType()
+    if field_type in ["timestamp", "timestamp_tz"]:
+        return types.TimestamptzType()
+    if field_type == "timestamp_ntz":
+        return types.TimestampType()
+    if field_type == "date":
+        return types.DateType()
+    if field_type == "bytes":
+        return types.BinaryType()
+    return types.BinaryType()

datacontract/export/jsonschema_converter.py

@@ -1,9 +1,8 @@
 import json
 from typing import Dict
 
-from datacontract.model.data_contract_specification import DataContractSpecification, Model, Field
-
 from datacontract.export.exporter import Exporter, _check_models_for_export
+from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model
 
 
 class JsonSchemaExporter(Exporter):
@@ -51,6 +50,8 @@ def to_property(field: Field) -> dict:
             property["type"] = json_type
     if json_format is not None:
         property["format"] = json_format
+    if field.primaryKey:
+        property["primaryKey"] = field.primaryKey
     if field.unique:
         property["unique"] = True
     if json_type == "object":

datacontract/export/odcs_v2_exporter.py

@@ -2,8 +2,8 @@ from typing import Dict
 
 import yaml
 
-from datacontract.model.data_contract_specification import DataContractSpecification, Model, Field
 from datacontract.export.exporter import Exporter
+from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model
 
 
 class OdcsV2Exporter(Exporter):

datacontract/export/odcs_v3_exporter.py

@@ -3,7 +3,7 @@ from typing import Dict
 import yaml
 
 from datacontract.export.exporter import Exporter
-from datacontract.model.data_contract_specification import DataContractSpecification, Model, Field
+from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model
 
 
 class OdcsV3Exporter(Exporter):
@@ -148,6 +148,10 @@ def to_odcs_schema(model_key, model_value: Model) -> dict:
     if properties:
         odcs_table["properties"] = properties
 
+    model_quality = to_odcs_quality_list(model_value.quality)
+    if len(model_quality) > 0:
+        odcs_table["quality"] = model_quality
+
     odcs_table["customProperties"] = []
     if model_value.model_extra is not None:
         for key, value in model_value.model_extra.items():
@@ -257,38 +261,48 @@ def to_property(field_name: str, field: Field) -> dict:
         del property["logicalTypeOptions"]
 
     if field.quality is not None:
-        quality_property = []
-        for quality in field.quality:
-            quality_dict = {"type": quality.type}
-            if quality.description is not None:
-                quality_dict["description"] = quality.description
-            if quality.query is not None:
-                quality_dict["query"] = quality.query
-            # dialect is not supported in v3.0.0
-            if quality.mustBe is not None:
-                quality_dict["mustBe"] = quality.mustBe
-            if quality.mustNotBe is not None:
-                quality_dict["mustNotBe"] = quality.mustNotBe
-            if quality.mustBeGreaterThan is not None:
-                quality_dict["mustBeGreaterThan"] = quality.mustBeGreaterThan
-            if quality.mustBeGreaterThanOrEqualTo is not None:
-                quality_dict["mustBeGreaterThanOrEqualTo"] = quality.mustBeGreaterThanOrEqualTo
-            if quality.mustBeLessThan is not None:
-                quality_dict["mustBeLessThan"] = quality.mustBeLessThan
-            if quality.mustBeLessThanOrEqualTo is not None:
-                quality_dict["mustBeLessThanOrEqualTo"] = quality.mustBeLessThanOrEqualTo
-            if quality.mustBeBetween is not None:
-                quality_dict["mustBeBetween"] = quality.mustBeBetween
-            if quality.mustNotBeBetween is not None:
-                quality_dict["mustNotBeBetween"] = quality.mustNotBeBetween
-            if quality.engine is not None:
-                quality_dict["engine"] = quality.engine
-            if quality.implementation is not None:
-                quality_dict["implementation"] = quality.implementation
-            quality_property.append(quality_dict)
+        quality_list = field.quality
+        quality_property = to_odcs_quality_list(quality_list)
         if len(quality_property) > 0:
             property["quality"] = quality_property
 
     # todo enum
 
     return property
+
+
+def to_odcs_quality_list(quality_list):
+    quality_property = []
+    for quality in quality_list:
+        quality_property.append(to_odcs_quality(quality))
+    return quality_property
+
+
+def to_odcs_quality(quality):
+    quality_dict = {"type": quality.type}
+    if quality.description is not None:
+        quality_dict["description"] = quality.description
+    if quality.query is not None:
+        quality_dict["query"] = quality.query
+    # dialect is not supported in v3.0.0
+    if quality.mustBe is not None:
+        quality_dict["mustBe"] = quality.mustBe
+    if quality.mustNotBe is not None:
+        quality_dict["mustNotBe"] = quality.mustNotBe
+    if quality.mustBeGreaterThan is not None:
+        quality_dict["mustBeGreaterThan"] = quality.mustBeGreaterThan
+    if quality.mustBeGreaterThanOrEqualTo is not None:
+        quality_dict["mustBeGreaterThanOrEqualTo"] = quality.mustBeGreaterThanOrEqualTo
+    if quality.mustBeLessThan is not None:
+        quality_dict["mustBeLessThan"] = quality.mustBeLessThan
+    if quality.mustBeLessThanOrEqualTo is not None:
+        quality_dict["mustBeLessThanOrEqualTo"] = quality.mustBeLessThanOrEqualTo
+    if quality.mustBeBetween is not None:
+        quality_dict["mustBeBetween"] = quality.mustBeBetween
+    if quality.mustNotBeBetween is not None:
+        quality_dict["mustNotBeBetween"] = quality.mustNotBeBetween
+    if quality.engine is not None:
+        quality_dict["engine"] = quality.engine
+    if quality.implementation is not None:
+        quality_dict["implementation"] = quality.implementation
+    return quality_dict