Sudo-free local runner for postgres benchmarks + large-v1 dump fix

[ZmeiGorynych]

What & why

Postgres benchmarks (livesqlbench-*, bird-interact-lite-exp) had no turnkey local path: the DB-load logic lives only in the cloud worker, and the system postgres isn't provisioned with the bird_interact role or the task DBs (and no sudo to it). This adds a private, sudo-free workflow so a local bird-interact --dataset <postgres benchmark> just runs — and fixes a benchmark-data staging bug that left two large-v1 DBs with only 1 table each.

Commits

1. Local runner (7d31931)

   • scripts/setup_local_postgres.py — initdb + run a separate PostgreSQL cluster owned by the current user under <main_checkout>/.local_pg/ (gitignored) on a spare port (5544). Creates bird_interact (LOGIN) + root (dump OWNER) roles under loopback trust auth; createdb + psql -f loads each benchmark's pg_dumps/<db>/*.sql. Idempotent (--stop, --recreate).
   • scripts/fetch_local_annotations.py — sync stable *.task.json from GCS into the local annotations store so --require-annotation passes.
   • scripts/run_local_postgres.py — one-shot orchestrator: load auth dotenv → provision+load → annotation sync → exec bird-interact with --data/--db-path auto-derived; everything after -- is passthrough.
   • CLAUDE.md workflow section; .gitignore /.local_pg/.

2. Surface partial dumps (feeaf7c) — the loader captured psql stderr, checks the resulting table count, and WARNs loudly + refuses to mark a DB done on any load error or a zero-table result (instead of silently "succeeding" and failing later as a cryptic per-task dry_run_error).

