fix(connection): drain PostgreSQL writes through a cursor instead of buffering - #1328
Conversation
…buffering
The write branch called client.query(), which materialises the entire result
set before the row limit is applied:
const result = await client.query(query, paramValues);
fetchedRows = result.rows; // whole set, then sliced later
So rowLimit bounded what was displayed, not what was resident. One Form
submit against a large table could exhaust the heap shared by every tenant
on the process.
The obvious fix is the dangerous one. readBoundedCursor performs ONE bounded
read and closes the cursor in its finally — right for a SELECT, but
PostgreSQL applies an UPDATE ... RETURNING incrementally, so rows never
pulled are never modified and closing the portal abandons them. Reusing it
here would silently turn "update 1,000 rows" into "update 11" and still
report success.
New drainBoundedCursor reads to exhaustion — every row produced, every side
effect run — while retaining only maxRows. Peak memory becomes the batch
size plus maxRows instead of the whole result.
One subtlety the integration test caught: the completing batch of a
RETURNING statement reports rowCount 0, which overwrote the real count and
made the module report NO_DATA for an UPDATE that changed a thousand rows.
affectedRowCount is now reported only when the statement returned no rows —
exactly the INSERT-without-RETURNING case it exists for.
Three integration tests against a real PostgreSQL, because a stubbed client
cannot prove any of this:
- UPDATE ... RETURNING over 1000 rows with rowLimit 10 returns <= 10 rows
AND leaves all 1000 updated. This is the regression guard: it fails on
any implementation that stops reading early.
- INSERT without RETURNING still reports its affected-row count, not
NO_DATA.
- a write exceeding the limit still flags COMPLETE_TRUNCATED.
postgres-error-handling.test.ts adapted: its failing-INSERT simulation drove
client.query, which the write path no longer uses. The failure now
originates in drainBoundedCursor; the test's subject — a failed query whose
ROLLBACK also fails must still surface the ORIGINAL error — is unchanged.
Verified: 8 pg suites, 72 tests.
Closes #1326
Refs #1298
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
alfredo1996
added
bug
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (4)
💤 Files with no reviewable changes (2)
🚧 Files skipped from review as they are similar to previous changes (2)
WalkthroughPostgreSQL write queries now drain cursor results to completion while retaining bounded returned rows. Affected-row reporting is preserved, and integration and error-handling tests cover write side effects, truncation, non-returning writes, and rollback failures. ChangesPostgreSQL write draining
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant PostgresConnectionModule
participant drainBoundedCursor
participant PostgreSQL
Client->>PostgresConnectionModule: submit write query with rowLimit
PostgresConnectionModule->>drainBoundedCursor: execute query with rowLimit + 1
drainBoundedCursor->>PostgreSQL: read cursor batches until exhaustion
PostgreSQL-->>drainBoundedCursor: returned rows and affected-row data
drainBoundedCursor-->>PostgresConnectionModule: bounded rows and truncation state
PostgresConnectionModule-->>Client: query status and returned rows
Possibly related issues
Possibly related PRs
Suggested labels: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (3)
connection/__tests__/postgresql/postgres-write-drain.ts (1)
112-120: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAssert
COMPLETEpositively rather than the absence ofNO_DATA.
not.toContain(NO_DATA)also passes onERRORor on no status at all. The stated contract for a non-returning INSERT isCOMPLETE.💚 Proposed fix
- expect(statuses).not.toContain(QueryStatus.NO_DATA); + expect(statuses).toContain(QueryStatus.COMPLETE); + expect(statuses).not.toContain(QueryStatus.NO_DATA);🤖 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 `@connection/__tests__/postgresql/postgres-write-drain.ts` around lines 112 - 120, Update the test "still reports the affected-row count for an INSERT without RETURNING" to assert that statuses positively contain QueryStatus.COMPLETE, replacing the weaker NO_DATA absence check while preserving the existing INSERT scenario and runWrite flow.connection/__tests__/postgresql/postgres-error-handling.test.ts (1)
211-218: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer a typed import over inline
require.
requirein a.tstest trips@typescript-eslint/no-var-requiresin most configs and loses typing. Import the mocked module at the top and usejest.mocked.♻️ Proposed refactor
- const { drainBoundedCursor } = require("../../src/postgresql/cursor-read"); - (drainBoundedCursor as jest.Mock).mockRejectedValueOnce( - new Error("insert exploded"), - ); + jest + .mocked(drainBoundedCursor) + .mockRejectedValueOnce(new Error("insert exploded"));Add at the top of the file:
import { drainBoundedCursor } from "../../src/postgresql/cursor-read";🤖 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 `@connection/__tests__/postgresql/postgres-error-handling.test.ts` around lines 211 - 218, Replace the inline require of drainBoundedCursor in the test with a top-level typed import from the cursor-read module, then configure the imported mock via jest.mocked before setting mockRejectedValueOnce. Preserve the existing rejection behavior and original-error assertion.connection/src/postgresql/cursor-read.ts (1)
163-177: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winDrain cost is unbounded by design — consider a larger batch or a drain ceiling.
Correctness requires exhausting the portal, but a 10M-row
UPDATE … RETURNINGnow costs 20k round trips atDRAIN_BATCH_SIZE = 500, all inside the write transaction holding locks. Thestatement_timeoutset by the caller bounds server-side execution per command, not the total client-driven drain wall time. Consider raising the batch size (rows are discarded pastmaxRows, so memory cost is one batch) and/or documenting the interaction with the transaction-level timeout.♻️ Suggested tweak
-/** Rows pulled per round-trip while draining. Bounds memory, not correctness. */ -const DRAIN_BATCH_SIZE = 500; +/** + * Rows pulled per round-trip while draining. Bounds memory, not correctness. + * Larger batches cut round-trips on very large writes; rows beyond `maxRows` + * are discarded immediately, so peak memory is one batch + `maxRows`. + */ +const DRAIN_BATCH_SIZE = 2000;🤖 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 `@connection/src/postgresql/cursor-read.ts` around lines 163 - 177, Reduce the client-driven drain overhead in the cursor-reading flow by increasing DRAIN_BATCH_SIZE while keeping discarded rows bounded to one batch beyond maxRows. Preserve full portal exhaustion and affectedRowCount correctness, and document that draining can extend transaction wall time beyond per-command statement_timeout.
🤖 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.
Inline comments:
In `@connection/__tests__/postgresql/postgres-write-drain.ts`:
- Around line 60-63: Guard the connectionModule cleanup in the afterAll hook so
close() is only called when construction succeeded, matching the existing
optional guard on container.stop(). Preserve the cleanup timeout and ensure
setup failures are not masked by a TypeError.
- Around line 91-97: Update the assertion in the “applies the UPDATE to EVERY
row while returning at most rowLimit” test to require exactly ROW_LIMIT returned
rows, replacing the weaker upper-bound check while preserving the existing
UPDATE and runWrite setup.
- Around line 70-86: Update the runQuery invocation in the postgres write-drain
test to use params instead of parameters and connectionType instead of type,
matching _runSqlQuery and DEFAULT_CONNECTION_CONFIG. Preserve the existing
query, rowLimit, accessMode, callbacks, and status tracking.
- Around line 28-35: Change the container variable declaration in the PostgreSQL
write-drain test to use the StartedPostgreSqlContainer type returned by
PostgreSqlContainer.start(), preserving the existing getUsername(), getPort(),
and stop() usage.
In `@connection/src/postgresql/PostgresConnectionModule.ts`:
- Around line 173-186: Update the preceding write-query block comment to
describe the current drainBoundedCursor behavior rather than claiming writes use
the direct affected-row-count path. In the caller around drainBoundedCursor,
consume batch.truncated for truncation handling instead of recomputing it from
fetchedRows.length, keeping the result consistent with the helper’s interface.
---
Nitpick comments:
In `@connection/__tests__/postgresql/postgres-error-handling.test.ts`:
- Around line 211-218: Replace the inline require of drainBoundedCursor in the
test with a top-level typed import from the cursor-read module, then configure
the imported mock via jest.mocked before setting mockRejectedValueOnce. Preserve
the existing rejection behavior and original-error assertion.
In `@connection/__tests__/postgresql/postgres-write-drain.ts`:
- Around line 112-120: Update the test "still reports the affected-row count for
an INSERT without RETURNING" to assert that statuses positively contain
QueryStatus.COMPLETE, replacing the weaker NO_DATA absence check while
preserving the existing INSERT scenario and runWrite flow.
In `@connection/src/postgresql/cursor-read.ts`:
- Around line 163-177: Reduce the client-driven drain overhead in the
cursor-reading flow by increasing DRAIN_BATCH_SIZE while keeping discarded rows
bounded to one batch beyond maxRows. Preserve full portal exhaustion and
affectedRowCount correctness, and document that draining can extend transaction
wall time beyond per-command statement_timeout.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 0457533c-088b-4a7b-b712-ef156c827b57
📒 Files selected for processing (4)
connection/__tests__/postgresql/postgres-error-handling.test.tsconnection/__tests__/postgresql/postgres-write-drain.tsconnection/src/postgresql/PostgresConnectionModule.tsconnection/src/postgresql/cursor-read.ts
|
@coderabbitai review |
✅ Action performedReview finished.
|
CodeRabbit on #1328, four findings — all correct: - the test passed `parameters: {}` where the module destructures `params`, and `type:` where the config declares `connectionType`. The second is the one that mattered: DEFAULT_CONNECTION_CONFIG.connectionType defaults to NEO4J, so the suite was configuring a Neo4j type against the Postgres module. Behaviour was still right, but the test would have started lying the moment anything branched on it. - `expect(rows.length).toBeLessThanOrEqual(ROW_LIMIT)` also passes when the drain returns ZERO rows — exactly the regression this suite exists to catch. Now asserts the exact bound. - afterAll called connectionModule.close() unguarded, so a failure in beforeAll before construction raised a TypeError that masked the real setup error. `container?.stop()` was already guarded. - the block comment above the branch still claimed writes "keep the direct path", the opposite of what it now does. Rewritten to say what actually differs: both stream, they differ in how they STOP. DrainedCursor.truncated is dropped rather than wired up. The caller derives truncation uniformly from the retained row count, so the field was dead on arrival; shipping an unused one invites a future reader to trust it. Verified: 8 pg suites, 72 tests. Refs #1326 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
All four addressed — thanks, one of these was a real hole in the test rather than a style point.
Unguarded Stale block comment + discarded On 72 tests across 8 Postgres suites still green with the stricter assertion. |
CodeRabbit on #1328: start() resolves to StartedPostgreSqlContainer, so getUsername/getPort/stop were being called on the wrong type. postgres-query.ts has the identical mistake and is left alone here — an unrelated file in a P1 PR. Flagged on the PR instead. Refs #1326 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Fifth finding addressed — Worth noting: The other four from this review are already fixed in the previous commit; those comments are anchored to pre-fix line numbers rather than still being open. |
|
2 participants
Closes #1326 · Refs #1298 — the PostgreSQL half. Neo4j is #1325.
Problem
The write branch buffered everything before the row limit applied:
So
rowLimitbounded what was displayed, not what was resident. One Form submit against a large table could exhaust the heap shared by every tenant on the process.Why the obvious fix would have been worse
readBoundedCursordoes one bounded read and closes the cursor in itsfinally. That is right for aSELECT, and wrong here: PostgreSQL applies anUPDATE … RETURNINGincrementally, so rows never pulled are never modified, and closing the portal abandons them. Swapping it in would have turned "update 1,000 rows" into "update 11" — silently, while still reporting success.drainBoundedCursorreads to exhaustion — every row produced, every side effect run — retaining onlymaxRows. Peak memory becomes the batch size plusmaxRowsrather than the whole result.A bug the integration test caught in my own fix
The completing batch of a
RETURNINGstatement reportsrowCount: 0. My first version let that overwrite the real count, so anUPDATEthat changed a thousand rows reportedNO_DATA.affectedRowCountis now surfaced only when the statement returned no rows — precisely theINSERT-without-RETURNINGcase it exists for.Worth stating plainly: a stubbed client would have happily passed the broken version.
Tests — against a real PostgreSQL
connection/__tests__/postgresql/postgres-write-drain.ts, Testcontainers:UPDATE … RETURNINGover 1,000 rows withrowLimit: 10returns ≤10 rows and leaves all 1,000 updated. This is the regression guard for the trap above — it fails on any implementation that stops reading early.INSERTwithoutRETURNINGstill reports its affected-row count rather thanNO_DATA— the behaviour the original buffered path existed to protect.COMPLETE_TRUNCATED.Note these pass on the old code too: they assert correctness properties the buffered version also satisfied. Their job is to make the fix's failure modes impossible to ship, not to demonstrate the memory bug — which is why the memory claim rests on the implementation change rather than on a flaky heap measurement.
One adapted test
postgres-error-handling.test.tssimulated a failingINSERTthroughclient.query, which the write path no longer uses — so the failure never occurred andonFailwas never called. The rejection now originates indrainBoundedCursor. The test's subject is unchanged: a failed query whoseROLLBACKalso fails must still surface the original error.Verification
8 Postgres suites, 72 tests.
🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes
Tests