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

datacontract/export/sqlalchemy_converter.py

@@ -0,0 +1,169 @@
+import ast
+import typing
+
+import datacontract.model.data_contract_specification as spec
+from datacontract.export.exporter import Exporter
+from datacontract.export.exporter import _determine_sql_server_type
+
+
+class SQLAlchemyExporter(Exporter):
+    def export(
+        self, data_contract: spec.DataContractSpecification, model, server, sql_server_type, export_args
+    ) -> dict:
+        sql_server_type = _determine_sql_server_type(data_contract, sql_server_type, server)
+        return to_sqlalchemy_model_str(data_contract, sql_server_type, server)
+
+
+DECLARATIVE_BASE = "Base"
+
+
+def to_sqlalchemy_model_str(contract: spec.DataContractSpecification, sql_server_type: str = "", server=None) -> str:
+    server_obj = contract.servers.get(server)
+    classdefs = [
+        generate_model_class(model_name, model, server_obj, sql_server_type)
+        for (model_name, model) in contract.models.items()
+    ]
+    documentation = (
+        [ast.Expr(ast.Constant(contract.info.description))] if (contract.info and contract.info.description) else []
+    )
+
+    declarative_base = ast.ClassDef(
+        name=DECLARATIVE_BASE,
+        bases=[ast.Name(id="DeclarativeBase", ctx=ast.Load())],
+        body=[ast.Pass()],
+        keywords=[],
+        decorator_list=[],
+    )
+
+    databricks_timestamp = ast.ImportFrom(
+        module="databricks.sqlalchemy", names=[ast.alias("TIMESTAMP"), ast.alias("TIMESTAMP_NTZ")]
+    )
+    timestamp = ast.ImportFrom(module="sqlalchemy", names=[ast.alias(name="TIMESTAMP")])
+    result = ast.Module(
+        body=[
+            ast.ImportFrom(module="sqlalchemy.orm", names=[ast.alias(name="DeclarativeBase")]),
+            ast.ImportFrom(
+                module="sqlalchemy",
+                names=[
+                    ast.alias("Column"),
+                    ast.alias("Date"),
+                    ast.alias("Integer"),
+                    ast.alias("Numeric"),
+                    ast.alias("String"),
+                    ast.alias("Text"),
+                    ast.alias("VARCHAR"),
+                    ast.alias("BigInteger"),
+                    ast.alias("Float"),
+                    ast.alias("Double"),
+                    ast.alias("Boolean"),
+                    ast.alias("Date"),
+                    ast.alias("ARRAY"),
+                    ast.alias("LargeBinary"),
+                ],
+            ),
+            databricks_timestamp if sql_server_type == "databricks" else timestamp,
+            *documentation,
+            declarative_base,
+            *classdefs,
+        ],
+        type_ignores=[],
+    )
+    return ast.unparse(result)
+
+
+def Call(name, *args, **kwargs) -> ast.Call:
+    return ast.Call(
+        ast.Name(name),
+        args=[v for v in args],
+        keywords=[ast.keyword(arg=f"{k}", value=ast.Constant(v)) for (k, v) in kwargs.items()],
+    )
+
+
+def Column(predicate, **kwargs) -> ast.Call:
+    return Call("Column", predicate, **kwargs)
+
+
+def sqlalchemy_primitive(field: spec.Field):
+    sqlalchemy_name = {
+        "string": Call("String", ast.Constant(field.maxLength)),
+        "text": Call("Text", ast.Constant(field.maxLength)),
+        "varchar": Call("VARCHAR", ast.Constant(field.maxLength)),
+        "number": Call("Numeric", ast.Constant(field.precision), ast.Constant(field.scale)),
+        "decimal": Call("Numeric", ast.Constant(field.precision), ast.Constant(field.scale)),
+        "numeric": Call("Numeric", ast.Constant(field.precision), ast.Constant(field.scale)),
+        "int": ast.Name("Integer"),
+        "integer": ast.Name("Integer"),
+        "long": ast.Name("BigInteger"),
+        "bigint": ast.Name("BigInteger"),
+        "float": ast.Name("Float"),
+        "double": ast.Name("Double"),
+        "boolean": ast.Name("Boolean"),
+        "timestamp": ast.Name("TIMESTAMP"),
+        "timestamp_tz": Call("TIMESTAMP", ast.Constant(True)),
+        "timestamp_ntz": ast.Name("TIMESTAMP_NTZ"),
+        "date": ast.Name("Date"),
+        "bytes": Call("LargeBinary", ast.Constant(field.maxLength)),
+    }
+    return sqlalchemy_name.get(field.type)
+
+
+def constant_field_value(field_name: str, field: spec.Field) -> tuple[ast.Call, typing.Optional[ast.ClassDef]]:
+    new_type = sqlalchemy_primitive(field)
+    match field.type:
+        case "array":
+            new_type = Call("ARRAY", sqlalchemy_primitive(field.items))
+    if new_type is None:
+        raise RuntimeError(f"Unsupported field type {field.type}.")
+
+    return Column(new_type, nullable=not field.required, comment=field.description, primary_key=field.primary), None
+
+
+def column_assignment(field_name: str, field: spec.Field) -> tuple[ast.Call, typing.Optional[ast.ClassDef]]:
+    return constant_field_value(field_name, field)
+
+
+def is_simple_field(field: spec.Field) -> bool:
+    return field.type not in set(["object", "record", "struct"])
+
+
+def field_definitions(fields: dict[str, spec.Field]) -> tuple[list[ast.Expr], list[ast.ClassDef]]:
+    annotations: list[ast.Expr] = []
+    classes: list[typing.Any] = []
+    for field_name, field in fields.items():
+        (ann, new_class) = column_assignment(field_name, field)
+        annotations.append(ast.Assign(targets=[ast.Name(id=field_name, ctx=ast.Store())], value=ann, lineno=0))
+    return (annotations, classes)
+
+
+def generate_model_class(
+    name: str, model_definition: spec.Model, server=None, sql_server_type: str = ""
+) -> ast.ClassDef:
+    (field_assignments, nested_classes) = field_definitions(model_definition.fields)
+    documentation = [ast.Expr(ast.Constant(model_definition.description))] if model_definition.description else []
+
+    schema = None if server is None else server.schema_
+    table_name = ast.Constant(name)
+    if sql_server_type == "databricks":
+        table_name = ast.Constant(name.lower())
+
+    result = ast.ClassDef(
+        name=name.capitalize(),
+        bases=[ast.Name(id=DECLARATIVE_BASE, ctx=ast.Load())],
+        body=[
+            *documentation,
+            ast.Assign(targets=[ast.Name("__tablename__")], value=table_name, lineno=0),
+            ast.Assign(
+                targets=[ast.Name("__table_args__")],
+                value=ast.Dict(
+                    keys=[ast.Constant("comment"), ast.Constant("schema")],
+                    values=[ast.Constant(model_definition.description), ast.Constant(schema)],
+                ),
+                lineno=0,
+            ),
+            *nested_classes,
+            *field_assignments,
+        ],
+        keywords=[],
+        decorator_list=[],
+    )
+    return result

