fix(postgres): keep the process alive when the server drops a checked-out connection (#480)

pg-pool only listens for a client's 'error' while the client is idle:
the listener is removed on checkout and re-attached on release. When the
server ends the connection while a query is in flight, pg emits 'error'
on the Client itself, and with nobody listening Node reports an
unhandled 'error' event and exits the whole DBHub process. #423 added
a pool-level listener, which only covers idle clients.

The reported trigger is a multi-statement batch whose second statement
runs pg_terminate_backend(pg_backend_pid()): the ROLLBACK sent by
rollbackQuietly is still queued when the socket closes, so the client
is checked out when the event fires. The same window exists for a
single statement under readonly (BEGIN READ ONLY takes the same path,
and the plugin/MCPB bundles are read-only by default), and for any
drop that arrives as a socket error without a FATAL message first
(failover, proxy reset), where even a plain single statement crashes.

Route every pool.connect() through acquireClient(), which attaches a
client-level 'error' listener for the whole checkout and passes the
captured error to release() so pg-pool closes and discards the client
instead of handing the dead connection to the next caller. This mirrors
what pool.query() already does for its own checkout.

Adds a loopback protocol-level test that ends the connection mid-batch,
inside a READ ONLY transaction, and by a bare socket reset, and checks
that the tool call errors, the process survives, and the next query gets
a fresh connection. The two pool mocks now expose the EventEmitter
surface a real pg.Client has.

Fixes #479


Claude-Session: https://claude.ai/code/session_017Eyu55gHK6LVo6am2nvCrz

Co-authored-by: Claude <noreply@anthropic.com>

Author: tianzhou

src/connectors/__tests__/pool-error-listener.test.ts

@@ -17,7 +17,9 @@ import { EventEmitter } from "events";
  */
 
 class FakePgPool extends EventEmitter {
-  connect = vi.fn().mockResolvedValue({ release: vi.fn() });
+  // The client is a real EventEmitter, like pg.Client: the connector attaches a
+  // per-checkout 'error' listener to it (bytebase/dbhub#479).
+  connect = vi.fn().mockResolvedValue(Object.assign(new EventEmitter(), { release: vi.fn() }));
   end = vi.fn().mockResolvedValue(undefined);
 }
 

src/connectors/__tests__/postgres-connection-loss.test.ts

@@ -0,0 +1,97 @@
+import { createServer, type Socket } from "node:net";
+import { once } from "node:events";
+import { afterEach, beforeEach, expect, it, vi } from "vitest";
+import { PostgresConnector } from "../postgres/index.js";
+
+// A loopback PostgreSQL protocol peer that ends the connection on demand while
+// a query is in flight, the way pg_terminate_backend, a failover or a proxy
+// does. pg emits 'error' on the Client for an unexpected disconnect; while the
+// client is checked out of the pool nothing but the connector is listening, so
+// without a listener the whole process exits (#479).
+function message(type: string, body: Buffer) {
+  const length = Buffer.alloc(4);
+  length.writeInt32BE(body.length + 4);
+  return Buffer.concat([Buffer.from(type), length, body]);
+}
+const ready = message("Z", Buffer.from("I"));
+const sockets = new Set<Socket>();
+let connections = 0;
+const server = createServer(socket => {
+  connections++;
+  sockets.add(socket);
+  socket.on("close", () => sockets.delete(socket));
+  let startup = true;
+  let buffer: Buffer = Buffer.alloc(0);
+  socket.on("data", data => {
+    buffer = Buffer.concat([buffer, data]);
+    while (buffer.length >= (startup ? 4 : 5)) {
+      const length = buffer.readInt32BE(startup ? 0 : 1) + (startup ? 0 : 1);
+      if (buffer.length < length) return;
+      const packet = buffer.subarray(0, length);
+      buffer = buffer.subarray(length);
+      if (startup) {
+        startup = false;
+        socket.write(Buffer.concat([message("R", Buffer.alloc(4)), ready]));
+      } else if (packet[0] === 81) { // simple query
+        const sql = packet.subarray(5, length - 1).toString();
+        if (sql.includes("pg_terminate_backend")) {
+          // FATAL, then the backend closes the socket.
+          socket.end(message("E", Buffer.from("SFATAL\0C57P01\0Mterminating connection due to administrator command\0\0")));
+        } else if (sql.includes("reset_socket")) {
+          // No error message first: the TCP connection just goes away.
+          socket.destroy();
+        } else {
+          socket.write(Buffer.concat([message("C", Buffer.from("SELECT 0\0")), ready]));
+        }
+      } else if (packet[0] === 88) { // terminate
+        socket.end();
+      }
+    }
+  });
+});
+
+let connector: PostgresConnector;
+
+beforeEach(async () => {
+  connections = 0;
+  server.listen(0, "127.0.0.1");
+  await once(server, "listening");
+  const port = (server.address() as { port: number }).port;
+  connector = new PostgresConnector();
+  await connector.connect(`postgres://user@127.0.0.1:${port}/db?sslmode=disable`);
+  expect(connections).toBe(1);
+});
+
+afterEach(async () => {
+  await connector.disconnect();
+  for (const socket of sockets) socket.destroy();
+  await new Promise<void>(resolve => server.close(() => resolve()));
+  vi.restoreAllMocks();
+});
+
+async function expectRecovery(): Promise<void> {
+  // Let the socket 'end'/'close' land while the process is still alive.
+  await new Promise(resolve => setTimeout(resolve, 50));
+  // The dropped client was discarded, not returned to the pool: the next
+  // query opens a fresh connection instead of drawing the dead one.
+  await expect(connector.executeSQL("SELECT 1", {})).resolves.toMatchObject({ resultSets: [{ rowCount: 0 }] });
+  expect(connections).toBe(2);
+}
+
+it("survives the server ending the connection mid-batch (multi-statement)", async () => {
+  await expect(connector.executeSQL("SELECT 1; SELECT pg_terminate_backend(pg_backend_pid())", {}))
+    .rejects.toThrow("terminating connection due to administrator command");
+  await expectRecovery();
+});
+
+it("survives the server ending the connection inside a READ ONLY transaction", async () => {
+  await expect(connector.executeSQL("SELECT pg_terminate_backend(pg_backend_pid())", { readonly: true }))
+    .rejects.toThrow("terminating connection due to administrator command");
+  await expectRecovery();
+});
+
+it("survives the connection being reset without an error message", async () => {
+  await expect(connector.executeSQL("SELECT reset_socket()", {}))
+    .rejects.toThrow("Connection terminated unexpectedly");
+  await expectRecovery();
+});

src/connectors/__tests__/postgres-query-timeout.test.ts

@@ -19,17 +19,28 @@ const clientTimeout = () => new Error("Query read timeout");
 class FakePgPool extends EventEmitter {
   statements: string[] = [];
   failOn: (sql: string) => Error | undefined = () => undefined;
-  client = {
+  /**
+   * pg-pool's release for the checked-out client. pg-pool reassigns
+   * `client.release` on every checkout, and the connector wraps it per
+   * checkout to attach an 'error' listener (bytebase/dbhub#479), so the
+   * underlying mock is reinstalled by connect() and asserted on directly.
+   */
+  release = vi.fn();
+  // A real EventEmitter, like pg.Client.
+  client = Object.assign(new EventEmitter(), {
     processID: BACKEND_PID,
     query: vi.fn(async (sql: string) => {
       this.statements.push(sql);
       const failure = this.failOn(sql);
       if (failure) throw failure;
       return { rows: [{ ok: 1 }], rowCount: 1 };
     }),
-    release: vi.fn(),
-  };
-  connect = vi.fn(async () => this.client);
+    release: this.release as (err?: Error | boolean) => void,
+  });
+  connect = vi.fn(async () => {
+    this.client.release = this.release;
+    return this.client;
+  });
   query = vi.fn(async (_config: unknown) => ({ rows: [{ pg_cancel_backend: true }], rowCount: 1 }));
   end = vi.fn().mockResolvedValue(undefined);
 }
@@ -53,7 +64,7 @@ async function connect() {
     queryTimeoutSeconds: 30,
   });
   // connect() checks out and releases a client to test the connection.
