datacontract-cli 0.10.11__py3-none-any.whl → 0.10.13__py3-none-any.whl

datacontract/imports/dbt_importer.py

@@ -1,117 +1,93 @@
 import json
-
-from typing import (
-    List,
-)
+from typing import TypedDict
 
 from datacontract.imports.importer import Importer
 from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model
+from dbt.artifacts.resources.v1.components import ColumnInfo
+from dbt.contracts.graph.manifest import Manifest
+
+
+class DBTImportArgs(TypedDict, total=False):
+    """
+    A dictionary representing arguments for importing DBT models.
+    Makes the DBT Importer more customizable by allowing for flexible filtering
+    of models and their properties, through wrapping or extending.
+
+    Attributes:
+        dbt_models: The keys of models to be used in contract. All as default.
+        resource_types: Nodes listed in resource_types are kept while importing. model as default.
+    """
+
+    dbt_nodes: list[str]
+    resource_types: list[str]
 
 
 class DbtManifestImporter(Importer):
     def import_source(
-        self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
-    ) -> dict:
-        data = read_dbt_manifest(manifest_path=source)
+        self,
+        data_contract_specification: DataContractSpecification,
+        source: str,
+        import_args: DBTImportArgs,
+    ) -> DataContractSpecification:
+        manifest = read_dbt_manifest(manifest_path=source)
         return import_dbt_manifest(
-            data_contract_specification, manifest_dict=data, dbt_models=import_args.get("dbt_model")
+            data_contract_specification=data_contract_specification,
+            manifest=manifest,
+            dbt_nodes=import_args.get("dbt_nodes", []),
+            resource_types=import_args.get("resource_types", ["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")
+def read_dbt_manifest(manifest_path: str) -> Manifest:
+    """Read a manifest from file."""
+    with open(file=manifest_path, mode="r", encoding="utf-8") as f:
+        manifest_dict: dict = json.load(f)
+    return Manifest.from_dict(manifest_dict)
 
-    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:
+def import_dbt_manifest(
+    data_contract_specification: DataContractSpecification,
+    manifest: Manifest,
+    dbt_nodes: list[str],
+    resource_types: list[str],
+) -> DataContractSpecification:
+    """
+    Extracts all relevant information from the manifest,
+    and puts it in a data contract specification.
+    """
+    data_contract_specification.info.title = manifest.metadata.project_name
+    data_contract_specification.info.dbt_version = manifest.metadata.dbt_version
+
+    data_contract_specification.models = data_contract_specification.models or {}
+    for model_contents in manifest.nodes.values():
+        # Only intressted in processing models.
+        if model_contents.resource_type not in resource_types:
+            continue
+
+        # To allow args stored in dbt_models to filter relevant models.
+        # If dbt_models is empty, use all models.
+        if dbt_nodes and model_contents.name not in dbt_nodes:
             continue
 
         dc_model = Model(
-            description=model.description,
-            tags=model.tags,
-            fields=create_fields(model.columns),
+            description=model_contents.description,
+            tags=model_contents.tags,
+            fields=create_fields(columns=model_contents.columns),
         )
 
-        data_contract_specification.models[model.name] = dc_model
+        data_contract_specification.models[model_contents.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
+def create_fields(columns: dict[str, ColumnInfo]) -> dict[str, Field]:
+    fields = {
+        column.name: Field(
+            description=column.description,
+            type=column.data_type if column.data_type else "",
+            tags=column.tags,
         )
-        fields[column.name] = field
+        for column in columns.values()
+    }
 
     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

datacontract/imports/glue_importer.py

@@ -1,6 +1,6 @@
 import boto3
-from typing import List
-
+from typing import List, Dict, Generator
+import re
 from datacontract.imports.importer import Importer
 from datacontract.model.data_contract_specification import (
     DataContractSpecification,
@@ -13,7 +13,7 @@ from datacontract.model.data_contract_specification import (
 class GlueImporter(Importer):
     def import_source(
         self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
-    ) -> dict:
+    ) -> DataContractSpecification:
         return import_glue(data_contract_specification, source, import_args.get("glue_table"))
 
 
@@ -39,7 +39,7 @@ def get_glue_database(database_name: str):
 
     return (
         response["Database"]["CatalogId"],
-        response["Database"].get("LocationUri", "None"),
+        response["Database"].get("LocationUri"),
     )
 
 
@@ -75,7 +75,7 @@ def get_glue_tables(database_name: str) -> List[str]:
     return table_names
 
 
-def get_glue_table_schema(database_name: str, table_name: str):
+def get_glue_table_schema(database_name: str, table_name: str) -> List[Dict]:
     """Get the schema of a Glue table.
 
     Args:
@@ -93,11 +93,11 @@ def get_glue_table_schema(database_name: str, table_name: str):
         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 {}
+        return []
     except Exception as e:
         # todo catch all
         print(f"Error: {e}")
-        return {}
+        return []
 
     table_schema = response["Table"]["StorageDescriptor"]["Columns"]
 
@@ -109,10 +109,9 @@ def get_glue_table_schema(database_name: str, table_name: str):
                     "Name": pk["Name"],
                     "Type": pk["Type"],
                     "Hive": True,
-                    "Comment": "Partition Key",
+                    "Comment": pk.get("Comment"),
                 }
             )
-
     return table_schema
 
 
@@ -120,7 +119,7 @@ def import_glue(
     data_contract_specification: DataContractSpecification,
     source: str,
     table_names: List[str],
-):
+) -> DataContractSpecification:
     """Import the schema of a Glue database.
 
     Args:
@@ -140,8 +139,13 @@ def import_glue(
     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(type="glue", account=catalogid, database=source, location=location_uri),
+        "production": Server(**server_kwargs),
     }
 
     for table_name in table_names:
@@ -161,12 +165,6 @@ def import_glue(
             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(",")
-                field.precision = int(perc_scale[0])
-                field.scale = int(perc_scale[1])
-
         data_contract_specification.models[table_name] = Model(
             type="table",
             fields=fields,
@@ -186,27 +184,43 @@ def create_typed_field(dtype: str) -> Field:
     """
     field = Field()
     dtype = dtype.strip().lower().replace(" ", "")
-    if dtype.startswith(("array", "struct", "map")):
-        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])
-        elif dtype.startswith("map"):
-            field.type = "map"
-            key_type = orig_dtype[4:-1].split(",", 1)[0]
-            value_type = orig_dtype[4:-1].split(",", 1)[1]
+    # 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):
+def split_fields(s: str) -> Generator[str, None, None]:
     """Split a string of fields considering nested structures.
 
     Args:
@@ -253,30 +267,20 @@ def map_type_from_sql(sql_type: str) -> str:
         return None
 
     sql_type = sql_type.lower()
-    if sql_type.startswith("varchar"):
-        return "varchar"
-    if sql_type.startswith("string"):
-        return "string"
-    if sql_type.startswith("text"):
-        return "text"
-    if sql_type.startswith("byte"):
-        return "byte"
-    if sql_type.startswith("short"):
-        return "short"
-    if sql_type.startswith("integer") or sql_type.startswith("int"):
-        return "integer"
-    if sql_type.startswith("long") or sql_type.startswith("bigint"):
-        return "long"
-    if sql_type.startswith("float"):
-        return "float"
-    if sql_type.startswith("double"):
-        return "double"
-    if sql_type.startswith("boolean"):
-        return "boolean"
-    if sql_type.startswith("timestamp"):
-        return "timestamp"
-    if sql_type.startswith("date"):
-        return "date"
-    if sql_type.startswith("decimal"):
-        return "decimal"
-    return "variant"
+
+    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/importer.py

@@ -14,7 +14,7 @@ class Importer(ABC):
         data_contract_specification: DataContractSpecification,
         source: str,
         import_args: dict,
-    ) -> dict:
+    ) -> DataContractSpecification:
         pass
 
 
@@ -22,6 +22,7 @@ class ImportFormat(str, Enum):
     sql = "sql"
     avro = "avro"
     dbt = "dbt"
+    dbml = "dbml"
     glue = "glue"
     jsonschema = "jsonschema"
     bigquery = "bigquery"

datacontract/imports/importer_factory.py

@@ -88,3 +88,8 @@ importer_factory.register_lazy_importer(
 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",
+)

datacontract/imports/odcs_importer.py

@@ -46,7 +46,7 @@ DATACONTRACT_TYPES = [
 class OdcsImporter(Importer):
     def import_source(
         self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
-    ) -> dict:
+    ) -> DataContractSpecification:
         return import_odcs(data_contract_specification, source)
 
 

datacontract/imports/spark_importer.py

@@ -14,7 +14,7 @@ class SparkImporter(Importer):
         data_contract_specification: DataContractSpecification,
         source: str,
         import_args: dict,
-    ) -> dict:
+    ) -> DataContractSpecification:
         """
         Imports data from a Spark source into the data contract specification.
 
@@ -63,12 +63,12 @@ def import_from_spark_df(df: DataFrame) -> Model:
     schema = df.schema
 
     for field in schema:
-        model.fields[field.name] = _field_from_spark(field)
+        model.fields[field.name] = _field_from_struct_type(field)
 
     return model
 
 
-def _field_from_spark(spark_field: types.StructField) -> Field:
+def _field_from_struct_type(spark_field: types.StructField) -> Field:
     """
     Converts a Spark StructField into a Field object for the data contract.
 
@@ -76,18 +76,37 @@ def _field_from_spark(spark_field: types.StructField) -> Field:
         spark_field: The Spark StructField to convert.
 
     Returns:
-        Field: The corresponding Field object.
+        Field: The generated Field object.
     """
-    field_type = _data_type_from_spark(spark_field.dataType)
     field = Field()
-    field.type = field_type
     field.required = not spark_field.nullable
+    field.description = spark_field.metadata.get("comment")
 
-    if field_type == "array":
-        field.items = _field_from_spark(spark_field.dataType.elementType)
+    return _type_from_data_type(field, spark_field.dataType)
 
-    if field_type == "struct":
-        field.fields = {sf.name: _field_from_spark(sf) for sf in spark_field.dataType.fields}
+
+def _type_from_data_type(field: Field, spark_type: types.DataType) -> Field:
+    """
+    Maps Spark data types to the Data Contract type system and updates the field.
+
+    Args:
+        field: The Field object to update.
+        spark_type: The Spark data type to map.
+
+    Returns:
+        Field: The updated Field object.
+    """
+    field.type = _data_type_from_spark(spark_type)
+
+    if field.type == "array":
+        field.items = _type_from_data_type(Field(required=not spark_type.containsNull), spark_type.elementType)
+
+    elif field.type == "map":
+        field.keys = _type_from_data_type(Field(required=True), spark_type.keyType)
+        field.values = _type_from_data_type(Field(required=not spark_type.valueContainsNull), spark_type.valueType)
+
+    elif field.type == "struct":
+        field.fields = {sf.name: _field_from_struct_type(sf) for sf in spark_type.fields}
 
     return field
 
@@ -104,7 +123,7 @@ def _data_type_from_spark(spark_type: types.DataType) -> str:
     """
     if isinstance(spark_type, types.StringType):
         return "string"
-    elif isinstance(spark_type, types.IntegerType):
+    elif isinstance(spark_type, (types.IntegerType, types.ShortType)):
         return "integer"
     elif isinstance(spark_type, types.LongType):
         return "long"
@@ -116,6 +135,8 @@ def _data_type_from_spark(spark_type: types.DataType) -> str:
         return "struct"
     elif isinstance(spark_type, types.ArrayType):
         return "array"
+    elif isinstance(spark_type, types.MapType):
+        return "map"
     elif isinstance(spark_type, types.TimestampType):
         return "timestamp"
     elif isinstance(spark_type, types.TimestampNTZType):
@@ -130,5 +151,7 @@ def _data_type_from_spark(spark_type: types.DataType) -> str:
         return "decimal"
     elif isinstance(spark_type, types.NullType):
         return "null"
+    elif isinstance(spark_type, types.VarcharType):
+        return "varchar"
     else:
         raise ValueError(f"Unsupported Spark type: {spark_type}")

datacontract/imports/sql_importer.py

@@ -7,7 +7,7 @@ from datacontract.model.data_contract_specification import DataContractSpecifica
 class SqlImporter(Importer):
     def import_source(
         self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
-    ) -> dict:
+    ) -> DataContractSpecification:
         return import_sql(data_contract_specification, self.import_format, source)