datacontract-cli 0.10.11__py3-none-any.whl → 0.10.13__py3-none-any.whl

datacontract/imports/unity_importer.py

@@ -1,17 +1,37 @@
 import json
-import requests
 import os
-import typing
+from typing import List, Optional
+
+from pyspark.sql import types
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.service.catalog import TableInfo, ColumnInfo
 
 from datacontract.imports.importer import Importer
+from datacontract.imports.spark_importer import _field_from_struct_type
 from datacontract.model.data_contract_specification import DataContractSpecification, Model, Field
 from datacontract.model.exceptions import DataContractException
 
 
 class UnityImporter(Importer):
+    """
+    UnityImporter class for importing data contract specifications from Unity Catalog.
+    """
+
     def import_source(
         self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
-    ) -> dict:
+    ) -> DataContractSpecification:
+        """
+        Import data contract specification from a source.
+
+        :param data_contract_specification: The data contract specification to be imported.
+        :type data_contract_specification: DataContractSpecification
+        :param source: The source from which to import the data contract specification.
+        :type source: str
+        :param import_args: Additional arguments for the import process.
+        :type import_args: dict
+        :return: The imported data contract specification.
+        :rtype: DataContractSpecification
+        """
         if source is not None:
             data_contract_specification = import_unity_from_json(data_contract_specification, source)
         else:
@@ -24,9 +44,21 @@ class UnityImporter(Importer):
 def import_unity_from_json(
     data_contract_specification: DataContractSpecification, source: str
 ) -> DataContractSpecification:
+    """
+    Import data contract specification from a JSON file.
+
+    :param data_contract_specification: The data contract specification to be imported.
+    :type data_contract_specification: DataContractSpecification
+    :param source: The path to the JSON file.
+    :type source: str
+    :return: The imported data contract specification.
+    :rtype: DataContractSpecification
+    :raises DataContractException: If there is an error parsing the JSON file.
+    """
     try:
         with open(source, "r") as file:
