datacontract-cli 0.10.1__py3-none-any.whl → 0.10.2__py3-none-any.whl

datacontract/engines/soda/connections/kafka.py

@@ -1,40 +1,36 @@
 import os
-
-import pyspark.sql.functions as fn
 from pyspark.sql import SparkSession
+from pyspark.sql.functions import col, expr, from_json
 from pyspark.sql.avro.functions import from_avro
-from pyspark.sql.functions import from_json, col
 from pyspark.sql.types import (
     StructType,
-    DataType,
-    NullType,
-    ArrayType,
-    BinaryType,
-    DateType,
-    TimestampNTZType,
-    TimestampType,
-    BooleanType,
-    LongType,
-    IntegerType,
-    DoubleType,
-    DecimalType,
-    StringType,
     StructField,
+    StringType,
+    DecimalType,
+    DoubleType,
+    IntegerType,
+    LongType,
+    BooleanType,
+    TimestampType,
+    TimestampNTZType,
+    DateType,
+    BinaryType,
+    ArrayType,
+    NullType,
+    DataType,
 )
 
 from datacontract.export.avro_converter import to_avro_schema_json
-from datacontract.model.data_contract_specification import \
-    DataContractSpecification, Server, Field
+from datacontract.model.data_contract_specification import DataContractSpecification, Server, Field
 from datacontract.model.exceptions import DataContractException
 
 
