datacontract-cli 0.10.0__py3-none-any.whl → 0.10.37__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

@@ -0,0 +1,262 @@
+import atexit
+import logging
+import tempfile
+
+from databricks.sdk import WorkspaceClient
+from pyspark.sql import DataFrame, SparkSession, types
+
+from datacontract.imports.importer import Importer
+from datacontract.model.data_contract_specification import (
+    DataContractSpecification,
+    Field,
+    Model,
+    Server,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class SparkImporter(Importer):
+    def import_source(
+        self,
+        data_contract_specification: DataContractSpecification,
+        source: str,
+        import_args: dict,
+    ) -> DataContractSpecification:
+        """
+        Imports data from a Spark source into the data contract specification.
+
+        Args:
+            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.
+        """
+        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,
+    dataframe: DataFrame | None = None,
+    description: str | None = None,
+) -> DataContractSpecification:
+    """
+    Imports schema(s) from Spark into a Data Contract Specification.
+
+    Args:
+        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 contract spec with imported models.
+    """
+
+    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")
+
+    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(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.
+    """
+    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)
+
+    return model
+
+
+def _field_from_struct_type(spark_field: types.StructField) -> Field:
+    """
+    Converts a Spark StructField into a Field object for the data contract.
+
+    Args:
+        spark_field: The Spark StructField to convert.
+
+    Returns:
+        Field: The generated Field object.
+    """
+    field = Field()
+    field.required = not spark_field.nullable
+    field.description = spark_field.metadata.get("comment")
+
+    return _type_from_data_type(field, spark_field.dataType)
+
+
+def _type_from_data_type(field: Field, spark_type: types.DataType) -> Field:
+    """
+    Maps Spark data types to the Data Contract type system and updates the field.
+
+    Args:
+        field: The Field object to update.
+        spark_type: The Spark data type to map.
+
+    Returns:
+        Field: The updated Field object.
+    """
+    field.type = _data_type_from_spark(spark_type)
+
+    if field.type == "array":
+        field.items = _type_from_data_type(Field(required=not spark_type.containsNull), spark_type.elementType)
+
+    elif field.type == "map":
+        field.keys = _type_from_data_type(Field(required=True), spark_type.keyType)
+        field.values = _type_from_data_type(Field(required=not spark_type.valueContainsNull), spark_type.valueType)
+
+    elif field.type == "struct":
+        field.fields = {sf.name: _field_from_struct_type(sf) for sf in spark_type.fields}
+
+    return field
+
+
+def _data_type_from_spark(spark_type: types.DataType) -> str:
+    """
+    Maps Spark data types to the Data Contract type system.
+
+    Args:
+        spark_type: The Spark data type to map.
+
+    Returns:
+        str: The corresponding Data Contract type.
+    """
+    if isinstance(spark_type, types.StringType):
+        return "string"
+    elif isinstance(spark_type, (types.IntegerType, types.ShortType)):
+        return "integer"
+    elif isinstance(spark_type, types.LongType):
+        return "long"
+    elif isinstance(spark_type, types.FloatType):
+        return "float"
+    elif isinstance(spark_type, types.DoubleType):
+        return "double"
+    elif isinstance(spark_type, types.StructType):
+        return "struct"
+    elif isinstance(spark_type, types.ArrayType):
+        return "array"
+    elif isinstance(spark_type, types.MapType):
+        return "map"
+    elif isinstance(spark_type, types.TimestampType):
+        return "timestamp"
+    elif isinstance(spark_type, types.TimestampNTZType):
+        return "timestamp_ntz"
+    elif isinstance(spark_type, types.DateType):
+        return "date"
+    elif isinstance(spark_type, types.BooleanType):
+        return "boolean"
+    elif isinstance(spark_type, types.BinaryType):
+        return "bytes"
+    elif isinstance(spark_type, types.DecimalType):
+        return "decimal"
+    elif isinstance(spark_type, types.NullType):
+        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