LineageGuard AI

Stops silent data changes before they break production ML—using DataHub context to score blast radius, propose a constrained repair, validate it in dbt, and write the decision back.

Comment

Inspiration

A schema change that looks harmless at the table level can silently invalidate features, models, and online deployments. Teams usually discover the real blast radius only after a model degrades or an API fails. DataHub already knows the relationships, owners, governance tags, and ML entities involved; the missing piece was an agent that could turn that context into a safe, tested response.

What it does

LineageGuard AI is an event-driven ML change-safety agent. When DataHub emits a schema change, it:

  1. reads the changed dataset and three-hop lineage through the official DataHub MCP server;
  2. traces the impact across datasets, MLFeature, MLModel, and MLModelDeployment entities;
  3. combines breaking-change, production, governance, ownership, and usage evidence into an explainable risk score;
  4. chooses a bounded response: observe, open a draft repair PR, or block and remediate;
  5. proposes the smallest compatibility repair;
  6. applies that patch only inside a temporary sandbox and runs a real dbt build;
  7. optionally opens a GitHub draft PR and writes a durable decision document back to DataHub.

The included incident renames customer_age to age_years without a compatibility alias. Against the real local DataHub graph, LineageGuard finds one affected feature, one model, one production deployment, PII and Tier1 governance context, and a technical owner. It scores the incident 90/100 (critical), restores the stable contract, and passes all nine dbt build/test steps.

How we built it

The demo runs on DataHub OSS 1.6, the official DataHub MCP server, DataHub Actions, the Python SDK, dbt-duckdb, FastAPI, and Docker.

A custom DataHub Action consumes real schemaMetadata Metadata Change Log events and normalizes current/previous aspects into a small versioned event contract. The MCP context provider uses official read-only tools for entities and lineage, with SDK enrichment for ML deployment relationships. The risk engine is deterministic and exposes every point of evidence, so the safety decision is reviewable rather than hidden inside an LLM prompt.

The repair planner is deliberately constrained. It can produce exact one-match compatibility edits, rejects ambiguous replacements and path traversal, and never writes to the working tree during validation. A temporary project copy receives the proposed patch and must pass an allowlisted dbt build with a timeout. External mutation is disabled by default; GitHub output is always a draft PR.

The full workflow was exercised with real DataHub ingestion, a real Action event, MCP context, ML metadata, a passing sandbox repair, and a DataHub decision-document writeback.

What makes it different

LineageGuard does not stop at displaying lineage or generating remediation prose. It closes the operational loop:

context graph → evidence-backed decision → minimal repair → executable proof → durable catalog context

That makes it useful to both people and future agents. The next responder can see what changed, why it was critical, which production ML assets were affected, what repair passed, and who owns the system.

Challenges we ran into

The OSS MCP response did not expose every ML deployment relationship needed for the demo, so we combined official MCP reads with a narrow SDK enrichment while preserving read-only MCP behavior. Running an actual dbt validation inside a DataHub Action also required tuning the Kafka consumer poll interval. Finally, we treated repair execution as a security boundary: exact paths, exact matches, temporary directories, command allowlists, and timeouts were implemented before enabling any external write path.

Accomplishments

  • Real dataset → feature → model → production deployment lineage in DataHub.
  • Real-time schema event consumption through DataHub Actions.
  • Explainable 90/100 risk result on the live OSS graph.
  • Passing 9/9 dbt repair validation in an isolated sandbox.
  • Durable DataHub decision-document writeback.
  • 10 passing unit tests and 5/5 deterministic evaluation cases.
  • Locked, reproducible Docker image and zero-credential demo mode.
  • Public Apache-2.0 repository with successful GitHub Actions CI.

What we learned

Reliable agents need more than tool access. They need high-quality context, explicit decision policies, and a proof step before taking action. DataHub's graph makes the impact legible; deterministic scoring and sandbox validation make the agent trustworthy.

What's next

Next steps include domain-specific policy packs, column-level repair strategies, model-specific smoke tests, retraining triggers, approval workflows, DataHub Cloud usage signals, and repository-scale multi-file draft PRs.

Try it live

The Cloudflare judge experience replays the verified DataHub OSS 1.6.0 and official MCP 0.6.0 result with no local setup. It is explicitly a zero-secret evidence replay; the full implementation and reproduction path remain in the public repository.

Source and setup: github.com/spectramaster/lineageguard-ai

Built With

Share this project:

Updates