-            unity_schema = json.loads(file.read())
+            json_contents = json.loads(file.read())
+            unity_schema = TableInfo.from_dict(json_contents)
     except json.JSONDecodeError as e:
         raise DataContractException(
             type="schema",
@@ -39,114 +71,103 @@ def import_unity_from_json(
 
 
 def import_unity_from_api(
-    data_contract_specification: DataContractSpecification, unity_table_full_name: typing.Optional[str] = None
+    data_contract_specification: DataContractSpecification, unity_table_full_name: Optional[str] = None
 ) -> DataContractSpecification:
-    databricks_instance = os.getenv("DATABRICKS_IMPORT_INSTANCE")
-    access_token = os.getenv("DATABRICKS_IMPORT_ACCESS_TOKEN")
-
-    if not databricks_instance or not access_token:
-        print("Missing environment variables for Databricks instance or access token.")
-        print("Both, $DATABRICKS_IMPORT_INSTANCE and $DATABRICKS_IMPORT_ACCESS_TOKEN must be set.")
-        exit(1)  # Exit if variables are not set
-
-    api_url = f"{databricks_instance}/api/2.1/unity-catalog/tables/{unity_table_full_name}"
-
-    headers = {"Authorization": f"Bearer {access_token}"}
-    response = requests.get(api_url, headers=headers)
-
-    if response.status_code != 200:
+    """
+    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]
+    :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)
+    except Exception as e:
         raise DataContractException(
             type="schema",
             name="Retrieve unity catalog schema",
-            reason=f"Failed to retrieve unity catalog schema from databricks instance: {response.status_code} {response.text}",
+            reason=f"Failed to retrieve unity catalog schema from databricks profile: {os.getenv('DATABRICKS_CONFIG_PROFILE')}",
             engine="datacontract",
+            original_exception=e,
         )
 
-    convert_unity_schema(data_contract_specification, response.json())
+    convert_unity_schema(data_contract_specification, unity_schema)
 
     return data_contract_specification
 
 
 def convert_unity_schema(
-    data_contract_specification: DataContractSpecification, unity_schema: dict
+    data_contract_specification: DataContractSpecification, unity_schema: TableInfo
 ) -> DataContractSpecification:
+    """
+    Convert Unity schema to data contract specification.
+
+    :param data_contract_specification: The data contract specification to be converted.
+    :type data_contract_specification: DataContractSpecification
+    :param unity_schema: The Unity schema to be converted.
+    :type unity_schema: TableInfo
+    :return: The converted data contract specification.
+    :rtype: DataContractSpecification
+    """
     if data_contract_specification.models is None:
         data_contract_specification.models = {}
 
-    fields = import_table_fields(unity_schema.get("columns"))
+    fields = import_table_fields(unity_schema.columns)
 
-    table_id = unity_schema.get("table_id")
+    table_id = unity_schema.name or unity_schema.table_id
 
     data_contract_specification.models[table_id] = Model(fields=fields, type="table")
 
-    if unity_schema.get("name") is not None:
-        data_contract_specification.models[table_id].title = unity_schema.get("name")
+    if unity_schema.name:
+        data_contract_specification.models[table_id].title = unity_schema.name
+
+    if unity_schema.comment:
+        data_contract_specification.models[table_id].description = unity_schema.comment
 
     return data_contract_specification
 
 
-def import_table_fields(table_fields):
+def import_table_fields(columns: List[ColumnInfo]) -> dict[str, Field]:
+    """
+    Import table fields from Unity schema columns.
+
+    Here we are first converting the `ColumnInfo.type_json` to a Spark StructField object
+    so we can leave the complexity of the Spark field types to the Spark JSON schema parser,
+    then re-use the logic in `datacontract.imports.spark_importer` to convert the StructField
+    into a Field object.
+
+    :param columns: The list of Unity schema columns.
+    :type columns: List[ColumnInfo]
+    :return: A dictionary of imported fields.
+    :rtype: dict[str, Field]
+    """
     imported_fields = {}
-    for field in table_fields:
-        field_name = field.get("name")
-        imported_fields[field_name] = Field()
-        imported_fields[field_name].required = field.get("nullable") == "false"
-        imported_fields[field_name].description = field.get("comment")
-
-        # databricks api 2.1 specifies that type_name can be any of:
-        # BOOLEAN | BYTE | SHORT | INT | LONG | FLOAT | DOUBLE | DATE | TIMESTAMP | TIMESTAMP_NTZ | STRING
-        # | BINARY | DECIMAL | INTERVAL | ARRAY | STRUCT | MAP | CHAR | NULL | USER_DEFINED_TYPE | TABLE_TYPE
-        if field.get("type_name") in ["INTERVAL", "ARRAY", "STRUCT", "MAP", "USER_DEFINED_TYPE", "TABLE_TYPE"]:
-            # complex types are not supported, yet
-            raise DataContractException(
-                type="schema",
-                result="failed",
-                name="Map unity type to data contract type",
-                reason=f"type ${field.get('type_name')} is not supported yet for unity import",
-                engine="datacontract",
-            )
 
-        imported_fields[field_name].type = map_type_from_unity(field.get("type_name"))
+    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)
 
     return imported_fields
 
 
-def map_type_from_unity(type_str: str):
-    if type_str == "BOOLEAN":
-        return "boolean"
-    elif type_str == "BYTE":
-        return "bytes"
-    elif type_str == "SHORT":
-        return "int"
-    elif type_str == "INT":
-        return "int"
-    elif type_str == "LONG":
-        return "long"
-    elif type_str == "FLOAT":
-        return "float"
-    elif type_str == "DOUBLE":
-        return "double"
-    elif type_str == "DATE":
-        return "date"
-    elif type_str == "TIMESTAMP":
-        return "timestamp"
-    elif type_str == "TIMESTAMP_NTZ":
-        return "timestamp_ntz"
-    elif type_str == "STRING":
-        return "string"
-    elif type_str == "BINARY":
-        return "bytes"
-    elif type_str == "DECIMAL":
-        return "decimal"
-    elif type_str == "CHAR":
-        return "varchar"
-    elif type_str == "NULL":
-        return "null"
-    else:
-        raise DataContractException(
-            type="schema",
-            result="failed",
-            name="Map unity type to data contract type",
-            reason=f"Unsupported type {type_str} in unity json definition.",
-            engine="datacontract",
-        )
+def _type_json_to_spark_field(type_json: str) -> types.StructField:
+    """
+    Parses a JSON string representing a Spark field and returns a StructField object.
+
+    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)

datacontract/integration/{publish_datamesh_manager.py → datamesh_manager.py}

@@ -2,28 +2,29 @@ import os
 
 import requests
 
+from datacontract.model.data_contract_specification import DataContractSpecification
 from datacontract.model.run import Run
 
 
-def publish_datamesh_manager(run: Run, publish_url: str):
+def publish_test_results_to_datamesh_manager(run: Run, publish_url: str):
     try:
         if publish_url is None:
             # this url supports Data Mesh Manager and Data Contract Manager
             url = "https://api.datamesh-manager.com/api/test-results"
         else:
             url = publish_url
+
         api_key = os.getenv("DATAMESH_MANAGER_API_KEY")
         if api_key is None:
             api_key = os.getenv("DATACONTRACT_MANAGER_API_KEY")
-
-        if run.dataContractId is None:
-            raise Exception("Cannot publish run results, as data contract ID is unknown")
-
         if api_key is None:
             raise Exception(
                 "Cannot publish run results, as DATAMESH_MANAGER_API_KEY nor DATACONTRACT_MANAGER_API_KEY are not set"
             )
 
+        if run.dataContractId is None:
+            raise Exception("Cannot publish run results, as data contract ID is unknown")
+
         headers = {"Content-Type": "application/json", "x-api-key": api_key}
         request_body = run.model_dump_json()
         # print("Request Body:", request_body)
@@ -36,3 +37,30 @@ def publish_datamesh_manager(run: Run, publish_url: str):
         run.log_info(f"Published test results to {url}")
     except Exception as e:
         run.log_error(f"Failed publishing test results. Error: {str(e)}")
+
+
+def publish_data_contract_to_datamesh_manager(data_contract_specification: DataContractSpecification):
+    try:
+        api_key = os.getenv("DATAMESH_MANAGER_API_KEY")
+        if api_key is None:
+            api_key = os.getenv("DATACONTRACT_MANAGER_API_KEY")
+        if api_key is None:
+            raise Exception(
+                "Cannot publish data contract, as neither DATAMESH_MANAGER_API_KEY nor DATACONTRACT_MANAGER_API_KEY is set"
+            )
+        headers = {"Content-Type": "application/json", "x-api-key": api_key}
+        spec = data_contract_specification
+        id = spec.id
+        url = "https://api.datamesh-manager.com/api/datacontracts/{0}".format(id)
+        request_body = spec.model_dump_json().encode("utf-8")
+        response = requests.put(
+            url=url,
+            data=request_body,
+            headers=headers,
+        )
+        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}")
+    except Exception as e:
+        print(f"Failed publishing data contract. Error: {str(e)}")

