Warden

Generates what's safe. Warns what's not.

Comment

Inspiration

Every day, engineers prompt LLMs to "write me a pipeline for this table" — and the model happily generates code that reads nothing about the data. It doesn't know the column is PII. It doesn't know the field is a restricted glossary term. It doesn't know the table is governed by a compliance policy. The result is the new normal: fluent, confident, and catastrophically unsafe data code.

We wanted to prove the opposite is possible. Not by adding more guardrails on the output side, but by starting on the right side — the catalog. What if an agent refused to write a single line before understanding the data it was about to touch? What if governance wasn't a review step at the end, but the first thing the agent reads?

That's WARDEN: governance-aware code generation, where the DataHub catalog is not a nice-to-have lookup — it is the source of truth the agent is built around.

What it does

WARDEN is a DataHub-powered AI agent that generates compliance-aware data pipelines and delivers them as ready-to-merge GitHub pull requests.

Give it a plain-English request:

Create an airflow_dag that aggregates order amount from the orders dataset by order_date

and WARDEN will:

  1. Resolve intent — turn your sentence into a structured request (dataset, fields, operation, pipeline type).
  2. Read the catalog first — resolve a real DataHub URN, then pull its schema, semantic tags, glossary terms, and lineage from the DataHub GMS GraphQL API.
  3. Classify every field — attach tags like pii, restricted, or warden_generated to each column.
  4. Enforce deterministically — run a rules-based compliance engine that returns one of three verdicts:
    • ALLOW → generate the DAG and open a pull request
    • MASK → wrap sensitive columns in a real, canonical masking function, then open the PR
    • BLOCK → generate nothing, open nothing, and explain exactly which metadata rule made it unsafe
  5. Validate the code — Python AST validation for correctness and a leakage validator to guarantee masked columns never appear in outputs.
  6. Write the audit back to DataHub — tag the source dataset (warden_generated, warden_blocked) and append an audit note with the PR link, so the next human or agent inherits the knowledge.

Every PR carries the full governance story: the original request, datasets used, per-field decisions, risk scores, and validation verdict.

How we built it

Stack: FastAPI backend + vanilla-JS UI, DataHub GMS GraphQL for catalog reads and write-back, GitHub REST API for PR automation, Apache Airflow (via airflow-dag-parser) for generated code, edge-tts for narration, and headless Chrome driven over CDP for the demo screenshots.

The compliance core. Decisions come from a rules engine over catalog metadata. Each field gets a risk score and a verdict; the engine is deterministic and explainable — no LLM in the hot path for whether to block:

$$R(\text{field}) = \max_{\text{rule } r} \big( w_r \cdot \mathbb{1}[\text{field matches rule } r] \big)$$

The catalog-first generator. The LLM is only invoked after catalog resolution, and every generated field references the resolved, classified schema. Validation is layered: an AST validator catches broken code, a leakage validator proves masked/restricted fields can't leak into outputs.

Real write-back. Tags land on real DataHub datasets via GMS GraphQL; audit notes are appended to dataset documentation. It's not a demo loop — the state actually persists in the graph.

Hardening the PR path. Writing files to a real repo surfaced a classic bug: overwriting an existing file needs the current SHA. We added get_file_sha() + SHA payloads to upload_file(), then locked it in with regression tests — the suite grew to 122 passing tests.

Demo production. We made the whole pipeline runnable in an offline demo mode and produced the demo video with real, live PRs — one ALLOW, one MASK — and a genuine BLOCK that created no PR.

Challenges we ran into

  • The catalog is a graph, not a spreadsheet. Datasets have URNs, owners, sibling sets, tag attachments, glossary terms, and lineage edges — each a different GraphQL entity. Teaching the agent to resolve a URN and read all of it correctly was most of the battle.
  • GitHub's SHA check. create on an existing file returns 422. Diagnosing and fixing it required deep-API debugging (get_file_sha() + SHA-aware uploads) and tests to keep it fixed.
  • PII masking that's actually safe. Masking in a string template is easy; proving (and demoing) that a masked field can't leak through downstream operators required a dedicated leakage validator.
  • Demoing honesty. It's tempting to fake a BLOCK. We demoed a real one against live metadata — and you can see in the video that no PR exists.
  • Determinism vs. model noise. If a compliance verdict depends on the model, it's not compliance. We kept the LLM out of the decision path entirely.

Accomplishments that we're proud of

  • A compliance engine whose verdicts are explainable — every BLOCK/MASK tells you the exact metadata rule that triggered it.
  • Real, persistent governance write-back to DataHub, not a mocked graph.
  • 122 passing tests on governance, code generation, masking, validation, and PR flows.
  • A one-command offline demo mode so anyone can run it without a cluster.
  • A demo video with live PRs and a genuinely blocked case — no actors, no fake screenshots.

What we learned

  • Context is a governance decision. What an agent "knows" about the data changes what it's allowed to write — so the agent's catalog access is the control plane.
  • Deterministic beats clever for safety-critical paths. Encode the rules; let the LLM do the fluent part.
  • The write-back is what makes it a system. An audit that nobody can see (or that isn't persisted) is not governance.
  • Demo every claim. "No PR was opened" is only convincing when the video shows the graph with no PR.

What's next for Warden

  • Policy-as-config: YAML-published compliance policies teams can tune without touching code.
  • Multi-system lineage: act on lineage edges, not just the source dataset — e.g., re-mask consumers when a producer column changes.
  • Post-merge enforcement: CI checks that block merges when a policy decision would change, closing the loop.
  • Human-in-the-loop modes: configurable override workflows for MASK→ALLOW escalation with full audit trail.
  • More generators: dbt models, Spark jobs, and SQL-embedded lineage.
  • Trust signals: surface per-field risk scores and coverage metrics in the DataHub UI itself.

Open source (Apache 2.0) — github.com/AlvinGeorge-AG/Warden

Built With

Share this project:

Updates