DEV Community

DatanestDigital
DatanestDigital

Posted on

A read-only SQL guard for LLM-generated queries with sqlglot

The fastest way to learn that "the model would never do that" is a false statement is to give an LLM a database connection. A text-to-SQL agent will, sooner or later, emit a DROP, a DELETE without a WHERE, or a cartesian join over your biggest table. Sometimes because a user asked for it, sometimes because a prompt injection hid in a column value, and sometimes just because it felt creative.

The usual first defence is if sql.strip().lower().startswith("select"). It fails in at least three fun ways: SELECT 1; DROP TABLE orders, a CTE that hides an INSERT, and (in DuckDB) SELECT * FROM read_csv('/etc/passwd'), which is technically a SELECT that reads your filesystem.

The robust fix: parse the SQL, then reason about the tree. sqlglot gives you a dialect-aware AST in pure Python. Here is the shape of the guard I use, with a DuckDB dialect.

Step 1: exactly one statement, and its root must be a SELECT

import sqlglot
from sqlglot import exp
from sqlglot.errors import SqlglotError

DIALECT = "duckdb"

class GuardError(ValueError):
    def __init__(self, code: str, message: str) -> None:
        super().__init__(f"{code}: {message}")
        self.code = code   # machine-readable, so you can eval on it

def validate_sql(sql: str, max_rows: int = 1000) -> str:
    if not sql or not sql.strip():
        raise GuardError("empty", "no SQL provided")
    try:
        stmts = [s for s in sqlglot.parse(sql, read=DIALECT) if s is not None]
    except SqlglotError as e:
        raise GuardError("parse_error", str(e).splitlines()[0]) from e
    if len(stmts) != 1:
        raise GuardError("multiple_statements", f"expected 1 statement, got {len(stmts)}")
    tree = stmts[0]
    if not isinstance(tree, (exp.Select, exp.SetOperation, exp.Subquery)):
        raise GuardError("not_select", f"only SELECT is allowed, got {type(tree).__name__}")
    ...
Enter fullscreen mode Exit fullscreen mode

Note the error code. You will want it later: a good guard is also a source of labelled failure data.

Unparseable SQL is rejected too. That is a feature: if sqlglot cannot understand it, you do not want to run it.

Step 2: walk the whole tree, not just the root

A root SELECT is not enough, because data-modifying CTEs and SELECT INTO exist. So we walk every node and refuse anything on a forbidden list:

_FORBIDDEN = tuple(
    getattr(exp, n) for n in (
        "Insert", "Update", "Delete", "Merge", "Create", "Drop", "Alter",
        "TruncateTable", "Command", "Copy", "Set", "Use", "Pragma",
        "Attach", "Detach", "Transaction", "Commit", "Rollback",
        "Grant", "Revoke", "Install", "LoadData", "Export",
    ) if hasattr(exp, n)   # class names differ between sqlglot versions
)

for node in tree.walk():
    if isinstance(node, _FORBIDDEN):
        raise GuardError("forbidden_statement", f"{type(node).__name__} is not allowed")
    if isinstance(node, exp.Into):
        raise GuardError("forbidden_statement", "SELECT INTO is not allowed")
Enter fullscreen mode Exit fullscreen mode

The hasattr filter is not decoration. sqlglot renames node classes across releases, so pin the version you tested against and re-run your guard tests when you bump it.

Step 3: block file and network readers

DuckDB table functions like read_csv, read_parquet and glob are the sneaky ones: pure SELECTs with filesystem side channels. Block them by name, and also block path-like table references such as FROM 'data.csv':

_BLOCKED_FUNCS = {"read_csv", "read_csv_auto", "read_parquet", "read_json",
                  "read_text", "read_blob", "glob", "parquet_scan", "query"}

# inside the tree walk
if isinstance(node, (exp.Anonymous, exp.Func)):
    name = (node.name or node.sql_name() or "").lower()
    if name in _BLOCKED_FUNCS:
        raise GuardError("forbidden_function", f"{name}() is not allowed")
if isinstance(node, exp.Table) and isinstance(node.this, exp.Literal):
    raise GuardError("forbidden_function", "file-path table references are not allowed")
Enter fullscreen mode Exit fullscreen mode

If you target Postgres or Snowflake, your blocklist will differ (pg_read_file, COPY, stages, and so on). The principle holds: enumerate side-effecting entry points for your engine, and prefer an allowlist of tables on top.

Step 4: table allowlist, with CTE awareness

