datacontract-cli 0.10.8__py3-none-any.whl → 0.10.10__py3-none-any.whl

datacontract/export/sodacl_converter.py

@@ -41,27 +41,32 @@ def to_checks(model_key, model_value, server_type: str, check_types: bool):
         if field.unique:
             checks.append(check_field_unique(field_name, quote_field_name))
         if field.minLength is not None:
-            checks.append(check_field_min_length(field_name, field.minLength))
+            checks.append(check_field_min_length(field_name, field.minLength, quote_field_name))
         if field.maxLength is not None:
-            checks.append(check_field_max_length(field_name, field.maxLength))
+            checks.append(check_field_max_length(field_name, field.maxLength, quote_field_name))
         if field.minimum is not None:
-            checks.append(check_field_minimum(field_name, field.minimum))
+            checks.append(check_field_minimum(field_name, field.minimum, quote_field_name))
         if field.maximum is not None:
-            checks.append(check_field_maximum(field_name, field.maximum))
+            checks.append(check_field_maximum(field_name, field.maximum, quote_field_name))
         if field.exclusiveMinimum is not None:
-            checks.append(check_field_minimum(field_name, field.exclusiveMinimum))
-            checks.append(check_field_not_equal(field_name, field.exclusiveMinimum))
+            checks.append(check_field_minimum(field_name, field.exclusiveMinimum, quote_field_name))
+            checks.append(check_field_not_equal(field_name, field.exclusiveMinimum, quote_field_name))
         if field.exclusiveMaximum is not None:
-            checks.append(check_field_maximum(field_name, field.exclusiveMaximum))
-            checks.append(check_field_not_equal(field_name, field.exclusiveMaximum))
+            checks.append(check_field_maximum(field_name, field.exclusiveMaximum, quote_field_name))
+            checks.append(check_field_not_equal(field_name, field.exclusiveMaximum, quote_field_name))
         if field.pattern is not None:
-            checks.append(check_field_regex(field_name, field.pattern))
+            checks.append(check_field_regex(field_name, field.pattern, quote_field_name))
         if field.enum is not None and len(field.enum) > 0:
-            checks.append(check_field_enum(field_name, field.enum))
+            checks.append(check_field_enum(field_name, field.enum, quote_field_name))
         # TODO references: str = None
         # TODO format
 
-    return f"checks for {model_key}", checks
+    checks_for_model_key = f"checks for {model_key}"
+
+    if quote_field_name:
+        checks_for_model_key = f'checks for "{model_key}"'
+
+    return checks_for_model_key, checks
 
 
 def check_field_is_present(field_name):
@@ -126,7 +131,7 @@ def check_field_minimum(field_name, minimum, quote_field_name: bool = False):
         field_name = f'"{field_name}"'
     return {
         f"invalid_count({field_name}) = 0": {
-            "name": f"Check that field {field_name} has a minimum of {min}",
+            "name": f"Check that field {field_name} has a minimum of {minimum}",
             "valid min": minimum,
         }
     }

datacontract/export/spark_converter.py