3. Fix large-v1 dump staging (402d96d) — the large-v1 Google Drive zip ships one dump file per table plus partial *_full.sql aggregates (disaster_relief's largest is 29/49 tables). The old "pick largest file" heuristic grabbed a single big single-table file → a 1-table DB. download_pg_dumps.combine_per_table_dumps now concatenates the single-table files, hoists CREATE TYPE first and defers FOREIGN KEYs last → one clean load. Verified: residential 48/48, disaster 50/50, 0 errors (was 1/48, 1/50).

Tests

• tests/scripts/test_local_postgres_helpers.py — BIRD_PG_* export shape, cluster location, required roles, dotenv parser, marker-skip-on-error.
• tests/scripts/test_download_pg_dumps_combine.py — classifiers vs pg_dump comment banners, type-first/FK-last ordering, single-table selection over aggregates, COPY/meta-command splitting.
• Full non-integration suite: 3862 passed.

Validation

End-to-end local run of livesqlbench-large (opus-4-8, raw, one-shot) works against the private cluster; grading produces correct pass/fail verdicts (validated a passing task too).

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features

  • Added a one-command local PostgreSQL benchmark runner.
  • Added local setup and database-loading support for isolated benchmark runs.
  • Added automatic syncing of required local task annotations.

• Bug Fixes

  • Improved dump handling so larger datasets load more reliably.
  • Excluded local generated PostgreSQL data from version control.

• Documentation

  • Expanded local benchmark instructions and troubleshooting notes.

• Tests

  • Added coverage for dump combining and local Postgres helper behavior.

[coderabbitai]

Warning

Review limit reached

You’ve reached a temporary PR review limit under our Fair Usage Limits Policy.

Your recent review volume is higher than typical usage, so adaptive limits are currently applied.

Next review available in: 21 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews.

How do review limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window.

Please refer docs for additional details.

Review details

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: b4964364-46ab-4fec-b8a3-2ee2b6e4f384

📥 Commits

Reviewing files that changed from the base of the PR and between 402d96d and 7e7e4c2.

📒 Files selected for processing (6)

• CLAUDE.md
• scripts/download_pg_dumps.py
• scripts/fetch_local_annotations.py
• scripts/setup_local_postgres.py
• tests/scripts/test_download_pg_dumps_combine.py
• tests/scripts/test_local_postgres_helpers.py

📝 Walkthrough

Walkthrough

This PR adds tooling to run Postgres-backed benchmarks locally without sudo: a cluster provisioning/loading script, a GCS annotation sync script, an orchestrator that ties them together with bird-interact, improved per-table pg_dump combining logic, documentation, and unit tests.

Changes

Local Postgres Benchmark Workflow

Layer / File(s)	Summary
Per-table dump combining scripts/download_pg_dumps.py, tests/scripts/test_download_pg_dumps_combine.py	Adds statement classification/combination helpers to merge per-table SQL dumps in type-then-table-then-FK order, reworks stage_dumps to group sources by target and combine when multiple exist, and adds tests for ordering and statement parsing.
Cluster provisioning and database loading scripts/setup_local_postgres.py, tests/scripts/test_local_postgres_helpers.py	Adds a CLI to init/start/stop a local Postgres cluster, create required roles, resolve target databases, load benchmark dumps with error-aware marker tracking, write env exports, and unit tests for helpers and marker behavior.
Annotation syncing scripts/fetch_local_annotations.py	Adds a script to resolve benchmark instance-to-database mappings and download missing annotation JSON files from GCS, reporting counts and exit status.
End-to-end orchestrator and docs scripts/run_local_postgres.py, .gitignore, CLAUDE.md	Adds an orchestrator loading auth env, provisioning the cluster, optionally syncing annotations, and invoking bird-interact; ignores the local cluster data directory and documents the workflow.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant User
  participant RunLocalPostgres
  participant SetupLocalPostgres
  participant FetchLocalAnnotations
  participant BirdInteract
  User->>RunLocalPostgres: run_local_postgres.py --benchmark ...
  RunLocalPostgres->>RunLocalPostgres: load_env_file(auth env)
  RunLocalPostgres->>SetupLocalPostgres: ensure_cluster, load_databases
  SetupLocalPostgres-->>RunLocalPostgres: BIRD_PG_* env vars
  RunLocalPostgres->>FetchLocalAnnotations: main(benchmark, instance filter)
  FetchLocalAnnotations-->>RunLocalPostgres: fetch status
  RunLocalPostgres->>BirdInteract: subprocess.call(bird-interact ...)
  BirdInteract-->>User: benchmark results

Loading

Possibly related PRs

• MotleyAI/bird-agents#30: Both PRs modify the staging logic in scripts/download_pg_dumps.py, with this PR extending it to combine per-table dumps.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main changes: a sudo-free local Postgres benchmark runner and the large-v1 dump staging fix.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (6)

scripts/fetch_local_annotations.py (1)

41-42: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Empty instance-id tokens not filtered.

Same pattern as setup_local_postgres.py: a trailing comma in --instance-ids produces an empty-string id, leading to a confusing "not found" error rather than a clear validation message.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@scripts/fetch_local_annotations.py` around lines 41 - 42, The instance ID
parsing in fetch_local_annotations.py does not filter out empty tokens, so a
trailing comma can produce an empty string and a misleading lookup failure.
Update the argument handling around args.instance_ids to strip and discard blank
entries before assigning ids, matching the validation pattern used in
setup_local_postgres.py, and ensure the code paths that use ids receive only
non-empty instance IDs.

scripts/setup_local_postgres.py (3)

319-320: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Empty instance-id tokens not filtered.

args.instance_ids.split(",") with e.g. a trailing comma yields an empty string in the list, which then fails the missing check with a confusing instance_ids not found in ...: [''] error instead of a clear "empty id" message. run_local_postgres.py's equivalent parsing already filters with if s.strip().

🐛 Suggested fix

     ids = ([s.strip() for s in args.instance_ids.split(",")]
-           if args.instance_ids else None)
+           if args.instance_ids else None)
+    if ids:
+        ids = [i for i in ids if i]

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@scripts/setup_local_postgres.py` around lines 319 - 320, The instance ID
parsing in the local Postgres setup still keeps empty tokens from
args.instance_ids, which leads to confusing downstream missing-ID errors. Update
the parsing logic around the ids assignment to filter out blank entries after
splitting, matching the existing behavior used in run_local_postgres.py. Keep
the fix localized to the args.instance_ids handling so the rest of the
validation flow sees only real IDs.

---

259-266: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win

"ERROR:" detection is locale-dependent.

errs = [ln for ln in (r.stderr or "").splitlines() if "ERROR:" in ln] relies on psql emitting the English ERROR: prefix. If the process env has a non-English locale, libpq/psql may localize this prefix (e.g. ERREUR:), causing had_errors to stay False for a genuinely broken load — silently defeating the core hardening this PR adds. Pin the locale for the load subprocess.

♻️ Suggested fix

             r = _psql(bindir, port, p["sock"], "-q", "-o", os.devnull,
                       "-v", "ON_ERROR_STOP=0", "-f", str(sql), db=db,
-                      capture_output=True, text=True, check=False)
+                      capture_output=True, text=True, check=False,
+                      env={**os.environ, "LC_ALL": "C", "LANG": "C"})

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@scripts/setup_local_postgres.py` around lines 259 - 266, The error detection
in the _psql load path is relying on the localized "ERROR:" prefix, so pin the
subprocess locale to a stable English setting when calling _psql from the
setup_local_postgres script. Update the call site that captures r.stderr and
computes errs so psql/libpq emits consistent error text regardless of the host
locale, preserving had_errors for broken loads.

---

208-230: 🔒 Security & Privacy | 🔵 Trivial | 💤 Low value

Raw string interpolation for SQL identifiers/literals.

_db_exists/_table_count build queries via f-strings with unescaped db/role values (flagged by static analysis as SQL-injection-style construction). db currently only comes from trusted local directory names / benchmark metadata, so exploitation risk is minimal, but a stray quote in a db name would still break the query silently.

♻️ Suggested hardening

 def _db_exists(bindir: Path, port: int, db: str) -> bool:
     p = _paths()
     r = _psql(
         bindir, port, p["sock"], "-tAc",
-        f"SELECT 1 FROM pg_database WHERE datname='{db}'",
+        "SELECT 1 FROM pg_database WHERE datname=%s".replace("%s", f"'{db.replace(chr(39), chr(39)*2)}'"),
         capture_output=True, text=True,
     )
     return r.stdout.strip() == "1"

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@scripts/setup_local_postgres.py` around lines 208 - 230, _raw SQL string
interpolation in _db_exists and _table_count should be hardened by avoiding
direct f-string construction for SQL identifiers/literals. Update the
query-building in these helpers to use a safe quoting/escaping approach or
parameterized execution via _psql so db values cannot break the statement. Keep
the changes localized to _db_exists and _table_count and preserve their current
return behavior.

