diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 7e4a836327..e766da0896 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -3445,6 +3445,9 @@ importers: makage: specifier: ^0.3.0 version: 0.3.0 + pg-cache: + specifier: workspace:^ + version: link:../pg-cache/dist pgsql-test: specifier: workspace:^ version: link:../pgsql-test/dist diff --git a/postgres/pg-cache/README.md b/postgres/pg-cache/README.md index 8888375971..8fac82be64 100644 --- a/postgres/pg-cache/README.md +++ b/postgres/pg-cache/README.md @@ -23,6 +23,7 @@ npm install pg-cache ## Features - LRU cache for PostgreSQL connection pools +- Checkout sanitation for reused node-postgres clients - Automatic pool cleanup and disposal - Extensible cleanup callback system - Service cache for general use @@ -126,6 +127,11 @@ The main PostgreSQL pool cache instance. ### getPgPool(config: Partial): Pool Get or create a cached PostgreSQL pool using the provided configuration. +Clients from the default node-postgres factory run `DISCARD ALL` before every +checkout, and stale client-side prepared-statement bookkeeping is cleared to +match the server. If sanitation fails, the client is destroyed and the checkout +fails. Alternate registered pool factories retain ownership of backend-specific +checkout sanitation. ### svcCache diff --git a/postgres/pg-cache/src/__tests__/driver.test.ts b/postgres/pg-cache/src/__tests__/driver.test.ts index dd155d538f..105e3a669a 100644 --- a/postgres/pg-cache/src/__tests__/driver.test.ts +++ b/postgres/pg-cache/src/__tests__/driver.test.ts @@ -47,6 +47,7 @@ describe('pg-cache pool-factory seam', () => { it('getPgPool builds via the registered factory (no real pg connection)', () => { const cfg = freshConfig(); const mock = createMockPool(); + const alternateConnect = mock.connect; const factory = jest.fn(() => mock); registerPgPoolFactory(factory); @@ -54,6 +55,7 @@ describe('pg-cache pool-factory seam', () => { expect(factory).toHaveBeenCalledTimes(1); expect(pool).toBe(mock); + expect(pool.connect).toBe(alternateConnect); pgCache.delete(cfg.database); }); diff --git a/postgres/pg-cache/src/__tests__/sanitizer.test.ts b/postgres/pg-cache/src/__tests__/sanitizer.test.ts new file mode 100644 index 0000000000..0c38c47206 --- /dev/null +++ b/postgres/pg-cache/src/__tests__/sanitizer.test.ts @@ -0,0 +1,129 @@ +import type { Pool, PoolClient } from 'pg'; + +import { + installCheckoutSanitizer, + sanitizePgClient, +} from '../sanitizer'; + +type PreparedConnection = { + parsedStatements: Record; + _graphilePreparedStatementCache?: { dispose: jest.Mock }; +}; + +const createClient = ( + connection: PreparedConnection, + query = jest.fn().mockResolvedValue({ rows: [] }) +): PoolClient => + ({ + connection, + query, + release: jest.fn(), + }) as unknown as PoolClient; + +const createPool = (connect: jest.Mock): Pool => + ({ + connect, + }) as unknown as Pool; + +describe('PostgreSQL checkout sanitation', () => { + it('discards session state and clears client-side prepared-statement bookkeeping', async () => { + const graphileCache = { dispose: jest.fn() }; + const connection: PreparedConnection = { + parsedStatements: { + tenantLookup: 'select 1', + tenantMutation: 'select 2', + }, + _graphilePreparedStatementCache: graphileCache, + }; + const client = createClient(connection); + + await expect(sanitizePgClient(client)).resolves.toBe(client); + + expect(client.query).toHaveBeenCalledWith('DISCARD ALL'); + expect(connection.parsedStatements).toEqual({}); + expect(connection).not.toHaveProperty('_graphilePreparedStatementCache'); + expect(graphileCache.dispose).not.toHaveBeenCalled(); + expect(client.release).not.toHaveBeenCalled(); + }); + + it('destroys a client and preserves the original sanitation error', async () => { + const error = new Error('DISCARD ALL failed'); + const connection: PreparedConnection = { + parsedStatements: { tenantLookup: 'select 1' }, + }; + const client = createClient( + connection, + jest.fn().mockRejectedValue(error) + ); + + await expect(sanitizePgClient(client)).rejects.toBe(error); + + expect(client.release).toHaveBeenCalledWith(true); + expect(connection.parsedStatements).toEqual({ + tenantLookup: 'select 1', + }); + }); + + it('sanitizes every promise-based checkout and installs only once', async () => { + const connection: PreparedConnection = { parsedStatements: {} }; + const client = createClient(connection); + const connect = jest.fn().mockResolvedValue(client); + const pool = createPool(connect); + + expect(installCheckoutSanitizer(pool)).toBe(pool); + expect(installCheckoutSanitizer(pool)).toBe(pool); + + await expect(pool.connect()).resolves.toBe(client); + await expect(pool.connect()).resolves.toBe(client); + + expect(connect).toHaveBeenCalledTimes(2); + expect(client.query).toHaveBeenNthCalledWith(1, 'DISCARD ALL'); + expect(client.query).toHaveBeenNthCalledWith(2, 'DISCARD ALL'); + }); + + it('preserves node-postgres callback checkout semantics', async () => { + const connection: PreparedConnection = { parsedStatements: {} }; + const client = createClient(connection); + const pool = createPool(jest.fn().mockResolvedValue(client)); + installCheckoutSanitizer(pool); + + await new Promise((resolve, reject) => { + pool.connect((error, checkedOutClient, done) => { + try { + expect(error).toBeUndefined(); + expect(checkedOutClient).toBe(client); + expect(done).toEqual(expect.any(Function)); + done(); + expect(client.release).toHaveBeenCalledWith(); + resolve(); + } catch (assertionError) { + reject(assertionError); + } + }); + }); + }); + + it('reports callback checkout failures only after destroying the client', async () => { + const error = new Error('cannot sanitize client'); + const connection: PreparedConnection = { parsedStatements: {} }; + const client = createClient( + connection, + jest.fn().mockRejectedValue(error) + ); + const pool = createPool(jest.fn().mockResolvedValue(client)); + installCheckoutSanitizer(pool); + + await new Promise((resolve, reject) => { + pool.connect((checkoutError, checkedOutClient) => { + try { + expect(checkoutError).toBe(error); + expect(checkedOutClient).toBeUndefined(); + expect(client.release).toHaveBeenCalledWith(true); + resolve(); + } catch (assertionError) { + reject(assertionError); + } + }); + }); + }); +}); diff --git a/postgres/pg-cache/src/pg.ts b/postgres/pg-cache/src/pg.ts index 08920e924d..920a900ddd 100644 --- a/postgres/pg-cache/src/pg.ts +++ b/postgres/pg-cache/src/pg.ts @@ -5,6 +5,7 @@ import { getPgEnvOptions, PgConfig, PgPoolConfig } from 'pg-env'; import { getActivePgPoolFactory, PgPoolFactory } from './driver'; import { pgCache } from './lru'; +import { installCheckoutSanitizer } from './sanitizer'; const log = new Logger('pg-cache'); @@ -97,7 +98,7 @@ export const defaultPgPoolFactory: PgPoolFactory = (pgConfig): pg.Pool => { } }); - return pgPool; + return installCheckoutSanitizer(pgPool); }; export const getPgPool = (pgConfig: Partial & { pool?: PgPoolConfig }): pg.Pool => { diff --git a/postgres/pg-cache/src/sanitizer.ts b/postgres/pg-cache/src/sanitizer.ts new file mode 100644 index 0000000000..379763e3e9 --- /dev/null +++ b/postgres/pg-cache/src/sanitizer.ts @@ -0,0 +1,79 @@ +import type pg from 'pg'; + +type PgConnectionWithPreparedState = { + parsedStatements?: Record; + _graphilePreparedStatementCache?: unknown; +}; + +type PgClientWithPreparedState = pg.PoolClient & { + connection?: PgConnectionWithPreparedState; +}; + +const checkoutSanitizedPools = new WeakSet(); + +/** + * Forget prepared statements that PostgreSQL removed during `DISCARD ALL`. + * + * Graphile's cache is deliberately deleted rather than disposed: its disposer + * issues asynchronous `DEALLOCATE` queries, which would duplicate `DISCARD ALL` + * and could race with the next owner of the checked-out client. + */ +export function clearPreparedStatementBookkeeping(client: pg.PoolClient): void { + const connection = (client as PgClientWithPreparedState).connection; + if (!connection) return; + + if (connection.parsedStatements) { + for (const statementName of Object.keys(connection.parsedStatements)) { + delete connection.parsedStatements[statementName]; + } + } + + delete connection._graphilePreparedStatementCache; +} + +/** + * Restore a checked-out PostgreSQL client to server defaults before reuse. + * A client that cannot be sanitized is destroyed instead of being returned to + * application code with unknown session state. + */ +export async function sanitizePgClient(client: pg.PoolClient): Promise { + try { + await client.query('DISCARD ALL'); + clearPreparedStatementBookkeeping(client); + return client; + } catch (error) { + client.release(true); + throw error; + } +} + +/** + * Sanitize every client obtained from a node-postgres pool. `pool.query()` also + * goes through `connect()`, so both checkout APIs share the same boundary. + */ +export function installCheckoutSanitizer(pool: pg.Pool): pg.Pool { + if (checkoutSanitizedPools.has(pool)) return pool; + + const connect = pool.connect.bind(pool); + const sanitizedConnect = async (): Promise => { + const client = await connect(); + return sanitizePgClient(client); + }; + + pool.connect = ((callback?: ( + err: Error | undefined, + client: pg.PoolClient | undefined, + done: (release?: boolean | Error) => void + ) => void): Promise | void => { + const pendingClient = sanitizedConnect(); + if (!callback) return pendingClient; + + pendingClient.then( + (client) => callback(undefined, client, client.release.bind(client)), + (error: Error) => callback(error, undefined, () => undefined) + ); + }) as typeof pool.connect; + + checkoutSanitizedPools.add(pool); + return pool; +} diff --git a/postgres/pg-query-context/package.json b/postgres/pg-query-context/package.json index ce07870dc9..70c5c21625 100644 --- a/postgres/pg-query-context/package.json +++ b/postgres/pg-query-context/package.json @@ -34,6 +34,7 @@ "devDependencies": { "@types/pg": "^8.20.4", "makage": "^0.3.0", + "pg-cache": "workspace:^", "pgsql-test": "workspace:^" }, "keywords": [ diff --git a/postgres/pg-query-context/src/__tests__/postgres.integration.test.ts b/postgres/pg-query-context/src/__tests__/postgres.integration.test.ts index 1b2e935aef..3de42a8413 100644 --- a/postgres/pg-query-context/src/__tests__/postgres.integration.test.ts +++ b/postgres/pg-query-context/src/__tests__/postgres.integration.test.ts @@ -1,4 +1,5 @@ import { Pool, type PoolClient } from 'pg'; +import { defaultPgPoolFactory } from 'pg-cache'; import { getConnections } from 'pgsql-test'; import pgQueryContext, { withPgClient } from '../index'; @@ -11,6 +12,23 @@ interface SessionState { user_id: string; } +interface CheckoutState { + application_name: string; + backend_pid: number; + default_transaction_read_only: string; + search_path: string; + row_security: string; +} + +const CHECKOUT_STATE_QUERY = ` + SELECT + current_setting('application_name') AS application_name, + pg_backend_pid() AS backend_pid, + current_setting('default_transaction_read_only') AS default_transaction_read_only, + current_setting('search_path') AS search_path, + current_setting('row_security') AS row_security +`; + async function readSessionState(client: PoolClient): Promise { const result = await client.query(` SELECT @@ -23,17 +41,28 @@ async function readSessionState(client: PoolClient): Promise { return result.rows[0]; } +async function readCheckoutState(client: PoolClient): Promise { + const result = await client.query(CHECKOUT_STATE_QUERY); + return result.rows[0]; +} + describe('pg-query-context transaction-local integration', () => { let db: Awaited>['db']; let teardown: Awaited>['teardown']; let singleClientPool: Pool; + let sanitizedPool: Pool; beforeAll(async () => { ({ db, teardown } = await getConnections({}, [])); singleClientPool = new Pool({ ...db.config, max: 1 }); + sanitizedPool = defaultPgPoolFactory({ + ...db.config, + pool: { max: 1 }, + }) as Pool; }); afterAll(async () => { + if (sanitizedPool) await sanitizedPool.end(); if (singleClientPool) await singleClientPool.end(); if (teardown) await teardown(); }); @@ -198,4 +227,99 @@ describe('pg-query-context transaction-local integration', () => { afterRollbackClient.release(); } }); + + it('sanitizes a reused checkout before applying the complete request context', async () => { + const firstClient = await sanitizedPool.connect(); + let baseline: CheckoutState; + try { + baseline = await readCheckoutState(firstClient); + await firstClient.query("SET application_name TO 'f10-tenant-poison'"); + await firstClient.query('SET default_transaction_read_only TO on'); + await firstClient.query('SET search_path TO pg_catalog'); + await firstClient.query('SET row_security TO off'); + await firstClient.query({ + name: 'f10-checkout-canary', + text: 'SELECT 1 AS value', + }); + } finally { + firstClient.release(); + } + + const insideContext = await withPgClient( + sanitizedPool, + { + role: 'none', + transaction_read_only: 'on', + search_path: 'pg_catalog', + row_security: 'on', + 'jwt.claims.user_id': 'f10-request-user', + }, + async (client) => { + const state = await readSessionState(client); + const statement = await client.query<{ value: number }>({ + name: 'f10-checkout-canary', + text: 'SELECT 2 AS value', + }); + const checkout = await readCheckoutState(client); + return { checkout, state, value: statement.rows[0].value }; + } + ); + + expect(insideContext.checkout.backend_pid).toBe(baseline.backend_pid); + expect(insideContext.checkout.application_name).toBe( + baseline.application_name + ); + expect(insideContext.state).toEqual({ + role: 'none', + transaction_read_only: 'on', + search_path: 'pg_catalog', + row_security: 'on', + user_id: 'f10-request-user', + }); + expect(insideContext.value).toBe(2); + + const afterRequest = await sanitizedPool.connect(); + try { + await expect(readCheckoutState(afterRequest)).resolves.toEqual(baseline); + } finally { + afterRequest.release(); + } + }); + + it('sanitizes direct pool.query checkouts through the same boundary', async () => { + const baseline = await sanitizedPool.query( + CHECKOUT_STATE_QUERY + ); + + await sanitizedPool.query("SET application_name TO 'f10-pool-query-poison'"); + + const restored = await sanitizedPool.query( + CHECKOUT_STATE_QUERY + ); + + expect(restored.rows[0]).toEqual(baseline.rows[0]); + }); + + it('destroys a checkout whose open transaction prevents sanitation', async () => { + const dirtyClient = await sanitizedPool.connect(); + const dirtyState = await readCheckoutState(dirtyClient); + await dirtyClient.query('BEGIN'); + await dirtyClient.query("SET application_name TO 'f10-open-transaction'"); + dirtyClient.release(); + + await expect(sanitizedPool.connect()).rejects.toMatchObject({ + code: '25001', + }); + + const replacementClient = await sanitizedPool.connect(); + try { + const replacementState = await readCheckoutState(replacementClient); + expect(replacementState.backend_pid).not.toBe(dirtyState.backend_pid); + expect(replacementState.application_name).not.toBe( + 'f10-open-transaction' + ); + } finally { + replacementClient.release(); + } + }); });