@@ -0,0 +1,211 @@
+from pyspark.sql import types
+from datacontract.model.data_contract_specification import (
+    DataContractSpecification,
+    Model,
+    Field,
+)
+from datacontract.export.exporter import Exporter
+
+
+class SparkExporter(Exporter):
+    """
+    Exporter class for exporting data contracts to Spark schemas.
+    """
+
+    def export(
+        self,
+        data_contract: DataContractSpecification,
+        model,
+        server,
+        sql_server_type,
+        export_args,
+    ) -> dict[str, types.StructType]:
+        """
+        Export the given data contract to Spark schemas.
+
+        Args:
+            data_contract (DataContractSpecification): The data contract specification.
+            model: Not used in this implementation.
+            server: Not used in this implementation.
+            sql_server_type: Not used in this implementation.
+            export_args: Additional arguments for export.
+
+        Returns:
+            dict[str, types.StructType]: A dictionary mapping model names to their corresponding Spark schemas.
+        """
+        return to_spark(data_contract)
+
+
+def to_spark(contract: DataContractSpecification) -> str:
+    """
+    Converts a DataContractSpecification into a Spark schema string.
+
+    Args:
+        contract (DataContractSpecification): The data contract specification containing models.
+
+    Returns:
+        str: A string representation of the Spark schema for each model in the contract.
+    """
+    return "\n\n".join(
+        f"{model_name} = {print_schema(to_spark_schema(model))}" for model_name, model in contract.models.items()
+    )
+
+
+def to_spark_dict(contract: DataContractSpecification) -> dict[str, types.StructType]:
+    """
+    Convert a data contract specification to Spark schemas.
+
+    Args:
+        contract (DataContractSpecification): The data contract specification.
+
+    Returns:
+        dict[str, types.StructType]: A dictionary mapping model names to their corresponding Spark schemas.
+    """
+    return {model_name: to_spark_schema(model) for model_name, model in contract.models.items()}
+
+
+def to_spark_schema(model: Model) -> types.StructType:
+    """
+    Convert a model to a Spark schema.
+
+    Args:
+        model (Model): The model to convert.
+
+    Returns:
+        types.StructType: The corresponding Spark schema.
+    """
+    return to_struct_type(model.fields)
+
+
+def to_struct_type(fields: dict[str, Field]) -> types.StructType:
+    """
+    Convert a dictionary of fields to a Spark StructType.
+
+    Args:
+        fields (dict[str, Field]): The fields to convert.
+
+    Returns:
+        types.StructType: The corresponding Spark StructType.
+    """
+    struct_fields = [to_struct_field(field, field_name) for field_name, field in fields.items()]
+    return types.StructType(struct_fields)
+
+
+def to_struct_field(field: Field, field_name: str) -> types.StructField:
+    """
+    Convert a field to a Spark StructField.
+
+    Args:
+        field (Field): The field to convert.
+        field_name (str): The name of the field.
+
+    Returns:
+        types.StructField: The corresponding Spark StructField.
+    """
+    data_type = to_data_type(field)
+    return types.StructField(name=field_name, dataType=data_type, nullable=not field.required)
+
+
+def to_data_type(field: Field) -> types.DataType:
+    """
+    Convert a field to a Spark DataType.
+
+    Args:
+        field (Field): The field to convert.
+
+    Returns:
+        types.DataType: The corresponding Spark DataType.
+    """
+    field_type = field.type
+    if field_type is None or field_type in ["null"]:
+        return types.NullType()
+    if field_type == "array":
+        return types.ArrayType(to_data_type(field.items))
+    if field_type in ["object", "record", "struct"]:
+        return types.StructType(to_struct_type(field.fields))
+    if field_type in ["string", "varchar", "text"]:
+        return types.StringType()
+    if field_type in ["number", "decimal", "numeric"]:
+        return types.DecimalType()
+    if field_type in ["integer", "int"]:
+        return types.IntegerType()
+    if field_type == "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.TimestampType()
+    if field_type == "timestamp_ntz":
+        return types.TimestampNTZType()
+    if field_type == "date":
+        return types.DateType()
+    if field_type == "bytes":
+        return types.BinaryType()
+    return types.BinaryType()
+
+
+def print_schema(dtype: types.DataType) -> str:
+    """
+    Converts a PySpark DataType schema to its equivalent code representation.
+
+    Args:
+        dtype (types.DataType): The PySpark DataType schema to be converted.
+
+    Returns:
+        str: The code representation of the PySpark DataType schema.
+    """
+
+    def indent(text: str, level: int) -> str:
+        """
+        Indents each line of the given text by a specified number of levels.
+
+        Args:
+            text (str): The text to be indented.
+            level (int): The number of indentation levels.
+
+        Returns:
+            str: The indented text.
+        """
+        return "\n".join([f'{"    " * level}{line}' for line in text.split("\n")])
+
+    def repr_column(column: types.StructField) -> str:
+        """
+        Converts a PySpark StructField to its code representation.
+
+        Args:
+            column (types.StructField): The StructField to be converted.
+
+        Returns:
+            str: The code representation of the StructField.
+        """
+        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)"
+
+    def format_struct_type(struct_type: types.StructType) -> str:
+        """
+        Converts a PySpark StructType to its code representation.
+
+        Args:
+            struct_type (types.StructType): The StructType to be converted.
+
+        Returns:
+            str: The code representation of the StructType.
+        """
+        fields = ",\n".join([indent(repr_column(field), 1) for field in struct_type.fields])
+        return f"StructType([\n{fields}\n])"
+
+    if isinstance(dtype, types.StructType):
+        return format_struct_type(dtype)
+    elif isinstance(dtype, types.ArrayType):
+        return f"ArrayType({print_schema(dtype.elementType)})"
+    elif isinstance(dtype, types.DecimalType):
+        return f"DecimalType({dtype.precision}, {dtype.scale})"
+    else:
+        dtype_str = str(dtype)
+        return dtype_str if dtype_str.endswith("()") else f"{dtype_str}()"

