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

datacontract/imports/csv_importer.py

@@ -0,0 +1,143 @@
+import os
+from typing import Any, Dict, List
+
+import duckdb
+
+from datacontract.imports.importer import Importer
+from datacontract.model.data_contract_specification import DataContractSpecification, Model, Server
+
+
+class CsvImporter(Importer):
+    def import_source(
+        self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
+    ) -> DataContractSpecification:
+        return import_csv(data_contract_specification, source)
+
+
+def import_csv(
+    data_contract_specification: DataContractSpecification, source: str, include_examples: bool = False
+) -> DataContractSpecification:
+    # use the file name as table name
+    table_name = os.path.splitext(os.path.basename(source))[0]
+
+    # use duckdb to auto detect format, columns, etc.
+    con = duckdb.connect(database=":memory:")
+    con.sql(
+        f"""CREATE VIEW "{table_name}" AS SELECT * FROM read_csv_auto('{source}', hive_partitioning=1, auto_type_candidates = ['BOOLEAN', 'INTEGER', 'BIGINT', 'DOUBLE', 'VARCHAR']);"""
+    )
+    dialect = con.sql(f"SELECT * FROM sniff_csv('{source}', sample_size = 1000);").fetchnumpy()
+    tbl = con.table(table_name)
+
+    if data_contract_specification.servers is None:
+        data_contract_specification.servers = {}
+
+    delimiter = None if dialect is None else dialect["Delimiter"][0]
+
+    if dialect is not None:
+        dc_types = [map_type_from_duckdb(x["type"]) for x in dialect["Columns"][0]]
+    else:
+        dc_types = [map_type_from_duckdb(str(x)) for x in tbl.dtypes]
+
+    data_contract_specification.servers["production"] = Server(
+        type="local", path=source, format="csv", delimiter=delimiter
+    )
+
+    rowcount = tbl.shape[0]
+
+    tallies = dict()
+    for row in tbl.describe().fetchall():
+        if row[0] not in ["count", "max", "min"]:
+            continue
+        for i in range(tbl.shape[1]):
+            tallies[(row[0], tbl.columns[i])] = row[i + 1] if row[0] != "count" else int(row[i + 1])
+
+    samples: Dict[str, List] = dict()
+    for i in range(tbl.shape[1]):
+        field_name = tbl.columns[i]
+        if tallies[("count", field_name)] > 0 and tbl.dtypes[i] not in ["BOOLEAN", "BLOB"]:
+            sql = f"""SELECT DISTINCT "{field_name}" FROM "{table_name}" WHERE "{field_name}" IS NOT NULL USING SAMPLE 5 ROWS;"""
+            samples[field_name] = [x[0] for x in con.sql(sql).fetchall()]
+
+    formats: Dict[str, str] = dict()
+    for i in range(tbl.shape[1]):
+        field_name = tbl.columns[i]
+        if tallies[("count", field_name)] > 0 and tbl.dtypes[i] == "VARCHAR":
+            sql = f"""SELECT
+              count_if("{field_name}" IS NOT NULL) as count,
+              count_if(regexp_matches("{field_name}", '^[\\w-\\.]+@([\\w-]+\\.)+[\\w-]{{2,4}}$')) as email,
+              count_if(regexp_matches("{field_name}", '^[[a-z0-9]{{8}}-?[a-z0-9]{{4}}-?[a-z0-9]{{4}}-?[a-z0-9]{{4}}-?[a-z0-9]{{12}}]')) as uuid
+              FROM "{table_name}";
+              """
+            res = con.sql(sql).fetchone()
+            if res[1] == res[0]:
+                formats[field_name] = "email"
+            elif res[2] == res[0]:
+                formats[field_name] = "uuid"
+
+    fields = {}
+    for i in range(tbl.shape[1]):
+        field_name = tbl.columns[i]
+        dc_type = dc_types[i]
+
+        ## specifying "integer" rather than "bigint" looks nicer
+        if (
+            dc_type == "bigint"
+            and tallies[("max", field_name)] <= 2147483647
+            and tallies[("min", field_name)] >= -2147483648
+        ):
+            dc_type = "integer"
+
+        field: Dict[str, Any] = {"type": dc_type, "format": formats.get(field_name, None)}
+
+        if tallies[("count", field_name)] == rowcount:
+            field["required"] = True
+        if dc_type not in ["boolean", "bytes"]:
+            distinct_values = tbl.count(f'DISTINCT "{field_name}"').fetchone()[0]  # type: ignore
+            if distinct_values > 0 and distinct_values == tallies[("count", field_name)]:
+                field["unique"] = True
+        s = samples.get(field_name, None)
+        if s is not None:
+            field["examples"] = s
+        if dc_type in ["integer", "bigint", "float", "double"]:
+            field["minimum"] = tallies[("min", field_name)]
+            field["maximum"] = tallies[("max", field_name)]
+
+        fields[field_name] = field
+
+    model_examples = None
+    if include_examples:
+        model_examples = con.sql(f"""SELECT DISTINCT * FROM "{table_name}" USING SAMPLE 5 ROWS;""").fetchall()
+
+    data_contract_specification.models[table_name] = Model(
+        type="table", description="Generated model of " + source, fields=fields, examples=model_examples
+    )
+
+    return data_contract_specification
+
+
+_duck_db_types = {
+    "BOOLEAN": "boolean",
+    "BLOB": "bytes",
+    "TINYINT": "integer",
+    "SMALLINT": "integer",
+    "INTEGER": "integer",
+    "BIGINT": "bigint",
+    "UTINYINT": "integer",
+    "USMALLINT": "integer",
+    "UINTEGER": "integer",
+    "UBIGINT": "bigint",
+    "FLOAT": "float",
+    "DOUBLE": "double",
+    "VARCHAR": "string",
+    "TIMESTAMP": "timestamp",
+    "DATE": "date",
+    # TODO: Add support for NULL
+}
+
+
+def map_type_from_duckdb(sql_type: None | str):
+    if sql_type is None:
+        return None
+
+    sql_type_normed = sql_type.upper().strip()
+    return _duck_db_types.get(sql_type_normed, "string")