datacontract/imports/avro_importer.py

@@ -133,6 +133,7 @@ def import_record_fields(record_fields: List[avro.schema.Field]) -> Dict[str, Fi
         elif field.type.type == "enum":
             imported_field.type = "string"
             imported_field.enum = field.type.symbols
+            imported_field.title = field.type.name
             if not imported_field.config:
                 imported_field.config = {}
             imported_field.config["avroType"] = "enum"

datacontract/imports/bigquery_importer.py

@@ -10,13 +10,13 @@ from datacontract.model.exceptions import DataContractException
 class BigQueryImporter(Importer):
     def import_source(
         self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
-    ) -> dict:
+    ) -> DataContractSpecification:
         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_table"),
                 import_args.get("bigquery_project"),
                 import_args.get("bigquery_dataset"),
             )

datacontract/imports/dbml_importer.py

@@ -0,0 +1,112 @@
+from pydbml import PyDBML, Database
+from typing import List
+
+from pyparsing import ParseException
+
+from datacontract.imports.importer import Importer
+from datacontract.imports.sql_importer import map_type_from_sql
+from datacontract.model.data_contract_specification import DataContractSpecification, Model, Field
+from datacontract.model.exceptions import DataContractException
+
+
+class DBMLImporter(Importer):
+    def import_source(
+        self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
+    ) -> DataContractSpecification:
+        data_contract_specification = import_dbml_from_source(
+            data_contract_specification,
+            source,
+            import_args.get("dbml_schema"),
+            import_args.get("dbml_table"),
+        )
+        return data_contract_specification
+
+
+def import_dbml_from_source(
+    data_contract_specification: DataContractSpecification,
+    source: str,
+    import_schemas: List[str],
+    import_tables: List[str],
+) -> DataContractSpecification:
+    try:
+        with open(source, "r") as file:
+            dbml_schema = PyDBML(file)
+    except ParseException as e:
+        raise DataContractException(
+            type="schema",
+            name="Parse DBML schema",
+            reason=f"Failed to parse DBML schema from {source}",
+            engine="datacontract",
+            original_exception=e,
+        )
+
+    return convert_dbml(data_contract_specification, dbml_schema, import_schemas, import_tables)
+
+
+def convert_dbml(
+    data_contract_specification: DataContractSpecification,
+    dbml_schema: Database,
+    import_schemas: List[str],
+    import_tables: List[str],
+) -> DataContractSpecification:
+    if dbml_schema.project is not None:
+        data_contract_specification.info.title = dbml_schema.project.name
+
+    if data_contract_specification.models is None:
+        data_contract_specification.models = {}
+
+    for table in dbml_schema.tables:
+        schema_name = table.schema
+        table_name = table.name
+
+        # Skip if import schemas or table names are defined
+        # and the current table doesn't match
+        # if empty no filtering is done
+        if import_schemas and schema_name not in import_schemas:
+            continue
+
+        if import_tables and table_name not in import_tables:
+            continue
+
+        fields = import_table_fields(table, dbml_schema.refs)
+
+        data_contract_specification.models[table_name] = Model(
+            fields=fields, namespace=schema_name, description=table.note.text
+        )
+
+    return data_contract_specification
+
+
+def import_table_fields(table, references) -> dict[str, Field]:
+    imported_fields = {}
+    for field in table.columns:
+        field_name = field.name
+        imported_fields[field_name] = Field()
+        imported_fields[field_name].required = field.not_null
+        imported_fields[field_name].description = field.note.text
+        imported_fields[field_name].primary = field.pk
+        imported_fields[field_name].unique = field.unique
+        # This is an assumption, that these might be valid SQL Types, since
+        # DBML doesn't really enforce anything other than 'no spaces' in column types
+        imported_fields[field_name].type = map_type_from_sql(field.type)
+
+        ref = get_reference(field, references)
+        if ref is not None:
+            imported_fields[field_name].references = ref
+
+    return imported_fields
+
+
+def get_reference(field, references):
+    result = None
+    for ref in references:
+        ref_table_name = ref.col1[0].table.name
+        ref_col_name = ref.col1[0].name
+        field_table_name = field.table.name
+        field_name = field.name
+
+        if ref_table_name == field_table_name and ref_col_name == field_name:
+            result = f"{ref.col2[0].table.name}.{ref.col2[0].name}"
+            return result
+
+    return result

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