DataHub Guard for DataHub

Inspiration Data platform teams frequently face silent compliance leaks: an upstream developer adds a column containing sensitive customer data (like email_address or credit_card_num) to a raw database, and within days, that column silently flows downstream through dbt models, PySpark pipelines, Feast ML feature stores, and into production ML models.

Catching these leaks manually requires tedious lineage tracing across multiple tools and manual cross-team coordination. We built DataHub Guard to answer a crucial question: What happens when AI agents have complete, bi-directional context over the modern data stack? By using DataHub as our central context engine, we created an autonomous system that detects leaks, calculates impact, generates code fixes, and updates catalog governance end-to-end.

What it does DataHub Guard is an autonomous multi-agent platform engineering system that operates through a four-agent workflow:

PIIAuditor Agent: Performs deterministic schema audits and pattern scanning to identify unflagged sensitive attributes (Email, Phone, Credit Cards, SSN, IP Addresses).

LineageTracer Agent: Queries DataHub's GraphQL and REST Lineage APIs to calculate the total downstream blast radius, mapping affected analytics tables, dbt transformations, ML feature stores, and deployed models.

DataHubWriter Agent: Performs bi-directional metadata synchronization back to DataHub—applying PII-Exposed and Quarantined tags, updating schema descriptions, flagging degraded ML models, and alerting dataset owners on Slack.

RemediationEngine Agent: Generates production-ready dbt SQL (using sha256() hashing and masking via modern SELECT * REPLACE syntax) and PySpark scripts, automatically opening a Pull Request via the GitHub API.

How we built it We built DataHub Guard around a modular Python architecture designed for high reliability and zero binary dependency friction:

Core Engine & APIs: Built in Python 3.10+ using direct REST and GraphQL calls against DataHub's native API surfaces (/graphql and /api/v2).

Agent Architecture: Orchestrated four specialized agents (PIIAuditor, LineageTracer, DataHubWriter, RemediationEngine) with strict typing enforced via mypy and pydantic.

Integrations: Leveraged the GitHub API for automated PR creation, Slack Webhooks for real-time owner notifications, and dbt/PySpark templating engines for remediation code generation.

Open Source: Full codebase released under the Apache License 2.0 with complete unit tests and mock artifacts in /examples.

Challenges we ran into Environment & Runtime Build Friction: When developing on bleeding-edge Python runtimes (such as Python 3.14), heavy native C-extension SDK dependencies can cause compilation blocks. We solved this by implementing a lightweight, zero-dependency DataHub client wrapper communicating directly with DataHub's native GraphQL and REST API surfaces.

Preventing Schema Drift in Remediation Code: Generating raw SQL transformations can easily break downstream queries if schema structures change. We utilized SQL SELECT * REPLACE (...) syntax to preserve non-sensitive column ordering while applying hashing functions strictly to flagged PII fields.

Recursive Lineage Traversal: Downstream data graphs often contain complex branching and circular references. We designed our lineage tracer to recursively traverse DataHub's downstreamLineage GraphQL fields while maintaining a visit-set to avoid infinite loops.

Accomplishments that we're proud of True Bi-Directional Metadata Usage: Rather than just reading metadata from DataHub, our agents write state back to the catalog—tagging entities as Quarantined so other tools and agents immediately inherit the operational context.

Zero-Hallucination Code Generation: Because the RemediationEngine fetches the exact column metadata and data types directly from DataHub before generating dbt/PySpark diffs, the resulting code compiles on the first try.

Production-Grade Codebase: Achieved 100% pass rates across ruff check ., mypy datahub_guard, and pytest with SPDX-compliant headers across all source files.

What we learned We learned that DataHub's context graph is the missing link for autonomous platform engineering. Without reliable metadata regarding schemas, lineage, ownership, and ML entities, AI agents hallucinate schema definitions or fail on basic transformation tasks. When agents can query DataHub's MCP and GraphQL APIs, complex cross-platform governance tasks that usually take hours can be completed deterministically in seconds.

What's next for DataHub Guard for DataHub DataHub Skills & MCP Integration: Expand the agent capabilities into a full DataHub Skill so data engineers can trigger mitigation directly from Slack or conversational interfaces using MCP tools.

Automated Rollbacks & Circuit Breakers: Integrate directly with Airflow and Dagster APIs to automatically pause downstream DAG runs the moment an unencrypted PII column is flagged in DataHub.

Fine-Grained Column-Level Lineage: Extend the lineage tracer to leverage DataHub's fine-grained column lineage to mask only the specific sub-fields propagated into downstream views and feature stores.

Built With

• apache
• cli
• datahub-graphql-api
• datahub-rest-api
• dbt-core
• github-rest-api-(pygithub)
• lineage-aspects
• mypy
• pytest
• python-3.10+
• ruff
• slack-webhooks