datacontract-cli 0.10.25__py3-none-any.whl → 0.10.27__py3-none-any.whl

datacontract/catalog/catalog.py

@@ -6,7 +6,7 @@ import pytz
 from jinja2 import Environment, PackageLoader, select_autoescape
 
 from datacontract.data_contract import DataContract
-from datacontract.export.html_export import get_version
+from datacontract.export.html_exporter import get_version
 from datacontract.model.data_contract_specification import DataContractSpecification
 
 

datacontract/cli.py

@@ -297,6 +297,14 @@ def import_(
         str,
         typer.Option(help="The location (url or path) of the Data Contract Specification JSON Schema"),
     ] = None,
+    owner: Annotated[
+        Optional[str],
+        typer.Option(help="The owner or team responsible for managing the data contract."),
+    ] = None,
+    id: Annotated[
+        Optional[str],
+        typer.Option(help="The identifier for the the data contract."),
+    ] = None,
 ):
     """
     Create a data contract from the given source location. Saves to file specified by `output` option if present, otherwise prints to stdout.
@@ -316,6 +324,8 @@ def import_(
         dbml_schema=dbml_schema,
         dbml_table=dbml_table,
         iceberg_table=iceberg_table,
+        owner=owner,
+        id=id,
     )
     if output is None:
         console.print(result.to_yaml(), markup=False, soft_wrap=True)

datacontract/data_contract.py

@@ -25,7 +25,7 @@ from datacontract.lint.linters.field_pattern_linter import FieldPatternLinter
 from datacontract.lint.linters.field_reference_linter import FieldReferenceLinter
 from datacontract.lint.linters.notice_period_linter import NoticePeriodLinter
 from datacontract.lint.linters.valid_constraints_linter import ValidFieldConstraintsLinter
-from datacontract.model.data_contract_specification import DataContractSpecification
+from datacontract.model.data_contract_specification import DataContractSpecification, Info
 from datacontract.model.exceptions import DataContractException
 from datacontract.model.run import Check, ResultEnum, Run
 
@@ -270,6 +270,16 @@ class DataContract:
     ) -> DataContractSpecification:
         data_contract_specification_initial = DataContract.init(template=template, schema=schema)
 
-        return importer_factory.create(format).import_source(
+        imported_data_contract_specification = importer_factory.create(format).import_source(
             data_contract_specification=data_contract_specification_initial, source=source, import_args=kwargs
         )
+
+        # Set id and owner if provided
+        if kwargs.get("id"):
+            data_contract_specification_initial.id = kwargs["id"]
+        if kwargs.get("owner"):
+            if data_contract_specification_initial.info is None:
+                data_contract_specification_initial.info = Info()
+            data_contract_specification_initial.info.owner = kwargs["owner"]
+
+        return imported_data_contract_specification

datacontract/engines/soda/connections/kafka.py

@@ -27,6 +27,7 @@ def create_spark_session():
     tmp_dir = tempfile.TemporaryDirectory(prefix="datacontract-cli-spark")
     atexit.register(tmp_dir.cleanup)
 
+    pyspark_version = "3.5.5"  # MUST be the same as in the pyproject.toml
     spark = (
         SparkSession.builder.appName("datacontract")
         .config("spark.sql.warehouse.dir", f"{tmp_dir}/spark-warehouse")
@@ -34,7 +35,7 @@ def create_spark_session():
         .config("spark.ui.enabled", "false")
         .config(
             "spark.jars.packages",
-            "org.apache.spark:spark-sql-kafka-0-10_2.12:3.5.5,org.apache.spark:spark-avro_2.12:3.5.5",
+            f"org.apache.spark:spark-sql-kafka-0-10_2.12:{pyspark_version},org.apache.spark:spark-avro_2.12:{pyspark_version}",
         )
         .getOrCreate()
     )

datacontract/export/dbt_converter.py

@@ -27,7 +27,7 @@ class DbtStageExporter(Exporter):
         )
 
 
-def to_dbt_models_yaml(data_contract_spec: DataContractSpecification, server: str = None):
+def to_dbt_models_yaml(data_contract_spec: DataContractSpecification, server: str = None) -> str:
     dbt = {
         "version": 2,
         "models": [],
@@ -102,8 +102,11 @@ def _to_dbt_model(
         "name": model_key,
     }
     model_type = _to_dbt_model_type(model_value.type)
+
     dbt_model["config"] = {"meta": {"data_contract": data_contract_spec.id}}
-    dbt_model["config"]["materialized"] = model_type
+
+    if model_type:
+        dbt_model["config"]["materialized"] = model_type
 
     if data_contract_spec.info.owner is not None:
         dbt_model["config"]["meta"]["owner"] = data_contract_spec.info.owner
@@ -123,7 +126,7 @@ def _to_dbt_model_type(model_type):
     # Allowed values: table, view, incremental, ephemeral, materialized view
     # Custom values also possible
     if model_type is None:
-        return "table"
+        return None
     if model_type.lower() == "table":
         return "table"
     if model_type.lower() == "view":

datacontract/export/exporter.py

@@ -33,6 +33,7 @@ class ExportFormat(str, Enum):
     avro_idl = "avro-idl"
     sql = "sql"
     sql_query = "sql-query"
+    mermaid = "mermaid"
     html = "html"
     go = "go"
     bigquery = "bigquery"

datacontract/export/exporter_factory.py

@@ -89,6 +89,12 @@ exporter_factory.register_lazy_exporter(
     class_name="DbtExporter",
 )
 
+exporter_factory.register_lazy_exporter(
+    name=ExportFormat.mermaid,
+    module_path="datacontract.export.mermaid_exporter",
+    class_name="MermaidExporter",
+)
+
 exporter_factory.register_lazy_exporter(
     name=ExportFormat.dbt_sources,
     module_path="datacontract.export.dbt_converter",
@@ -127,7 +133,7 @@ exporter_factory.register_lazy_exporter(
 
 exporter_factory.register_lazy_exporter(
     name=ExportFormat.html,
-    module_path="datacontract.export.html_export",
+    module_path="datacontract.export.html_exporter",
     class_name="HtmlExporter",
 )
 

datacontract/export/{html_export.py → html_exporter.py}

@@ -17,6 +17,8 @@ class HtmlExporter(Exporter):
 
 
 def to_html(data_contract_spec: DataContractSpecification) -> str:
+    from datacontract.export.mermaid_exporter import to_mermaid
+
     # Load templates from templates folder
     package_loader = PackageLoader("datacontract", "templates")
     env = Environment(
@@ -54,6 +56,9 @@ def to_html(data_contract_spec: DataContractSpecification) -> str:
     formatted_date = now.strftime("%d %b %Y %H:%M:%S UTC")
     datacontract_cli_version = get_version()
 
+    # Get the mermaid diagram
+    mermaid_diagram = to_mermaid(data_contract_spec)
+
     # Render the template with necessary data
     html_string = template.render(
         datacontract=data_contract_spec,
@@ -62,6 +67,7 @@ def to_html(data_contract_spec: DataContractSpecification) -> str:
         datacontract_yaml=datacontract_yaml,
         formatted_date=formatted_date,
         datacontract_cli_version=datacontract_cli_version,
+        mermaid_diagram=mermaid_diagram,
     )
 
     return html_string

datacontract/export/mermaid_exporter.py

@@ -0,0 +1,32 @@
+from datacontract.export.exporter import Exporter
+from datacontract.model.data_contract_specification import DataContractSpecification
+
+
+class MermaidExporter(Exporter):
+    def export(self, data_contract, model, server, sql_server_type, export_args) -> dict:
+        return to_mermaid(data_contract)
+
+
+def to_mermaid(data_contract_spec: DataContractSpecification) -> str | None:
+    mmd_entity = "erDiagram\n\t"
+    mmd_references = []
+    try:
+        for model_name, model in data_contract_spec.models.items():
+            entity_block = ""
+            for field_name, field in model.fields.items():
+                entity_block += f"\t{field_name.replace('#', 'Nb').replace(' ', '_').replace('/', 'by')}{'🔑' if field.primaryKey or (field.unique and field.required) else ''}{'⌘' if field.references else ''} {field.type}\n"
+                if field.references:
+                    mmd_references.append(
+                        f'"📑{field.references.split(".")[0] if "." in field.references else ""}"'
+                        + "}o--{ ||"
+                        + f'"📑{model_name}"'
+                    )
+            mmd_entity += f'\t"**{model_name}**"' + "{\n" + entity_block + "}\n"
+
+        if mmd_entity == "":
+            return None
+        else:
+            return f"{mmd_entity}\n"
+    except Exception as e:
+        print(f"error : {e}")
+        return None

datacontract/export/odcs_v3_exporter.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Any, Dict
 
 from open_data_contract_standard.model import (
     CustomProperty,
@@ -202,13 +202,31 @@ def to_logical_type(type: str) -> str | None:
         return "array"
     if type.lower() in ["array"]:
         return "array"
+    if type.lower() in ["variant"]:
+        return "variant"
     if type.lower() in ["null"]:
         return None
     return None
 
 
-def to_physical_type(type: str) -> str | None:
-    return type
+def to_physical_type(config: Dict[str, Any]) -> str | None:
+    if config is None:
+        return None
+    if "postgresType" in config:
+        return config["postgresType"]
+    elif "bigqueryType" in config:
+        return config["bigqueryType"]
+    elif "snowflakeType" in config:
+        return config["snowflakeType"]
+    elif "redshiftType" in config:
+        return config["redshiftType"]
+    elif "sqlserverType" in config:
+        return config["sqlserverType"]
+    elif "databricksType" in config:
+        return config["databricksType"]
+    elif "physicalType" in config:
+        return config["physicalType"]
+    return None
 
 
 def to_property(field_name: str, field: Field) -> SchemaProperty:
@@ -231,7 +249,7 @@ def to_property(field_name: str, field: Field) -> SchemaProperty:
 
     if field.type is not None:
         property.logicalType = to_logical_type(field.type)
-        property.physicalType = to_physical_type(field.type)
+        property.physicalType = to_physical_type(field.config) or field.type
 
     if field.description is not None:
         property.description = field.description

datacontract/export/sodacl_converter.py

@@ -2,12 +2,14 @@ import yaml
 
 from datacontract.engines.data_contract_checks import create_checks
 from datacontract.export.exporter import Exporter
+from datacontract.model.data_contract_specification import DataContractSpecification, Server
 from datacontract.model.run import Run
 
 
 class SodaExporter(Exporter):
-    def export(self, data_contract, model, server, sql_server_type, export_args) -> dict:
+    def export(self, data_contract, model, server, sql_server_type, export_args) -> str:
         run = Run.create_run()
+        server = get_server(data_contract, server)
         run.checks.extend(create_checks(data_contract, server))
         return to_sodacl_yaml(run)
 
@@ -28,3 +30,9 @@ def to_sodacl_yaml(run: Run) -> str:
             else:
                 sodacl_dict[key] = value
     return yaml.dump(sodacl_dict)
+
+
+def get_server(data_contract_specification: DataContractSpecification, server_name: str = None) -> Server | None:
+    if server_name is None:
+        return None
+    return data_contract_specification.servers.get(server_name)

datacontract/export/sql_converter.py

@@ -4,7 +4,7 @@ from datacontract.model.data_contract_specification import DataContractSpecifica
 
 
 class SqlExporter(Exporter):
-    def export(self, data_contract, model, server, sql_server_type, export_args) -> dict:
+    def export(self, data_contract, model, server, sql_server_type, export_args) -> str:
         server_type = _determine_sql_server_type(
             data_contract,
             sql_server_type,
@@ -13,7 +13,7 @@ class SqlExporter(Exporter):
 
 
 class SqlQueryExporter(Exporter):
-    def export(self, data_contract, model, server, sql_server_type, export_args) -> dict:
+    def export(self, data_contract, model, server, sql_server_type, export_args) -> str:
         model_name, model_value = _check_models_for_export(data_contract, model, self.export_format)
         server_type = _determine_sql_server_type(data_contract, sql_server_type, export_args.get("server"))
         return to_sql_query(
@@ -117,6 +117,8 @@ def _to_sql_table(model_name, model, server_type="snowflake"):
             result += " primary key"
         if server_type == "databricks" and field.description is not None:
             result += f' COMMENT "{_escape(field.description)}"'
+        if server_type == "snowflake" and field.description is not None:
+            result += f" COMMENT '{_escape(field.description)}'"
         if current_field_index < fields:
             result += ","
         result += "\n"
@@ -124,6 +126,8 @@ def _to_sql_table(model_name, model, server_type="snowflake"):
     result += ")"
     if server_type == "databricks" and model.description is not None:
         result += f' COMMENT "{_escape(model.description)}"'
+    if server_type == "snowflake" and model.description is not None:
+        result += f" COMMENT='{_escape(model.description)}'"
     result += ";\n"
     return result
 

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"]:
@@ -197,6 +201,8 @@ def convert_to_databricks(field: Field) -> None | str:
     if type.lower() in ["array"]:
         item_type = convert_to_databricks(field.items)
         return f"ARRAY<{item_type}>"
+    if type.lower() in ["variant"]:
+        return "VARIANT"
     return None
 
 

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(
@@ -46,15 +51,17 @@ def import_spark(data_contract_specification: DataContractSpecification, source:
     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)
+        data_contract_specification.models[temp_view] = import_from_spark_df(spark, source, df)
     return data_contract_specification
 
 
-def import_from_spark_df(df: DataFrame) -> Model:
+def import_from_spark_df(spark: SparkSession, source: str, df: DataFrame) -> 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.
 
     Returns:
@@ -63,6 +70,8 @@ def import_from_spark_df(df: DataFrame) -> Model:
     model = Model()
     schema = df.schema
 
+    model.description = _table_comment_from_spark(spark, source)
+
     for field in schema:
         model.fields[field.name] = _field_from_struct_type(field)
 
@@ -154,5 +163,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
+        print(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
+        print(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
+        print(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/templates/datacontract.html

@@ -5,6 +5,8 @@
   <meta charset="UTF-8">
   <meta name="viewport" content="width=device-width, initial-scale=1.0">
   {# <script src="https://cdn.tailwindcss.com"></script> #}
+  <script src="https://unpkg.com/@panzoom/panzoom@4.6.0/dist/panzoom.min.js"></script>
+
   <style>
     {{ style | safe }}
   </style>
@@ -29,7 +31,6 @@
       </div>
     </div>
   </nav>
-
   <main class="pb-7">
 
     <div class="pt-5 mx-auto max-w-7xl sm:px-6 lg:px-8">
@@ -77,7 +78,6 @@
             {{ render_partial('partials/datacontract_information.html', datacontract = datacontract) }}
           </section>
 
-
           {% if datacontract.servers %}
           <section id="servers">
             <div class="px-4 sm:px-0">
@@ -103,6 +103,52 @@
           </section>
           {% endif %}
 
+          <section id="diagram" class="mt-6">
+            <div class="px-4 sm:px-0">
+              <h1 class="text-base font-semibold leading-6 text-gray-900">Entity Relationship
+                Diagram</h1>
+              <p class="text-sm text-gray-500">Visual representation of data model relationships</p>
+            </div>
+            <div class="mt-3">
+              <div class="overflow-hidden bg-white shadow-sm ring-1 ring-gray-900/5 sm:rounded-lg">
+                <div class="diagram-container p-4 w-full" id="diagram-container">
+                  <pre class="mermaid">
+                    {{ mermaid_diagram }}
+                  </pre>
+                </div>
+              </div>
+            </div>
+            <script type="module">
+              import mermaid from 'https://cdn.jsdelivr.net/npm/mermaid/dist/mermaid.esm.min.mjs';
+
+              mermaid.initialize({
+                startOnLoad: false,
+                theme: 'neutral'
+              });
+              await mermaid.run({
+                querySelector: '.mermaid',
+                postRenderCallback: (id) => {
+                  const container = document.getElementById("diagram-container");
+                  const svgElement = container.querySelector("svg");
+
+                  if (svgElement) {
+                    // Initialize Panzoom
+                    const panzoomInstance = Panzoom(svgElement, {
+                      maxScale: 5,
+                      minScale: 0.5,
+                      step: 0.1,
+                    });
+
+                    // Mouse wheel zoom
+                    container.addEventListener("wheel", (event) => {
+                      event.preventDefault();
+                      panzoomInstance.zoomWithWheel(event);
+                    });
+                  }
+                }
+              });
+            </script>
+          </section>
 
           <section id="models">
             <div class="flex justify-between">