From a00de6abe594e8789cb3bd83505c00a9a0f45383 Mon Sep 17 00:00:00 2001 From: Kevin Costner <120246174+kevincostner17@users.noreply.github.com> Date: Tue, 15 Sep 2026 15:11:23 +0530 Subject: [PATCH] fix(dbt): honour on_low_score="fail" in FreshDataDbtTransform; unique audit files for same-alias models on_low_score="fail" was ignored by FreshDataDbtTransform (#343): - run() only raised TrustGateError when fail_on_low_score=True, so with on_low_score="fail" a failing gate returned should_fail=True and the pipeline carried on. run() now raises when result.should_fail or when fail_on_low_score is set and the gate did not pass. The audit file is still written before raising. - run() takes a keyword-only raise_on_fail=True. gate_manifest calls run(raise_on_fail=False), so a failing model under on_low_score="fail" is recorded as failed and the run continues. The summary shape and the skipped/all_passed semantics are unchanged. Same-alias models overwrote each other's audit file (#344): - FreshDataDbtTransform gains audit_name, used as the audit file stem (_audit.json) instead of the table name. It is checked with _validate_audit_table_name when the transform is configured. - When output_dir is set, gate_manifest counts aliases across the models it gates (case-insensitively, since audit files may land on a case-insensitive filesystem). Models whose alias is shared are written to ._audit.json, or to _audit.json when the schema is missing or that name is still not unique. Other models keep _audit.json. An unsafe schema is rejected by the validator and recorded as that model's error. Closes #343 Closes #344 --- docs/integrations.md | 10 +- src/freshdata/integrations/dbt/__init__.py | 68 +++- .../test_dbt_transform_fail_audit.py | 292 ++++++++++++++++++ 3 files changed, 357 insertions(+), 13 deletions(-) create mode 100644 tests/test_integrations/test_dbt_transform_fail_audit.py diff --git a/docs/integrations.md b/docs/integrations.md index 54fa401f..56c6de2c 100644 --- a/docs/integrations.md +++ b/docs/integrations.md @@ -111,10 +111,18 @@ result = FreshDataDbtTransform( model_name="analytics.orders", output_dir="target/freshdata", trust_score_threshold=80.0, - fail_on_low_score=True, + on_low_score="fail", # raise TrustGateError on a failing gate ).run() ``` +With `on_low_score="fail"` (or the older `fail_on_low_score=True`), `run()` writes +the audit file and then raises `TrustGateError` on a failing gate; pass +`run(raise_on_fail=False)` to get the failing result back instead. When +`dbt-gate` / `gate_manifest` gates several models that share an alias in different +schemas, their audit files are named `._audit.json` (or +`_audit.json` when the schema is missing), so no model's audit +overwrites another's. Other models keep `_audit.json`. + A bundled Jinja macro, `freshdata_trust_gate`, documents the recommended `on-run-end` invocation; see `freshdata/integrations/dbt/macros/freshdata_trust_gate.sql`. diff --git a/src/freshdata/integrations/dbt/__init__.py b/src/freshdata/integrations/dbt/__init__.py index 02700ca4..f1ae931b 100644 --- a/src/freshdata/integrations/dbt/__init__.py +++ b/src/freshdata/integrations/dbt/__init__.py @@ -19,6 +19,7 @@ import json import logging import os +from collections import Counter from dataclasses import dataclass from pathlib import Path, PurePath from typing import TYPE_CHECKING, Any @@ -99,10 +100,15 @@ class FreshDataDbtTransform: clean_config: CleanConfig | None = None system_actor: str = "freshdata" fail_on_low_score: bool = False + #: File stem for the audit (``_audit.json``); defaults to the table + #: name. :func:`gate_manifest` sets it when two gated models share an alias. + audit_name: str | None = None def __post_init__(self) -> None: - """Reject invalid gate policies when the transform is configured.""" + """Reject invalid gate policies and unsafe audit names at configuration.""" self.on_low_score = validate_on_low_score(self.on_low_score) + if self.audit_name is not None: + self.audit_name = _validate_audit_table_name(self.audit_name) def _split_table(self) -> tuple[str | None, str]: if self.schema: @@ -112,16 +118,24 @@ def _split_table(self) -> tuple[str | None, str]: return ".".join(prefix) or None, table return None, self.model_name + def _audit_stem(self, table: str) -> str: + return self.audit_name if self.audit_name is not None else table + def _write_audit(self, table: str, result: TrustGateResult) -> Path: - table = _validate_audit_table_name(table) + stem = _validate_audit_table_name(self._audit_stem(table)) out_dir = Path(self.output_dir) # type: ignore[arg-type] out_dir.mkdir(parents=True, exist_ok=True) - path = out_dir / f"{table}_audit.json" + path = out_dir / f"{stem}_audit.json" path.write_text(json.dumps(result.to_dict(), indent=2, default=str)) return path - def run(self) -> TrustGateResult: - """Read the model table, gate it, optionally write an audit, return the result.""" + def run(self, *, raise_on_fail: bool = True) -> TrustGateResult: + """Read the model table, gate it, optionally write an audit, return the result. + + A failing gate raises :class:`TrustGateError` (after the audit is written) when + ``on_low_score="fail"`` or ``fail_on_low_score=True``. Pass + ``raise_on_fail=False`` to get the failing result back instead. + """ conn = self.conn_str or os.environ.get("FRESHDATA_WAREHOUSE_CONN") if not conn: raise ValueError( @@ -129,7 +143,7 @@ def run(self) -> TrustGateResult: ) schema, table = self._split_table() if self.output_dir: - _validate_audit_table_name(table) + _validate_audit_table_name(self._audit_stem(table)) df = _read_table(conn, schema, table) _, result = evaluate_trust_gate( df, @@ -141,11 +155,39 @@ def run(self) -> TrustGateResult: ) if self.output_dir: self._write_audit(table, result) - if self.fail_on_low_score and not result.passed: + if raise_on_fail and ( + result.should_fail or (self.fail_on_low_score and not result.passed) + ): raise TrustGateError(result.message) return result +def _audit_names(models: list[tuple[str, dict[str, Any]]]) -> list[str | None]: + """Return an audit file stem per model, or ``None`` to keep the table name. + + Audit files are named after the model's alias, which dbt only requires to be + unique within a schema. When gated models share an alias (compared + case-insensitively, as audit files may land on a case-insensitive filesystem), + each of them is named ``"."`` instead, or after its manifest + ``unique_id`` when it has no schema or that name is still not unique. + """ + tables = [node.get("alias") or node.get("name") for _, node in models] + alias_counts = Counter(t.casefold() for t in tables if isinstance(t, str)) + names: list[str | None] = [] + for (node_id, node), table in zip(models, tables): + if not isinstance(table, str) or alias_counts[table.casefold()] < 2: + names.append(None) + continue + schema = node.get("schema") + names.append(f"{schema}.{table}" if isinstance(schema, str) and schema else node_id) + stems = [name if name is not None else table for name, table in zip(names, tables)] + stem_counts = Counter(s.casefold() for s in stems if isinstance(s, str)) + return [ + node_id if isinstance(stem, str) and stem_counts[stem.casefold()] > 1 else name + for (node_id, _), name, stem in zip(models, names, stems) + ] + + def gate_manifest( manifest_path: str | Path, *, @@ -178,9 +220,9 @@ def gate_manifest( if not isinstance(nodes, dict): raise ValueError(f"{manifest_path} is not a dbt manifest: no 'nodes' mapping") - models: list[Any] = [] # raw manifest nodes (untyped JSON) + models: list[tuple[str, Any]] = [] # (unique_id, raw manifest node) skipped: list[dict[str, Any]] = [] - for node in nodes.values(): + for node_id, node in nodes.items(): if not isinstance(node, dict) or node.get("resource_type") != "model": continue config = node.get("config") @@ -190,11 +232,12 @@ def gate_manifest( elif config.get("enabled") is False: skipped.append({"model": node.get("name"), "reason": "disabled"}) else: - models.append(node) + models.append((node_id, node)) + audit_names = _audit_names(models) if output_dir else [None] * len(models) summaries: list[dict[str, Any]] = [] failed = 0 - for node in models: + for (_, node), audit_name in zip(models, audit_names): name = node.get("name") schema = node.get("schema") table = node.get("alias") or name @@ -208,7 +251,8 @@ def gate_manifest( output_dir=output_dir, clean_config=clean_config, system_actor=system_actor, - ).run() + audit_name=audit_name, + ).run(raise_on_fail=False) except Exception as exc: # noqa: BLE001 - one bad model must not abort the run logger.warning("freshdata: gating model %r failed: %s", name, exc) summaries.append({"model": name, "error": str(exc)}) diff --git a/tests/test_integrations/test_dbt_transform_fail_audit.py b/tests/test_integrations/test_dbt_transform_fail_audit.py new file mode 100644 index 00000000..9cc9943e --- /dev/null +++ b/tests/test_integrations/test_dbt_transform_fail_audit.py @@ -0,0 +1,292 @@ +"""dbt transform: ``on_low_score="fail"`` raises (#343); unique audit files (#344).""" + +from __future__ import annotations + +import json +import sqlite3 +from contextlib import closing + +import pandas as pd +import pytest +import sqlalchemy as sa +from sqlalchemy import event + +from freshdata.integrations import TrustGateError +from freshdata.integrations.dbt import FreshDataDbtTransform, gate_manifest + + +@pytest.fixture +def warehouse(tmp_path, sample_df): + conn = f"sqlite:///{tmp_path / 'wh.db'}" + engine = sa.create_engine(conn) + with engine.begin() as connection: + sample_df.to_sql("orders", connection, index=False) + sample_df.to_sql("customers", connection, index=False) + engine.dispose() + return conn + + +@pytest.fixture +def schemas(tmp_path, monkeypatch): + """A sqlite warehouse exposing ``staging``, ``marts`` and ``archive`` via ATTACH. + + ``staging.orders`` (3 rows), ``marts.orders`` (4 rows) and ``archive.ORDERS`` + (5 rows) differ in row count, so an audit file's ``row_count_in`` shows which + model wrote it. The main database holds ``customers`` and a 2-row ``orders``. + """ + frames = { + "staging": ("orders", pd.DataFrame({"a": [1, None, None], "b": [None, None, "z"]})), + "marts": ("orders", pd.DataFrame({"a": [1, 2, 3, 4]})), + "archive": ("ORDERS", pd.DataFrame({"a": [1, 2, 3, 4, 5]})), + } + paths = {} + for schema, (table, frame) in frames.items(): + paths[schema] = tmp_path / f"{schema}.db" + # A stdlib connection: pandas 1.5's case-sensitivity check for a + # non-lowercase table name does not work on a SQLAlchemy 2 Connection. + with closing(sqlite3.connect(paths[schema])) as connection: + frame.to_sql(table, connection, index=False) + connection.commit() + main = f"sqlite:///{tmp_path / 'main.db'}" + engine = sa.create_engine(main) + with engine.begin() as connection: + pd.DataFrame({"id": [1, 2]}).to_sql("customers", connection, index=False) + pd.DataFrame({"a": [5, 6]}).to_sql("orders", connection, index=False) + engine.dispose() + + original = sa.create_engine + + def create_engine(url, **kwargs): # noqa: ANN001, ANN003, ANN202 + eng = original(url, **kwargs) + + @event.listens_for(eng, "connect") + def _attach(dbapi_connection, _record): # noqa: ANN001, ANN202 + for schema, path in paths.items(): + dbapi_connection.execute(f"ATTACH DATABASE '{path}' AS {schema}") + + return eng + + monkeypatch.setattr(sa, "create_engine", create_engine) + return main + + +def _write_manifest(tmp_path, nodes): + path = tmp_path / "manifest.json" + manifest = { + "nodes": { + unique_id: {"resource_type": "model", **node} for unique_id, node in nodes.items() + } + } + path.write_text(json.dumps(manifest)) + return path + + +def _audit_files(directory): + return sorted(p.name for p in directory.iterdir()) + + +# --------------------------------------------------------------------------- # +# #343: on_low_score="fail" # +# --------------------------------------------------------------------------- # +def test_on_low_score_fail_raises_and_still_writes_audit(warehouse, tmp_path): + out = tmp_path / "audit" + with pytest.raises(TrustGateError, match="trust gate failed"): + FreshDataDbtTransform( + model_name="orders", + conn_str=warehouse, + trust_score_threshold=999.0, + on_low_score="fail", + output_dir=str(out), + ).run() + audit = json.loads((out / "orders_audit.json").read_text()) + assert audit["passed"] is False + assert audit["on_low_score"] == "fail" + + +def test_run_raise_on_fail_false_returns_failing_result(warehouse): + result = FreshDataDbtTransform( + model_name="orders", + conn_str=warehouse, + trust_score_threshold=999.0, + on_low_score="fail", + fail_on_low_score=True, + ).run(raise_on_fail=False) + assert result.passed is False + assert result.should_fail is True + + +@pytest.mark.parametrize("on_low_score", ["warn", "skip"]) +def test_non_fail_policies_do_not_raise(warehouse, on_low_score): + result = FreshDataDbtTransform( + model_name="orders", + conn_str=warehouse, + trust_score_threshold=999.0, + on_low_score=on_low_score, + ).run() + assert result.passed is False + assert result.should_fail is False + + +@pytest.mark.parametrize("kwargs", [{"on_low_score": "fail"}, {"fail_on_low_score": True}]) +def test_passing_gate_never_raises(warehouse, kwargs): + result = FreshDataDbtTransform( + model_name="orders", conn_str=warehouse, trust_score_threshold=0.0, **kwargs + ).run() + assert result.passed is True + + +def test_gate_manifest_fail_policy_records_failure_and_continues(warehouse, tmp_path): + manifest = _write_manifest( + tmp_path, + { + "model.proj.orders": {"name": "orders", "schema": None, "alias": "orders"}, + "model.proj.customers": { + "name": "customers", + "schema": None, + "alias": "customers", + }, + }, + ) + summary = gate_manifest( + str(manifest), + conn_str=warehouse, + trust_score_threshold=999.0, + on_low_score="fail", + ) + assert summary["models_processed"] == 2 + assert summary["failed_models"] == 2 + assert summary["all_passed"] is False + assert [m["model"] for m in summary["models"]] == ["orders", "customers"] + for model in summary["models"]: + assert "error" not in model + assert model["passed"] is False + + +# --------------------------------------------------------------------------- # +# #344: unique audit files for same-alias models # +# --------------------------------------------------------------------------- # +def test_same_alias_in_different_schemas_writes_distinct_audits(schemas, tmp_path): + manifest = _write_manifest( + tmp_path, + { + "model.proj.staging_orders": { + "name": "staging_orders", + "schema": "staging", + "alias": "orders", + }, + "model.proj.marts_orders": { + "name": "marts_orders", + "schema": "marts", + "alias": "orders", + }, + "model.proj.customers": {"name": "customers", "schema": None}, + }, + ) + out = tmp_path / "audit" + summary = gate_manifest( + str(manifest), conn_str=schemas, output_dir=str(out), trust_score_threshold=0.0 + ) + assert summary["models_processed"] == 3 + assert all("error" not in m for m in summary["models"]) + assert _audit_files(out) == [ + "customers_audit.json", + "marts.orders_audit.json", + "staging.orders_audit.json", + ] + staging = json.loads((out / "staging.orders_audit.json").read_text()) + marts = json.loads((out / "marts.orders_audit.json").read_text()) + assert staging["row_count_in"] == 3 + assert marts["row_count_in"] == 4 + + +def test_non_colliding_models_keep_alias_audit_names(warehouse, tmp_path): + manifest = _write_manifest( + tmp_path, + { + "model.proj.orders": {"name": "orders", "schema": None, "alias": "orders"}, + "model.proj.customers": {"name": "customers", "schema": None}, + }, + ) + out = tmp_path / "audit" + summary = gate_manifest( + str(manifest), conn_str=warehouse, output_dir=str(out), trust_score_threshold=0.0 + ) + assert summary["all_passed"] is True + assert _audit_files(out) == ["customers_audit.json", "orders_audit.json"] + + +def test_colliding_alias_without_schema_uses_unique_id(schemas, tmp_path): + manifest = _write_manifest( + tmp_path, + { + "model.proj.main_orders": {"name": "main_orders", "schema": None, "alias": "orders"}, + "model.proj.archive_orders": { + "name": "archive_orders", + "schema": "archive", + "alias": "ORDERS", + }, + "model.proj.customers": {"name": "customers", "schema": None}, + }, + ) + out = tmp_path / "audit" + summary = gate_manifest( + str(manifest), conn_str=schemas, output_dir=str(out), trust_score_threshold=0.0 + ) + assert summary["models_processed"] == 3 + assert all("error" not in m for m in summary["models"]) + # "orders" and "ORDERS" collide case-insensitively: the schema-less model falls + # back to its unique_id, the other gets ".". "customers" is unique. + assert _audit_files(out) == [ + "archive.ORDERS_audit.json", + "customers_audit.json", + "model.proj.main_orders_audit.json", + ] + main = json.loads((out / "model.proj.main_orders_audit.json").read_text()) + archive = json.loads((out / "archive.ORDERS_audit.json").read_text()) + assert main["row_count_in"] == 2 + assert archive["row_count_in"] == 5 + + +def test_unsafe_schema_in_colliding_alias_is_rejected(schemas, tmp_path): + manifest = _write_manifest( + tmp_path, + { + "model.proj.evil_orders": { + "name": "evil_orders", + "schema": "../evil", + "alias": "orders", + }, + "model.proj.marts_orders": { + "name": "marts_orders", + "schema": "marts", + "alias": "orders", + }, + }, + ) + out = tmp_path / "audit" + summary = gate_manifest( + str(manifest), conn_str=schemas, output_dir=str(out), trust_score_threshold=0.0 + ) + evil, marts = summary["models"] + assert "safe dbt model name" in evil["error"] + assert "error" not in marts + assert summary["failed_models"] == 1 + assert _audit_files(out) == ["marts.orders_audit.json"] + assert not list(tmp_path.glob("evil*")) + + +@pytest.mark.parametrize("audit_name", ["", ".", "..", "../orders", "a/b", "a\\b", "/abs"]) +def test_transform_rejects_unsafe_audit_name(audit_name): + with pytest.raises(ValueError, match="safe dbt model name"): + FreshDataDbtTransform(model_name="orders", audit_name=audit_name) + + +def test_transform_audit_name_overrides_file_name(warehouse, tmp_path): + FreshDataDbtTransform( + model_name="orders", + conn_str=warehouse, + output_dir=str(tmp_path / "audit"), + trust_score_threshold=0.0, + audit_name="main.orders", + ).run() + assert _audit_files(tmp_path / "audit") == ["main.orders_audit.json"]