datacontract-cli 0.9.7__py3-none-any.whl → 0.9.8__py3-none-any.whl

datacontract/engines/soda/check_soda_execute.py

@@ -16,8 +16,7 @@ from datacontract.engines.soda.connections.snowflake import \
 from datacontract.export.sodacl_converter import to_sodacl_yaml
 from datacontract.model.data_contract_specification import \
     DataContractSpecification, Server
-from datacontract.model.run import \
-    Run, Check, Log
+from datacontract.model.run import Run, Check, Log
 
 
 def check_soda_execute(run: Run, data_contract: DataContractSpecification, server: Server, spark, tmp_dir):
@@ -34,13 +33,15 @@ def check_soda_execute(run: Run, data_contract: DataContractSpecification, serve
             scan.add_duckdb_connection(duckdb_connection=con, data_source_name=server.type)
             scan.set_data_source_name(server.type)
         else:
-            run.checks.append(Check(
-                type="general",
-                name="Check that format is supported",
-                result="warning",
-                reason=f"Format {server.format} not yet supported by datacontract CLI",
-                engine="datacontract",
-            ))
+            run.checks.append(
+                Check(
+                    type="general",
+                    name="Check that format is supported",
+                    result="warning",
+                    reason=f"Format {server.format} not yet supported by datacontract CLI",
+                    engine="datacontract",
+                )
+            )
             run.log_warn(f"Format {server.format} not yet supported by datacontract CLI")
             return
     elif server.type == "snowflake":
@@ -73,21 +74,24 @@ def check_soda_execute(run: Run, data_contract: DataContractSpecification, serve
         scan.set_data_source_name(server.type)
 
     else:
-        run.checks.append(Check(
-            type="general",
-            name="Check that server type is supported",
-            result="warning",
-            reason=f"Server type {server.type} not yet supported by datacontract CLI",
-            engine="datacontract-cli",
-        ))
+        run.checks.append(
+            Check(
+                type="general",
+                name="Check that server type is supported",
+                result="warning",
+                reason=f"Server type {server.type} not yet supported by datacontract CLI",
+                engine="datacontract-cli",
+            )
+        )
         run.log_warn(f"Server type {server.type} not yet supported by datacontract CLI")
         return
 
     # Don't check types for json format, as they are checked with json schema
     # Don't check types for avro format, as they are checked with avro schema
     # Don't check types for csv format, as they are hard to detect
+    server_type = server.type
     check_types = server.format != "json" and server.format != "csv" and server.format != "avro"
-    sodacl_yaml_str = to_sodacl_yaml(data_contract, check_types)
+    sodacl_yaml_str = to_sodacl_yaml(data_contract, server_type, check_types)
     # print("sodacl_yaml_str:\n" + sodacl_yaml_str)
     scan.add_sodacl_yaml_str(sodacl_yaml_str)
 
@@ -102,9 +106,12 @@ def check_soda_execute(run: Run, data_contract: DataContractSpecification, serve
     for c in scan_results.get("checks"):
         check = Check(
             type="schema",
-            result="passed" if c.get("outcome") == "pass" else "failed" if c.get("outcome") == "fail" else c.get(
-                "outcome"),
-            reason=', '.join(c.get("outcomeReasons")),
+            result="passed"
+            if c.get("outcome") == "pass"
+            else "failed"
+            if c.get("outcome") == "fail"
+            else c.get("outcome"),
+            reason=", ".join(c.get("outcomeReasons")),
             name=c.get("name"),
             model=c.get("table"),
             field=c.get("column"),
@@ -114,21 +121,25 @@ def check_soda_execute(run: Run, data_contract: DataContractSpecification, serve
         run.checks.append(check)
 
     for log in scan_results.get("logs"):
-        run.logs.append(Log(
-            timestamp=log.get("timestamp"),
-            level=log.get("level"),
-            message=log.get("message"),
-        ))
+        run.logs.append(
+            Log(
+                timestamp=log.get("timestamp"),
+                level=log.get("level"),
+                message=log.get("message"),
+            )
+        )
 
     if scan.has_error_logs():
         run.log_warn("Engine soda-core has errors. See the logs for details.")
-        run.checks.append(Check(
-            type="general",
-            name="Execute quality checks",
-            result="warning",
-            reason=f"Engine soda-core has errors. See the logs for details.",
-            engine="soda-core",
-        ))
+        run.checks.append(
+            Check(
+                type="general",
+                name="Execute quality checks",
+                result="warning",
+                reason="Engine soda-core has errors. See the logs for details.",
+                engine="soda-core",
+            )
+        )
         return
 
 
@@ -138,10 +149,10 @@ def update_reason(check, c):
         return
     if check.reason is not None and check.reason != "":
         return
-    for block in c['diagnostics']['blocks']:
-        if block['title'] == 'Diagnostics':
+    for block in c["diagnostics"]["blocks"]:
+        if block["title"] == "Diagnostics":
             # Extract and print the 'text' value
-            diagnostics_text = block['text']
+            diagnostics_text = block["text"]
             print(diagnostics_text)
             diagnostics_text_split = diagnostics_text.split(":icon-fail: ")
             if len(diagnostics_text_split) > 1:

datacontract/engines/soda/connections/bigquery.py

@@ -1,18 +1,20 @@
 import os
+
 import yaml
 
+
 # https://docs.soda.io/soda/connect-bigquery.html#authentication-methods
 def to_bigquery_soda_configuration(server):
     # with service account key, using an external json file
     soda_configuration = {
         f"data_source {server.type}": {
             "type": "bigquery",
-            "account_info_json_path": os.getenv('DATACONTRACT_BIGQUERY_ACCOUNT_INFO_JSON_PATH'),
+            "account_info_json_path": os.getenv("DATACONTRACT_BIGQUERY_ACCOUNT_INFO_JSON_PATH"),
             "auth_scopes": ["https://www.googleapis.com/auth/bigquery"],
             "project_id": server.project,
-            "dataset": server.dataset
+            "dataset": server.dataset,
         }
     }
 
     soda_configuration_str = yaml.dump(soda_configuration)
-    return soda_configuration_str
+    return soda_configuration_str

datacontract/engines/soda/connections/dask.py

@@ -1,4 +1,3 @@
-
 # def add_s3_connection_dask_json(data_contract, scan, server):
 #     s3_access_key_id = os.getenv('DATACONTRACT_S3_ACCESS_KEY_ID')
 #     s3_secret_access_key = os.getenv('DATACONTRACT_S3_SECRET_ACCESS_KEY')

datacontract/engines/soda/connections/databricks.py

@@ -11,8 +11,8 @@ def to_databricks_soda_configuration(server):
             "host": server.host,
             "catalog": server.catalog,
             "schema": server.schema_,
-            "http_path": os.getenv('DATACONTRACT_DATABRICKS_HTTP_PATH'),
-            "token": os.getenv('DATACONTRACT_DATABRICKS_TOKEN'),
+            "http_path": os.getenv("DATACONTRACT_DATABRICKS_HTTP_PATH"),
+            "token": os.getenv("DATACONTRACT_DATABRICKS_TOKEN"),
         }
     }
 

datacontract/engines/soda/connections/duckdb.py

@@ -15,7 +15,7 @@ def get_duckdb_connection(data_contract, server):
     for model_name in data_contract.models:
         model_path = path
         if "{model}" in model_path:
-            model_path = model_path.format(model = model_name)
+            model_path = model_path.format(model=model_name)
         logging.info(f"Creating table {model_name} for {model_path}")
 
         if server.format == "json":
@@ -39,9 +39,9 @@ def get_duckdb_connection(data_contract, server):
 
 
 def setup_s3_connection(con, server):
-    s3_region = os.getenv('DATACONTRACT_S3_REGION')
-    s3_access_key_id = os.getenv('DATACONTRACT_S3_ACCESS_KEY_ID')
-    s3_secret_access_key = os.getenv('DATACONTRACT_S3_SECRET_ACCESS_KEY')
+    s3_region = os.getenv("DATACONTRACT_S3_REGION")
+    s3_access_key_id = os.getenv("DATACONTRACT_S3_ACCESS_KEY_ID")
+    s3_secret_access_key = os.getenv("DATACONTRACT_S3_SECRET_ACCESS_KEY")
     # con.install_extension("httpfs")
     # con.load_extension("httpfs")
     if server.endpointUrl is not None:

datacontract/engines/soda/connections/kafka.py

@@ -4,7 +4,23 @@ import pyspark.sql.functions as fn
 from pyspark.sql import SparkSession
 from pyspark.sql.avro.functions import from_avro
 from pyspark.sql.functions import from_json, col
-from pyspark.sql.types import *
+from pyspark.sql.types import (
+    StructType,
+    DataType,
+    NullType,
+    ArrayType,
+    BinaryType,
+    DateType,
+    TimestampNTZType,
+    TimestampType,
+    BooleanType,
+    LongType,
+    IntegerType,
+    DoubleType,
+    DecimalType,
+    StringType,
+    StructField,
+)
 
 from datacontract.export.avro_converter import to_avro_schema_json
 from datacontract.model.data_contract_specification import \
@@ -15,14 +31,18 @@ from datacontract.model.exceptions import DataContractException
 def create_spark_session(tmp_dir) -> SparkSession:
     # TODO: Update dependency versions when updating pyspark
     # TODO: add protobuf library
-    spark = SparkSession.builder.appName("datacontract") \
-        .config("spark.sql.warehouse.dir", tmp_dir + "/spark-warehouse") \
-        .config("spark.streaming.stopGracefullyOnShutdown", True) \
-        .config('spark.jars.packages',
-                'org.apache.spark:spark-sql-kafka-0-10_2.12:3.5.0,org.apache.spark:spark-avro_2.12:3.5.0') \
+    spark = (
+        SparkSession.builder.appName("datacontract")
+        .config("spark.sql.warehouse.dir", tmp_dir + "/spark-warehouse")
+        .config("spark.streaming.stopGracefullyOnShutdown", True)
+        .config(
+            "spark.jars.packages",
+            "org.apache.spark:spark-sql-kafka-0-10_2.12:3.5.0,org.apache.spark:spark-avro_2.12:3.5.0",
+        )
         .getOrCreate()
+    )
     spark.sparkContext.setLogLevel("WARN")
-    print(f'Using PySpark version {spark.version}')
+    print(f"Using PySpark version {spark.version}")
     return spark
 
 
@@ -32,14 +52,14 @@ def read_kafka_topic(spark: SparkSession, data_contract: DataContractSpecificati
     auth_options = get_auth_options()
 
     # read full kafka topic
-    df = spark \
-        .read \
-        .format("kafka") \
-        .options(**auth_options) \
-        .option("kafka.bootstrap.servers", host) \
-        .option("subscribe", topic) \
-        .option("startingOffsets", "earliest") \
+    df = (
+        spark.read.format("kafka")
+        .options(**auth_options)
+        .option("kafka.bootstrap.servers", host)
+        .option("subscribe", topic)
+        .option("startingOffsets", "earliest")
         .load()
+    )
     # TODO a warning if none or multiple models
     model_name, model = next(iter(data_contract.models.items()))
     if server.format == "avro":
@@ -73,8 +93,8 @@ def read_kafka_topic(spark: SparkSession, data_contract: DataContractSpecificati
 
 
 def get_auth_options():
-    kafka_sasl_username = os.getenv('DATACONTRACT_KAFKA_SASL_USERNAME')
-    kafka_sasl_password = os.getenv('DATACONTRACT_KAFKA_SASL_PASSWORD')
+    kafka_sasl_username = os.getenv("DATACONTRACT_KAFKA_SASL_USERNAME")
+    kafka_sasl_password = os.getenv("DATACONTRACT_KAFKA_SASL_PASSWORD")
     if kafka_sasl_username is None:
         auth_options = {}
     else:
@@ -130,4 +150,3 @@ def to_struct_field(field_name: str, field: Field) -> StructField:
         data_type = DataType()
 
     return StructField(field_name, data_type, nullable=not field.required)
-

datacontract/engines/soda/connections/postgres.py

@@ -10,12 +10,12 @@ def to_postgres_soda_configuration(server):
             "type": "postgres",
             "host": server.host,
             "port": str(server.port),
-            "username": os.getenv('DATACONTRACT_POSTGRES_USERNAME'),
-            "password": os.getenv('DATACONTRACT_POSTGRES_PASSWORD'),
+            "username": os.getenv("DATACONTRACT_POSTGRES_USERNAME"),
+            "password": os.getenv("DATACONTRACT_POSTGRES_PASSWORD"),
             "database": server.database,
             "schema": server.schema_,
         }
     }
 
     soda_configuration_str = yaml.dump(soda_configuration)
-    return soda_configuration_str
+    return soda_configuration_str

datacontract/engines/soda/connections/snowflake.py

@@ -7,13 +7,13 @@ def to_snowflake_soda_configuration(server):
     soda_configuration = {
         f"data_source {server.type}": {
             "type": "snowflake",
-            "username": os.getenv('DATACONTRACT_SNOWFLAKE_USERNAME'),
-            "password": os.getenv('DATACONTRACT_SNOWFLAKE_PASSWORD'),
-            "role": os.getenv('DATACONTRACT_SNOWFLAKE_ROLE'),
+            "username": os.getenv("DATACONTRACT_SNOWFLAKE_USERNAME"),
+            "password": os.getenv("DATACONTRACT_SNOWFLAKE_PASSWORD"),
+            "role": os.getenv("DATACONTRACT_SNOWFLAKE_ROLE"),
             "account": server.account,
             "database": server.database,
             "schema": server.schema_,
-            "warehouse": os.getenv('DATACONTRACT_SNOWFLAKE_WAREHOUSE'),
+            "warehouse": os.getenv("DATACONTRACT_SNOWFLAKE_WAREHOUSE"),
             "connection_timeout": 5,  # minutes
         }
     }

datacontract/export/avro_converter.py

@@ -6,16 +6,14 @@ from datacontract.model.data_contract_specification import Field
 def to_avro_schema(model_name, model) -> dict:
     return to_avro_record(model_name, model.fields, model.description)
 
+
 def to_avro_schema_json(model_name, model) -> str:
     schema = to_avro_schema(model_name, model)
     return json.dumps(schema, indent=2, sort_keys=False)
 
 
 def to_avro_record(name, fields, description) -> dict:
-    schema = {
-        "type": "record",
-        "name": name
-    }
+    schema = {"type": "record", "name": name}
     if description is not None:
         schema["doc"] = description
     schema["fields"] = to_avro_fields(fields)
@@ -30,9 +28,7 @@ def to_avro_fields(fields):
 
 
 def to_avro_field(field, field_name):
-    avro_field = {
-        "name": field_name
-    }
+    avro_field = {"name": field_name}
     if field.description is not None:
         avro_field["doc"] = field.description
     avro_field["type"] = to_avro_type(field, field_name)

datacontract/export/avro_idl_converter.py

@@ -1,12 +1,14 @@
-from datacontract.model.data_contract_specification import DataContractSpecification, Field
-from datacontract.lint.resolve import inline_definitions_into_data_contract
+import typing
 from dataclasses import dataclass
 from enum import Enum
-import typing
 from io import StringIO
 
+from datacontract.lint.resolve import inline_definitions_into_data_contract
+from datacontract.model.data_contract_specification import \
+    DataContractSpecification, Field
 from datacontract.model.exceptions import DataContractException
 
+
 def to_avro_idl(contract: DataContractSpecification) -> str:
     """Serialize the provided data contract specification into an Avro IDL string.
 
@@ -18,6 +20,7 @@ def to_avro_idl(contract: DataContractSpecification) -> str:
     to_avro_idl_stream(contract, stream)
     return stream.getvalue()
 
+
 def to_avro_idl_stream(contract: DataContractSpecification, stream: typing.TextIO):
     """Serialize the provided data contract specification into Avro IDL."""
     ir = _contract_to_avro_idl_ir(contract)
@@ -28,6 +31,7 @@ def to_avro_idl_stream(contract: DataContractSpecification, stream: typing.TextI
         _write_model_type(model_type, stream)
     stream.write("}\n")
 
+
 class AvroPrimitiveType(Enum):
     int = "int"
     long = "long"
@@ -38,49 +42,71 @@ class AvroPrimitiveType(Enum):
     null = "null"
     bytes = "bytes"
 
+
 class AvroLogicalType(Enum):
     decimal = "decimal"
     date = "date"
     time_ms = "time_ms"
     timestamp_ms = "timestamp_ms"
 
+
 @dataclass
 class AvroField:
     name: str
     required: bool
     description: typing.Optional[str]
 
+
 @dataclass
 class AvroPrimitiveField(AvroField):
     type: typing.Union[AvroPrimitiveType, AvroLogicalType]
 
+
 @dataclass
 class AvroComplexField(AvroField):
     subfields: list[AvroField]
 
+
 @dataclass
 class AvroArrayField(AvroField):
     type: AvroField
 
+
 @dataclass
 class AvroModelType:
     name: str
     description: typing.Optional[str]
     fields: list[AvroField]
 
+
 @dataclass
 class AvroIDLProtocol:
     name: typing.Optional[str]
     description: typing.Optional[str]
     model_types: list[AvroModelType]
 
-avro_primitive_types = set(["string", "text", "varchar",
-                            "float", "double", "int",
-                            "integer", "long", "bigint",
-                            "boolean", "timestamp_ntz",
-                            "timestamp", "timestamp_tz",
-                            "date", "bytes",
-                            "null"])
+
+avro_primitive_types = set(
+    [
+        "string",
+        "text",
+        "varchar",
+        "float",
+        "double",
+        "int",
+        "integer",
+        "long",
+        "bigint",
+        "boolean",
+        "timestamp_ntz",
+        "timestamp",
+        "timestamp_tz",
+        "date",
+        "bytes",
+        "null",
+    ]
+)
+
 
 def _to_avro_primitive_logical_type(field_name: str, field: Field) -> AvroPrimitiveField:
     result = AvroPrimitiveField(field_name, field.required, field.description, AvroPrimitiveType.string)
@@ -114,10 +140,11 @@ def _to_avro_primitive_logical_type(field_name: str, field: Field) -> AvroPrimit
                 model=field,
                 reason="Unknown field type {field.type}",
                 result="failed",
-                message="Avro IDL type conversion failed."
+                message="Avro IDL type conversion failed.",
             )
     return result
 
+
 def _to_avro_idl_type(field_name: str, field: Field) -> AvroField:
     if field.type in avro_primitive_types:
         return _to_avro_primitive_logical_type(field_name, field)
@@ -125,17 +152,14 @@ def _to_avro_idl_type(field_name: str, field: Field) -> AvroField:
         match field.type:
             case "array":
                 return AvroArrayField(
-                    field_name,
-                    field.required,
-                    field.description,
-                    _to_avro_idl_type(field_name, field.items)
+                    field_name, field.required, field.description, _to_avro_idl_type(field_name, field.items)
                 )
             case "object" | "record" | "struct":
                 return AvroComplexField(
                     field_name,
                     field.required,
                     field.description,
-                    [_to_avro_idl_type(field_name, field) for (field_name, field) in field.fields.items()]
+                    [_to_avro_idl_type(field_name, field) for (field_name, field) in field.fields.items()],
                 )
             case _:
                 raise DataContractException(
@@ -144,56 +168,55 @@ def _to_avro_idl_type(field_name: str, field: Field) -> AvroField:
                     model=type,
                     reason="Unknown Data Contract field type",
                     result="failed",
-                    message="Avro IDL type conversion failed."
+                    message="Avro IDL type conversion failed.",
                 )
 
 
 def _generate_field_types(contract: DataContractSpecification) -> list[AvroField]:
     result = []
-    for (_, model) in contract.models.items():
-        for (field_name, field) in model.fields.items():
+    for _, model in contract.models.items():
+        for field_name, field in model.fields.items():
             result.append(_to_avro_idl_type(field_name, field))
     return result
 
+
 def generate_model_types(contract: DataContractSpecification) -> list[AvroModelType]:
     result = []
-    for (model_name, model) in contract.models.items():
-        result.append(AvroModelType(
-            name=model_name,
-            description=model.description,
-            fields=_generate_field_types(contract)
-        ))
+    for model_name, model in contract.models.items():
+        result.append(
+            AvroModelType(name=model_name, description=model.description, fields=_generate_field_types(contract))
+        )
     return result
 
+
 def _model_name_to_identifier(model_name: str):
-    return "".join([word.title() for word in  model_name.split()])
+    return "".join([word.title() for word in model_name.split()])
 
-def _contract_to_avro_idl_ir(contract: DataContractSpecification) -> AvroIDLProtocol:
 
+def _contract_to_avro_idl_ir(contract: DataContractSpecification) -> AvroIDLProtocol:
     """Convert models into an intermediate representation for later serialization into Avro IDL.
 
-      Each model is converted to a record containing a field for each model field.
-      """
+    Each model is converted to a record containing a field for each model field.
+    """
     inlined_contract = contract.model_copy()
     inline_definitions_into_data_contract(inlined_contract)
-    protocol_name = (_model_name_to_identifier(contract.info.title)
-                     if contract.info and contract.info.title
-                     else None)
-    description = (contract.info.description if
-                   contract.info and contract.info.description
-                   else None)
-    return AvroIDLProtocol(name=protocol_name,
-                           description=description,
-                           model_types=generate_model_types(inlined_contract))
+    protocol_name = _model_name_to_identifier(contract.info.title) if contract.info and contract.info.title else None
+    description = contract.info.description if contract.info and contract.info.description else None
+    return AvroIDLProtocol(
+        name=protocol_name, description=description, model_types=generate_model_types(inlined_contract)
+    )
+
 
 def _write_indent(indent: int, stream: typing.TextIO):
     stream.write("    " * indent)
 
+
 def _write_field_description(field: AvroField, indent: int, stream: typing.TextIO):
     if field.description:
         _write_indent(indent, stream)
         stream.write(f"/** {field.description} */\n")
 
+
 def _write_field_type_definition(field: AvroField, indent: int, stream: typing.TextIO) -> str:
     # Write any extra information (such as record type definition) and return
     # the name of the generated type. Writes descriptions only for record
@@ -215,7 +238,7 @@ def _write_field_type_definition(field: AvroField, indent: int, stream: typing.T
             for subfield in subfields:
                 subfield_types.append(_write_field_type_definition(subfield, indent + 1, stream))
             # Reference all defined record types.
-            for (field, subfield_type) in zip(field.subfields, subfield_types):
+            for field, subfield_type in zip(field.subfields, subfield_types):
                 _write_field_description(field, indent + 1, stream)
                 _write_indent(indent + 1, stream)
                 stream.write(f"{subfield_type} {field.name};\n")
@@ -234,15 +257,15 @@ def _write_field_type_definition(field: AvroField, indent: int, stream: typing.T
         case _:
             raise RuntimeError("Unknown Avro field type {field}")
 
-def _write_field(field: AvroField,
-                indent,
-                stream: typing.TextIO):
+
+def _write_field(field: AvroField, indent, stream: typing.TextIO):
     # Start of recursion.
     typename = _write_field_type_definition(field, indent, stream)
     _write_field_description(field, indent, stream)
     _write_indent(indent, stream)
     stream.write(f"{typename} {field.name};\n")
 
+
 def _write_model_type(model: AvroModelType, stream: typing.TextIO):
     # Called once for each model
     if model.description: