DataHub-Grounded dbt Model Generator

Describe a data transformation in plain English. Get a DataHub-grounded dbt model, validated, registered in the catalog, and shipped as a GitHub PR.

Comment

Inspiration

Data teams constantly rebuild the same kinds of dbt models - joins, filters, lifetime value calculations - by hand, every time someone needs a new view of the data. Meanwhile, if you just ask an LLM to write that SQL for you, it often hallucinates table or column names that don't actually exist in your warehouse, because it's guessing from general training knowledge instead of your real schema.

We wanted to see if we could fix that specific failure mode using DataHub - not just as a place to look up metadata, but as the source of truth an AI agent is forced to work within. If the agent can only ever reference tables and columns that DataHub confirms are real, hallucination stops being possible by construction, not just by hoping the model behaves.

What it does

You describe a data transformation in plain English - for example, "join orders with customer lifetime value, filtered to active accounts". The agent then:

  1. Searches DataHub for the relevant tables, their exact schemas, and upstream/downstream lineage
  2. Resolves business terms (like "active account") against DataHub's business glossary, so ambiguous language maps to the right column
  3. Flags any PII-tagged columns or failing data quality assertions before generation even happens
  4. Asks Groq (Llama 3.3) to generate the dbt SQL model and schema.yml, strictly grounded in that retrieved schema
  5. Validates every ref(), source(), and column reference against the real DataHub catalog, automatically asking the LLM to repair anything invalid, up to two retries
  6. Opens a GitHub PR with the finished files and a rich description covering lineage, glossary terms used, and governance warnings
  7. Writes back to DataHub - registering the new model as a dataset and adding upstream lineage edges, so the catalog reflects the new model just like it would after a real dbt run

The whole pipeline streams live to the frontend, so you watch each step happen in real time.

How we built it

The backend is Python + FastAPI, streaming progress events over Server-Sent Events. The frontend is React + Tailwind, rendering that stream as a live progress panel with SQL/YAML previews and the final PR link.

We built a full mock DataHub catalog early on - a realistic e-commerce data lake with bronze/silver/gold layers, lineage, glossary terms, and PII tags — so the whole pipeline could be developed and demoed without needing a live DataHub instance. Every DataHub-facing module (read and write) follows the same USE_MOCK_MCP switch, so the exact same code path works against real DataHub once one is available.

Our team split work by module: catalog grounding, LLM generation and validation, GitHub integration, and the DataHub write-back loop that closes the pipeline — registering generated models and lineage back into the catalog.

Challenges we ran into

The trickiest part wasn't generation - it was validation. Making sure every ref() call, every source() call, and every qualified column reference in generated SQL actually existed in the DataHub schema (not just "looked plausible") took real regex and parsing work, plus a repair loop that re-prompts the LLM with specific validation errors instead of just failing.

On the write-back side, we hit a subtle bug where the pipeline skipped registering the model back into DataHub whenever the GitHub PR step was skipped (e.g. no GitHub credentials configured) - an early return was exiting the whole pipeline before write-back ever ran. Catching that during testing, and building in idempotency (so re-running the same request doesn't create duplicate lineage edges) and disk persistence (so the mock catalog survives a server restart), took real end-to-end testing, not just unit tests in isolation.

Accomplishments that we're proud of

  • Zero hallucinated schema references - every model is validated against DataHub before it reaches GitHub, with automatic repair on failure
  • A genuinely closed loop - most demos only read from the catalog; we also write back, registering generated models and lineage into DataHub
  • Idempotent, non-fatal write-back - no duplicate lineage edges on re-runs, and DataHub outages never break the pipeline
  • A fully working mock DataHub environment, so the whole thing demos end-to-end with zero external dependencies
  • Every step streams live to the UI - nothing is a black box

What we learned

Grounding isn't just "add a system prompt with context" - it only works if you validate the output, not just condition the input. And a pipeline that reads from a catalog but never writes back is only half a loop; closing it - so the system's own output becomes part of the catalog for the next request - is what actually makes it feel like a real part of a data platform, not just a code generator.

What's next for DataHub-Grounded dbt Model Generator

  • Validate the write-back path against a real, live DataHub instance
  • Support multi-model generation from a single request
  • Add a "recent runs" history view
  • Smarter, embedding-based glossary and lineage matching
  • A reviewer workflow to approve/reject governance warnings before a PR opens

Built With

Share this project:

Updates