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

datacontract/imports/protobuf_importer.py

@@ -0,0 +1,264 @@
+import os
+import re
+import tempfile
+
+from google.protobuf import descriptor_pb2
+from grpc_tools import protoc
+
+from datacontract.imports.importer import Importer
+from datacontract.model.data_contract_specification import DataContractSpecification
+from datacontract.model.exceptions import DataContractException
+
+
+def map_type_from_protobuf(field_type: int):
+    protobuf_type_mapping = {
+        1: "double",
+        2: "float",
+        3: "long",
+        4: "long",  # uint64 mapped to long
+        5: "integer",  # int32 mapped to integer
+        6: "string",  # fixed64 mapped to string
+        7: "string",  # fixed32 mapped to string
+        8: "boolean",
+        9: "string",
+        12: "bytes",
+        13: "integer",  # uint32 mapped to integer
+        15: "integer",  # sfixed32 mapped to integer
+        16: "long",  # sfixed64 mapped to long
+        17: "integer",  # sint32 mapped to integer
+        18: "long",  # sint64 mapped to long
+    }
+    return protobuf_type_mapping.get(field_type, "string")
+
+
+def parse_imports(proto_file: str) -> list:
+    """
+    Parse import statements from a .proto file and return a list of imported file paths.
+    """
+    try:
+        with open(proto_file, "r") as f:
+            content = f.read()
+    except Exception as e:
+        raise DataContractException(
+            type="file",
+            name="Parse proto imports",
+            reason=f"Failed to read proto file: {proto_file}",
+            engine="datacontract",
+            original_exception=e,
+        )
+    imported_files = re.findall(r'import\s+"(.+?)";', content)
+    proto_dir = os.path.dirname(proto_file)
+    return [os.path.join(proto_dir, imp) for imp in imported_files]
+
+
+def compile_proto_to_binary(proto_files: list, output_file: str):
+    """
+    Compile the provided proto files into a single descriptor set using grpc_tools.protoc.
+    """
+    proto_dirs = set(os.path.dirname(proto) for proto in proto_files)
+    proto_paths = [f"--proto_path={d}" for d in proto_dirs]
+
+    args = [""] + proto_paths + [f"--descriptor_set_out={output_file}"] + proto_files
+    ret = protoc.main(args)
+    if ret != 0:
+        raise DataContractException(
+            type="schema",
+            name="Compile proto files",
+            reason=f"grpc_tools.protoc failed with exit code {ret}",
+            engine="datacontract",
+            original_exception=None,
+        )
+
+
+def extract_enum_values_from_fds(fds: descriptor_pb2.FileDescriptorSet, enum_name: str) -> dict:
+    """
+    Search the FileDescriptorSet for an enum definition with the given name
+    and return a dictionary of its values (name to number).
+    """
+    for file_descriptor in fds.file:
+        # Check top-level enums.
+        for enum in file_descriptor.enum_type:
+            if enum.name == enum_name:
+                return {value.name: value.number for value in enum.value}
+        # Check enums defined inside messages.
+        for message in file_descriptor.message_type:
+            for enum in message.enum_type:
+                if enum.name == enum_name:
+                    return {value.name: value.number for value in enum.value}
+    return {}
+
+
+def extract_message_fields_from_fds(fds: descriptor_pb2.FileDescriptorSet, message_name: str) -> dict:
+    """
+    Given a FileDescriptorSet and a message name, return a dict with its field definitions.
+    This function recurses for nested messages and handles enums.
+    """
+    for file_descriptor in fds.file:
+        for msg in file_descriptor.message_type:
+            if msg.name == message_name:
+                fields = {}
+                for field in msg.field:
+                    if field.type == 11:  # TYPE_MESSAGE
+                        nested_msg_name = field.type_name.split(".")[-1]
+                        nested_fields = extract_message_fields_from_fds(fds, nested_msg_name)
+                        if field.label == 3:  # repeated field
+                            field_info = {
+                                "description": f"List of {nested_msg_name}",
+                                "type": "array",
+                                "items": {"type": "object", "fields": nested_fields},
+                            }
+                        else:
+                            field_info = {
+                                "description": f"Nested object of {nested_msg_name}",
+                                "type": "object",
+                                "fields": nested_fields,
+                            }
+                    elif field.type == 14:  # TYPE_ENUM
+                        enum_name = field.type_name.split(".")[-1]
+                        enum_values = extract_enum_values_from_fds(fds, enum_name)
+                        field_info = {
+                            "description": f"Enum field {field.name}",
+                            "type": "string",
+                            "values": enum_values,
+                            "required": (field.label == 2),
+                        }
+                    else:
+                        field_info = {
+                            "description": f"Field {field.name}",
+                            "type": map_type_from_protobuf(field.type),
+                            "required": (field.label == 2),
+                        }
+                    fields[field.name] = field_info
+                return fields
+    return {}
+
+
+def import_protobuf(
+    data_contract_specification: DataContractSpecification, sources: list, import_args: dict = None
+) -> DataContractSpecification:
+    """
+    Gather all proto files (including those imported), compile them into one descriptor,
+    then generate models with nested fields and enums resolved.
+
+    The generated data contract uses generic defaults instead of specific hardcoded ones.
+    """
+
+    # --- Step 1: Gather all proto files (main and imported)
+    proto_files_set = set()
+    queue = list(sources)
+    while queue:
+        proto = queue.pop(0)
+        if proto not in proto_files_set:
+            proto_files_set.add(proto)
+            for imp in parse_imports(proto):
+                if os.path.exists(imp) and imp not in proto_files_set:
+                    queue.append(imp)
+    all_proto_files = list(proto_files_set)
+
+    # --- Step 2: Compile all proto files into a single descriptor set.
+    temp_descriptor = tempfile.NamedTemporaryFile(suffix=".pb", delete=False)
+    descriptor_file = temp_descriptor.name
+    temp_descriptor.close()  # Allow protoc to write to the file
+    try:
+        compile_proto_to_binary(all_proto_files, descriptor_file)
+
+        with open(descriptor_file, "rb") as f:
+            proto_data = f.read()
+        fds = descriptor_pb2.FileDescriptorSet()
+        try:
+            fds.ParseFromString(proto_data)
+        except Exception as e:
+            raise DataContractException(
+                type="schema",
+                name="Parse descriptor set",
+                reason="Failed to parse descriptor set from compiled proto files",
+                engine="datacontract",
+                original_exception=e,
+            )
+
+        # --- Step 3: Build models from the descriptor set.
+        all_models = {}
+        # Create a set of the main proto file basenames.
+        source_proto_basenames = {os.path.basename(proto) for proto in sources}
+
+        for file_descriptor in fds.file:
+            # Only process file descriptors that correspond to your main proto files.
+            if os.path.basename(file_descriptor.name) not in source_proto_basenames:
+                continue
+
+            for message in file_descriptor.message_type:
+                fields = {}
+                for field in message.field:
+                    if field.type == 11:  # TYPE_MESSAGE
+                        nested_msg_name = field.type_name.split(".")[-1]
+                        nested_fields = extract_message_fields_from_fds(fds, nested_msg_name)
+                        if field.label == 3:
+                            field_info = {
+                                "description": f"List of {nested_msg_name}",
+                                "type": "array",
+                                "items": {"type": "object", "fields": nested_fields},
+                            }
+                        else:
+                            field_info = {
+                                "description": f"Nested object of {nested_msg_name}",
+                                "type": "object",
+                                "fields": nested_fields,
+                            }
+                        fields[field.name] = field_info
+                    elif field.type == 14:  # TYPE_ENUM
+                        enum_name = field.type_name.split(".")[-1]
+                        enum_values = extract_enum_values_from_fds(fds, enum_name)
+                        field_info = {
+                            "description": f"Enum field {field.name}",
+                            "type": "string",
+                            "values": enum_values,
+                            "required": (field.label == 2),
+                        }
+                        fields[field.name] = field_info
+                    else:
+                        field_info = {
+                            "description": f"Field {field.name}",
+                            "type": map_type_from_protobuf(field.type),
+                            "required": (field.label == 2),
+                        }
+                        fields[field.name] = field_info
+
+                all_models[message.name] = {
+                    "description": f"Details of {message.name}.",
+                    "type": "table",
+                    "fields": fields,
+                }
+
+        data_contract_specification.models = all_models
+
+        return data_contract_specification
+    finally:
+        # Clean up the temporary descriptor file.
+        if os.path.exists(descriptor_file):
+            os.remove(descriptor_file)
+
+
+class ProtoBufImporter(Importer):
+    def __init__(self, name):
+        # 'name' is passed by the importer factory.
+        self.name = name
+
+    def import_source(
+        self,
+        data_contract_specification: DataContractSpecification,
+        source: str,
+        import_args: dict = None,
+    ) -> DataContractSpecification:
+        """
+        Import a protobuf file (and its imports) into the given DataContractSpecification.
+
+        Parameters:
+          - data_contract_specification: the initial specification to update.
+          - source: the protobuf file path.
+          - import_args: optional dictionary with additional arguments (e.g. 'output_dir').
+
+        Returns:
+          The updated DataContractSpecification.
+        """
+        # Wrap the source in a list because import_protobuf expects a list of sources.
+        return import_protobuf(data_contract_specification, [source], import_args)