datacontract/integration/{publish_opentelemetry.py → opentelemetry.py}

@@ -34,7 +34,7 @@ from datacontract.model.run import Run
 # - Metrics only, no logs yet (but loosely planned)
 
 
-def publish_opentelemetry(run: Run):
+def publish_test_results_to_opentelemetry(run: Run):
     try:
         if run.dataContractId is None:
             raise Exception("Cannot publish run results, as data contract ID is unknown")

datacontract/lint/resolve.py

@@ -114,7 +114,16 @@ def _resolve_definition_ref(ref, spec) -> Definition:
 
 def _find_by_path_in_spec(definition_path: str, spec: DataContractSpecification):
     path_elements = definition_path.split("/")
-    definition = spec.definitions[path_elements[2]]
+    definition_key = path_elements[2]
+    if definition_key not in spec.definitions:
+        raise DataContractException(
+            type="lint",
+            result="failed",
+            name="Check that data contract YAML is valid",
+            reason=f"Cannot resolve definition {definition_key}",
+            engine="datacontract",
+        )
+    definition = spec.definitions[definition_key]
     definition = _find_subfield_in_definition(definition, path_elements[3:])
     return definition
 

datacontract/lint/urls.py

@@ -25,16 +25,30 @@ def fetch_resource(url: str):
 
 
 def _set_api_key(headers, url):
