fix(postgres): cancel the running query when the request is aborted

src/connectors/__tests__/postgres.integration.test.ts

@@ -699,6 +699,125 @@ describe('PostgreSQL Connector Integration Tests', () => {
     });
   });
 
+  describe('Query cancellation (options.signal)', () => {
+    // A cancelled MCP tool call must stop the query on the *server*. Without
+    // that, aborting only stops us waiting for the answer while the backend
+    // keeps burning IO to produce a result nobody will read.
+    const PROBE = 'dbhub_cancel_probe';
+
+    /** Backends the server currently reports as running the probe query. */
+    const runningProbes = async (observer: PostgresConnector): Promise<number> => {
+      const result = await observer.executeSQL(
+        `SELECT count(*)::int AS n FROM pg_stat_activity
+         WHERE state = 'active'
+           AND query LIKE '%${PROBE}%'
+           AND query NOT LIKE '%pg_stat_activity%'`,
+        {}
+      );
+      return result.resultSets[0].rows[0].n;
+    };
+
+    const waitFor = async (predicate: () => Promise<boolean>, timeoutMs: number) => {
+      const deadline = Date.now() + timeoutMs;
+      while (Date.now() < deadline) {
+        if (await predicate()) return true;
+        await new Promise((resolve) => setTimeout(resolve, 50));
+      }
+      return false;
+    };
+
+    it.each([
+      ['plain execution', {} as const],
+      ['inside a READ ONLY transaction', { readonly: true } as const],
+    ])('kills the running backend when the request is aborted (%s)', async (_label, options) => {
+      const connector = new PostgresConnector();
+      const observer = new PostgresConnector();
+      try {
+        await connector.connect(postgresTest.connectionString);
+        await observer.connect(postgresTest.connectionString);
+
+        const controller = new AbortController();
+        const query = connector.executeSQL(`SELECT pg_sleep(30), '${PROBE}'`, {
+          ...options,
+          signal: controller.signal,
+        });
+        // Capture the rejection now so polling below can't trip an unhandled one.
+        const settled = query.then(
+          () => null,
+          (error) => error as NodeJS.ErrnoException
+        );
+
+        expect(await waitFor(async () => (await runningProbes(observer)) === 1, 10000)).toBe(true);
+
+        controller.abort();
+
+        const error = await settled;
+        expect(error).toBeInstanceOf(Error);
+        // 57014 = query_canceled: the server stopped it, we didn't just walk away.
+        expect((error as any).code).toBe('57014');
+
+        expect(await waitFor(async () => (await runningProbes(observer)) === 0, 10000)).toBe(true);
+      } finally {
+        await connector.disconnect();
+        await observer.disconnect();
+      }
+    }, 60000);
+
+    it('leaves the connector usable after a cancellation', async () => {
+      // The cancelled connection is destroyed rather than returned to the pool:
+      // a CancelRequest races the statement it targets, so it can land on the
+      // cleanup ROLLBACK instead and strand the session in an aborted
+      // transaction that the next borrower would inherit.
+      const connector = new PostgresConnector();
+      try {
+        await connector.connect(postgresTest.connectionString);
+
+        const controller = new AbortController();
+        const cancelled = connector
+          .executeSQL(`SELECT pg_sleep(30), '${PROBE}'`, {
+            readonly: true,
+            signal: controller.signal,
+          })
+          .catch(() => 'rejected');
+        await new Promise((resolve) => setTimeout(resolve, 500));
+        controller.abort();
+        expect(await cancelled).toBe('rejected');
+
+        const after = await connector.executeSQL('SELECT 1 AS ok', { readonly: true });
+        expect(after.resultSets[0].rows[0].ok).toBe(1);
+      } finally {
+        await connector.disconnect();
+      }
+    }, 60000);
+
+    it('does not run the query at all when the signal is already aborted', async () => {
+      const connector = new PostgresConnector();
+      try {
+        await connector.connect(postgresTest.connectionString);
+
+        const controller = new AbortController();
+        controller.abort();
+
+        await expect(
+          connector.executeSQL('SELECT 1', { signal: controller.signal })
+        ).rejects.toThrow();
+      } finally {
+        await connector.disconnect();
+      }
+    }, 30000);
+
+    it('runs normally when no signal is supplied', async () => {
+      const connector = new PostgresConnector();
+      try {
+        await connector.connect(postgresTest.connectionString);
+        const result = await connector.executeSQL('SELECT 1 AS ok', {});
+        expect(result.resultSets[0].rows[0].ok).toBe(1);
+      } finally {
+        await connector.disconnect();
+      }
+    }, 30000);
+  });
+
   describe('Search Path Configuration Tests', () => {
     it('should use first schema in search_path as default for discovery', async () => {
       const connector = new PostgresConnector();

src/connectors/interface.ts

@@ -121,6 +121,14 @@ export interface ExecuteOptions {
    * Note: SDK-level readonly enforcement is set via ConnectorConfig.readonly
    */
   readonly?: boolean;
+  /**
+   * Cancellation signal for the request this execution serves — the MCP
+   * request's own AbortSignal, so a tool call the client cancels stops the
+   * query on the *database* instead of leaving it running while nobody waits
+   * for the result. Honoured by the PostgreSQL connector; a connector that
+   * ignores it runs the query to completion as before.
+   */
+  signal?: AbortSignal;
 }
 
 /**

src/connectors/postgres/index.ts

@@ -155,6 +155,10 @@ export class PostgresConnector implements Connector {
 
   private pool: pg.Pool | null = null;
 
+  // Kept so cancelBackend can open its own connection to the same server
+  // without borrowing from the pool
+  private poolConfig: pg.PoolConfig | null = null;
+
   // Source ID is set by ConnectorManager after cloning
   private sourceId: string = "default";
 
@@ -193,6 +197,7 @@ export class PostgresConnector implements Connector {
         }
       }
 
+      this.poolConfig = poolConfig;
       this.pool = new Pool(poolConfig);
 
       // Test the connection
@@ -206,6 +211,7 @@ export class PostgresConnector implements Connector {
         this.pool = null;
         await closeQuietly(() => pool.end());
       }
+      this.poolConfig = null;
       console.error("Failed to connect to PostgreSQL database:", err);
       throw err;
     }
@@ -216,6 +222,7 @@ export class PostgresConnector implements Connector {
       await this.pool.end();
       this.pool = null;
     }
+    this.poolConfig = null;
   }
 
   async getSchemas(): Promise<string[]> {
@@ -668,13 +675,46 @@ export class PostgresConnector implements Connector {
   }
 
 
+  /**
+   * How long to give the out-of-band cancellation before abandoning it.
+   * Connecting and asking the server to signal a backend is metadata-only work
+   * that a healthy server answers in milliseconds; the bound exists only so
+   * cleanup cannot outlive the query it is cleaning up after, independently of
+   * the user's (possibly long, possibly unset) query timeout.
+   */
+  private static readonly CANCEL_TIMEOUT_MS = 5000;
+
   async executeSQL(sql: string, options: ExecuteOptions, parameters?: any[]): Promise<SQLResult> {
     if (!this.pool) {
       throw new Error("Not connected to database");
     }
 
+    const signal = options.signal;
+    // Nothing to run for a caller that is already gone - and no reason to take
+    // a connection from the pool to find that out.
+    signal?.throwIfAborted();
+
     const client = await this.pool.connect();
+    // Captured up front: after the cancellation lands, reading this off the
+    // client races its teardown.
+    const backendPid = backendPidOf(client);
+    let wasCancelled = false;
+    const onAbort = () => {
+      wasCancelled = true;
+      // Fire-and-forget: an abort listener cannot be awaited, and the
+      // cancellation reaches the caller through the in-flight query rejecting
+      // with 57014 (query_canceled), not through this call. cancelBackend
+      // never rejects.
+      void this.cancelBackend(backendPid);
+    };
+
     try {
+      // Re-checked now that we hold a connection: an AbortSignal dispatches
+      // "abort" exactly once, so one that fired while we waited for the pool
+      // would never reach a listener registered afterwards.
+      signal?.throwIfAborted();
+      signal?.addEventListener("abort", onAbort, { once: true });
+
       // Check if this is a multi-statement query
       const statements = splitSQLStatements(sql, "postgres");
 
@@ -765,9 +805,70 @@ export class PostgresConnector implements Connector {
         return { resultSets };
       }
     } finally {
-      client.release();
+      // Removed explicitly rather than left to `once`: a signal outlives a
+      // single executeSQL call, so listeners would otherwise pile up pointing
+      // at connections that have already gone back to the pool.
+      signal?.removeEventListener("abort", onAbort);
+      // release(true) destroys the connection instead of returning it to the
+      // pool. After a cancellation the session's state is not trustworthy: the
+      // CancelRequest races the statement it targets, so it can instead land
+      // on the ROLLBACK issued during cleanup and leave the session sitting in
+      // an aborted transaction that the next borrower would inherit. Same
+      // reasoning as the MySQL connector's isConnectionPoisoned handling.
+      client.release(wasCancelled);
     }
   }
+
+  /**
+   * Best-effort server-side cancellation of whatever `backendPid` is running,
+   * for when the MCP client aborts the tool call.
+   *
+   * PostgreSQL cancellation is out-of-band by design: the request has to arrive
+   * on a *different* connection, because the one running the query is blocked
+   * waiting for its result. This opens a dedicated short-lived connection
+   * rather than borrowing from the pool - a pool saturated with the very
+   * long-running queries being cancelled has nothing left to lend, which is
+   * exactly the situation this runs in. That mirrors what `pg`'s own
+   * Client.cancel does at the protocol level; that entry point is not reusable
+   * here because it needs the internal Query object, which the promise-based
+   * query API never exposes.
+   *
+   * Cancellation is advisory: the backend acts on it only at an interrupt
+   * point, and a statement that has already finished ignores it. A resolved
+   * call is therefore not a guarantee that anything stopped.
+   */
+  private async cancelBackend(backendPid: number | null): Promise<void> {
+    if (backendPid === null || !this.poolConfig) {
+      return;
+    }
+    const cancelClient = new pg.Client({
+      ...this.poolConfig,
+      connectionTimeoutMillis: PostgresConnector.CANCEL_TIMEOUT_MS,
+      query_timeout: PostgresConnector.CANCEL_TIMEOUT_MS,
+    });
+    try {
+      await cancelClient.connect();
+      await cancelClient.query("SELECT pg_cancel_backend($1)", [backendPid]);
+    } catch (error) {
+      // Unconfirmed cancellation: the statement may still be running. There is
+      // nothing further to do from here - the caller already sees the abort -
+      // but it must be visible, since it means a query outlived its request.
+      console.error(`Failed to cancel PostgreSQL backend ${backendPid}:`, error);
+    } finally {
+      await closeQuietly(() => cancelClient.end());
+    }
+  }
+}
+
+/**
+ * The server-side PID of the backend serving this connection. `pg` fills it in
+ * from the BackendKeyData message at connection time, but does not declare it
+ * on its published types, hence the narrow cast. Null when it is unavailable,
+ * which leaves the connection simply uncancellable rather than erroring.
+ */
+function backendPidOf(client: pg.PoolClient): number | null {
+  const processID = (client as unknown as { processID?: number | null }).processID;
+  return typeof processID === "number" ? processID : null;
 }
 
 // Create and register the connector

src/tools/__tests__/execute-sql.test.ts

@@ -164,6 +164,53 @@ describe('execute-sql tool', () => {
     });
   });
 
+  describe('cancellation', () => {
+    // The connector needs the request's AbortSignal to stop the query on the
+    // database. Without it a cancelled tool call only stops us waiting for the
+    // answer, and the query runs to completion server-side.
+    it('forwards the request AbortSignal to the connector', async () => {
+      vi.mocked(mockConnector.executeSQL).mockResolvedValue({ resultSets: [] });
+      const controller = new AbortController();
+
+      const handler = createExecuteSqlToolHandler('test_source');
+      // The shape the SDK actually hands a tool handler: `{ sessionId, mcpReq, http }`.
+      await handler({ sql: 'SELECT 1' }, { sessionId: 's', mcpReq: { signal: controller.signal } });
+
+      expect(mockConnector.executeSQL).toHaveBeenCalledWith(
+        'SELECT 1',
+        expect.objectContaining({ signal: controller.signal })
+      );
+    });
+
+    // Regression guard: reading `extra.signal` instead of `extra.mcpReq.signal` type-checks and
+    // passes a hand-rolled `{ signal }` fixture, while silently never cancelling anything in
+    // production — the signal is simply always undefined.
+    it('does not read a signal off the handler extra itself', async () => {
+      vi.mocked(mockConnector.executeSQL).mockResolvedValue({ resultSets: [] });
+      const controller = new AbortController();
+
+      const handler = createExecuteSqlToolHandler('test_source');
+      await handler({ sql: 'SELECT 1' }, { signal: controller.signal });
+
+      expect(mockConnector.executeSQL).toHaveBeenCalledWith(
+        'SELECT 1',
+        expect.objectContaining({ signal: undefined })
+      );
+    });
+
+    it('passes an undefined signal when the caller provides no extra', async () => {
+      vi.mocked(mockConnector.executeSQL).mockResolvedValue({ resultSets: [] });
+
+      const handler = createExecuteSqlToolHandler('test_source');
+      await handler({ sql: 'SELECT 1' }, null);
+
+      expect(mockConnector.executeSQL).toHaveBeenCalledWith(
+        'SELECT 1',
+        expect.objectContaining({ signal: undefined })
+      );
+    });
+  });
+
   describe('read-only mode enforcement', () => {
     // Statement-classification coverage (write keywords, comment stripping,
     // dialect-specific bypasses, ...) is pinned in

src/tools/execute-sql.ts

@@ -63,6 +63,12 @@ export function createExecuteSqlToolHandler(sourceId?: string) {
       const executeOptions = {
         readonly: isReadOnlyPolicy(policy),
         maxRows: toolConfig?.max_rows,
+        // The MCP request's cancellation signal, which the SDK hangs off the
+        // request context (`extra.mcpReq`) rather than the handler extra
+        // itself. Without it a cancelled tool call only stops us waiting for
+        // the answer: the query keeps running on the database, still burning
+        // the IO nobody is going to read.
+        signal: extra?.mcpReq?.signal as AbortSignal | undefined,
       };
       result = await connector.executeSQL(sql, executeOptions);