datacontract/imports/spark_importer.py

@@ -1,3 +1,8 @@
+import atexit
+import logging
+import tempfile
+
+from databricks.sdk import WorkspaceClient
 from pyspark.sql import DataFrame, SparkSession, types
 
 from datacontract.imports.importer import Importer
@@ -8,6 +13,8 @@ from datacontract.model.data_contract_specification import (
     Server,
 )
 
+logger = logging.getLogger(__name__)
+
 
 class SparkImporter(Importer):
     def import_source(
@@ -23,39 +30,69 @@ class SparkImporter(Importer):
             data_contract_specification: The data contract specification object.
             source: The source string indicating the Spark tables to read.
             import_args: Additional arguments for the import process.
-
         Returns:
             dict: The updated data contract specification.
         """
-        return import_spark(data_contract_specification, source)
+        dataframe = import_args.get("dataframe", None)
+        description = import_args.get("description", None)
+        return import_spark(data_contract_specification, source, dataframe, description)
 
 
-def import_spark(data_contract_specification: DataContractSpecification, source: str) -> DataContractSpecification:
+def import_spark(
+    data_contract_specification: DataContractSpecification,
+    source: str,
+    dataframe: DataFrame | None = None,
+    description: str | None = None,
+) -> DataContractSpecification:
     """
-    Reads Spark tables and updates the data contract specification with their schemas.
+    Imports schema(s) from Spark into a Data Contract Specification.
 
     Args:
-        data_contract_specification: The data contract specification to update.
-        source: A comma-separated string of Spark temporary views to read.
+        data_contract_specification (DataContractSpecification): The contract spec to update.
+        source (str): Comma-separated Spark table/view names.
+        dataframe (DataFrame | None): Optional Spark DataFrame to import.
+        description (str | None): Optional table-level description.
 
     Returns:
-        DataContractSpecification: The updated data contract specification.
+        DataContractSpecification: The updated contract spec with imported models.
     """
-    spark = SparkSession.builder.getOrCreate()
+
+    tmp_dir = tempfile.TemporaryDirectory(prefix="datacontract-cli-spark")
+    atexit.register(tmp_dir.cleanup)
+
+    spark = (
+        SparkSession.builder.config("spark.sql.warehouse.dir", f"{tmp_dir}/spark-warehouse")
+        .config("spark.streaming.stopGracefullyOnShutdown", "true")
+        .config("spark.ui.enabled", "false")
+        .getOrCreate()
+    )
     data_contract_specification.servers["local"] = Server(type="dataframe")
-    for temp_view in source.split(","):
-        temp_view = temp_view.strip()
-        df = spark.read.table(temp_view)
-        data_contract_specification.models[temp_view] = import_from_spark_df(df)
+
+    if dataframe is not None:
+        if not isinstance(dataframe, DataFrame):
+            raise TypeError("Expected 'dataframe' to be a pyspark.sql.DataFrame")
+        data_contract_specification.models[source] = import_from_spark_df(spark, source, dataframe, description)
+        return data_contract_specification
+
+    if not source:
+        raise ValueError("Either 'dataframe' or a valid 'source' must be provided")
+
+    for table_name in map(str.strip, source.split(",")):
+        df = spark.read.table(table_name)
+        data_contract_specification.models[table_name] = import_from_spark_df(spark, table_name, df, description)
+
     return data_contract_specification
 
 
-def import_from_spark_df(df: DataFrame) -> Model:
+def import_from_spark_df(spark: SparkSession, source: str, df: DataFrame, description: str) -> Model:
     """
     Converts a Spark DataFrame into a Model.
 
     Args:
+        spark: SparkSession
+        source: A comma-separated string of Spark temporary views to read.
         df: The Spark DataFrame to convert.
+        description: Table level comment
 
     Returns:
         Model: The generated data contract model.
@@ -63,6 +100,11 @@ def import_from_spark_df(df: DataFrame) -> Model:
     model = Model()
     schema = df.schema
 
+    if description is None:
+        model.description = _table_comment_from_spark(spark, source)
+    else:
+        model.description = description
+
     for field in schema:
         model.fields[field.name] = _field_from_struct_type(field)
 
@@ -154,5 +196,67 @@ def _data_type_from_spark(spark_type: types.DataType) -> str:
         return "null"
     elif isinstance(spark_type, types.VarcharType):
         return "varchar"
+    elif isinstance(spark_type, types.VariantType):
+        return "variant"
     else:
         raise ValueError(f"Unsupported Spark type: {spark_type}")
+
+
+def _table_comment_from_spark(spark: SparkSession, source: str):
+    """
+    Attempts to retrieve the table-level comment from a Spark table using multiple fallback methods.
+
+    Args:
+        spark (SparkSession): The active Spark session.
+        source (str): The name of the table (without catalog or schema).
+
+    Returns:
+        str or None: The table-level comment, if found.
+    """
+
+    # Get Current Catalog and Schema from Spark Session
+    try:
+        current_catalog = spark.sql("SELECT current_catalog()").collect()[0][0]
+    except Exception:
+        current_catalog = "hive_metastore"  # Fallback for non-Unity Catalog clusters
+    try:
+        current_schema = spark.catalog.currentDatabase()
+    except Exception:
+        current_schema = spark.sql("SELECT current_database()").collect()[0][0]
+
+    # Get table comment if it exists
+    table_comment = ""
+    source = f"{current_catalog}.{current_schema}.{source}"
+    try:
+        # Initialize WorkspaceClient for Unity Catalog API calls
+        workspace_client = WorkspaceClient()
+        created_table = workspace_client.tables.get(full_name=f"{source}")
+        table_comment = created_table.comment
+        logger.info(f"'{source}' table comment retrieved using 'WorkspaceClient.tables.get({source})'")
+        return table_comment
+    except Exception:
+        pass
+
+    # Fallback to Spark Catalog API for Hive Metastore or Non-UC Tables
+    try:
+        table_comment = spark.catalog.getTable(f"{source}").description
+        logger.info(f"'{source}' table comment retrieved using 'spark.catalog.getTable({source}).description'")
+        return table_comment
+    except Exception:
+        pass
+
+    # Final Fallback Using DESCRIBE TABLE EXTENDED
+    try:
+        rows = spark.sql(f"DESCRIBE TABLE EXTENDED {source}").collect()
+        for row in rows:
+            if row.col_name.strip().lower() == "comment":
+                table_comment = row.data_type
+                break
+        logger.info(f"'{source}' table comment retrieved using 'DESCRIBE TABLE EXTENDED {source}'")
+        return table_comment
+    except Exception:
+        pass
+
+    logger.info(f"{source} table comment could not be retrieved")
+
+    return None

datacontract/imports/sql_importer.py

@@ -105,7 +105,7 @@ def to_dialect(import_args: dict) -> Dialects | None:
     return None
 
 
-def to_physical_type_key(dialect: Dialects | None) -> str:
+def to_physical_type_key(dialect: Dialects | str | None) -> str:
     dialect_map = {
         Dialects.TSQL: "sqlserverType",
         Dialects.POSTGRES: "postgresType",
@@ -116,6 +116,8 @@ def to_physical_type_key(dialect: Dialects | None) -> str:
         Dialects.MYSQL: "mysqlType",
         Dialects.DATABRICKS: "databricksType",
     }
+    if isinstance(dialect, str):
+        dialect = Dialects[dialect.upper()] if dialect.upper() in Dialects.__members__ else None
     return dialect_map.get(dialect, "physicalType")
 
 
@@ -198,7 +200,7 @@ def get_precision_scale(column):
     return None, None
 
 
-def map_type_from_sql(sql_type: str):
+def map_type_from_sql(sql_type: str) -> str | None:
     if sql_type is None:
         return None
 
@@ -218,7 +220,7 @@ def map_type_from_sql(sql_type: str):
         return "string"
     elif sql_type_normed.startswith("ntext"):
         return "string"
-    elif sql_type_normed.startswith("int"):
+    elif sql_type_normed.startswith("int") and not sql_type_normed.startswith("interval"):
         return "int"
     elif sql_type_normed.startswith("bigint"):
         return "long"
@@ -228,6 +230,8 @@ def map_type_from_sql(sql_type: str):
         return "int"
     elif sql_type_normed.startswith("float"):
         return "float"
+    elif sql_type_normed.startswith("double"):
+        return "double"
     elif sql_type_normed.startswith("decimal"):
         return "decimal"
     elif sql_type_normed.startswith("numeric"):
@@ -240,26 +244,20 @@ def map_type_from_sql(sql_type: str):
         return "bytes"
     elif sql_type_normed.startswith("varbinary"):
         return "bytes"
+    elif sql_type_normed.startswith("raw"):
+        return "bytes"
+    elif sql_type_normed == "blob" or sql_type_normed == "bfile":
+        return "bytes"
     elif sql_type_normed == "date":
         return "date"
     elif sql_type_normed == "time":
         return "string"
-    elif sql_type_normed == "timestamp":
-        return "timestamp_ntz"
-    elif (
-        sql_type_normed == "timestamptz"
-        or sql_type_normed == "timestamp_tz"
-        or sql_type_normed == "timestamp with time zone"
-    ):
-        return "timestamp_tz"
-    elif sql_type_normed == "timestampntz" or sql_type_normed == "timestamp_ntz":
+    elif sql_type_normed.startswith("timestamp"):
+        return map_timestamp(sql_type_normed)
+    elif sql_type_normed == "datetime" or sql_type_normed == "datetime2":
         return "timestamp_ntz"
     elif sql_type_normed == "smalldatetime":
         return "timestamp_ntz"
-    elif sql_type_normed == "datetime":
-        return "timestamp_ntz"
-    elif sql_type_normed == "datetime2":
-        return "timestamp_ntz"
     elif sql_type_normed == "datetimeoffset":
         return "timestamp_tz"
     elif sql_type_normed == "uniqueidentifier":  # tsql
@@ -268,10 +266,28 @@ def map_type_from_sql(sql_type: str):
         return "string"
     elif sql_type_normed == "xml":  # tsql
         return "string"
+    elif sql_type_normed.startswith("number"):
+        return "number"
+    elif sql_type_normed == "clob" or sql_type_normed == "nclob":
+        return "text"
     else:
         return "variant"
 
 
+def map_timestamp(timestamp_type: str) -> str:
+    match timestamp_type:
+        case "timestamp" | "timestampntz" | "timestamp_ntz":
+            return "timestamp_ntz"
+        case "timestamptz" | "timestamp_tz" | "timestamp with time zone":
+            return "timestamp_tz"
+        case localTimezone if localTimezone.startswith("timestampltz"):
+            return "timestamp_tz"
+        case timezoneWrittenOut if timezoneWrittenOut.endswith("time zone"):
+            return "timestamp_tz"
+        case _:
+            return "timestamp"
+
+
 def read_file(path):
     if not os.path.exists(path):
         raise DataContractException(