DataHub Metadata-Aware Code Gen Agent

An AI agent that reads DataHub metadata to generate production-ready Airflow DAGs, dbt models with automatic PII masking, and quality assertions directly onto Git feature branches.

Comment

Inspiration

Every data engineer knows the pain: a production pipeline crashes at 2 AM because someone changed a column upstream without telling anyone, or sensitive customer data accidentally lands in a reporting table unmasked.

We wanted to fix this by connecting data governance directly to code generation. Instead of using DataHub as just a passive metadata catalog you check once a month, what if an AI agent could read live DataHub context before writing code? That way, every generated pipeline is schema-aware, compliant with security policies, and ready to deploy on day one.

What it does

Our Metadata-Aware Agent acts like an automated data engineer on your team:

  1. Reads Live Metadata: It queries DataHub to fetch table schemas, primary keys, security classifications, and downstream lineage graphs.
  2. Generates Multi-File Bundles: Instead of just outputting a single script, it generates a complete pipeline bundle containing an Airflow orchestration DAG, a dbt SQL transformation model, and dbt data quality assertions.
  3. Auto-Enforces Governance: If DataHub flags a column like customer_email with a PII tag, the agent automatically injects SHA256 cryptographic hashing into the dbt SQL code.
  4. Automates Git Workflows: It spins up a dedicated Git feature branch, writes the generated code, commits the files, and outputs a formatted Pull Request summary showing downstream lineage impact.

How we built it

We built the core automation in Python. The script extracts schema context and lineage URNs from DataHub and feeds that structural metadata to the Google Gemini API with strict output constraints.

Gemini parses the metadata and generates the synchronized multi-file bundle (dags/orders_ingestion_dag.py, models/stg_orders.sql, and models/schema.yml). The Python script then uses Git automation to manage local workspace branches, write the files to disk, and prepare everything for GitHub.

Challenges we faced

Getting an LLM to reliably output multi-file architectures without syntax mistakes or invented imports took careful prompt engineering and retry logic. We also ran into real-world Git edge cases along the way—like configuring environment security so GitHub secret scanning wouldn't block pushes when handling API keys, and keeping local feature branches clean.

Accomplishments that we're proud of

We're super proud that the agent doesn't just print code to a terminal—it builds real, working pipeline bundles. Seeing it automatically catch a PII tag on an email field in DataHub and autonomously write cryptographic hashing directly into the generated dbt SQL model was a huge win.

What we learned

We learned how powerful AI becomes when it's grounded in real-time metadata rather than raw guesses. Combining active metadata platforms like DataHub with generative models makes automated code generation drastically more reliable and production-safe.

What's next for DataHub Metadata-Aware Code Gen Agent

Next, we want to integrate GitHub Actions so that whenever a schema change event is published in DataHub, the agent automatically triggers in the background, writes the migration code, and opens a real Pull Request for human review.

Built With

Share this project:

Updates