datacontract-cli 0.10.0__py3-none-any.whl → 0.10.37__py3-none-any.whl

datacontract/imports/glue_importer.py

@@ -0,0 +1,288 @@
+import re
+from typing import Dict, Generator, List
+
+import boto3
+
+from datacontract.imports.importer import Importer
+from datacontract.model.data_contract_specification import (
+    DataContractSpecification,
+    Field,
+    Model,
+    Server,
+)
+
+
+class GlueImporter(Importer):
+    def import_source(
+        self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
+    ) -> DataContractSpecification:
+        return import_glue(data_contract_specification, source, import_args.get("glue_table"))
+
+
+def get_glue_database(database_name: str):
+    """Get the details Glue database.
+
+    Args:
+        database_name (str): glue database to request.
+
+    Returns:
+        set: catalogid and locationUri
+    """
+    glue = boto3.client("glue")
+    try:
+        response = glue.get_database(Name=database_name)
+    except glue.exceptions.EntityNotFoundException:
+        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"),
+    )
+
+
+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.
+
+    Returns:
+        List[str]: List of table names
+    """
+    glue = boto3.client("glue")
+
+    # Set the paginator
+    paginator = glue.get_paginator("get_tables")
+
+    # Initialize an empty list to store the table names
+    table_names = []
+    try:
+        # Paginate through the tables
+        for page in paginator.paginate(DatabaseName=database_name, PaginationConfig={"PageSize": 100}):
+            # Add the tables from the current page to the list
+            table_names.extend([table["Name"] for table in page["TableList"] if "Name" in table])
+    except glue.exceptions.EntityNotFoundException:
+        print(f"Database {database_name} not found.")
+        return []
+    except Exception as e:
+        # todo catch all
+        print(f"Error: {e}")
+        return []
+
+    return table_names
+
+
+def get_glue_table_schema(database_name: str, table_name: str) -> List[Dict]:
+    """Get the schema of a Glue table.
+
+    Args:
+        database_name (str): Glue database name.
+        table_name (str): Glue table name.
+
+    Returns:
+        dict: Table schema
+    """
+
+    glue = boto3.client("glue")
+
+    # Get the table schema
+    try:
+        response = glue.get_table(DatabaseName=database_name, Name=table_name)
+    except glue.exceptions.EntityNotFoundException:
+        print(f"Table {table_name} not found in database {database_name}.")
+        return []
+    except Exception as e:
+        # todo catch all
+        print(f"Error: {e}")
+        return []
+
+    table_schema = response["Table"]["StorageDescriptor"]["Columns"]
+
+    # when using hive partition keys, the schema is stored in the PartitionKeys field
+    if response["Table"].get("PartitionKeys") is not None:
+        for pk in response["Table"]["PartitionKeys"]:
+            table_schema.append(
+                {
+                    "Name": pk["Name"],
+                    "Type": pk["Type"],
+                    "Hive": True,
+                    "Comment": pk.get("Comment"),
+                }
+            )
+    return table_schema
+
+
+def import_glue(
+    data_contract_specification: DataContractSpecification,
+    source: str,
+    table_names: List[str],
+) -> DataContractSpecification:
+    """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
+    if catalogid is None:
+        return data_contract_specification
+
+    if table_names is None:
+        table_names = get_glue_tables(source)
+
+    server_kwargs = {"type": "glue", "account": catalogid, "database": source}
+
+    if location_uri:
+        server_kwargs["location"] = location_uri
+
+    data_contract_specification.servers = {
+        "production": Server(**server_kwargs),
+    }
+
+    for table_name in table_names:
+        if data_contract_specification.models is None:
+            data_contract_specification.models = {}
+
+        table_schema = get_glue_table_schema(source, table_name)
+
+        fields = {}
+        for column in table_schema:
+            field = create_typed_field(column["Type"])
+
+            # hive partitions are required, but are not primary keys
+            if column.get("Hive"):
+                field.required = True
+
+            field.description = column.get("Comment")
+            fields[column["Name"]] = field
+
+        data_contract_specification.models[table_name] = Model(
+            type="table",
+            fields=fields,
+        )
+
+    return data_contract_specification
+
+
+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(" ", "")
+    # Example: array<string>
+    if dtype.startswith("array"):
+        field.type = "array"
+        field.items = create_typed_field(dtype[6:-1])
+    # Example: struct<field1:float,field2:string>
+    elif dtype.startswith("struct"):
+        field.type = "struct"
+        for f in split_struct(dtype[7:-1]):
+            field_name, field_key = f.split(":", 1)
+            field.fields[field_name] = create_typed_field(field_key)
+    # Example: map<string,int>
+    elif dtype.startswith("map"):
+        field.type = "map"
+        map_match = re.match(r"map<(.+?),\s*(.+)>", dtype)
+        if map_match:
+            key_type = map_match.group(1)
+            value_type = map_match.group(2)
+            field.keys = create_typed_field(key_type)
+            field.values = create_typed_field(value_type)
+    # Example: decimal(38, 6) or decimal
+    elif dtype.startswith("decimal"):
+        field.type = "decimal"
+        decimal_match = re.match(r"decimal\((\d+),\s*(\d+)\)", dtype)
+        if decimal_match:  # if precision specified
+            field.precision = int(decimal_match.group(1))
+            field.scale = int(decimal_match.group(2))
+    # Example: varchar(255) or varchar
+    elif dtype.startswith("varchar"):
+        field.type = "varchar"
+        if len(dtype) > 7:
+            field.maxLength = int(dtype[8:-1])
+    else:
+        field.type = map_type_from_sql(dtype)
+    return field
+
+
+def split_fields(s: str) -> Generator[str, None, None]:
+    """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
+
+    sql_type = sql_type.lower()
+
+    type_mapping = {
+        "string": "string",
+        "int": "int",
+        "bigint": "bigint",
+        "float": "float",
+        "double": "double",
+        "boolean": "boolean",
+        "timestamp": "timestamp",
+        "date": "date",
+    }
+
+    for prefix, mapped_type in type_mapping.items():
+        if sql_type.startswith(prefix):
+            return mapped_type
+
+    return "unknown"

datacontract/imports/iceberg_importer.py

@@ -0,0 +1,172 @@
+from typing import Any, Dict
+
+from pydantic import ValidationError
+from pyiceberg import types as iceberg_types
+from pyiceberg.schema import Schema
+
+from datacontract.imports.importer import Importer
+from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model
+from datacontract.model.exceptions import DataContractException
+
+
+class IcebergImporter(Importer):
+    def import_source(
+        self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
+    ) -> DataContractSpecification:
+        schema = load_and_validate_iceberg_schema(source)
+        return import_iceberg(
+            data_contract_specification,
+            schema,
+            import_args.get("iceberg_table"),
+        )
+
+
+def load_and_validate_iceberg_schema(source: str) -> Schema:
+    with open(source, "r") as file:
+        try:
+            return Schema.model_validate_json(file.read())
+        except ValidationError as e:
+            raise DataContractException(
+                type="schema",
+                name="Parse iceberg schema",
+                reason=f"Failed to validate iceberg schema from {source}: {e}",
+                engine="datacontract",
+            )
+
+
+def import_iceberg(
+    data_contract_specification: DataContractSpecification, schema: Schema, table_name: str
+) -> DataContractSpecification:
+    if data_contract_specification.models is None:
+        data_contract_specification.models = {}
+
+    model = Model(type="table", title=table_name)
+
+    # Iceberg identifier_fields aren't technically primary keys since Iceberg doesn't support primary keys,
+    # but they are close enough that we can probably treat them as primary keys on the conversion.
+    # ref: https://iceberg.apache.org/spec/#identifier-field-ids
+    # this code WILL NOT support finding nested primary key fields.
+    identifier_fields_ids = schema.identifier_field_ids
+
+    for field in schema.fields:
+        model_field = _field_from_nested_field(field)
+
+        if field.field_id in identifier_fields_ids:
+            model_field.primaryKey = True
+
+        model.fields[field.name] = model_field
+
+    data_contract_specification.models[table_name] = model
+    return data_contract_specification
+
+
+def _field_from_nested_field(nested_field: iceberg_types.NestedField) -> Field:
+    """
+    Converts an Iceberg NestedField into a Field object for the data contract.
+
+    Args:
+        nested_field: The Iceberg NestedField to convert.
+
+    Returns:
+        Field: The generated Field object.
+    """
+    field = Field(
+        title=nested_field.name,
+        required=nested_field.required,
+        config=build_field_config(nested_field),
+    )
+
+    if nested_field.doc is not None:
+        field.description = nested_field.doc
+
+    return _type_from_iceberg_type(field, nested_field.field_type)
+
+
+def _type_from_iceberg_type(field: Field, iceberg_type: iceberg_types.IcebergType) -> Field:
+    """
+    Maps Iceberg data types to the Data Contract type system and updates the field.
+
+    Args:
+        field: The Field object to update.
+        iceberg_type: The Iceberg data type to map.
+
+    Returns:
+        Field: The updated Field object.
+    """
+    field.type = _data_type_from_iceberg(iceberg_type)
+
+    if field.type == "array":
+        field.items = _type_from_iceberg_type(Field(required=iceberg_type.element_required), iceberg_type.element_type)
+
+    elif field.type == "map":
+        field.keys = _type_from_iceberg_type(Field(required=True), iceberg_type.key_type)
+        field.values = _type_from_iceberg_type(Field(required=iceberg_type.value_required), iceberg_type.value_type)
+
+    elif field.type == "object":
+        field.fields = {nf.name: _field_from_nested_field(nf) for nf in iceberg_type.fields}
+
+    return field
+
+
+def build_field_config(iceberg_field: iceberg_types.NestedField) -> Dict[str, Any]:
+    config = {}
+
+    if iceberg_field.field_id > 0:
+        config["icebergFieldId"] = iceberg_field.field_id
+
+    if iceberg_field.initial_default is not None:
+        config["icebergInitialDefault"] = iceberg_field.initial_default
+
+    if iceberg_field.write_default is not None:
+        config["icebergWriteDefault"] = iceberg_field.write_default
+
+    return config
+
+
+def _data_type_from_iceberg(type: iceberg_types.IcebergType) -> str:
+    """
+    Convert an Iceberg field type to a datacontract field type
+
+    Args:
+        type: The Iceberg field type
+
+    Returns:
+        str: The datacontract field type
+    """
+    if isinstance(type, iceberg_types.BooleanType):
+        return "boolean"
+    if isinstance(type, iceberg_types.IntegerType):
+        return "integer"
+    if isinstance(type, iceberg_types.LongType):
+        return "long"
+    if isinstance(type, iceberg_types.FloatType):
+        return "float"
+    if isinstance(type, iceberg_types.DoubleType):
+        return "double"
+    if isinstance(type, iceberg_types.DecimalType):
+        return "decimal"
+    if isinstance(type, iceberg_types.DateType):
+        return "date"
+    if isinstance(type, iceberg_types.TimeType):
+        # there isn't a great mapping for the iceberg type "time", just map to string for now
+        return "string"
+    if isinstance(type, iceberg_types.TimestampType):
+        return "timestamp_ntz"
+    if isinstance(type, iceberg_types.TimestamptzType):
+        return "timestamp_tz"
+    if isinstance(type, iceberg_types.StringType):
+        return "string"
+    if isinstance(type, iceberg_types.UUIDType):
+        return "string"
+    if isinstance(type, iceberg_types.BinaryType):
+        return "bytes"
+    if isinstance(type, iceberg_types.FixedType):
+        return "bytes"
+    if isinstance(type, iceberg_types.MapType):
+        return "map"
+    if isinstance(type, iceberg_types.ListType):
+        return "array"
+    if isinstance(type, iceberg_types.StructType):
+        return "object"
+
+    raise ValueError(f"Unknown Iceberg type: {type}")

datacontract/imports/importer.py

@@ -0,0 +1,51 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+
+from datacontract_specification.model import DataContractSpecification
+from open_data_contract_standard.model import OpenDataContractStandard
+
+
+class Importer(ABC):
+    def __init__(self, import_format) -> None:
+        self.import_format = import_format
+
+    @abstractmethod
+    def import_source(
+        self,
+        data_contract_specification: DataContractSpecification | OpenDataContractStandard,
+        source: str,
+        import_args: dict,
+    ) -> DataContractSpecification | OpenDataContractStandard:
+        pass
+
+
+class ImportFormat(str, Enum):
+    sql = "sql"
+    avro = "avro"
+    dbt = "dbt"
+    dbml = "dbml"
+    glue = "glue"
+    jsonschema = "jsonschema"
+    json = "json"
+    bigquery = "bigquery"
+    odcs = "odcs"
+    unity = "unity"
+    spark = "spark"
+    iceberg = "iceberg"
+    parquet = "parquet"
+    csv = "csv"
+    protobuf = "protobuf"
+    excel = "excel"
+
+    @classmethod
+    def get_supported_formats(cls):
+        return list(map(lambda c: c.value, cls))
+
+
+class Spec(str, Enum):
+    datacontract_specification = "datacontract_specification"
+    odcs = "odcs"
+
+    @classmethod
+    def get_supported_types(cls):
+        return list(map(lambda c: c.value, cls))

datacontract/imports/importer_factory.py

@@ -0,0 +1,128 @@
+import importlib.util
+import sys
+
+from datacontract.imports.importer import Importer, ImportFormat
+
+
+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 supported.")
+        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",
+)
+importer_factory.register_lazy_importer(
+    name=ImportFormat.spark,
+    module_path="datacontract.imports.spark_importer",
+    class_name="SparkImporter",
+)
+importer_factory.register_lazy_importer(
+    name=ImportFormat.dbt, module_path="datacontract.imports.dbt_importer", class_name="DbtManifestImporter"
+)
+importer_factory.register_lazy_importer(
+    name=ImportFormat.dbml,
+    module_path="datacontract.imports.dbml_importer",
+    class_name="DBMLImporter",
+)
+importer_factory.register_lazy_importer(
+    name=ImportFormat.iceberg,
+    module_path="datacontract.imports.iceberg_importer",
+    class_name="IcebergImporter",
+)
+importer_factory.register_lazy_importer(
+    name=ImportFormat.parquet,
+    module_path="datacontract.imports.parquet_importer",
+    class_name="ParquetImporter",
+)
+importer_factory.register_lazy_importer(
+    name=ImportFormat.csv,
+    module_path="datacontract.imports.csv_importer",
+    class_name="CsvImporter",
+)
+importer_factory.register_lazy_importer(
+    name=ImportFormat.protobuf,
+    module_path="datacontract.imports.protobuf_importer",
+    class_name="ProtoBufImporter",
+)
+importer_factory.register_lazy_importer(
+    name=ImportFormat.excel,
+    module_path="datacontract.imports.excel_importer",
+    class_name="ExcelImporter",
+)
+
+
+importer_factory.register_lazy_importer(
+    name=ImportFormat.json,
+    module_path="datacontract.imports.json_importer",
+    class_name="JsonImporter",
+)