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

datacontract/imports/avro_importer.py

@@ -1,11 +1,46 @@
+from typing import Dict, List
+
 import avro.schema
 
-from datacontract.model.data_contract_specification import \
-    DataContractSpecification, Model, Field
+from datacontract.imports.importer import Importer
+from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model
 from datacontract.model.exceptions import DataContractException
 
 
+class AvroImporter(Importer):
+    """Class to import Avro Schema file"""
+
+    def import_source(
+        self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
+    ) -> DataContractSpecification:
+        """
+        Import Avro schema from a source file.
+
+        Args:
+            data_contract_specification: The data contract specification to update.
+            source: The path to the Avro schema file.
+            import_args: Additional import arguments.
+
+        Returns:
+            The updated data contract specification.
+        """
+        return import_avro(data_contract_specification, source)
+
+
 def import_avro(data_contract_specification: DataContractSpecification, source: str) -> DataContractSpecification:
+    """
+    Import an Avro schema from a file and update the data contract specification.
+
+    Args:
+        data_contract_specification: The data contract specification to update.
+        source: The path to the Avro schema file.
+
+    Returns:
+        DataContractSpecification: The updated data contract specification.
+
+    Raises:
+        DataContractException: If there's an error parsing the Avro schema.
+    """
     if data_contract_specification.models is None:
         data_contract_specification.models = {}
 
