diff --git a/CHANGELOG.md b/CHANGELOG.md index 27f89a4a7..de35f2966 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +### Added + +- Databricks backend now supports recursive array and struct checks via CTE-based virtual models, enabling read-only contract validation on warehouses with limited CREATE permissions (#1278). + ## [1.0.1] - 2026-06-10 ### Added diff --git a/README.md b/README.md index 4e54df029..113f1832d 100644 --- a/README.md +++ b/README.md @@ -2523,6 +2523,20 @@ Ensure you have a JDK 17 or 21 installed. Java 25 causes issues. java --version ``` +#### Linux system package for postgres/psycopg-based tests + +**Ubuntu/Debian:** +```bash +sudo apt-get update +sudo apt-get install -y libpq-dev +``` +**Fedora/RHEL:** +```bash +# Fedora/RHEL: +sudo dnf install -y postgresql-devel +# Arch: +sudo pacman -S postgresql-libs +``` ### Docker Build diff --git a/datacontract/engines/checks/create_checks.py b/datacontract/engines/checks/create_checks.py index 20eaeba01..ff4c7590b 100644 --- a/datacontract/engines/checks/create_checks.py +++ b/datacontract/engines/checks/create_checks.py @@ -23,10 +23,14 @@ from datacontract.engines.checks.check_spec import CheckSpec, MetricType, Op, Threshold from datacontract.engines.checks.type_normalize import normalize_type_name +from datacontract.model.server import get_server_type logger = logging.getLogger(__name__) _FILE_SERVER_TYPES = {"local", "s3", "gcs", "azure"} +_VERIFIED_NESTED_SQL_SERVER_TYPES = {"dataframe", "databricks"} +_SUPPORTED_NESTED_STRUCT_SERVER_TYPES = {"dataframe", "databricks"} +_SUPPORTED_NESTED_ARRAY_SERVER_TYPES = {"dataframe", "databricks"} # --------------------------------------------------------------------------- @@ -60,6 +64,39 @@ def expected_type_category(prop: SchemaProperty) -> tuple[str, str]: return normalize_type_name(label), (label or "") +def _property_type(prop: SchemaProperty) -> str: + return normalize_type_name(prop.physicalType or prop.logicalType) + + +def _iter_property_paths( + model: str, + properties: list[SchemaProperty] | None, + server_type: str | None, + prefix: str | None = None, + nested: bool = False, +): + for prop in properties or []: + field = prop.physicalName or prop.name + field_path = f"{prefix}.{field}" if prefix else field + yield model, field_path, prop, nested + + prop_type = _property_type(prop) + if ( + server_type in _SUPPORTED_NESTED_STRUCT_SERVER_TYPES + and prop_type in {"object", "record", "struct"} + and prop.properties + ): + yield from _iter_property_paths(model, prop.properties, server_type, field_path, True) + elif ( + server_type in _SUPPORTED_NESTED_ARRAY_SERVER_TYPES + and prop_type == "array" + and prop.items + and prop.items.properties + ): + nested_model = f"{model}__{field}" + yield from _iter_property_paths(nested_model, prop.items.properties, server_type, None, True) + + _PERCENT_UNITS = {"percent", "percentage", "%"} @@ -146,25 +183,24 @@ def create_checks( def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> List[CheckSpec]: checks: List[CheckSpec] = [] - server_type = server.type if server and server.type else None + server_type = get_server_type(server) if server is not None else None model = to_schema_name(schema_object, server_type) properties = schema_object.properties or [] check_types = is_check_types(server) uses_raw_view = ( - server is not None and server.type in _FILE_SERVER_TYPES and server.format in ("csv", "parquet", "json") + server is not None and server_type in _FILE_SERVER_TYPES and server.format in ("csv", "parquet", "json") ) - for prop in properties: + for item_model, field, prop, is_nested in _iter_property_paths(model, properties, server_type): # ODCS physicalName is the real column; mirror to_schema_name at field level. - field = prop.physicalName or prop.name checks.append( CheckSpec( - key=f"{model}__{field}__field_is_present", + key=f"{item_model}__{field}__field_is_present", category="schema", type="field_is_present", name=f"Check that field '{field}' is present", - model=model, + model=item_model, field=field, metric=MetricType.FIELD_PRESENT, uses_raw_view=uses_raw_view, @@ -175,11 +211,11 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> category, label = expected_type_category(prop) checks.append( CheckSpec( - key=f"{model}__{field}__field_type", + key=f"{item_model}__{field}__field_type", category="schema", type="field_type", name=f"Check that field {field} has type {label}", - model=model, + model=item_model, field=field, metric=MetricType.FIELD_TYPE, expected_category=category, @@ -190,7 +226,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> if prop.required: checks.append( _missing_count_check( - model, + item_model, field, "field_required", Threshold(Op.EQ, 0), @@ -201,7 +237,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> if prop.unique: checks.append( _duplicate_count_check( - model, + item_model, field, "field_unique", Threshold(Op.EQ, 0), @@ -214,7 +250,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> if min_length is not None: checks.append( _invalid_count_check( - model, + item_model, field, "field_min_length", name=f"Check that field {field} has a min length of {min_length}", @@ -226,7 +262,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> if max_length is not None: checks.append( _invalid_count_check( - model, + item_model, field, "field_max_length", name=f"Check that field {field} has a max length of {max_length}", @@ -238,7 +274,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> if minimum is not None: checks.append( _invalid_count_check( - model, + item_model, field, "field_minimum", name=f"Check that field {field} has a minimum of {minimum}", @@ -250,7 +286,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> if maximum is not None: checks.append( _invalid_count_check( - model, + item_model, field, "field_maximum", name=f"Check that field {field} has a maximum of {maximum}", @@ -262,7 +298,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> if exclusive_minimum is not None: checks.append( _invalid_count_check( - model, + item_model, field, "field_minimum", name=f"Check that field {field} has a minimum of {exclusive_minimum}", @@ -271,7 +307,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> ) checks.append( _invalid_count_check( - model, + item_model, field, "field_not_equal", name=f"Check that field {field} is not equal to {exclusive_minimum}", @@ -283,7 +319,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> if exclusive_maximum is not None: checks.append( _invalid_count_check( - model, + item_model, field, "field_maximum", name=f"Check that field {field} has a maximum of {exclusive_maximum}", @@ -292,7 +328,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> ) checks.append( _invalid_count_check( - model, + item_model, field, "field_not_equal", name=f"Check that field {field} is not equal to {exclusive_maximum}", @@ -304,7 +340,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> if pattern is not None: checks.append( _invalid_count_check( - model, + item_model, field, "field_regex", name=f"Check that field {field} matches regex pattern {pattern}", @@ -316,7 +352,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> if enum_values: checks.append( _invalid_count_check( - model, + item_model, field, "field_enum", name=f"Check that field {field} only contains enum values {enum_values}", @@ -325,7 +361,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> ) if prop.quality: - checks.extend(_quality_checks(model, field, prop.quality, server)) + checks.extend(_quality_checks(item_model, field, prop.quality, server, is_nested=is_nested)) if schema_object.quality: checks.extend(_quality_checks(model, None, schema_object.quality, server)) @@ -421,7 +457,7 @@ def _row_count_check(model, threshold: Threshold, severity=None) -> CheckSpec: # quality list # --------------------------------------------------------------------------- def _quality_checks( - model: str, field: Optional[str], quality_list: List[DataQuality], server: Optional[Server] + model: str, field: Optional[str], quality_list: List[DataQuality], server: Optional[Server], is_nested: bool = False ) -> List[CheckSpec]: checks: List[CheckSpec] = [] count = 0 @@ -444,6 +480,31 @@ def _quality_checks( ) ) elif quality.type == "sql": + server_type = get_server_type(server) if server is not None else None + if is_nested and server_type not in _VERIFIED_NESTED_SQL_SERVER_TYPES: + if field is None: + check_key = f"{model}__quality_sql_{count}" + check_type = "model_quality_sql" + else: + check_key = f"{model}__{field}__quality_sql_{count}" + check_type = "field_quality_sql" + checks.append( + CheckSpec( + key=check_key, + category="quality", + type=check_type, + name=quality.description or "Quality Check", + model=model, + field=field, + metric=MetricType.UNSUPPORTED, + preset_result="warning", + preset_reason=( + "Nested SQL quality checks are only verified for Spark (dataframe) and Databricks." + ), + ) + ) + count += 1 + continue if field is None: check_key = f"{model}__quality_sql_{count}" check_type = "model_quality_sql" diff --git a/datacontract/engines/ibis/connections/connect.py b/datacontract/engines/ibis/connections/connect.py index d70743918..eebe25d9a 100644 --- a/datacontract/engines/ibis/connections/connect.py +++ b/datacontract/engines/ibis/connections/connect.py @@ -82,6 +82,9 @@ def connect_ibis( "please provide one with the DataContract class" ) return None + from datacontract.engines.ibis.connections.kafka import add_spark_nested_views_for_contract + + add_spark_nested_views_for_contract(spark, data_contract, schema_name=schema_name) return ibis.pyspark.connect(session=spark) if server_type == "databricks": @@ -90,8 +93,26 @@ def connect_ibis( database_name = ".".join(filter(None, [server.catalog, server.schema_])) if database_name: spark.sql(f"USE {database_name}") + from datacontract.engines.ibis.connections.kafka import add_spark_nested_views_for_contract + + add_spark_nested_views_for_contract(spark, data_contract, schema_name=schema_name) return ibis.pyspark.connect(session=spark) - return _connect_databricks(ibis, server, run) + backend = _connect_databricks(ibis, server, run) + # Wire in CTE-based virtual models for nested array checks (read-only, no CREATE TABLE). + from datacontract.engines.ibis.connections.databricks_nested_models import ( + build_databricks_virtual_model_queries_for_contract, + ) + + if backend and data_contract: + virtual_queries = build_databricks_virtual_model_queries_for_contract( + data_contract, schema_name=schema_name + ) + if virtual_queries: + try: + setattr(backend, "_dc_virtual_model_queries", virtual_queries) + except Exception: + logger.debug("Could not attach databricks virtual model queries", exc_info=True) + return backend if server_type == "postgres": return ibis.postgres.connect( @@ -204,22 +225,34 @@ def connect_ibis( def _connect_databricks(ibis, server: Server, run: Run): """Connect to Databricks SQL directly, selecting the auth method from env vars. - + Uses a _NoVolumeBackend subclass to skip ibis' hardcoded CREATE VOLUME call, + enabling read-only contract checks on Databricks warehouses. Auth is resolved in priority order, so an existing token-based setup keeps working unchanged: - - 1. personal access token (``DATACONTRACT_DATABRICKS_TOKEN``) — the default + 1. personal access token (DATACONTRACT_DATABRICKS_TOKEN) - the default 2. OAuth machine-to-machine / service principal, from - ``DATACONTRACT_DATABRICKS_CLIENT_ID`` + ``DATACONTRACT_DATABRICKS_CLIENT_SECRET`` + DATACONTRACT_DATABRICKS_CLIENT_ID + DATACONTRACT_DATABRICKS_CLIENT_SECRET (the usual choice for CI/CD) - 3. a local Databricks config profile (``DATACONTRACT_DATABRICKS_PROFILE``), + 3. a local Databricks config profile (DATACONTRACT_DATABRICKS_PROFILE), delegating to the Databricks SDK's unified auth (also covers Azure CLI/MSI) - 4. an explicit connector ``auth_type`` (``DATACONTRACT_DATABRICKS_AUTH_TYPE``), - e.g. ``databricks-oauth`` for the interactive user-to-machine browser flow - - The OAuth credential providers build their SDK ``Config`` lazily, so token + 4. an explicit connector auth_type (DATACONTRACT_DATABRICKS_AUTH_TYPE), + e.g. databricks-oauth for the interactive user-to-machine browser flow + The OAuth credential providers build their SDK Config lazily, so token exchange happens when the connection is opened rather than while reading env. """ + from ibis.backends.databricks import Backend as DatabricksBackend + + class _NoVolumeBackend(DatabricksBackend): + """Databricks ibis backend that skips CREATE VOLUME on connect. + ibis calls _post_connect() during do_connect(), which by default issues + CREATE VOLUME IF NOT EXISTS. For read-only contract checks (the common case), + this volume is unnecessary and often forbidden by warehouse permissions. + Overriding _post_connect as a no-op allows the connection to succeed. + """ + + def _post_connect(self, *, memtable_volume) -> None: + pass # No-op: skip CREATE VOLUME, not needed for read-only checks. + host = server.host or require_env("DATACONTRACT_DATABRICKS_SERVER_HOSTNAME", server_type="databricks") kwargs = dict( server_hostname=host, @@ -227,37 +260,32 @@ def _connect_databricks(ibis, server: Server, run: Run): catalog=server.catalog, schema=server.schema_, ) - token = os.getenv("DATACONTRACT_DATABRICKS_TOKEN") client_id = os.getenv("DATACONTRACT_DATABRICKS_CLIENT_ID") client_secret = os.getenv("DATACONTRACT_DATABRICKS_CLIENT_SECRET") profile = os.getenv("DATACONTRACT_DATABRICKS_PROFILE") auth_type = os.getenv("DATACONTRACT_DATABRICKS_AUTH_TYPE") - + backend = _NoVolumeBackend() if token: run.log_info("Connecting to databricks with a personal access token") - return ibis.databricks.connect(access_token=token, **kwargs) - + return backend.connect(access_token=token, **kwargs) if client_id and client_secret: run.log_info("Connecting to databricks with an OAuth service principal (M2M)") sdk_host = host if host.startswith("http") else f"https://{host}" kwargs["credentials_provider"] = _databricks_credentials_provider( host=sdk_host, client_id=client_id, client_secret=client_secret ) - return ibis.databricks.connect(**kwargs) - + return backend.connect(**kwargs) if profile: run.log_info(f"Connecting to databricks with config profile '{profile}'") kwargs["credentials_provider"] = _databricks_credentials_provider(profile=profile) - return ibis.databricks.connect(**kwargs) - + return backend.connect(**kwargs) if auth_type: run.log_info(f"Connecting to databricks with auth_type '{auth_type}'") - return ibis.databricks.connect(auth_type=auth_type, **kwargs) - + return backend.connect(auth_type=auth_type, **kwargs) # Nothing configured: fail with the same clear message as before. token = require_env("DATACONTRACT_DATABRICKS_TOKEN", server_type="databricks") - return ibis.databricks.connect(access_token=token, **kwargs) + return backend.connect(access_token=token, **kwargs) def _databricks_credentials_provider(**config_kwargs): diff --git a/datacontract/engines/ibis/connections/databricks_nested_models.py b/datacontract/engines/ibis/connections/databricks_nested_models.py new file mode 100644 index 000000000..1b1de78db --- /dev/null +++ b/datacontract/engines/ibis/connections/databricks_nested_models.py @@ -0,0 +1,98 @@ +"""Build CTE-based virtual models for recursive Databricks nested checks. + +For array item checks targeting `{model}__{array_field}`, this module generates +SQL WITH clauses that explode the array and expose nested properties as columns, +allowing read-only recursive checks without CREATE TABLE/CREATE VOLUME. + +Example: For an `orders` table with `items ARRAY>`, +the virtual model `orders__items` resolves to: + + WITH __dc_source__ AS (SELECT * FROM orders) + SELECT __dc_nested__.* FROM __dc_source__ + LATERAL VIEW OUTER explode_outer(`items`) AS __dc_nested__ + +This is then available to checks targeting `orders__items` (nested array model). +""" + +from __future__ import annotations + +from open_data_contract_standard.model import OpenDataContractStandard, SchemaProperty + + +def build_databricks_virtual_model_queries_for_contract( + data_contract: OpenDataContractStandard, + schema_name: str = "all", +) -> dict[str, str]: + """Build a dict of {model_name: CTE_query} for all nested array models in the contract. + + Only models matching the optional schema filter are included. + """ + queries: dict[str, str] = {} + if data_contract.schema_ is None: + return queries + + for schema_obj in data_contract.schema_: + if schema_name != "all" and schema_obj.name != schema_name: + continue + model = schema_obj.physicalName or schema_obj.name + specs: dict[str, dict[str, str]] = {} + _collect_databricks_virtual_model_specs(model, schema_obj.properties, specs) + for virtual_model, query in _render_virtual_models(model, specs).items(): + queries[virtual_model] = query + + return queries + + +def _collect_databricks_virtual_model_specs( + parent_model: str, + properties: list[SchemaProperty] | None, + specs: dict[str, dict[str, str]], + prefix: str | None = None, +): + """Recursively collect array specs that need virtual models. + + For each array field found, record (parent_model, field_name, item_type_properties) + so we can render its CTE later. + """ + for prop in properties or []: + field = prop.physicalName or prop.name + field_path = f"{prefix}.{field}" if prefix else field + prop_type = (prop.physicalType or prop.logicalType or "").lower() + + # Array of struct: create a virtual model for this array. + if prop_type == "array" and prop.items and prop.items.properties: + virtual_model = f"{parent_model}__{field}" + specs[virtual_model] = { + "parent": parent_model, + "field": field, + "template": "SELECT __dc_nested__.* FROM {source} LATERAL VIEW OUTER explode_outer(`{field}`) AS __dc_nested__", + } + + # Struct with nested properties: recurse for any nested arrays. + if prop_type in {"object", "record", "struct"} and prop.properties: + _collect_databricks_virtual_model_specs(parent_model, prop.properties, specs, field_path) + + +def _render_virtual_models(model: str, specs: dict[str, dict[str, str]]) -> dict[str, str]: + """Render CTE SQL for all virtual models, handling dependencies.""" + queries: dict[str, str] = {} + for virtual_model in specs: + queries[virtual_model] = _render_databricks_virtual_model_query(virtual_model, specs) + return queries + + +def _render_databricks_virtual_model_query(model: str, specs: dict[str, dict[str, str]]) -> str: + """Render a single CTE query for a virtual model, recursively handling parent deps.""" + spec = specs[model] + parent = spec["parent"] + + # If parent is also virtual, render its CTE first; otherwise, SELECT from the real table. + if parent in specs: + parent_query = _render_databricks_virtual_model_query(parent, specs) + else: + parent_query = f"SELECT * FROM {parent}" + + source = "__dc_source__" + template = spec["template"] + field = spec["field"] + return f"WITH {source} AS ({parent_query}) {template.format(source=source, field=field)}" diff --git a/datacontract/engines/ibis/connections/kafka.py b/datacontract/engines/ibis/connections/kafka.py index 2b2234f41..e3fd055ce 100644 --- a/datacontract/engines/ibis/connections/kafka.py +++ b/datacontract/engines/ibis/connections/kafka.py @@ -222,6 +222,66 @@ def _get_type(prop: SchemaProperty) -> Optional[str]: return None +def _field_name(prop: SchemaProperty) -> str: + return prop.physicalName or prop.name + + +def add_spark_nested_views(spark, model_name: str, properties: List[SchemaProperty] | None): + """Create Spark temp views for nested struct fields and array-of-struct items. + + The engine-neutral recursive check builder targets array item checks at + ``{model}__{array_field}``, mirroring the DuckDB nested-view convention. + For struct fields, the executor resolves dotted paths against the parent + model directly, but we still recurse here so arrays nested under structs can + materialize their own item views. + """ + if not properties: + return + + try: + from pyspark.sql import functions as F + except ImportError as e: + raise DataContractException( + type="schema", + result="failed", + name="pyspark is missing", + reason="Install the extra datacontract-cli[kafka] to use kafka", + engine="datacontract", + original_exception=e, + ) + + parent = spark.table(model_name) + nested_alias = "__dc_nested__" + for prop in properties: + field_name = _field_name(prop) + field_type = (_get_type(prop) or "").lower() + + if field_type in {"object", "record", "struct"} and prop.properties: + child = parent.select(F.col(f"`{field_name}`").alias(nested_alias)) + if not prop.required: + child = child.where(F.col(nested_alias).isNotNull()) + child.select(f"{nested_alias}.*").createOrReplaceTempView(f"{model_name}__{field_name}") + add_spark_nested_views(spark, f"{model_name}__{field_name}", prop.properties) + + elif field_type == "array" and prop.items and prop.items.properties: + child = parent + if not prop.required: + child = child.where(F.col(f"`{field_name}`").isNotNull()) + child = child.select(F.explode_outer(F.col(f"`{field_name}`")).alias(nested_alias)) + child.select(f"{nested_alias}.*").createOrReplaceTempView(f"{model_name}__{field_name}") + add_spark_nested_views(spark, f"{model_name}__{field_name}", prop.items.properties) + + +def add_spark_nested_views_for_contract(spark, data_contract: OpenDataContractStandard, schema_name: str = "all"): + if not data_contract.schema_: + return + for schema_obj in data_contract.schema_: + model_name = schema_obj.physicalName or schema_obj.name + if schema_name != "all" and schema_obj.name != schema_name: + continue + add_spark_nested_views(spark, model_name, schema_obj.properties) + + def to_struct_type(properties: List[SchemaProperty]): try: from pyspark.sql.types import StructType diff --git a/datacontract/engines/ibis/ibis_check_execute.py b/datacontract/engines/ibis/ibis_check_execute.py index 77bcdd739..a658f075e 100644 --- a/datacontract/engines/ibis/ibis_check_execute.py +++ b/datacontract/engines/ibis/ibis_check_execute.py @@ -193,11 +193,11 @@ def _run_model( if spec.metric == MetricType.ROW_COUNT: named = t.count().name(spec.key) elif spec.metric == MetricType.MISSING_COUNT: - col = _resolve_col(columns, spec.field) - named = _count_true(_missing_expr(t, col, spec.missing_values)).name(spec.key) + col = _resolve_expr(t, columns, spec.field) + named = _count_true(_missing_expr(col, spec.missing_values)).name(spec.key) elif spec.metric == MetricType.INVALID_COUNT: - col = _resolve_col(columns, spec.field) - expr = _invalid_expr(t, col, schema[col], spec) + col = _resolve_expr(t, columns, spec.field) + expr = _invalid_expr(t, col, _resolve_dtype(schema, spec.field), spec) if expr is None: # No validity constraints => nothing can be invalid. _set_impl(run, spec.key, "invalid_count = 0 (no validity constraints configured)", None) @@ -207,7 +207,7 @@ def _run_model( elif spec.metric == MetricType.DUPLICATE_COUNT: _run_duplicate(run, t, columns, spec) elif spec.metric == MetricType.FIELD_PRESENT: - _run_present(run, con, model, columns, spec) + _run_present(run, con, model, columns, schema, spec) elif spec.metric == MetricType.FIELD_TYPE: _run_type(run, schema, columns, spec) elif spec.metric in (MetricType.FRESHNESS, MetricType.RETENTION): @@ -345,11 +345,11 @@ def _samples_for(t, columns, schema, spec: CheckSpec, identifiers, sensitive): if spec.metric == MetricType.DUPLICATE_COUNT: return _duplicate_samples(t, columns, sensitive, spec) - col = _resolve_col(columns, spec.field) + col = _resolve_expr(t, columns, spec.field) if spec.metric == MetricType.MISSING_COUNT: - predicate = _missing_expr(t, col, spec.missing_values) + predicate = _missing_expr(col, spec.missing_values) else: # INVALID_COUNT - predicate = _invalid_expr(t, col, schema[col], spec) + predicate = _invalid_expr(t, col, _resolve_dtype(schema, spec.field), spec) if predicate is None: return None @@ -405,12 +405,12 @@ def _count_true(bool_expr): return bool_expr.ifelse(1, 0).sum() -def _missing_expr(t, col, missing_values): - cond = t[col].isnull() +def _missing_expr(col, missing_values): + cond = col.isnull() if missing_values: non_null = [v for v in missing_values if v is not None] if non_null: - cond = cond | t[col].isin(non_null) + cond = cond | col.isin(non_null) return cond @@ -506,17 +506,17 @@ def _valid_expr(t, col, dtype, spec: CheckSpec): """Boolean: a non-missing value satisfies all configured validity constraints.""" conds = [] if spec.valid_values is not None: - conds.append(t[col].isin(spec.valid_values)) + conds.append(col.isin(spec.valid_values)) if spec.valid_regex is not None: - conds.append(_regex_search_expr(t, _as_string(t[col], dtype), spec.valid_regex)) + conds.append(_regex_search_expr(t, _as_string(col, dtype), spec.valid_regex)) if spec.valid_min is not None: - conds.append(t[col] >= spec.valid_min) + conds.append(col >= spec.valid_min) if spec.valid_max is not None: - conds.append(t[col] <= spec.valid_max) + conds.append(col <= spec.valid_max) if spec.valid_min_length is not None: - conds.append(_as_string(t[col], dtype).length() >= spec.valid_min_length) + conds.append(_as_string(col, dtype).length() >= spec.valid_min_length) if spec.valid_max_length is not None: - conds.append(_as_string(t[col], dtype).length() <= spec.valid_max_length) + conds.append(_as_string(col, dtype).length() <= spec.valid_max_length) if not conds: return None expr = conds[0] @@ -527,13 +527,13 @@ def _valid_expr(t, col, dtype, spec: CheckSpec): def _invalid_expr(t, col, dtype, spec: CheckSpec): """Reproduce soda's invalid_count: NOT missing AND (NOT valid OR in invalid_values).""" - missing = _missing_expr(t, col, spec.missing_values) + missing = _missing_expr(col, spec.missing_values) valid = _valid_expr(t, col, dtype, spec) invalid_terms = [] if valid is not None: invalid_terms.append(~valid) if spec.invalid_values: - invalid_terms.append(t[col].isin(spec.invalid_values)) + invalid_terms.append(col.isin(spec.invalid_values)) if not invalid_terms: return None invalid_any = invalid_terms[0] @@ -572,7 +572,7 @@ def _constraint_info(spec: CheckSpec) -> dict: # dedicated check runners # --------------------------------------------------------------------------- def _run_duplicate(run: Run, t, columns, spec: CheckSpec): - cols = [_resolve_col(columns, c) for c in (spec.columns or [spec.field])] + cols = [_resolve_expr(t, columns, c) for c in (spec.columns or [spec.field])] grouped = t.group_by(cols).aggregate(_dup_n=t.count()) dup_groups = grouped.filter(grouped["_dup_n"] > 1) _record_sql(run, spec, dup_groups) @@ -583,17 +583,21 @@ def _run_duplicate(run: Run, t, columns, spec: CheckSpec): _update_diagnostics(run, spec.key, {"columns": cols}) -def _run_present(run: Run, con, model: str, columns, spec: CheckSpec): +def _run_present(run: Run, con, model: str, columns, schema, spec: CheckSpec): target = f"{model}__raw__" if spec.uses_raw_view else model _set_impl(run, spec.key, f"column '{spec.field}' exists in {target}", "introspection") - present = set(columns.keys()) if spec.uses_raw_view: try: - raw = con.table(f"{model}__raw__") - present = {c.lower() for c in raw.columns} + raw = _resolve_table(con, f"{model}__raw__") + table = raw except Exception: - pass - ok = spec.field.lower() in present + table = _resolve_table(con, model) + target_schema = table.schema() + else: + # Reuse the already-resolved model schema to avoid an extra lookup that + # can fail on case-sensitive backends (for example Oracle). + target_schema = schema + ok = _field_present(target_schema, spec.field) _set_diagnostics(run, spec.key, _diag(metric="field_present", field=spec.field, present=ok)) _set_result( run, @@ -611,12 +615,11 @@ def _run_type(run: Run, schema, columns, spec: CheckSpec): "introspection", ) expected_label = f"{spec.expected_type_label} ({spec.expected_category})" - actual_col = columns.get(spec.field.lower()) - if actual_col is None: + dtype = _resolve_dtype(schema, spec.field) + if dtype is None: _set_diagnostics(run, spec.key, _diag(metric="field_type", field=spec.field, expected=expected_label)) _set_result(run, spec.key, ResultEnum.failed, f"Column '{spec.field}' is missing") return - dtype = schema[actual_col] actual_category = ibis_dtype_category(dtype) _set_diagnostics( run, @@ -638,8 +641,8 @@ def _run_type(run: Run, schema, columns, spec: CheckSpec): def _run_freshness(run: Run, t, columns, spec: CheckSpec): import pandas as pd - col = _resolve_col(columns, spec.field) - reduction = t[col].min() if spec.metric == MetricType.RETENTION else t[col].max() + col = _resolve_expr(t, columns, spec.field) + reduction = col.min() if spec.metric == MetricType.RETENTION else col.max() _record_sql(run, spec, t.aggregate(value=reduction)) raw = reduction.execute() if raw is None or pd.isna(raw): @@ -846,11 +849,67 @@ def _resolve_col(columns: dict, field: str) -> str: return actual +def _resolve_expr(t, columns: dict, field: str): + if field is None: + raise _ColumnNotFound("Column 'None' not found") + if "." not in field: + return t[_resolve_col(columns, field)] + return _resolve_nested_expr(t, field, columns) + + +def _resolve_nested_expr(t, field: str, columns: dict): + expr = t[_resolve_col(columns, field.split(".", 1)[0])] + for part in field.split(".")[1:]: + expr = expr[part] + return expr + + +def _resolve_dtype(schema, field: str): + if field is None: + return None + current = schema + parts = field.split(".") + dtype = None + for idx, part in enumerate(parts): + try: + dtype = current[part] + except Exception: + return None + if idx < len(parts) - 1: + try: + current = dtype.fields + except Exception: + return None + return dtype + + +def _field_present(schema, field: str) -> bool: + if field is None: + return False + return _resolve_dtype(schema, field) is not None + + def _resolve_table(con, model: str): - """Resolve a table by name, tolerating case differences across dialects.""" + """Resolve a table by name, tolerating case differences and virtual models. + + Falls back to CTE-based virtual models (e.g. for Databricks nested array + checks) before trying list_tables(). Virtual models are stored on the + connection object as _dc_virtual_model_queries. + """ try: return con.table(model) except Exception: + # Try virtual models (Databricks nested array CTEs). + virtual_queries = getattr(con, "_dc_virtual_model_queries", None) + if isinstance(virtual_queries, dict): + query = virtual_queries.get(model) + if query is None: + # Case-insensitive match. + match = next((name for name in virtual_queries if name.lower() == model.lower()), None) + query = virtual_queries.get(match) if match else None + if query: + return con.sql(query) + # Fall back to list_tables for case-insensitive real table lookup. try: available = con.list_tables() except Exception: diff --git a/tests/fixtures/dataframe/datacontract_nested.yaml b/tests/fixtures/dataframe/datacontract_nested.yaml new file mode 100644 index 000000000..fe3af15e8 --- /dev/null +++ b/tests/fixtures/dataframe/datacontract_nested.yaml @@ -0,0 +1,44 @@ +dataContractSpecification: 1.2.1 +id: dataframetestnested +info: + title: dataframetestnested + version: 0.0.1 + owner: my-domain-team +servers: + unittest: + type: dataframe +models: + my_nested_table: + type: table + fields: + id: + type: varchar + required: true + unique: true + user: + type: struct + required: true + fields: + email: + type: varchar + required: true + pattern: "^.+@.+$" + status: + type: varchar + quality: + - type: sql + query: SELECT COUNT(*) FROM {model} WHERE {field} NOT IN ('active', 'inactive') + mustBe: 0 + line_items: + type: array + items: + type: struct + fields: + sku: + type: varchar + required: true + status: + type: varchar + enum: + - ok + - hold diff --git a/tests/test_connect_databricks.py b/tests/test_connect_databricks.py index d42ebb147..aeb639c0f 100644 --- a/tests/test_connect_databricks.py +++ b/tests/test_connect_databricks.py @@ -1,11 +1,11 @@ """Unit tests for Databricks auth-method selection in connect_ibis. -These do not hit Databricks: ``ibis.databricks.connect`` is patched and we only +These do not hit Databricks: we patch the _NoVolumeBackend.connect method and only assert which auth kwargs the dispatch passes for a given set of env vars. """ -import ibis import pytest +from ibis.backends.databricks import Backend as DatabricksBackend from open_data_contract_standard.model import Server from datacontract.engines.ibis.connections.connect import connect_ibis @@ -31,14 +31,23 @@ def clean_databricks_env(monkeypatch): @pytest.fixture def captured_connect(monkeypatch): - """Patch ibis.databricks.connect to record the kwargs it is called with.""" + """Patch DatabricksBackend.connect to record the kwargs it is called with.""" calls = {} - def fake_connect(**kwargs): + def fake_connect(self, **kwargs): calls.update(kwargs) - return "connection" - - monkeypatch.setattr(ibis.databricks, "connect", fake_connect) + # Return a minimal mock backend with the required attributes for downstream code. + mock = type( + "MockBackend", + (), + { + "name": "databricks", + "_dc_virtual_model_queries": {}, + }, + )() + return mock + + monkeypatch.setattr(DatabricksBackend, "connect", fake_connect) return calls @@ -56,7 +65,7 @@ def test_personal_access_token_is_default(clean_databricks_env, captured_connect result = _connect() - assert result == "connection" + assert result is not None assert captured_connect["access_token"] == "dapiTOKEN" assert captured_connect["http_path"] == "/sql/1.0/warehouses/abc" assert captured_connect["server_hostname"] == "dbc-x.cloud.databricks.com" @@ -119,3 +128,38 @@ def test_missing_auth_raises(clean_databricks_env, captured_connect): with pytest.raises(DataContractException): _connect() + + +def test_no_create_volume_on_connect(clean_databricks_env, monkeypatch): + """_post_connect must never execute (no CREATE VOLUME).""" + clean_databricks_env.setenv("DATACONTRACT_DATABRICKS_TOKEN", "dapiTOKEN") + clean_databricks_env.setenv("DATACONTRACT_DATABRICKS_HTTP_PATH", "/sql/1.0/warehouses/abc") + + post_connect_called = [] + + original_post_connect = DatabricksBackend._post_connect + + def spy_post_connect(self, *, memtable_volume): + post_connect_called.append(memtable_volume) + original_post_connect(self, memtable_volume=memtable_volume) + + try: + # Patch to spy on all calls (including _NoVolumeBackend overrides via MRO). + DatabricksBackend._post_connect = spy_post_connect + + # Also patch connect to return early so we don't hit actual Databricks. + def fake_connect(self, **kwargs): + self.con = None + self._memtable_volume = kwargs.get("memtable_volume") + return self + + monkeypatch.setattr(DatabricksBackend, "connect", fake_connect) + + _connect() + + # Since _NoVolumeBackend overrides _post_connect with a no-op, the spy + # should never be invoked by the do_connect flow (it goes to _NoVolumeBackend's + # override first via MRO). + assert post_connect_called == [] + finally: + DatabricksBackend._post_connect = original_post_connect diff --git a/tests/test_connect_databricks_virtual_models.py b/tests/test_connect_databricks_virtual_models.py new file mode 100644 index 000000000..fd62520c2 --- /dev/null +++ b/tests/test_connect_databricks_virtual_models.py @@ -0,0 +1,81 @@ +"""Unit tests for Databricks CTE virtual model generation.""" + +from datacontract.data_contract import DataContract +from datacontract.engines.ibis.connections.databricks_nested_models import ( + build_databricks_virtual_model_queries_for_contract, +) + + +def test_builds_virtual_queries_for_structs_and_arrays(): + """Virtual models are generated for array items, not struct fields.""" + contract = """ +apiVersion: v3.0.2 +kind: DataContract +id: test-nested +version: 1.0.0 +status: active +schema: + - name: orders + properties: + - name: customer + logicalType: object + properties: + - name: email + logicalType: string + - name: items + logicalType: array + items: + logicalType: object + properties: + - name: item_id + logicalType: string + - name: qty + logicalType: integer +""" + odcs = DataContract(data_contract_str=contract).get_data_contract() + + queries = build_databricks_virtual_model_queries_for_contract(odcs) + + # Struct fields don't get virtual models; they resolve via dotted paths. + assert "orders__customer" not in queries + # Array items get virtual models with LATERAL VIEW OUTER explode_outer. + assert "orders__items" in queries + assert "LATERAL VIEW OUTER explode_outer(`items`)" in queries["orders__items"] + assert "SELECT __dc_nested__.* FROM" in queries["orders__items"] + + +def test_build_virtual_queries_respects_schema_filter(): + """Only models matching the schema_name filter are included.""" + contract = """ +apiVersion: v3.0.2 +kind: DataContract +id: test-nested +version: 1.0.0 +status: active +schema: + - name: orders + properties: + - name: items + logicalType: array + items: + logicalType: object + properties: + - name: item_id + logicalType: string + - name: shipments + properties: + - name: tracking_events + logicalType: array + items: + logicalType: object + properties: + - name: status + logicalType: string +""" + odcs = DataContract(data_contract_str=contract).get_data_contract() + + # Only include "orders" schema. + queries = build_databricks_virtual_model_queries_for_contract(odcs, schema_name="orders") + + assert "orders__items" in queries + assert "shipments__tracking_events" not in queries diff --git a/tests/test_create_checks_nested.py b/tests/test_create_checks_nested.py new file mode 100644 index 000000000..186ae9eb1 --- /dev/null +++ b/tests/test_create_checks_nested.py @@ -0,0 +1,67 @@ +from open_data_contract_standard.model import Server + +from datacontract.data_contract import DataContract +from datacontract.engines.checks.check_spec import MetricType +from datacontract.engines.checks.create_checks import create_checks + +CONTRACT = """ +apiVersion: v3.0.2 +kind: DataContract +id: nested-checks +version: 1.0.0 +status: active +schema: + - name: orders + properties: + - name: id + logicalType: string + required: true + - name: user + logicalType: object + properties: + - name: email + logicalType: string + required: true + logicalTypeOptions: + pattern: ^.+@.+$ + - name: status + logicalType: string + quality: + - type: sql + query: SELECT COUNT(*) FROM {model} WHERE {field} NOT IN ('active', 'inactive') + mustBe: 0 + - name: line_items + logicalType: array + items: + logicalType: object + properties: + - name: sku + logicalType: string + required: true +""" + + +def _checks(server_type: str): + odcs = DataContract(data_contract_str=CONTRACT).get_data_contract() + return create_checks(odcs, Server(type=server_type)) + + +def test_create_checks_recurses_for_dataframe_nested_structs_and_arrays(): + checks = _checks("dataframe") + + assert any(c.field == "user.email" and c.type == "field_required" and c.model == "orders" for c in checks) + assert any(c.field == "user.email" and c.type == "field_regex" and c.model == "orders" for c in checks) + assert any(c.field == "sku" and c.type == "field_required" and c.model == "orders__line_items" for c in checks) + + nested_sql = next(c for c in checks if c.type == "field_quality_sql") + assert nested_sql.field == "user.status" + assert nested_sql.model == "orders" + assert nested_sql.metric == MetricType.CUSTOM_SQL + assert "user.status" in (nested_sql.query or "") + + +def test_create_checks_skips_nested_checks_for_unverified_backends(): + checks = _checks("postgres") + + assert not any(c.field == "user.email" for c in checks) + assert not any(c.model == "orders__line_items" for c in checks) diff --git a/tests/test_ibis_check_execute.py b/tests/test_ibis_check_execute.py new file mode 100644 index 000000000..8fca369f3 --- /dev/null +++ b/tests/test_ibis_check_execute.py @@ -0,0 +1,71 @@ +from datacontract.engines.checks.check_spec import CheckSpec, MetricType +from datacontract.engines.ibis.ibis_check_execute import _run_present +from datacontract.model.run import Check, ResultEnum, Run + + +class _FakeTable: + def __init__(self, schema): + self._schema = schema + + def schema(self): + return self._schema + + +class _NoLookupConnection: + def table(self, _name): + raise AssertionError("table() should not be called for non-raw field presence checks") + + +class _CaseSensitiveConnection: + def __init__(self, tables): + self._tables = tables + + def table(self, name): + if name in self._tables: + return self._tables[name] + raise KeyError(name) + + def list_tables(self): + return list(self._tables.keys()) + + +def _run_with_stubbed_check(key: str = "k") -> Run: + run = Run.create_run() + run.checks = [Check(type="field_is_present", key=key)] + return run + + +def test_run_present_uses_resolved_schema_without_extra_lookup(): + run = _run_with_stubbed_check() + spec = CheckSpec( + key="k", + category="schema", + type="field_is_present", + name="field is present", + model="checks_testcase", + field="CTC_ID", + metric=MetricType.FIELD_PRESENT, + ) + + _run_present(run, _NoLookupConnection(), "checks_testcase", {"ctc_id": "CTC_ID"}, {"CTC_ID": "int64"}, spec) + + assert run.checks[0].result == ResultEnum.passed + + +def test_run_present_raw_view_falls_back_to_model_with_case_insensitive_resolution(): + run = _run_with_stubbed_check() + spec = CheckSpec( + key="k", + category="schema", + type="field_is_present", + name="field is present", + model="checks_testcase", + field="CTC_ID", + metric=MetricType.FIELD_PRESENT, + uses_raw_view=True, + ) + con = _CaseSensitiveConnection({"CHECKS_TESTCASE": _FakeTable({"CTC_ID": "int64"})}) + + _run_present(run, con, "checks_testcase", {"ctc_id": "CTC_ID"}, {"IGNORED": "int64"}, spec) + + assert run.checks[0].result == ResultEnum.passed diff --git a/tests/test_test_databricks.py b/tests/test_test_databricks.py index 6d2630830..4bda3d800 100644 --- a/tests/test_test_databricks.py +++ b/tests/test_test_databricks.py @@ -2,8 +2,11 @@ import pytest from dotenv import load_dotenv +from open_data_contract_standard.model import Server from datacontract.data_contract import DataContract +from datacontract.engines.checks.check_spec import MetricType +from datacontract.engines.checks.create_checks import create_checks # logging.basicConfig(level=logging.DEBUG, force=True) @@ -12,6 +15,47 @@ load_dotenv(override=True) +def test_nested_struct_and_array_checks_enabled_for_databricks(): + contract = """ +apiVersion: v3.0.2 +kind: DataContract +id: databricks-nested +version: 1.0.0 +status: active +schema: + - name: orders + properties: + - name: customer + logicalType: object + properties: + - name: email + logicalType: string + quality: + - type: sql + query: SELECT COUNT(*) FROM {model} WHERE {field} IS NULL + mustBe: 0 + - name: discounts + logicalType: array + items: + logicalType: object + properties: + - name: discount_code + logicalType: string + required: true +""" + odcs = DataContract(data_contract_str=contract).get_data_contract() + + checks = create_checks(odcs, Server(type="databricks")) + + nested_sql = next(c for c in checks if c.type == "field_quality_sql") + assert nested_sql.field == "customer.email" + assert nested_sql.metric == MetricType.CUSTOM_SQL + assert nested_sql.model == "orders" + assert "customer.email" in (nested_sql.query or "") + # Array models must also be generated for Databricks (via virtual CTE models). + assert any(c.model == "orders__discounts" for c in checks) + + @pytest.mark.skipif( os.environ.get("DATACONTRACT_DATABRICKS_TOKEN") is None, reason="Requires DATACONTRACT_DATABRICKS_TOKEN to be set" ) diff --git a/tests/test_test_dataframe.py b/tests/test_test_dataframe.py index c77813a74..4c26d97c7 100644 --- a/tests/test_test_dataframe.py +++ b/tests/test_test_dataframe.py @@ -21,6 +21,7 @@ # logging.basicConfig(level=logging.INFO, force=True) datacontract = "fixtures/dataframe/datacontract.yaml" +datacontract_nested = "fixtures/dataframe/datacontract_nested.yaml" load_dotenv(override=True) @@ -78,6 +79,38 @@ def test_test_dataframe_fail(spark: SparkSession): assert len(failed) == 3 +def test_test_dataframe_nested_checks_pass(spark: SparkSession): + _prepare_nested_dataframe(spark) + data_contract = DataContract( + data_contract_file=datacontract_nested, + spark=spark, + ) + + run = data_contract.test() + + print(run.pretty()) + assert run.has_passed() + assert all(check.result == "passed" for check in run.checks) + + +def test_test_dataframe_nested_checks_fail(spark: SparkSession): + _prepare_nested_fail_dataframe(spark) + data_contract = DataContract( + data_contract_file=datacontract_nested, + spark=spark, + ) + + run = data_contract.test() + + print(run.pretty()) + assert not run.has_passed() + failed = {(check.model, check.type, check.field) for check in run.checks if check.result == "failed"} + assert ("my_nested_table", "field_regex", "user.email") in failed + assert ("my_nested_table", "field_quality_sql", "user.status") in failed + assert ("my_nested_table__line_items", "field_required", "sku") in failed + assert ("my_nested_table__line_items", "field_invalid_values", "status") in failed + + schema = StructType( [ StructField("field_one", StringType(), nullable=False), @@ -102,6 +135,35 @@ def test_test_dataframe_fail(spark: SparkSession): ) +nested_schema = StructType( + [ + StructField("id", StringType(), nullable=False), + StructField( + "user", + StructType( + [ + StructField("email", StringType(), nullable=True), + StructField("status", StringType(), nullable=True), + ] + ), + nullable=True, + ), + StructField( + "line_items", + ArrayType( + StructType( + [ + StructField("sku", StringType(), nullable=True), + StructField("status", StringType(), nullable=True), + ] + ) + ), + nullable=True, + ), + ] +) + + def _prepare_dataframe(spark): data = [ Row( @@ -172,3 +234,37 @@ def _prepare_fail_dataframe(spark): # Create temporary view # Name must match the model name in the data contract df.createOrReplaceTempView("my_table") + + +def _prepare_nested_dataframe(spark): + data = [ + Row( + id="1", + user=Row(email="alice@example.com", status="active"), + line_items=[Row(sku="SKU-1", status="ok"), Row(sku="SKU-2", status="hold")], + ), + Row( + id="2", + user=Row(email="bob@example.com", status="inactive"), + line_items=[Row(sku="SKU-3", status="ok")], + ), + ] + df = spark.createDataFrame(data, schema=nested_schema) + df.createOrReplaceTempView("my_nested_table") + + +def _prepare_nested_fail_dataframe(spark): + data = [ + Row( + id="1", + user=Row(email="alice-at-example.com", status="blocked"), + line_items=[Row(sku=None, status="bad"), Row(sku="SKU-2", status="hold")], + ), + Row( + id="2", + user=Row(email="bob@example.com", status="inactive"), + line_items=[Row(sku="SKU-3", status="ok")], + ), + ] + df = spark.createDataFrame(data, schema=nested_schema) + df.createOrReplaceTempView("my_nested_table")