allowed = {t.lower() for t in allowed_tables}
ctes = {c.alias.lower() for c in tree.find_all(exp.CTE)}
for t in tree.find_all(exp.Table):
    if t.name.lower() not in allowed and t.name.lower() not in ctes:
        raise GuardError("table_not_allowed", f"table '{t.name}' is not in the allowlist")
Enter fullscreen mode Exit fullscreen mode

CTE names look like tables to the tree, so exclude them or you will reject every WITH query.

Step 5: a row cap the model cannot raise

Do not append LIMIT to the user's SQL; they may already have one, and LIMIT 10000000 would beat yours. Wrap instead:

def _cap_rows(tree: exp.Expression, max_rows: int) -> str:
    inner = tree.sql(dialect=DIALECT)
    return f"SELECT * FROM ({inner}) AS _capped LIMIT {int(max_rows)}"
Enter fullscreen mode Exit fullscreen mode

The inner query can lower the cap, never raise it. Executing the regenerated SQL (tree.sql()), not the original string, also means you run what you validated.

Step 6: a cost proxy that catches cartesian joins

You cannot always trust the engine's EXPLAIN to be cheap, comparable or available. A deterministic proxy is surprisingly effective: count rows in each referenced table, add them up, and multiply whenever a join has no ON or USING.

def table_refs(sql):
    tree = sqlglot.parse_one(sql, read=DIALECT)
    ctes = {c.alias.lower() for c in tree.find_all(exp.CTE)}
    out = []
    for t in tree.find_all(exp.Table):
        if t.name.lower() in ctes:
            continue
        j = t.parent if isinstance(t.parent, exp.Join) else None
        out.append((t.name, bool(j and not j.args.get("on") and not j.args.get("using"))))
    return out

total = 0.0
for table, cartesian in table_refs(sql):
    n = con.execute(f'SELECT count(*) FROM "{table}"').fetchone()[0]
    total = total * max(n, 1) if cartesian else total + n
if max_cost is not None and total > max_cost:
    raise GuardError("cost_cap", f"estimated cost {total:.0f} exceeds cap {max_cost:.0f}")
Enter fullscreen mode Exit fullscreen mode

It is crude by design. It is not a query planner, it will over-estimate some queries and under-estimate others, but it turns "join everything to everything" into a refusal rather than an outage.

Step 7: measure it, with execution accuracy

A guard stops damage; it does not tell you whether the agent is right. For that, compare result sets, not SQL text. Two very different queries can both be correct:

def result_sets_equal(a, b, ordered=False):
    norm = lambda rows: [tuple(round(v, 6) if isinstance(v, float) else v for v in r) for r in rows]
    na, nb = norm(a), norm(b)
    return na == nb if ordered else sorted(na, key=repr) == sorted(nb, key=repr)
Enter fullscreen mode Exit fullscreen mode

A golden case is just a question plus a reference query (and ordered: true when order matters):

cases:
  - id: customer_count
    question: how many customers do we have
    golden_sql: "SELECT count(*) FROM customers"
  - id: top_products
    question: top 5 products by revenue
    golden_sql: "SELECT p.name, sum(o.amount) FROM orders o JOIN products p ON o.product_id=p.product_id GROUP BY p.name ORDER BY 2 DESC LIMIT 5"
    ordered: true
Enter fullscreen mode Exit fullscreen mode

Run the agent, run the golden SQL through the same guarded warehouse, compare. Guard rejections show up as failures with their code, so you can tell "model wrote DDL" from "model got the join wrong". Floats get rounded before comparison because otherwise you will chase 1e-12 ghosts.

What a guard does not do

Be honest about limits: it is no substitute for a read-only database role. Use both. It does not understand row-level security, it cannot tell that a legitimate-looking query reveals data the user should not see, and a cost proxy is not a planner. Treat the guard as one layer, and add a tiny test file of known-bad queries (; chains, CTE-wrapped inserts, read_csv) that you run on every sqlglot upgrade.

The packaged version

I built this into a Python kit, the Data Analyst Agent Kit: the guard above (with the full forbidden list), a DuckDB demo warehouse, Vega-Lite chart output, a deterministic mock LLM so tests run offline, and the execution-accuracy eval harness. It is one of the Agent Forge Pro kits ($147; the core runtime is $297, the SDR and support kits are $197 each, doc-review is $247, and there are $497 and $597 bundles). MIT licensed for your own use (no resale), with a 30-day money-back guarantee. There is also a free sample if you want to look before buying: free sample.

Store: Agent Forge Pro

If you only want the idea, everything above is enough to build your own in an afternoon. Pin sqlglot, write the bad-query tests first, and keep the read-only DB role.

Top comments (0)