Text-to-SQL agent returns wrong numbers or runs runaway queries: schema grounding, read-only roles, statement timeouts and AST validation

#Symptoms

You gave an agent a run_sql tool over a reporting database. Two kinds of complaint come in:

• Wrong answers that look right. "Revenue last quarter" is 3x what finance reports. Active customers include deleted test accounts. A join fans out and counts every order line instead of every order.
• Runaway queries. A question produces a cross join or a scan of a billion-row table without a date filter. The query runs for 40 minutes, holds a snapshot open, and replication lag or vacuum suffers.

Both come from the same root cause: the model is guessing at semantics it was never given, and nothing between the model and the database checks what it sends.

#Quick diagnosis

Start from the trace, not the final answer. For each bad answer, pull the question, the SQL the model generated, the row count and the time taken. If you don't log the SQL text and the database's query ID for every tool call, fix that first.

-- What is the agent running right now, and for how long?
SELECT pid, now() - query_start AS runtime, state, left(query, 200) AS query
FROM pg_stat_activity
WHERE usename = 'agent_reader' AND state <> 'idle'
ORDER BY runtime DESC;

#Fix 1: lock down the database role

Prompt instructions are not a control. Enforce limits in the database, where the model can't get around them:

CREATE ROLE agent_reader LOGIN PASSWORD :'pw' CONNECTION LIMIT 10;
GRANT CONNECT ON DATABASE reporting TO agent_reader;
GRANT USAGE ON SCHEMA analytics TO agent_reader;
GRANT SELECT ON ALL TABLES IN SCHEMA analytics TO agent_reader;   -- curated views only

ALTER ROLE agent_reader SET default_transaction_read_only = on;
ALTER ROLE agent_reader SET statement_timeout = '15s';
ALTER ROLE agent_reader SET idle_in_transaction_session_timeout = '30s';
ALTER ROLE agent_reader SET work_mem = '32MB';
ALTER ROLE agent_reader SET temp_file_limit = '1GB';

• Grant on a schema of curated views, not on the raw tables. pg_read_all_data is convenient, but it gives read access to everything, including tables the agent should never see.
• default_transaction_read_only can be turned off by the session (SET), so it isn't a security boundary on its own. The real boundary is that the role has no write grants. The setting adds a clear error message if a write is ever attempted.
• Use row-level security if the agent answers on behalf of users in different tenants, and set the tenant in the session from your code, never from the model.

Caution

Point the agent at a read replica or a reporting copy, not the primary. Even a read-only query can hurt the primary: long snapshots block vacuum cleanup, and a large sort can fill temp space. statement_timeout limits the damage. It doesn't remove it.

On a hot standby, long queries can also be cancelled by replication conflicts. Return that error to the model as "query cancelled, narrow the date range", not as a generic failure.

#Fix 2: give the model a semantic layer, not a schema dump

Wrong-but-plausible answers come from missing meaning. Replace "here are 300 tables" with a short, curated context:

• Views that encode the business rules: analytics.orders_clean already excludes test accounts, soft-deleted rows and internal tenants. The model can't forget a filter that's built in.
• One row of description per column that matters, with units and grain: net_amount_eur numeric -- order total after discounts, excl. VAT, one row per order.
• Metric definitions as named, tested SQL: "revenue = sum(net_amount_eur) from orders_clean where status = 'completed'". Let the agent call a get_metric(name, period, dimensions) tool for the top metrics, and use free-form SQL only for the long tail.
• Few-shot pairs of question and verified SQL for the joins people get wrong.

Choose which tables to include per question with retrieval over the table and column descriptions, so the context stays small even when the warehouse has thousands of tables.

#Fix 3: validate the SQL before it runs

Parse every statement and reject anything outside a narrow allowlist:

import sqlglot
from sqlglot import exp

ALLOWED_TABLES = {"orders_clean", "customers_clean", "products", "calendar"}
BIG_TABLES = {"orders_clean"}

def validate(sql: str) -> str:
    stmts = sqlglot.parse(sql, read="postgres")
    if len(stmts) != 1 or stmts[0] is None:
        raise ValueError("Send exactly one SQL statement.")
    tree = stmts[0]
    if not isinstance(tree, (exp.Select, exp.Union)):
        raise ValueError("Only SELECT queries are allowed.")
    for node in tree.walk():
        if isinstance(node, (exp.Insert, exp.Update, exp.Delete, exp.Create, exp.Drop, exp.Command)):
            raise ValueError("Only read queries are allowed.")
    cte_names = {c.alias_or_name for c in tree.find_all(exp.CTE)}
    tables = {t.name for t in tree.find_all(exp.Table)} - cte_names
    unknown = tables - ALLOWED_TABLES
    if unknown:
        raise ValueError(f"Unknown tables {sorted(unknown)}. Allowed: {sorted(ALLOWED_TABLES)}")
    if tables & BIG_TABLES and not any(
        "order_date" in c.name for c in tree.find_all(exp.Column)
        if c.find_ancestor(exp.Where)):
        raise ValueError("orders_clean needs a WHERE filter on order_date.")
    if not tree.args.get("limit"):
        tree = tree.limit(1000)
    return tree.sql(dialect="postgres")

Error messages are written for the model to act on. A clear "needs a WHERE filter on order_date" gets a corrected query on the next turn. A bare "invalid query" gets the same SQL again.

Tip

The validator is defence in depth, not the boundary. Parsers miss things: functions with side effects, dialect quirks. The database role in Fix 1 is what actually stops writes. The validator exists to give fast, specific feedback and to stop expensive shapes before they reach the database.

#Fix 4: an EXPLAIN cost gate

Ask the planner before you run the query. It costs milliseconds:

import json

def cost_check(conn, sql, max_cost=5_000_000):
    with conn.cursor() as cur:
        cur.execute("EXPLAIN (FORMAT JSON) " + sql)
        plan = cur.fetchone()[0]
        if isinstance(plan, str):
            plan = json.loads(plan)
        top = plan[0]["Plan"]
        if top["Total Cost"] > max_cost:
            raise ValueError(
                f"Query too expensive (planner cost {top['Total Cost']:.0f}, "
                f"~{top['Plan Rows']:.0f} rows). Add filters or aggregate further.")

Planner cost is in arbitrary units, not milliseconds. Calibrate max_cost by running EXPLAIN on a set of known-good and known-bad agent queries against your data, and recheck it after big data growth or a statistics change. Plain EXPLAIN doesn't execute the query. EXPLAIN ANALYZE does, so never use that in the gate.

#Fix 5: check the answer, not just the query

For questions that matter, cross-check before the agent replies:

• Run a sanity query alongside: the same filters with count(*) and count(DISTINCT <grain key>). If they differ, the join fanned out.
• Compare headline numbers with a metric tool or a known dashboard value when one exists, and flag a mismatch to the user instead of hiding it.
• Show the SQL and the assumptions ("completed orders, EUR, excluding VAT, UTC dates") with the answer, so users can catch a wrong definition.

#Test it

Build a golden set of 50-100 questions with verified SQL results. Score the agent on result equality (same rows and values), not SQL text equality: two different queries can both be right. Run the set on every change to the prompt, model, views or validator, and track:

#Prevent it next time

• Keep business rules in views and metric definitions that the agent consumes, and own them like any other data product.
• Review the slowest and the rejected agent queries every week. They show which views or descriptions are missing.
• Never widen the agent role's grants to fix a single question. Add a curated view instead.