Skip to content

Commit cc1ebfa

Browse files
committed
feat(nl2sql): EXPLAIN dry-run pre-flight for translate_nl_to_sql
Second of the two remaining security-hardening.md items. Before returning or executing the generated SQL, run a non-executing EXPLAIN (reusing mcpg.query.explain_query, so it's SafeSqlDriver-validated) to catch queries that pass the structural allowlist but reference a non-existent column/table or mistype something: - Planner rejection -> structured error='query invalid: <message>' with the SQL returned for inspection; nothing executed. - Pre-flight statement_timeout -> degrade to 'skip, proceed' (don't block a valid translation). - New explain_preflight tool arg (default on); bypassable per call. Surface snapshot regenerated (translate_nl_to_sql gains explain_preflight). 2 new unit tests (invalid-query surfaced; bypass when disabled) + an EXPLAIN route in the nl2sql test schema so the default-on pre-flight passes. Completes security-hardening.md — both NL->SQL items now shipped. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0122yLZLJ8t4W43sdN6BmTZc
1 parent 69ffe6c commit cc1ebfa

6 files changed

Lines changed: 127 additions & 2 deletions

File tree

CHANGELOG.md

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,13 @@ adheres to [Semantic Versioning](https://semver.org/).
88

99
### Added
1010

11+
- **NL→SQL EXPLAIN dry-run pre-flight.** `translate_nl_to_sql` runs a
12+
non-executing `EXPLAIN` on the generated SQL before returning/executing,
13+
catching queries that pass the structural allowlist but reference a
14+
non-existent column/table. A planner rejection surfaces as
15+
`error="query invalid: …"` (SQL returned, nothing executed); a pre-flight
16+
timeout degrades to "proceed". Toggle via the `explain_preflight` arg
17+
(default on).
1118
- **NL→SQL prompt-injection hardening (boundary defence).**
1219
`translate_nl_to_sql` now fences the user question in
1320
`<user_request>``</user_request>` delimiters and instructs the model

docs/security-hardening.md

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -275,7 +275,15 @@ default).
275275
**Effort:** small (system-prompt edit + refusal-path handling + 3-5
276276
tests). Surfaced 2026-06-30 from a security-feature review.
277277

278-
### ⬜ NL→SQL EXPLAIN dry-run pre-flight
278+
### ✅ NL→SQL EXPLAIN dry-run pre-flight
279+
**Shipped.** `translate_nl_to_sql` now runs a non-executing
280+
`EXPLAIN (FORMAT JSON)` (reusing `mcpg.query.explain_query`, so the SQL is
281+
`SafeSqlDriver`-validated first) before returning or executing. A planner
282+
rejection surfaces as a structured `error="query invalid: <message>"` with
283+
the SQL returned for inspection and nothing executed; a pre-flight
284+
`statement_timeout` degrades to "skip, proceed" rather than blocking a valid
285+
translation. Controlled by the `explain_preflight` tool arg (default on).
286+
279287
**Problem.** Generated SQL is validated structurally (AST allowlist +
280288
single-statement assertion) but not **semantically** — a query that
281289
references a non-existent column/table or has a type mismatch is

src/mcpg/nl2sql.py

Lines changed: 54 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@
4040

4141
from mcpg._vendor.sql import SqlDriver, obfuscate_password
4242
from mcpg.introspection import describe_table, list_foreign_keys, list_tables
43-
from mcpg.query import DEFAULT_MAX_ROWS, QueryError, run_select
43+
from mcpg.query import DEFAULT_MAX_ROWS, QueryError, explain_query, run_select
4444

4545
if TYPE_CHECKING:
4646
from mcpg.config import Settings
@@ -1006,6 +1006,35 @@ def _detect_refusal(sql: str, raw: str) -> str | None:
10061006
return None
10071007

10081008

1009+
async def _explain_preflight(driver: SqlDriver, sql: str) -> str | None:
1010+
"""Non-executing ``EXPLAIN`` of the generated SQL — semantic pre-flight.
1011+
1012+
Returns a ``"query invalid: …"`` message when the planner rejects the
1013+
query (a non-existent column/table or type mismatch passes the structural
1014+
allowlist but fails here), or ``None`` to proceed. ``EXPLAIN`` without
1015+
``ANALYZE`` does not run the query, so this is cheap and side-effect free
1016+
(it reuses :func:`mcpg.query.explain_query`, which validates through
1017+
``SafeSqlDriver`` first). A pre-flight ``statement_timeout`` — possible on
1018+
a query over many partitions or a deep view — degrades to "skip, proceed"
1019+
rather than blocking a valid translation.
1020+
"""
1021+
try:
1022+
_assert_single_statement(sql)
1023+
await explain_query(driver, sql)
1024+
except NL2SQLError as exc:
1025+
return f"query invalid: {exc}"
1026+
except QueryError as exc:
1027+
msg = str(exc)
1028+
lowered = msg.lower()
1029+
if "statement timeout" in lowered or "canceling statement" in lowered:
1030+
logger.warning("NL→SQL EXPLAIN pre-flight timed out; proceeding: %s", obfuscate_password(msg))
1031+
return None
1032+
return f"query invalid: {obfuscate_password(msg)}"
1033+
except Exception as exc: # pragma: no cover - unexpected; degrade to skip
1034+
logger.warning("NL→SQL EXPLAIN pre-flight skipped: %s", obfuscate_password(str(exc)))
1035+
return None
1036+
1037+
10091038
async def translate_nl_to_sql(
10101039
driver: SqlDriver,
10111040
*,
@@ -1014,6 +1043,7 @@ async def translate_nl_to_sql(
10141043
question: str,
10151044
schema: str,
10161045
execute: bool = False,
1046+
explain_preflight: bool = True,
10171047
table_filter: tuple[str, ...] | None = None,
10181048
max_tokens: int = DEFAULT_MAX_TOKENS,
10191049
max_rows: int = DEFAULT_MAX_ROWS,
@@ -1118,6 +1148,15 @@ async def translate_nl_to_sql(
11181148
sql, explanation = _parse_response(raw)
11191149

11201150
refusal_reason = _detect_refusal(sql, raw)
1151+
1152+
# Semantic pre-flight: a non-executing EXPLAIN catches queries that pass
1153+
# the structural allowlist but reference a non-existent column/table, so
1154+
# the caller gets a clean "query invalid" instead of a raw runtime error.
1155+
# Runs before both the generation-only return and the execute path.
1156+
preflight_error = None
1157+
if refusal_reason is None and explain_preflight and sql:
1158+
preflight_error = await _explain_preflight(driver, sql)
1159+
11211160
if refusal_reason is not None:
11221161
# The model declined (embedded instructions / out-of-scope). Surface
11231162
# a structured refusal and never forward the sentinel to execution.
@@ -1134,6 +1173,20 @@ async def translate_nl_to_sql(
11341173
refused=True,
11351174
refusal_reason=refusal_reason,
11361175
)
1176+
elif preflight_error is not None:
1177+
# EXPLAIN rejected the generated SQL — return it with the planner
1178+
# message so the agent can see what's wrong, but don't execute.
1179+
translation = TranslationResult(
1180+
sql=sql,
1181+
explanation=explanation,
1182+
model=model,
1183+
provider=provider.name,
1184+
executed=False,
1185+
rows=[],
1186+
columns=[],
1187+
row_count=0,
1188+
error=preflight_error,
1189+
)
11371190
elif not execute or not sql:
11381191
translation = TranslationResult(
11391192
sql=sql,

src/mcpg/tools.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3161,6 +3161,7 @@ async def translate_nl_to_sql(
31613161
schema: str,
31623162
provider: str | None = None,
31633163
execute: bool = False,
3164+
explain_preflight: bool = True,
31643165
table_filter: list[str] | None = None,
31653166
max_rows: int = query.DEFAULT_MAX_ROWS,
31663167
database: _DatabaseArg = None,
@@ -3179,6 +3180,7 @@ async def translate_nl_to_sql(
31793180
question=question,
31803181
schema=schema,
31813182
execute=execute,
3183+
explain_preflight=explain_preflight,
31823184
table_filter=tuple(table_filter) if table_filter else None,
31833185
max_tokens=settings.nl2sql_max_tokens,
31843186
max_rows=max_rows,

tests/contract/tool_surface.snapshot.json

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8325,6 +8325,11 @@
83258325
"title": "Execute",
83268326
"type": "boolean"
83278327
},
8328+
"explain_preflight": {
8329+
"default": true,
8330+
"title": "Explain Preflight",
8331+
"type": "boolean"
8332+
},
83288333
"max_rows": {
83298334
"default": 1000,
83308335
"title": "Max Rows",

tests/unit/test_nl2sql.py

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,9 @@ def _routes_for_simple_schema() -> dict[str, list[dict[str, object]]]:
9595
],
9696
# list_foreign_keys — empty for the simple schema.
9797
"WHERE con.contype = 'f'": [],
98+
# EXPLAIN pre-flight — return a minimal valid plan so the default-on
99+
# semantic pre-flight passes for generated SELECTs in these tests.
100+
"EXPLAIN (FORMAT JSON)": [{"QUERY PLAN": [{"Plan": {"Node Type": "Result"}}]}],
98101
}
99102

100103

@@ -460,6 +463,53 @@ async def test_translate_nl_to_sql_detects_bare_refusal_sentinel() -> None:
460463
assert result.sql == ""
461464

462465

466+
def _routes_with_failing_explain() -> dict[str, list[dict[str, object]]]:
467+
"""Simple-schema routes but with EXPLAIN returning no plan — simulates a
468+
query that passes the structural allowlist yet the planner rejects."""
469+
routes = _routes_for_simple_schema()
470+
routes["EXPLAIN (FORMAT JSON)"] = [] # explain_query raises "produced no plan"
471+
return routes
472+
473+
474+
async def test_translate_nl_to_sql_preflight_reports_invalid_query() -> None:
475+
# The generated SQL is well-formed structurally but the EXPLAIN pre-flight
476+
# rejects it — surface a structured "query invalid" instead of executing.
477+
provider = _StubProvider(response='{"sql": "SELECT nope FROM public.widget", "explanation": "x"}')
478+
479+
result = await translate_nl_to_sql(
480+
FakeRoutingDriver(_routes_with_failing_explain()), # type: ignore[arg-type]
481+
provider=provider, # type: ignore[arg-type]
482+
model="m",
483+
question="give me the nope column",
484+
schema="public",
485+
execute=True,
486+
)
487+
488+
assert result.executed is False
489+
assert result.error is not None
490+
assert result.error.startswith("query invalid")
491+
assert result.sql == "SELECT nope FROM public.widget" # returned for inspection
492+
493+
494+
async def test_translate_nl_to_sql_preflight_can_be_disabled() -> None:
495+
# explain_preflight=False skips the EXPLAIN round-trip entirely, so even a
496+
# driver whose EXPLAIN would fail yields a plain generation result.
497+
provider = _StubProvider(response='{"sql": "SELECT nope FROM public.widget", "explanation": "x"}')
498+
499+
result = await translate_nl_to_sql(
500+
FakeRoutingDriver(_routes_with_failing_explain()), # type: ignore[arg-type]
501+
provider=provider, # type: ignore[arg-type]
502+
model="m",
503+
question="give me the nope column",
504+
schema="public",
505+
execute=False,
506+
explain_preflight=False,
507+
)
508+
509+
assert result.error is None
510+
assert result.sql == "SELECT nope FROM public.widget"
511+
512+
463513
async def test_translate_nl_to_sql_records_provider_and_model_on_the_result() -> None:
464514
provider = _StubProvider(response='{"sql": "SELECT 1", "explanation": "smoke"}')
465515

0 commit comments

Comments
 (0)