From 78e4cc05513d5b773b6d9403d4c8a27a75e69529 Mon Sep 17 00:00:00 2001 From: Sadha Chilukoori Date: Mon, 13 Jul 2026 16:16:56 -0700 Subject: [PATCH] Add guides for signal store, DQ checks, and SQL lineage parser --- docs/dq-checks.md | 178 +++++++++++++++++++++++++++++++++++++++++++ docs/signal-store.md | 162 +++++++++++++++++++++++++++++++++++++++ docs/sql-lineage.md | 121 +++++++++++++++++++++++++++++ mkdocs.yml | 3 + 4 files changed, 464 insertions(+) create mode 100644 docs/dq-checks.md create mode 100644 docs/signal-store.md create mode 100644 docs/sql-lineage.md diff --git a/docs/dq-checks.md b/docs/dq-checks.md new file mode 100644 index 0000000..767c0d1 --- /dev/null +++ b/docs/dq-checks.md @@ -0,0 +1,178 @@ +# Data Quality Checks + +Tollkeeper runs DQ checks against staged data before publishing. If any check fails, the staged version is rolled back and production stays untouched. + +## Built-in checks (Polars) + +Install the `polars` extra: + +```bash +pip install tollkeeper[polars] +``` + +Five checks ship with tollkeeper: + +### NullCheck + +Fails if a column contains any null values. + +```python +from tollkeeper import NullCheck + +NullCheck("customer_id") +# Result: "3 nulls in 'customer_id'" -> FAIL +``` + +### RowCountCheck + +Fails if the row count is below a threshold. + +```python +from tollkeeper import RowCountCheck + +RowCountCheck(min_rows=100) +# Result: "47 rows, minimum 100" -> FAIL +``` + +### UniqueCheck + +Fails if any combination of columns has duplicates. + +```python +from tollkeeper import UniqueCheck + +UniqueCheck(["region", "date"]) +# Result: "2 duplicate groups on ['region', 'date']" -> FAIL +``` + +### ExpressionCheck + +Fails if any row does not satisfy a Polars expression. + +```python +import polars as pl +from tollkeeper import ExpressionCheck + +ExpressionCheck("positive_revenue", pl.col("revenue") > 0) +# Result: "5 rows violate 'positive_revenue'" -> FAIL +``` + +### SqlCheck + +Fails if any row does not satisfy a SQL WHERE condition. Uses Polars' in-memory SQL engine. + +```python +from tollkeeper import SqlCheck + +SqlCheck("valid_age", "age > 0 AND age < 150") +# Result: "1 rows violate 'valid_age'" -> FAIL +``` + +The staged data is registered as a table named `data` in the SQL context. + +## Using checks in a pipeline + +```python +from tollkeeper import Tollkeeper, CsvBackend, NullCheck, RowCountCheck, UniqueCheck + +backend = CsvBackend("/data/staging", "/data/prod") + +( + Tollkeeper(backend) + .table("orders") + .audit([ + NullCheck("order_id"), + RowCountCheck(min_rows=1), + UniqueCheck(["order_id"]), + ]) + .publish() +) +``` + +## Failure modes + +### Hard failure (default) + +All checks run, then if any failed, the staged version is rolled back and `AuditFailedError` is raised: + +```python +from tollkeeper import AuditFailedError + +try: + Tollkeeper(backend).table("orders").audit(checks, on_failure="stop").publish() +except AuditFailedError as e: + print(f"Table: {e.table}") + print(f"Version: {e.version_ref}") + for result in e.failed_checks: + print(f" {result.check_name}: {result.details}") +``` + +### Soft failure + +Publish proceeds despite failures. Use `on_notify` to handle the failures: + +```python +def alert(table, version_ref, failed_checks): + for check in failed_checks: + send_alert(f"{table}: {check.check_name} failed - {check.details}") + +( + Tollkeeper(backend) + .table("orders") + .audit(checks, on_failure="continue", on_notify=alert) + .publish() +) +``` + +## Check report + +After auditing, the session's report shows what passed and what failed: + +```python +session = Tollkeeper(backend).table("orders").audit(checks) + +session.report.passed # list of CheckResult where passed=True +session.report.failed # list of CheckResult where passed=False +session.report.results # all results +``` + +## Remote check execution + +Pass a connection to run checks against a remote engine (e.g., Trino, Presto): + +```python +Tollkeeper(backend).table("orders").audit(checks, conn=trino_connection) +``` + +The `conn` is forwarded to each check's `run()` method. Built-in Polars checks ignore it, but custom checks can use it. + +## Writing a custom check + +Subclass `BaseCheck` and implement `run()`: + +```python +from tollkeeper.checks.base import BaseCheck, CheckResult + +class FreshnessCheck(BaseCheck): + """Fails if the most recent row is older than max_age_hours.""" + + def __init__(self, timestamp_col: str, max_age_hours: int = 24) -> None: + self._col = timestamp_col + self._max_age = max_age_hours + + def run(self, version_ref, *, conn=None): + import polars as pl + from datetime import datetime, timedelta + + df = pl.scan_csv(version_ref).collect() + newest = df[self._col].cast(pl.Datetime).max() + cutoff = datetime.now() - timedelta(hours=self._max_age) + fresh = newest >= cutoff + return CheckResult( + check_name=self.name, + passed=fresh, + details=f"newest row: {newest}, cutoff: {cutoff}", + ) +``` + +The `name` property defaults to the class name. Override it for a custom display name. diff --git a/docs/signal-store.md b/docs/signal-store.md new file mode 100644 index 0000000..434101d --- /dev/null +++ b/docs/signal-store.md @@ -0,0 +1,162 @@ +# Signal Store + +Tollkeeper's signal store tracks which tables have passed their audit. Pipelines that depend on a table can poll or wait for its signal before proceeding. + +## Setup + +Two implementations ship with tollkeeper: + +```python +from tollkeeper import SqliteSignalStore, DbApiSignalStore + +# SQLite (zero config, good for single-machine pipelines) +store = SqliteSignalStore("/var/data/signals.db") + +# DB-API 2.0 (PostgreSQL, MySQL, any PEP 249 driver) +import psycopg2 +conn = psycopg2.connect("dbname=pipeline") +store = DbApiSignalStore(conn, paramstyle="format") # PostgreSQL uses %s +``` + +`SqliteSignalStore` defaults to `:memory:` if no path is given. + +`DbApiSignalStore` accepts `paramstyle="qmark"` (SQLite, default) or `paramstyle="format"` (PostgreSQL, MySQL). + +Both create the required tables on first use. + +## Wiring into Tollkeeper + +Pass the store to the `Tollkeeper` constructor. When an audit passes, tollkeeper writes a signal automatically. + +```python +from tollkeeper import Tollkeeper, CsvBackend, SqliteSignalStore + +store = SqliteSignalStore("signals.db") +backend = CsvBackend("/data/staging", "/data/prod") + +( + Tollkeeper(backend, signal_store=store) + .table("orders") + .audit([RowCountCheck(min_rows=1)]) + .publish() +) + +# Signal written automatically on audit pass +assert store.check("orders") is not None +``` + +If the audit fails, no signal is written. + +## Reading and waiting for signals + +```python +# Non-blocking check +signal = store.check("orders") +if signal: + print(f"orders passed at {signal.execution_ts}") + +# Blocking wait (raises TimeoutError after 300s by default) +signal = store.wait("orders", timeout_s=60, poll_s=5) +``` + +## Execution context + +Partition signals by execution context (date, region, etc.) so daily runs don't overwrite each other: + +```python +ctx = {"ds": "2026-07-13"} + +( + Tollkeeper(backend, signal_store=store) + .table("orders") + .audit([RowCountCheck(min_rows=1)], execution_ctx=ctx) + .publish() +) + +# Only returns the signal for this specific date +store.check("orders", {"ds": "2026-07-13"}) + +# Different date returns None +store.check("orders", {"ds": "2026-07-12"}) # None +``` + +## Dependencies and cascading deletes + +Register that one table depends on another. When the upstream signal is deleted (e.g., a re-run), downstream signals are automatically invalidated. + +```python +store.register_dep("raw_events", "dim_users", cascade_policy="cascade") +store.register_dep("dim_users", "fact_sessions", cascade_policy="cascade") + +# Write signals for all three +store.write(Signal(table_name="raw_events")) +store.write(Signal(table_name="dim_users")) +store.write(Signal(table_name="fact_sessions")) + +# Deleting raw_events cascades through the chain +store.delete("raw_events") +assert store.check("dim_users") is None +assert store.check("fact_sessions") is None +``` + +Use `cascade_policy="notify"` to get a callback instead of automatic deletion: + +```python +def on_upstream_invalidated(upstream, downstream, ctx): + print(f"{upstream} was invalidated, {downstream} may be stale") + +store = SqliteSignalStore("signals.db", on_delete_callback=on_upstream_invalidated) +store.register_dep("raw_events", "dim_users", cascade_policy="notify") +``` + +## DQ result storage + +The signal store also persists individual check results for auditing and debugging: + +```python +from tollkeeper.signals.base import DqResult + +result = DqResult(table_name="orders", check_name="NullCheck", passed=True, details="0 nulls in 'id'") +store.write_dq_result(result) + +results = store.get_dq_results("orders") +for r in results: + print(f"{r.check_name}: {'PASS' if r.passed else 'FAIL'} - {r.details}") +``` + +## Resource management + +Signal stores hold database connections. Close them when done: + +```python +# Explicit close +store = SqliteSignalStore("signals.db") +# ... use store ... +store.close() + +# Context manager (preferred) +with SqliteSignalStore("signals.db") as store: + Tollkeeper(backend, signal_store=store).table("orders").audit(checks).publish() +# connection closed automatically +``` + +## Custom signal store + +Implement the `SignalStore` ABC to back signals with Redis, DynamoDB, or anything else: + +```python +from tollkeeper.signals.base import SignalStore, Signal, DqResult + +class RedisSignalStore(SignalStore): + def write(self, signal: Signal) -> None: ... + def delete(self, table: str, execution_ctx: dict | None = None) -> None: ... + def check(self, table: str, execution_ctx: dict | None = None) -> Signal | None: ... + def write_dq_result(self, result: DqResult) -> None: ... + def get_dq_results(self, table: str, execution_ctx: dict | None = None) -> list[DqResult]: ... + def delete_dq_results(self, table: str, execution_ctx: dict | None = None) -> None: ... + def register_dep(self, upstream_table: str, downstream_table: str, ...) -> None: ... + def get_downstream(self, table: str, execution_ctx: dict | None = None) -> list[tuple[str, dict, str]]: ... + def close(self) -> None: ... +``` + +The `wait()` method is inherited and works with any implementation that provides `check()`. diff --git a/docs/sql-lineage.md b/docs/sql-lineage.md new file mode 100644 index 0000000..b5b06de --- /dev/null +++ b/docs/sql-lineage.md @@ -0,0 +1,121 @@ +# SQL Lineage Parser + +Tollkeeper's lineage parser extracts source and sink tables from SQL statements using sqlglot. This powers automatic dependency detection in the Airflow integration, but you can use it standalone. + +## Setup + +Install the `sqlglot` extra: + +```bash +pip install tollkeeper[sqlglot] +``` + +## Basic usage + +```python +from tollkeeper import extract_lineage + +result = extract_lineage(""" + INSERT INTO analytics.fact_orders + SELECT o.*, c.name + FROM raw.orders o + JOIN raw.customers c ON o.customer_id = c.id +""") + +result.sources # frozenset({'raw.orders', 'raw.customers'}) +result.sinks # frozenset({'analytics.fact_orders'}) +``` + +## Supported statements + +| Statement | Sources | Sinks | +|-----------|---------|-------| +| `SELECT ... FROM a JOIN b` | `{a, b}` | `{}` | +| `INSERT INTO t SELECT ... FROM a` | `{a}` | `{t}` | +| `CREATE TABLE t AS SELECT ... FROM a` | `{a}` | `{t}` | +| `MERGE INTO t USING s ON ...` | `{s}` | `{t}` | + +Multi-statement SQL is supported. Sources and sinks are accumulated across all statements. + +## CTE handling + +Common Table Expressions are excluded from sources. Only the real tables they reference are reported: + +```python +result = extract_lineage(""" + WITH enriched AS ( + SELECT * FROM raw.events + ) + INSERT INTO analytics.sessions + SELECT * FROM enriched +""") + +result.sources # frozenset({'raw.events'}), not {'enriched'} +result.sinks # frozenset({'analytics.sessions'}) +``` + +## Fully qualified names + +Catalog, schema, and table names are preserved: + +```python +result = extract_lineage("SELECT * FROM my_catalog.my_schema.my_table") +result.sources # frozenset({'my_catalog.my_schema.my_table'}) +``` + +## Dialect support + +Pass a sqlglot dialect for engine-specific SQL syntax: + +```python +result = extract_lineage( + "SELECT * FROM `project.dataset.table`", + dialect="bigquery", +) +``` + +See [sqlglot dialects](https://sqlglot.com/sqlglot/dialects.html) for the full list. + +## Error handling + +```python +# Empty SQL +extract_lineage("") # raises ValueError("SQL string is empty") + +# Unparseable SQL +extract_lineage("NOT SQL") # raises ValueError("Failed to parse SQL: ...") + +# Jinja templates (common in Airflow) +extract_lineage("SELECT * FROM {{ params.table }}") +# raises ValueError("SQL contains Jinja template expressions; provide explicit sources/sinks") +``` + +Jinja-templated SQL cannot be statically parsed. In the Airflow integration, provide `sources` and `sinks` explicitly when your SQL uses templates. + +## Use with Airflow task groups + +The parser is used automatically by `TollkeeperTaskGroup` to wire up sensor dependencies: + +```python +from airflow_tollkeeper import TollkeeperTaskGroup + +group = TollkeeperTaskGroup( + table="fact_orders", + sql="INSERT INTO fact_orders SELECT * FROM dim_customers JOIN raw_events ...", + backend=backend, + signal_store=store, +) +# Sensors are created for dim_customers and raw_events automatically +``` + +If the SQL contains Jinja or the parser can't extract what you need, pass sources explicitly: + +```python +group = TollkeeperTaskGroup( + table="fact_orders", + sql="{{ params.insert_sql }}", + sources=["dim_customers", "raw_events"], + backend=backend, + signal_store=store, +) +``` diff --git a/mkdocs.yml b/mkdocs.yml index 86e6637..f248871 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -31,6 +31,9 @@ nav: - Home: index.md - Getting Started: getting-started.md - Guides: + - Data Quality Checks: dq-checks.md + - Signal Store: signal-store.md + - SQL Lineage Parser: sql-lineage.md - SQL Passthrough Strategy: sql-strategy-flow.md - Docker Testing: docker-testing.md - API Reference: api.md