datacontract-cli 0.10.26__py3-none-any.whl → 0.10.28__py3-none-any.whl

datacontract/export/sql_type_converter.py

@@ -158,9 +158,13 @@ def convert_to_dataframe(field: Field) -> None | str:
 # databricks data types:
 # https://docs.databricks.com/en/sql/language-manual/sql-ref-datatypes.html
 def convert_to_databricks(field: Field) -> None | str:
-    if field.config and "databricksType" in field.config:
-        return field.config["databricksType"]
     type = field.type
+    if (
+        field.config
+        and "databricksType" in field.config
+        and type.lower() not in ["array", "object", "record", "struct"]
+    ):
+        return field.config["databricksType"]
     if type is None:
         return None
     if type.lower() in ["string", "varchar", "text"]:

datacontract/imports/excel_importer.py

@@ -31,8 +31,11 @@ logger = logging.getLogger(__name__)
 
 class ExcelImporter(Importer):
     def import_source(
-        self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
-    ) -> OpenDataContractStandard:
+        self,
+        data_contract_specification: DataContractSpecification | OpenDataContractStandard,
+        source: str,
+        import_args: dict,
+    ) -> DataContractSpecification | OpenDataContractStandard:
         return import_excel_as_odcs(source)
 
 

datacontract/imports/importer.py

@@ -12,7 +12,7 @@ class Importer(ABC):
     @abstractmethod
     def import_source(
         self,
-        data_contract_specification: DataContractSpecification,
+        data_contract_specification: DataContractSpecification | OpenDataContractStandard,
         source: str,
         import_args: dict,
     ) -> DataContractSpecification | OpenDataContractStandard:
@@ -39,3 +39,12 @@ class ImportFormat(str, Enum):
     @classmethod
     def get_supported_formats(cls):
         return list(map(lambda c: c.value, cls))
+
+
+class Spec(str, Enum):
+    datacontract_specification = "datacontract_specification"
+    odcs = "odcs"
+
+    @classmethod
+    def get_supported_types(cls):
+        return list(map(lambda c: c.value, cls))

datacontract/imports/odcs_importer.py

@@ -48,9 +48,9 @@ def import_odcs(data_contract_specification: DataContractSpecification, source:
             engine="datacontract",
         )
     elif odcs_api_version.startswith("v3."):
-        from datacontract.imports.odcs_v3_importer import import_odcs_v3
+        from datacontract.imports.odcs_v3_importer import import_odcs_v3_as_dcs
 
