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

datacontract/imports/sql_importer.py

@@ -1,37 +1,76 @@
-from simple_ddl_parser import parse_from_file
+import logging
+import os
 
-from datacontract.model.data_contract_specification import \
-    DataContractSpecification, Model, Field
+import sqlglot
+from sqlglot.dialects.dialect import Dialects
 
+from datacontract.imports.importer import Importer
+from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model, Server
+from datacontract.model.exceptions import DataContractException
+from datacontract.model.run import ResultEnum
 
-def import_sql(data_contract_specification: DataContractSpecification, format: str, source: str):
-    ddl = parse_from_file(source, group_by_type=True)
-    tables = ddl["tables"]
+
+class SqlImporter(Importer):
+    def import_source(
+        self, data_contract_specification: DataContractSpecification, source: str, import_args: dict
+    ) -> DataContractSpecification:
+        return import_sql(data_contract_specification, self.import_format, source, import_args)
+
+
+def import_sql(
+    data_contract_specification: DataContractSpecification, format: str, source: str, import_args: dict = None
+) -> DataContractSpecification:
+    sql = read_file(source)
+
+    dialect = to_dialect(import_args)
+
+    try:
+        parsed = sqlglot.parse_one(sql=sql, read=dialect)
+    except Exception as e:
+        logging.error(f"Error parsing SQL: {str(e)}")
+        raise DataContractException(
+            type="import",
+            name=f"Reading source from {source}",
+            reason=f"Error parsing SQL: {str(e)}",
+            engine="datacontract",
+            result=ResultEnum.error,
+        )
+
+    server_type: str | None = to_server_type(source, dialect)
+    if server_type is not None:
+        data_contract_specification.servers[server_type] = Server(type=server_type)
+
+    tables = parsed.find_all(sqlglot.expressions.Table)
 
     for table in tables:
         if data_contract_specification.models is None:
             data_contract_specification.models = {}
 
-        table_name = table["table_name"]
+        table_name = table.this.name
 
         fields = {}
-        for column in table["columns"]:
+        for column in parsed.find_all(sqlglot.exp.ColumnDef):
+            if column.parent.this.name != table_name:
+                continue
+
             field = Field()
