DataHub Architect AI

Autonomous metadata-driven AI agent powered by Google Gemini that inspects DataHub schemas to synthesize dbt models, Airflow DAGs & Data Contracts.

Comment

💡 Inspiration & Problem Statement

Building AI agents that generate data pipelines or SQL queries often fails in real-world production environments. Conventional LLM agents lack real-time awareness of underlying database schemas, column data types, foreign key relationships, PII privacy tags, and governance SLAs. Without complete context, agents hallucinate non-existent columns, produce invalid joins, or violate data privacy contracts.

We built DataHub Architect AI to solve this fundamental gap. By bridging Google Gemini 1.5 Flash with DataHub's Open-Source Context Platform & GMS GraphQL API, our agent inspects catalog metadata before writing code.

⚡ What It Does

DataHub Architect AI is an autonomous AI copilot & interactive studio that:

  1. Queries DataHub GMS & MCP Context: Connects directly to DataHub GMS GraphQL API (/api/graphql) or DataHub Cloud to inspect real-time entity schemas, column data types, primary/foreign keys, ownership, and PII tags across datasets (orders, charges, users, fct_sales). Includes an interactive UI settings modal for custom GMS URLs and access tokens.
  2. Synthesizes 4 Production-Grade Artifacts in 1 Pass:
    • 📜 dbt Model (.sql): Clean transformation code with exact column aliases and joins.
    • 📋 dbt Schema (.yml): Comprehensive documentation with DataHub metadata tags (domain, owner) and quality tests (unique, not_null).
    • 🐍 Airflow DAG (.py): Production Python DAG featuring an automated DataHub metadata lineage emitter task (datahub emit lineage).
    • 🛡️ Data Contract (.yaml): Declarative Data Contract enforcing schema nullability rules, freshness SLAs, and upstream lineage URNs.
  3. Displays Interactive Visual Lineage: Renders a dynamic visual DAG showing upstream raw tables connecting through the AI transformation model down to Gold analytics layer datasets.
  4. Validates Quality SLAs: Runs real-time integrity checks (0 type mismatches, freshness lag enforcement, zero breaking changes).

🛠️ How We Built It

  • Frontend & UI: Built with React 18, Vite 5, and Vanilla CSS with modern Glassmorphism design tokens, interactive GMS connection modal, and dual theme support (Dark/Light mode with localStorage persistence).
  • AI / LLM Core: Integrated with the official @google/generative-ai SDK to interface directly with Google Gemini 1.5 Flash. We configured structured JSON outputs (responseMimeType: "application/json") with a multi-model fallback loop to guarantee 100% parseable artifacts.
  • Context Platform Integration: Native DataHub GraphQL API client (fetchDataHubGraphQL) querying GMS search and lineage APIs, paired with a hybrid Catalog Sandbox fallback for offline evaluation.
  • Deployment: Hosted live on Vercel with automated GitHub CI/CD integration.

🧠 Challenges We Ran Into

  • LLM Output Structuring: Ensuring the model outputs 4 distinct valid code files (SQL, YML, Python, YAML) in a single response without surrounding markdown text required precise JSON schema definition and system prompt engineering.
  • Context Size & Precision: Formatting DataHub catalog metadata so Gemini receives complete schema types without overloading the prompt token window.

🏅 Accomplishments That We're Proud Of

  • Zero Schema Breaking: The agent guarantees 100% schema-valid dbt models and DAGs that respect existing column data types and foreign key constraints.
  • Live DataHub GMS & AI Synthesis: Direct GraphQL integration with DataHub GMS combined with real-time Google Gemini LLM synthesis.
  • Production-Ready Live App: Successfully deployed and running live on Vercel.

📚 What We Learned

  • How combining DataHub's Context Platform with LLMs eliminates hallucination in automated data engineering.
  • Best practices for structuring multi-artifact outputs using Google Gemini JSON response schemas.

🔮 What's Next for DataHub Architect AI

  • Add direct automated PR generation (GitHub/GitLab API) for synthesized dbt models.
  • Support real-time DataHub GraphQL API webhooks for automated contract breach alerts.
  • Expand support for Dagster and Prefect pipeline orchestration targets.

Built With

Share this project:

Updates