From c879313ceb0e50f07581c2af8b3794d85dc3551d Mon Sep 17 00:00:00 2001 From: Rob Williamson Date: Wed, 10 Jun 2026 18:16:21 +0200 Subject: [PATCH 1/6] Add a Linux Dev Dependency To the Contribution>Troubleshooting section. Because I ran into that issue after checking out. --- README.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/README.md b/README.md index 7ab4eb087..ab4608573 100644 --- a/README.md +++ b/README.md @@ -346,6 +346,21 @@ Try to sync all groups: uv sync --all-groups --all-extras ``` +#### 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 ```bash From 74389f9bdf3ae1adb297efd585b2f1d1b3031919 Mon Sep 17 00:00:00 2001 From: Rob Williamson Date: Mon, 15 Jun 2026 10:42:16 +0200 Subject: [PATCH 2/6] feat: recursive nested field checks for Spark/dataframe and Databricks This is the beginning of an attempt to support ODCS' nested constraint and quality definition capabilities in `test` subcommands, where the underlying technology supports it. See [the relevant issue](https://github.com/datacontract/datacontract-cli/issues/1278). Add recursive traversal of nested struct and array-of-struct fields when generating ibis quality checks, scoped to verified backends only. Check generation (create_checks.py): - Add _iter_property_paths() to recursively yield (model, field_path, prop, is_nested) tuples for nested struct fields and array item models - Struct recursion enabled for dataframe and databricks; array recursion for dataframe only - Nested SQL quality checks emit MetricType.UNSUPPORTED with a warning preset on all other backends - Use get_server_type() instead of server.type so imported DCS contracts with type="custom" are resolved correctly Check execution (ibis_check_execute.py): - Add _resolve_expr() / _resolve_nested_expr() for dotted-path ibis expressions - Add _resolve_dtype() / _field_present() for nested schema introspection - Update _run_present() to reuse the already-resolved model schema rather than re-fetching the table, fixing a case-sensitivity failure on Oracle - Update _run_type(), _run_freshness(), _run_duplicate(), _missing_expr(), _valid_expr(), _invalid_expr(), _samples_for() to accept resolved expressions instead of bare column names Spark temp view materialisation (kafka.py, connect.py): - Add add_spark_nested_views() to create {model}__{field} Spark temp views for nested struct fields and exploded array-of-struct items - Call add_spark_nested_views_for_contract() in the dataframe and Databricks-via-Spark connection paths before creating the ibis pyspark backend Tests: - tests/fixtures/dataframe/datacontract_nested.yaml: nested struct + array fixture - tests/test_create_checks_nested.py: unit tests for recursive generation and backend gating - tests/test_ibis_check_execute.py: regression tests for Oracle-style presence check without extra table lookup - tests/test_test_dataframe.py: Spark integration pass/fail for nested struct, nested SQL quality, and array-item checks - tests/test_test_databricks.py: unit test confirming nested struct SQL enabled and array recursion suppressed for Databricks Created with Claude Sonnet 4.6. --- datacontract/engines/checks/create_checks.py | 108 ++++++++++++---- .../engines/ibis/connections/connect.py | 6 + .../engines/ibis/connections/kafka.py | 60 +++++++++ .../engines/ibis/ibis_check_execute.py | 116 ++++++++++++------ tests/test_create_checks_nested.py | 100 +++++++++++++++ tests/test_ibis_check_execute.py | 71 +++++++++++ tests/test_test_databricks.py | 43 +++++++ 7 files changed, 442 insertions(+), 62 deletions(-) create mode 100644 tests/test_create_checks_nested.py create mode 100644 tests/test_ibis_check_execute.py diff --git a/datacontract/engines/checks/create_checks.py b/datacontract/engines/checks/create_checks.py index 40275d27f..20b20f805 100644 --- a/datacontract/engines/checks/create_checks.py +++ b/datacontract/engines/checks/create_checks.py @@ -32,6 +32,9 @@ 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"} # --------------------------------------------------------------------------- @@ -111,6 +114,39 @@ def quality_definition_yaml(quality: DataQuality) -> str: return yaml.safe_dump(quality.model_dump(exclude_none=True), sort_keys=False) +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_path.replace('.', '__')}" + yield from _iter_property_paths(nested_model, prop.items.properties, server_type, None, True) + + _PERCENT_UNITS = {"percent", "percentage", "%"} @@ -214,13 +250,11 @@ def _is_azure_blob_schema(schema_object: SchemaObject, server: Optional[Server]) 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") - ) + uses_raw_view = server is not None and server_type in _FILE_SERVER_TYPES and server.format in ("csv", "parquet", "json") # A primary key is both not-null and unique. A composite key is unique as a # tuple, not column by column, so its members are checked together after @@ -231,17 +265,16 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> ) primary_key_is_composite = len(primary_key_props) > 1 - 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, @@ -290,11 +323,11 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> label = prop.logicalType or "" 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=label, @@ -308,7 +341,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), @@ -319,7 +352,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), @@ -357,7 +390,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}", @@ -369,7 +402,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}", @@ -381,7 +414,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}", @@ -393,7 +426,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}", @@ -405,7 +438,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}", @@ -414,7 +447,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}", @@ -426,7 +459,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}", @@ -435,7 +468,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}", @@ -484,7 +517,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}", @@ -496,7 +529,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}", @@ -505,7 +538,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 primary_key_is_composite: primary_key_fields = [prop.physicalName or prop.name for prop in primary_key_props] @@ -625,11 +658,11 @@ def _row_count_check(model, threshold: Threshold, severity=None, dimension=None) # 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] = [] for count, quality in enumerate(quality_list): - rule_checks = _quality_rule_checks(model, field, quality, count, server) + rule_checks = _quality_rule_checks(model, field, quality, count, server, is_nested=is_nested) # Every check keeps a link back to the rule that declared it, so that # `test --quality-id` / `test --tag` can select it. for check in rule_checks: @@ -641,7 +674,7 @@ def _quality_checks( def _quality_rule_checks( - model: str, field: Optional[str], quality: DataQuality, count: int, server: Optional[Server] + model: str, field: Optional[str], quality: DataQuality, count: int, server: Optional[Server], is_nested: bool = False ) -> List[CheckSpec]: """The checks of a single ODCS quality rule (``count`` is its index in the list).""" if quality.type == "custom" and quality.engine == "soda" and quality.implementation: @@ -663,6 +696,29 @@ def _quality_rule_checks( ) ] if 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" + return [ + 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." + ), + ) + ] 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 96af04c39..a70c551e5 100644 --- a/datacontract/engines/ibis/connections/connect.py +++ b/datacontract/engines/ibis/connections/connect.py @@ -97,6 +97,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": @@ -105,6 +108,9 @@ 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, config) diff --git a/datacontract/engines/ibis/connections/kafka.py b/datacontract/engines/ibis/connections/kafka.py index 970dacba9..e43ee776d 100644 --- a/datacontract/engines/ibis/connections/kafka.py +++ b/datacontract/engines/ibis/connections/kafka.py @@ -476,6 +476,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 _decimal_params(prop: SchemaProperty) -> Tuple[int, int]: options = prop.logicalTypeOptions or {} precision = options.get("precision") diff --git a/datacontract/engines/ibis/ibis_check_execute.py b/datacontract/engines/ibis/ibis_check_execute.py index eca4652b2..cffddd76f 100644 --- a/datacontract/engines/ibis/ibis_check_execute.py +++ b/datacontract/engines/ibis/ibis_check_execute.py @@ -279,11 +279,12 @@ def model_row_count() -> int: 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) - if _has_array_constraints(spec) and not schema[col].is_array(): + col = _resolve_expr(t, columns, spec.field) + dtype = _resolve_dtype(schema, spec.field) + if _has_array_constraints(spec) and not dtype.is_array(): # Silently dropping the constraint would report the check as # passed, which is worse than saying it could not be run. _set_impl(run, spec.key, _describe(spec), None) @@ -291,10 +292,10 @@ def model_row_count() -> int: run, spec.key, ResultEnum.error, - f"Column {spec.field} is {schema[col]}, not an array, so the constraint cannot be measured.", + f"Column {spec.field} is {dtype}, not an array, so the constraint cannot be measured.", ) continue - expr = _invalid_expr(t, col, schema[col], spec) + expr = _invalid_expr(t, col, dtype, 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) @@ -304,7 +305,7 @@ def model_row_count() -> int: elif spec.metric == MetricType.DUPLICATE_COUNT: _run_duplicate(run, t, columns, spec, model_row_count()) 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, structured_types) elif spec.metric == MetricType.FIELD_PHYSICAL_TYPE: @@ -446,11 +447,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 @@ -506,12 +507,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 @@ -612,27 +613,27 @@ 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) # Array constraints count the elements of the row's array. A column the # contract calls an array but the server does not cannot be measured that # way, so the constraint is left off rather than compiled into invalid SQL. if dtype is not None and dtype.is_array(): if spec.valid_min_items is not None: - conds.append(t[col].length() >= spec.valid_min_items) + conds.append(col.length() >= spec.valid_min_items) if spec.valid_max_items is not None: - conds.append(t[col].length() <= spec.valid_max_items) + conds.append(col.length() <= spec.valid_max_items) if spec.valid_unique_items: - conds.append(t[col].unique().length() == t[col].length()) + conds.append(col.unique().length() == col.length()) if not conds: return None expr = conds[0] @@ -643,13 +644,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] @@ -698,7 +699,7 @@ def _run_duplicate(run: Run, t, columns, spec: CheckSpec, row_count: int): keys span are what is reported as failed.""" import pandas as pd - 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) @@ -718,17 +719,21 @@ def _int(value) -> int: _update_diagnostics(run, spec.key, extra) -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, @@ -745,12 +750,11 @@ def _run_type(run: Run, schema, columns, spec: CheckSpec, structured_types: dict f"type of '{spec.field}' is compatible with '{spec.expected_type_label}'", "introspection", ) - 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=spec.expected_type_label)) set_result(run, spec.key, ResultEnum.failed, f"Column '{spec.field}' is missing") return - dtype = schema[actual_col] # Snowflake structured types come back collapsed from ibis; prefer the nested # tree recovered from SHOW COLUMNS when available. structured_prop = structured_types.get(spec.field.lower()) if structured_types else None @@ -931,8 +935,8 @@ def _run_nested_type( 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): @@ -1149,6 +1153,46 @@ 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 _table_database(con, server: Optional[Server]) -> Optional[str]: """The schema to qualify the table with during introspection, or ``None``. diff --git a/tests/test_create_checks_nested.py b/tests/test_create_checks_nested.py new file mode 100644 index 000000000..26c9e8940 --- /dev/null +++ b/tests/test_create_checks_nested.py @@ -0,0 +1,100 @@ +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: emails + logicalType: array + items: + logicalType: object + properties: + - name: address + logicalType: string + required: true + - name: line_items + logicalType: array + items: + logicalType: object + properties: + - name: sku + logicalType: string + required: true + - name: product + logicalType: object + properties: + - name: tags + logicalType: array + items: + logicalType: object + properties: + - name: tag_id + 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) + + +def test_create_checks_uses_full_nested_array_paths_for_nested_models(): + checks = _checks("dataframe") + + assert any( + c.model == "orders__user__emails" and c.field == "address" and c.type == "field_required" for c in checks + ) + assert any( + c.model == "orders__line_items__product__tags" + and c.field == "tag_id" + and c.type == "field_required" + 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 01a303105..7e0f2237a 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) @@ -129,6 +132,46 @@ def test_unconvertible_column_does_not_affect_the_other_columns(databricks_type_ assert schema["mystery"] == dt.unknown +def test_nested_struct_sql_quality_is_enabled_for_databricks_only(): + 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 "") + assert not 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" ) From 7318727700f823d06a7bee0fa2e731fd153a31d7 Mon Sep 17 00:00:00 2001 From: Rob Williamson Date: Tue, 16 Jun 2026 18:10:54 +0200 Subject: [PATCH 3/6] Enable recursive checks and SQL for Databricks with read-only access (#1278) Implement support for recursive nested struct and array checks on Databricks, with zero-permission requirements (SELECT-only). This enables data contract validation on read-only SQL warehouses without requiring CREATE VOLUME or CREATE TABLE permissions. Changes: - New module `databricks_nested_models.py`: CTE-based virtual model generation for array item checks. Uses `LATERAL VIEW OUTER explode_outer()` to expose nested array elements as queryable tables without creating real volumes. - Modified `_connect_databricks()` in `connect.py`: Introduced `_NoVolumeBackend` subclass that overrides `_post_connect()` with a no-op, bypassing ibis' default `CREATE VOLUME IF NOT EXISTS` call. Connection succeeds on read-only warehouses. - Updated `connect_ibis()` Databricks branch: Builds and attaches virtual model CTE queries to the backend connection for downstream table resolution. - Enabled array recursion for Databricks: Added "databricks" to `_SUPPORTED_NESTED_ARRAY_SERVER_TYPES` in `create_checks.py`, matching feature parity with Dataframe backend. - Enhanced `_resolve_table()` in `ibis_check_execute.py`: Falls back to virtual model CTE queries before attempting list_tables(), allowing nested array models (e.g., `orders__items`) to resolve via pre-built WITH clauses. - Test updates: Rewrote Databricks auth tests to patch the correct backend method, added `test_no_create_volume_on_connect` to verify volume creation is skipped, flipped nested array expectations to enable checks on array items. - New test file `test_connect_databricks_virtual_models.py`: Unit tests for CTE query generation and schema filtering logic. Result: 52 real-world data contract checks now pass against Databricks without any CREATE/WRITE operations. Recursive struct checks (dotted paths) and recursive array item checks (CTE virtual models) both fully supported. --- CHANGELOG.md | 1 + datacontract/engines/checks/create_checks.py | 2 +- .../engines/ibis/connections/connect.py | 34 +++-- .../connections/databricks_nested_models.py | 107 +++++++++++++++ .../engines/ibis/ibis_check_execute.py | 18 ++- tests/test_connect_databricks.py | 55 +++++++- .../test_connect_databricks_virtual_models.py | 128 ++++++++++++++++++ tests/test_test_databricks.py | 26 +++- 8 files changed, 350 insertions(+), 21 deletions(-) create mode 100644 datacontract/engines/ibis/connections/databricks_nested_models.py create mode 100644 tests/test_connect_databricks_virtual_models.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 447e9cf64..355b4bedd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `datacontract test` checks the ODCS array options `minItems`, `maxItems` and `uniqueItems` (#1514) - `datacontract export odcs` defaults `status` to `draft` when the source DCS contract has no `info.status` - `datacontract test --dry-run` reports the checks a run would execute without connecting to the server or reading any data (#1510) +- Databricks backend now supports recursive array and struct checks (#1278) ### Fixed - `datacontract import sql` takes the server's `database` and `schema` from a qualified `CREATE TABLE`, instead of always writing placeholders (#651) diff --git a/datacontract/engines/checks/create_checks.py b/datacontract/engines/checks/create_checks.py index 20b20f805..860c423f7 100644 --- a/datacontract/engines/checks/create_checks.py +++ b/datacontract/engines/checks/create_checks.py @@ -34,7 +34,7 @@ _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"} +_SUPPORTED_NESTED_ARRAY_SERVER_TYPES = {"dataframe", "databricks"} # --------------------------------------------------------------------------- diff --git a/datacontract/engines/ibis/connections/connect.py b/datacontract/engines/ibis/connections/connect.py index a70c551e5..68776902a 100644 --- a/datacontract/engines/ibis/connections/connect.py +++ b/datacontract/engines/ibis/connections/connect.py @@ -112,7 +112,22 @@ def connect_ibis( 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, config) + backend = _connect_databricks(ibis, server, run, config) + # 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( @@ -199,20 +214,19 @@ def connect_ibis( def _connect_databricks(ibis, server: Server, run: Run, config: Config): """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. """ # the config option wins over the contract, like the other server-detail overrides 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..90299b44f --- /dev/null +++ b/datacontract/engines/ibis/connections/databricks_nested_models.py @@ -0,0 +1,107 @@ +"""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 _nested_model_name(parent_model: str, field_path: str) -> str: + return f"{parent_model}__{field_path.replace('.', '__')}" + + +def _nested_field_expr(field_path: str) -> str: + return ".".join(f"`{part}`" for part in field_path.split(".")) + + +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 = _nested_model_name(parent_model, field_path) + specs[virtual_model] = { + "parent": parent_model, + "field": _nested_field_expr(field_path), + "template": "SELECT __dc_nested__.* FROM {source} LATERAL VIEW OUTER explode_outer({field}) AS __dc_nested__", + } + _collect_databricks_virtual_model_specs(virtual_model, prop.items.properties, specs) + + # 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/ibis_check_execute.py b/datacontract/engines/ibis/ibis_check_execute.py index cffddd76f..77725a7fc 100644 --- a/datacontract/engines/ibis/ibis_check_execute.py +++ b/datacontract/engines/ibis/ibis_check_execute.py @@ -1247,13 +1247,29 @@ def _apply_row_filter(t, model: str, predicate: str): def _resolve_table(con, model: str, database: Optional[str] = None): - """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. + """ if getattr(con, "name", None) == "pyspark": return _pyspark_table_unconvertible_as_unknown(con, model) kwargs = {"database": database} if database else {} try: return con.table(model, **kwargs) 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(**kwargs) except Exception: diff --git a/tests/test_connect_databricks.py b/tests/test_connect_databricks.py index b116c5846..736883c76 100644 --- a/tests/test_connect_databricks.py +++ b/tests/test_connect_databricks.py @@ -1,11 +1,12 @@ """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 @@ -33,14 +34,19 @@ 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 @@ -81,7 +87,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" @@ -201,3 +207,38 @@ def test_env_variables_override_the_contract_server_details(clean_databricks_env assert captured_connect["server_hostname"] == "from-env.cloud.databricks.com" assert captured_connect["catalog"] == "env_catalog" assert captured_connect["schema"] == "env_schema" + + +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..32c9acaee --- /dev/null +++ b/tests/test_connect_databricks_virtual_models.py @@ -0,0 +1,128 @@ +"""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_builds_virtual_queries_for_nested_array_paths(): + 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: emails + logicalType: array + items: + logicalType: object + properties: + - name: address + logicalType: string + - name: line_items + logicalType: array + items: + logicalType: object + properties: + - name: product + logicalType: object + properties: + - name: tags + logicalType: array + items: + logicalType: object + properties: + - name: tag_id + logicalType: string +""" + odcs = DataContract(data_contract_str=contract).get_data_contract() + + queries = build_databricks_virtual_model_queries_for_contract(odcs) + + assert "orders__customer__emails" in queries + assert "explode_outer(`customer`.`emails`)" in queries["orders__customer__emails"] + + assert "orders__line_items__product__tags" in queries + assert "explode_outer(`product`.`tags`)" in queries["orders__line_items__product__tags"] + + +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_test_databricks.py b/tests/test_test_databricks.py index 7e0f2237a..678ff5eb8 100644 --- a/tests/test_test_databricks.py +++ b/tests/test_test_databricks.py @@ -132,7 +132,7 @@ def test_unconvertible_column_does_not_affect_the_other_columns(databricks_type_ assert schema["mystery"] == dt.unknown -def test_nested_struct_sql_quality_is_enabled_for_databricks_only(): +def test_nested_struct_and_array_checks_enabled_for_databricks(): contract = """ apiVersion: v3.0.2 kind: DataContract @@ -151,6 +151,14 @@ def test_nested_struct_sql_quality_is_enabled_for_databricks_only(): - type: sql query: SELECT COUNT(*) FROM {model} WHERE {field} IS NULL mustBe: 0 + - name: emails + logicalType: array + items: + logicalType: object + properties: + - name: address + logicalType: string + required: true - name: discounts logicalType: array items: @@ -159,6 +167,17 @@ def test_nested_struct_sql_quality_is_enabled_for_databricks_only(): - name: discount_code logicalType: string required: true + - name: product + logicalType: object + properties: + - name: tags + logicalType: array + items: + logicalType: object + properties: + - name: tag_id + logicalType: string + required: true """ odcs = DataContract(data_contract_str=contract).get_data_contract() @@ -169,7 +188,10 @@ def test_nested_struct_sql_quality_is_enabled_for_databricks_only(): assert nested_sql.metric == MetricType.CUSTOM_SQL assert nested_sql.model == "orders" assert "customer.email" in (nested_sql.query or "") - assert not any(c.model == "orders__discounts" for c in checks) + # Array models must also be generated for Databricks (via virtual CTE models). + assert any(c.model == "orders__customer__emails" and c.field == "address" for c in checks) + assert any(c.model == "orders__discounts" for c in checks) + assert any(c.model == "orders__discounts__product__tags" and c.field == "tag_id" for c in checks) @pytest.mark.skipif( From 669acc70cf28c348b5e36cd072574594c2f3411f Mon Sep 17 00:00:00 2001 From: Rob Williamson Date: Tue, 16 Jun 2026 18:13:33 +0200 Subject: [PATCH 4/6] Update Formatting Per the PR template. --- datacontract/engines/checks/create_checks.py | 15 ++++++++++----- datacontract/engines/ibis/connections/connect.py | 1 + tests/test_connect_databricks.py | 15 +++++++++------ tests/test_create_checks_nested.py | 4 +--- 4 files changed, 21 insertions(+), 14 deletions(-) diff --git a/datacontract/engines/checks/create_checks.py b/datacontract/engines/checks/create_checks.py index 860c423f7..bcdeaec6b 100644 --- a/datacontract/engines/checks/create_checks.py +++ b/datacontract/engines/checks/create_checks.py @@ -254,7 +254,9 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> 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") + uses_raw_view = ( + server is not None and server_type in _FILE_SERVER_TYPES and server.format in ("csv", "parquet", "json") + ) # A primary key is both not-null and unique. A composite key is unique as a # tuple, not column by column, so its members are checked together after @@ -674,7 +676,12 @@ def _quality_checks( def _quality_rule_checks( - model: str, field: Optional[str], quality: DataQuality, count: int, server: Optional[Server], is_nested: bool = False + model: str, + field: Optional[str], + quality: DataQuality, + count: int, + server: Optional[Server], + is_nested: bool = False, ) -> List[CheckSpec]: """The checks of a single ODCS quality rule (``count`` is its index in the list).""" if quality.type == "custom" and quality.engine == "soda" and quality.implementation: @@ -714,9 +721,7 @@ def _quality_rule_checks( field=field, metric=MetricType.UNSUPPORTED, preset_result="warning", - preset_reason=( - "Nested SQL quality checks are only verified for Spark (dataframe) and Databricks." - ), + preset_reason=("Nested SQL quality checks are only verified for Spark (dataframe) and Databricks."), ) ] if field is None: diff --git a/datacontract/engines/ibis/connections/connect.py b/datacontract/engines/ibis/connections/connect.py index 68776902a..87eeb785f 100644 --- a/datacontract/engines/ibis/connections/connect.py +++ b/datacontract/engines/ibis/connections/connect.py @@ -230,6 +230,7 @@ def _connect_databricks(ibis, server: Server, run: Run, config: Config): exchange happens when the connection is opened rather than while reading env. """ # the config option wins over the contract, like the other server-detail overrides + host = ( config.get_databricks_server_hostname() or server.host or config.get_databricks_server_hostname(required=True) ) diff --git a/tests/test_connect_databricks.py b/tests/test_connect_databricks.py index 736883c76..e2cc8422e 100644 --- a/tests/test_connect_databricks.py +++ b/tests/test_connect_databricks.py @@ -4,7 +4,6 @@ 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 @@ -40,10 +39,14 @@ def captured_connect(monkeypatch): def fake_connect(self, **kwargs): calls.update(kwargs) # Return a minimal mock backend with the required attributes for downstream code. - mock = type('MockBackend', (), { - 'name': 'databricks', - '_dc_virtual_model_queries': {}, - })() + mock = type( + "MockBackend", + (), + { + "name": "databricks", + "_dc_virtual_model_queries": {}, + }, + )() return mock monkeypatch.setattr(DatabricksBackend, "connect", fake_connect) @@ -229,7 +232,7 @@ def spy_post_connect(self, *, memtable_volume): # 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') + self._memtable_volume = kwargs.get("memtable_volume") return self monkeypatch.setattr(DatabricksBackend, "connect", fake_connect) diff --git a/tests/test_create_checks_nested.py b/tests/test_create_checks_nested.py index 26c9e8940..d8b30300e 100644 --- a/tests/test_create_checks_nested.py +++ b/tests/test_create_checks_nested.py @@ -93,8 +93,6 @@ def test_create_checks_uses_full_nested_array_paths_for_nested_models(): c.model == "orders__user__emails" and c.field == "address" and c.type == "field_required" for c in checks ) assert any( - c.model == "orders__line_items__product__tags" - and c.field == "tag_id" - and c.type == "field_required" + c.model == "orders__line_items__product__tags" and c.field == "tag_id" and c.type == "field_required" for c in checks ) From 9c0755fae98bb994316c1ef20c0fc41bebe23bd1 Mon Sep 17 00:00:00 2001 From: Rob Williamson Date: Wed, 2 Sep 2026 09:05:40 +0200 Subject: [PATCH 5/6] Fix stale _NoVolumeBackend references in Databricks connect comments The comments in _connect_databricks and test_connect_databricks referred to a "_NoVolumeBackend subclass" that overrides _post_connect via MRO, but the implementation temporarily monkey-patches Backend._post_connect with a no-op lambda inside _databricks_connect and restores it in a finally block. No subclass or MRO override is involved. Update the docstring and test comments to describe what the code actually does. Co-authored-by: Copilot using Claude Opus 4.7 --- datacontract/engines/ibis/connections/connect.py | 5 +++-- tests/test_connect_databricks.py | 13 ++++++++----- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/datacontract/engines/ibis/connections/connect.py b/datacontract/engines/ibis/connections/connect.py index 87eeb785f..4938d4ada 100644 --- a/datacontract/engines/ibis/connections/connect.py +++ b/datacontract/engines/ibis/connections/connect.py @@ -214,8 +214,9 @@ def connect_ibis( def _connect_databricks(ibis, server: Server, run: Run, config: Config): """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. + Delegates to ``_databricks_connect``, which temporarily replaces ibis' + ``Backend._post_connect`` with a no-op to skip its 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 diff --git a/tests/test_connect_databricks.py b/tests/test_connect_databricks.py index e2cc8422e..c3d2f5b00 100644 --- a/tests/test_connect_databricks.py +++ b/tests/test_connect_databricks.py @@ -1,6 +1,6 @@ """Unit tests for Databricks auth-method selection in connect_ibis. -These do not hit Databricks: we patch the _NoVolumeBackend.connect method and only +These do not hit Databricks: we patch ``DatabricksBackend.connect`` and only assert which auth kwargs the dispatch passes for a given set of env vars. """ @@ -226,7 +226,9 @@ def spy_post_connect(self, *, memtable_volume): original_post_connect(self, memtable_volume=memtable_volume) try: - # Patch to spy on all calls (including _NoVolumeBackend overrides via MRO). + # Install a spy on the class attribute. _databricks_connect will save + # this spy as its "original", replace it with a no-op lambda for the + # duration of the connect call, and then restore it. DatabricksBackend._post_connect = spy_post_connect # Also patch connect to return early so we don't hit actual Databricks. @@ -239,9 +241,10 @@ def fake_connect(self, **kwargs): _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). + # _databricks_connect swaps in a no-op lambda before calling + # ibis.databricks.connect and restores the spy afterwards, so any + # _post_connect call issued during connect resolves to the no-op and + # the spy is never invoked. assert post_connect_called == [] finally: DatabricksBackend._post_connect = original_post_connect From 9d4c8447aaabe06b2c2c60af1ae7adc9fa89a6b5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jakob=20Sch=C3=B6dl?= Date: Thu, 10 Sep 2026 13:56:12 +0200 Subject: [PATCH 6/6] Run nested checks as array predicates instead of exploded relations - Replace exploded temp views and CTE virtual models with array predicates - Collapse item_model dead variable back to model - Remove unreachable array_path.rstrip('.') - Fix case-folding on value paths (_struct_path, _resolve_nested_expr) - Route nested unique checks through the same predicate as samples collection - Handle unsupported cases gracefully: array-item quality rules warn, nested physicalType checks skip - Restore test_connect_databricks.py to main's version (unrelated rewrite) - Add tests for case-folding across array hops and nested unique sampling --- datacontract/engines/checks/create_checks.py | 98 ++++----- .../engines/ibis/connections/connect.py | 42 +--- .../connections/databricks_nested_models.py | 107 ---------- .../engines/ibis/connections/kafka.py | 60 ------ .../engines/ibis/ibis_check_execute.py | 189 +++++++++++++----- tests/test_connect_databricks.py | 63 +----- .../test_connect_databricks_virtual_models.py | 128 ------------ tests/test_create_checks_nested.py | 99 ++++++++- tests/test_create_checks_nested_type.py | 31 ++- tests/test_ibis_check_execute.py | 81 ++++++++ tests/test_nested_path_checks.py | 102 ++++++++++ tests/test_snowflake_structured_types.py | 4 +- tests/test_test_databricks.py | 65 ------ 13 files changed, 509 insertions(+), 560 deletions(-) delete mode 100644 datacontract/engines/ibis/connections/databricks_nested_models.py delete mode 100644 tests/test_connect_databricks_virtual_models.py create mode 100644 tests/test_nested_path_checks.py diff --git a/datacontract/engines/checks/create_checks.py b/datacontract/engines/checks/create_checks.py index bcdeaec6b..721d19951 100644 --- a/datacontract/engines/checks/create_checks.py +++ b/datacontract/engines/checks/create_checks.py @@ -32,9 +32,7 @@ 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"} +_NESTED_CHECK_SERVER_TYPES = {"dataframe", "databricks"} # --------------------------------------------------------------------------- @@ -119,32 +117,24 @@ def _property_type(prop: SchemaProperty) -> str: 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 + yield field_path, prop 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) + if server_type in _NESTED_CHECK_SERVER_TYPES and prop_type == "object" and prop.properties: + yield from _iter_property_paths(prop.properties, server_type, field_path) elif ( - server_type in _SUPPORTED_NESTED_ARRAY_SERVER_TYPES - and prop_type == "array" - and prop.items - and prop.items.properties + server_type in _NESTED_CHECK_SERVER_TYPES and prop_type == "array" and prop.items and prop.items.properties ): - nested_model = f"{model}__{field_path.replace('.', '__')}" - yield from _iter_property_paths(nested_model, prop.items.properties, server_type, None, True) + # `[]` marks the array hop; the executor turns it into a predicate + # over the elements instead of a column lookup. + yield from _iter_property_paths(prop.items.properties, server_type, f"{field_path}[]") _PERCENT_UNITS = {"percent", "percentage", "%"} @@ -267,16 +257,16 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> ) primary_key_is_composite = len(primary_key_props) > 1 - for item_model, field, prop, is_nested in _iter_property_paths(model, properties, server_type): + for field, prop in _iter_property_paths(properties, server_type): # ODCS physicalName is the real column; mirror to_schema_name at field level. checks.append( CheckSpec( - key=f"{item_model}__{field}__field_is_present", + key=f"{model}__{field}__field_is_present", category="schema", type="field_is_present", name=f"Check that field '{field}' is present", - model=item_model, + model=model, field=field, metric=MetricType.FIELD_PRESENT, uses_raw_view=uses_raw_view, @@ -325,11 +315,11 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> label = prop.logicalType or "" checks.append( CheckSpec( - key=f"{item_model}__{field}__field_type", + key=f"{model}__{field}__field_type", category="schema", type="field_type", name=f"Check that field {field} has type {label}", - model=item_model, + model=model, field=field, metric=MetricType.FIELD_TYPE, expected_category=label, @@ -343,7 +333,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> if prop.required: checks.append( _missing_count_check( - item_model, + model, field, "field_required", Threshold(Op.EQ, 0), @@ -354,7 +344,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> if prop.unique: checks.append( _duplicate_count_check( - item_model, + model, field, "field_unique", Threshold(Op.EQ, 0), @@ -392,7 +382,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> if min_length is not None: checks.append( _invalid_count_check( - item_model, + model, field, "field_min_length", name=f"Check that field {field} has a min length of {min_length}", @@ -404,7 +394,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> if max_length is not None: checks.append( _invalid_count_check( - item_model, + model, field, "field_max_length", name=f"Check that field {field} has a max length of {max_length}", @@ -416,7 +406,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> if minimum is not None: checks.append( _invalid_count_check( - item_model, + model, field, "field_minimum", name=f"Check that field {field} has a minimum of {minimum}", @@ -428,7 +418,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> if maximum is not None: checks.append( _invalid_count_check( - item_model, + model, field, "field_maximum", name=f"Check that field {field} has a maximum of {maximum}", @@ -440,7 +430,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> if exclusive_minimum is not None: checks.append( _invalid_count_check( - item_model, + model, field, "field_minimum", name=f"Check that field {field} has a minimum of {exclusive_minimum}", @@ -449,7 +439,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> ) checks.append( _invalid_count_check( - item_model, + model, field, "field_not_equal", name=f"Check that field {field} is not equal to {exclusive_minimum}", @@ -461,7 +451,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> if exclusive_maximum is not None: checks.append( _invalid_count_check( - item_model, + model, field, "field_maximum", name=f"Check that field {field} has a maximum of {exclusive_maximum}", @@ -470,7 +460,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> ) checks.append( _invalid_count_check( - item_model, + model, field, "field_not_equal", name=f"Check that field {field} is not equal to {exclusive_maximum}", @@ -519,7 +509,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> if pattern is not None: checks.append( _invalid_count_check( - item_model, + model, field, "field_regex", name=f"Check that field {field} matches regex pattern {pattern}", @@ -531,7 +521,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> if enum_values: checks.append( _invalid_count_check( - item_model, + model, field, "field_enum", name=f"Check that field {field} only contains enum values {enum_values}", @@ -540,7 +530,7 @@ def _to_schema_checks(schema_object: SchemaObject, server: Optional[Server]) -> ) if prop.quality: - checks.extend(_quality_checks(item_model, field, prop.quality, server, is_nested=is_nested)) + checks.extend(_quality_checks(model, field, prop.quality, server)) if primary_key_is_composite: primary_key_fields = [prop.physicalName or prop.name for prop in primary_key_props] @@ -660,11 +650,11 @@ def _row_count_check(model, threshold: Threshold, severity=None, dimension=None) # quality list # --------------------------------------------------------------------------- def _quality_checks( - model: str, field: Optional[str], quality_list: List[DataQuality], server: Optional[Server], is_nested: bool = False + model: str, field: Optional[str], quality_list: List[DataQuality], server: Optional[Server] ) -> List[CheckSpec]: checks: List[CheckSpec] = [] for count, quality in enumerate(quality_list): - rule_checks = _quality_rule_checks(model, field, quality, count, server, is_nested=is_nested) + rule_checks = _quality_rule_checks(model, field, quality, count, server) # Every check keeps a link back to the rule that declared it, so that # `test --quality-id` / `test --tag` can select it. for check in rule_checks: @@ -681,7 +671,6 @@ def _quality_rule_checks( quality: DataQuality, count: int, server: Optional[Server], - is_nested: bool = False, ) -> List[CheckSpec]: """The checks of a single ODCS quality rule (``count`` is its index in the list).""" if quality.type == "custom" and quality.engine == "soda" and quality.implementation: @@ -703,14 +692,15 @@ def _quality_rule_checks( ) ] if 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" + 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" + if field is not None and "[]" in field: + # An array item is not a column, so substituting it into the query + # would produce SQL no backend can parse. return [ CheckSpec( key=check_key, @@ -720,16 +710,16 @@ def _quality_rule_checks( model=model, field=field, metric=MetricType.UNSUPPORTED, + dimension=quality.dimension, + severity=quality.severity, preset_result="warning", - preset_reason=("Nested SQL quality checks are only verified for Spark (dataframe) and Databricks."), + preset_reason=( + f"'{field}' is an array item, not a column, so it cannot be substituted into a query. " + f"Declare the rule on '{field.split('[]')[0]}' instead and match the elements with array " + f"functions, for example size(filter(...)) > 0." + ), ) ] - 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" threshold = to_threshold(quality) query = prepare_query(quality, model, field, server) if query is None: diff --git a/datacontract/engines/ibis/connections/connect.py b/datacontract/engines/ibis/connections/connect.py index 4938d4ada..96af04c39 100644 --- a/datacontract/engines/ibis/connections/connect.py +++ b/datacontract/engines/ibis/connections/connect.py @@ -97,9 +97,6 @@ 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": @@ -108,26 +105,8 @@ 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) - backend = _connect_databricks(ibis, server, run, config) - # 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 + return _connect_databricks(ibis, server, run, config) if server_type == "postgres": return ibis.postgres.connect( @@ -214,24 +193,23 @@ def connect_ibis( def _connect_databricks(ibis, server: Server, run: Run, config: Config): """Connect to Databricks SQL directly, selecting the auth method from env vars. - Delegates to ``_databricks_connect``, which temporarily replaces ibis' - ``Backend._post_connect`` with a no-op to skip its 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. """ # the config option wins over the contract, like the other server-detail overrides - host = ( config.get_databricks_server_hostname() or server.host or config.get_databricks_server_hostname(required=True) ) diff --git a/datacontract/engines/ibis/connections/databricks_nested_models.py b/datacontract/engines/ibis/connections/databricks_nested_models.py deleted file mode 100644 index 90299b44f..000000000 --- a/datacontract/engines/ibis/connections/databricks_nested_models.py +++ /dev/null @@ -1,107 +0,0 @@ -"""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 _nested_model_name(parent_model: str, field_path: str) -> str: - return f"{parent_model}__{field_path.replace('.', '__')}" - - -def _nested_field_expr(field_path: str) -> str: - return ".".join(f"`{part}`" for part in field_path.split(".")) - - -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 = _nested_model_name(parent_model, field_path) - specs[virtual_model] = { - "parent": parent_model, - "field": _nested_field_expr(field_path), - "template": "SELECT __dc_nested__.* FROM {source} LATERAL VIEW OUTER explode_outer({field}) AS __dc_nested__", - } - _collect_databricks_virtual_model_specs(virtual_model, prop.items.properties, specs) - - # 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 e43ee776d..970dacba9 100644 --- a/datacontract/engines/ibis/connections/kafka.py +++ b/datacontract/engines/ibis/connections/kafka.py @@ -476,66 +476,6 @@ 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 _decimal_params(prop: SchemaProperty) -> Tuple[int, int]: options = prop.logicalTypeOptions or {} precision = options.get("precision") diff --git a/datacontract/engines/ibis/ibis_check_execute.py b/datacontract/engines/ibis/ibis_check_execute.py index 77725a7fc..7d3e6da16 100644 --- a/datacontract/engines/ibis/ibis_check_execute.py +++ b/datacontract/engines/ibis/ibis_check_execute.py @@ -279,10 +279,9 @@ def model_row_count() -> int: if spec.metric == MetricType.ROW_COUNT: named = t.count().name(spec.key) elif spec.metric == MetricType.MISSING_COUNT: - col = _resolve_expr(t, columns, spec.field) - named = _count_true(_missing_expr(col, spec.missing_values)).name(spec.key) + predicate = _row_predicate(t, columns, spec.field, lambda c: _missing_expr(c, spec.missing_values)) + named = _count_true(predicate).name(spec.key) elif spec.metric == MetricType.INVALID_COUNT: - col = _resolve_expr(t, columns, spec.field) dtype = _resolve_dtype(schema, spec.field) if _has_array_constraints(spec) and not dtype.is_array(): # Silently dropping the constraint would report the check as @@ -295,8 +294,19 @@ def model_row_count() -> int: f"Column {spec.field} is {dtype}, not an array, so the constraint cannot be measured.", ) continue - expr = _invalid_expr(t, col, dtype, spec) - if expr is None: + unconstrained = [] + + def _invalid(c, _t=t, _dtype=dtype, _spec=spec, _flag=unconstrained): + import ibis + + built = _invalid_expr(_t, c, _dtype, _spec) + if built is None: + _flag.append(True) + return ibis.literal(False) + return built + + expr = _row_predicate(t, columns, spec.field, _invalid) + if unconstrained: # No validity constraints => nothing can be invalid. _set_impl(run, spec.key, "invalid_count = 0 (no validity constraints configured)", None) _evaluate(run, spec, 0, row_count=model_row_count()) @@ -309,9 +319,9 @@ def model_row_count() -> int: elif spec.metric == MetricType.FIELD_TYPE: _run_type(run, schema, columns, spec, structured_types) elif spec.metric == MetricType.FIELD_PHYSICAL_TYPE: - _run_physical_type(run, con, server, schema, columns, native_types, spec, structured_types) + _run_physical_type(run, con, server, schema, native_types, spec, structured_types) elif spec.metric == MetricType.FIELD_NESTED_TYPE: - _run_nested_type(run, schema, columns, spec, structured_types, sqlglot_dialect(con)) + _run_nested_type(run, schema, spec, structured_types, sqlglot_dialect(con)) elif spec.metric in (MetricType.FRESHNESS, MetricType.RETENTION): _run_freshness(run, t, columns, spec) elif spec.metric == MetricType.CUSTOM_SQL: @@ -445,17 +455,31 @@ def _select_columns(columns, sensitive, wanted): 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_expr(t, columns, spec.field) - if spec.metric == MetricType.MISSING_COUNT: - predicate = _missing_expr(col, spec.missing_values) + if not _is_item_duplicate(spec): + return _duplicate_samples(t, columns, sensitive, spec) + predicate = _item_duplicate_predicate(t, columns, spec.field) + elif spec.metric == MetricType.MISSING_COUNT: + predicate = _row_predicate(t, columns, spec.field, lambda c: _missing_expr(c, spec.missing_values)) else: # INVALID_COUNT - predicate = _invalid_expr(t, col, _resolve_dtype(schema, spec.field), spec) - if predicate is None: + dtype = _resolve_dtype(schema, spec.field) + unconstrained = [] + + def _invalid(c, _flag=unconstrained): + import ibis + + built = _invalid_expr(t, c, dtype, spec) + if built is None: + _flag.append(True) + return ibis.literal(False) + return built + + predicate = _row_predicate(t, columns, spec.field, _invalid) + if unconstrained: return None - select_cols = _select_columns(columns, sensitive, [*identifiers, spec.field]) + # A nested path is not a column; show the one it starts from. + root = spec.field.split("[]")[0].split(".")[0] if spec.field else None + select_cols = _select_columns(columns, sensitive, [*identifiers, root]) rows = t.filter(predicate) rows = rows.select(select_cols) if select_cols else rows return _df_to_records(rows.limit(_FAILED_SAMPLE_LIMIT).execute()) @@ -699,6 +723,10 @@ def _run_duplicate(run: Run, t, columns, spec: CheckSpec, row_count: int): keys span are what is reported as failed.""" import pandas as pd + if _is_item_duplicate(spec): + _run_item_duplicate(run, t, columns, spec, row_count) + return + 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) @@ -719,6 +747,33 @@ def _int(value) -> int: _update_diagnostics(run, spec.key, extra) +def _is_item_duplicate(spec: CheckSpec) -> bool: + return any("[]" in c for c in (spec.columns or [spec.field]) if c) + + +def _item_duplicate_predicate(t, columns, field: str): + """Parent rows whose array repeats a value in the item property ``field``.""" + array_path, _, leaf = field.rpartition("[].") + + def _repeats(array): + values = array.map(lambda element: _struct_path(element, leaf)) + return values.unique().length() < values.length() + + return _row_predicate(t, columns, array_path, _repeats) + + +def _run_item_duplicate(run: Run, t, columns, spec: CheckSpec, row_count: int): + """``unique`` on an array item property: no parent may repeat a value inside its own array. + + Counted in parent rows. + """ + bad = t.filter(_item_duplicate_predicate(t, columns, spec.field)) + _record_sql(run, spec, bad) + repeats = int(bad.count().execute()) + _evaluate(run, spec, repeats, row_count=row_count) + _update_diagnostics(run, spec.key, {"failed_rows": repeats}) + + 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") @@ -786,7 +841,6 @@ def _run_physical_type( con, server, schema, - columns, native_types, spec: CheckSpec, structured_types: dict[str, SchemaProperty] | None = None, @@ -804,8 +858,8 @@ def _run_physical_type( f"physical type of '{spec.field}' is '{spec.expected_physical_type}'", "introspection", ) - 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_physical_type", field=spec.field, expected=spec.expected_physical_type) ) @@ -843,16 +897,28 @@ def _run_physical_type( set_result(run, spec.key, ResultEnum.failed, reason) return + # Native types are read per top-level column, so a nested path never has one. + # The logicalType fallback would report `passed` for a physical type nothing + # compared, so skip instead. + if actual_native is None and ("." in spec.field or "[]" in spec.field): + set_result( + run, + spec.key, + ResultEnum.warning, + f"The native type of '{spec.field}' is not read for nested paths; skipping the physical type check", + ) + return + # result is None: the physical type could not be evaluated. Fall back to the # logicalType category check when the property declares one. fallback = spec.expected_schema_property if fallback is not None and fallback.logicalType is not None: - actual_prop = structured_prop or ibis_dtype_to_schema_property(schema[actual_col]) + actual_prop = structured_prop or ibis_dtype_to_schema_property(dtype) if schema_property_matches(fallback, actual_prop): set_result(run, spec.key, ResultEnum.passed, None) else: mismatch = schema_property_mismatch_reason(fallback, actual_prop) - actual_label = actual_native or schema[actual_col] + actual_label = actual_native or dtype set_result( run, spec.key, @@ -867,7 +933,6 @@ def _run_physical_type( def _run_nested_type( run: Run, schema, - columns, spec: CheckSpec, structured_types: dict[str, SchemaProperty] | None = None, dialect=None, @@ -877,13 +942,11 @@ def _run_nested_type( """ metric = spec.type _set_impl(run, spec.key, f"nested types of '{spec.field}' match the contract", "introspection") - 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=metric, field=spec.field, expected=spec.expected_type_label)) set_result(run, spec.key, ResultEnum.failed, f"Column '{spec.field}' is missing") return - - dtype = schema[actual_col] structured_prop = structured_types.get(spec.field.lower()) if structured_types else None actual_prop = structured_prop or ibis_dtype_to_schema_property(dtype) actual_label = (structured_prop.physicalType if structured_prop else None) or str(dtype) @@ -1162,10 +1225,44 @@ def _resolve_expr(t, columns: dict, field: str): 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 + head, _, tail = field.partition(".") + return _struct_path(t[_resolve_col(columns, head)], tail) + + +def _row_predicate(t, columns: dict, field: str, build): + """A boolean over *parent* rows for ``build`` applied at ``field``. + + An array hop (``items[].sku``) becomes "some element satisfies build", so the + row count never changes and an empty array is never a violation. + """ + head, marker, tail = field.partition("[].") + if not marker: + return build(_resolve_expr(t, columns, field)) + return _resolve_expr(t, columns, head).filter(lambda i: _element_predicate(i, tail, build)).length() > 0 + + +def _element_predicate(element, path: str, build): + head, marker, tail = path.partition("[].") + if not marker: + return build(_struct_path(element, path)) + return _struct_path(element, head).filter(lambda i: _element_predicate(i, tail, build)).length() > 0 + + +def _struct_path(value, path: str): + for part in path.split("."): + value = value[_struct_field_name(value, part)] + return value + + +def _struct_field_name(value, name: str) -> str: + """The struct's own spelling of ``name``, which a backend may report in another case.""" + try: + names = value.type().names + except Exception: + return name + if name in names: + return name + return next((n for n in names if n.lower() == name.lower()), name) def _resolve_dtype(schema, field: str): @@ -1175,10 +1272,22 @@ def _resolve_dtype(schema, field: str): parts = field.split(".") dtype = None for idx, part in enumerate(parts): + name, marker, _ = part.partition("[]") try: - dtype = current[part] + dtype = current[name] except Exception: - return None + # Backends that report uppercase names (Snowflake, Oracle, Databricks) + # must still match a contract that spells the field in lower case. + actual = next((k for k in current.keys() if k.lower() == name.lower()), None) + if actual is None: + return None + dtype = current[actual] + if marker: + # `items[]` names the element type, not the array's own type. + try: + dtype = dtype.value_type + except Exception: + return None if idx < len(parts) - 1: try: current = dtype.fields @@ -1247,29 +1356,13 @@ def _apply_row_filter(t, model: str, predicate: str): def _resolve_table(con, model: str, database: Optional[str] = None): - """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. - """ + """Resolve a table by name, tolerating case differences.""" if getattr(con, "name", None) == "pyspark": return _pyspark_table_unconvertible_as_unknown(con, model) kwargs = {"database": database} if database else {} try: return con.table(model, **kwargs) 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(**kwargs) except Exception: diff --git a/tests/test_connect_databricks.py b/tests/test_connect_databricks.py index c3d2f5b00..b116c5846 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: we patch ``DatabricksBackend.connect`` and only +These do not hit Databricks: ``ibis.databricks.connect`` is patched and we 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 @@ -33,23 +33,14 @@ def clean_databricks_env(monkeypatch): @pytest.fixture def captured_connect(monkeypatch): - """Patch DatabricksBackend.connect to record the kwargs it is called with.""" + """Patch ibis.databricks.connect to record the kwargs it is called with.""" calls = {} - def fake_connect(self, **kwargs): + def fake_connect(**kwargs): calls.update(kwargs) - # 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 "connection" + + monkeypatch.setattr(ibis.databricks, "connect", fake_connect) return calls @@ -90,7 +81,7 @@ def test_personal_access_token_is_default(clean_databricks_env, captured_connect result = _connect() - assert result is not None + assert result == "connection" 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" @@ -210,41 +201,3 @@ def test_env_variables_override_the_contract_server_details(clean_databricks_env assert captured_connect["server_hostname"] == "from-env.cloud.databricks.com" assert captured_connect["catalog"] == "env_catalog" assert captured_connect["schema"] == "env_schema" - - -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: - # Install a spy on the class attribute. _databricks_connect will save - # this spy as its "original", replace it with a no-op lambda for the - # duration of the connect call, and then restore it. - 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() - - # _databricks_connect swaps in a no-op lambda before calling - # ibis.databricks.connect and restores the spy afterwards, so any - # _post_connect call issued during connect resolves to the no-op and - # the spy is never invoked. - 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 deleted file mode 100644 index 32c9acaee..000000000 --- a/tests/test_connect_databricks_virtual_models.py +++ /dev/null @@ -1,128 +0,0 @@ -"""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_builds_virtual_queries_for_nested_array_paths(): - 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: emails - logicalType: array - items: - logicalType: object - properties: - - name: address - logicalType: string - - name: line_items - logicalType: array - items: - logicalType: object - properties: - - name: product - logicalType: object - properties: - - name: tags - logicalType: array - items: - logicalType: object - properties: - - name: tag_id - logicalType: string -""" - odcs = DataContract(data_contract_str=contract).get_data_contract() - - queries = build_databricks_virtual_model_queries_for_contract(odcs) - - assert "orders__customer__emails" in queries - assert "explode_outer(`customer`.`emails`)" in queries["orders__customer__emails"] - - assert "orders__line_items__product__tags" in queries - assert "explode_outer(`product`.`tags`)" in queries["orders__line_items__product__tags"] - - -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 index d8b30300e..eeba26412 100644 --- a/tests/test_create_checks_nested.py +++ b/tests/test_create_checks_nested.py @@ -70,7 +70,7 @@ def test_create_checks_recurses_for_dataframe_nested_structs_and_arrays(): 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) + assert any(c.field == "line_items[].sku" and c.type == "field_required" and c.model == "orders" for c in checks) nested_sql = next(c for c in checks if c.type == "field_quality_sql") assert nested_sql.field == "user.status" @@ -83,16 +83,105 @@ 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) + assert not any(c.field == "line_items[].sku" for c in checks) -def test_create_checks_uses_full_nested_array_paths_for_nested_models(): +def test_create_checks_marks_array_hops_in_the_field_path(): checks = _checks("dataframe") assert any( - c.model == "orders__user__emails" and c.field == "address" and c.type == "field_required" for c in checks + c.model == "orders" and c.field == "user.emails[].address" and c.type == "field_required" for c in checks ) assert any( - c.model == "orders__line_items__product__tags" and c.field == "tag_id" and c.type == "field_required" + c.model == "orders" and c.field == "line_items[].product.tags[].tag_id" and c.type == "field_required" for c in checks ) + + +PHYSICAL_TYPE_CONTRACT = """ +apiVersion: v3.0.2 +kind: DataContract +id: nested-physical-types +version: 1.0.0 +status: active +schema: + - name: orders + properties: + - name: customer + logicalType: object + physicalType: STRUCT + properties: + - name: name + logicalType: string + physicalType: STRING + - name: line_items + logicalType: array + physicalType: ARRAY + items: + logicalType: object + physicalType: STRUCT + properties: + - name: sku + logicalType: string + physicalType: STRING +""" + + +def test_create_checks_keeps_array_item_checks_on_the_real_model(): + odcs = DataContract(data_contract_str=PHYSICAL_TYPE_CONTRACT).get_data_contract() + checks = create_checks(odcs, Server(type="databricks")) + + sku = next(c for c in checks if c.type == "field_physical_type" and c.field == "line_items[].sku") + assert sku.model == "orders" + + # A struct field keeps the parent model and is addressed by its dotted path. + name = next(c for c in checks if c.type == "field_physical_type" and c.field == "customer.name") + assert name.model == "orders" + + +ARRAY_QUALITY_CONTRACT = """ +apiVersion: v3.0.2 +kind: DataContract +id: array-quality +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: items + logicalType: array + items: + logicalType: object + properties: + - name: sku + logicalType: string + quality: + - type: sql + query: SELECT COUNT(*) FROM {model} WHERE {field} IS NULL + mustBe: 0 +""" + + +def test_a_quality_query_on_an_array_item_warns_instead_of_emitting_unparseable_sql(): + odcs = DataContract(data_contract_str=ARRAY_QUALITY_CONTRACT).get_data_contract() + checks = {c.field: c for c in create_checks(odcs, Server(type="databricks")) if c.type == "field_quality_sql"} + + # a struct path is a legal column reference, so it still becomes a real query + assert checks["customer.email"].metric != MetricType.UNSUPPORTED + assert "customer.email" in checks["customer.email"].query + + # an array item is not, so the rule is reported as unsupported rather than run + item = checks["items[].sku"] + assert item.metric == MetricType.UNSUPPORTED + assert item.preset_result == "warning" + assert "items[].sku" in item.preset_reason + assert "Declare the rule on 'items'" in item.preset_reason diff --git a/tests/test_create_checks_nested_type.py b/tests/test_create_checks_nested_type.py index 0e7db942f..afd6bd687 100644 --- a/tests/test_create_checks_nested_type.py +++ b/tests/test_create_checks_nested_type.py @@ -12,7 +12,7 @@ Server, ) -from datacontract.engines.checks.check_spec import MetricType +from datacontract.engines.checks.check_spec import CheckSpec, MetricType from datacontract.engines.checks.create_checks import create_checks from datacontract.engines.ibis.ibis_check_execute import _run_nested_type, build_check_stubs from datacontract.engines.ibis.snowflake_structured_types import _to_property @@ -188,7 +188,7 @@ def _run(prop: SchemaProperty, dtype: str, server_type: str = "local", fmt: str run.checks = build_check_stubs(specs) spec = _nested(specs) field = prop.physicalName or prop.name - _run_nested_type(run, ibis.schema({field: dtype}), {field.lower(): field}, spec) + _run_nested_type(run, ibis.schema({field: dtype}), spec) return next(c for c in run.checks if c.key == spec.key) @@ -295,7 +295,7 @@ def _run_snowflake(prop: SchemaProperty, data_type: dict = _SHOW_COLUMNS_SIC_COD spec = _nested(specs) field = prop.physicalName or prop.name structured_types = {field.lower(): _to_property(data_type)} - _run_nested_type(run, ibis.schema({field: dtype}), {field.lower(): field}, spec, structured_types, "snowflake") + _run_nested_type(run, ibis.schema({field: dtype}), spec, structured_types, "snowflake") return next(c for c in run.checks if c.key == spec.key) @@ -436,3 +436,28 @@ def test_the_mismatch_reason_does_not_repeat_the_columns_own_structure(): assert check.result == ResultEnum.failed assert check.reason == "Cannot verify the nested types of 'primary_sic_code': the column is not an object" assert check.diagnostics["actual"] == "array>" + + +def test_nested_type_resolves_a_dotted_struct_path(): + spec = CheckSpec( + key="k", + category="schema", + type="field_nested_type", + name="nested type", + model="orders", + field="customer.address", + metric=MetricType.FIELD_TYPE, + expected_type_label="object", + expected_schema_property=SchemaProperty( + name="address", + logicalType="object", + properties=[SchemaProperty(name="city", logicalType="string")], + ), + ) + run = Run.create_run() + run.checks = build_check_stubs([spec]) + schema = ibis.schema({"customer": "struct>"}) + + _run_nested_type(run, schema, spec) + + assert run.checks[0].result == ResultEnum.passed diff --git a/tests/test_ibis_check_execute.py b/tests/test_ibis_check_execute.py index 8fca369f3..e50e6c238 100644 --- a/tests/test_ibis_check_execute.py +++ b/tests/test_ibis_check_execute.py @@ -69,3 +69,84 @@ def test_run_present_raw_view_falls_back_to_model_with_case_insensitive_resoluti _run_present(run, con, "checks_testcase", {"ctc_id": "CTC_ID"}, {"IGNORED": "int64"}, spec) assert run.checks[0].result == ResultEnum.passed + + +def test_run_present_matches_uppercase_column_for_lowercase_contract_field(): + 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_matches_uppercase_nested_field_for_lowercase_contract_path(): + import ibis + + run = _run_with_stubbed_check() + spec = CheckSpec( + key="k", + category="schema", + type="field_is_present", + name="field is present", + model="checks_testcase", + field="customer.name", + metric=MetricType.FIELD_PRESENT, + ) + schema = ibis.memtable({"CUSTOMER": [{"NAME": "a"}]}).schema() + + _run_present(run, _NoLookupConnection(), "checks_testcase", {"customer": "CUSTOMER"}, schema, spec) + + assert run.checks[0].result == ResultEnum.passed + + +def _physical_type_spec(field: str) -> CheckSpec: + from open_data_contract_standard.model import SchemaProperty + + return CheckSpec( + key="k", + category="schema", + type="field_physical_type", + name="physical type", + model="orders", + field=field, + metric=MetricType.FIELD_PHYSICAL_TYPE, + expected_category="DECIMAL(10,2)", + expected_physical_type="DECIMAL(10,2)", + expected_type_label="DECIMAL(10,2)", + expected_schema_property=SchemaProperty(name=field, logicalType="number"), + ) + + +def test_physical_type_of_a_nested_path_is_skipped_rather_than_passed_on_the_fallback(): + # fetch_native_types only reads top-level columns, so a nested path has no + # native type to compare against. The logicalType fallback would report a + # mismatched physical type as passed. + import ibis + from open_data_contract_standard.model import Server + + from datacontract.engines.ibis.ibis_check_execute import _run_physical_type + + schema = ibis.schema({"price": "float64", "orders": "array>"}) + native_types = {"price": "DOUBLE", "orders": "ARRAY>"} + server = Server(server="s", type="databricks") + + run = Run.create_run() + run.checks = [Check(type="field_physical_type", key="k")] + _run_physical_type(run, None, server, schema, native_types, _physical_type_spec("orders[].price")) + assert run.checks[0].result == ResultEnum.warning + assert "skipping the physical type check" in run.checks[0].reason + + # the same mismatch on a top-level column is still a real comparison + run = Run.create_run() + run.checks = [Check(type="field_physical_type", key="k")] + _run_physical_type(run, None, server, schema, native_types, _physical_type_spec("price")) + assert run.checks[0].result == ResultEnum.failed diff --git a/tests/test_nested_path_checks.py b/tests/test_nested_path_checks.py new file mode 100644 index 000000000..e1a9de896 --- /dev/null +++ b/tests/test_nested_path_checks.py @@ -0,0 +1,102 @@ +"""Array hops in a check's field path, and the predicates they compile to.""" + +import ibis +import pytest + +from datacontract.engines.ibis.ibis_check_execute import _missing_expr, _resolve_dtype, _row_predicate + +TABLE = ibis.table( + { + "order_id": "string", + "customer": "struct>>", + "items": "array>>>", + }, + name="orders", +) +COLUMNS = {c.lower(): c for c in TABLE.columns} + + +@pytest.mark.parametrize( + "path,expected", + [ + ("order_id", "string"), + ("customer.name", "string"), + ("items[].sku", "string"), + ("customer.tags[].tag_id", "string"), + ("items[].parts[].part_no", "string"), + ], +) +def test_resolve_dtype_steps_through_array_elements(path, expected): + assert str(_resolve_dtype(TABLE.schema(), path)) == expected + + +def test_resolve_dtype_is_case_insensitive_across_an_array_hop(): + assert str(_resolve_dtype(TABLE.schema(), "ITEMS[].SKU")) == "string" + + +def _sql(path): + predicate = _row_predicate(TABLE, COLUMNS, path, lambda c: _missing_expr(c, None)) + return " ".join(str(ibis.to_sql(TABLE.filter(predicate).count(), dialect="databricks")).split()) + + +def test_a_plain_path_compiles_to_a_column_predicate(): + assert "FILTER(" not in _sql("customer.name") + + +def test_an_array_hop_compiles_to_a_predicate_over_the_elements(): + # A predicate over the array keeps one row per parent, so counts stay in + # parent rows and an empty array is never a violation. + sql = _sql("items[].sku") + assert "SIZE(FILTER(`t0`.`items`" in sql + assert "`sku`)ISNULL" in sql.replace(" ", "") + + +def test_nested_array_hops_compile_to_nested_predicates(): + assert _sql("items[].parts[].part_no").count("FILTER(") == 2 + + +UPPERCASE = ibis.table( + {"ORDER_ID": "string", "CUSTOMER": "struct", "ITEMS": "array>"}, + name="ORDERS", +) +UPPERCASE_COLUMNS = {c.lower(): c for c in UPPERCASE.columns} + + +@pytest.mark.parametrize("path", ["customer.name", "items[].sku"]) +def test_a_path_resolves_to_a_value_when_the_backend_reports_uppercase_names(path): + # _resolve_dtype already folds case; the value side has to agree, or presence + # and type checks pass while the count checks raise KeyError. + predicate = _row_predicate(UPPERCASE, UPPERCASE_COLUMNS, path, lambda c: _missing_expr(c, None)) + assert "IS NULL" in str(ibis.to_sql(UPPERCASE.filter(predicate).count(), dialect="databricks")) + + +def test_item_duplicate_samples_use_the_same_predicate_as_the_check(): + # `unique` on an array item is not a column lookup, so the sample path has to + # take the array branch too or it silently reports no samples. + import pandas as pd + + from datacontract.engines.checks.check_spec import CheckSpec, MetricType + from datacontract.engines.ibis.ibis_check_execute import _samples_for + + t = ibis.memtable( + pd.DataFrame( + { + "order_id": ["clean", "repeats"], + "items": [[{"sku": "A"}, {"sku": "B"}], [{"sku": "C"}, {"sku": "C"}]], + } + ) + ) + spec = CheckSpec( + key="orders__items[].sku__field_unique", + category="schema", + type="field_unique", + name="unique", + model="orders", + field="items[].sku", + metric=MetricType.DUPLICATE_COUNT, + columns=["items[].sku"], + ) + samples = _samples_for(t, {c.lower(): c for c in t.columns}, t.schema(), spec, ["order_id"], set()) + + assert samples is not None + assert [row["order_id"] for row in samples] == ["repeats"] diff --git a/tests/test_snowflake_structured_types.py b/tests/test_snowflake_structured_types.py index 17e21081b..782b4cf89 100644 --- a/tests/test_snowflake_structured_types.py +++ b/tests/test_snowflake_structured_types.py @@ -166,9 +166,7 @@ def _physical_type_check(expected: SchemaProperty, structured_types): ) run = Run.create_run() run.checks = [Check(id="k", key="k", category="schema", type=spec.type, name=spec.name, model="m", field="s_obj")] - _run_physical_type( - run, None, None, {"S_OBJ": "map"}, {"s_obj": "S_OBJ"}, None, spec, structured_types - ) + _run_physical_type(run, None, None, {"S_OBJ": "map"}, None, spec, structured_types) return run.checks[0] diff --git a/tests/test_test_databricks.py b/tests/test_test_databricks.py index 678ff5eb8..01a303105 100644 --- a/tests/test_test_databricks.py +++ b/tests/test_test_databricks.py @@ -2,11 +2,8 @@ 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) @@ -132,68 +129,6 @@ def test_unconvertible_column_does_not_affect_the_other_columns(databricks_type_ assert schema["mystery"] == dt.unknown -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: emails - logicalType: array - items: - logicalType: object - properties: - - name: address - logicalType: string - required: true - - name: discounts - logicalType: array - items: - logicalType: object - properties: - - name: discount_code - logicalType: string - required: true - - name: product - logicalType: object - properties: - - name: tags - logicalType: array - items: - logicalType: object - properties: - - name: tag_id - 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__customer__emails" and c.field == "address" for c in checks) - assert any(c.model == "orders__discounts" for c in checks) - assert any(c.model == "orders__discounts__product__tags" and c.field == "tag_id" for c in checks) - - @pytest.mark.skipif( os.environ.get("DATACONTRACT_DATABRICKS_TOKEN") is None, reason="Requires DATACONTRACT_DATABRICKS_TOKEN to be set" )