-  pool.client.release.mockClear();
+  pool.release.mockClear();
   pool.connect.mockClear();
   return connector;
 }
@@ -78,8 +89,8 @@ describe("PostgreSQL client-side query timeout", () => {
     expect(pool.statements).not.toContain("ROLLBACK");
     // Released exactly once, with a truthy argument: pg-pool closes the socket
     // and drops the client instead of returning it to the pool.
-    expect(pool.client.release).toHaveBeenCalledTimes(1);
-    expect(pool.client.release).toHaveBeenCalledWith(true);
+    expect(pool.release).toHaveBeenCalledTimes(1);
+    expect(pool.release).toHaveBeenCalledWith(true);
     // The statement is cancelled on the server over another pooled connection,
     // bounded by its own short timeout...
     expect(pool.query).toHaveBeenCalledWith({
@@ -89,7 +100,7 @@ describe("PostgreSQL client-side query timeout", () => {
     });
     // ...after the client was discarded, so the cancel can use its pool slot
     // (pool_max_connections = 1 would otherwise wait on itself).
-    expect(pool.client.release.mock.invocationCallOrder[0]).toBeLessThan(
+    expect(pool.release.mock.invocationCallOrder[0]).toBeLessThan(
       pool.query.mock.invocationCallOrder[0]
     );
   });
@@ -102,7 +113,7 @@ describe("PostgreSQL client-side query timeout", () => {
     await expect(connector.executeSQL("SELECT pg_sleep(600)", { readonly: true })).rejects.toThrow(
       "Query read timeout"
     );
-    expect(pool.client.release).toHaveBeenCalledWith(true);
+    expect(pool.release).toHaveBeenCalledWith(true);
   });
 
   it("keeps the client for a server-side statement_timeout cancellation", async () => {
@@ -119,8 +130,8 @@ describe("PostgreSQL client-side query timeout", () => {
     // The server ended the statement itself: the client is clean, so it is
     // rolled back and returned to the pool as usual, and nothing is cancelled.
     expect(pool.statements[pool.statements.length - 1]).toBe("ROLLBACK");
-    expect(pool.client.release).toHaveBeenCalledTimes(1);
-    expect(pool.client.release).toHaveBeenCalledWith();
+    expect(pool.release).toHaveBeenCalledTimes(1);
+    expect(pool.release).toHaveBeenCalledWith();
     expect(pool.query).not.toHaveBeenCalled();
   });
 });

src/connectors/postgres/index.ts

@@ -291,7 +291,7 @@ export class PostgresConnector implements Connector {
       });
 
       // Test the connection
-      const client = await this.pool.connect();
+      const client = await this.acquireClient();
       client.release();
     } catch (err) {
       // Tear down the pool if it was created before the failure, otherwise it
@@ -318,7 +318,7 @@ export class PostgresConnector implements Connector {
       throw new Error("Not connected to database");
     }
 
-    const client = await this.pool.connect();
+    const client = await this.acquireClient();
     try {
       const result = await client.query(`
         SELECT schema_name
@@ -338,7 +338,7 @@ export class PostgresConnector implements Connector {
       throw new Error("Not connected to database");
     }
 
-    const client = await this.pool.connect();
+    const client = await this.acquireClient();
     try {
       // Use the configured default schema (from search_path config, defaults to 'public')
       const schemaToUse = schema || this.defaultSchema;
@@ -368,7 +368,7 @@ export class PostgresConnector implements Connector {
       throw new Error("Not connected to database");
     }
 
-    const client = await this.pool.connect();
+    const client = await this.acquireClient();
     try {
       const schemaToUse = schema || this.defaultSchema;
 
@@ -394,7 +394,7 @@ export class PostgresConnector implements Connector {
       throw new Error("Not connected to database");
     }
 
-    const client = await this.pool.connect();
+    const client = await this.acquireClient();
     try {
       // Use the configured default schema (from search_path config, defaults to 'public')
       const schemaToUse = schema || this.defaultSchema;
@@ -421,7 +421,7 @@ export class PostgresConnector implements Connector {
       throw new Error("Not connected to database");
     }
 
-    const client = await this.pool.connect();
+    const client = await this.acquireClient();
     try {
       // Use the configured default schema (from search_path config, defaults to 'public')
       const schemaToUse = schema || this.defaultSchema;
@@ -477,7 +477,7 @@ export class PostgresConnector implements Connector {
       throw new Error("Not connected to database");
     }
 
-    const client = await this.pool.connect();
+    const client = await this.acquireClient();
     try {
       // Use the configured default schema (from search_path config, defaults to 'public')
       const schemaToUse = schema || this.defaultSchema;
@@ -519,7 +519,7 @@ export class PostgresConnector implements Connector {
       throw new Error("Not connected to database");
     }
 
-    const client = await this.pool.connect();
+    const client = await this.acquireClient();
     try {
       // Use the configured default schema (from search_path config, defaults to 'public')
       const schemaToUse = schema || this.defaultSchema;
@@ -551,7 +551,7 @@ export class PostgresConnector implements Connector {
       throw new Error("Not connected to database");
     }
 
-    const client = await this.pool.connect();
+    const client = await this.acquireClient();
     try {
       const schemaToUse = schema || this.defaultSchema;
 
@@ -581,7 +581,7 @@ export class PostgresConnector implements Connector {
       throw new Error("Not connected to database");
     }
 
-    const client = await this.pool.connect();
+    const client = await this.acquireClient();
     try {
       const [connectionsResult, maxConnectionsResult, bufferCacheResult] = await Promise.all([
         client.query(`
@@ -643,7 +643,7 @@ export class PostgresConnector implements Connector {
       throw new Error("Not connected to database");
     }
 
-    const client = await this.pool.connect();
+    const client = await this.acquireClient();
     try {
       // Use the configured default schema (from search_path config, defaults to 'public')
       const schemaToUse = schema || this.defaultSchema;
@@ -680,7 +680,7 @@ export class PostgresConnector implements Connector {
       throw new Error("Not connected to database");
     }
 
-    const client = await this.pool.connect();
+    const client = await this.acquireClient();
     try {
       // Use the configured default schema (from search_path config, defaults to 'public')
       const schemaToUse = schema || this.defaultSchema;
@@ -773,7 +773,7 @@ export class PostgresConnector implements Connector {
       throw new Error("Not connected to database");
     }
 
-    const client = await this.pool.connect();
+    const client = await this.acquireClient();
     // Captured up front, before a timeout can lead to the client being discarded.
     const backendPid = (client as unknown as { processID: number | null }).processID;
     let isClientDiscarded = false;
@@ -891,6 +891,44 @@ export class PostgresConnector implements Connector {
     }
   }
 
+  /**
+   * Checks a client out of the pool with an 'error' listener attached for the
+   * whole checkout, and releases it with that error so the pool discards it.
+   *
+   * pg-pool only listens for a client's 'error' while the client is idle: the
+   * listener is removed on checkout and re-attached on release. If the server
+   * ends the connection while a query is in flight (pg_terminate_backend, a
+   * failover, idle_in_transaction_session_timeout, a proxy reset), pg emits
+   * 'error' on the Client itself, and with nobody listening Node exits the
+   * whole process (#479). The pool-level listener in connect() cannot help:
+   * it only covers idle clients. pool.query() attaches exactly this kind of
+   * per-checkout listener; this does the same for the pool.connect() path.
+   */
+  private async acquireClient(): Promise<pg.PoolClient> {
+    if (!this.pool) {
+      throw new Error("Not connected to database");
+    }
+    const client = await this.pool.connect();
+    let connectionError: Error | undefined;
+    const onError = (err: Error) => {
+      connectionError = err;
+    };
+    client.on("error", onError);
+    const release = client.release;
+    client.release = (err?: Error | boolean) => {
+      client.removeListener("error", onError);
+      // Releasing with a truthy argument makes the pool close and drop the
+      // client instead of handing the dead connection to the next caller.
+      const reason = err ?? connectionError;
+      if (reason === undefined) {
+        release();
+      } else {
+        release(reason);
+      }
+    };
+    return client;
+  }
+
   /**
    * Best-effort rollback after a failed statement, so a failed ROLLBACK (e.g.
    * a dropped connection) cannot mask the original error. Skipped for a