datacontract/export/sql_type_converter.py

@@ -15,6 +15,8 @@ def convert_to_sql_type(field: Field, server_type: str) -> str:
         return convert_type_to_sqlserver(field)
     elif server_type == "bigquery":
         return convert_type_to_bigquery(field)
+    elif server_type == "trino":
+        return convert_type_to_trino(field)
     return field.type
 
 
@@ -249,3 +251,29 @@ def get_type_config(field: Field, config_attr: str) -> dict[str, str] | None:
     if not field.config:
         return None
     return field.config.get(config_attr, None)
+
+
+def convert_type_to_trino(field: Field) -> None | str:
+    """Convert from supported datacontract types to equivalent trino types"""
+    field_type = field.type
+
+    if field_type.lower() in ["string", "text", "varchar"]:
+        return "varchar"
+    # tinyint, smallint not supported by data contract
+    if field_type.lower() in ["number", "decimal", "numeric"]:
+        # precision and scale not supported by data contract
+        return "decimal"
+    if field_type.lower() in ["int", "integer"]:
+        return "integer"
+    if field_type.lower() in ["long", "bigint"]:
+        return "bigint"
+    if field_type.lower() in ["float"]:
+        return "real"
+    if field_type.lower() in ["timestamp", "timestamp_tz"]:
+        return "timestamp(3) with time zone"
+    if field_type.lower() in ["timestamp_ntz"]:
+        return "timestamp(3)"
+    if field_type.lower() in ["bytes"]:
+        return "varbinary"
+    if field_type.lower() in ["object", "record", "struct"]:
+        return "json"

datacontract/imports/avro_importer.py

@@ -1,10 +1,46 @@
+from typing import Dict, List
+
 import avro.schema
 
+from datacontract.imports.importer import Importer
 from datacontract.model.data_contract_specification import DataContractSpecification, Model, Field
 from datacontract.model.exceptions import DataContractException
 
 
+class AvroImporter(Importer):
+    """Class to import Avro Schema file"""
+
+    def import_source(
+        self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
+    ) -> DataContractSpecification:
+        """
+        Import Avro schema from a source file.
+
+        Args:
+            data_contract_specification: The data contract specification to update.
+            source: The path to the Avro schema file.
+            import_args: Additional import arguments.
+
+        Returns:
+            The updated data contract specification.
+        """
+        return import_avro(data_contract_specification, source)
+
+
 def import_avro(data_contract_specification: DataContractSpecification, source: str) -> DataContractSpecification:
+    """
+    Import an Avro schema from a file and update the data contract specification.
+
+    Args:
+        data_contract_specification: The data contract specification to update.
+        source: The path to the Avro schema file.
+
+    Returns:
+        DataContractSpecification: The updated data contract specification.
+
+    Raises:
+        DataContractException: If there's an error parsing the Avro schema.
+    """
     if data_contract_specification.models is None:
         data_contract_specification.models = {}
 
