datacontract-cli 0.10.23__py3-none-any.whl → 0.10.40__py3-none-any.whl

datacontract/imports/avro_importer.py

@@ -55,7 +55,6 @@ def import_avro(data_contract_specification: DataContractSpecification, source:
             engine="datacontract",
             original_exception=e,
         )
-
     # type record is being used for both the table and the object types in data contract
     # -> CONSTRAINT: one table per .avsc input, all nested records are interpreted as objects
     fields = import_record_fields(avro_schema.fields)
@@ -92,6 +91,20 @@ def handle_config_avro_custom_properties(field: avro.schema.Field, imported_fiel
         imported_field.config["avroDefault"] = field.default
 
 
+LOGICAL_TYPE_MAPPING = {
+    "decimal": "decimal",
+    "date": "date",
+    "time-millis": "time",
+    "time-micros": "time",
+    "timestamp-millis": "timestamp_tz",
+    "timestamp-micros": "timestamp_tz",
+    "local-timestamp-micros": "timestamp_ntz",
+    "local-timestamp-millis": "timestamp_ntz",
+    "duration": "string",
+    "uuid": "string",
+}
+
+
 def import_record_fields(record_fields: List[avro.schema.Field]) -> Dict[str, Field]:
     """
     Import Avro record fields and convert them to data contract fields.
@@ -117,13 +130,23 @@ def import_record_fields(record_fields: List[avro.schema.Field]) -> Dict[str, Fi
             imported_field.fields = import_record_fields(field.type.fields)
         elif field.type.type == "union":
             imported_field.required = False
-            type = import_type_of_optional_field(field)
-            imported_field.type = type
-            if type == "record":
-                imported_field.fields = import_record_fields(get_record_from_union_field(field).fields)
-            elif type == "array":
-                imported_field.type = "array"
-                imported_field.items = import_avro_array_items(get_array_from_union_field(field))
+            # Check for enum in union first, since it needs special handling
+            enum_schema = get_enum_from_union_field(field)
+            if enum_schema:
+                imported_field.type = "string"
+                imported_field.enum = enum_schema.symbols
+                imported_field.title = enum_schema.name
+                if not imported_field.config:
+                    imported_field.config = {}
+                imported_field.config["avroType"] = "enum"
+            else:
+                type = import_type_of_optional_field(field)
+                imported_field.type = type
+                if type == "record":
+                    imported_field.fields = import_record_fields(get_record_from_union_field(field).fields)
+                elif type == "array":
+                    imported_field.type = "array"
+                    imported_field.items = import_avro_array_items(get_array_from_union_field(field))
         elif field.type.type == "array":
             imported_field.type = "array"
             imported_field.items = import_avro_array_items(field.type)
@@ -137,9 +160,15 @@ def import_record_fields(record_fields: List[avro.schema.Field]) -> Dict[str, Fi
             if not imported_field.config:
                 imported_field.config = {}
             imported_field.config["avroType"] = "enum"
-        else:  # primitive type
-            imported_field.type = map_type_from_avro(field.type.type)
-
+        else:
+            logical_type = field.type.get_prop("logicalType")
+            if logical_type in LOGICAL_TYPE_MAPPING:
+                imported_field.type = LOGICAL_TYPE_MAPPING[logical_type]
+                if logical_type == "decimal":
+                    imported_field.precision = field.type.precision
+                    imported_field.scale = field.type.scale
+            else:
+                imported_field.type = map_type_from_avro(field.type.type)
         imported_fields[field.name] = imported_field
 
     return imported_fields
@@ -212,7 +241,11 @@ def import_type_of_optional_field(field: avro.schema.Field) -> str:
     """
     for field_type in field.type.schemas:
         if field_type.type != "null":
-            return map_type_from_avro(field_type.type)
+            logical_type = field_type.get_prop("logicalType")
+            if logical_type and logical_type in LOGICAL_TYPE_MAPPING:
+                return LOGICAL_TYPE_MAPPING[logical_type]
+            else:
+                return map_type_from_avro(field_type.type)
     raise DataContractException(
         type="schema",
         result="failed",
@@ -254,6 +287,22 @@ def get_array_from_union_field(field: avro.schema.Field) -> avro.schema.ArraySch
     return None
 
 
+def get_enum_from_union_field(field: avro.schema.Field) -> avro.schema.EnumSchema | None:
+    """
+    Get the enum schema from a union field.
+
+    Args:
+        field: The Avro field with a union type.
+
+    Returns:
+        The enum schema if found, None otherwise.
+    """
+    for field_type in field.type.schemas:
+        if field_type.type == "enum":
+            return field_type
+    return None
+
+
 def map_type_from_avro(avro_type_str: str) -> str:
     """
     Map Avro type strings to data contract type strings.
@@ -276,6 +325,8 @@ def map_type_from_avro(avro_type_str: str) -> str:
         return "binary"
     elif avro_type_str == "double":
         return "double"
+    elif avro_type_str == "float":
+        return "float"
     elif avro_type_str == "int":
         return "int"
     elif avro_type_str == "long":

datacontract/imports/csv_importer.py

@@ -1,89 +1,143 @@
 import os
+from typing import Any, Dict, List
 
-import clevercsv
+import duckdb
 
 from datacontract.imports.importer import Importer
-from datacontract.model.data_contract_specification import DataContractSpecification, Example, Field, Model, Server
+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, self.import_format, source)
+        return import_csv(data_contract_specification, source)
 
 
-def import_csv(data_contract_specification: DataContractSpecification, format: str, source: str):
-    include_example = False
-
-    # detect encoding and dialect
-    encoding = clevercsv.encoding.get_encoding(source)
-    with open(source, "r", newline="") as fp:
-        dialect = clevercsv.Sniffer().sniff(fp.read(10000))
-
-    # using auto detecting of the format and encoding
-    df = clevercsv.read_dataframe(source)
-
-    if data_contract_specification.models is None:
-        data_contract_specification.models = {}
-
+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=dialect.delimiter
+        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 column, dtype in df.dtypes.items():
-        field = Field()
-        field.type = map_type_from_pandas(dtype.name)
-        fields[column] = field
+    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=f"Csv file with encoding {encoding}",
-        fields=fields,
+        type="table", description="Generated model of " + source, fields=fields, examples=model_examples
     )
 
-    # multiline data is not correctly handled by yaml dump
-    if include_example:
-        if data_contract_specification.examples is None:
-            data_contract_specification.examples = []
-
-        # read first 10 lines with the detected encoding
-        with open(source, "r", encoding=encoding) as csvfile:
-            lines = csvfile.readlines()[:10]
-
-        data_contract_specification.examples.append(Example(type="csv", model=table_name, data="".join(lines)))
-
     return data_contract_specification
 
 
-def map_type_from_pandas(sql_type: str):
+_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.lower().strip()
-
-    if sql_type_normed == "object":
-        return "string"
-    elif sql_type_normed.startswith("str"):
-        return "string"
-    elif sql_type_normed.startswith("int"):
-        return "integer"
-    elif sql_type_normed.startswith("float"):
-        return "float"
-    elif sql_type_normed.startswith("bool"):
-        return "boolean"
-    elif sql_type_normed.startswith("timestamp"):
-        return "timestamp"
-    elif sql_type_normed == "datetime64":
-        return "date"
-    elif sql_type_normed == "timedelta[ns]":
-        return "timestamp_ntz"
-    else:
-        return "variant"
+    sql_type_normed = sql_type.upper().strip()
+    return _duck_db_types.get(sql_type_normed, "string")