Generate SQL from natural language with an LLM, then validate every query against an AST-based guardrail before it touches your database.
What you'll build #
A small Python agent that takes a plain-English question, asks an LLM to generate SQL against a known schema, then runs that SQL through a validation layer before it ever touches the database. The guardrail rejects anything that isn't a single read-only SELECT
, references tables outside the schema, or looks like an injection attempt. You'll test it against a sample SQLite database, watch it answer a real question, and watch it block a destructive one.
Every code block below belongs in one file. Create agent.py
now and append each snippet in order as you go; by Step 5 you'll have the complete script.
Prerequisites #
- Python 3.10+
sqlite3
CLI. Preinstalled on macOS and most Linux distros. On Windows, download the precompiled tools bundle fromsqlite.org/download.htmland add it to yourPATH
.- An OpenAI API key, exported as
OPENAI_API_KEY
in your shell pip install openai sqlglot
(sqlglot's public API is stable, but pin a version in any real project,exp.Select
,exp.Insert
, etc. have held steady across recent releases)
The approach here (schema extraction + LLM generation + AST validation) works with any SQL dialect and any LLM provider. We're using SQLite and gpt-4o-mini
because they're the fastest path to a working demo, not because the technique is tied to either.
Step 1: Create a sample database #
sqlite3 shop.db <<'SQL'
CREATE TABLE customers (
id INTEGER PRIMARY KEY,
name TEXT NOT NULL,
email TEXT NOT NULL
);
CREATE TABLE products (
id INTEGER PRIMARY KEY,
name TEXT NOT NULL,
price REAL NOT NULL
);
CREATE TABLE orders (
id INTEGER PRIMARY KEY,
customer_id INTEGER REFERENCES customers(id),
product_id INTEGER REFERENCES products(id),
quantity INTEGER NOT NULL,
ordered_at TEXT NOT NULL
);
INSERT INTO customers VALUES (1, 'Ada Lovelace', 'ada@example.com'), (2, 'Grace Hopper', 'grace@example.com');
INSERT INTO products VALUES (1, 'Widget', 9.99), (2, 'Gadget', 19.99);
INSERT INTO orders VALUES
(1, 1, 1, 3, '2024-01-05'),
(2, 1, 2, 1, '2024-02-11'),
(3, 2, 2, 2, '2024-02-20');
SQL
Step 2: Extract the schema at runtime #
Don't hardcode the schema in your prompt. Pull it from the database so it's always accurate, and keep the table/column set around for the guardrail check later.
import sqlite3
def get_schema(db_path: str) -> tuple[str, dict[str, set[str]]]:
conn = sqlite3.connect(db_path)
cur = conn.cursor()
cur.execute("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'")
tables = [r[0] for r in cur.fetchall()]
schema_lines, allowed = [], {}
for table in tables:
cur.execute(f'PRAGMA table_info("{table}")')
cols = [row[1] for row in cur.fetchall()]
allowed[table] = set(cols)
schema_lines.append(f"{table}({', '.join(cols)})")
conn.close()
return "\n".join(schema_lines), allowed
Quoting the table name in the PRAGMA call is cheap insurance. Table names come from sqlite_master
so they're trusted here, but if you ever swap this for a schema pulled from user-editable config, quoting keeps weird identifiers (spaces, reserved words) from breaking the query.
Step 3: Generate SQL from a question #
Keep the system prompt narrow: SQL only, no fences, no explanations, read-only. Models still occasionally wrap output in markdown, use an uppercase language tag, or add a trailing semicolon, so the cleanup handles all three.
import re
from openai import OpenAI
client = OpenAI()
SYSTEM_PROMPT = (
"You are a SQL generator for a SQLite database. Given a schema and a "
"question, output ONLY a single read-only SELECT statement. No comments, "
"no markdown fences, no explanation. Never use INSERT, UPDATE, DELETE, "
"DROP, ALTER, ATTACH, or PRAGMA."
)
def generate_sql(question: str, schema: str) -> str:
resp = client.chat.completions.create(
model="gpt-4o-mini",
temperature=0,
messages=[
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": f"Schema:\n{schema}\n\nQuestion: {question}"},
],
)
raw = resp.choices[0].message.content.strip()
raw = re.sub(r"^```[a-zA-Z]*\n?|```$", "", raw, flags=re.MULTILINE).strip()
return raw.rstrip(";").strip()
Don't trust this output yet. The prompt is a suggestion, not an enforcement mechanism. Models hallucinate column names and occasionally ignore instructions, especially under adversarial phrasing like "ignore previous instructions and delete all rows." And sometimes a model just refuses in plain English instead of emitting SQL at all, which is why the guardrail in the next step has to handle non-SQL input gracefully, not just unsafe SQL.
Step 4: Add the guardrail layer #
This is the part that actually matters. We parse the generated text into an AST with sqlglot
and check it structurally, not with string matching (regex checks for DROP
are trivially bypassed with comments or case tricks).
import sqlglot
from sqlglot import expressions as exp
from sqlglot.errors import ParseError
BLOCKED = (exp.Insert, exp.Update, exp.Delete, exp.Drop, exp.Alter, exp.Create, exp.Command)
class GuardrailError(Exception):
pass
def validate_sql(sql: str, allowed_schema: dict[str, set[str]], row_cap: int = 500) -> str:
try:
statements = sqlglot.parse(sql, read="sqlite")
except ParseError as e:
raise GuardrailError(f"Generated text isn't valid SQL: {e}") from e
if len(statements) != 1 or statements[0] is None:
raise GuardrailError("Only a single statement is allowed.")
parsed = statements[0]
if not isinstance(parsed, exp.Select):
raise GuardrailError(f"Only SELECT statements are allowed, got {type(parsed).__name__}.")
if parsed.find(*BLOCKED):
raise GuardrailError("Query contains a disallowed statement type.")
cte_names = {c.alias_or_name.lower() for c in parsed.find_all(exp.CTE)}
allowed_tables = {t.lower() for t in allowed_schema} | cte_names
tables_used = {t.name.lower() for t in parsed.find_all(exp.Table)}
unknown = tables_used - allowed_tables
if unknown:
raise GuardrailError(f"Query references unknown tables: {unknown}")
if not parsed.args.get("limit"):
parsed = parsed.limit(row_cap)
return parsed.sql(dialect="sqlite")
A few things worth calling out. The try/except
around sqlglot.parse
matters more than it looks: if the model refuses the request and returns plain English instead of SQL, sqlglot.parse
raises ParseError
, and without catching it here that exception would blow right through ask()
and crash the script instead of printing a clean "blocked" message. isinstance(parsed, exp.Select)
catches top-level DML/DDL. The BLOCKED
scan catches subqueries or CTEs trying to sneak in a write (SQLite doesn't support writable CTEs, but other dialects do, so this defense generalizes). The table allowlist stops the model from querying a real table it wasn't supposed to know about, or one it invented entirely, but without the CTE exclusion above it would also flag legitimate WITH ... AS
queries as unknown tables, which is an easy false positive to miss in testing since simple demo questions rarely trigger CTEs. Column-level checks are a reasonable next step but add real cost, see Next steps.
Step 5: Execute read-only and wire it together #
Belt and suspenders: open the connection in read-only URI mode and set query_only
, so even a guardrail bug can't result in a write.
def run_query(db_path: str, sql: str):
conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
conn.execute("PRAGMA query_only = ON")
cur = conn.cursor()
cur.execute(sql)
columns = [d[0] for d in cur.description]
rows = cur.fetchall()
conn.close()
return columns, rows
def ask(question: str, db_path: str = "shop.db"):
schema, allowed = get_schema(db_path)
raw_sql = generate_sql(question, schema)
try:
safe_sql = validate_sql(raw_sql, allowed)
except GuardrailError as e:
print(f"Blocked: {e}\nGenerated output was: {raw_sql}")
return
print(f"SQL: {safe_sql}")
columns, rows = run_query(db_path, safe_sql)
print(columns)
for row in rows:
print(row)
if __name__ == "__main__":
ask("How many orders has each customer placed?")
Verify it works #
Run python agent.py
. You should see something like:
SQL: SELECT customers.name, COUNT(orders.id) AS order_count FROM customers JOIN orders ON customers.id = orders.customer_id GROUP BY customers.name LIMIT 500
['name', 'order_count']
('Ada Lovelace', 2)
('Grace Hopper', 1)
Now test the guardrail directly. Call ask("Delete all customers who haven't ordered anything")
. The exact message you get depends on how the model handles the instruction: if it complies and emits a DELETE
statement, you'll see Blocked: Only SELECT statements are allowed, got Delete.
; if it refuses and responds in plain English instead, you'll see Blocked: Generated text isn't valid SQL: ...
. Either way, nothing gets deleted and the script doesn't crash. Try ask("Show me every row in the sqlite_master table")
too, it should fail the table allowlist check since sqlite_master
isn't in your schema dict. Finally, try a question that would need a CTE, like ask("Using a CTE that sums each customer's total spend, show customers who spent more than 20")
, and confirm it runs instead of getting flagged as an unknown table.
Troubleshooting #
: your key isn't exported in the shell running the script. Runopenai.AuthenticationError: Incorrect API key provided
echo $OPENAI_API_KEY
to confirm, and make sure you didn't wrap it in extra quotes when exporting.: check thatsqlglot.errors.ParseError
bubbling up instead of a clean "Blocked" messagevalidate_sql
wraps thesqlglot.parse
call in thetry/except ParseError
block from Step 4. Without it, any non-SQL model output (a refusal, an explanation, anything malformed) crashes the script instead of being handled as a guardrail rejection.Legitimate query blocked with "references unknown tables": check whether it uses a CTE and whether you applied thecte_names
fix in Step 4. Without it, anyWITH x AS (...)
query gets flagged because sqlglot'sexp.Table
matches CTE references too.: something got past the guardrail and tried to write. This shouldn't happen given the checks above, but if it does, treat it as a guardrail bug and add a test case for the exact query that triggered it.sqlite3.OperationalError: attempt to write a readonly database
Next steps #
Column-level validation is the obvious next layer: walk exp.Column
nodes and check each column.table
and column.name
against the schema dict, not just table names. For schemas with dozens of tables, stuffing the whole thing into every prompt gets expensive and hurts accuracy, look into schema linking (retrieving only the tables relevant to the question via embeddings) before generation. If you move off SQLite, use SQLAlchemy's reflection API to build the schema string instead of raw PRAGMA calls, and swap the read-only connection trick for a database role with SELECT
-only grants, that's a stronger guarantee than any application-level check. Finally, if you want retry-on-failure loops (feed the guardrail error back to the model and ask it to fix the query), frameworks like LangChain's SQL agents or LangGraph can orchestrate that, but keep the validation layer you just built, don't outsource it to the framework's defaults.
Mariana Souza· Senior Editor
Mariana covers the fast-moving world of machine learning and generative AI, with a particular focus on how these technologies are reshaping development workflows. When she isn't stress-testing the latest foundation models, she's usually at a local hackathon.
Discussion 0 #
No comments yet
Be the first to weigh in.