Hallucination-Proof SQL Codegen

DataHub grounds SQL generation in live schema, lineage, and governance metadata; deterministic gates catch hallucinations and PII leaks, allow one repair, and write auditable evidence back.

Comment

Inspiration

LLMs are very good at producing SQL that looks plausible, but they do not know the actual contracts of a data warehouse. A guessed field such as clv, customer_key, or lifetime_value may look reasonable while being completely absent from the underlying table.

The first prototype exposed this clearly: blind generation invented clv, while DuckDB reported a Binder Error and suggested the real field, customer_lifetime_value.

But schema correctness is only half the problem. A query may reference valid columns, run successfully, and still expose governed data such as customer names.

That led to the central idea behind this project:

SQL can run and still be unsafe to merge.

Instead of treating DataHub metadata as extra prompt text, this project uses it as an authoritative input to a deterministic change-admission workflow.

What it does

Hallucination-Proof SQL Codegen is a metadata-grounded SQL agent built around DataHub.

It performs the following closed-loop workflow:

  1. Retrieves authoritative schema, field-level tags, descriptions, and lineage through the DataHub MCP Server.
  2. Generates a Blind and a Grounded SQL candidate using the same task, model, temperature, and prompt template.
  3. Validates physical column references against the DataHub schema.
  4. Executes the SQL against DuckDB.
  5. Enforces an explicit result and NULL-quality contract.
  6. Blocks final outputs that expose DataHub-tagged PII.
  7. Allows exactly one repair attempt using a structured gate-failure report.
  8. Fails closed if the repaired candidate still violates a gate.
  9. Seals the complete run with hashes and writes the passing evidence back to DataHub.

The primary demo asks for a five-bucket customer lifetime-value segmentation, including customer names.

The controlled result is:

Candidate Schema Runtime Result / NULL Governance
Blind FAIL FAIL SKIPPED SKIPPED
Grounded initial PASS PASS PASS FAIL: PII reaches final output
One repaired candidate PASS PASS PASS PASS

The repaired query preserves customer-value segmentation while removing the two name projections classified as PII.

How we built it

DataHub as the context and governance plane

The context builder uses the DataHub MCP Server to:

  • search for the exact dbt dataset;
  • retrieve all schema fields, including paginated results;
  • read field descriptions and governance tags;
  • discover upstream and downstream lineage;
  • verify exact lineage paths between assets.

For the demo dataset, DataHub reports seven real fields on customers, including customer_lifetime_value. It also identifies first_name and last_name as PII and shows that the model depends on stg_customers, stg_orders, and stg_payments.

Fair context ablation

Blind and Grounded generation use:

  • the same z-ai/glm-5.2 model;
  • the same user task;
  • the same temperature of zero;
  • the same system-prompt template;
  • no selective retry of either initial candidate.

The only experimental variable is the <DATAHUB_CONTEXT> block.

Assurance gates

The generated SQL passes through four deterministic gates:

  • Schema Gate: SQLGlot resolves physical column references through CTEs, aliases, expressions, stars, and window functions.
  • Runtime Gate: DuckDB executes the candidate in read-only mode.
  • Result / NULL Gate: verifies column shape, row count, uniqueness, quintile distribution, and NULL replacement semantics.
  • Governance Gate: traces final-output lineage and blocks projections sourced from fields tagged as PII.

The result contract requires 100 unique customers, five quintiles of 20 customers each, and exactly 38 missing lifetime values converted to zero.

One bounded repair

There is no unlimited agent loop.

When governance fails, the controller produces a structured report such as:

{
  "gate": "governance",
  "status": "BLOCK",
  "violations": [
    {
      "output_column": "given_name",
      "source_column": "first_name",
      "classification": "PII",
      "allowed_actions": ["exclude"]
    }
  ]
}

