datacontract-cli 0.9.6.post2__py3-none-any.whl → 0.9.8__py3-none-any.whl

datacontract/export/avro_idl_converter.py

@@ -0,0 +1,280 @@
+import typing
+from dataclasses import dataclass
+from enum import Enum
+from io import StringIO
+
+from datacontract.lint.resolve import inline_definitions_into_data_contract
+from datacontract.model.data_contract_specification import \
+    DataContractSpecification, Field
+from datacontract.model.exceptions import DataContractException
+
+
+def to_avro_idl(contract: DataContractSpecification) -> str:
+    """Serialize the provided data contract specification into an Avro IDL string.
+
+    The data contract will be serialized as a protocol, with one record type
+    for each contained model. Model fields are mapped one-to-one to Avro IDL
+    record fields.
+    """
+    stream = StringIO()
+    to_avro_idl_stream(contract, stream)
+    return stream.getvalue()
+
+
+def to_avro_idl_stream(contract: DataContractSpecification, stream: typing.TextIO):
+    """Serialize the provided data contract specification into Avro IDL."""
+    ir = _contract_to_avro_idl_ir(contract)
+    if ir.description:
+        stream.write(f"/** {contract.info.description} */\n")
+    stream.write(f"protocol {ir.name or 'Unnamed'} {{\n")
+    for model_type in ir.model_types:
+        _write_model_type(model_type, stream)
+    stream.write("}\n")
+
+
+class AvroPrimitiveType(Enum):
+    int = "int"
+    long = "long"
+    string = "string"
+    boolean = "boolean"
+    float = "float"
+    double = "double"
+    null = "null"
+    bytes = "bytes"
+
+
+class AvroLogicalType(Enum):
+    decimal = "decimal"
+    date = "date"
+    time_ms = "time_ms"
+    timestamp_ms = "timestamp_ms"
+
+
+@dataclass
+class AvroField:
+    name: str
+    required: bool
+    description: typing.Optional[str]
+
+
+@dataclass
+class AvroPrimitiveField(AvroField):
+    type: typing.Union[AvroPrimitiveType, AvroLogicalType]
+
+
+@dataclass
+class AvroComplexField(AvroField):
+    subfields: list[AvroField]
+
+
+@dataclass
+class AvroArrayField(AvroField):
+    type: AvroField
+
+
+@dataclass
+class AvroModelType:
+    name: str
+    description: typing.Optional[str]
+    fields: list[AvroField]
+
+
+@dataclass
+class AvroIDLProtocol:
+    name: typing.Optional[str]
+    description: typing.Optional[str]
+    model_types: list[AvroModelType]
+
+
+avro_primitive_types = set(
+    [
+        "string",
+        "text",
+        "varchar",
+        "float",
+        "double",
+        "int",
+        "integer",
+        "long",
+        "bigint",
+        "boolean",
+        "timestamp_ntz",
+        "timestamp",
+        "timestamp_tz",
+        "date",
+        "bytes",
+        "null",
+    ]
+)
+
+
+def _to_avro_primitive_logical_type(field_name: str, field: Field) -> AvroPrimitiveField:
+    result = AvroPrimitiveField(field_name, field.required, field.description, AvroPrimitiveType.string)
+    match field.type:
+        case "string" | "text" | "varchar":
+            result.type = AvroPrimitiveType.string
+        case "float":
+            result.type = AvroPrimitiveType.float
+        case "double":
+            result.type = AvroPrimitiveType.double
+        case "int" | "integer":
+            result.type = AvroPrimitiveType.int
+        case "long" | "bigint":
+            result.type = AvroPrimitiveType.long
+        case "boolean":
+            result.type = AvroPrimitiveType.boolean
+        case "timestamp" | "timestamp_tz":
+            result.type = AvroPrimitiveType.string
+        case "timestamp_ntz":
+            result.type = AvroLogicalType.timestamp_ms
+        case "date":
+            result.type = AvroLogicalType.date
+        case "bytes":
+            result.type = AvroPrimitiveType.bytes
+        case "null":
+            result.type = AvroPrimitiveType.null
+        case _:
+            raise DataContractException(
+                type="general",
+                name="avro-idl-export",
+                model=field,
+                reason="Unknown field type {field.type}",
+                result="failed",
+                message="Avro IDL type conversion failed.",
+            )
+    return result
+
+
+def _to_avro_idl_type(field_name: str, field: Field) -> AvroField:
+    if field.type in avro_primitive_types:
+        return _to_avro_primitive_logical_type(field_name, field)
+    else:
+        match field.type:
+            case "array":
+                return AvroArrayField(
+                    field_name, field.required, field.description, _to_avro_idl_type(field_name, field.items)
+                )
+            case "object" | "record" | "struct":
+                return AvroComplexField(
+                    field_name,
+                    field.required,
+                    field.description,
+                    [_to_avro_idl_type(field_name, field) for (field_name, field) in field.fields.items()],
+                )
+            case _:
+                raise DataContractException(
+                    type="general",
+                    name="avro-idl-export",
+                    model=type,
+                    reason="Unknown Data Contract field type",
+                    result="failed",
+                    message="Avro IDL type conversion failed.",
+                )
+
+
+def _generate_field_types(contract: DataContractSpecification) -> list[AvroField]:
+    result = []
+    for _, model in contract.models.items():
+        for field_name, field in model.fields.items():
+            result.append(_to_avro_idl_type(field_name, field))
+    return result
+
+
+def generate_model_types(contract: DataContractSpecification) -> list[AvroModelType]:
+    result = []
+    for model_name, model in contract.models.items():
+        result.append(
+            AvroModelType(name=model_name, description=model.description, fields=_generate_field_types(contract))
+        )
+    return result
+
+
+def _model_name_to_identifier(model_name: str):
+    return "".join([word.title() for word in model_name.split()])
+
+
+def _contract_to_avro_idl_ir(contract: DataContractSpecification) -> AvroIDLProtocol:
+    """Convert models into an intermediate representation for later serialization into Avro IDL.
+
+    Each model is converted to a record containing a field for each model field.
+    """
+    inlined_contract = contract.model_copy()
+    inline_definitions_into_data_contract(inlined_contract)
+    protocol_name = _model_name_to_identifier(contract.info.title) if contract.info and contract.info.title else None
+    description = contract.info.description if contract.info and contract.info.description else None
+    return AvroIDLProtocol(
+        name=protocol_name, description=description, model_types=generate_model_types(inlined_contract)
+    )
+
+
+def _write_indent(indent: int, stream: typing.TextIO):
+    stream.write("    " * indent)
+
+
+def _write_field_description(field: AvroField, indent: int, stream: typing.TextIO):
+    if field.description:
+        _write_indent(indent, stream)
+        stream.write(f"/** {field.description} */\n")
+
+
+def _write_field_type_definition(field: AvroField, indent: int, stream: typing.TextIO) -> str:
+    # Write any extra information (such as record type definition) and return
+    # the name of the generated type. Writes descriptions only for record
+    # types. This leads to record types being described twice, once on the
+    # record definition, and once on use. The alternative (detect when the
+    # complex field type is not used in an array or another complex type) is
+    # significantly more complex to implement.
+    match field:
+        case AvroPrimitiveField(name, required, _, typ) if required is True:
+            return typ.value
+        case AvroPrimitiveField(name, required, _, typ):
+            return typ.value + "?"
+        case AvroComplexField(name, required, _, subfields):
+            _write_field_description(field, indent, stream)
+            _write_indent(indent, stream)
+            stream.write(f"record {name}_type {{\n")
+            subfield_types = []
+            # Recursively define records for all subfields if necessary
+            for subfield in subfields:
+                subfield_types.append(_write_field_type_definition(subfield, indent + 1, stream))
+            # Reference all defined record types.
+            for field, subfield_type in zip(field.subfields, subfield_types):
+                _write_field_description(field, indent + 1, stream)
+                _write_indent(indent + 1, stream)
+                stream.write(f"{subfield_type} {field.name};\n")
+            _write_indent(indent, stream)
+            stream.write("}\n")
+            if required is True:
+                return f"{name}_type"
+            else:
+                return f"{name}_type?"
+        case AvroArrayField(name, required, _, item_type):
+            subfield_type = _write_field_type_definition(item_type, indent, stream)
+            if required is True:
+                return f"array<{subfield_type}>"
+            else:
+                return f"array<{subfield_type}>?"
+        case _:
+            raise RuntimeError("Unknown Avro field type {field}")
+
+
+def _write_field(field: AvroField, indent, stream: typing.TextIO):
+    # Start of recursion.
+    typename = _write_field_type_definition(field, indent, stream)
+    _write_field_description(field, indent, stream)
+    _write_indent(indent, stream)
+    stream.write(f"{typename} {field.name};\n")
+
+
+def _write_model_type(model: AvroModelType, stream: typing.TextIO):
+    # Called once for each model
+    if model.description:
+        _write_indent(1, stream)
+        stream.write(f"/** {model.description} */\n")
+    _write_indent(1, stream)
+    stream.write(f"record {model.name} {{\n")
+    # Called for each model field
+    for field in model.fields:
+        _write_field(field, 2, stream)
+    _write_indent(1, stream)
+    stream.write("}\n")

