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

datacontract/engines/soda/check_soda_execute.py

@@ -1,28 +1,39 @@
 import logging
+import typing
+import uuid
 
-from pyspark.sql import SparkSession
-from soda.scan import Scan
-
-from datacontract.engines.soda.connections.bigquery import \
-    to_bigquery_soda_configuration
-from datacontract.engines.soda.connections.databricks import \
-    to_databricks_soda_configuration
-from datacontract.engines.soda.connections.duckdb import get_duckdb_connection
-from datacontract.engines.soda.connections.kafka import create_spark_session, \
-    read_kafka_topic
-from datacontract.engines.soda.connections.postgres import \
-    to_postgres_soda_configuration
-from datacontract.engines.soda.connections.snowflake import \
-    to_snowflake_soda_configuration
+from datacontract.engines.soda.connections.athena import to_athena_soda_configuration
+
+if typing.TYPE_CHECKING:
+    from pyspark.sql import SparkSession
+
+from duckdb.duckdb import DuckDBPyConnection
+
+from datacontract.engines.soda.connections.bigquery import to_bigquery_soda_configuration
+from datacontract.engines.soda.connections.databricks import to_databricks_soda_configuration
+from datacontract.engines.soda.connections.duckdb_connection import get_duckdb_connection
+from datacontract.engines.soda.connections.kafka import create_spark_session, read_kafka_topic
+from datacontract.engines.soda.connections.postgres import to_postgres_soda_configuration
+from datacontract.engines.soda.connections.snowflake import to_snowflake_soda_configuration
+from datacontract.engines.soda.connections.sqlserver import to_sqlserver_soda_configuration
+from datacontract.engines.soda.connections.trino import to_trino_soda_configuration
 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.data_contract_specification import DataContractSpecification, Server
+from datacontract.model.run import Check, Log, ResultEnum, Run
 
 
 def check_soda_execute(
-    run: Run, data_contract: DataContractSpecification, server: Server, spark: SparkSession, tmp_dir
+    run: Run,
+    data_contract: DataContractSpecification,
+    server: Server,
+    spark: "SparkSession" = None,
+    duckdb_connection: DuckDBPyConnection = None,
 ):
+    from soda.common.config_helper import ConfigHelper
+
+    ConfigHelper.get_instance().upsert_value("send_anonymous_usage_stats", False)
+    from soda.scan import Scan
+
     if data_contract is None:
         run.log_warn("Cannot run engine soda-core, as data contract is invalid")
         return
