Rifqi Hafizuddin Claude Opus 4.8 commited on
Commit
6974c5b
·
1 Parent(s): 4c62146

[NOTICKET][DB] fix(query): R2 - always compile a bounded LIMIT

Browse files

An IR with no limit compiled to a LIMIT-less SELECT, and DbExecutor materialized
the entire result set before capping at 10k rows - an unbounded fetch that can OOM
on a large user table.

- SqlCompiler now bounds every query: an explicit IR limit is honored (clamped to
MAX_RESULT_ROWS=10000); an unbounded query gets LIMIT cap+1 so the executor can
distinguish exactly-the-cap from more-rows-existed and flag truncation.
- CompiledSql.row_cap carries the effective cap; DbExecutor caps + derives truncated
from it and drops its own _ROW_HARD_CAP constant.
- test_sql.py: no-limit golden strings get the safety bound; +3 cases (no-limit
bounded, explicit-limit clamp, row_cap). Restored S608 to tests/** ruff ignore
(documented in PR3-DB but missing from pyproject) - golden tests assert literal SQL.

Also logs R13: a pre-existing, unrelated test_prompt failure found while running R2.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

PROGRESS.md CHANGED
@@ -38,7 +38,7 @@ Verified against code before logging. Severity: **critical** / important / nice-
38
  | # | Fix | Severity | Owner | Status |
39
  |---|---|---|---|---|
40
  | R1 | **AuthN/AuthZ** on data endpoints — derive `user_id` from JWT (`security/auth.py` helpers exist), reject body-supplied IDs. `/chat/stream` has none (`chat.py:40,128`); tenant isolation is client honesty. **Gates the engine-cache work.** | **critical** | DB/B | `[ ]` |
41
- | R2 | **Always compile a LIMIT** — `sql.py:281` emits no LIMIT when `ir.limit is None`; `db.py:93` materializes ALL rows before the 10k cap. Compile `LIMIT min(ir.limit or CAP, CAP)+1`. | **critical** | DB | `[ ]` |
42
  | R3 | **Commit `tests/` + minimal CI** — `tests/` is gitignored; the 200+ tests cited as done exist only on laptops (already caused rename rot). GitHub origin carries tests; HF Space gets the Docker build (already doesn't COPY tests). | **critical (process)** | shared | `[ ]` |
43
  | DB1 | **In-memory `describe_source`** (request-scoped `MemoizingCatalogReader`, `reader.py`) + **LLM-client hoist** (shared module-level `ChatHandler` in `chat.py`). Measured live: `describe_source` 3.5s→~2.0s (structured read now served from the planner's cached snapshot; only the unstructured read remains a round-trip), catalog reads/request ~5→~2. External `query_structured` handshake unchanged (DB2's job) so total slow path is ~flat until DB2. Tests: `tests/catalog/test_reader.py`. | important | agent | `[x]` |
44
  | DB2 | **Keyed engine cache** (LRU, small pool, `pool_recycle`, `pool_pre_ping`) replacing per-call `engine_scope`; read-only + statement_timeout via `connect_args options` (zero SET round-trips, read-only-at-birth); invalidate on client update/delete. Extract to `src/database_client/engine.py`. **Land with/after R1.** | important | DB | `[ ]` |
@@ -52,6 +52,7 @@ Verified against code before logging. Severity: **critical** / important / nice-
52
  | R10 | **Read-only enforcement is session-state, not a server role.** `REPO_CONTEXT.md` counts "read-only DB credentials" as a defense layer but nothing requests/verifies a read-only role. Either request read-only creds at registration (verify via `SELECT current_setting(...)`) or drop the claim. | important | DB | `[ ]` |
53
  | R11 | **De-duplicate** `_PLACEHOLDER_RE` (`task_runner.py:31` vs validator) and `_DATA_ACCESS_TOOLS` (invoker vs planner registry) — import one from the other; comments aren't a sync mechanism. | nice-to-have | agent/tool | `[ ]` |
54
  | R12 | **Doc/process hygiene** — gitignored canonical docs (`AGENT_ARCHITECTURE_CONTEXT_new.md`, `PROJECT_SUMMARY.md`) are cited by code docstrings but absent on disk; `CLAUDE.md` lists deleted modules (enricher, `pipeline/orchestrator.py`); `main` is 38 commits behind on a dead architecture. | nice-to-have | agent | `[ ]` |
 
55
 
56
  **Architecture verdict:** fundamentally sound (catalog-driven IR + deterministic compiler
57
  + static plan is the right call). Debt is transitional duplication (two planners/registries/
 
38
  | # | Fix | Severity | Owner | Status |
39
  |---|---|---|---|---|
40
  | R1 | **AuthN/AuthZ** on data endpoints — derive `user_id` from JWT (`security/auth.py` helpers exist), reject body-supplied IDs. `/chat/stream` has none (`chat.py:40,128`); tenant isolation is client honesty. **Gates the engine-cache work.** | **critical** | DB/B | `[ ]` |
41
+ | R2 | **Always compile a LIMIT** — `sql.py` now emits a bound for every query: explicit limit honored (clamped to `MAX_RESULT_ROWS=10000`), unbounded queries get `LIMIT cap+1` so an unbounded SELECT can't stream a whole table into memory. `CompiledSql.row_cap` carries the cap; `DbExecutor` caps + flags truncation from it (dropped its own `_ROW_HARD_CAP`). Tests updated (`test_sql.py`, +3 cases); `S608` restored to `tests/**` ruff ignore (was dropped). | **critical** | DB | `[x]` |
42
  | R3 | **Commit `tests/` + minimal CI** — `tests/` is gitignored; the 200+ tests cited as done exist only on laptops (already caused rename rot). GitHub origin carries tests; HF Space gets the Docker build (already doesn't COPY tests). | **critical (process)** | shared | `[ ]` |
43
  | DB1 | **In-memory `describe_source`** (request-scoped `MemoizingCatalogReader`, `reader.py`) + **LLM-client hoist** (shared module-level `ChatHandler` in `chat.py`). Measured live: `describe_source` 3.5s→~2.0s (structured read now served from the planner's cached snapshot; only the unstructured read remains a round-trip), catalog reads/request ~5→~2. External `query_structured` handshake unchanged (DB2's job) so total slow path is ~flat until DB2. Tests: `tests/catalog/test_reader.py`. | important | agent | `[x]` |
44
  | DB2 | **Keyed engine cache** (LRU, small pool, `pool_recycle`, `pool_pre_ping`) replacing per-call `engine_scope`; read-only + statement_timeout via `connect_args options` (zero SET round-trips, read-only-at-birth); invalidate on client update/delete. Extract to `src/database_client/engine.py`. **Land with/after R1.** | important | DB | `[ ]` |
 
52
  | R10 | **Read-only enforcement is session-state, not a server role.** `REPO_CONTEXT.md` counts "read-only DB credentials" as a defense layer but nothing requests/verifies a read-only role. Either request read-only creds at registration (verify via `SELECT current_setting(...)`) or drop the claim. | important | DB | `[ ]` |
53
  | R11 | **De-duplicate** `_PLACEHOLDER_RE` (`task_runner.py:31` vs validator) and `_DATA_ACCESS_TOOLS` (invoker vs planner registry) — import one from the other; comments aren't a sync mechanism. | nice-to-have | agent/tool | `[ ]` |
54
  | R12 | **Doc/process hygiene** — gitignored canonical docs (`AGENT_ARCHITECTURE_CONTEXT_new.md`, `PROJECT_SUMMARY.md`) are cited by code docstrings but absent on disk; `CLAUDE.md` lists deleted modules (enricher, `pipeline/orchestrator.py`); `main` is 38 commits behind on a dead architecture. | nice-to-have | agent | `[ ]` |
55
+ | R13 | **Pre-existing test failure** (found during R2, NOT caused by it): `tests/query/planner/test_prompt.py::test_render_catalog_with_sources` fails — `query/planner/prompt.py::render_catalog` now renders stable IDs (`src_test_db`) the test asserts are absent. Old query-planner path; confirmed failing on a clean tree. | nice-to-have | DB | `[ ]` |
56
 
57
  **Architecture verdict:** fundamentally sound (catalog-driven IR + deterministic compiler
58
  + static plan is the right call). Debt is transitional duplication (two planners/registries/
pyproject.toml CHANGED
@@ -120,7 +120,9 @@ ignore = [
120
  ]
121
 
122
  [tool.ruff.lint.per-file-ignores]
123
- "tests/**" = ["S101", "S105", "S106"]
 
 
124
 
125
  [tool.mypy]
126
  python_version = "3.12"
 
120
  ]
121
 
122
  [tool.ruff.lint.per-file-ignores]
123
+ # S608: golden compiler tests assert literal SQL strings (incl. concatenated
124
+ # suffixes) — they never execute against a DB, so it's a false positive here.
125
+ "tests/**" = ["S101", "S105", "S106", "S608"]
126
 
127
  [tool.mypy]
128
  python_version = "3.12"
src/query/compiler/sql.py CHANGED
@@ -30,11 +30,18 @@ from ..ir.models import (
30
  )
31
  from .base import BaseCompiler
32
 
 
 
 
 
 
 
33
 
34
  @dataclass
35
  class CompiledSql:
36
  sql: str
37
  params: dict[str, Any] = field(default_factory=dict)
 
38
 
39
 
40
  class SqlCompilerError(Exception):
@@ -69,14 +76,14 @@ class SqlCompiler(BaseCompiler):
69
  orderby_clause = self._build_orderby(
70
  ir.order_by, table, cols_by_id, select_aliases
71
  )
72
- limit_clause = self._build_limit(ir.limit)
73
 
74
  parts: list[str] = [select_clause, from_clause]
75
  for clause in (where_clause, groupby_clause, orderby_clause, limit_clause):
76
  if clause:
77
  parts.append(clause)
78
 
79
- return CompiledSql(sql=" ".join(parts), params=params)
80
 
81
  # ------------------------------------------------------------------
82
  # Catalog lookup
@@ -277,10 +284,18 @@ class SqlCompiler(BaseCompiler):
277
  parts.append(f"{ref} {ob.dir.upper()}")
278
  return "ORDER BY " + ", ".join(parts)
279
 
280
- def _build_limit(self, limit: int | None) -> str:
 
 
 
 
 
 
 
281
  if limit is None:
282
- return ""
283
- return f"LIMIT {int(limit)}"
 
284
 
285
  # ------------------------------------------------------------------
286
  # Helpers
 
30
  )
31
  from .base import BaseCompiler
32
 
33
+ # Hard ceiling on rows returned to the agent layer. Every compiled query is
34
+ # bounded by this even when the IR sets no limit, so an unbounded SELECT can never
35
+ # stream an entire user table over the wire / into memory. The executor caps to
36
+ # `row_cap` and flags truncation.
37
+ MAX_RESULT_ROWS = 10_000
38
+
39
 
40
  @dataclass
41
  class CompiledSql:
42
  sql: str
43
  params: dict[str, Any] = field(default_factory=dict)
44
+ row_cap: int = MAX_RESULT_ROWS # executor caps rows to this; flags truncation
45
 
46
 
47
  class SqlCompilerError(Exception):
 
76
  orderby_clause = self._build_orderby(
77
  ir.order_by, table, cols_by_id, select_aliases
78
  )
79
+ limit_clause, row_cap = self._build_limit(ir.limit)
80
 
81
  parts: list[str] = [select_clause, from_clause]
82
  for clause in (where_clause, groupby_clause, orderby_clause, limit_clause):
83
  if clause:
84
  parts.append(clause)
85
 
86
+ return CompiledSql(sql=" ".join(parts), params=params, row_cap=row_cap)
87
 
88
  # ------------------------------------------------------------------
89
  # Catalog lookup
 
284
  parts.append(f"{ref} {ob.dir.upper()}")
285
  return "ORDER BY " + ", ".join(parts)
286
 
287
+ def _build_limit(self, limit: int | None) -> tuple[str, int]:
288
+ """Return (LIMIT clause, row_cap).
289
+
290
+ Always bounded. An explicit IR limit is honored exactly (capped at
291
+ MAX_RESULT_ROWS). When the IR has no limit we still emit
292
+ `LIMIT MAX_RESULT_ROWS + 1` — the extra row lets the executor tell
293
+ "exactly the cap" from "more rows existed" and flag truncation.
294
+ """
295
  if limit is None:
296
+ return f"LIMIT {MAX_RESULT_ROWS + 1}", MAX_RESULT_ROWS
297
+ row_cap = min(int(limit), MAX_RESULT_ROWS)
298
+ return f"LIMIT {row_cap}", row_cap
299
 
300
  # ------------------------------------------------------------------
301
  # Helpers
src/query/executor/db.py CHANGED
@@ -40,7 +40,6 @@ from .base import BaseExecutor, QueryResult
40
  logger = get_logger("db_executor")
41
 
42
  _QUERY_TIMEOUT_SECONDS = 30
43
- _ROW_HARD_CAP = 10_000 # belt-and-suspenders cap regardless of LIMIT
44
  _DBCLIENT_PREFIX = "dbclient://"
45
  _POSTGRES_LIKE = frozenset({"postgres", "supabase"})
46
 
@@ -90,8 +89,10 @@ class DbExecutor(BaseExecutor):
90
  timeout=_QUERY_TIMEOUT_SECONDS,
91
  )
92
 
93
- truncated = len(rows) > _ROW_HARD_CAP
94
- capped = rows[:_ROW_HARD_CAP]
 
 
95
  elapsed_ms = int((time.perf_counter() - started) * 1000)
96
  logger.info(
97
  "db query complete",
 
40
  logger = get_logger("db_executor")
41
 
42
  _QUERY_TIMEOUT_SECONDS = 30
 
43
  _DBCLIENT_PREFIX = "dbclient://"
44
  _POSTGRES_LIKE = frozenset({"postgres", "supabase"})
45
 
 
89
  timeout=_QUERY_TIMEOUT_SECONDS,
90
  )
91
 
92
+ # The compiler bounded the SQL to `row_cap` (+1 when the IR was
93
+ # unbounded). More than row_cap rows means the result was truncated.
94
+ truncated = len(rows) > compiled.row_cap
95
+ capped = rows[:compiled.row_cap]
96
  elapsed_ms = int((time.perf_counter() - started) * 1000)
97
  logger.info(
98
  "db query complete",