Skip to content

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

Draft
Elrendio wants to merge 2 commits into
bytebase:mainfrom
Elrendio:fix/postgres-cancel-query-on-abort
Draft

fix(postgres): cancel the running query when the request is aborted#401
Elrendio wants to merge 2 commits into
bytebase:mainfrom
Elrendio:fix/postgres-cancel-query-on-abort

Conversation

@Elrendio

@Elrendio Elrendio commented Aug 10, 2026

Copy link
Copy Markdown
Contributor

src/tools/execute-sql.ts receives extra, which carries the MCP SDK's AbortSignal, and uses it
only for trackToolRequest(...). There is no AbortSignal/AbortController anywhere in src/, and
ExecuteOptions has no field to carry one. Cancelling a tool call therefore only stops us waiting
for the answer
: the backend goes on producing a result set nobody will ever read.

This is the same class of bug as #384 ("MySQL queries can continue running on the server after
query timeout"
) and this change mirrors the fix accepted for it in #386. MySQL's trigger is its
own client-side timeout; here the trigger is the client cancelling. Both end the same way: kill the
statement server-side over a second connection, and don't hand the connection back to the pool in a
state the next borrower inherits. PostgreSQL was simply left behind.

It matters most where it hurts most — read replicas serving agents, where an abandoned analytical
query keeps consuming IO for its full runtime and drives replication lag long after the agent has
given up and moved on. We hit this running DBHub against ~24 production PostgreSQL read replicas;
abandoned agent queries contributed to replication lag of up to 10.9 h and one production incident
caused by a single query that ran for 1 h 30.

Reproduction

Added as the integration test Query cancellation (options.signal): start SELECT pg_sleep(30),
abort the signal, then watch pg_stat_activity.

  • Before: the test fails after 30,042 ms with expected null to be an instance of Error
    the query resolved. The abort changed nothing and the backend slept the full 30 s.
  • After: the statement rejects with SQLSTATE 57014 (query_canceled) within milliseconds,
    and the backend is gone from pg_stat_activity.

Fix

  • ExecuteOptions gains an optional signal; execute_sql passes the request's own, read from
    extra.mcpReq.signal — a tool handler receives { sessionId, mcpReq, http }, and the AbortSignal
    hangs off the request context, not off the extra itself. Connectors
    that ignore it are unaffected — MySQL/MariaDB/SQLite/SQL Server behaviour is untouched.
  • PostgreSQL cancellation is out-of-band by design: the request must arrive on a different
    connection, since the one running the query is blocked waiting for it. cancelBackend 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.
  • The connection is destroyed instead of released after a cancellation. A 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. Mirrors MySQL's isConnectionPoisoned. Covered by a
    test asserting the connector still works after a cancellation.
  • The abort listener is removed in finally, so listeners don't accumulate across the statements of
    one request.
  • An already-aborted signal is refused before a connection is taken, and re-checked once one is held:
    an AbortSignal dispatches abort exactly once, so a signal that fired while we waited for the
    pool would never reach a listener registered afterwards.

Cancellation stays advisory — the backend acts on it at an interrupt point, and a statement that
already finished ignores it — so cancelBackend is best-effort and logs rather than throwing.

Testing

  • pnpm run test:unit961 passed (958 on main; +3 new).
  • pnpm run test:integration336 passed (331 on main; +5) across PostgreSQL, MySQL, MariaDB,
    SQL Server and SQLite containers.
  • pnpm run build:backend — clean.
  • Backed by a negative control: reverting only the connector change makes the new tests fail after
    30 s, as shown above.

New coverage: signal forwarding from the tool handler (including a regression guard that a
top-level extra.signal is not read), and four PostgreSQL cancellation integration tests (plain,
inside BEGIN READ ONLY, pool reuse after cancellation, already-aborted signal).

The handler-to-connector wiring is covered only by unit tests with a hand-written extra, which is
how the wrong path survived the first round here. An end-to-end test would need to drive a real
server over stdio against PostgreSQL; the existing JSON-RPC integration harness is HTTP + SQLite, so
that is not a small addition. I verified this path manually instead (see Reproduction).

Notes for reviewers

  • explain_sql also runs user SQL and could take the same signal — one line — but is left out to
    keep this diff to the reported bug. Happy to add it.
  • Cancellation is implemented for PostgreSQL only. The same treatment for SQLite/SQL Server would be
    a natural follow-up; MySQL already has the timeout-triggered half from fix(mysql): kill server-side query on timeout instead of leaking it into the pool #386.
  • npx tsc --noEmit reports 132 pre-existing errors on main; this branch reports the same 132,
    none in the changed files.
  • Found alongside a separate max_rows coverage bug, sent as its own PR — the two are independent
    and share no files but one integration test file.

The MCP request's AbortSignal reached `execute_sql` (as `extra`) but was
only ever read for telemetry. Nothing in `src/` referenced an
AbortSignal, and `ExecuteOptions` had no field to carry one, so
cancelling a tool call merely stopped us waiting for the answer: the
backend went on producing a result set nobody would ever read.

That is the same class of bug as bytebase#384 on MySQL — a query outliving the
client that asked for it — and this mirrors the fix accepted for it in
bytebase#386. MySQL's trigger is its own client-side timeout; here the trigger
is the client cancelling. Both end the same way: kill the statement
server-side over a second connection, and don't hand the connection
back to the pool in a state the next borrower would inherit.

It matters most exactly where it hurts most: on read replicas serving
agents, where an abandoned analytical query keeps consuming IO for its
full runtime and drives replication lag long after the agent has gone.

Reproduction (integration test `Query cancellation (options.signal)`):
start `SELECT pg_sleep(30)`, abort the signal, then watch
`pg_stat_activity`. Before this change the test fails after 30,042ms
with the query having *resolved* — the abort changed nothing and the
backend slept the full 30s. After it, the statement rejects with
SQLSTATE 57014 (query_canceled) within milliseconds and the backend is
gone from `pg_stat_activity`.

- `ExecuteOptions` gains an optional `signal`; `execute_sql` passes the
  request's own. Connectors that ignore it are unaffected.
- PostgreSQL cancellation is out-of-band: the request must arrive on a
  different connection, since the one running the query is blocked
  waiting for it. `cancelBackend` opens a dedicated short-lived
  connection rather than borrowing from the pool — a pool saturated
  with the very queries being cancelled has nothing left to lend.
  `pg`'s own `Client.cancel` is not reusable here: it needs the
  internal Query object, which the promise-based query API never
  exposes.
- The connection is destroyed instead of released after a cancellation.
  A CancelRequest races the statement it targets, so it can instead
  land on the ROLLBACK issued during cleanup and leave the session in
  an aborted transaction. Mirrors MySQL's `isConnectionPoisoned`.
- The abort listener is removed in `finally`, so listeners don't
  accumulate across the statements of one request.
- An already-aborted signal is refused before a connection is taken,
  and re-checked once one is held: an AbortSignal dispatches "abort"
  exactly once, so a signal that fired while we waited for the pool
  would never reach a listener registered afterwards.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Elrendio added a commit to Elrendio/dbhub that referenced this pull request Aug 10, 2026
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The signal was read as `extra.signal`, which the SDK never populates: a tool
handler receives `{ sessionId, mcpReq, http }`, and the request's AbortSignal
hangs off `extra.mcpReq`. `options.signal` was therefore always undefined, so
cancelling a tool call still left the query running on the database — the exact
failure the connector change was written to prevent.

Proven end to end against a real PostgreSQL, driving the server over stdio and
watching pg_stat_activity: with the fix the backend disappears 201 ms after
`notifications/cancelled`, and reverting only this line leaves it running past
10 s.

The unit test could not catch this: it hand-rolled `{ signal }` as the handler
extra, encoding the same wrong assumption as the code. It now uses the shape the
SDK actually passes, plus a regression guard asserting that a top-level
`extra.signal` is NOT read.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@Elrendio

Copy link
Copy Markdown
Contributor Author

Pushed 5354636: the signal was being read as extra.signal, which the SDK never populates — a tool handler gets { sessionId, mcpReq, http }, and the AbortSignal hangs off extra.mcpReq. So options.signal was always undefined and the connector-side cancellation never fired through the MCP path, which is the whole point of the PR.

Caught by driving a real server over stdio against PostgreSQL and watching pg_stat_activity: the backend disappears 201 ms after notifications/cancelled, and reverting only that line leaves it running past 10 s.

The unit test did not catch it because it hand-rolled { signal } as the handler extra, encoding the same wrong assumption as the code — it now uses the real shape, plus a guard asserting a top-level extra.signal is not read.

Elrendio added a commit to Elrendio/dbhub that referenced this pull request Aug 10, 2026
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant