datacontract-cli 0.10.14__py3-none-any.whl → 0.10.15__py3-none-any.whl

datacontract/export/great_expectations_converter.py

@@ -1,49 +1,118 @@
+"""
+This module provides functionalities to export data contracts to Great Expectations suites.
+It includes definitions for exporting different types of data (pandas, Spark, SQL) into
+Great Expectations expectations format.
+"""
+
 import json
-from typing import Dict, List, Any
+from enum import Enum
+from typing import Any, Dict, List
 
 import yaml
 
-from datacontract.model.data_contract_specification import DataContractSpecification, Field, Quality
-from datacontract.export.exporter import Exporter, _check_models_for_export
+from datacontract.export.exporter import (
+    Exporter,
+    _check_models_for_export,
+)
+from datacontract.export.pandas_type_converter import convert_to_pandas_type
+from datacontract.export.spark_converter import to_spark_data_type
+from datacontract.export.sql_type_converter import convert_to_sql_type
+from datacontract.model.data_contract_specification import (
+    DataContractSpecification,
+    Field,
+    Quality,
+)
+
+
+class GreatExpectationsEngine(Enum):
+    """Enum to represent the type of data engine for expectations.
+
+    Attributes:
+        pandas (str): Represents the Pandas engine type.
+        spark (str): Represents the Spark engine type.
+        sql (str): Represents the SQL engine type.
+    """
+
+    pandas = "pandas"
+    spark = "spark"
+    sql = "sql"
+
 
+class GreatExpectationsExporter(Exporter):
+    """Exporter class to convert data contracts to Great Expectations suites.
+
+    Methods:
+        export: Converts a data contract model to a Great Expectations suite.
+
+    """
 
-class GreateExpectationsExporter(Exporter):
     def export(self, data_contract, model, server, sql_server_type, export_args) -> dict:
+        """Exports a data contract model to a Great Expectations suite.
+
+        Args:
+            data_contract (DataContractSpecification): The data contract specification.
+            model (str): The model name to export.
+            server (str): The server information.
+            sql_server_type (str): Type of SQL server (e.g., "snowflake").
+            export_args (dict): Additional arguments for export, such as "suite_name" and "engine".
+
+        Returns:
+            dict: A dictionary representation of the Great Expectations suite.
+        """
+        expectation_suite_name = export_args.get("suite_name")
+        engine = export_args.get("engine")
         model_name, model_value = _check_models_for_export(data_contract, model, self.export_format)
-        return to_great_expectations(
-            data_contract,
-            model_name,
-        )
+        sql_server_type = "snowflake" if sql_server_type == "auto" else sql_server_type
+        return to_great_expectations(data_contract, model_name, expectation_suite_name, engine, sql_server_type)
 
 
-def to_great_expectations(data_contract_spec: DataContractSpecification, model_key: str) -> str:
-    """
-    Convert each model in the contract to a Great Expectation suite
-    @param data_contract_spec: data contract to export to great expectations
-    @param model_key: model to great expectations to
-    @return: a dictionary of great expectation suites
+def to_great_expectations(
+    data_contract_spec: DataContractSpecification,
+    model_key: str,
+    expectation_suite_name: str | None = None,
+    engine: str | None = None,
+    sql_server_type: str = "snowflake",
+) -> str:
+    """Converts a data contract model to a Great Expectations suite.
+
+    Args:
+        data_contract_spec (DataContractSpecification): The data contract specification.
+        model_key (str): The model key.
+        expectation_suite_name (str | None): Optional suite name for the expectations.
+        engine (str | None): Optional engine type (e.g., "pandas", "spark").
+        sql_server_type (str): The type of SQL server (default is "snowflake").
+
+    Returns:
+        str: JSON string of the Great Expectations suite.
     """
     expectations = []
+    if not expectation_suite_name:
+        expectation_suite_name = "{model_key}.{contract_version}".format(
+            model_key=model_key, contract_version=data_contract_spec.info.version
+        )
     model_value = data_contract_spec.models.get(model_key)
     quality_checks = get_quality_checks(data_contract_spec.quality)
-    expectations.extend(model_to_expectations(model_value.fields))
+    expectations.extend(model_to_expectations(model_value.fields, engine, sql_server_type))
     expectations.extend(checks_to_expectations(quality_checks, model_key))
-    model_expectation_suite = to_suite(model_key, data_contract_spec.info.version, expectations)
+    model_expectation_suite = to_suite(expectations, expectation_suite_name)
 
     return model_expectation_suite
 
 
-def to_suite(
-    model_key: str,
-    contract_version: str,
-    expectations: List[Dict[str, Any]],
-) -> str:
+def to_suite(expectations: List[Dict[str, Any]], expectation_suite_name: str) -> str:
+    """Converts a list of expectations to a JSON-formatted suite.
+
+    Args:
+        expectations (List[Dict[str, Any]]): List of expectations.
+        expectation_suite_name (str): Name of the expectation suite.
+
+    Returns:
+        str: JSON string of the expectation suite.
+    """
     return json.dumps(
         {
             "data_asset_type": "null",
-            "expectation_suite_name": "user-defined.{model_key}.{contract_version}".format(
-                model_key=model_key, contract_version=contract_version
-            ),
+            "expectation_suite_name": expectation_suite_name,
             "expectations": expectations,
             "meta": {},
         },
@@ -51,22 +120,53 @@ def to_suite(
     )
 
 
-def model_to_expectations(fields: Dict[str, Field]) -> List[Dict[str, Any]]:
-    """
-    Convert the model information to expectations
-    @param fields: model field
-    @return: list of expectations
+def model_to_expectations(fields: Dict[str, Field], engine: str | None, sql_server_type: str) -> List[Dict[str, Any]]:
+    """Converts model fields to a list of expectations.
+
+    Args:
+        fields (Dict[str, Field]): Dictionary of model fields.
+        engine (str | None): Engine type (e.g., "pandas", "spark").
+        sql_server_type (str): SQL server type.
+
+    Returns:
+        List[Dict[str, Any]]: List of expectations.
     """
     expectations = []
     add_column_order_exp(fields, expectations)
     for field_name, field in fields.items():
-        add_field_expectations(field_name, field, expectations)
+        add_field_expectations(field_name, field, expectations, engine, sql_server_type)
     return expectations
 
 
-def add_field_expectations(field_name, field: Field, expectations: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+def add_field_expectations(
+    field_name,
+    field: Field,
+    expectations: List[Dict[str, Any]],
+    engine: str | None,
+    sql_server_type: str,
+) -> List[Dict[str, Any]]:
+    """Adds expectations for a specific field based on its properties.
+
+    Args:
+        field_name (str): The name of the field.
+        field (Field): The field object.
+        expectations (List[Dict[str, Any]]): The expectations list to update.
+        engine (str | None): Engine type (e.g., "pandas", "spark").
+        sql_server_type (str): SQL server type.
+
+    Returns:
+        List[Dict[str, Any]]: Updated list of expectations.
+    """
     if field.type is not None:
-        expectations.append(to_column_types_exp(field_name, field.type))
+        if engine == GreatExpectationsEngine.spark.value:
+            field_type = to_spark_data_type(field).__class__.__name__
+        elif engine == GreatExpectationsEngine.pandas.value:
+            field_type = convert_to_pandas_type(field)
+        elif engine == GreatExpectationsEngine.sql.value:
+            field_type = convert_to_sql_type(field, sql_server_type)
+        else:
+            field_type = field.type
+        expectations.append(to_column_types_exp(field_name, field_type))
     if field.unique:
         expectations.append(to_column_unique_exp(field_name))
     if field.maxLength is not None or field.minLength is not None:
@@ -74,11 +174,16 @@ def add_field_expectations(field_name, field: Field, expectations: List[Dict[str
     if field.minimum is not None or field.maximum is not None:
         expectations.append(to_column_min_max_exp(field_name, field.minimum, field.maximum))
 
-    # TODO: all constraints
     return expectations
 
 
 def add_column_order_exp(fields: Dict[str, Field], expectations: List[Dict[str, Any]]):
+    """Adds expectation for column ordering.
+
+    Args:
+        fields (Dict[str, Field]): Dictionary of fields.
+        expectations (List[Dict[str, Any]]): The expectations list to update.
+    """
     expectations.append(
         {
             "expectation_type": "expect_table_columns_to_match_ordered_list",
@@ -89,6 +194,15 @@ def add_column_order_exp(fields: Dict[str, Field], expectations: List[Dict[str,
 
 
 def to_column_types_exp(field_name, field_type) -> Dict[str, Any]:
+    """Creates a column type expectation.
+
+    Args:
+        field_name (str): The name of the field.
+        field_type (str): The type of the field.
+
+    Returns:
+        Dict[str, Any]: Column type expectation.
+    """
     return {
         "expectation_type": "expect_column_values_to_be_of_type",
         "kwargs": {"column": field_name, "type_": field_type},
@@ -97,18 +211,54 @@ def to_column_types_exp(field_name, field_type) -> Dict[str, Any]:
 
 
 def to_column_unique_exp(field_name) -> Dict[str, Any]:
-    return {"expectation_type": "expect_column_values_to_be_unique", "kwargs": {"column": field_name}, "meta": {}}
+    """Creates a column uniqueness expectation.
+
+    Args:
+        field_name (str): The name of the field.
+
+    Returns:
+        Dict[str, Any]: Column uniqueness expectation.
+    """
+    return {
+        "expectation_type": "expect_column_values_to_be_unique",
+        "kwargs": {"column": field_name},
+        "meta": {},
+    }
 
 
 def to_column_length_exp(field_name, min_length, max_length) -> Dict[str, Any]:
+    """Creates a column length expectation.
+
+    Args:
+        field_name (str): The name of the field.
+        min_length (int | None): Minimum length.
+        max_length (int | None): Maximum length.
+
+    Returns:
+        Dict[str, Any]: Column length expectation.
+    """
     return {
         "expectation_type": "expect_column_value_lengths_to_be_between",
-        "kwargs": {"column": field_name, "min_value": min_length, "max_value": max_length},
+        "kwargs": {
+            "column": field_name,
+            "min_value": min_length,
+            "max_value": max_length,
+        },
         "meta": {},
     }
 
 
 def to_column_min_max_exp(field_name, minimum, maximum) -> Dict[str, Any]:
+    """Creates a column min-max value expectation.
+
+    Args:
+        field_name (str): The name of the field.
+        minimum (float | None): Minimum value.
+        maximum (float | None): Maximum value.
+
+    Returns:
+        Dict[str, Any]: Column min-max value expectation.
+    """
     return {
         "expectation_type": "expect_column_values_to_be_between",
         "kwargs": {"column": field_name, "min_value": minimum, "max_value": maximum},
@@ -117,6 +267,14 @@ def to_column_min_max_exp(field_name, minimum, maximum) -> Dict[str, Any]:
 
 
 def get_quality_checks(quality: Quality) -> Dict[str, Any]:
+    """Retrieves quality checks defined in a data contract.
+
+    Args:
+        quality (Quality): Quality object from the data contract.
+
+    Returns:
+        Dict[str, Any]: Dictionary of quality checks.
+    """
     if quality is None:
         return {}
     if quality.type is None:
@@ -131,11 +289,14 @@ def get_quality_checks(quality: Quality) -> Dict[str, Any]:
 
 
 def checks_to_expectations(quality_checks: Dict[str, Any], model_key: str) -> List[Dict[str, Any]]:
-    """
-    Get the quality definition for each model to the model expectation list
-    @param quality_checks: dictionary of quality checks by model
-    @param model_key: id of the model
-    @return: the list of expectations for that model
+    """Converts quality checks to a list of expectations.
+
+    Args:
+        quality_checks (Dict[str, Any]): Dictionary of quality checks by model.
+        model_key (str): The model key.
+
+    Returns:
+        List[Dict[str, Any]]: List of expectations for the model.
     """
     if quality_checks is None or model_key not in quality_checks:
         return []
@@ -148,3 +309,4 @@ def checks_to_expectations(quality_checks: Dict[str, Any], model_key: str) -> Li
     if isinstance(model_quality_checks, str):
         expectation_list = json.loads(model_quality_checks)
         return expectation_list
+    return []

datacontract/export/html_export.py

@@ -7,8 +7,8 @@ import pytz
 import yaml
 from jinja2 import Environment, PackageLoader, select_autoescape
 
-from datacontract.model.data_contract_specification import DataContractSpecification
 from datacontract.export.exporter import Exporter
+from datacontract.model.data_contract_specification import DataContractSpecification
 
 
 class HtmlExporter(Exporter):

datacontract/export/jsonschema_converter.py

@@ -1,9 +1,8 @@
 import json
 from typing import Dict
 
-from datacontract.model.data_contract_specification import DataContractSpecification, Model, Field
-
 from datacontract.export.exporter import Exporter, _check_models_for_export
+from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model
 
 
 class JsonSchemaExporter(Exporter):
@@ -51,6 +50,8 @@ def to_property(field: Field) -> dict:
             property["type"] = json_type
     if json_format is not None:
         property["format"] = json_format
+    if field.primaryKey:
+        property["primaryKey"] = field.primaryKey
     if field.unique:
         property["unique"] = True
     if json_type == "object":

datacontract/export/odcs_v2_exporter.py

@@ -2,8 +2,8 @@ from typing import Dict
 
 import yaml
 
-from datacontract.model.data_contract_specification import DataContractSpecification, Model, Field
 from datacontract.export.exporter import Exporter
+from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model
 
 
 class OdcsV2Exporter(Exporter):

datacontract/export/odcs_v3_exporter.py

@@ -3,7 +3,7 @@ from typing import Dict
 import yaml
 
 from datacontract.export.exporter import Exporter
-from datacontract.model.data_contract_specification import DataContractSpecification, Model, Field
+from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model
 
 
 class OdcsV3Exporter(Exporter):

datacontract/export/pandas_type_converter.py

@@ -0,0 +1,40 @@
+"""
+Module for converting data contract field types to corresponding pandas data types.
+"""
+
+from datacontract.model.data_contract_specification import Field
+
+
+def convert_to_pandas_type(field: Field) -> str:
+    """
+    Convert a data contract field type to the equivalent pandas data type.
+
+    Parameters:
+    ----------
+    field : Field
+        A Field object containing metadata about the data type of the field.
+
+    Returns:
+    -------
+    str
+        The corresponding pandas data type as a string.
+    """
+    field_type = field.type
+
+    if field_type in ["string", "varchar", "text"]:
+        return "str"
+    if field_type in ["integer", "int"]:
+        return "int32"
+    if field_type == "long":
+        return "int64"
+    if field_type == "float":
+        return "float32"
+    if field_type in ["number", "decimal", "numeric", "double"]:
+        return "float64"
+    if field_type == "boolean":
+        return "bool"
+    if field_type in ["timestamp", "timestamp_tz", "timestamp_ntz", "date"]:
+        return "datetime64[ns]"
+    if field_type == "bytes":
+        return "object"
+    return "object"

datacontract/export/protobuf_converter.py

@@ -1,5 +1,5 @@
-from datacontract.model.data_contract_specification import DataContractSpecification
 from datacontract.export.exporter import Exporter
+from datacontract.model.data_contract_specification import DataContractSpecification
 
 
 class ProtoBufExporter(Exporter):

datacontract/export/rdf_converter.py

@@ -1,9 +1,8 @@
 from pydantic import BaseModel
-from rdflib import Graph, Literal, BNode, RDF, URIRef, Namespace
-
-from datacontract.model.data_contract_specification import DataContractSpecification
+from rdflib import RDF, BNode, Graph, Literal, Namespace, URIRef
 
 from datacontract.export.exporter import Exporter
+from datacontract.model.data_contract_specification import DataContractSpecification
 
 
 class RdfExporter(Exporter):
@@ -58,8 +57,8 @@ def to_rdf(data_contract_spec: DataContractSpecification, base) -> Graph:
     else:
         g = Graph(base=Namespace(""))
 
-    dc = Namespace("https://datacontract.com/DataContractSpecification/0.9.2/")
-    dcx = Namespace("https://datacontract.com/DataContractSpecification/0.9.2/Extension/")
+    dc = Namespace("https://datacontract.com/DataContractSpecification/1.1.0/")
+    dcx = Namespace("https://datacontract.com/DataContractSpecification/1.1.0/Extension/")
 
     g.bind("dc", dc)
     g.bind("dcx", dcx)

datacontract/export/sodacl_converter.py

@@ -62,12 +62,16 @@ def to_checks(model_key, model_value, server_type: str, check_types: bool):
         if field.enum is not None and len(field.enum) > 0:
             checks.append(check_field_enum(field_name, field.enum, quote_field_name))
         if field.quality is not None and len(field.quality) > 0:
-            checks.append(check_quality_list(model_key, field_name, field.quality))
+            quality_list = check_quality_list(model_key, field_name, field.quality)
+            if (quality_list is not None) and len(quality_list) > 0:
+                checks.append(quality_list)
         # TODO references: str = None
         # TODO format
 
     if model_value.quality is not None and len(model_value.quality) > 0:
-        checks.append(check_quality_list(model_key, None, model_value.quality))
+        quality_list = check_quality_list(model_key, None, model_value.quality)
+        if (quality_list is not None) and len(quality_list) > 0:
+            checks.append(quality_list)
 
     checks_for_model_key = f"checks for {model_key}"
 

datacontract/export/spark_converter.py

@@ -1,10 +1,11 @@
 from pyspark.sql import types
+
+from datacontract.export.exporter import Exporter
 from datacontract.model.data_contract_specification import (
     DataContractSpecification,
-    Model,
     Field,
+    Model,
 )
-from datacontract.export.exporter import Exporter
 
 
 class SparkExporter(Exporter):
@@ -102,11 +103,11 @@ def to_struct_field(field: Field, field_name: str) -> types.StructField:
     Returns:
         types.StructField: The corresponding Spark StructField.
     """
-    data_type = to_data_type(field)
+    data_type = to_spark_data_type(field)
     return types.StructField(name=field_name, dataType=data_type, nullable=not field.required)
 
 
-def to_data_type(field: Field) -> types.DataType:
+def to_spark_data_type(field: Field) -> types.DataType:
     """
     Convert a field to a Spark DataType.
 
@@ -120,11 +121,11 @@ def to_data_type(field: Field) -> types.DataType:
     if field_type is None or field_type in ["null"]:
         return types.NullType()
     if field_type == "array":
-        return types.ArrayType(to_data_type(field.items))
+        return types.ArrayType(to_spark_data_type(field.items))
     if field_type in ["object", "record", "struct"]:
         return types.StructType(to_struct_type(field.fields))
     if field_type == "map":
-        return types.MapType(to_data_type(field.keys), to_data_type(field.values))
+        return types.MapType(to_spark_data_type(field.keys), to_spark_data_type(field.values))
     if field_type in ["string", "varchar", "text"]:
         return types.StringType()
     if field_type in ["number", "decimal", "numeric"]:

datacontract/export/sql_converter.py

@@ -1,8 +1,7 @@
+from datacontract.export.exporter import Exporter, _check_models_for_export, _determine_sql_server_type
 from datacontract.export.sql_type_converter import convert_to_sql_type
 from datacontract.model.data_contract_specification import DataContractSpecification, Model
 
-from datacontract.export.exporter import Exporter, _check_models_for_export, _determine_sql_server_type
-
 
 class SqlExporter(Exporter):
     def export(self, data_contract, model, server, sql_server_type, export_args) -> dict:

datacontract/export/sqlalchemy_converter.py

@@ -2,8 +2,7 @@ import ast
 import typing
 
 import datacontract.model.data_contract_specification as spec
-from datacontract.export.exporter import Exporter
-from datacontract.export.exporter import _determine_sql_server_type
+from datacontract.export.exporter import Exporter, _determine_sql_server_type
 
 
 class SQLAlchemyExporter(Exporter):

datacontract/export/terraform_converter.py

@@ -1,7 +1,7 @@
 import re
 
-from datacontract.model.data_contract_specification import DataContractSpecification, Server
 from datacontract.export.exporter import Exporter
+from datacontract.model.data_contract_specification import DataContractSpecification, Server
 
 
 class TerraformExporter(Exporter):

datacontract/imports/avro_importer.py

@@ -3,7 +3,7 @@ from typing import Dict, List
 import avro.schema
 
 from datacontract.imports.importer import Importer
-from datacontract.model.data_contract_specification import DataContractSpecification, Model, Field
+from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model
 from datacontract.model.exceptions import DataContractException
 
 

datacontract/imports/bigquery_importer.py

@@ -3,7 +3,7 @@ import logging
 from typing import List
 
 from datacontract.imports.importer import Importer
-from datacontract.model.data_contract_specification import DataContractSpecification, Model, Field
+from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model
 from datacontract.model.exceptions import DataContractException
 
 

datacontract/imports/dbml_importer.py

@@ -1,11 +1,11 @@
-from pydbml import PyDBML, Database
 from typing import List
 
+from pydbml import Database, PyDBML
 from pyparsing import ParseException
 
 from datacontract.imports.importer import Importer
 from datacontract.imports.sql_importer import map_type_from_sql
-from datacontract.model.data_contract_specification import DataContractSpecification, Model, Field
+from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model
 from datacontract.model.exceptions import DataContractException
 
 

datacontract/imports/dbt_importer.py

@@ -1,11 +1,12 @@
 import json
 from typing import TypedDict
 
-from datacontract.imports.importer import Importer
-from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model
 from dbt.artifacts.resources.v1.components import ColumnInfo
 from dbt.contracts.graph.manifest import Manifest
 
+from datacontract.imports.importer import Importer
+from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model
+
 
 class DBTImportArgs(TypedDict, total=False):
     """

datacontract/imports/glue_importer.py

@@ -1,11 +1,13 @@
-import boto3
-from typing import List, Dict, Generator
 import re
+from typing import Dict, Generator, List
+
+import boto3
+
 from datacontract.imports.importer import Importer
 from datacontract.model.data_contract_specification import (
     DataContractSpecification,
-    Model,
     Field,
+    Model,
     Server,
 )
 

datacontract/imports/iceberg_importer.py

@@ -1,12 +1,11 @@
-from typing import Dict, Any
+from typing import Any, Dict
 
-from datacontract.imports.importer import Importer
-from datacontract.model.data_contract_specification import DataContractSpecification, Model, Field
-
-from pyiceberg.schema import Schema
-from pyiceberg import types as iceberg_types
 from pydantic import ValidationError
+from pyiceberg import types as iceberg_types
+from pyiceberg.schema import Schema
 
+from datacontract.imports.importer import Importer
+from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model
 from datacontract.model.exceptions import DataContractException
 
 

datacontract/imports/importer.py

@@ -30,6 +30,7 @@ class ImportFormat(str, Enum):
     unity = "unity"
     spark = "spark"
     iceberg = "iceberg"
+    parquet = "parquet"
 
     @classmethod
     def get_supported_formats(cls):

datacontract/imports/importer_factory.py

@@ -1,6 +1,7 @@
 import importlib.util
 import sys
-from datacontract.imports.importer import ImportFormat, Importer
+
+from datacontract.imports.importer import Importer, ImportFormat
 
 
 class ImporterFactory:
@@ -98,3 +99,8 @@ importer_factory.register_lazy_importer(
     module_path="datacontract.imports.iceberg_importer",
     class_name="IcebergImporter",
 )
+importer_factory.register_lazy_importer(
+    name=ImportFormat.parquet,
+    module_path="datacontract.imports.parquet_importer",
+    class_name="ParquetImporter",
+)