Sanitize reused PostgreSQL connections at checkout

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<PgConfig>): 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
 

postgres/pg-cache/src/__tests__/driver.test.ts

@@ -47,13 +47,15 @@ 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<pg.Pool, [any]>(() => mock);
     registerPgPoolFactory(factory);
 
     const pool = getPgPool(cfg);
 
     expect(factory).toHaveBeenCalledTimes(1);
     expect(pool).toBe(mock);
+    expect(pool.connect).toBe(alternateConnect);
 
     pgCache.delete(cfg.database);
   });

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<string, string>;
+  _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<void>((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<void>((resolve, reject) => {
+      pool.connect((checkoutError, checkedOutClient) => {
+        try {
+          expect(checkoutError).toBe(error);
+          expect(checkedOutClient).toBeUndefined();
+          expect(client.release).toHaveBeenCalledWith(true);
+          resolve();
+        } catch (assertionError) {
+          reject(assertionError);
+        }
+      });
+    });
+  });
+});

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<PgConfig> & { pool?: PgPoolConfig }): pg.Pool => {

postgres/pg-cache/src/sanitizer.ts

@@ -0,0 +1,79 @@
+import type pg from 'pg';
+
+type PgConnectionWithPreparedState = {
+  parsedStatements?: Record<string, string>;
+  _graphilePreparedStatementCache?: unknown;
+};
+
+type PgClientWithPreparedState = pg.PoolClient & {
+  connection?: PgConnectionWithPreparedState;
+};
+
+const checkoutSanitizedPools = new WeakSet<pg.Pool>();
+
+/**
+ * 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<pg.PoolClient> {
+  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<pg.PoolClient> => {
+    const client = await connect();
+    return sanitizePgClient(client);
+  };
+
+  pool.connect = ((callback?: (
+    err: Error | undefined,
+    client: pg.PoolClient | undefined,
+    done: (release?: boolean | Error) => void
+  ) => void): Promise<pg.PoolClient> | 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;
+}

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": [