DataHub DAG Generator

An AI agent that reads your DataHub lineage and metadata to generate Airflow 3 DAGs - with quality checks, freshness gates, and PII audits wired in from your tags and glossary.

Comment

Inspiration

Data engineers already have everything they need to build a correct pipeline — it's all in DataHub: the lineage graph, freshness SLAs, PII tags, and quality glossary terms. But turning that institutional knowledge into an Airflow DAG still means hours of manual work: copy-paste a template, re-derive the execution order, remember which tables need freshness gates, add PII audits by hand.

We wanted to close that loop. If the metadata already says a table is tagged pii or carries a FreshnessSLA glossary term, a tool should be able to read those signals and wire the right tasks into the DAG automatically — without a human translating them.

What it does

datahub-dag is a CLI tool that takes a target table name and generates a production-ready Airflow 3 DAG in seconds:

  1. Discovers lineage — traverses the upstream graph in DataHub via BFS to find every table the target depends on
  2. Reads metadata signals — tags (pii, daily_refresh) and glossary terms (FreshnessSLA, EmptyLoad) on every node
  3. Maps signals to tasks — pii → data_audit_<table>, FreshnessSLA → freshness_check_<table>, EmptyLoad → validate_row_count_<table>
  4. Renders a clean DAG — topological sort + transitive reduction, deterministic output, Airflow 3 syntax
  5. Optionally opens a PR — commits the file to a branch and creates a GitHub PR with a lineage summary table

How we built it

The tool runs in two modes:

Agent mode (default): An LLM agent (Claude or any OpenRouter model) connects to a live DataHub instance via the MCP protocol. It calls search, get_lineage, and get_entities tools to explore the graph, then invokes a custom render_airflow_dag tool with a validated plan. The renderer is fully deterministic — the LLM plans, but never writes code directly.

Script mode: Skips the LLM entirely. DataHub SDK calls + policies.py rule engine produce the same output deterministically. Useful for CI or environments without an API key.

Key design decisions:

  • No LLM code injection — all BashOperator commands are generated by airflow.py from templates. The LLM cannot inject arbitrary shell commands.
  • Transitive reduction — removes redundant lineage edges before rendering so the DAG only shows direct dependencies.
  • Kahn's algorithm with a min-heap for a deterministic, reproducible task order across runs.

Challenges we ran into

  • MCP + async agent loop — bridging the MCP session lifecycle with a bounded async tool-call loop required careful state management to avoid dangling connections and runaway agents.
  • LLM safety boundary — deciding exactly where the LLM's authority ends and the deterministic renderer begins. We settled on: the LLM produces a plan (structured JSON), the renderer produces code.
  • Transitive edges in lineage — DataHub sometimes returns both direct and transitive edges. Without reduction, the generated DAG had redundant dependencies that would confuse Airflow's scheduler.

Accomplishments that we're proud of

  • A genuine end-to-end demo: point the tool at a real DataHub instance with the NYC Taxi dataset and get a working, runnable Airflow DAG in under 30 seconds.
  • The metadata → quality check mapping is table-driven and extensible — adding a new signal is a one-line change in policies.py.
  • Two modes, one output contract: GenerationResult is identical whether the LLM or the script engine produced it, so all downstream logic (file emit, PR creation, write-back) is shared.

What we learned

  • The MCP protocol is a surprisingly clean interface for giving LLM agents structured access to data catalogs — much better than prompt-stuffing raw API responses.
  • Metadata that looks decorative (tags, glossary terms) becomes load-bearing the moment you have a system that acts on it. Data quality improves as a side effect of better tagging.

What's next

  • Support for dbt and Spark operators alongside the current shell-based stubs
  • Schema change detection: if upstream schema drifts, flag it as a DAG comment
  • DataHub write-back of execution results to close the observability loop

Built With

  • airflow
  • claude
  • datahub
  • mcp
  • openrouter
  • python
Share this project:

Updates