@@ -37,7 +73,14 @@ def import_avro(data_contract_specification: DataContractSpecification, source:
     return data_contract_specification
 
 
-def handle_config_avro_custom_properties(field, imported_field):
+def handle_config_avro_custom_properties(field: avro.schema.Field, imported_field: Field) -> None:
+    """
+    Handle custom Avro properties and add them to the imported field's config.
+
+    Args:
+        field: The Avro field.
+        imported_field: The imported field to update.
+    """
     if field.get_prop("logicalType") is not None:
         if imported_field.config is None:
             imported_field.config = {}
@@ -49,7 +92,16 @@ def handle_config_avro_custom_properties(field, imported_field):
         imported_field.config["avroDefault"] = field.default
 
 
-def import_record_fields(record_fields):
+def import_record_fields(record_fields: List[avro.schema.Field]) -> Dict[str, Field]:
+    """
+    Import Avro record fields and convert them to data contract fields.
+
+    Args:
+        record_fields: List of Avro record fields.
+
+    Returns:
+        A dictionary of imported fields.
+    """
     imported_fields = {}
     for field in record_fields:
         imported_field = Field()
@@ -75,6 +127,15 @@ def import_record_fields(record_fields):
         elif field.type.type == "array":
             imported_field.type = "array"
             imported_field.items = import_avro_array_items(field.type)
+        elif field.type.type == "map":
+            imported_field.type = "map"
+            imported_field.values = import_avro_map_values(field.type)
+        elif field.type.type == "enum":
+            imported_field.type = "string"
+            imported_field.enum = field.type.symbols
+            if not imported_field.config:
+                imported_field.config = {}
+            imported_field.config["avroType"] = "enum"
         else:  # primitive type
             imported_field.type = map_type_from_avro(field.type.type)
 
@@ -83,7 +144,16 @@ def import_record_fields(record_fields):
     return imported_fields
 
 
-def import_avro_array_items(array_schema):
+def import_avro_array_items(array_schema: avro.schema.ArraySchema) -> Field:
+    """
+    Import Avro array items and convert them to a data contract field.
+
+    Args:
+        array_schema: The Avro array schema.
+
+    Returns:
+        Field: The imported field representing the array items.
+    """
     items = Field()
     for prop in array_schema.other_props:
         items.__setattr__(prop, array_schema.other_props[prop])
@@ -100,7 +170,45 @@ def import_avro_array_items(array_schema):
     return items
 
 
-def import_type_of_optional_field(field):
+def import_avro_map_values(map_schema: avro.schema.MapSchema) -> Field:
+    """
+    Import Avro map values and convert them to a data contract field.
+
+    Args:
+        map_schema: The Avro map schema.
+
+    Returns:
+        Field: The imported field representing the map values.
+    """
+    values = Field()
+    for prop in map_schema.other_props:
+        values.__setattr__(prop, map_schema.other_props[prop])
+
+    if map_schema.values.type == "record":
+        values.type = "object"
+        values.fields = import_record_fields(map_schema.values.fields)
+    elif map_schema.values.type == "array":
+        values.type = "array"
+        values.items = import_avro_array_items(map_schema.values)
+    else:  # primitive type
+        values.type = map_type_from_avro(map_schema.values.type)
+
+    return values
+
+
+def import_type_of_optional_field(field: avro.schema.Field) -> str:
+    """
+    Determine the type of optional field in an Avro union.
+
+    Args:
+        field: The Avro field with a union type.
+
+    Returns:
+        str: The mapped type of the non-null field in the union.
+
+    Raises:
+        DataContractException: If no non-null type is found in the union.
+    """
     for field_type in field.type.schemas:
         if field_type.type != "null":
             return map_type_from_avro(field_type.type)
@@ -113,21 +221,51 @@ def import_type_of_optional_field(field):
     )
 
 
-def get_record_from_union_field(field):
+def get_record_from_union_field(field: avro.schema.Field) -> avro.schema.RecordSchema | None:
+    """
+    Get the record schema from a union field.
+
+    Args:
+        field: The Avro field with a union type.
+
+    Returns:
+        The record schema if found, None otherwise.
+    """
     for field_type in field.type.schemas:
         if field_type.type == "record":
             return field_type
     return None
 
 
-def get_array_from_union_field(field):
+def get_array_from_union_field(field: avro.schema.Field) -> avro.schema.ArraySchema | None:
+    """
+    Get the array schema from a union field.
+
+    Args:
+        field: The Avro field with a union type.
+
+    Returns:
+        The array schema if found, None otherwise.
+    """
     for field_type in field.type.schemas:
         if field_type.type == "array":
             return field_type
     return None
 
 
-def map_type_from_avro(avro_type_str: str):
+def map_type_from_avro(avro_type_str: str) -> str:
+    """
+    Map Avro type strings to data contract type strings.
+
+    Args:
+        avro_type_str (str): The Avro type string.
+
+    Returns:
+        str: The corresponding data contract type string.
+
+    Raises:
+        DataContractException: If the Avro type is unsupported.
+    """
     # TODO: ambiguous mapping in the export
     if avro_type_str == "null":
         return "null"
@@ -147,6 +285,10 @@ def map_type_from_avro(avro_type_str: str):
         return "record"
     elif avro_type_str == "array":
         return "array"
