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

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,9 +1,17 @@
 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):
+    def import_source(
+        self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
+    ) -> dict:
+        return import_avro(data_contract_specification, source)
+
+
 def import_avro(data_contract_specification: DataContractSpecification, source: str) -> DataContractSpecification:
     if data_contract_specification.models is None:
         data_contract_specification.models = {}

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/glue_importer.py

@@ -1,6 +1,7 @@
 import boto3
 from typing import List
 
+from datacontract.imports.importer import Importer
 from datacontract.model.data_contract_specification import (
     DataContractSpecification,
     Model,
@@ -9,7 +10,14 @@ from datacontract.model.data_contract_specification import (
 )
 
 
-def get_glue_database(datebase_name: str):
+class GlueImporter(Importer):
+    def import_source(
+        self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
+    ) -> dict:
+        return import_glue(data_contract_specification, source, import_args.get("glue_tables"))
+
+
+def get_glue_database(database_name: str):
     """Get the details Glue database.
 
     Args:
@@ -18,31 +26,32 @@ def get_glue_database(datebase_name: str):
     Returns:
         set: catalogid and locationUri
     """
-
     glue = boto3.client("glue")
     try:
-        response = glue.get_database(Name=datebase_name)
+        response = glue.get_database(Name=database_name)
     except glue.exceptions.EntityNotFoundException:
-        print(f"Database not found {datebase_name}.")
+        print(f"Database not found {database_name}.")
         return (None, None)
     except Exception as e:
         # todo catch all
         print(f"Error: {e}")
         return (None, None)
 
-    return (response["Database"]["CatalogId"], response["Database"].get("LocationUri", "None"))
+    return (
+        response["Database"]["CatalogId"],
+        response["Database"].get("LocationUri", "None"),
+    )
 
 
 def get_glue_tables(database_name: str) -> List[str]:
     """Get the list of tables in a Glue database.
 
     Args:
-        database_name (str): glue database to request.
+        database_name (str): Glue database to request.
 
     Returns:
-        List[string]: List of table names
+        List[str]: List of table names
     """
-
     glue = boto3.client("glue")
 
     # Set the paginator
@@ -107,9 +116,21 @@ def get_glue_table_schema(database_name: str, table_name: str):
     return table_schema
 
 
-def import_glue(data_contract_specification: DataContractSpecification, source: str, table_names: List[str]):
-    """Import the schema of a Glue database."""
+def import_glue(
+    data_contract_specification: DataContractSpecification,
+    source: str,
+    table_names: List[str],
+):
+    """Import the schema of a Glue database.
 
+    Args:
+        data_contract_specification (DataContractSpecification): The data contract specification to update.
+        source (str): The name of the Glue database.
+        table_names (List[str]): List of table names to import. If None, all tables in the database are imported.
+
+    Returns:
+        DataContractSpecification: The updated data contract specification.
+    """
     catalogid, location_uri = get_glue_database(source)
 
     # something went wrong
@@ -131,21 +152,18 @@ def import_glue(data_contract_specification: DataContractSpecification, source:
 
         fields = {}
         for column in table_schema:
-            field = Field()
-            field.type = map_type_from_sql(column["Type"])
+            field = create_typed_field(column["Type"])
 
             # hive partitons are required, but are not primary keys
             if column.get("Hive"):
                 field.required = True
 
             field.description = column.get("Comment")
-
             fields[column["Name"]] = field
 
             if "decimal" in column["Type"]:
                 # Extract precision and scale from the string
                 perc_scale = column["Type"][8:-1].split(",")
-                print(perc_scale)
                 field.precision = int(perc_scale[0])
                 field.scale = int(perc_scale[1])
 
@@ -157,37 +175,102 @@ def import_glue(data_contract_specification: DataContractSpecification, source:
     return data_contract_specification
 
 
-def map_type_from_sql(sql_type: str):
+def create_typed_field(dtype: str) -> Field:
+    """Create a typed field based on the given data type.
+
+    Args:
+        dtype (str): The data type of the field.
+
+    Returns:
+        Field: The created field with the appropriate type.
+    """
+    field = Field()
+    dtype = dtype.strip().lower().replace(" ", "")
+    if dtype.startswith(("array", "struct")):
+        orig_dtype: str = dtype
+        if dtype.startswith("array"):
+            field.type = "array"
+            field.items = create_typed_field(orig_dtype[6:-1])
+        elif dtype.startswith("struct"):
+            field.type = "struct"
+            for f in split_struct(orig_dtype[7:-1]):
+                field.fields[f.split(":", 1)[0].strip()] = create_typed_field(f.split(":", 1)[1])
+    else:
+        field.type = map_type_from_sql(dtype)
+    return field
+
+
+def split_fields(s: str):
+    """Split a string of fields considering nested structures.
+
+    Args:
+        s (str): The string to split.
+
+    Yields:
+        str: The next field in the string.
+    """
+    counter: int = 0
+    last: int = 0
+    for i, x in enumerate(s):
+        if x in ("<", "("):
+            counter += 1
+        elif x in (">", ")"):
+            counter -= 1
+        elif x == "," and counter == 0:
+            yield s[last:i]
+            last = i + 1
+    yield s[last:]
+
+
+def split_struct(s: str) -> List[str]:
+    """Split a struct string into individual fields.
+
+    Args:
+        s (str): The struct string to split.
+
+    Returns:
+        List[str]: List of individual fields in the struct.
+    """
+    return list(split_fields(s=s))
+
+
+def map_type_from_sql(sql_type: str) -> str:
+    """Map an SQL type to a corresponding field type.
+
+    Args:
+        sql_type (str): The SQL type to map.
+
+    Returns:
+        str: The corresponding field type.
+    """
     if sql_type is None:
         return None
 
-    if sql_type.lower().startswith("varchar"):
+    sql_type = sql_type.lower()
+    if sql_type.startswith("varchar"):
         return "varchar"
-    if sql_type.lower().startswith("string"):
+    if sql_type.startswith("string"):
         return "string"
-    if sql_type.lower().startswith("text"):
+    if sql_type.startswith("text"):
         return "text"
-    elif sql_type.lower().startswith("byte"):
+    if sql_type.startswith("byte"):
         return "byte"
-    elif sql_type.lower().startswith("short"):
+    if sql_type.startswith("short"):
         return "short"
-    elif sql_type.lower().startswith("integer"):
+    if sql_type.startswith("integer") or sql_type.startswith("int"):
         return "integer"
-    elif sql_type.lower().startswith("long"):
+    if sql_type.startswith("long") or sql_type.startswith("bigint"):
         return "long"
-    elif sql_type.lower().startswith("bigint"):
-        return "long"
-    elif sql_type.lower().startswith("float"):
+    if sql_type.startswith("float"):
         return "float"
-    elif sql_type.lower().startswith("double"):
+    if sql_type.startswith("double"):
         return "double"
-    elif sql_type.lower().startswith("boolean"):
+    if sql_type.startswith("boolean"):
         return "boolean"
-    elif sql_type.lower().startswith("timestamp"):
+    if sql_type.startswith("timestamp"):
         return "timestamp"
-    elif sql_type.lower().startswith("date"):
+    if sql_type.startswith("date"):
         return "date"
-    elif sql_type.lower().startswith("decimal"):
+    if sql_type.startswith("decimal"):
         return "decimal"
-    else:
-        return "variant"
+    return "variant"

datacontract/imports/importer.py

@@ -0,0 +1,29 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+
+from datacontract.model.data_contract_specification import DataContractSpecification
+
+
+class Importer(ABC):
+    def __init__(self, import_format) -> None:
+        self.import_format = import_format
+
+    @abstractmethod
+    def import_source(
+        self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
+    ) -> dict:
+        pass
+
+
+class ImportFormat(str, Enum):
+    sql = "sql"
+    avro = "avro"
+    glue = "glue"
+    jsonschema = "jsonschema"
+    bigquery = "bigquery"
+    odcs = "odcs"
+    unity = "unity"
+
+    @classmethod
+    def get_suported_formats(cls):
+        return list(map(lambda c: c.value, cls))

datacontract/imports/importer_factory.py

@@ -0,0 +1,72 @@
+import importlib.util
+import sys
+from datacontract.imports.importer import ImportFormat, Importer
+
+
+class ImporterFactory:
+    def __init__(self):
+        self.dict_importer = {}
+        self.dict_lazy_importer = {}
+
+    def register_importer(self, name, importer: Importer):
+        self.dict_importer.update({name: importer})
+
+    def register_lazy_importer(self, name: str, module_path: str, class_name: str):
+        self.dict_lazy_importer.update({name: (module_path, class_name)})
+
+    def create(self, name) -> Importer:
+        importers = self.dict_importer.copy()
+        importers.update(self.dict_lazy_importer.copy())
+        if name not in importers.keys():
+            raise ValueError(f"The '{name}' format is not suportted.")
+        importer_class = importers[name]
+        if type(importers[name]) is tuple:
+            importer_class = load_module_class(module_path=importers[name][0], class_name=importers[name][1])
+        if not importer_class:
+            raise ValueError(f"Module {name} could not be loaded.")
+        return importer_class(name)
+
+
+def import_module(module_path):
+    if importlib.util.find_spec(module_path) is not None:
+        try:
+            module = importlib.import_module(module_path)
+        except ModuleNotFoundError:
+            return None
+        sys.modules[module_path] = module
+        return module
+
+
+def load_module_class(module_path, class_name):
+    module = import_module(module_path)
+    if not module:
+        return None
+    return getattr(module, class_name)
+
+
+importer_factory = ImporterFactory()
+importer_factory.register_lazy_importer(
+    name=ImportFormat.avro, module_path="datacontract.imports.avro_importer", class_name="AvroImporter"
+)
+importer_factory.register_lazy_importer(
+    name=ImportFormat.bigquery,
+    module_path="datacontract.imports.bigquery_importer",
+    class_name="BigQueryImporter",
+)
+importer_factory.register_lazy_importer(
+    name=ImportFormat.glue, module_path="datacontract.imports.glue_importer", class_name="GlueImporter"
+)
+importer_factory.register_lazy_importer(
+    name=ImportFormat.jsonschema,
+    module_path="datacontract.imports.jsonschema_importer",
+    class_name="JsonSchemaImporter",
+)
+importer_factory.register_lazy_importer(
+    name=ImportFormat.odcs, module_path="datacontract.imports.odcs_importer", class_name="OdcsImporter"
+)
+importer_factory.register_lazy_importer(
+    name=ImportFormat.sql, module_path="datacontract.imports.sql_importer", class_name="SqlImporter"
+)
+importer_factory.register_lazy_importer(
+    name=ImportFormat.unity, module_path="datacontract.imports.unity_importer", class_name="UnityImporter"
+)

datacontract/imports/jsonschema_importer.py

@@ -2,10 +2,18 @@ import json
 
 import fastjsonschema
 
+from datacontract.imports.importer import Importer
 from datacontract.model.data_contract_specification import DataContractSpecification, Model, Field, Definition
 from datacontract.model.exceptions import DataContractException
 
 
+class JsonSchemaImporter(Importer):
+    def import_source(
+        self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
+    ) -> dict:
+        return import_jsonschema(data_contract_specification, source)
+
+
 def convert_json_schema_properties(properties, is_definition=False):
     fields = {}
     for field_name, field_schema in properties.items():

datacontract/imports/odcs_importer.py

@@ -2,6 +2,7 @@ import datetime
 import logging
 from typing import Any, Dict, List
 import yaml
+from datacontract.imports.importer import Importer
 from datacontract.model.data_contract_specification import (
     Availability,
     Contact,
@@ -42,6 +43,13 @@ DATACONTRACT_TYPES = [
 ]
 
 
+class OdcsImporter(Importer):
+    def import_source(
+        self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
+    ) -> dict:
+        return import_odcs(data_contract_specification, source)
+
+
 def import_odcs(data_contract_specification: DataContractSpecification, source: str) -> DataContractSpecification:
     try:
         with open(source, "r") as file:

datacontract/imports/sql_importer.py

@@ -1,8 +1,16 @@
 from simple_ddl_parser import parse_from_file
 
+from datacontract.imports.importer import Importer
 from datacontract.model.data_contract_specification import DataContractSpecification, Model, Field
 
 
+class SqlImporter(Importer):
+    def import_source(
+        self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
+    ) -> dict:
+        return import_sql(data_contract_specification, self.import_format, source)
+
+
 def import_sql(data_contract_specification: DataContractSpecification, format: str, source: str):
     ddl = parse_from_file(source, group_by_type=True)
     tables = ddl["tables"]