Automated PII Classification Agent

PII Classification report
Agent analysis
Conversation with chat agent
Report sent to email
Report summary

Inspiration

PII data being unintentionally exposed is a serious issue faced by many organizations across industries. We wanted to provide a solution that offers guidance to data stewards and owners on tagging and masking strategies, with the potential to extend to perform code generation.

What it does

The PII classification agent helps data stewards identify data columns with potential PII data and suggests tagging and masking strategies. Additionally, there is an autonomous agent which uses DataHub action listeners to capture any changes applied on DataHub, synthesize a PII analysis for upstream and downstream data assets and emails the information to the owner of that data asset.

How we built it

The core is a fully agentic pipeline, not a scripted classifier: Agent is given DataHub's MCP tools directly and decides for itself what to search, which columns to trace, and how many lineage hops to follow. We built two custom tools that DataHub's standard MCP server doesn't provide — check_masking_property, which reads DataHub's real, human-set "Masking" structured property, and get_sample_values, which pulls actual column sample data from DataHub's profiling so masking decisions are grounded in evidence, not guesses.

The two modes share one agent graph. Chat mode is a FastAPI + vanilla JS web UI for ad-hoc questions. Automated mode is a DataHub Actions listener consuming platform-change events over Kafka; when several changes land close together, they're collected into a session-window batch and handed to the agent as one investigation — the agent itself decides, using real lineage-connectivity checks, whether that batch represents one causally-related incident or several unrelated ones, and generates a separate Excel report + email per incident.

We built out a real BFSI pipeline end to end: raw consumer/credit data in MySQL, loaded into Snowflake via a custom script that also emits column-level lineage straight to DataHub, flowing through a semantic transformation layer into a Tableau-facing view — so the agent traces a genuine multi-hop, multi-platform lineage chain instead of a toy example.

Challenges we ran into

Challenge: When multiple changes happened close together, the automated agent sent a separate email report for each one — flooding the owner's inbox instead of giving them one clear picture. Solution: We redesigned it to collect related changes into a short waiting window, then let the agent decide whether they're actually connected. If they are, it sends one combined report instead of many — if they're unrelated, it still keeps them separate so nothing gets mixed together.

Accomplishments that we're proud of

Using the datahub actions listener to trigger an agent to synthesize and email reports to the owner.

What we learned

We learned how powerful DataHub can be as a context layer for an AI agent, not just a place to store metadata. Instead of the AI guessing which columns might be sensitive, it could ask DataHub real questions — what's the lineage, what tags already exist, what do the actual sample values look like — and use those real answers to reason with. That grounding is what made the difference between an AI that sounds confident and one that's actually trustworthy.