-        return import_odcs_v3(data_contract_specification, source)
+        return import_odcs_v3_as_dcs(data_contract_specification, source)
     else:
         raise DataContractException(
             type="schema",

datacontract/imports/odcs_v3_importer.py

@@ -29,17 +29,18 @@ class OdcsImporter(Importer):
     def import_source(
         self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
     ) -> DataContractSpecification:
-        return import_odcs_v3(data_contract_specification, source)
+        return import_odcs_v3_as_dcs(data_contract_specification, source)
 
 
-def import_odcs_v3(data_contract_specification: DataContractSpecification, source: str) -> DataContractSpecification:
+def import_odcs_v3_as_dcs(
+    data_contract_specification: DataContractSpecification, source: str
+) -> DataContractSpecification:
     source_str = read_resource(source)
-    return import_odcs_v3_from_str(data_contract_specification, source_str)
+    odcs = parse_odcs_v3_from_str(source_str)
+    return import_from_odcs(data_contract_specification, odcs)
 
 
-def import_odcs_v3_from_str(
-    data_contract_specification: DataContractSpecification, source_str: str
-) -> DataContractSpecification:
+def parse_odcs_v3_from_str(source_str):
     try:
         odcs = OpenDataContractStandard.from_string(source_str)
     except Exception as e:
@@ -50,11 +51,10 @@ def import_odcs_v3_from_str(
             engine="datacontract",
             original_exception=e,
         )
-
-    return import_from_odcs_model(data_contract_specification, odcs)
+    return odcs
 
 
-def import_from_odcs_model(data_contract_specification, odcs):
+def import_from_odcs(data_contract_specification: DataContractSpecification, odcs: OpenDataContractStandard):
     data_contract_specification.id = odcs.id
     data_contract_specification.info = import_info(odcs)
     data_contract_specification.servers = import_servers(odcs)

datacontract/imports/spark_importer.py

@@ -1,3 +1,6 @@
+import logging
+
+from databricks.sdk import WorkspaceClient
 from pyspark.sql import DataFrame, SparkSession, types
 
 from datacontract.imports.importer import Importer
@@ -8,6 +11,8 @@ from datacontract.model.data_contract_specification import (
     Server,
 )
 
+logger = logging.getLogger(__name__)
+
 
 class SparkImporter(Importer):
     def import_source(
@@ -23,39 +28,60 @@ 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()
     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 +89,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)
 
@@ -158,3 +189,63 @@ def _data_type_from_spark(spark_type: types.DataType) -> str:
         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
 

datacontract/imports/unity_importer.py

@@ -1,14 +1,14 @@
 import json
 import os
-from typing import List, Optional
+from typing import List
 
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.service.catalog import ColumnInfo, TableInfo
-from pyspark.sql import types
+from open_data_contract_standard.model import OpenDataContractStandard
 
 from datacontract.imports.importer import Importer
-from datacontract.imports.spark_importer import _field_from_struct_type
-from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model
+from datacontract.imports.sql_importer import map_type_from_sql, to_physical_type_key
+from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model, Server
 from datacontract.model.exceptions import DataContractException
 
 
@@ -18,8 +18,11 @@ class UnityImporter(Importer):
     """
 
     def import_source(
-        self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
-    ) -> DataContractSpecification:
+        self,
+        data_contract_specification: DataContractSpecification | OpenDataContractStandard,
+        source: str,
+        import_args: dict,
+    ) -> DataContractSpecification | OpenDataContractStandard:
         """
         Import data contract specification from a source.
 
@@ -35,15 +38,14 @@ class UnityImporter(Importer):
         if source is not None:
             data_contract_specification = import_unity_from_json(data_contract_specification, source)
         else:
-            data_contract_specification = import_unity_from_api(
-                data_contract_specification, import_args.get("unity_table_full_name")
-            )
+            unity_table_full_name_list = import_args.get("unity_table_full_name")
+            data_contract_specification = import_unity_from_api(data_contract_specification, unity_table_full_name_list)
         return data_contract_specification
 
 
 def import_unity_from_json(
-    data_contract_specification: DataContractSpecification, source: str
-) -> DataContractSpecification:
+    data_contract_specification: DataContractSpecification | OpenDataContractStandard, source: str
+) -> DataContractSpecification | OpenDataContractStandard:
     """
     Import data contract specification from a JSON file.
 
@@ -71,39 +73,66 @@ def import_unity_from_json(
 
 
 def import_unity_from_api(
-    data_contract_specification: DataContractSpecification, unity_table_full_name: Optional[str] = None
+    data_contract_specification: DataContractSpecification, unity_table_full_name_list: List[str] = None
 ) -> DataContractSpecification:
     """
     Import data contract specification from Unity Catalog API.
 
     :param data_contract_specification: The data contract specification to be imported.
     :type data_contract_specification: DataContractSpecification
-    :param unity_table_full_name: The full name of the Unity table.
-    :type unity_table_full_name: Optional[str]
+    :param unity_table_full_name_list: The full name of the Unity table.
+    :type unity_table_full_name_list: list[str]
     :return: The imported data contract specification.
     :rtype: DataContractSpecification
     :raises DataContractException: If there is an error retrieving the schema from the API.
     """
     try:
-        workspace_client = WorkspaceClient()
-        unity_schema: TableInfo = workspace_client.tables.get(unity_table_full_name)
+        # print(f"Retrieving Unity Catalog schema for table: {unity_table_full_name}")
+        host, token = os.getenv("DATACONTRACT_DATABRICKS_SERVER_HOSTNAME"), os.getenv("DATACONTRACT_DATABRICKS_TOKEN")
+        # print(f"Databricks host: {host}, token: {'***' if token else 'not set'}")
+        if not host:
+            raise DataContractException(
+                type="configuration",
+                name="Databricks configuration",
+                reason="DATACONTRACT_DATABRICKS_SERVER_HOSTNAME environment variable is not set",
+                engine="datacontract",
+            )
+        if not token:
+            raise DataContractException(
+                type="configuration",
+                name="Databricks configuration",
+                reason="DATACONTRACT_DATABRICKS_TOKEN environment variable is not set",
+                engine="datacontract",
+            )
+        workspace_client = WorkspaceClient(host=host, token=token)
     except Exception as e:
         raise DataContractException(
             type="schema",
             name="Retrieve unity catalog schema",
-            reason=f"Failed to retrieve unity catalog schema from databricks profile: {os.getenv('DATABRICKS_CONFIG_PROFILE')}",
+            reason="Failed to connect to unity catalog schema",
             engine="datacontract",
             original_exception=e,
         )
 
-    convert_unity_schema(data_contract_specification, unity_schema)
+    for unity_table_full_name in unity_table_full_name_list:
+        try:
+            unity_schema: TableInfo = workspace_client.tables.get(unity_table_full_name)
+        except Exception as e:
+            raise DataContractException(
+                type="schema",
+                name="Retrieve unity catalog schema",
+                reason=f"Unity table {unity_table_full_name} not found",
+                engine="datacontract",
+                original_exception=e,
+            )
+        data_contract_specification = convert_unity_schema(data_contract_specification, unity_schema)
 
     return data_contract_specification
 
 
 def convert_unity_schema(
-    data_contract_specification: DataContractSpecification, unity_schema: TableInfo
-) -> DataContractSpecification:
+    data_contract_specification: DataContractSpecification | OpenDataContractStandard, unity_schema: TableInfo
+) -> DataContractSpecification | OpenDataContractStandard:
     """
     Convert Unity schema to data contract specification.
 
