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

datacontract/export/protobuf_converter.py

@@ -4,102 +4,196 @@ from datacontract.model.data_contract_specification import DataContractSpecifica
 
 class ProtoBufExporter(Exporter):
     def export(self, data_contract, model, server, sql_server_type, export_args) -> dict:
-        return to_protobuf(data_contract)
+        # Returns a dict containing the protobuf representation.
+        proto = to_protobuf(data_contract)
+        return {"protobuf": proto}
 
 
-def to_protobuf(data_contract_spec: DataContractSpecification):
+def to_protobuf(data_contract_spec: DataContractSpecification) -> str:
+    """
+    Generates a Protobuf file from the data contract specification.
+    Scans all models for enum fields (even if the type is "string") by checking for a "values" property.
+    """
     messages = ""
+    enum_definitions = {}
+
+    # Iterate over all models to generate messages and collect enum definitions.
     for model_name, model in data_contract_spec.models.items():
-        messages += to_protobuf_message(model_name, model.fields, model.description, 0)
+        for field_name, field in model.fields.items():
+            # If the field has enum values, collect them.
+            if _is_enum_field(field):
+                enum_name = _get_enum_name(field, field_name)
+                enum_values = _get_enum_values(field)
+                if enum_values and enum_name not in enum_definitions:
+                    enum_definitions[enum_name] = enum_values
+
+        messages += to_protobuf_message(model_name, model.fields, getattr(model, "description", ""), 0)
         messages += "\n"
 
-    result = f"""syntax = "proto3";
-
-{messages}
-"""
-
-    return result
-
-
-def _to_protobuf_message_name(model_name):
-    return model_name[0].upper() + model_name[1:]
-
-
-def to_protobuf_message(model_name, fields, description, indent_level: int = 0):
+    # Build header with syntax and package declarations.
+    header = 'syntax = "proto3";\n\n'
+    package = getattr(data_contract_spec, "package", "example")
+    header += f"package {package};\n\n"
+
+    # Append enum definitions.
+    for enum_name, enum_values in enum_definitions.items():
+        header += f"// Enum for {enum_name}\n"
+        header += f"enum {enum_name} {{\n"
+        # Only iterate if enum_values is a dictionary.
+        if isinstance(enum_values, dict):
+            for enum_const, value in sorted(enum_values.items(), key=lambda item: item[1]):
+                normalized_const = enum_const.upper().replace(" ", "_")
+                header += f"  {normalized_const} = {value};\n"
+        else:
+            header += f"  // Warning: Enum values for {enum_name} are not a dictionary\n"
+        header += "}\n\n"
+    return header + messages
+
+
+def _is_enum_field(field) -> bool:
+    """
+    Returns True if the field (dict or object) has a non-empty "values" property.
+    """
+    if isinstance(field, dict):
+        return bool(field.get("values"))
+    return bool(getattr(field, "values", None))
+
+
+def _get_enum_name(field, field_name: str) -> str:
+    """
+    Returns the enum name either from the field's "enum_name" or derived from the field name.
+    """
+    if isinstance(field, dict):
+        return field.get("enum_name", _to_protobuf_message_name(field_name))
+    return getattr(field, "enum_name", None) or _to_protobuf_message_name(field_name)
+
+
+def _get_enum_values(field) -> dict:
+    """
+    Returns the enum values from the field.
+    If the values are not a dictionary, attempts to extract enum attributes.
+    """
+    if isinstance(field, dict):
+        values = field.get("values", {})
+    else:
+        values = getattr(field, "values", {})
+
+    if not isinstance(values, dict):
+        # If values is a BaseModel (or similar) with a .dict() method, use it.
+        if hasattr(values, "dict") and callable(values.dict):
+            values_dict = values.dict()
+            return {k: v for k, v in values_dict.items() if k.isupper() and isinstance(v, int)}
+        else:
+            # Otherwise, iterate over attributes that look like enums.
+            return {
+                key: getattr(values, key)
+                for key in dir(values)
+                if key.isupper() and isinstance(getattr(values, key), int)
+            }
+    return values
+
+
+def _to_protobuf_message_name(name: str) -> str:
+    """
+    Returns a valid Protobuf message/enum name by capitalizing the first letter.
+    """
+    return name[0].upper() + name[1:] if name else name
+
+
+def to_protobuf_message(model_name: str, fields: dict, description: str, indent_level: int = 0) -> str:
+    """
+    Generates a Protobuf message definition from the model's fields.
+    Handles nested messages for complex types.
+    """
     result = ""
+    if description:
+        result += f"{indent(indent_level)}// {description}\n"
 
-    if description is not None:
-        result += f"""{indent(indent_level)}/* {description} */\n"""
-
-    fields_protobuf = ""
+    result += f"message {_to_protobuf_message_name(model_name)} {{\n"
     number = 1
     for field_name, field in fields.items():