@@ -30,9 +41,10 @@ def check_soda_execute(
     run.log_info("Running engine soda-core")
     scan = Scan()
 
-    if server.type == "s3" or server.type == "local":
-        if server.format in ["json", "parquet", "csv"]:
-            con = get_duckdb_connection(data_contract, server)
+    if server.type in ["s3", "gcs", "azure", "local"]:
+        if server.format in ["json", "parquet", "csv", "delta"]:
+            run.log_info(f"Configuring engine soda-core to connect to {server.type} {server.format} with duckdb")
+            con = get_duckdb_connection(data_contract, server, run, duckdb_connection)
             scan.add_duckdb_connection(duckdb_connection=con, data_source_name=server.type)
             scan.set_data_source_name(server.type)
         else:
@@ -40,7 +52,7 @@ def check_soda_execute(
                 Check(
                     type="general",
                     name="Check that format is supported",
-                    result="warning",
+                    result=ResultEnum.warning,
                     reason=f"Format {server.format} not yet supported by datacontract CLI",
                     engine="datacontract",
                 )
@@ -61,27 +73,52 @@ def check_soda_execute(
         scan.set_data_source_name(server.type)
     elif server.type == "databricks":
         if spark is not None:
-            logging.info("Use Spark to connect to data source")
+            run.log_info("Connecting to databricks via spark")
             scan.add_spark_session(spark, data_source_name=server.type)
             scan.set_data_source_name(server.type)
-            spark.sql(f"USE {server.catalog}.{server.schema_}")
+            database_name = ".".join(filter(None, [server.catalog, server.schema_]))
+            spark.sql(f"USE {database_name}")
         else:
+            run.log_info("Connecting to databricks directly")
             soda_configuration_str = to_databricks_soda_configuration(server)
             scan.add_configuration_yaml_str(soda_configuration_str)
             scan.set_data_source_name(server.type)
+    elif server.type == "dataframe":
+        if spark is None:
+            run.log_warn(
+                "Server type dataframe only works with the Python library and requires a Spark session, "
+                "please provide one with the DataContract class"
+            )
+            return
+        else:
+            logging.info("Use Spark to connect to data source")
+            scan.add_spark_session(spark, data_source_name="datacontract-cli")
+            scan.set_data_source_name("datacontract-cli")
     elif server.type == "kafka":
         if spark is None:
-            spark = create_spark_session(tmp_dir)
-        read_kafka_topic(spark, data_contract, server, tmp_dir)
+            spark = create_spark_session()
+        read_kafka_topic(spark, data_contract, server)
         scan.add_spark_session(spark, data_source_name=server.type)
         scan.set_data_source_name(server.type)
+    elif server.type == "sqlserver":
+        soda_configuration_str = to_sqlserver_soda_configuration(server)
+        scan.add_configuration_yaml_str(soda_configuration_str)
+        scan.set_data_source_name(server.type)
+    elif server.type == "trino":
+        soda_configuration_str = to_trino_soda_configuration(server)
+        scan.add_configuration_yaml_str(soda_configuration_str)
+        scan.set_data_source_name(server.type)
+    elif server.type == "athena":
+        soda_configuration_str = to_athena_soda_configuration(server)
+        scan.add_configuration_yaml_str(soda_configuration_str)
+        scan.set_data_source_name(server.type)
 
     else:
         run.checks.append(
             Check(
                 type="general",
                 name="Check that server type is supported",
-                result="warning",
+                result=ResultEnum.warning,
                 reason=f"Server type {server.type} not yet supported by datacontract CLI",
                 engine="datacontract-cli",
             )
@@ -89,37 +126,34 @@ def check_soda_execute(
         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, server_type, check_types)
+    sodacl_yaml_str = to_sodacl_yaml(run)
     # print("sodacl_yaml_str:\n" + sodacl_yaml_str)
     scan.add_sodacl_yaml_str(sodacl_yaml_str)
 
     # Execute the scan
-    logging.info("Starting soda scan")
+    logging.info("Starting soda scan with checks:\n" + sodacl_yaml_str)
     scan.execute()
     logging.info("Finished soda scan")
 
     # pprint.PrettyPrinter(indent=2).pprint(scan.build_scan_results())
 
     scan_results = scan.get_scan_results()
-    for c in scan_results.get("checks"):
-        check = Check(
-            type="schema",
-            result=to_result(c),
-            reason=", ".join(c.get("outcomeReasons")),
-            name=c.get("name"),
-            model=c.get("table"),
-            field=c.get("column"),
-            engine="soda-core",
-            diagnostics=c.get("diagnostics"),
-        )
-        update_reason(check, c)
-        run.checks.append(check)
+    for scan_result in scan_results.get("checks"):
+        name = scan_result.get("name")
+        check = get_check(run, scan_result)
+        if check is None:
+            check = Check(
+                id=str(uuid.uuid4()),
+                category="custom",
+                type="custom",
+                name=name,
+                engine="soda-core",
+            )
+            run.checks.append(check)
+        check.result = to_result(scan_result)
+        check.reason = ", ".join(scan_result.get("outcomeReasons"))
+        check.diagnostics = scan_result.get("diagnostics")
+        update_reason(check, scan_result)
 
     for log in scan_results.get("logs"):
         run.logs.append(
@@ -135,8 +169,8 @@ def check_soda_execute(
         run.checks.append(
             Check(
                 type="general",
-                name="Execute quality checks",
-                result="warning",
+                name="Data Contract Tests",
+                result=ResultEnum.warning,
                 reason="Engine soda-core has errors. See the logs for details.",
                 engine="soda-core",
             )
@@ -144,14 +178,22 @@ def check_soda_execute(
         return
 
 
-def to_result(c) -> str:
+def get_check(run, scan_result) -> Check | None:
+    check_by_name = next((c for c in run.checks if c.key == scan_result.get("name")), None)
+    if check_by_name is not None:
+        return check_by_name
+
+    return None
+
+
+def to_result(c) -> ResultEnum:
     soda_outcome = c.get("outcome")
     if soda_outcome == "pass":
-        return "passed"
+        return ResultEnum.passed
     elif soda_outcome == "fail":
-        return "failed"
+        return ResultEnum.failed
     else:
-        return soda_outcome
+        return ResultEnum.unknown
 
 
 def update_reason(check, c):
@@ -164,9 +206,11 @@ def update_reason(check, c):
         if block["title"] == "Diagnostics":
             # Extract and print the 'text' value
             diagnostics_text = block["text"]
-            print(diagnostics_text)
+            # print(diagnostics_text)
             diagnostics_text_split = diagnostics_text.split(":icon-fail: ")
             if len(diagnostics_text_split) > 1:
                 check.reason = diagnostics_text_split[1].strip()
-                print(check.reason)
+                # print(check.reason)
             break  # Exit the loop once the desired block is found
+    if "fail" in c["diagnostics"]:
+        check.reason = f"Value: {c['diagnostics']['value']} Fail: {c['diagnostics']['fail']}"

datacontract/engines/soda/connections/athena.py

@@ -0,0 +1,79 @@
+import os
+
+import yaml
+
+from datacontract.model.exceptions import DataContractException
+
+
+def to_athena_soda_configuration(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_session_token = os.getenv("DATACONTRACT_S3_SESSION_TOKEN")
+
+    # Validate required parameters
+    if not s3_access_key_id:
+        raise DataContractException(
+            type="athena-connection",
+            name="missing_access_key_id",
+            reason="AWS access key ID is required. Set the DATACONTRACT_S3_ACCESS_KEY_ID environment variable.",
+            engine="datacontract",
+        )
+
+    if not s3_secret_access_key:
+        raise DataContractException(
+            type="athena-connection",
+            name="missing_secret_access_key",
+            reason="AWS secret access key is required. Set the DATACONTRACT_S3_SECRET_ACCESS_KEY environment variable.",
+            engine="datacontract",
+        )
+
+    if not hasattr(server, "schema_") or not server.schema_:
+        raise DataContractException(
+            type="athena-connection",
+            name="missing_schema",
+            reason="Schema is required for Athena connection. Specify the schema where your tables exist in the server configuration.",
+            engine="datacontract",
+        )
+
+    if not hasattr(server, "stagingDir") or not server.stagingDir:
+        raise DataContractException(
+            type="athena-connection",
+            name="missing_s3_staging_dir",
+            reason="S3 staging directory is required for Athena connection. This should be the Amazon S3 Query Result Location (e.g., 's3://my-bucket/athena-results/').",
+            engine="datacontract",
+        )
+
+    # Validate S3 staging directory format
+    if not server.stagingDir.startswith("s3://"):
+        raise DataContractException(
+            type="athena-connection",
+            name="invalid_s3_staging_dir",
+            reason=f"S3 staging directory must start with 's3://'. Got: {server.s3_staging_dir}. Example: 's3://my-bucket/athena-results/'",
+            engine="datacontract",
+        )
+
+    data_source = {
+        "type": "athena",
+        "access_key_id": s3_access_key_id,
+        "secret_access_key": s3_secret_access_key,
+        "schema": server.schema_,
+        "staging_dir": server.stagingDir,
+    }
+
+    if s3_region:
+        data_source["region_name"] = s3_region
+    elif server.region_name:
+        data_source["region_name"] = server.region_name
+
+    if server.catalog:
+        # Optional, Identify the name of the Data Source, also referred to as a Catalog. The default value is `awsdatacatalog`.
+        data_source["catalog"] = server.catalog
+
+    if s3_session_token:
+        data_source["aws_session_token"] = s3_session_token
+
+    soda_configuration = {f"data_source {server.type}": data_source}
+
+    soda_configuration_str = yaml.dump(soda_configuration)
+    return soda_configuration_str

datacontract/engines/soda/connections/bigquery.py

@@ -6,10 +6,17 @@ 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
+
+    # check for our own environment variable first
+    account_info = os.getenv("DATACONTRACT_BIGQUERY_ACCOUNT_INFO_JSON_PATH")
+    if account_info is None:
+        # but as a fallback look for the default google one
+        account_info = os.getenv("GOOGLE_APPLICATION_CREDENTIALS")
+
     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": account_info,
             "auth_scopes": ["https://www.googleapis.com/auth/bigquery"],
             "project_id": server.project,
             "dataset": server.dataset,

datacontract/engines/soda/connections/databricks.py

@@ -4,15 +4,24 @@ import yaml
 
 
 def to_databricks_soda_configuration(server):
+    token = os.getenv("DATACONTRACT_DATABRICKS_TOKEN")
+    if token is None:
+        raise ValueError("DATACONTRACT_DATABRICKS_TOKEN environment variable is not set")
+    http_path = os.getenv("DATACONTRACT_DATABRICKS_HTTP_PATH")
+    host = server.host
+    if host is None:
+        host = os.getenv("DATACONTRACT_DATABRICKS_SERVER_HOSTNAME")
+    if host is None:
+        raise ValueError("DATACONTRACT_DATABRICKS_SERVER_HOSTNAME environment variable is not set")
     soda_configuration = {
         f"data_source {server.type}": {
             "type": "spark",
             "method": "databricks",
-            "host": server.host,
+            "host": host,
             "catalog": server.catalog,
             "schema": server.schema_,
-            "http_path": os.getenv("DATACONTRACT_DATABRICKS_HTTP_PATH"),
-            "token": os.getenv("DATACONTRACT_DATABRICKS_TOKEN"),
+            "http_path": http_path,
+            "token": token,
         }
     }
 

datacontract/engines/soda/connections/duckdb_connection.py

@@ -0,0 +1,241 @@
+import os
+from typing import Any, Dict
+
+import duckdb
+
+from datacontract.export.duckdb_type_converter import convert_to_duckdb_csv_type, convert_to_duckdb_json_type
+from datacontract.model.data_contract_specification import DataContractSpecification, Field, Model, Server
+from datacontract.model.run import Run
+
+
+def get_duckdb_connection(
+    data_contract: DataContractSpecification,
+    server: Server,
+    run: Run,
+    duckdb_connection: duckdb.DuckDBPyConnection | None = None,
+) -> duckdb.DuckDBPyConnection:
+    if duckdb_connection is None:
+        con = duckdb.connect(database=":memory:")
+    else:
+        con = duckdb_connection
+
+    path: str = ""
+    if server.type == "local":
+        path = server.path
+    if server.type == "s3":
+        path = server.location
+        setup_s3_connection(con, server)
+    if server.type == "gcs":
+        path = server.location
+        setup_gcs_connection(con, server)
+    if server.type == "azure":
+        path = server.location
+        setup_azure_connection(con, server)
+    for model_name, model in data_contract.models.items():
+        model_path = path
+        if "{model}" in model_path:
+            model_path = model_path.format(model=model_name)
+        run.log_info(f"Creating table {model_name} for {model_path}")
+
+        if server.format == "json":
+            json_format = "auto"
+            if server.delimiter == "new_line":
+                json_format = "newline_delimited"
+            elif server.delimiter == "array":
+                json_format = "array"
+            columns = to_json_types(model)
+            if columns is None:
+                con.sql(f"""
+                        CREATE VIEW "{model_name}" AS SELECT * FROM read_json_auto('{model_path}', format='{json_format}', hive_partitioning=1);
+                        """)
+            else:
+                con.sql(
+                    f"""CREATE VIEW "{model_name}" AS SELECT * FROM read_json_auto('{model_path}', format='{json_format}', columns={columns}, hive_partitioning=1);"""
+                )
+                add_nested_views(con, model_name, model.fields)
+        elif server.format == "parquet":
+            con.sql(f"""
+                        CREATE VIEW "{model_name}" AS SELECT * FROM read_parquet('{model_path}', hive_partitioning=1);
+                        """)
+        elif server.format == "csv":
+            columns = to_csv_types(model)
+            run.log_info("Using columns: " + str(columns))
+            if columns is None:
+                con.sql(
+                    f"""CREATE VIEW "{model_name}" AS SELECT * FROM read_csv('{model_path}', hive_partitioning=1);"""
+                )
+            else:
+                con.sql(
+                    f"""CREATE VIEW "{model_name}" AS SELECT * FROM read_csv('{model_path}', hive_partitioning=1, columns={columns});"""
+                )
+        elif server.format == "delta":
+            con.sql("update extensions;")  # Make sure we have the latest delta extension
+            con.sql(f"""CREATE VIEW "{model_name}" AS SELECT * FROM delta_scan('{model_path}');""")
+        table_info = con.sql(f"PRAGMA table_info('{model_name}');").fetchdf()
+        if table_info is not None and not table_info.empty:
+            run.log_info(f"DuckDB Table Info: {table_info.to_string(index=False)}")
+    return con
+
+
+def to_csv_types(model) -> dict[Any, str | None] | None:
+    if model is None:
+        return None
+    columns = {}
+    # ['SQLNULL', 'BOOLEAN', 'BIGINT', 'DOUBLE', 'TIME', 'DATE', 'TIMESTAMP', 'VARCHAR']
+    for field_name, field in model.fields.items():
+        columns[field_name] = convert_to_duckdb_csv_type(field)
+    return columns
+
+
+def to_json_types(model: Model) -> dict[Any, str | None] | None:
+    if model is None:
+        return None
+    columns = {}
+    for field_name, field in model.fields.items():
+        columns[field_name] = convert_to_duckdb_json_type(field)
+    return columns
+
+
+def add_nested_views(con: duckdb.DuckDBPyConnection, model_name: str, fields: Dict[str, Field] | None):
+    model_name = model_name.strip('"')
+    if fields is None:
+        return
+    for field_name, field in fields.items():
+        if field.type is None or field.type.lower() not in ["array", "object"]:
+            continue
+        field_type = field.type.lower()
+        if field_type == "array" and field.items is None:
+            continue
+        elif field_type == "object" and field.fields is None:
+            continue
+
+        nested_model_name = f"{model_name}__{field_name}"
+        max_depth = 2 if field_type == "array" else 1
+
+        ## if parent field is not required, the nested objects may respolve
+        ## to a row of NULLs -- but if the objects themselves have required
+        ## fields, this will fail the check.
+        where = "" if field.required else f" WHERE {field_name} IS NOT NULL"
+        con.sql(f"""
+            CREATE VIEW IF NOT EXISTS "{nested_model_name}" AS
+            SELECT unnest({field_name}, max_depth := {max_depth}) as {field_name} FROM "{model_name}" {where}
+            """)
+        if field_type == "array":
+            add_nested_views(con, nested_model_name, field.items.fields)
+        elif field_type == "object":
+            add_nested_views(con, nested_model_name, field.fields)
+
+
+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_session_token = os.getenv("DATACONTRACT_S3_SESSION_TOKEN")
+    s3_endpoint = "s3.amazonaws.com"
+    use_ssl = "true"
+    url_style = "vhost"
+    if server.endpointUrl is not None:
+        url_style = "path"
+        s3_endpoint = server.endpointUrl.removeprefix("http://").removeprefix("https://")
+        if server.endpointUrl.startswith("http://"):
+            use_ssl = "false"
+
+    if s3_access_key_id is not None:
+        if s3_session_token is not None:
+            con.sql(f"""
+                CREATE OR REPLACE SECRET s3_secret (
+                    TYPE S3,
+                    PROVIDER CREDENTIAL_CHAIN,
+                    REGION '{s3_region}',
+                    KEY_ID '{s3_access_key_id}',
+                    SECRET '{s3_secret_access_key}',
+                    SESSION_TOKEN '{s3_session_token}',
+                    ENDPOINT '{s3_endpoint}',
+                    USE_SSL '{use_ssl}',
+                    URL_STYLE '{url_style}'
+                );
+            """)
+        else:
+            con.sql(f"""
+                CREATE OR REPLACE SECRET s3_secret (
+                    TYPE S3,
+                    PROVIDER CREDENTIAL_CHAIN,
+                    REGION '{s3_region}',
+                    KEY_ID '{s3_access_key_id}',
+                    SECRET '{s3_secret_access_key}',
+                    ENDPOINT '{s3_endpoint}',
+                    USE_SSL '{use_ssl}',
+                    URL_STYLE '{url_style}'
+                );
+            """)
+
+    #     con.sql(f"""
+    #                 SET s3_region = '{s3_region}';
+    #                 SET s3_access_key_id = '{s3_access_key_id}';
+    #                 SET s3_secret_access_key = '{s3_secret_access_key}';
+    #                 """)
+    # else:
+    #     con.sql("""
+    #                 RESET s3_region;
+    #                 RESET s3_access_key_id;
+    #                 RESET s3_secret_access_key;
+    #     """)
+    # con.sql("RESET s3_session_token")
+    # print(con.sql("SELECT * FROM duckdb_settings() WHERE name like 's3%'"))
+
+
+def setup_gcs_connection(con, server):
+    key_id = os.getenv("DATACONTRACT_GCS_KEY_ID")
+    secret = os.getenv("DATACONTRACT_GCS_SECRET")
+
+    if key_id is None:
+        raise ValueError("Error: Environment variable DATACONTRACT_GCS_KEY_ID is not set")
+    if secret is None:
+        raise ValueError("Error: Environment variable DATACONTRACT_GCS_SECRET is not set")
+
+    con.sql(f"""
+    CREATE SECRET gcs_secret (
+        TYPE GCS,
+        KEY_ID '{key_id}',
+        SECRET '{secret}'
+    );
+    """)
+
+
+def setup_azure_connection(con, server):
+    tenant_id = os.getenv("DATACONTRACT_AZURE_TENANT_ID")
+    client_id = os.getenv("DATACONTRACT_AZURE_CLIENT_ID")
+    client_secret = os.getenv("DATACONTRACT_AZURE_CLIENT_SECRET")
+    storage_account = server.storageAccount
+
+    if tenant_id is None:
+        raise ValueError("Error: Environment variable DATACONTRACT_AZURE_TENANT_ID is not set")
+    if client_id is None:
+        raise ValueError("Error: Environment variable DATACONTRACT_AZURE_CLIENT_ID is not set")
+    if client_secret is None:
+        raise ValueError("Error: Environment variable DATACONTRACT_AZURE_CLIENT_SECRET is not set")
+
+    con.install_extension("azure")
+    con.load_extension("azure")
+
+    if storage_account is not None:
+        con.sql(f"""
+        CREATE SECRET azure_spn (
+            TYPE AZURE,
+            PROVIDER SERVICE_PRINCIPAL,
+            TENANT_ID '{tenant_id}',
+            CLIENT_ID '{client_id}',
+            CLIENT_SECRET '{client_secret}',
+            ACCOUNT_NAME '{storage_account}'
+        );
+        """)
+    else:
+        con.sql(f"""
+        CREATE SECRET azure_spn (
+            TYPE AZURE,
+            PROVIDER SERVICE_PRINCIPAL,
+            TENANT_ID '{tenant_id}',
+            CLIENT_ID '{client_id}',
+            CLIENT_SECRET '{client_secret}'
+        );
+        """)