Schema-Drift Auto-Repair Agent

When an upstream column is renamed, retyped or dropped, this agent uses DataHub's column-level lineage to find exactly what breaks, repairs the code, and ships the fix as a reviewed pull request.

Comment

Inspiration

Schema drift is a weekly tax on every data team, and it fails silently. An upstream team renames a column, nothing throws an error, and the downstream models keep running while quietly producing wrong numbers. You find out days later from someone staring at a bad dashboard.

The tooling gap is specific. Grep tells you which files contain a string. Table-level lineage tells you everything downstream. Neither answers the question that matters: which code actually breaks, and which code is fine? That question is answerable, but only with column-level lineage, which is exactly what DataHub has.

Most agents in this space generate new code from English. I wanted the inverse: an agent that repairs existing code when the world changes underneath it.

What it does

Given a renamed, retyped, or dropped source column, the agent:

  1. Detects the drift by diffing a committed schema snapshot against the live catalog, inferring a rename from type, ordinal position and name similarity, with a confidence score and a written rationale.
  2. Computes the blast radius from column-level lineage and sorts every downstream asset into three buckets, each with a stated reason. In the filmed run: 3 assets require a patch (two dbt models and an Airflow task), 2 marts are downstream but insulated by an alias created one hop up, so they are flagged for review and deliberately not touched, and 7 models are correctly skipped, each reported with the columns it actually consumes.
  3. Rewrites the code. Four files across those three assets: two dbt SQL models, the schema.yml (including a new test), and the Airflow DAG. sqlglot parses the SQL to locate references, then surgical edits touch only the changed tokens. The language model never writes code. A test asserts the generated patches are byte-identical with and without it; the model's job is the PR narrative and the migration note.
  4. Validates every column reference against the catalog before anything ships. In the filmed run that gate is 23 of 23 resolved: 15 against live DataHub schemas, 6 against the repaired schema this change creates, and 2 inside the query itself. One unresolvable reference blocks the entire pull request.
  5. Opens a real GitHub PR whose body carries the column-lineage path as a Mermaid diagram, the captured queries that prove usage, and the reason every untouched model was left alone.
  6. Writes the repair back to DataHub: corrected fine-grained lineage, column documentation, tags on every repaired asset, an incident raised through the incidents API and moved to triage, an institutional-memory link, and a process-instance run record. The next engineer opens the catalog and finds the answer already there.

Why the three buckets matter

Reporting the skipped models with reasons is the point of the project. "This mart really is downstream, but it reads order_id, quantity and unit_price, none of which are on the changed column's path" is what proves the agent understood the graph instead of pattern-matching a string. Listing everything downstream is easy. Knowing what actually breaks is not.

How I built it

Python 3.11, the OpenAI Agents SDK, sqlglot, FastAPI, and a Vite + React + Tailwind UI.

DataHub is the reasoning substrate, used through two surfaces on purpose. The agent's reasoning reads go through the DataHub MCP server: search, list_schema_fields, get_lineage on the changed column, get_dataset_queries, and get_lineage_paths_between to confirm each path. The filmed run makes 14 tool calls, 8 of them DataHub MCP reads, all visible live in the UI timeline. The engine's deterministic reads (blast radius, validation schemas) and all write-back go through the DataHub Python SDK and GraphQL.

The demo warehouse is a real dbt project (5 raw sources, 5 staging models, 6 marts) plus an Airflow DAG, living in the same repo the agent opens PRs against. One clone, no Snowflake account, fully reproducible, and every finished run can be reopened read-only in the UI with ?run=<run-id>.

Challenges I ran into

Column-level lineage has an undocumented rule. A fine-grained edge only renders if both endpoints carry a schemaMetadata aspect with exactly matching field paths, and table-level upstreams coexist in the same aspect. Miss either and edges vanish with no error at any layer. The seed script now regression-tests this contract.

Soft-delete corrupts the lineage index. Re-seeding after a soft delete brought the datasets back and GraphQL read their lineage correctly, but degree-1 column edges never returned to lineage search, while hop-2 and hop-3 results still did. The graph looked merely incomplete rather than broken, which is the worst kind of wrong. Hard delete plus re-emit fixes it.

The most valuable bug I fixed was a silent one. An unlucky run once classified every asset as "correctly skipped", produced zero patches, and rendered as a plausible success, because the lineage index had returned nothing and the engine accepted that as an answer. Saying "nothing is affected" when the truth is "I could not tell" is the worst thing a tool like this can do. It now refuses instead: when the catalog's lineage aspects declare edges that search cannot see, the run fails loudly and says why.

Accomplishments that I'm proud of

  • The three-bucket blast radius with a written reason for every skip, backed by per-column lineage evidence a reviewer can click through.
  • Zero hallucinated columns, enforced rather than hoped: AST-located edits plus a hard validation gate, with a test proving the model cannot alter generated code.
  • Surgical diffs a human will actually review. No whole-file reformatting, ever.
  • A structural guarantee you can check from a clone in ten seconds: runtime code never emits a schemaMetadata aspect, so no repair can redefine a dataset behind your back. A test enforces it.

What I learned

Metadata is only load-bearing if you let it block you. The interesting design decisions were all refusals: refusing to emit a patch with an unresolved reference, refusing to report a blast radius the index could not confirm, refusing to let the model touch code. An agent that always produces an answer is easy. One that knows when it does not have one is the useful one.

What's next for Schema-Drift Auto-Repair Agent

Prefect and Dagster support alongside Airflow, multi-repo PR fan-out for drift that crosses repository boundaries, and contributing a schema-drift-repair skill upstream to datahub-project/datahub-skills carrying the lineage-contract gotchas this project paid for. UPDATE: The packaged datahub-incident-triage skill is now proposed upstream as https://github.com/datahub-project/datahub-skills/pull/141

Built With

Share this project:

Updates