diff --git a/connection/__tests__/postgresql/postgres-error-handling.test.ts b/connection/__tests__/postgresql/postgres-error-handling.test.ts index 17d2ac69..6ca52b9a 100644 --- a/connection/__tests__/postgresql/postgres-error-handling.test.ts +++ b/connection/__tests__/postgresql/postgres-error-handling.test.ts @@ -17,6 +17,14 @@ import { // (which can't back a real pg-cursor) still exercises the transaction logic. jest.mock("../../src/postgresql/cursor-read", () => ({ readBoundedCursor: jest.fn().mockResolvedValue({ rows: [], fields: [] }), + // Writes drain rather than stopping early (#1298); the write path calls + // this one, so the double has to provide it or every write test fails on + // "drainBoundedCursor is not a function" rather than on its own assertion. + drainBoundedCursor: jest.fn().mockResolvedValue({ + rows: [], + fields: [], + affectedRowCount: 0, + }), })); function makeModule(): PostgresConnectionModule { @@ -199,6 +207,15 @@ describe("PostgresConnectionModule — error-path routing", () => { jest.spyOn(mod.authModule, "getPool").mockReturnValue(pool as any); const errSpy = jest.spyOn(console, "error").mockImplementation(() => {}); + // Writes now stream through drainBoundedCursor rather than client.query + // (#1298), so that is where the statement failure originates. The test's + // subject is unchanged: a failing query whose ROLLBACK also fails must + // still surface the ORIGINAL error through onFail. + const { drainBoundedCursor } = require("../../src/postgresql/cursor-read"); + (drainBoundedCursor as jest.Mock).mockRejectedValueOnce( + new Error("insert exploded"), + ); + const onFail = jest.fn(); await mod.runQuery( { query: "INSERT INTO t VALUES (1)", params: {} }, diff --git a/connection/__tests__/postgresql/postgres-write-drain.ts b/connection/__tests__/postgresql/postgres-write-drain.ts new file mode 100644 index 00000000..fa73df3b --- /dev/null +++ b/connection/__tests__/postgresql/postgres-write-drain.ts @@ -0,0 +1,147 @@ +import { PostgresConnectionModule } from "../../src/postgresql/PostgresConnectionModule"; +import { + DEFAULT_CONNECTION_CONFIG, + QueryStatus, + AuthType, + ConnectionTypes, +} from "@neoboard/connector-sdk"; +import { + PostgreSqlContainer, + type StartedPostgreSqlContainer, +} from "@testcontainers/postgresql"; + +/** + * Write-path row limiting against a REAL PostgreSQL (#1298 / #1326). + * + * The write branch buffered the entire result set and sliced afterwards, 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. + * + * The obvious fix is the dangerous one. `readBoundedCursor` performs ONE + * bounded `cursor.read()` and closes the cursor in its `finally`. PostgreSQL + * executes a portal incrementally, so an `UPDATE … RETURNING` suspended after + * `rowLimit + 1` rows has NOT applied the rest, and closing the portal + * abandons that work — turning "update 1,000 rows" into "update 11" while + * still reporting success. + * + * These tests exist to make that failure impossible to ship. A stubbed client + * cannot prove any of it; only a real database can. + */ +describe("PostgreSQL write path — row limit must not truncate side effects", () => { + let container: StartedPostgreSqlContainer; + let connectionModule: PostgresConnectionModule; + + const ROW_COUNT = 1000; + const ROW_LIMIT = 10; + + beforeAll(async () => { + container = await new PostgreSqlContainer("postgres:16-alpine").start(); + + connectionModule = new PostgresConnectionModule({ + username: container.getUsername(), + password: container.getPassword(), + authType: AuthType.NATIVE, + uri: `postgresql://${container.getHost()}:${container.getPort()}/${container.getDatabase()}`, + }); + + expect(await connectionModule.authModule.verifyAuthentication()).toBe(true); + + const client = await connectionModule.getPool()!.connect(); + try { + await client.query( + `CREATE TABLE counters (id SERIAL PRIMARY KEY, n INT NOT NULL)`, + ); + await client.query( + `INSERT INTO counters (n) SELECT 0 FROM generate_series(1, $1)`, + [ROW_COUNT], + ); + } finally { + client.release(); + } + }, 120_000); + + afterAll(async () => { + // Guarded: if beforeAll throws before construction, an unguarded close() + // raises a TypeError that masks the real setup failure. + await connectionModule?.close(); + await container?.stop(); + }, 60_000); + + /** Run a write query through the module and collect what the caller sees. */ + function runWrite(query: string) { + return new Promise<{ rows: unknown[]; statuses: QueryStatus[] }>( + (resolve, reject) => { + const statuses: QueryStatus[] = []; + connectionModule.runQuery( + { query, params: {} }, + { + onSuccess: (rows: unknown) => + resolve({ rows: rows as unknown[], statuses }), + onFail: reject, + setStatus: (s: QueryStatus) => statuses.push(s), + setFields: () => {}, + setSchema: () => {}, + }, + { + ...DEFAULT_CONNECTION_CONFIG, + connectionType: ConnectionTypes.POSTGRESQL, + rowLimit: ROW_LIMIT, + accessMode: "WRITE", + }, + ); + }, + ); + } + + it("applies the UPDATE to EVERY row while returning at most rowLimit", async () => { + const { rows } = await runWrite( + `UPDATE counters SET n = n + 1 RETURNING *`, + ); + + // Exactly the cap, not merely "no more than" — `<=` would also pass if + // the drain returned nothing, which is the regression this suite exists + // to catch. + expect(rows).toHaveLength(ROW_LIMIT); + + // ...but every row must have been updated. This is the assertion that + // fails on any implementation which stops reading the portal early. + const client = await connectionModule.getPool()!.connect(); + try { + const { rows: check } = await client.query( + `SELECT count(*)::int AS updated FROM counters WHERE n = 1`, + ); + expect(check[0].updated).toBe(ROW_COUNT); + } finally { + client.release(); + } + }, 60_000); + + it("still reports the affected-row count for an INSERT without RETURNING", async () => { + // The buffered path was kept originally because result.rowCount is what + // makes a non-returning write report COMPLETE rather than NO_DATA. Any + // cursor-based rewrite has to preserve that. + const { statuses } = await runWrite( + `INSERT INTO counters (n) SELECT 99 FROM generate_series(1, 5)`, + ); + + expect(statuses).not.toContain(QueryStatus.NO_DATA); + + const client = await connectionModule.getPool()!.connect(); + try { + const { rows: check } = await client.query( + `SELECT count(*)::int AS inserted FROM counters WHERE n = 99`, + ); + expect(check[0].inserted).toBe(5); + } finally { + client.release(); + } + }, 60_000); + + it("flags truncation when a write returns more rows than the limit", async () => { + const { statuses } = await runWrite( + `UPDATE counters SET n = n WHERE n <> 99 RETURNING *`, + ); + + expect(statuses).toContain(QueryStatus.COMPLETE_TRUNCATED); + }, 60_000); +}); diff --git a/connection/src/postgresql/PostgresConnectionModule.ts b/connection/src/postgresql/PostgresConnectionModule.ts index a2d65693..76b82f4a 100644 --- a/connection/src/postgresql/PostgresConnectionModule.ts +++ b/connection/src/postgresql/PostgresConnectionModule.ts @@ -11,7 +11,7 @@ import { } from "@neoboard/connector-sdk"; import { PostgresRecordParser } from "./PostgresRecordParser"; import { Pool, PoolClient, FieldDef } from "pg"; -import { readBoundedCursor } from "./cursor-read"; +import { readBoundedCursor, drainBoundedCursor } from "./cursor-read"; import { extractTableSchemaFromFields, isAuthenticationError } from "./utils"; import { determineQueryStatus } from "@neoboard/connector-sdk"; import { wrapError, ConnectorErrorType } from "@neoboard/connector-sdk"; @@ -150,12 +150,15 @@ export class PostgresConnectionModule extends ConnectionModule { .map((k) => params[k]) : []; - // Fetch rows. READ queries stream through a server-side cursor so a - // huge result set never buffers in memory — we pull at most rowLimit + 1 - // rows (the MAX_ROWS+1 truncation probe). WRITE queries (Form widgets) - // keep the direct path: their result sets are small and we need the - // driver's affected-row count so an INSERT without RETURNING still - // reports COMPLETE rather than NO_DATA. + // Both paths stream through a server-side cursor so a huge result set + // never buffers in memory; each pulls at most rowLimit + 1 rows for the + // MAX_ROWS+1 truncation probe. They differ in how they STOP: + // READ stops as soon as truncation is known, releasing the portal. + // WRITE drains to exhaustion, because PostgreSQL applies an + // UPDATE ... RETURNING incrementally — rows never pulled are + // never modified — and then reports the driver's affected-row + // count so an INSERT without RETURNING still reads as COMPLETE + // rather than NO_DATA (#1298, #1326). let fetchedRows: Record[]; let fields: FieldDef[] | undefined; let affectedRowCount: number | undefined; @@ -170,10 +173,20 @@ export class PostgresConnectionModule extends ConnectionModule { fetchedRows = batch.rows; fields = batch.fields; } else { - const result = await client.query(query, paramValues); - fetchedRows = result.rows; - fields = result.fields; - affectedRowCount = result.rowCount ?? undefined; + // Writes stream too, but they must be DRAINED rather than stopped + // early: PostgreSQL applies an UPDATE ... RETURNING incrementally, so + // rows never pulled are never modified. readBoundedCursor stops after + // one bounded read and closes the portal — correct for a SELECT, + // silently partially-applied for a write (#1298, #1326). + const batch = await drainBoundedCursor( + client, + query, + paramValues, + config.rowLimit + 1, + ); + fetchedRows = batch.rows; + fields = batch.fields; + affectedRowCount = batch.affectedRowCount; } // Commit transaction diff --git a/connection/src/postgresql/cursor-read.ts b/connection/src/postgresql/cursor-read.ts index 7290f581..afad8a00 100644 --- a/connection/src/postgresql/cursor-read.ts +++ b/connection/src/postgresql/cursor-read.ts @@ -77,3 +77,99 @@ function closeCursorSafely(cursor: Cursor): Promise { } }); } + +/** A drained cursor: retained rows, field descriptors, and the true row count. */ +export interface DrainedCursor extends CursorBatch { + /** + * Rows the statement actually produced or affected — not the retained + * count. Feeds the COMPLETE / NO_DATA decision for writes that return + * nothing, e.g. an INSERT without RETURNING. + */ + affectedRowCount: number | undefined; +} + +/** Rows pulled per round-trip while draining. Bounds memory, not correctness. */ +const DRAIN_BATCH_SIZE = 500; + +/** + * Executes `query` through a server-side cursor and reads it **to exhaustion**, + * retaining at most `maxRows` rows. + * + * This is the WRITE-path counterpart to `readBoundedCursor`, and the difference + * is correctness rather than performance (#1298, #1326). + * + * PostgreSQL executes a portal incrementally. `readBoundedCursor` does one + * bounded read and then closes the cursor, which is exactly right for a SELECT + * — but on an `UPDATE … RETURNING` the rows never pulled are never modified, + * and closing the portal abandons them. Reusing it for writes would silently + * turn "update 1,000,000 rows" into "update 26" and still report success. + * + * So this keeps reading until a batch comes back empty — every row is produced, + * every side effect runs — while retaining only `maxRows`. Peak memory becomes + * the batch size plus `maxRows` instead of the whole result set. + * + * The user's query text is passed unmodified; parameters stay positional. + */ +export async function drainBoundedCursor( + client: PoolClient, + query: string, + values: unknown[], + maxRows: number, +): Promise { + const cursor = client.query(new Cursor(query, values)); + const rows: Record[] = []; + let fields: FieldDef[] = []; + let affectedRowCount: number | undefined; + let total = 0; + + try { + for (;;) { + const batch = await new Promise<{ + rows: Record[]; + fields: FieldDef[]; + rowCount: number | undefined; + }>((resolve, reject) => { + cursor.read(DRAIN_BATCH_SIZE, (err, batchRows, result) => { + if (err) { + reject(err); + return; + } + resolve({ + rows: batchRows as Record[], + fields: result?.fields ?? [], + rowCount: result?.rowCount ?? undefined, + }); + }); + }); + + if (batch.fields.length > 0 && fields.length === 0) { + fields = batch.fields; + } + // A non-returning statement (INSERT without RETURNING) yields no rows, + // so its affected count only ever arrives on the completing batch. + if (batch.rowCount !== undefined && batch.rowCount !== null) { + affectedRowCount = batch.rowCount; + } + + if (batch.rows.length === 0) break; + + total += batch.rows.length; + for (const row of batch.rows) { + if (rows.length < maxRows) rows.push(row); + } + } + } finally { + await closeCursorSafely(cursor); + } + + return { + rows, + fields, + // Only meaningful for statements that RETURN nothing. A returning + // statement's completing batch reports rowCount 0, which would otherwise + // overwrite the real count and make the caller report NO_DATA for an + // UPDATE that changed a thousand rows. When rows were produced, let the + // caller derive the count from them. + affectedRowCount: total > 0 ? undefined : affectedRowCount, + }; +}