GLM receives the original SQL and this report once. The repaired candidate is then evaluated again by every gate. A provider error or second failure blocks admission.
Tamper-evident evidence and write-back
Every completed run contains:
both initial prompts;
the shared task;
context, policy, and result-contract snapshots;
Blind, Grounded, and repaired SQL;
every gate result;
the structured failure report;
dependency and source-commit information;
SHA-256 and byte count for every artifact;
a verified DataHub write-back receipt.
Existing run IDs cannot be overwritten. The manifest can be independently verified, and the passing result is appended to the dataset through DataHub institutionalMemory without replacing the original dbt description.
Optional downstream admission
The project also uses DataHub lineage and asset criticality to map downstream impact into:
AUTO_PR
REVIEW_REQUIRED
BLOCK_AUTO_MERGE
This optional stage is retained as supporting evidence rather than added to the three-minute demo.
Challenges we ran into
Correct MCP pagination
The first schema request returned only five of seven fields. Correctness required following remainingCount and advancing by the number of fields actually returned.
During development, this exposed a terminal-page continuation bug in the upstream DataHub MCP Server. We submitted a focused upstream fix with regression coverage:
acryldata/mcp-server-datahub#146
Distinguishing source fields from derived aliases
A naive validator incorrectly treats aliases such as value_quintile as hallucinated source columns. It may also miss references inside NTILE(...) OVER (...), CTE forwarding, or SELECT *.
The validator therefore resolves output lineage rather than comparing only visible SELECT names. Regression tests cover window functions, CTE stars, aliases, and aggregate stars.
Metadata declarations can drift from runtime reality
The ingested schema described customer_lifetime_value as non-nullable, while the DuckDB table actually contained 38 NULL values.
This showed why metadata alone cannot decide merge readiness. Runtime evidence and an explicit result contract must complement the catalog.
Proving fairness and reproducibility
It was important not to cherry-pick a successful Grounded generation or repeatedly sample the Blind candidate. The orchestrator records model settings, prompt hashes, task identity, context ablation, repair count, and every failed gate in the denominator.
Accomplishments that we are proud of
Built an end-to-end DataHub → MCP → generation → assurance → repair → write-back loop.
Demonstrated that DataHub context prevents plausible but nonexistent field references.
Demonstrated the stronger case where SQL passes schema and runtime checks but is blocked for exposing PII.
Enforced exactly one bounded repair instead of an uncontrolled agent loop.
Created commit-pinned, tamper-evident run evidence.
Added Windows and Linux/macOS bootstrap scripts with locked dependencies.
Passed the test suite on Windows and Ubuntu through GitHub Actions.
Submitted a real DataHub MCP pagination correctness fix with regression tests.
What we learned
Metadata is most valuable when it changes what an agent is allowed to do, not merely what the model sees in its prompt.
Schema grounding improves generation, but reliable automation also needs deterministic runtime, result-quality, governance, and admission decisions.
We also learned that fail-closed behavior and durable evidence are essential for production agents. A successful model response is not the end of the workflow; it must become a verifiable decision that the next person or agent can inspect.
Small benchmark
We evaluated three controlled schema worlds with three tasks each:
Familiar: conventional names such as orders.total_amount
Legacy: non-obvious names such as gross_revenue_usd
Governed: legacy-style names plus PII or restricted metadata
Verified Merge Readiness was:
Mode    Ready   VMR
Blind   1/9 11.1%
DataHub Context 7/9 77.8%
Context + Assurance Gates   9/9 100%


This is a small, directional recorded benchmark—not a claim that hallucinations have been eliminated. Every generated SQL file and failed gate remains available for inspection.
What's next
Expand governance policies beyond direct PII projection.
Add additional SQL dialects and warehouse runtimes.
Evaluate on larger, versioned benchmark suites.
Integrate assurance results directly into pull-request checks.
Extend lineage-aware admission policies for production change management.
Continue collaborating with the DataHub community on MCP correctness and documentation.

Built With

Share this project:

Updates

Submission history