-            field.type = map_type_from_sql(map_type_from_sql(column["type"]))
-            if not column["nullable"]:
-                field.required = True
-            if column["unique"]:
-                field.unique = True
-            fields[column["name"]] = field
-            if column["size"] is not None:
-                field.maxLength = column["size"]
-
-        if len(table["primary_key"]) == 1:
-            primary_key = table["primary_key"][0]
-            if primary_key in fields:
-                fields[primary_key].unique = True
-                fields[primary_key].required = True
-                fields[primary_key].primary = True
+            col_name = column.this.name
+            col_type = to_col_type(column, dialect)
+            field.type = map_type_from_sql(col_type)
+            col_description = get_description(column)
+            field.description = col_description
+            field.maxLength = get_max_length(column)
+            precision, scale = get_precision_scale(column)
+            field.precision = precision
+            field.scale = scale
+            field.primaryKey = get_primary_key(column)
+            field.required = column.find(sqlglot.exp.NotNullColumnConstraint) is not None or None
+            physical_type_key = to_physical_type_key(dialect)
+            field.config = {
+                physical_type_key: col_type,
+            }
+
+            fields[col_name] = field
 
         data_contract_specification.models[table_name] = Model(
             type="table",
@@ -41,23 +80,223 @@ def import_sql(data_contract_specification: DataContractSpecification, format: s
     return data_contract_specification
 
 
-def map_type_from_sql(sql_type: str):
+def get_primary_key(column) -> bool | None:
+    if column.find(sqlglot.exp.PrimaryKeyColumnConstraint) is not None:
+        return True
+    if column.find(sqlglot.exp.PrimaryKey) is not None:
+        return True
+    return None
+
+
+def to_dialect(import_args: dict) -> Dialects | None:
+    if import_args is None:
+        return None
+    if "dialect" not in import_args:
+        return None
+    dialect = import_args.get("dialect")
+    if dialect is None:
+        return None
+    if dialect == "sqlserver":
+        return Dialects.TSQL
+    if dialect.upper() in Dialects.__members__:
+        return Dialects[dialect.upper()]
+    if dialect == "sqlserver":
+        return Dialects.TSQL
+    return None
+
+
+def to_physical_type_key(dialect: Dialects | str | None) -> str:
+    dialect_map = {
+        Dialects.TSQL: "sqlserverType",
+        Dialects.POSTGRES: "postgresType",
+        Dialects.BIGQUERY: "bigqueryType",
+        Dialects.SNOWFLAKE: "snowflakeType",
+        Dialects.REDSHIFT: "redshiftType",
+        Dialects.ORACLE: "oracleType",
+        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")
+
+
+def to_server_type(source, dialect: Dialects | None) -> str | None:
+    if dialect is None:
+        return None
+    dialect_map = {
+        Dialects.TSQL: "sqlserver",
+        Dialects.POSTGRES: "postgres",
+        Dialects.BIGQUERY: "bigquery",
+        Dialects.SNOWFLAKE: "snowflake",
+        Dialects.REDSHIFT: "redshift",
+        Dialects.ORACLE: "oracle",
+        Dialects.MYSQL: "mysql",
+        Dialects.DATABRICKS: "databricks",
+    }
+    return dialect_map.get(dialect, None)
+
+
+def to_col_type(column, dialect):
+    col_type_kind = column.args["kind"]
+    if col_type_kind is None:
+        return None
+
+    return col_type_kind.sql(dialect)
+
+
+def to_col_type_normalized(column):
+    col_type = column.args["kind"].this.name
+    if col_type is None:
+        return None
+    return col_type.lower()
+
+
+def get_description(column: sqlglot.expressions.ColumnDef) -> str | None:
+    if column.comments is None:
+        return None
+    return " ".join(comment.strip() for comment in column.comments)
+
+
+def get_max_length(column: sqlglot.expressions.ColumnDef) -> int | None:
+    col_type = to_col_type_normalized(column)
+    if col_type is None:
+        return None
+    if col_type not in ["varchar", "char", "nvarchar", "nchar"]:
+        return None
+    col_params = list(column.args["kind"].find_all(sqlglot.expressions.DataTypeParam))
+    max_length_str = None
+    if len(col_params) == 0:
+        return None
+    if len(col_params) == 1:
+        max_length_str = col_params[0].name
+    if len(col_params) == 2:
+        max_length_str = col_params[1].name
+    if max_length_str is not None:
+        return int(max_length_str) if max_length_str.isdigit() else None
+
+
+def get_precision_scale(column):
+    col_type = to_col_type_normalized(column)
+    if col_type is None:
+        return None, None
+    if col_type not in ["decimal", "numeric", "float", "number"]:
+        return None, None
+    col_params = list(column.args["kind"].find_all(sqlglot.expressions.DataTypeParam))
+    if len(col_params) == 0:
+        return None, None
+    if len(col_params) == 1:
+        if not col_params[0].name.isdigit():
+            return None, None
+        precision = int(col_params[0].name)
+        scale = 0
+        return precision, scale
+    if len(col_params) == 2:
+        if not col_params[0].name.isdigit() or not col_params[1].name.isdigit():
+            return None, None
+        precision = int(col_params[0].name)
+        scale = int(col_params[1].name)
+        return precision, scale
+    return None, None
+
+
+def map_type_from_sql(sql_type: str) -> str | None:
     if sql_type is None:
         return None
 
-    if sql_type.lower().startswith("varchar"):
-        return "varchar"
-    if sql_type.lower().startswith("string"):
+    sql_type_normed = sql_type.lower().strip()
+
+    if sql_type_normed.startswith("varchar"):
         return "string"
-    if sql_type.lower().startswith("text"):
-        return "text"
-    elif sql_type.lower().startswith("int"):
-        return "integer"
-    elif sql_type.lower().startswith("float"):
+    elif sql_type_normed.startswith("char"):
+        return "string"
+    elif sql_type_normed.startswith("string"):
+        return "string"
+    elif sql_type_normed.startswith("nchar"):
+        return "string"
+    elif sql_type_normed.startswith("text"):
+        return "string"
+    elif sql_type_normed.startswith("nvarchar"):
+        return "string"
+    elif sql_type_normed.startswith("ntext"):
+        return "string"
+    elif sql_type_normed.startswith("int") and not sql_type_normed.startswith("interval"):
+        return "int"
+    elif sql_type_normed.startswith("bigint"):
+        return "long"
+    elif sql_type_normed.startswith("tinyint"):
+        return "int"
+    elif sql_type_normed.startswith("smallint"):
+        return "int"
+    elif sql_type_normed.startswith("float"):
         return "float"
-    elif sql_type.lower().startswith("bool"):
+    elif sql_type_normed.startswith("double"):
+        return "double"
+    elif sql_type_normed.startswith("decimal"):
+        return "decimal"
+    elif sql_type_normed.startswith("numeric"):
+        return "decimal"
+    elif sql_type_normed.startswith("bool"):
+        return "boolean"
+    elif sql_type_normed.startswith("bit"):
         return "boolean"
-    elif sql_type.lower().startswith("timestamp"):
-        return "timestamp"
+    elif sql_type_normed.startswith("binary"):
+        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.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 == "datetimeoffset":
+        return "timestamp_tz"
+    elif sql_type_normed == "uniqueidentifier":  # tsql
+        return "string"
+    elif sql_type_normed == "json":
+        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(
+            type="import",
+            name=f"Reading source from {path}",
+            reason=f"The file '{path}' does not exist.",
+            engine="datacontract",
+            result=ResultEnum.error,
+        )
+    with open(path, "r") as file:
+        file_content = file.read()
+    return file_content

datacontract/imports/unity_importer.py

@@ -0,0 +1,219 @@
+import json
+import os
+from typing import List
+
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.service.catalog import ColumnInfo, TableInfo
+from open_data_contract_standard.model import OpenDataContractStandard
+
+from datacontract.imports.importer import Importer
+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
+
+
+class UnityImporter(Importer):
+    """
+    UnityImporter class for importing data contract specifications from Unity Catalog.
+    """
+
+    def import_source(
+        self,
+        data_contract_specification: DataContractSpecification | OpenDataContractStandard,
+        source: str,
+        import_args: dict,
+    ) -> DataContractSpecification | OpenDataContractStandard:
+        """
+        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:
+            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 | OpenDataContractStandard, source: str
+) -> DataContractSpecification | OpenDataContractStandard:
+    """
+    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:
+            json_contents = json.loads(file.read())
+            unity_schema = TableInfo.from_dict(json_contents)
+    except json.JSONDecodeError as e:
+        raise DataContractException(
+            type="schema",
+            name="Parse unity schema",
+            reason=f"Failed to parse unity schema from {source}",
+            engine="datacontract",
+            original_exception=e,
+        )
+    return convert_unity_schema(data_contract_specification, unity_schema)
+
+
+def import_unity_from_api(
+    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_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:
+        # print(f"Retrieving Unity Catalog schema for table: {unity_table_full_name}")
+        profile = os.getenv("DATACONTRACT_DATABRICKS_PROFILE")
+        host, token = os.getenv("DATACONTRACT_DATABRICKS_SERVER_HOSTNAME"), os.getenv("DATACONTRACT_DATABRICKS_TOKEN")
+        # print(f"Databricks host: {host}, token: {'***' if token else 'not set'}")
+        exception = DataContractException(
+            type="configuration",
+            name="Databricks configuration",
+            reason="",
+            engine="datacontract",
+        )
+        if not profile and not host and not token:
+            reason = "Either DATACONTRACT_DATABRICKS_PROFILE or both DATACONTRACT_DATABRICKS_SERVER_HOSTNAME and DATACONTRACT_DATABRICKS_TOKEN environment variables must be set"
+            exception.reason = reason
+            raise exception
+        if token and not host:
+            reason = "DATACONTRACT_DATABRICKS_SERVER_HOSTNAME environment variable is not set"
+            exception.reason = reason
+            raise exception
+        if host and not token:
+            reason = "DATACONTRACT_DATABRICKS_TOKEN environment variable is not set"
+            exception.reason = reason
+            raise exception
+        workspace_client = WorkspaceClient(profile=profile) if profile else WorkspaceClient(host=host, token=token)
+    except Exception as e:
+        raise DataContractException(
+            type="schema",
+            name="Retrieve unity catalog schema",
+            reason="Failed to connect to unity catalog schema",
+            engine="datacontract",
+            original_exception=e,
+        )
+
+    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 | OpenDataContractStandard, unity_schema: TableInfo
+) -> DataContractSpecification | OpenDataContractStandard:
+    """
+    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 = {}
+
+    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
+
+    data_contract_specification.models[table_id] = Model(fields=fields, type="table")
+
+    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(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 column in columns:
+        imported_fields[column.name] = _to_field(column)
+
+    return imported_fields
+
+
+def _to_field(column: ColumnInfo) -> Field:
+    field = Field()
+    # The second condition evaluates for complex types (e.g. variant)
+    if column.type_name is not None or (column.type_name is None and column.type_text 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
+
+    return field

datacontract/init/init_template.py

@@ -0,0 +1,20 @@
+import importlib.resources as resources
+import logging
+
+import requests
+
+DEFAULT_DATA_CONTRACT_INIT_TEMPLATE = "datacontract-1.2.1.init.yaml"
+
+
+def get_init_template(location: str = None) -> str:
+    if location is None:
+        logging.info("Use default bundled template " + DEFAULT_DATA_CONTRACT_INIT_TEMPLATE)
+        schemas = resources.files("datacontract")
+        template = schemas.joinpath("schemas", DEFAULT_DATA_CONTRACT_INIT_TEMPLATE)
+        with template.open("r") as file:
+            return file.read()
+    elif location.startswith("http://") or location.startswith("https://"):
+        return requests.get(location).text
+    else:
+        with open(location, "r") as file:
+            return file.read()

datacontract/integration/datamesh_manager.py

@@ -0,0 +1,86 @@
+import os
+
+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:
+        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 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 for unknown data contract ID")
+
+        headers = {"Content-Type": "application/json", "x-api-key": api_key}
+        request_body = run.model_dump_json()
+        # print("Request Body:", request_body)
+        response = requests.post(
+            url,
+            data=request_body,
+            headers=headers,
+            verify=ssl_verification,
+        )
+        # print("Status Code:", response.status_code)
+        # print("Response Body:", response.text)
+        if response.status_code != 200:
+            run.log_error(f"Error publishing test results to Data Mesh Manager: {response.text}")
+            return
+        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)}")
+
+
+def publish_data_contract_to_datamesh_manager(data_contract_dict: dict, ssl_verification: bool):
+    try:
+        api_key = os.getenv("DATAMESH_MANAGER_API_KEY")
+        host = "https://api.datamesh-manager.com"
+        if os.getenv("DATAMESH_MANAGER_HOST") is not None:
+            host = os.getenv("DATAMESH_MANAGER_HOST")
+        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}
+        id = data_contract_dict["id"]
+        url = f"{host}/api/datacontracts/{id}"
+        response = requests.put(
+            url=url,
+            json=data_contract_dict,
+            headers=headers,
+            verify=ssl_verification,
+        )
+        if response.status_code != 200:
+            print(f"Error publishing data contract to Data Mesh Manager: {response.text}")
+            exit(1)
+
+        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)}")