datacontract/export/dbt_converter.py

@@ -2,14 +2,11 @@ from typing import Dict
 
 import yaml
 
+from datacontract.export.sql_type_converter import convert_to_sql_type
 from datacontract.model.data_contract_specification import \
     DataContractSpecification, Model, Field
 
 
-# snowflake data types:
-# https://docs.snowflake.com/en/sql-reference/data-types.html
-
-
 def to_dbt_models_yaml(data_contract_spec: DataContractSpecification):
     dbt = {
         "version": 2,
@@ -18,18 +15,17 @@ def to_dbt_models_yaml(data_contract_spec: DataContractSpecification):
     for model_key, model_value in data_contract_spec.models.items():
         dbt_model = _to_dbt_model(model_key, model_value, data_contract_spec)
         dbt["models"].append(dbt_model)
-    return yaml.dump(dbt, indent=2, sort_keys=False)
+    return yaml.dump(dbt, indent=2, sort_keys=False, allow_unicode=True)
 
 
-def to_dbt_staging_sql(data_contract_spec: DataContractSpecification):
+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(f"Export to dbt-staging-sql currently only works with exactly one model in the data contract.")
+        print("Export to dbt-staging-sql currently only works with exactly one model in the data contract.")
         return ""
 
     id = data_contract_spec.id
-    model_name, model = next(iter(data_contract_spec.models.items()))
     columns = []
-    for field_name, field in model.fields.items():
+    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"""
@@ -40,15 +36,10 @@ def to_dbt_staging_sql(data_contract_spec: DataContractSpecification):
 
 
 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, "tables": []}
     dbt = {
         "version": 2,
-        "sources": [
-            source
-        ],
+        "sources": [source],
     }
     if data_contract_spec.info.owner is not None:
         source["meta"] = {"owner": data_contract_spec.info.owner}
@@ -62,7 +53,7 @@ def to_dbt_sources_yaml(data_contract_spec: DataContractSpecification, server: s
     for model_key, model_value in data_contract_spec.models.items():
         dbt_model = _to_dbt_source_table(model_key, model_value)
         source["tables"].append(dbt_model)
-    return yaml.dump(dbt, indent=2, sort_keys=False)
+    return yaml.dump(dbt, indent=2, sort_keys=False, allow_unicode=True)
 
 
 def _to_dbt_source_table(model_key, model_value: Model) -> dict:
@@ -83,20 +74,14 @@ def _to_dbt_model(model_key, model_value: Model, data_contract_spec: DataContrac
         "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"] = {"meta": {"data_contract": data_contract_spec.id}}
     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
 
     if _supports_constraints(model_type):
-        dbt_model["config"]["contract"] = {
-            "enforced": True
-        }
+        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)
@@ -133,14 +118,14 @@ def _to_columns(fields: Dict[str, Field], supports_constraints: bool, supports_d
 
 def _to_column(field: Field, supports_constraints: bool, supports_datatype: bool) -> dict:
     column = {}
-    dbt_type = _convert_type_to_snowflake(field.type)
+    dbt_type = convert_to_sql_type(field, "snowflake")
     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}})
+                {"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:
@@ -162,7 +147,8 @@ def _to_column(field: Field, supports_constraints: bool, supports_datatype: bool
         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})
+            {"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:
@@ -172,71 +158,60 @@ def _to_column(field: Field, supports_constraints: bool, supports_datatype: bool
     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}})
-    if field.minimum is not None or field.maximum is not None and field.minimumExclusive is None and field.maximumExclusive is None:
+            {"dbt_expectations.expect_column_values_to_match_regex": {"regex": field.pattern}}
+        )
+    if (
+        field.minimum is not None
+        or field.maximum is not None
+        and field.exclusiveMinimum is None
+        and field.exclusiveMaximum is None
+    ):
         range_test = {}
         if field.minimum is not None:
             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})
-    elif field.minimumExclusive is not None or field.maximumExclusive is not None and field.minimum is None and field.maximum is None:
+    elif (
+        field.exclusiveMinimum is not None
+        or field.exclusiveMaximum is not None
+        and field.minimum is None
+        and field.maximum is None
+    ):
         range_test = {}
-        if field.minimumExclusive is not None:
-            range_test["min_value"] = field.minimumExclusive
-        if field.maximumExclusive is not None:
-            range_test["max_value"] = field.maximumExclusive
+        if field.exclusiveMinimum is not None:
+            range_test["min_value"] = field.exclusiveMinimum
+        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})
     else:
         if field.minimum is not None:
             column.setdefault("tests", []).append(
-                {"dbt_expectations.expect_column_values_to_be_between": {"min_value": field.minimum}})
+                {"dbt_expectations.expect_column_values_to_be_between": {"min_value": field.minimum}}
+            )
         if field.maximum is not None:
             column.setdefault("tests", []).append(
-                {"dbt_expectations.expect_column_values_to_be_between": {"max_value": field.maximum}})
-        if field.minimumExclusive is not None:
-            column.setdefault("tests", []).append({"dbt_expectations.expect_column_values_to_be_between": {
-                "min_value": field.minimumExclusive, "strictly": True}})
-        if field.maximumExclusive is not None:
-            column.setdefault("tests", []).append({"dbt_expectations.expect_column_values_to_be_between": {
-                "max_value": field.maximumExclusive, "strictly": True}})
+                {"dbt_expectations.expect_column_values_to_be_between": {"max_value": field.maximum}}
+            )
+        if field.exclusiveMinimum is not None:
+            column.setdefault("tests", []).append(
+                {
+                    "dbt_expectations.expect_column_values_to_be_between": {
+                        "min_value": field.exclusiveMinimum,
+                        "strictly": True,
+                    }
+                }
+            )
+        if field.exclusiveMaximum is not None:
+            column.setdefault("tests", []).append(
+                {
+                    "dbt_expectations.expect_column_values_to_be_between": {
+                        "max_value": field.exclusiveMaximum,
+                        "strictly": True,
+                    }
+                }
+            )
 
     # TODO: all constraints
     return column
-
-
-def _convert_type_to_snowflake(type) -> None | str:
-    # currently optimized for snowflake
-    # LEARNING: data contract has no direct support for CHAR,CHARACTER
-    # LEARNING: data contract has no support for "date-time", "datetime", "time"
-    # LEARNING: No precision and scale support in data contract
-    # LEARNING: no support for any
-    # GEOGRAPHY and GEOMETRY are not supported by the mapping
-    if type is None:
-        return None
-    if type.lower() in ["string", "varchar", "text"]:
-        return type.upper()  # STRING, TEXT, VARCHAR are all the same in snowflake
-    if type.lower() in ["timestamp", "timestamp_tz"]:
-        return "TIMESTAMP_TZ"
-    if type.lower() in ["timestamp_ntz"]:
-        return "TIMESTAMP_NTZ"
-    if type.lower() in ["date"]:
-        return "DATE"
-    if type.lower() in ["time"]:
-        return "TIME"
-    if type.lower() in ["number", "decimal", "numeric"]:
-        return "NUMBER"  # precision and scale not supported by data contract
-    if type.lower() in ["float", "double"]:
-        return "FLOAT"
-    if type.lower() in ["integer", "int", "long", "bigint"]:
-        return "NUMBER"  # always NUMBER(38,0)
-    if type.lower() in ["boolean"]:
-        return "BOOLEAN"
-    if type.lower() in ["object", "record", "struct"]:
-        return "OBJECT"
-    if type.lower() in ["bytes"]:
-        return "BINARY"
-    if type.lower() in ["array"]:
-        return "ARRAY"
-    return None

datacontract/export/great_expectations_converter.py

@@ -0,0 +1,141 @@
+import json
+from typing import Dict, List, Any
+
+import yaml
+
+from datacontract.model.data_contract_specification import \
+    DataContractSpecification, Field, Quality
+
+
+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
+    """
+    expectations = []
+    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(checks_to_expectations(quality_checks, model_key))
+    model_expectation_suite = to_suite(model_key, data_contract_spec.info.version, expectations)
+
+    return model_expectation_suite
+
+
+def to_suite(
+    model_key: str,
+    contract_version: str,
+    expectations: List[Dict[str, Any]],
+) -> str:
+    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
+            ),
+            "expectations": expectations,
+            "meta": {},
+        },
+        indent=2,
+    )
+
+
+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
+    """
+    expectations = []
+    add_column_order_exp(fields, expectations)
+    for field_name, field in fields.items():
+        add_field_expectations(field_name, field, expectations)
+    return expectations
+
+
+def add_field_expectations(field_name, field: Field, expectations: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    if field.type is not None:
+        expectations.append(to_column_types_exp(field_name, field.type))
+    if field.unique is not None:
+        expectations.append(to_column_unique_exp(field_name))
+    if field.maxLength is not None or field.minLength is not None:
+        expectations.append(to_column_length_exp(field_name, field.minLength, field.maxLength))
+    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]]):
+    expectations.append(
+        {
+            "expectation_type": "expect_table_columns_to_match_ordered_list",
+            "kwargs": {"column_list": list(fields.keys())},
+            "meta": {},
+        }
+    )
+
+
+def to_column_types_exp(field_name, field_type) -> Dict[str, Any]:
+    return {
+        "expectation_type": "expect_column_values_to_be_of_type",
+        "kwargs": {"column": field_name, "type_": field_type},
+        "meta": {},
+    }
+
+
+def to_column_unique_exp(field_name) -> Dict[str, Any]:
+    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]:
+    return {
+        "expectation_type": "expect_column_value_lengths_to_be_between",
+        "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]:
+    return {
+        "expectation_type": "expect_column_values_to_be_between",
+        "kwargs": {"column": field_name, "min_value": minimum, "max_value": maximum},
+        "meta": {},
+    }
+
+
+def get_quality_checks(quality: Quality) -> Dict[str, Any]:
+    if quality is None:
+        return {}
+    if quality.type is None:
+        return {}
+    if quality.type.lower() != "great-expectations":
+        return {}
+    if isinstance(quality.specification, str):
+        quality_specification = yaml.safe_load(quality.specification)
+    else:
+        quality_specification = quality.specification
+    return quality_specification
+
+
+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
+    """
+    if quality_checks is None or model_key not in quality_checks:
+        return []
+
+    model_quality_checks = quality_checks[model_key]
+
+    if model_quality_checks is None:
+        return []
+
+    if isinstance(model_quality_checks, str):
+        expectation_list = json.loads(model_quality_checks)
+        return expectation_list

datacontract/export/jsonschema_converter.py

@@ -12,16 +12,18 @@ def to_jsonschemas(data_contract_spec: DataContractSpecification):
         jsonschmemas[model_key] = jsonschema
     return jsonschmemas
 
+
 def to_jsonschema_json(model_key, model_value: Model) -> str:
     jsonschema = to_jsonschema(model_key, model_value)
     return json.dumps(jsonschema, indent=2)
 
+
 def to_jsonschema(model_key, model_value: Model) -> dict:
     return {
         "$schema": "http://json-schema.org/draft-07/schema#",
         "type": "object",
         "properties": to_properties(model_value.fields),
-        "required": to_required(model_value.fields)
+        "required": to_required(model_value.fields),
     }