scripts/run_local_postgres.py (1)

47-50: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Personal machine path as default.

_DEFAULT_ENV_FILE hardcodes one developer's home directory layout. Harmless since a missing file is a non-fatal no-op and it's overridable via --env-file/BIRD_ENV_FILE, but a generic default (e.g. <repo>/.env or ~/.bird-interact/.env) would avoid baking a personal path into shared source.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@scripts/run_local_postgres.py` around lines 47 - 50, The default env file
path in the _DEFAULT_ENV_FILE assignment is hardcoded to a personal
Dropbox/PyCharm directory, which should be replaced with a generic shared
default. Update the run_local_postgres.py configuration so the default points to
a repo-relative or user-neutral location, while still honoring BIRD_ENV_FILE and
the --env-file override. Keep the change limited to the _DEFAULT_ENV_FILE
initialization and related path fallback behavior.

tests/scripts/test_local_postgres_helpers.py (1)

50-52: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Test doesn't verify "under main checkout".

The assertion only checks cluster_dir().name == ".local_pg"; it doesn't verify the path is actually rooted at paths.main_checkout_root() as the test name claims.

♻️ Suggested strengthening

 def test_cluster_dir_is_under_main_checkout():
     # Worktree-safe: cluster lives in the main checkout, name is stable.
-    assert setup_local_postgres.cluster_dir().name == ".local_pg"
+    from bird_interact_agents import paths
+    assert setup_local_postgres.cluster_dir() == paths.main_checkout_root() / ".local_pg"

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/scripts/test_local_postgres_helpers.py` around lines 50 - 52, The test
test_cluster_dir_is_under_main_checkout only checks the directory name, so it
does not prove the cluster path is rooted under the main checkout. Strengthen
this assertion by also verifying setup_local_postgres.cluster_dir() resolves
relative to paths.main_checkout_root(), using the existing cluster_dir() and
main_checkout_root() helpers so the test matches its name and intent.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@scripts/fetch_local_annotations.py`:
- Around line 41-42: The instance ID parsing in fetch_local_annotations.py does
not filter out empty tokens, so a trailing comma can produce an empty string and
a misleading lookup failure. Update the argument handling around
args.instance_ids to strip and discard blank entries before assigning ids,
matching the validation pattern used in setup_local_postgres.py, and ensure the
code paths that use ids receive only non-empty instance IDs.

In `@scripts/run_local_postgres.py`:
- Around line 47-50: The default env file path in the _DEFAULT_ENV_FILE
assignment is hardcoded to a personal Dropbox/PyCharm directory, which should be
replaced with a generic shared default. Update the run_local_postgres.py
configuration so the default points to a repo-relative or user-neutral location,
while still honoring BIRD_ENV_FILE and the --env-file override. Keep the change
limited to the _DEFAULT_ENV_FILE initialization and related path fallback
behavior.

In `@scripts/setup_local_postgres.py`:
- Around line 319-320: The instance ID parsing in the local Postgres setup still
keeps empty tokens from args.instance_ids, which leads to confusing downstream
missing-ID errors. Update the parsing logic around the ids assignment to filter
out blank entries after splitting, matching the existing behavior used in
run_local_postgres.py. Keep the fix localized to the args.instance_ids handling
so the rest of the validation flow sees only real IDs.
- Around line 259-266: The error detection in the _psql load path is relying on
the localized "ERROR:" prefix, so pin the subprocess locale to a stable English
setting when calling _psql from the setup_local_postgres script. Update the call
site that captures r.stderr and computes errs so psql/libpq emits consistent
error text regardless of the host locale, preserving had_errors for broken
loads.
- Around line 208-230: _raw SQL string interpolation in _db_exists and
_table_count should be hardened by avoiding direct f-string construction for SQL
identifiers/literals. Update the query-building in these helpers to use a safe
quoting/escaping approach or parameterized execution via _psql so db values
cannot break the statement. Keep the changes localized to _db_exists and
_table_count and preserve their current return behavior.

In `@tests/scripts/test_local_postgres_helpers.py`:
- Around line 50-52: The test test_cluster_dir_is_under_main_checkout only
checks the directory name, so it does not prove the cluster path is rooted under
the main checkout. Strengthen this assertion by also verifying
setup_local_postgres.cluster_dir() resolves relative to
paths.main_checkout_root(), using the existing cluster_dir() and
main_checkout_root() helpers so the test matches its name and intent.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: a4e3edd8-765e-4dc0-8276-0c6e42a707ce

📥 Commits

Reviewing files that changed from the base of the PR and between 51ab9b3 and 402d96d.

📒 Files selected for processing (8)

• .gitignore
• CLAUDE.md
• scripts/download_pg_dumps.py
• scripts/fetch_local_annotations.py
• scripts/run_local_postgres.py
• scripts/setup_local_postgres.py
• tests/scripts/test_download_pg_dumps_combine.py
• tests/scripts/test_local_postgres_helpers.py