datacontract/imports/dbml_importer.py

@@ -0,0 +1,112 @@
+from typing import List
+
+from pydbml import Database, PyDBML
+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, Field, Model
+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].primaryKey = 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

@@ -0,0 +1,240 @@
+import json
+from typing import TypedDict
+
+from dbt.artifacts.resources.v1.components import ColumnInfo
+from dbt.contracts.graph.manifest import Manifest
+from dbt.contracts.graph.nodes import GenericTestNode, ManifestNode, ModelNode
+from dbt_common.contracts.constraints import ConstraintType
+
+from datacontract.imports.bigquery_importer import map_type_from_bigquery
+from datacontract.imports.importer import Importer
+from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model
+
+
+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: DBTImportArgs,
+    ) -> DataContractSpecification:
+        manifest = read_dbt_manifest(manifest_path=source)
+        return import_dbt_manifest(
+            data_contract_specification=data_contract_specification,
+            manifest=manifest,
+            dbt_nodes=import_args.get("dbt_model", []),
+            resource_types=import_args.get("resource_types", ["model"]),
+        )
+
+
+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)
+    manifest = Manifest.from_dict(manifest_dict)
+    manifest.build_parent_and_child_maps()
+    return manifest
+
+
+def _get_primary_keys(manifest: Manifest, node: ManifestNode) -> list[str]:
+    node_unique_id = node.unique_id
+    if isinstance(node, ModelNode):
+        test_nodes = []
+        for node_id in manifest.child_map.get(node_unique_id, []):
+            test_node = manifest.nodes.get(node_id)
+            if not test_node or test_node.resource_type != "test":
+                continue
+            if not isinstance(test_node, GenericTestNode):
+                continue
+            if test_node.config.where is not None:
+                continue
+            test_nodes.append(test_node)
+        return node.infer_primary_key(test_nodes)
+    return []
+
+
+def _get_references(manifest: Manifest, node: ManifestNode) -> dict[str, str]:
+    node_unique_id = node.unique_id
+    references = {}
+    for node_id in manifest.child_map.get(node_unique_id, []):
+        test_node = manifest.nodes.get(node_id)
+        if not test_node or test_node.resource_type != "test":
+            continue
+        if not isinstance(test_node, GenericTestNode):
+            continue
+        if test_node.test_metadata.name != "relationships":
+            continue
+        if test_node.config.where is not None:
+            continue
+        if test_node.attached_node != node_unique_id:
+            continue
+        relationship_target_node_id = [n for n in test_node.depends_on.nodes if n != node_unique_id][0]
+        relationship_target_node = manifest.nodes.get(relationship_target_node_id)
+        references[f"{node.name}.{test_node.column_name}"] = (
+            f"""{relationship_target_node.name}.{test_node.test_metadata.kwargs["field"]}"""
+        )
+    return references
+
+
+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
+    adapter_type = manifest.metadata.adapter_type
+    data_contract_specification.models = data_contract_specification.models or {}
+    for node in manifest.nodes.values():
+        # Only intressted in processing models.
+        if node.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 node.name not in dbt_nodes:
+            continue
+
+        model_unique_id = node.unique_id
+        primary_keys = _get_primary_keys(manifest, node)
+        references = _get_references(manifest, node)
+
+        primary_key = None
+        if len(primary_keys) == 1:
+            primary_key = primary_keys[0]
+
+        dc_model = Model(
+            description=node.description,
+            tags=node.tags,
+            fields=create_fields(
+                manifest,
+                model_unique_id=model_unique_id,
+                columns=node.columns,
+                primary_key_name=primary_key,
+                references=references,
+                adapter_type=adapter_type,
+            ),
+        )
+        if len(primary_keys) > 1:
+            dc_model.primaryKey = primary_keys
+
+        data_contract_specification.models[node.name] = dc_model
+
+    return data_contract_specification
+
+
+def convert_data_type_by_adapter_type(data_type: str, adapter_type: str) -> str:
+    if adapter_type == "bigquery":
+        return map_type_from_bigquery(data_type)
+    return data_type
+
+
+def create_fields(
+    manifest: Manifest,
+    model_unique_id: str,
+    columns: dict[str, ColumnInfo],
+    primary_key_name: str,
+    references: dict[str, str],
+    adapter_type: str,
+) -> dict[str, Field]:
+    fields = {
+        column.name: create_field(manifest, model_unique_id, column, primary_key_name, references, adapter_type)
+        for column in columns.values()
+    }
+    return fields
+
+
+def get_column_tests(manifest: Manifest, model_name: str, column_name: str) -> list[dict[str, str]]:
+    column_tests = []
+    model_node = manifest.nodes.get(model_name)
+    if not model_node:
+        raise ValueError(f"Model {model_name} not found in manifest.")
+
+    model_unique_id = model_node.unique_id
+    test_ids = manifest.child_map.get(model_unique_id, [])
+
+    for test_id in test_ids:
+        test_node = manifest.nodes.get(test_id)
+        if not test_node or test_node.resource_type != "test":
+            continue
+
+        if not isinstance(test_node, GenericTestNode):
+            continue
+
+        if test_node.column_name != column_name:
+            continue
+
+        if test_node.config.where is not None:
+            continue
+
+        column_tests.append(
+            {
+                "test_name": test_node.name,
+                "test_type": test_node.test_metadata.name,
+                "column": test_node.column_name,
+            }
+        )
+    return column_tests
+
+
+def create_field(
+    manifest: Manifest,
+    model_unique_id: str,
+    column: ColumnInfo,
+    primary_key_name: str,
+    references: dict[str, str],
+    adapter_type: str,
+) -> Field:
+    column_type = convert_data_type_by_adapter_type(column.data_type, adapter_type) if column.data_type else ""
+    field = Field(
+        description=column.description,
+        type=column_type,
+        tags=column.tags,
+    )
+
+    all_tests = get_column_tests(manifest, model_unique_id, column.name)
+
+    required = False
+    if any(constraint.type == ConstraintType.not_null for constraint in column.constraints):
+        required = True
+    if [test for test in all_tests if test["test_type"] == "not_null"]:
+        required = True
+    if required:
+        field.required = required
+
+    unique = False
+    if any(constraint.type == ConstraintType.unique for constraint in column.constraints):
+        unique = True
+    if [test for test in all_tests if test["test_type"] == "unique"]:
+        unique = True
+    if unique:
+        field.unique = unique
+
+    if column.name == primary_key_name:
+        field.primaryKey = True
+
+    references_key = f"{manifest.nodes[model_unique_id].name}.{column.name}"
+    if references_key in references:
+        field.references = references[references_key]
+
+    return field