+    elif avro_type_str == "map":
+        return "map"
+    elif avro_type_str == "enum":
+        return "string"
     else:
         raise DataContractException(
             type="schema",

datacontract/imports/bigquery_importer.py

@@ -2,10 +2,27 @@ import json
 import logging
 from typing import List
 
+from datacontract.imports.importer import Importer
 from datacontract.model.data_contract_specification import DataContractSpecification, Model, Field
 from datacontract.model.exceptions import DataContractException
 
 
+class BigQueryImporter(Importer):
+    def import_source(
+        self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
+    ) -> dict:
+        if source is not None:
+            data_contract_specification = import_bigquery_from_json(data_contract_specification, source)
+        else:
+            data_contract_specification = import_bigquery_from_api(
+                data_contract_specification,
+                import_args.get("bigquery_tables"),
+                import_args.get("bigquery_project"),
+                import_args.get("bigquery_dataset"),
+            )
+        return data_contract_specification
+
+
 def import_bigquery_from_json(
     data_contract_specification: DataContractSpecification, source: str
 ) -> DataContractSpecification:

datacontract/imports/dbt_importer.py

@@ -0,0 +1,117 @@
+import json
+
+from typing import (
+    List,
+)
+
+from datacontract.imports.importer import Importer
+from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model
+
+
+class DbtManifestImporter(Importer):
+    def import_source(
+        self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
+    ) -> dict:
+        data = read_dbt_manifest(manifest_path=source)
+        return import_dbt_manifest(
+            data_contract_specification, manifest_dict=data, dbt_models=import_args.get("dbt_model")
+        )
+
+
+def import_dbt_manifest(
+    data_contract_specification: DataContractSpecification, manifest_dict: dict, dbt_models: List[str]
+):
+    data_contract_specification.info.title = manifest_dict.get("info").get("project_name")
+    data_contract_specification.info.dbt_version = manifest_dict.get("info").get("dbt_version")
+
+    if data_contract_specification.models is None:
+        data_contract_specification.models = {}
+
+    for model in manifest_dict.get("models", []):
+        if dbt_models and model.name not in dbt_models:
+            continue
+
+        dc_model = Model(
+            description=model.description,
+            tags=model.tags,
+            fields=create_fields(model.columns),
+        )
+
+        data_contract_specification.models[model.name] = dc_model
+
+    return data_contract_specification
+
+
+def create_fields(columns: List):
+    fields = {}
+    for column in columns:
+        field = Field(
+            description=column.description, type=column.data_type if column.data_type else "", tags=column.tags
+        )
+        fields[column.name] = field
+
+    return fields
+
+
+def read_dbt_manifest(manifest_path: str):
+    with open(manifest_path, "r", encoding="utf-8") as f:
+        manifest = json.load(f)
+    return {"info": manifest.get("metadata"), "models": create_manifest_models(manifest)}
+
+
+def create_manifest_models(manifest: dict) -> List:
+    models = []
+    nodes = manifest.get("nodes")
+
+    for node in nodes.values():
+        if node["resource_type"] != "model":
+            continue
+
+        models.append(DbtModel(node))
+    return models
+
+
+class DbtColumn:
+    name: str
+    description: str
+    data_type: str
+    meta: dict
+    tags: List
+
+    def __init__(self, node_column: dict):
+        self.name = node_column.get("name")
+        self.description = node_column.get("description")
+        self.data_type = node_column.get("data_type")
+        self.meta = node_column.get("meta", {})
+        self.tags = node_column.get("tags", [])
+
+    def __repr__(self) -> str:
+        return self.name
+
+
+class DbtModel:
+    name: str
+    database: str
+    schema: str
+    description: str
+    unique_id: str
+    tags: List
+
+    def __init__(self, node: dict):
+        self.name = node.get("name")
+        self.database = node.get("database")
+        self.schema = node.get("schema")
+        self.description = node.get("description")
+        self.display_name = node.get("display_name")
+        self.unique_id = node.get("unique_id")
+        self.columns = []
+        self.tags = node.get("tags")
+        if node.get("columns"):
+            self.add_columns(node.get("columns").values())
+
+    def add_columns(self, model_columns: List):
+        for column in model_columns:
+            self.columns.append(DbtColumn(column))
+
+    def __repr__(self) -> str:
+        return self.name