-def create_spark_session(tmp_dir) -> SparkSession:
-    # TODO: Update dependency versions when updating pyspark
-    # TODO: add protobuf library
+def create_spark_session(tmp_dir: str) -> SparkSession:
+    """Create and configure a Spark session."""
     spark = (
         SparkSession.builder.appName("datacontract")
-        .config("spark.sql.warehouse.dir", tmp_dir + "/spark-warehouse")
-        .config("spark.streaming.stopGracefullyOnShutdown", True)
+        .config("spark.sql.warehouse.dir", f"{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",
@@ -47,106 +43,113 @@ def create_spark_session(tmp_dir) -> SparkSession:
 
 
 def read_kafka_topic(spark: SparkSession, data_contract: DataContractSpecification, server: Server, tmp_dir):
-    host = server.host
-    topic = server.topic
-    auth_options = get_auth_options()
-
-    # read full kafka topic
+    """Read and process data from a Kafka topic based on the server configuration."""
     df = (
         spark.read.format("kafka")
-        .options(**auth_options)
-        .option("kafka.bootstrap.servers", host)
-        .option("subscribe", topic)
+        .options(**get_auth_options())
+        .option("kafka.bootstrap.servers", server.host)
+        .option("subscribe", server.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":
-        avro_schema = to_avro_schema_json(model_name, model)
-
-        # Parse out the extra bytes from the Avro data
-        # A Kafka message contains a key and a value. Data going through a Kafka topic in Confluent Cloud has five bytes added to the beginning of every Avro value. If you are using Avro format keys, then five bytes will be added to the beginning of those as well. For this example, we’re assuming string keys. These bytes consist of one magic byte and four bytes representing the schema ID of the schema in the registry that is needed to decode that data. The bytes need to be removed so that the schema ID can be determined and the Avro data can be parsed. To manipulate the data, we need a couple of imports:
-        df2 = df.withColumn("fixedValue", fn.expr("substring(value, 6, length(value)-5)"))
-
-        options = {"mode": "PERMISSIVE"}
-        df3 = df2.select(from_avro(col("fixedValue"), avro_schema, options).alias("avro")).select(col("avro.*"))
-    elif server.format == "json":
-        # TODO A good warning when the conversion to json fails
-        struct_type = to_struct_type(model.fields)
-        df2 = df.selectExpr("CAST(key AS STRING)", "CAST(value AS STRING)")
-
-        options = {"mode": "PERMISSIVE"}
-        df3 = df2.select(from_json(df2.value, struct_type, options).alias("json")).select(col("json.*"))
-    else:
-        raise DataContractException(
-            type="test",
-            name="Configuring Kafka checks",
-            result="warning",
-            reason=f"Kafka format '{server.format}' is not supported. Skip executing tests.",
-            engine="datacontract",
-        )
 
-    # df3.writeStream.toTable(model_name, checkpointLocation=tmp_dir + "/checkpoint")
-    df3.createOrReplaceTempView(model_name)
-    # print(spark.sql(f"select * from {model_name}").show())
+    match server.format:
+        case "avro":
+            process_avro_format(df, model_name, model)
+        case "json":
+            process_json_format(df, model_name, model)
+        case _:
+            raise DataContractException(
+                type="test",
+                name="Configuring Kafka checks",
+                result="warning",
+                reason=f"Kafka format '{server.format}' is not supported. " f"Skip executing tests.",
+                engine="datacontract",
+            )
+
+
+def process_avro_format(df, model_name, model):
+    avro_schema = to_avro_schema_json(model_name, model)
+    df2 = df.withColumn("fixedValue", expr("substring(value, 6, length(value)-5)"))
+    options = {"mode": "PERMISSIVE"}
+    df2.select(from_avro(col("fixedValue"), avro_schema, options).alias("avro")).select(
+        col("avro.*")
+    ).createOrReplaceTempView(model_name)
+
+
+def process_json_format(df, model_name, model):
+    struct_type = to_struct_type(model.fields)
+    df.selectExpr("CAST(key AS STRING)", "CAST(value AS STRING)").select(
+        from_json(col("value"), struct_type, {"mode": "PERMISSIVE"}).alias("json")
+    ).select(col("json.*")).createOrReplaceTempView(model_name)
 
 
 def get_auth_options():
+    """Retrieve Kafka authentication options from environment variables."""
     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:
-        kafka_sasl_jaas_config = f'org.apache.kafka.common.security.plain.PlainLoginModule required username="{kafka_sasl_username}" password="{kafka_sasl_password}";'
-        auth_options = {
-            "kafka.sasl.mechanism": "PLAIN",
-            "kafka.security.protocol": "SASL_SSL",
-            "kafka.sasl.jaas.config": kafka_sasl_jaas_config,
-        }
-    return auth_options
+        return {}
+
+    return {
+        "kafka.sasl.mechanism": "PLAIN",
+        "kafka.security.protocol": "SASL_SSL",
+        "kafka.sasl.jaas.config": (
+            f"org.apache.kafka.common.security.plain.PlainLoginModule required "
+            f'username="{kafka_sasl_username}" password="{kafka_sasl_password}";'
+        ),
+    }
 
 
 def to_struct_type(fields):
-    struct_fields = []
-    for field_name, field in fields.items():
-        struct_fields.append(to_struct_field(field_name, field))
-    return StructType(struct_fields)
+    """Convert field definitions to Spark StructType."""
+    return StructType([to_struct_field(field_name, field) for field_name, field in fields.items()])
 
 
 def to_struct_field(field_name: str, field: Field) -> StructField:
-    if field.type is None:
-        data_type = DataType()
-    if field.type in ["string", "varchar", "text"]:
-        data_type = StringType()
-    elif field.type in ["number", "decimal", "numeric"]:
-        data_type = DecimalType()
-    elif field.type in ["float", "double"]:
-        data_type = DoubleType()
-    elif field.type in ["integer", "int"]:
-        data_type = IntegerType()
-    elif field.type in ["long", "bigint"]:
-        data_type = LongType()
-    elif field.type in ["boolean"]:
-        data_type = BooleanType()
-    elif field.type in ["timestamp", "timestamp_tz"]:
-        data_type = TimestampType()
-    elif field.type in ["timestamp_ntz"]:
-        data_type = TimestampNTZType()
-    elif field.type in ["date"]:
-        data_type = DateType()
-    elif field.type in ["time"]:
-        data_type = DataType()
-    elif field.type in ["object", "record", "struct"]:
-        data_type = to_struct_type(field.fields)
-    elif field.type in ["binary"]:
-        data_type = BinaryType()
-    elif field.type in ["array"]:
-        # TODO support array structs
-        data_type = ArrayType()
-    elif field.type in ["null"]:
-        data_type = NullType()
-    else:
-        data_type = DataType()
+    """Map field definitions to Spark StructField using match-case."""
+    match field.type:
+        case "string" | "varchar" | "text":
+            data_type = StringType()
+        case "number" | "decimal" | "numeric":
+            data_type = DecimalType()
+        case "float" | "double":
+            data_type = DoubleType()
+        case "integer" | "int":
+            data_type = IntegerType()
+        case "long" | "bigint":
+            data_type = LongType()
+        case "boolean":
+            data_type = BooleanType()
+        case "timestamp" | "timestamp_tz":
+            data_type = TimestampType()
+        case "timestamp_ntz":
+            data_type = TimestampNTZType()
+        case "date":
+            data_type = DateType()
+        case "time":
+            data_type = DataType()  # Specific handling for time type
+        case "object" | "record" | "struct":
+            data_type = StructType(
+                [to_struct_field(sub_field_name, sub_field) for sub_field_name, sub_field in field.fields.items()]
+            )
+        case "binary":
+            data_type = BinaryType()
+        case "array":
+            element_type = (
+                StructType(
+                    [to_struct_field(sub_field_name, sub_field) for sub_field_name, sub_field in field.fields.items()]
+                )
+                if field.fields
+                else DataType()
+            )
+            data_type = ArrayType(element_type)
+        case "null":
+            data_type = NullType()
+        case _:
+            data_type = DataType()  # Fallback generic DataType
 
     return StructField(field_name, data_type, nullable=not field.required)

datacontract/export/avro_idl_converter.py

@@ -4,8 +4,7 @@ from enum import Enum
 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.data_contract_specification import DataContractSpecification, Field
 from datacontract.model.exceptions import DataContractException
 
 

datacontract/export/dbt_converter.py

@@ -3,8 +3,7 @@ from typing import Dict
 import yaml
 
 from datacontract.export.sql_type_converter import convert_to_sql_type
-from datacontract.model.data_contract_specification import \
-    DataContractSpecification, Model, Field
+from datacontract.model.data_contract_specification import DataContractSpecification, Model, Field
 
 
 def to_dbt_models_yaml(data_contract_spec: DataContractSpecification):

datacontract/export/great_expectations_converter.py

@@ -3,8 +3,7 @@ from typing import Dict, List, Any
 
 import yaml
 
-from datacontract.model.data_contract_specification import \
-    DataContractSpecification, Field, Quality
+from datacontract.model.data_contract_specification import DataContractSpecification, Field, Quality
 
 
 def to_great_expectations(data_contract_spec: DataContractSpecification, model_key: str) -> str:

datacontract/export/html_export.py

@@ -6,8 +6,7 @@ import pytz
 import yaml
 from jinja2 import Environment, PackageLoader, select_autoescape
 
-from datacontract.model.data_contract_specification import \
-    DataContractSpecification
+from datacontract.model.data_contract_specification import DataContractSpecification
 
 
 def to_html(data_contract_spec: DataContractSpecification) -> str:
@@ -40,9 +39,9 @@ def to_html(data_contract_spec: DataContractSpecification) -> str:
 
     datacontract_yaml = data_contract_spec.to_yaml()
 
-    tz = pytz.timezone('UTC')
+    tz = pytz.timezone("UTC")
     now = datetime.datetime.now(tz)
-    formatted_date = now.strftime('%d %b %Y %H:%M:%S UTC')
+    formatted_date = now.strftime("%d %b %Y %H:%M:%S UTC")
     datacontract_cli_version = get_version()
 
     # Render the template with necessary data

datacontract/export/jsonschema_converter.py

@@ -1,8 +1,7 @@
 import json
 from typing import Dict
 
-from datacontract.model.data_contract_specification import \
-    DataContractSpecification, Model, Field
+from datacontract.model.data_contract_specification import DataContractSpecification, Model, Field
 
 
 def to_jsonschemas(data_contract_spec: DataContractSpecification):

datacontract/export/odcs_converter.py

@@ -2,8 +2,7 @@ from typing import Dict
 
 import yaml
 
-from datacontract.model.data_contract_specification import \
-    DataContractSpecification, Model, Field
+from datacontract.model.data_contract_specification import DataContractSpecification, Model, Field
 
 
 def to_odcs_yaml(data_contract_spec: DataContractSpecification):

datacontract/export/rdf_converter.py

@@ -141,7 +141,7 @@ def add_info(contract, info, graph, dc, dcx):
     graph.add((bnode_info, dc.version, Literal(info.version)))
 
     # add owner
-    owner = URIRef(info.owner)
+    owner = Literal(info.owner)
     graph.add((bnode_info, dc.owner, owner))
 
     # add contact

datacontract/export/sodacl_converter.py

@@ -1,8 +1,7 @@
 import yaml
 
 from datacontract.export.sql_type_converter import convert_to_sql_type
-from datacontract.model.data_contract_specification import \
-    DataContractSpecification
+from datacontract.model.data_contract_specification import DataContractSpecification
 
 
 def to_sodacl_yaml(

datacontract/export/terraform_converter.py

@@ -1,7 +1,6 @@
 import re
 
-from datacontract.model.data_contract_specification import \
-    DataContractSpecification, Server
+from datacontract.model.data_contract_specification import DataContractSpecification, Server
 
 
 def to_terraform(data_contract_spec: DataContractSpecification, server_id: str = None) -> str:

datacontract/imports/avro_importer.py

@@ -1,7 +1,6 @@
 import avro.schema
 
-from datacontract.model.data_contract_specification import \
-    DataContractSpecification, Model, Field
+from datacontract.model.data_contract_specification import DataContractSpecification, Model, Field
 from datacontract.model.exceptions import DataContractException
 
 

datacontract/imports/glue_importer.py

@@ -0,0 +1,183 @@
+import boto3
+from typing import List
+
+from datacontract.model.data_contract_specification import (
+    DataContractSpecification,
+    Model,
+    Field,
+    Server,
+)
+
+
+def get_glue_database(datebase_name: str):
+    """Get the details Glue database.
+
+    Args:
+        database_name (str): glue database to request.
+
+    Returns:
+        set: catalogid and locationUri
+    """
+
+    glue = boto3.client("glue")
+    try:
+        response = glue.get_database(Name=datebase_name)
+    except glue.exceptions.EntityNotFoundException:
+        print(f"Database not found {datebase_name}.")
+        return (None, None)
+    except Exception as e:
+        # todo catch all
+        print(f"Error: {e}")
+        return (None, None)
+
+    return (response["Database"]["CatalogId"], response["Database"].get("LocationUri", "None"))
+
+
+def get_glue_tables(database_name: str) -> List[str]:
+    """Get the list of tables in a Glue database.
+
+    Args:
+        database_name (str): glue database to request.
+
+    Returns:
+        List[string]: List of table names
+    """
+
+    glue = boto3.client("glue")
+
+    # Set the paginator
+    paginator = glue.get_paginator("get_tables")
+
+    # Initialize an empty list to store the table names
+    table_names = []
+    try:
+        # Paginate through the tables
+        for page in paginator.paginate(DatabaseName=database_name, PaginationConfig={"PageSize": 100}):
+            # Add the tables from the current page to the list
+            table_names.extend([table["Name"] for table in page["TableList"] if "Name" in table])
+    except glue.exceptions.EntityNotFoundException:
+        print(f"Database {database_name} not found.")
+        return []
+    except Exception as e:
+        # todo catch all
+        print(f"Error: {e}")
+        return []
+
+    return table_names
+
+
+def get_glue_table_schema(database_name: str, table_name: str):
+    """Get the schema of a Glue table.
+
+    Args:
+        database_name (str): Glue database name.
+        table_name (str): Glue table name.
+
+    Returns:
+        dict: Table schema
+    """
+
+    glue = boto3.client("glue")
+
+    # Get the table schema
+    try:
+        response = glue.get_table(DatabaseName=database_name, Name=table_name)
+    except glue.exceptions.EntityNotFoundException:
+        print(f"Table {table_name} not found in database {database_name}.")
+        return {}
+    except Exception as e:
+        # todo catch all
+        print(f"Error: {e}")
+        return {}
+
+    table_schema = response["Table"]["StorageDescriptor"]["Columns"]
+
+    # when using hive partition keys, the schema is stored in the PartitionKeys field
+    if response["Table"].get("PartitionKeys") is not None:
+        for pk in response["Table"]["PartitionKeys"]:
+            table_schema.append(
+                {
+                    "Name": pk["Name"],
+                    "Type": pk["Type"],
+                    "Hive": True,
+                    "Comment": "Partition Key",
+                }
+            )
+
+    return table_schema
+
+
+def import_glue(data_contract_specification: DataContractSpecification, source: str):
+    """Import the schema of a Glue database."""
+
+    catalogid, location_uri = get_glue_database(source)
+
+    # something went wrong
+    if catalogid is None:
+        return data_contract_specification
+
+    tables = get_glue_tables(source)
+
+    data_contract_specification.servers = {
+        "production": Server(type="glue", account=catalogid, database=source, location=location_uri),
+    }
+
+    for table_name in tables:
+        if data_contract_specification.models is None:
+            data_contract_specification.models = {}
+
+        table_schema = get_glue_table_schema(source, table_name)
+
+        fields = {}
+        for column in table_schema:
+            field = Field()
+            field.type = map_type_from_sql(column["Type"])
+
+            # hive partitons are required, but are not primary keys
+            if column.get("Hive"):
+                field.required = True
+
+            field.description = column.get("Comment")
+
+            fields[column["Name"]] = field
+
+        data_contract_specification.models[table_name] = Model(
+            type="table",
+            fields=fields,
+        )
+
+    return data_contract_specification
+
+
+def map_type_from_sql(sql_type: str):
+    if sql_type is None:
+        return None
+
+    if sql_type.lower().startswith("varchar"):
+        return "varchar"
+    if sql_type.lower().startswith("string"):
+        return "string"
+    if sql_type.lower().startswith("text"):
+        return "text"
+    elif sql_type.lower().startswith("byte"):
+        return "byte"
+    elif sql_type.lower().startswith("short"):
+        return "short"
+    elif sql_type.lower().startswith("integer"):
+        return "integer"
+    elif sql_type.lower().startswith("long"):
+        return "long"
+    elif sql_type.lower().startswith("bigint"):
+        return "long"
+    elif sql_type.lower().startswith("float"):
+        return "float"
+    elif sql_type.lower().startswith("double"):
+        return "double"
+    elif sql_type.lower().startswith("boolean"):
+        return "boolean"
+    elif sql_type.lower().startswith("timestamp"):
+        return "timestamp"
+    elif sql_type.lower().startswith("date"):
+        return "date"
+    else:
+        return "variant"

datacontract/imports/sql_importer.py

@@ -1,7 +1,6 @@
 from simple_ddl_parser import parse_from_file
 
-from datacontract.model.data_contract_specification import \
-    DataContractSpecification, Model, Field
+from datacontract.model.data_contract_specification import DataContractSpecification, Model, Field
 
 
 def import_sql(data_contract_specification: DataContractSpecification, format: str, source: str):
@@ -45,19 +44,31 @@ def map_type_from_sql(sql_type: str):
     if sql_type is None:
         return None
 
-    if sql_type.lower().startswith("varchar"):
+    sql_type_normed = sql_type.lower().strip()
+    
+    if sql_type_normed.startswith("varchar"):
         return "varchar"
-    if sql_type.lower().startswith("string"):
+    elif sql_type_normed.startswith("string"):
         return "string"
-    if sql_type.lower().startswith("text"):
+    elif sql_type_normed.startswith("text"):
         return "text"
-    elif sql_type.lower().startswith("int"):
+    elif sql_type_normed.startswith("int"):
         return "integer"
-    elif sql_type.lower().startswith("float"):
+    elif sql_type_normed.startswith("float"):
         return "float"
-    elif sql_type.lower().startswith("bool"):
+    elif sql_type_normed.startswith("bool"):
         return "boolean"
-    elif sql_type.lower().startswith("timestamp"):
+    elif sql_type_normed.startswith("timestamp"):
         return "timestamp"
+    elif sql_type_normed == "date":
+        return "date"
+    elif sql_type_normed == "smalldatetime":
+        return "timestamp_ntz"
+    elif sql_type_normed == "datetime":
+        return "timestamp_ntz"
+    elif sql_type_normed == "datetime2":
+        return "timestamp_ntz"
+    elif sql_type_normed == "datetimeoffset":
+        return "timestamp_tz"        
     else:
         return "variant"

datacontract/integration/publish_opentelemetry.py

@@ -4,14 +4,11 @@ import os
 from importlib import metadata
 
 from opentelemetry import metrics
-from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import \
-    OTLPMetricExporter as OTLPgRPCMetricExporter
-from opentelemetry.exporter.otlp.proto.http.metric_exporter import \
-    OTLPMetricExporter
+from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import OTLPMetricExporter as OTLPgRPCMetricExporter
+from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
 from opentelemetry.metrics import Observation
 from opentelemetry.sdk.metrics import MeterProvider
-from opentelemetry.sdk.metrics.export import ConsoleMetricExporter, \
-    PeriodicExportingMetricReader
+from opentelemetry.sdk.metrics.export import ConsoleMetricExporter, PeriodicExportingMetricReader
 
 from datacontract.model.run import Run
 

datacontract/lint/linters/example_model_linter.py

@@ -4,8 +4,7 @@ import json
 
 import yaml
 
-from datacontract.model.data_contract_specification import \
-    DataContractSpecification, Example
+from datacontract.model.data_contract_specification import DataContractSpecification, Example
 from ..lint import Linter, LinterResult
 
 

datacontract/lint/linters/field_pattern_linter.py

@@ -1,7 +1,6 @@
 import re
 
-from datacontract.model.data_contract_specification import \
-    DataContractSpecification
+from datacontract.model.data_contract_specification import DataContractSpecification
 from ..lint import Linter, LinterResult
 
 

datacontract/lint/linters/notice_period_linter.py

@@ -1,7 +1,6 @@
 import re
 
-from datacontract.model.data_contract_specification import \
-    DataContractSpecification
+from datacontract.model.data_contract_specification import DataContractSpecification
 from ..lint import Linter, LinterResult
 
 

datacontract/lint/linters/quality_schema_linter.py

@@ -1,7 +1,6 @@
 import yaml
 
-from datacontract.model.data_contract_specification import \
-    DataContractSpecification, Model
+from datacontract.model.data_contract_specification import DataContractSpecification, Model
 from ..lint import Linter, LinterResult