diff --git a/src/connectors/__tests__/postgres.integration.test.ts b/src/connectors/__tests__/postgres.integration.test.ts index c632fd8..f7dcfe9 100644 --- a/src/connectors/__tests__/postgres.integration.test.ts +++ b/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 => { + 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, 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(); diff --git a/src/connectors/interface.ts b/src/connectors/interface.ts index 151ba2b..335400a 100644 --- a/src/connectors/interface.ts +++ b/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; } /** diff --git a/src/connectors/postgres/index.ts b/src/connectors/postgres/index.ts index d8b9139..40c6b4e 100644 --- a/src/connectors/postgres/index.ts +++ b/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 { @@ -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 { 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 { + 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 diff --git a/src/tools/__tests__/execute-sql.test.ts b/src/tools/__tests__/execute-sql.test.ts index 12e7664..39ae846 100644 --- a/src/tools/__tests__/execute-sql.test.ts +++ b/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 diff --git a/src/tools/execute-sql.ts b/src/tools/execute-sql.ts index c68fcda..cc71769 100644 --- a/src/tools/execute-sql.ts +++ b/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);