feat(sql-optim-env): add execution-grounded DuckDB SQL query optimization environment#1006
feat(sql-optim-env): add execution-grounded DuckDB SQL query optimization environment#1006OfficialAbhinavSingh wants to merge 3 commits into
Conversation
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Pull request overview
Adds a new community environment, sql_optim_env, where agents rewrite slow SQL and are rewarded based on execution-grounded measurements from an in-memory DuckDB (speedup + result correctness), plus auxiliary rubric components (issue detection, approval, summary, severity). The PR follows the standard OpenEnv env layout (models/client/server), ships Docker + openenv.yaml, and wires the environment into the docs catalog.
Changes:
- Introduces
sql_optim_envserver runtime (tasks, DuckDB executor, graders, FastAPI app) and client/models for typed transport. - Adds pytest coverage for reset/step/state and client parsing helpers.
- Adds documentation page + environments index/toctree entries for the new environment.
Reviewed changes
Copilot reviewed 18 out of 18 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/envs/test_sql_optim_environment.py | Adds unit tests for env reset/step/state and client parsing helpers. |
| envs/sql_optim_env/server/tasks.py | Defines five SQL anti-pattern tasks and metadata used by the environment. |
| envs/sql_optim_env/server/sql_optim_environment.py | Implements the server-side Environment reset/step/state logic. |
| envs/sql_optim_env/server/scoring_models.py | Adds internal Reward model returned by the grader. |
| envs/sql_optim_env/server/graders.py | Implements execution-grounded grading and score breakdown. |
| envs/sql_optim_env/server/executor.py | Implements DuckDB-backed query execution, timing, and correctness checks. |
| envs/sql_optim_env/server/Dockerfile | Provides container build/run path for the environment server. |
| envs/sql_optim_env/server/app.py | Creates the FastAPI app via create_app for WebSocket sessions. |
| envs/sql_optim_env/server/init.py | Exposes SQLOptimEnvironment from the server package. |
| envs/sql_optim_env/README.md | Environment README with quick start, schema, reward breakdown, and usage. |
| envs/sql_optim_env/pyproject.toml | Declares the env package, dependencies (incl. duckdb), and server entrypoint. |
| envs/sql_optim_env/openenv.yaml | Registers the environment spec for deployment/runtime. |
| envs/sql_optim_env/models.py | Defines wire Action/Observation/State models. |
| envs/sql_optim_env/client.py | Adds SQLOptimEnv client with payload/parse helpers. |
| envs/sql_optim_env/init.py | Package exports and top-level docs/example. |
| docs/source/environments/sql_optim.md | Adds docs page for the environment. |
| docs/source/environments.md | Adds sql_optim_env to the environments catalog page. |
| docs/source/_toctree.yml | Adds the environment docs page to the docs toctree. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| # Real execution comparison, carried back to the agent for its next step. | ||
| opt_q = (action.optimized_query or "").strip() | ||
| if opt_q: | ||
| try: | ||
| self._last_execution = get_executor().compare( | ||
| self._task_data["sql_query"], opt_q | ||
| ) | ||
| except Exception: | ||
| self._last_execution = None | ||
|
|
| # Grade (runs DuckDB internally). | ||
| reward = grade(self._task_data, action) | ||
| self._cumulative_reward += reward.score | ||
|
|
||
| # Real execution comparison, carried back to the agent for its next step. | ||
| opt_q = (action.optimized_query or "").strip() | ||
| if opt_q: | ||
| try: | ||
| self._last_execution = get_executor().compare( | ||
| self._task_data["sql_query"], opt_q | ||
| ) |
| BIT_XOR is commutative+associative — order-independent fingerprint. | ||
| Falls back to count-only if the DuckDB version doesn't support the function. | ||
| """ | ||
| # Try BIT_XOR of a numeric hash (portable across DuckDB versions) |
| env.reset( | ||
| task_id="task_2_join_optimization" | ||
| if "task_2_join_optimization" in TASKS | ||
| else DEFAULT_TASK_ID | ||
| ) |
73116e3 to
34c1564
Compare
|
Thanks for the thorough review — all seven were real. Addressed in Execution safety / correctness
Task definitions
Verified: 21 passed (13 existing + 8 new regression tests), |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 18 out of 18 changed files in this pull request and generated 3 comments.
Comments suppressed due to low confidence (1)
envs/sql_optim_env/server/executor.py:230
_checksum()executes queries on the shared DuckDB connection without any locking. Even if_run()is serialized, concurrent calls to_checksum()(from another session) can still race with_run()/compare()and corrupt results or raise driver errors.
Wrap these conn.execute(...) calls with the same lock used elsewhere so the singleton connection is always accessed serially.
try:
wrapped = sql_template.format(query=query)
result = self.conn.execute(wrapped).fetchone()
return result[0], result[1], None
except Exception:
continue
# Final fallback: count only
try:
cnt = self.conn.execute(f"SELECT COUNT(*) FROM ({query}) t").fetchone()[0]
return cnt, None, None
| _FORBIDDEN_KEYWORDS = ( | ||
| "insert", | ||
| "update", | ||
| "delete", | ||
| "drop", | ||
| "create", | ||
| "alter", | ||
| "truncate", | ||
| "replace", | ||
| "attach", |
| if isolate: | ||
| self.conn.execute("BEGIN TRANSACTION") | ||
| try: | ||
| for _ in range(runs): | ||
| try: | ||
| t0 = time.perf_counter() | ||
| rows = self.conn.execute(query).fetchall() | ||
| timings.append((time.perf_counter() - t0) * 1000.0) | ||
| except Exception as exc: | ||
| return 99_999.0, None, str(exc) | ||
| finally: | ||
| if isolate: | ||
| try: | ||
| self.conn.execute("ROLLBACK") | ||
| except Exception: | ||
| pass |
| def _run( | ||
| self, query: str, runs: int = 3, isolate: bool = False | ||
| ) -> Tuple[float, Optional[List], Optional[str]]: |
34c1564 to
55432ec
Compare
|
Addressed the second round (Bugbot) plus @copilot's remaining points in
On @copilot's note about Verified: 23 passed (13 original + 10 regression), |
55432ec to
feeb579
Compare
|
Round 3 addressed in Execution engine
Scoring
Tasks
Verified: 28 passed (13 original + 15 regression), |
|
Addressed in Added a test that valid rewrites using Verified: 28 passed, |
feeb579 to
b1776b1
Compare
…izing rows _run() fetched every result row into Python via fetchall(), three times over for the timing median. For a query like SELECT * FROM events (1M rows) that held the whole result in memory and made the measurement report tuple construction rather than query execution: fetchall() took 783.6ms of which DuckDB execution was 25.5ms, so 97% of the reported time was client-side conversion. Timing now comes from DuckDB's profiler (EXPLAIN ANALYZE -> Total Time), which executes the query for real and discards the result inside the engine, so no rows cross into Python. Rows are materialized only up to the 50k cap compare() already used to pick its precise row-by-row comparison; above that the exact row count comes from an in-engine COUNT(*) and correctness from the existing order-independent checksum. A wall-clock streaming drain is kept as a fallback for a DuckDB version whose profile footer cannot be parsed, so timing degrades instead of breaking. The execution lock becomes reentrant so a whole measurement stays serialized as before. Comparing SELECT * FROM events goes from a 867.7MB peak and 4212ms reported to 14.5MB and 18.5ms; task_5_window_functions compare() drops from 13.6s to 5.7s. Small-result queries pay one extra profiled execution (task_1 101ms -> 233ms); the total across all five tasks is still ~45% lower. Behaviour is otherwise unchanged: writes are still rejected by the read-only connection, row counts stay exact, and results under the cap are still compared row by row.
|
Addressed the remaining
The finding was measurable. For
Worth flagging the tradeoff rather than burying it: small-result queries now pay one extra profiled execution, so Behaviour is otherwise unchanged: writes are still rejected by the read-only connection ( New Env suite: One note in case it comes up in a future review of this file: the over-cap probe leaves a partially-consumed result, and it deliberately isn't closed — in duckdb-python |
…verified match credit
Two High-severity review findings on the execution engine.
The read-only connection stops the agent rewrite from changing the database,
but not from reaching the filesystem. On the previous commit an agent-authored
query could run `COPY ... TO` to write an arbitrary file (verified: the file
was created), `read_csv('/etc/passwd')` to read one, plus `read_text`,
`read_blob`, `read_parquet`, `glob`, `COPY FROM`, `EXPORT DATABASE`, `ATTACH`
and `INSTALL`/`LOAD`. Both connections now open with
`enable_external_access=false`, and with `lock_configuration=true` so a
rewrite cannot `SET` or `PRAGMA` its way back out (DuckDB additionally refuses
to re-enable external access on a running database). The escape hatches are
blocked on the `EXPLAIN ANALYZE` timing path as well.
Second, a rewrite ending in a `--` comment silently bought full correctness
credit. The checksum wraps the query in a subquery, and on a single line the
trailing comment swallowed the closing parenthesis, so the checksum failed and
equal row counts alone were reported as a match: `SELECT id FROM events -- c`
against `SELECT id + 1 AS id FROM events -- c` matched on a million rows that
all differed, unlocking the speedup slice too. `_as_subquery()` now puts the
query on its own line, and an over-cap result with no usable checksum reports
`results_match=False` instead of guessing, as does the surrounding except
path.
Note that the review's stated mechanism for the second finding -- `str.format`
raising on braces in the query -- does not reproduce: `format()` does not
rescan substituted values, and a brace-containing query checksums normally.
The trailing-comment path is the one that was exploitable.
TestSandboxAndCorrectnessCredit covers both; five of its six tests fail on the
previous commit. 38 env tests pass, 1544 repo-wide, ruff clean.
|
Both round-5 findings addressed in 1. Untrusted SQL could reach the filesystem — confirmed, and broader than reported. Across earlier rounds I argued that the read-only connection was "the guarantee" and that no query-level checks were needed. That holds for database writes only. It is not a sandbox. On
So: arbitrary file read and arbitrary file write in the environment server process, driven by the model's own SQL. Both connections now open with 2. Checksum fallback could fake a result match — confirmed, but not by the mechanism described. The reported mechanism was that The conclusion was right anyway, through a different door: a trailing A million rows, every one of them different, awarded full correctness credit — and since correctness gates the speedup slice (added in round 3), that unlocked speedup credit as well. Two changes: Verification. New One inference rather than a verified fact, flagged as such: I have only tested this against duckdb 1.5.5 on Linux. The option names are stable across 1.x as far as I know, but I have not verified the behaviour on another version or platform. |
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 2bf81ee. Configure here.
| self.conn = duckdb.connect(self._path, read_only=True, config=_DUCKDB_CONFIG) | ||
| # Reentrant so `_run` can hold the lock across a whole measurement while | ||
| # the helpers it calls still take it individually. | ||
| self._exec_lock = threading.RLock() |
There was a problem hiding this comment.
Agent SQL can hang shared DB
High Severity
Agent-authored SQL runs on a process-wide singleton connection under a global lock, with no memory_limit or execution timeout in _DUCKDB_CONFIG. A rewrite such as a multi-way cross join can wedge or OOM the whole server and block every other session while _run holds _exec_lock.
Additional Locations (2)
Reviewed by Cursor Bugbot for commit 2bf81ee. Configure here.
Summary
Adds
sql_optim_env, an execution-grounded SQL query optimization environment. The agent receives a slow SQL query plus its schema and returns a rewritten query; reward is not a keyword heuristic — the environment executes both the original and the optimized query against an in-memory DuckDB seeded with synthetic data (10k users, 500k orders, 1k products, 1M events) and scores the agent on measured speedup and result-correctness across five anti-pattern tasks (basic anti-patterns → correlated subqueries → wildcard scans → implicit joins → window-function audits).Fills a gap in the current catalog: there are coding, reasoning, and finance-QA environments, but none that grounds reward in real database execution.
Type of Change
Alignment Checklist
Before submitting, verify:
.claude/docs/PRINCIPLES.mdand this PR aligns with our principles (reward is computed inside the environment from real execution; no domain logic leaks to the trainer).claude/docs/INVARIANTS.mdand no invariants are violated (Gym-likereset/step/state; simulation controls are training-orchestration only, not agent-facing; the client never imports fromserver/)bash .claude/hooks/lint.shand tests and addressed all issuesRFC Status
Test Plan
PYTHONPATH=src:envs uv run pytest tests/envs/test_sql_optim_environment.py— 13 passed (env driven against real DuckDB: reset loads each task, step scores speedup/correctness, state round-trips; model serialization; client parse helpers).SQLOptimEnv(base_url=...).sync()→reset(task_id="task_1_basic_antipatterns")loads the query,step(...)returns a graded reward (0.68),state()round-trips. Round-trip passes.docker build -f envs/sql_optim_env/server/Dockerfile .onghcr.io/huggingface/openenv-base:latest— builds successfully.python scripts/sync_env_docs.py --check,ruff check,ruff format --check— all clean.Reviewers can drive it with:
Claude Code Review
Self-check against the two-tier model: no bugs, uninitialized state, or debug code; reward stays inside the environment (
PRINCIPLES.md), and the Gymreset/step/stateboundary, agent/infra separation, and client↔server import rule (INVARIANTS.md) are all preserved. No alignment flags.Context: I built the original for the Meta PyTorch OpenEnv hackathon and adapted it to the current
envs/conventions (models/client/server,connect4_envshape). Follows on from the sync-bootstrap work in #935 / #959.Note
Medium Risk
Agent-authored SQL runs on the server, so sandbox correctness matters; the PR invests heavily in DuckDB isolation, read-only execution, and regression tests, but it remains a large new execution surface with no changes to OpenEnv core.
Overview
Introduces
sql_optim_env, a new OpenEnv environment for SQL query optimization. Agents get a slow query and schema, submit suggestions plus anoptimized_query, and receive a composite reward computed inside the server from actual DuckDB runs—not static heuristics alone.The server seeds a temp DuckDB with synthetic tables (up to 1M rows), runs original vs. rewrite on a read-only, sandboxed connection (
enable_external_access=false, locked config), measures speed viaEXPLAIN ANALYZE(with bounded materialization for large results), and checks equivalence with row/sorted or SQL checksum paths.grade()blends execution speedup (only when results match), correctness, keyword-based issue detection, approval, summary, and severity; episodes are multi-step withlast_executionfeedback and end at max steps or score ≥ 0.95.Five tasks (
task_1…task_5) progress from sargability/SELECT *through correlated subqueries, wildcard scans, implicit joins, and window-function audits. Packaging follows other envs: Pydantic models,SQLOptimEnvclient, FastAPIcreate_app,openenv.yaml, Dockerfile, and synced docs (environments/sql_optim, catalog card).tests/envs/test_sql_optim_environment.pycovers reset/step, sandbox escapes, grading edge cases, and large-result measurement.Reviewed by Cursor Bugbot for commit 2bf81ee. Bugbot is set up for automated code reviews on this repo. Configure here.