@@ -20,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)
@@ -38,35 +72,118 @@ def import_avro(data_contract_specification: DataContractSpecification, source:
     return data_contract_specification
 
 
-def import_record_fields(record_fields):
+def handle_config_avro_custom_properties(field: avro.schema.Field, imported_field: Field) -> None:
+    """
+    Handle custom Avro properties and add them to the imported field's config.
+
+    Args:
+        field: The Avro field.
+        imported_field: The imported field to update.
+    """
+    if field.get_prop("logicalType") is not None:
+        if imported_field.config is None:
+            imported_field.config = {}
+        imported_field.config["avroLogicalType"] = field.get_prop("logicalType")
+
+    if field.default is not None:
+        if imported_field.config is None:
+            imported_field.config = {}
+        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.
+
+    Args:
+        record_fields: List of Avro record fields.
+
+    Returns:
+        A dictionary of imported fields.
+    """
     imported_fields = {}
     for field in record_fields:
-        imported_fields[field.name] = Field()
-        imported_fields[field.name].required = True
-        imported_fields[field.name].description = field.doc
-        for prop in field.other_props:
-            imported_fields[field.name].__setattr__(prop, field.other_props[prop])
+        imported_field = Field()
+        imported_field.required = True
+        imported_field.description = field.doc
+
+        handle_config_avro_custom_properties(field, imported_field)
 
+        # Determine field type and handle nested structures
         if field.type.type == "record":
-            imported_fields[field.name].type = "object"
-            imported_fields[field.name].description = field.type.doc
-            imported_fields[field.name].fields = import_record_fields(field.type.fields)
+            imported_field.type = "object"
+            imported_field.description = field.type.doc
+            imported_field.fields = import_record_fields(field.type.fields)
         elif field.type.type == "union":
-            imported_fields[field.name].required = False
-            type = import_type_of_optional_field(field)
-            imported_fields[field.name].type = type
-            if type == "record":
-                imported_fields[field.name].fields = import_record_fields(get_record_from_union_field(field).fields)
+            imported_field.required = False
+            # 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_fields[field.name].type = "array"
-            imported_fields[field.name].items = import_avro_array_items(field.type)
-        else:  # primitive type
-            imported_fields[field.name].type = map_type_from_avro(field.type.type)
+            imported_field.type = "array"
+            imported_field.items = import_avro_array_items(field.type)
+        elif field.type.type == "map":
+            imported_field.type = "map"
+            imported_field.values = import_avro_map_values(field.type)
+        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"
+        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
 
 
-def import_avro_array_items(array_schema):
+def import_avro_array_items(array_schema: avro.schema.ArraySchema) -> Field:
+    """
+    Import Avro array items and convert them to a data contract field.
+
+    Args:
+        array_schema: The Avro array schema.
+
+    Returns:
+        Field: The imported field representing the array items.
+    """
     items = Field()
     for prop in array_schema.other_props:
         items.__setattr__(prop, array_schema.other_props[prop])
@@ -83,10 +200,52 @@ def import_avro_array_items(array_schema):
     return items
 
 
-def import_type_of_optional_field(field):
+def import_avro_map_values(map_schema: avro.schema.MapSchema) -> Field:
+    """
+    Import Avro map values and convert them to a data contract field.
+
+    Args:
+        map_schema: The Avro map schema.
+
+    Returns:
+        Field: The imported field representing the map values.
+    """
+    values = Field()
+    for prop in map_schema.other_props:
+        values.__setattr__(prop, map_schema.other_props[prop])
+
+    if map_schema.values.type == "record":
+        values.type = "object"
+        values.fields = import_record_fields(map_schema.values.fields)
+    elif map_schema.values.type == "array":
+        values.type = "array"
+        values.items = import_avro_array_items(map_schema.values)
+    else:  # primitive type
+        values.type = map_type_from_avro(map_schema.values.type)
+
+    return values
+
+
+def import_type_of_optional_field(field: avro.schema.Field) -> str:
+    """
+    Determine the type of optional field in an Avro union.
+
+    Args:
+        field: The Avro field with a union type.
+
+    Returns:
+        str: The mapped type of the non-null field in the union.
+
+    Raises:
+        DataContractException: If no non-null type is found in the union.
+    """
     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",
@@ -96,14 +255,67 @@ def import_type_of_optional_field(field):
     )
 
 
-def get_record_from_union_field(field):
+def get_record_from_union_field(field: avro.schema.Field) -> avro.schema.RecordSchema | None:
+    """
+    Get the record schema from a union field.
+
+    Args:
+        field: The Avro field with a union type.
+
+    Returns:
+        The record schema if found, None otherwise.
+    """
     for field_type in field.type.schemas:
         if field_type.type == "record":
             return field_type
     return None
 
 
-def map_type_from_avro(avro_type_str: str):
+def get_array_from_union_field(field: avro.schema.Field) -> avro.schema.ArraySchema | None:
+    """
+    Get the array schema from a union field.
+
+    Args:
+        field: The Avro field with a union type.
+
+    Returns:
+        The array schema if found, None otherwise.
+    """
+    for field_type in field.type.schemas:
+        if field_type.type == "array":
+            return field_type
+    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.
+
+    Args:
+        avro_type_str (str): The Avro type string.
+
+    Returns:
+        str: The corresponding data contract type string.
+
+    Raises:
+        DataContractException: If the Avro type is unsupported.
+    """
     # TODO: ambiguous mapping in the export
     if avro_type_str == "null":
         return "null"
@@ -113,6 +325,8 @@ def map_type_from_avro(avro_type_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":
@@ -121,6 +335,12 @@ def map_type_from_avro(avro_type_str: str):
         return "boolean"
     elif avro_type_str == "record":
         return "record"
+    elif avro_type_str == "array":
+        return "array"
+    elif avro_type_str == "map":
+        return "map"
+    elif avro_type_str == "enum":
+        return "string"
     else:
         raise DataContractException(
             type="schema",

datacontract/imports/bigquery_importer.py

@@ -0,0 +1,221 @@
+import json
+import logging
+from typing import List
+
+from datacontract.imports.importer import Importer
+from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model
+from datacontract.model.exceptions import DataContractException
+
+
+class BigQueryImporter(Importer):
+    def import_source(
+        self, data_contract_specification: DataContractSpecification, source: str, import_args: 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_table"),
+                import_args.get("bigquery_project"),
+                import_args.get("bigquery_dataset"),
+            )
+        return data_contract_specification
+
+
+def import_bigquery_from_json(
+    data_contract_specification: DataContractSpecification, source: str
+) -> DataContractSpecification:
+    try:
+        with open(source, "r") as file:
+            bigquery_schema = json.loads(file.read())
+    except json.JSONDecodeError as e:
+        raise DataContractException(
+            type="schema",
+            name="Parse bigquery schema",
+            reason=f"Failed to parse bigquery schema from {source}",
+            engine="datacontract",
+            original_exception=e,
+        )
+    return convert_bigquery_schema(data_contract_specification, bigquery_schema)
+
+
+def import_bigquery_from_api(
+    data_contract_specification: DataContractSpecification,
+    bigquery_tables: List[str],
+    bigquery_project: str,
+    bigquery_dataset: str,
+) -> DataContractSpecification:
+    try:
+        from google.cloud import bigquery
+    except ImportError as e:
+        raise DataContractException(
+            type="schema",
+            result="failed",
+            name="bigquery extra missing",
+            reason="Install the extra datacontract-cli[bigquery] to use bigquery",
+            engine="datacontract",
+            original_exception=e,
+        )
+
+    client = bigquery.Client(project=bigquery_project)
+
+    if bigquery_tables is None:
+        bigquery_tables = fetch_table_names(client, bigquery_dataset)
+
+    for table in bigquery_tables:
+        try:
+            api_table = client.get_table("{}.{}.{}".format(bigquery_project, bigquery_dataset, table))
+
+        except ValueError as e:
+            raise DataContractException(
+                type="schema",
+                result="failed",
+                name="Invalid table name for bigquery API",
+                reason=f"Tablename {table} is invalid for the bigquery API",
+                original_exception=e,
+                engine="datacontract",
+            )
+
+        if api_table is None:
+            raise DataContractException(
+                type="request",
+                result="failed",
+                name="Query bigtable Schema from API",
+                reason=f"Table {table} bnot found on bigtable schema Project {bigquery_project}, dataset {bigquery_dataset}.",
+                engine="datacontract",
+            )
+
+        convert_bigquery_schema(data_contract_specification, api_table.to_api_repr())
+
+    return data_contract_specification
+
+
+def fetch_table_names(client, dataset: str) -> List[str]:
+    table_names = []
+    api_tables = client.list_tables(dataset)
+    for api_table in api_tables:
+        table_names.append(api_table.table_id)
+
+    return table_names
+
+
+def convert_bigquery_schema(
+    data_contract_specification: DataContractSpecification, bigquery_schema: dict
+) -> DataContractSpecification:
+    if data_contract_specification.models is None:
+        data_contract_specification.models = {}
+
+    fields = import_table_fields(bigquery_schema.get("schema").get("fields"))
+
+    # Looking at actual export data, I guess this is always set and friendlyName isn't, though I couldn't say
+    # what exactly leads to friendlyName being set
+    table_id = bigquery_schema.get("tableReference").get("tableId")
+
+    data_contract_specification.models[table_id] = Model(
+        fields=fields, type=map_bigquery_type(bigquery_schema.get("type"))
+    )
+
+    # Copy the description, if it exists
+    if bigquery_schema.get("description") is not None:
+        data_contract_specification.models[table_id].description = bigquery_schema.get("description")
+
+    # Set the title from friendlyName if it exists
+    if bigquery_schema.get("friendlyName") is not None:
+        data_contract_specification.models[table_id].title = bigquery_schema.get("friendlyName")
+
+    return data_contract_specification
+
+
+def import_table_fields(table_fields):
+    imported_fields = {}
+    for field in table_fields:
+        field_name = field.get("name")
+        imported_fields[field_name] = Field()
+        imported_fields[field_name].required = field.get("mode") == "REQUIRED"
+        imported_fields[field_name].description = field.get("description")
+
+        if field.get("type") == "RECORD":
+            imported_fields[field_name].type = "object"
+            imported_fields[field_name].fields = import_table_fields(field.get("fields"))
+        elif field.get("type") == "STRUCT":
+            imported_fields[field_name].type = "struct"
+            imported_fields[field_name].fields = import_table_fields(field.get("fields"))
+        elif field.get("type") == "RANGE":
+            # This is a range of date/datetime/timestamp but multiple values
+            # So we map it to an array
+            imported_fields[field_name].type = "array"
+            imported_fields[field_name].items = Field(
+                type=map_type_from_bigquery(field["rangeElementType"].get("type"))
+            )
+        else:  # primitive type
+            imported_fields[field_name].type = map_type_from_bigquery(field.get("type"))
+
+        if field.get("type") == "STRING":
+            # in bigquery both string and bytes have maxLength but in the datacontracts
+            # spec it is only valid for strings
+            if field.get("maxLength") is not None:
+                imported_fields[field_name].maxLength = int(field.get("maxLength"))
+
+        if field.get("type") == "NUMERIC" or field.get("type") == "BIGNUMERIC":
+            if field.get("precision") is not None:
+                imported_fields[field_name].precision = int(field.get("precision"))
+
+            if field.get("scale") is not None:
+                imported_fields[field_name].scale = int(field.get("scale"))
+
+    return imported_fields
+
+
+def map_type_from_bigquery(bigquery_type_str: str):
+    if bigquery_type_str == "STRING":
+        return "string"
+    elif bigquery_type_str == "BYTES":
+        return "bytes"
+    elif bigquery_type_str == "INTEGER":
+        return "int"
+    elif bigquery_type_str == "INT64":
+        return "bigint"
+    elif bigquery_type_str == "FLOAT":
+        return "float"
+    elif bigquery_type_str == "FLOAT64":
+        return "double"
+    elif bigquery_type_str == "BOOLEAN" or bigquery_type_str == "BOOL":
+        return "boolean"
+    elif bigquery_type_str == "TIMESTAMP":
+        return "timestamp"
+    elif bigquery_type_str == "DATE":
+        return "date"
+    elif bigquery_type_str == "TIME":
+        return "timestamp_ntz"
+    elif bigquery_type_str == "DATETIME":
+        return "timestamp"
+    elif bigquery_type_str == "NUMERIC":
+        return "numeric"
+    elif bigquery_type_str == "BIGNUMERIC":
+        return "double"
+    elif bigquery_type_str == "GEOGRAPHY":
+        return "object"
+    elif bigquery_type_str == "JSON":
+        return "object"
+    else:
+        raise DataContractException(
+            type="schema",
+            result="failed",
+            name="Map bigquery type to data contract type",
+            reason=f"Unsupported type {bigquery_type_str} in bigquery json definition.",
+            engine="datacontract",
+        )
+
+
+def map_bigquery_type(bigquery_type: str) -> str:
+    if bigquery_type == "TABLE" or bigquery_type == "EXTERNAL" or bigquery_type == "SNAPSHOT":
+        return "table"
+    elif bigquery_type == "VIEW" or bigquery_type == "MATERIALIZED_VIEW":
+        return "view"
+    else:
+        logger = logging.getLogger(__name__)
+        logger.info(
+            f"Can't properly map bigquery table type '{bigquery_type}' to datacontracts model types. Mapping it to table."
+        )
+        return "table"