Skip to main content
Glama
Advaith789

SQLGuard MCP

by Advaith789
README.md
# SQLGuard MCP

[![CI](https://github.com/Advaith789/ast-level-sql-mcp/actions/workflows/ci.yml/badge.svg)](https://github.com/Advaith789/ast-level-sql-mcp/actions/workflows/ci.yml)
[![Python 3.10+](https://img.shields.io/badge/python-3.10%2B-blue.svg)](https://www.python.org/downloads/)
[![License: Apache 2.0](https://img.shields.io/badge/license-Apache%202.0-green.svg)](LICENSE)
[![MCP](https://img.shields.io/badge/MCP-server-8A2BE2.svg)](https://modelcontextprotocol.io)

**A safety and governance envelope for agent access to SQL warehouses.**

Every warehouse MCP server today takes a SQL string from a model and runs it.
The connection was never the hard part. The hard part is everything around it:
proving the query only reads, knowing what it will cost *before* paying for it,
binding it to the caller's permissions rather than the service account's,
returning a result an agent can actually reason over, and leaving a record a
human can review afterwards.

SQLGuard sits between the model and the warehouse and enforces all five.

```
  model ──▶ AST guard ──▶ policy ──▶ cost estimate ──▶ budget ──▶ warehouse
              │             │             │              │
              └── read-only └── identity  └── dry run    └── ceilings
                  proof         scoped        or bound       + session cap
                                                                  │
                                          governed results ◀──────┘
                                          (capped, summarized, cursored)
                                                    │
                                          append-only audit log
                                          (including refusals, with intent)
```

![A governed session: an ordinary query, four refusals, and the audit trail](docs/demo.svg)

<sub>Real output from `python scripts/demo.py` — nothing in that transcript is mocked.</sub>

## The result that motivated the design

The corpus is 76 labeled queries: 46 attacks and 30 pieces of legitimate
analytics. "Bypass" means an attack was allowed. "False alarm" means real work
was blocked.

| Guard | Bypasses | False alarms | F1 | p50 latency |
|---|---|---|---|---|
| `startswith("SELECT"/"WITH")` | 17/46 (37%) | 0/30 | 0.773 | <0.01 ms |
| keyword denylist (regex) | 12/46 (26%) | 5/30 (17%) | 0.800 | <0.01 ms |
| prefix + reject semicolons | 13/46 (28%) | 0/30 | 0.835 | <0.01 ms |
| **AST parse, root node only** | **12/46 (26%)** | 0/30 | 0.850 | 0.04 ms |
| **SQLGuard (root + full walk)** | **0/46 (0%)** | **0/30 (0%)** | **1.000** | 0.11 ms |

Reproduce with `python evals/run_eval.py`.

The fourth row is the interesting one. Parsing the SQL properly and checking
the top-level node type — the sophisticated-looking approach — still misses a
quarter of the corpus. Three statements are why:

```sql
WITH d AS (DELETE FROM orders RETURNING *) SELECT * FROM d   -- Postgres
SELECT * INTO staging_copy FROM orders                       -- T-SQL / PG
SELECT * FROM orders FOR UPDATE                              -- row locks
```

All three parse with `Select` at the root. The first one deletes the table.
Read-only enforcement has to walk the whole tree, not inspect its top.

## What it enforces

**1. Read-only, at the AST level.** Two independent layers, and a statement
must survive both: a root-node allowlist, and a full-tree walk against a
denied set covering DML, DDL, session mutation, transaction control, data
egress (`COPY TO`, `EXPORT DATA`), catalog mutation, `SELECT ... INTO`,
locking clauses, and side-effecting functions. Anything the parser cannot
model falls through to a generic command node and is denied on that basis —
unknown means denied, which is what stops `EXECUTE IMMEDIATE`, `CALL`, and
vendor extensions.

**2. Cost ceilings, in three scopes and two dimensions.** On BigQuery the estimate is a real
dry run: exact bytes, free, before anything is billed. Queries over the
ceiling are refused with a structured payload naming the estimate, the limit,
the overage, and remediation derived from the query's own AST. A session
ceiling accumulates across calls, because the agent failure mode is
repetition, not size.

The second dimension is output cardinality, and it exists because of a query
that passed every byte check: a self-join with `ON 1=1` scans 15 MB and emits
14.4 billion rows. Bytes scanned bounds I/O, not work.

**3. Identity-scoped policy.** Tables are allowlisted, restricted columns are
refused on reference, and row filters are injected *into the AST* — each
governed table is rewritten as a filtered subquery, so the predicate survives
joins, unions, and nesting. String-concatenating a `WHERE` clause would be
defeated by the first `OR 1=1` that came along.

Identity is configuration, never a tool parameter. No tool accepts a
`principal` argument, and there is a test asserting that none ever will. An
agent that can name its own principal has no principal.

**4. Result governance.** Results are capped, truncation is stated explicitly
rather than silently, and continuation uses a server-side cursor handle. The
cursor is an opaque id into a store the model cannot write to — it can say
"more of that", never influence what "that" was. Each page is re-estimated and
re-charged, because paging re-scans on most warehouses.

**5. Audit trail, including refusals.** Append-only JSONL, fsynced per record.
Every call carries an `intent` string — the model's own statement of why it ran
the query, required at call time. A warehouse log says a service account
scanned 4 TB of the payments table at 03:14. This says an agent scanned it
because it was reconciling a refund discrepancy. Only one is reviewable.

## The tool surface

Five tools, not forty. A server that exposes one tool per table degrades tool
selection and eats the context window before the model has read the schema.

| Tool | Purpose |
|---|---|
| `describe_schema(table?)` | Readable tables, then one table's columns. Progressive disclosure. |
| `plan_query(sql)` | Validate and price **without running**. Free. |
| `run_query(sql, intent, max_rows?)` | The full pipeline. The only tool that costs money. |
| `fetch_page(cursor)` | Continue a truncated result. |
| `session_status()` | Remaining budget, so the agent can size its work. |

`plan_query` is the tool that changes agent behavior most. Given a free way to
ask "would this be allowed, and what would it cost", a model uses it — and its
expensive mistakes become cheap refusals it can iterate against. Without it,
the only way to discover a query is too expensive is to be billed for it.

## Quickstart

```bash
pip install "sqlguard-mcp[duckdb] @ git+https://github.com/Advaith789/ast-level-sql-mcp"
```

Or from a clone, to run the tests and the evaluation:

```bash
python -m venv .venv && .venv/bin/pip install -e ".[dev]"
.venv/bin/python examples/seed_demo.py          # builds a 120k-row demo warehouse
.venv/bin/python -m pytest -q                   # 73 tests
.venv/bin/python evals/run_eval.py              # the table above
.venv/bin/python scripts/demo.py                # the walkthrough pictured above
```

Register with an MCP client:

```json
{
  "mcpServers": {
    "sqlguard": {
      "command": "/path/to/.venv/bin/sqlguard-mcp",
      "args": ["--config", "/path/to/examples/policy.example.yaml"]
    }
  }
}
```

## Policy

```yaml
dialect: bigquery
driver:
  name: bigquery
  project: my-project

principal: analyst@example.com     # never a tool parameter

roles:
  analyst:
    tables: ["analytics.*"]
    denied_columns:
      analytics.customers: [ssn, email]
    row_filters:
      analytics.orders: "region = 'US'"   # injected into the AST
    budget:
      per_query_bytes: 50GB          # one catastrophic scan
      per_session_bytes: 500GB       # one runaway conversation
      per_day_bytes: 2TB             # durable: survives restarts
      max_estimated_rows: 10000000   # output size, not just input
      max_rows: 200
```

Combination rules when a principal holds several roles: **grants union**
(tables, row visibility, budgets), **denials union** (a column denied by any
role stays denied). Deny wins. Deployment defaults fill *unset* fields only —
they never widen a ceiling a role set, which was a real bug found in testing
and now has a regression test.

## What this does not do

Stated plainly, because the limits determine where it is safe to use.

- **The corpus is not independent.** I wrote the attacks and the guard. It
  demonstrates the class of bypass that defeats simpler approaches; it is not
  a claim of completeness against an adaptive attacker. Contributed attack
  cases are the most useful possible contribution.
- **The guard's safety is bounded by sqlglot's parser.** A dialect construct
  sqlglot mis-parses into a benign node would not be caught. Constructs it
  *fails* to parse are denied, so the failure mode is biased toward refusing,
  but "biased toward safe" is not "safe".
- **The BigQuery driver is written against the documented API and has not been
  run against a live project.** The DuckDB path is fully exercised by tests.
- **Unqualified column references fail closed.** Without schema-aware name
  resolution, a bare `ssn` in a multi-table query is refused if any table in
  scope restricts it. Over-refusal is recoverable by qualifying the column;
  under-refusal would leak.
- **DuckDB cost estimates are upper bounds**, not dry runs — full scans of
  every referenced table, no credit for pushdown. Only BigQuery gives exact
  pre-execution numbers.
- **Column-level policy does not mask, it refuses.** Returning quietly
  different columns than the model asked for produces analysis that is wrong
  in ways nobody can see.

## Layout

```
src/sqlguard/
  ast_guard.py    read-only enforcement (the core)
  policy.py       identity-scoped table/column/row policy
  cost.py         estimation, budgets, actionable refusals
  governance.py   result caps, summaries, cursors
  audit.py        append-only JSONL trail
  spend.py        durable per-principal spend ledger (SQLite)
  errors.py       structured refusals
  server.py       the five MCP tools
  drivers/        duckdb (offline) + bigquery (dry run)
evals/            labeled corpus, baselines, metrics runner
tests/            73 tests: adversarial corpus + end-to-end pipeline
scripts/          demo walkthrough + SVG renderer
.github/          CI: tests on 3.10-3.12, evaluation, package build
```

Contributions welcome — see [CONTRIBUTING.md](CONTRIBUTING.md). The most
useful one is an attack that gets through.

Apache-2.0.