-    if ".datamesh-manager.com/" not in url:
-        return
-    datamesh_manager_api_key = os.getenv("DATAMESH_MANAGER_API_KEY")
-    if datamesh_manager_api_key is None or datamesh_manager_api_key == "":
-        print("Error: Data Mesh Manager API Key is not set. Set env variable DATAMESH_MANAGER_API_KEY.")
-        raise DataContractException(
-            type="lint",
-            name=f"Reading data contract from {url}",
-            reason="Error: Data Mesh Manager API Key is not set. Set env variable DATAMESH_MANAGER_API_KEY.",
-            engine="datacontract",
-            result="error",
-        )
-    headers["x-api-key"] = datamesh_manager_api_key
+    if ".datamesh-manager.com/" in url:
+        datamesh_manager_api_key = os.getenv("DATAMESH_MANAGER_API_KEY")
+        if datamesh_manager_api_key is None or datamesh_manager_api_key == "":
+            print("Error: Data Mesh Manager API Key is not set. Set env variable DATAMESH_MANAGER_API_KEY.")
+            raise DataContractException(
+                type="lint",
+                name=f"Reading data contract from {url}",
+                reason="Error: Data Mesh Manager API Key is not set. Set env variable DATAMESH_MANAGER_API_KEY.",
+                engine="datacontract",
+                result="error",
+            )
+        headers["x-api-key"] = datamesh_manager_api_key
+    elif ".datacontract-manager.com/" in url:
+        datacontract_manager_api_key = os.getenv("DATACONTRACT_MANAGER_API_KEY")
+        if datacontract_manager_api_key is None or datacontract_manager_api_key == "":
+            print("Error: Data Contract Manager API Key is not set. Set env variable DATACONTRACT_MANAGER_API_KEY.")
+            raise DataContractException(
+                type="lint",
+                name=f"Reading data contract from {url}",
+                reason="Error: Data Contract Manager API Key is not set. Set env variable DATACONTRACT_MANAGER_API_KEY.",
+                engine="datacontract",
+                result="error",
+            )
+        headers["x-api-key"] = datacontract_manager_api_key
+    else:
+        # do nothing
+        pass

datacontract/model/data_contract_specification.py

@@ -73,7 +73,7 @@ class Definition(pyd.BaseModel):
     exclusiveMaximum: int = None
     pii: bool = None
     classification: str = None
-    fields: Dict[str, "Definition"] = {}
+    fields: Dict[str, "Field"] = {}
     tags: List[str] = []
     links: Dict[str, str] = {}
     example: str = None
@@ -239,4 +239,8 @@ class DataContractSpecification(pyd.BaseModel):
         return DataContractSpecification(**data)
 
     def to_yaml(self):
-        return yaml.dump(self.model_dump(exclude_defaults=True, exclude_none=True), sort_keys=False, allow_unicode=True)
+        return yaml.dump(
+            self.model_dump(exclude_defaults=True, exclude_none=True, by_alias=True),
+            sort_keys=False,
+            allow_unicode=True,
+        )