Building AI Agents That Safely Query Enterprise Databases

The first time I let an agent run SQL against a production Postgres replica, it tried to SELECT * from a 400M-row events table to "check the schema." The query planner did what query planners do. The on-call channel lit up. Nothing was corrupted - it was a read replica - but I learned the lesson cheaply: an LLM with a database connection is a junior engineer with no fear and infinite confidence. You have to build the guardrails before you hand over the keys.
This post is the playbook I now use whenever I wire an agent into a real database - Postgres, Snowflake, or an internal API that fronts one. It's the same pattern I've shipped into RAG and tool-use pipelines for B2B SaaS systems, and it's the difference between a demo that wows a stakeholder and a system you can actually leave running.
Start with a read-only role, then take privileges away
The single highest-leverage thing you can do is make the database itself enforce your safety properties. Prompts lie, validators have bugs, models hallucinate tools. The database doesn't.
For Postgres, the role I create for any agent looks like this:
CREATE ROLE agent_ro LOGIN PASSWORD '...';
REVOKE ALL ON SCHEMA public FROM agent_ro;
GRANT USAGE ON SCHEMA analytics TO agent_ro;
GRANT SELECT ON ALL TABLES IN SCHEMA analytics TO agent_ro;
ALTER DEFAULT PRIVILEGES IN SCHEMA analytics
GRANT SELECT ON TABLES TO agent_ro;
ALTER ROLE agent_ro SET statement_timeout = '15s';
ALTER ROLE agent_ro SET idle_in_transaction_session_timeout = '30s';
ALTER ROLE agent_ro SET default_transaction_read_only = on;
ALTER ROLE agent_ro SET work_mem = '32MB';
Four things matter here beyond SELECT-only:
statement_timeout- the agent cannot stall a connection with a runaway scan. 15 seconds is my default for analyst-style agents; 3 seconds for anything user-facing.default_transaction_read_only- belt and suspenders. Even if a function call would mutate, the transaction refuses.- Schema isolation - the agent sees
analytics, notpublic, notauth, notbilling. PII lives behind a view layer or doesn't get exposed at all. - A pgbouncer pool dedicated to the agent, capped at a small number of connections (8-16). If the agent loses its mind, blast radius is bounded.
On Snowflake the equivalent is a dedicated role + warehouse with STATEMENT_TIMEOUT_IN_SECONDS and STATEMENT_QUEUED_TIMEOUT_IN_SECONDS set, plus row access policies on anything sensitive. Use a small warehouse - an agent firing exploratory queries on an X-Large will burn credits faster than you can read the Slack alert.
The principle: assume the agent will try every bad query you can imagine, and a few you can't. Make those queries impossible at the engine level, not at the prompt level.
Schema-aware prompting beats "here is the whole DDL"
The naive approach is to dump the entire schema into the system prompt. This breaks in three ways: you blow the context window on a real warehouse, you confuse the model with 200 tables it doesn't need, and you leak structural information about systems the agent has no business touching.
What works better is a two-stage retrieval pattern over the schema itself:
- Embed every table + column + description into a small vector index (I use pgvector for this; it's the same infra as the rest of the RAG stack). Each chunk is one table: name, columns with types, foreign keys, a one-line business description, and 2-3 sample values per column.
- At query time, retrieve the top 5-8 relevant tables based on the user's question, plus any tables joined to those via foreign keys (one hop). Inject only those into the prompt.
The "one-line business description" matters more than the schema. orders.status with values ('pending','authorized','captured','refunded','void') and a sentence saying "captured means the customer was charged; refunded means money returned; void means authorization released without capture" prevents the agent from writing WHERE status = 'completed' and returning zero rows with full confidence.
I keep these descriptions in a _agent_schema_docs table that data engineers maintain by hand. It's annoying. It's also the difference between an agent that gets the answer right 60% of the time and one that gets it right 92% of the time. Document the columns the agent will actually use, not all of them.
Validate the SQL before it ever hits the database
Even with a read-only role, you want a validation layer between the model's output and the connection. Three checks, in order:
1. Parse, don't regex. Use a real SQL parser - sqlglot in Python is what I reach for. Regex on SQL is a trap; the moment someone's column is named delete_at your "block DELETE" rule misfires.
import sqlglot
from sqlglot import exp
ALLOWED = {exp.Select, exp.With, exp.Union, exp.CTE, exp.Subquery}
BLOCKED = {exp.Insert, exp.Update, exp.Delete, exp.Drop,
exp.Create, exp.Alter, exp.TruncateTable, exp.Command}
def validate(sql: str, dialect="postgres"):
tree = sqlglot.parse_one(sql, dialect=dialect)
for node in tree.walk():
if type(node[0]) in BLOCKED:
raise ValueError(f"Disallowed statement: {type(node[0]).__name__}")
# Single statement only
if len(sqlglot.parse(sql, dialect=dialect)) != 1:
raise ValueError("Multiple statements not allowed")
return tree
2. Enforce a LIMIT. If the model didn't add one, add one - 1000 rows by default, configurable per tool. The agent never needs to see the full table; it needs to answer the question.
3. EXPLAIN before EXECUTE. This is the trick that has saved me the most pain. Run EXPLAIN (FORMAT JSON) first, parse the estimated rows and cost, and refuse to execute anything above a threshold. On Postgres:
plan = cur.execute(f"EXPLAIN (FORMAT JSON) {sql}").fetchone()[0]
total_cost = plan[0]["Plan"]["Total Cost"]
est_rows = plan[0]["Plan"]["Plan Rows"]
if total_cost > 50_000 or est_rows > 1_000_000:
raise QueryTooExpensive(...)
When the agent gets a QueryTooExpensive error back, a decent model will try again with a tighter WHERE clause or a sample. This feedback loop is most of what makes agentic SQL workable - the agent learns the shape of your data through bounded failures.
Tool design: small, typed, and honest about what they return
I've moved away from giving agents a single run_sql tool. It works, but it concentrates too much risk in one place and gives the model too many ways to be clever. A better pattern is a small set of typed tools that each do one thing:
| Tool | Purpose | Returns |
|---|---|---|
search_tables(query) |
Semantic search over schema docs | Top-k table descriptions |
describe_table(name) |
Full schema + sample rows for one table | Columns, types, 5 sample rows |
run_query(sql, limit) |
Execute validated read-only SQL | Rows + row count + truncation flag |
get_distinct_values(table, column) |
Cheap cardinality probe | Distinct values up to 100 |
get_distinct_values is the unsung hero. Half the agent's wrong answers come from guessing enum values. Give it a cheap, safe way to look them up and accuracy jumps.
Each tool returns structured JSON with a truncated boolean and an elapsed_ms field. The agent needs to know when it's only seeing part of the picture - models reason better when the tool surface tells the truth about its limits.
Observability: log the prompt, the plan, the query, the rows hash
You will be debugging agent behavior at 11pm on a Tuesday. Build for that day now.
For every tool call I log, at minimum:
trace_idlinking the whole agent run- The exact prompt + retrieved schema chunks
- The generated SQL, the parsed AST, the validation outcome
- The
EXPLAINplan - Query duration and rows returned
- A hash of the result set (not the rows themselves, for PII reasons) so I can tell when an "identical" query suddenly returns different data
- The model's final natural-language answer
This goes into a agent_audit table plus an OpenTelemetry trace. The audit table is what the security and compliance folks ask for; the traces are what I actually use to debug. When a stakeholder says "the agent told me revenue was X but the dashboard says Y," I need to walk the trace from question -> retrieved tables -> SQL -> rows -> answer in under five minutes. That only works if you instrumented it on day one.
One more thing I've learned the hard way: log token counts and tool-call counts per session, with a hard cap. An agent in a bad reasoning loop can rack up real money in OpenAI or Claude API spend before anyone notices. A cap of, say, 20 tool calls and 100k total tokens per user session catches almost every runaway without affecting legitimate work.
The PII problem nobody wants to talk about
A read-only role doesn't protect you from data exfiltration. An agent that can SELECT email, full_name, ssn FROM users is a leak waiting to happen the moment that result flows back into the LLM context and then into the final answer shown to a user who shouldn't see it.
Two patterns work:
1. Views, not tables. The agent's schema contains views like users_safe where sensitive columns are masked, hashed, or simply absent. The base tables aren't granted at all. This is annoying to maintain but it's the only approach I trust for genuinely sensitive data.
2. Output-side redaction. Before any tool result is fed back to the model, run it through a redactor that masks obvious PII patterns (emails, phone numbers, card-shaped numbers). Combine with column-level tags in your schema docs - if a retrieved column is tagged sensitive: true, the tool wrapper either refuses, aggregates, or masks.
If the use case is "let executives ask questions about revenue," the agent doesn't need row-level customer data at all - it needs aggregates. Design the view layer accordingly. The best safety property is the data isn't reachable.
What I'd do if I were starting tomorrow
If you're standing up an agent against a real database next week, here's the order I'd build in:
- Read replica + dedicated read-only role with timeouts and a small connection pool. Half a day.
- Schema docs table with hand-written business descriptions for the 20 tables that matter. One to two days of work with a data engineer who knows the warehouse.
- Four typed tools (
search_tables,describe_table,get_distinct_values,run_query) withsqlglotvalidation andEXPLAIN-gated execution. Two days. - Audit logging to a dedicated table + OTel traces from day one. Half a day if you already have tracing.
- Eval set of 40-60 real questions with known-correct answers, run nightly. This is what tells you whether your prompt changes are making things better or worse. Without it you're flying blind.
- Tighten: PII views, per-tool rate limits, cost caps, and a kill switch (a feature flag that disables the tool entirely).
Skip step 5 and you will never be able to upgrade the model, change the prompt, or refactor the tools with any confidence. The eval set is the asset that compounds.
The agents that work in production are boring on purpose. They have small tools, narrow permissions, loud telemetry, and a database that says "no" before they can hurt anything. The cleverness goes into the schema docs, the eval set, and the failure feedback - not into trusting the model to behave.
If you're wiring an LLM into your warehouse or internal APIs and want a second pair of eyes on the architecture, get in touch - or browse the rest of the blog for more on RAG, agents, and production LLM systems.
Frequently asked questions
How do I give an LLM agent safe access to a Postgres database?
I create a dedicated read-only role at the database level rather than relying on prompts or validators to enforce safety. The role gets SELECT-only privileges on a specific schema (never public, auth, or billing), a statement_timeout of 15 seconds, default_transaction_read_only set to on, and a capped pgbouncer pool of 8-16 connections. PII stays behind a view layer or isn't exposed at all. The principle is to assume the agent will attempt every bad query imaginable and make those queries impossible at the engine level, because prompts lie but the database doesn't.
Should I put my entire database schema in the agent's system prompt?
No, dumping the full DDL breaks in three ways: it blows the context window on real warehouses, it confuses the model with hundreds of irrelevant tables, and it leaks structural information about systems the agent shouldn't touch. Instead, I embed each table's name, columns, foreign keys, a business description, and sample values into a vector index (pgvector works well), then retrieve only the top 5-8 relevant tables plus one-hop foreign-key neighbors at query time. This pattern keeps prompts focused and dramatically improves accuracy compared to a schema dump.
How do I safely validate SQL generated by an LLM before executing it?
I run three checks in order: parse the SQL with a real parser like sqlglot (never regex, because column names like delete_at will break naive rules), block any node types other than SELECT/WITH/UNION/CTE/Subquery, and reject multi-statement queries. Then I enforce a LIMIT (1000 rows by default) if the model didn't add one. Finally, I run EXPLAIN (FORMAT JSON) first and refuse to execute if estimated cost or row count exceeds a threshold. This catches runaway queries before they ever hit the database engine.
Why do business descriptions of database columns matter for AI agents?
Schema structure alone isn't enough because LLMs guess at enum values and column meanings, then return confidently wrong results. For example, an agent might write WHERE status = 'completed' when the actual values are pending, authorized, captured, refunded, and void, returning zero rows with no error. I maintain a _agent_schema_docs table with one-line business descriptions and 2-3 sample values per column, hand-curated by data engineers. In my experience this single change moves accuracy from roughly 60% to 92% on analyst-style questions.
What database settings prevent an AI agent from running away with resources?
On Postgres I set statement_timeout (15 seconds for analyst agents, 3 seconds for user-facing ones), idle_in_transaction_session_timeout to 30 seconds, work_mem to 32MB, and cap the agent's pgbouncer pool at 8-16 connections to bound blast radius. On Snowflake the equivalent is a dedicated role plus a small warehouse with STATEMENT_TIMEOUT_IN_SECONDS and STATEMENT_QUEUED_TIMEOUT_IN_SECONDS configured, plus row access policies on sensitive tables. Never let an agent fire exploratory queries on an X-Large warehouse - it will burn credits faster than you can read the Slack alert.
Building something hard with AI or automation? I am open to talk.
Get in touch