@@ -117,6 +146,21 @@ def convert_unity_schema(
     if data_contract_specification.models is None:
         data_contract_specification.models = {}
 
+    if data_contract_specification.servers is None:
+        data_contract_specification.servers = {}
+
+    # Configure databricks server with catalog and schema from Unity table info
+    schema_name = unity_schema.schema_name
+    catalog_name = unity_schema.catalog_name
+    if catalog_name and schema_name:
+        server_name = "myserver"  # Default server name
+
+        data_contract_specification.servers[server_name] = Server(
+            type="databricks",
+            catalog=catalog_name,
+            schema=schema_name,
+        )
+
     fields = import_table_fields(unity_schema.columns)
 
     table_id = unity_schema.name or unity_schema.table_id
@@ -149,25 +193,21 @@ def import_table_fields(columns: List[ColumnInfo]) -> dict[str, Field]:
     imported_fields = {}
 
     for column in columns:
-        struct_field: types.StructField = _type_json_to_spark_field(column.type_json)
-        imported_fields[column.name] = _field_from_struct_type(struct_field)
+        imported_fields[column.name] = _to_field(column)
 
     return imported_fields
 
 
-def _type_json_to_spark_field(type_json: str) -> types.StructField:
-    """
-    Parses a JSON string representing a Spark field and returns a StructField object.
+def _to_field(column: ColumnInfo) -> Field:
+    field = Field()
+    if column.type_name is not None:
+        sql_type = str(column.type_text)
+        field.type = map_type_from_sql(sql_type)
+    physical_type_key = to_physical_type_key("databricks")
+    field.config = {
+        physical_type_key: sql_type,
+    }
+    field.required = column.nullable is None or not column.nullable
+    field.description = column.comment if column.comment else None
 
-    The reason we do this is to leverage the Spark JSON schema parser to handle the
-    complexity of the Spark field types. The field `type_json` in the Unity API is
-    the output of a `StructField.jsonValue()` call.
-
-    :param type_json: The JSON string representing the Spark field.
-    :type type_json: str
-
-    :return: The StructField object.
-    :rtype: types.StructField
-    """
-    type_dict = json.loads(type_json)
-    return types.StructField.fromJson(type_dict)
+    return field

datacontract/integration/datamesh_manager.py

@@ -4,6 +4,9 @@ import requests
 
 from datacontract.model.run import Run
 
+# used to retrieve the HTML location of the published data contract or test results
+RESPONSE_HEADER_LOCATION_HTML = "location-html"
+
 
 def publish_test_results_to_datamesh_manager(run: Run, publish_url: str, ssl_verification: bool):
     try:
@@ -38,7 +41,12 @@ def publish_test_results_to_datamesh_manager(run: Run, publish_url: str, ssl_ver
         if response.status_code != 200:
             run.log_error(f"Error publishing test results to Data Mesh Manager: {response.text}")
             return
-        run.log_info(f"Published test results to {url}")
+        run.log_info("Published test results successfully")
+
+        location_html = response.headers.get(RESPONSE_HEADER_LOCATION_HTML)
+        if location_html is not None and len(location_html) > 0:
+            print(f"🚀 Open {location_html}")
+
     except Exception as e:
         run.log_error(f"Failed publishing test results. Error: {str(e)}")
 
@@ -67,6 +75,12 @@ def publish_data_contract_to_datamesh_manager(data_contract_dict: dict, ssl_veri
         if response.status_code != 200:
             print(f"Error publishing data contract to Data Mesh Manager: {response.text}")
             exit(1)
-        print(f"Published data contract to {url}")
+
+        print("✅ Published data contract successfully")
+
+        location_html = response.headers.get(RESPONSE_HEADER_LOCATION_HTML)
+        if location_html is not None and len(location_html) > 0:
+            print(f"🚀 Open {location_html}")
+
     except Exception as e:
         print(f"Failed publishing data contract. Error: {str(e)}")

datacontract/lint/resolve.py

@@ -5,8 +5,9 @@ import warnings
 import fastjsonschema
 import yaml
 from fastjsonschema import JsonSchemaValueException
+from open_data_contract_standard.model import OpenDataContractStandard
 
-from datacontract.imports.odcs_v3_importer import import_odcs_v3_from_str
+from datacontract.imports.odcs_v3_importer import import_from_odcs, parse_odcs_v3_from_str
 from datacontract.lint.resources import read_resource
 from datacontract.lint.schema import fetch_schema
 from datacontract.lint.urls import fetch_resource
@@ -46,6 +47,34 @@ def resolve_data_contract(
         )
 
 
+def resolve_data_contract_v2(
+    data_contract_location: str = None,
+    data_contract_str: str = None,
+    data_contract: DataContractSpecification | OpenDataContractStandard = None,
+    schema_location: str = None,
+    inline_definitions: bool = False,
+    inline_quality: bool = False,
+) -> DataContractSpecification | OpenDataContractStandard:
+    if data_contract_location is not None:
+        return resolve_data_contract_from_location_v2(
+            data_contract_location, schema_location, inline_definitions, inline_quality
+        )
+    elif data_contract_str is not None:
+        return _resolve_data_contract_from_str_v2(
+            data_contract_str, schema_location, inline_definitions, inline_quality
+        )
+    elif data_contract is not None:
+        return data_contract
+    else:
+        raise DataContractException(
+            type="lint",
+            result=ResultEnum.failed,
+            name="Check that data contract YAML is valid",
+            reason="Data contract needs to be provided",
+            engine="datacontract",
+        )
+
+
 def resolve_data_contract_dict(
     data_contract_location: str = None,
     data_contract_str: str = None,
@@ -67,6 +96,13 @@ def resolve_data_contract_dict(
         )
 
 
+def resolve_data_contract_from_location_v2(
+    location, schema_location: str = None, inline_definitions: bool = False, inline_quality: bool = False
+) -> DataContractSpecification | OpenDataContractStandard:
+    data_contract_str = read_resource(location)
+    return _resolve_data_contract_from_str_v2(data_contract_str, schema_location, inline_definitions, inline_quality)
+
+
 def resolve_data_contract_from_location(
     location, schema_location: str = None, inline_definitions: bool = False, inline_quality: bool = False
 ) -> DataContractSpecification:
@@ -242,6 +278,21 @@ def _get_quality_ref_file(quality_spec: str | object) -> str | object:
     return quality_spec
 
 
+def _resolve_data_contract_from_str_v2(
+    data_contract_str, schema_location: str = None, inline_definitions: bool = False, inline_quality: bool = False
+) -> DataContractSpecification | OpenDataContractStandard:
+    yaml_dict = _to_yaml(data_contract_str)
+
+    if is_open_data_contract_standard(yaml_dict):
+        logging.info("Importing ODCS v3")
+        # if ODCS, then validate the ODCS schema and import to DataContractSpecification directly
+        odcs = parse_odcs_v3_from_str(data_contract_str)
+        return odcs
+
+    logging.info("Importing DCS")
+    return _resolve_dcs_from_yaml_dict(inline_definitions, inline_quality, schema_location, yaml_dict)
+
+
 def _resolve_data_contract_from_str(
     data_contract_str, schema_location: str = None, inline_definitions: bool = False, inline_quality: bool = False
 ) -> DataContractSpecification:
@@ -250,15 +301,19 @@ def _resolve_data_contract_from_str(
     if is_open_data_contract_standard(yaml_dict):
         logging.info("Importing ODCS v3")
         # if ODCS, then validate the ODCS schema and import to DataContractSpecification directly
+        odcs = parse_odcs_v3_from_str(data_contract_str)
+
         data_contract_specification = DataContractSpecification(dataContractSpecification="1.1.0")
-        return import_odcs_v3_from_str(data_contract_specification, source_str=data_contract_str)
-    else:
-        logging.info("Importing DCS")
+        return import_from_odcs(data_contract_specification, odcs)
+
+    logging.info("Importing DCS")
+    return _resolve_dcs_from_yaml_dict(inline_definitions, inline_quality, schema_location, yaml_dict)
+
 
+def _resolve_dcs_from_yaml_dict(inline_definitions, inline_quality, schema_location, yaml_dict):
     _validate_data_contract_specification_schema(yaml_dict, schema_location)
     data_contract_specification = yaml_dict
     spec = DataContractSpecification(**data_contract_specification)
-
     if inline_definitions:
         inline_definitions_into_data_contract(spec)
     ## Suppress DeprecationWarning when accessing spec.quality,
@@ -276,7 +331,6 @@ def _resolve_data_contract_from_str(
             )
     if spec_quality and inline_quality:
         _resolve_quality_ref(spec_quality)
-
     return spec