-        if field.type in ["object", "record", "struct"]:
-            fields_protobuf += (
-                "\n".join(
-                    map(
-                        lambda x: "  " + x,
-                        to_protobuf_message(field_name, field.fields, field.description, indent_level + 1).splitlines(),
-                    )
-                )
-                + "\n"
-            )
-
-        fields_protobuf += to_protobuf_field(field_name, field, field.description, number, 1) + "\n"
+        # For nested objects, generate a nested message.
+        field_type = _get_field_type(field)
+        if field_type in ["object", "record", "struct"]:
+            nested_desc = field.get("description", "") if isinstance(field, dict) else getattr(field, "description", "")
+            nested_fields = field.get("fields", {}) if isinstance(field, dict) else field.fields
+            nested_message = to_protobuf_message(field_name, nested_fields, nested_desc, indent_level + 1)
+            result += nested_message + "\n"
+
+        field_desc = field.get("description", "") if isinstance(field, dict) else getattr(field, "description", "")
+        result += to_protobuf_field(field_name, field, field_desc, number, indent_level + 1) + "\n"
         number += 1
-    result += f"message {_to_protobuf_message_name(model_name)} {{\n{fields_protobuf}}}\n"
 
+    result += f"{indent(indent_level)}}}\n"
     return result
 
 
-def to_protobuf_field(field_name, field, description, number: int, indent_level: int = 0):
-    optional = ""
-    if not field.required:
-        optional = "optional "
-
+def to_protobuf_field(field_name: str, field, description: str, number: int, indent_level: int = 0) -> str:
+    """
+    Generates a field definition within a Protobuf message.
+    """
     result = ""
-
-    if description is not None:
-        result += f"""{indent(indent_level)}/* {description} */\n"""
-
-    result += f"{indent(indent_level)}{optional}{_convert_type(field_name, field)} {field_name} = {number};"
-
+    if description:
+        result += f"{indent(indent_level)}// {description}\n"
+    result += f"{indent(indent_level)}{_convert_type(field_name, field)} {field_name} = {number};"
     return result
 
 
-def indent(indent_level):
+def indent(indent_level: int) -> str:
     return "  " * indent_level
 
 
-def _convert_type(field_name, field) -> None | str:
-    type = field.type
-    if type is None:
-        return None
-    if type.lower() in ["string", "varchar", "text"]:
-        return "string"
-    if type.lower() in ["timestamp", "timestamp_tz"]:
-        return "string"
-    if type.lower() in ["timestamp_ntz"]:
-        return "string"
-    if type.lower() in ["date"]:
+def _get_field_type(field) -> str:
+    """
+    Retrieves the field type from the field definition.
+    """
+    if isinstance(field, dict):
+        return field.get("type", "").lower()
+    return getattr(field, "type", "").lower()
+
+
+def _convert_type(field_name: str, field) -> str:
+    """
+    Converts a field's type (from the data contract) to a Protobuf type.
+    Prioritizes enum conversion if a non-empty "values" property exists.
+    """
+    # For debugging purposes
+    print("Converting field:", field_name)
+    # If the field should be treated as an enum, return its enum name.
+    if _is_enum_field(field):
+        return _get_enum_name(field, field_name)
+
+    lower_type = _get_field_type(field)
+    if lower_type in ["string", "varchar", "text"]:
         return "string"
-    if type.lower() in ["time"]:
+    if lower_type in ["timestamp", "timestamp_tz", "timestamp_ntz", "date", "time"]:
         return "string"
-    if type.lower() in ["number", "decimal", "numeric"]:
+    if lower_type in ["number", "decimal", "numeric"]:
         return "double"
-    if type.lower() in ["float", "double"]:
-        return type.lower()
-    if type.lower() in ["integer", "int"]:
+    if lower_type in ["float", "double"]:
+        return lower_type
+    if lower_type in ["integer", "int"]:
         return "int32"
-    if type.lower() in ["long", "bigint"]:
+    if lower_type in ["long", "bigint"]:
         return "int64"
-    if type.lower() in ["boolean"]:
+    if lower_type in ["boolean"]:
         return "bool"
-    if type.lower() in ["bytes"]:
+    if lower_type in ["bytes"]:
         return "bytes"
-    if type.lower() in ["object", "record", "struct"]:
+    if lower_type in ["object", "record", "struct"]:
         return _to_protobuf_message_name(field_name)
-    if type.lower() in ["array"]:
-        # TODO spec is missing arrays
-        return "repeated string"
-    return None
+    if lower_type == "array":
+        # Handle array types. Check for an "items" property.
+        items = field.get("items") if isinstance(field, dict) else getattr(field, "items", None)
+        if items and isinstance(items, dict) and items.get("type"):
+            item_type = items.get("type", "").lower()
+            if item_type in ["object", "record", "struct"]:
+                # Singularize the field name (a simple approach).
+                singular = field_name[:-1] if field_name.endswith("s") else field_name
+                return "repeated " + _to_protobuf_message_name(singular)
+            else:
+                return "repeated " + _convert_type(field_name, items)
+        else:
+            return "repeated string"
+    # Fallback for unrecognized types.
+    return "string"

datacontract/export/spark_converter.py

@@ -175,7 +175,7 @@ def print_schema(dtype: types.DataType) -> str:
         Returns:
             str: The indented text.
         """
-        return "\n".join([f'{"    " * level}{line}' for line in text.split("\n")])
+        return "\n".join([f"{'    ' * level}{line}" for line in text.split("\n")])
 
     def repr_column(column: types.StructField) -> str:
         """

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.
@@ -137,9 +150,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 +231,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",
@@ -276,6 +299,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")