Repository navigation
feat: re-attach to a running execution after a lost connection (DB-760) - #80
Merged
Merged
Conversation
When the WebSocket to a SQL session drops mid-query, the session keeps running the execution, but the driver had no way to re-claim it (2026-10-14: an MCP query job's socket hit a keepalive timeout and its Spark job ran on for 52 minutes, uncancellable). Add a recovery path: reconnect to the same session, then re-attach to the execution by its ID. - ConnectionLostError(OperationalError), with session_id and execution_id, replaces the plain OperationalError delivered to pending queries on transport loss. Same message, so `except OperationalError` keeps working. - Connection.reconnect() opens a new connection straight to the same session WebSocket, with connect_direct()'s headers and settings, via a factory connect_direct() hands the Connection. Going around the REST API means a dead session fails fast instead of being replaced. A Connection built without a factory raises NotSupportedError. - Cursor.execution_id exposes the current or last execution's ID. - Cursor.attach(execution_id, store=None) resets like execute(), then sends retrieve_results with attach: true. A session that supports it makes this connection the owner and answers with attached: true; an older one ignores the flag and only answers. Reply handling for an attached query: - running/pending: keep waiting. Acknowledged, events are pushed here; otherwise the stale-query watchdog polls it with its usual backoff. - succeeded with results or result_uri: complete as usual. - succeeded, store, no result_uri: complete empty if acknowledged, otherwise fail with an OperationalError, since the location went to the lost connection (or the result is empty) and nothing else will arrive. - "Execution not found" for the attach: ExecutionNotFoundError, without the watchdog's evicted-keep-waiting guard. A later not-found is read by acknowledgement: keep waiting if acknowledged, else ExecutionNotFoundError. - Unacknowledged and still running with probes disabled: fail now, rather than wait forever. The attach is itself a retrieve whose reply may carry the full result, so while it's unanswered the watchdog allows one probe after 4N, as after the normal retrieve. That probe repeats the attach: the session dispatches each request as its own task, so the two replies can arrive in either order, and either then reports the acknowledgement. Attaching an ID this connection already tracks raises ProgrammingError (replacing the registration would strand the first handler), as does attaching the cursor's own in-flight execution, which the reset would otherwise cancel. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…-760) An attached query whose outcome was cancelled completed like an executed one: an empty result, no error. With sql-session's orphan lease, a re-attach that arrives after the grace period has run out (reconnecting can take longer than 120 s, and a grace of 0 is a supported setting) then reads as a successful 0-row result. For a write cancelled midway, that reports it as succeeded. Add ExecutionCancelledError(OperationalError), with session_id and execution_id, exported like the other two. An attached query that reaches cancelled (from the attach reply, a probe reply or a pushed state_updated) without this connection having cancelled it now fails with it, saying the execution was cancelled before or while re-attaching and that its outcome, for a write whether it committed, is unknown. The connection records that it sent a cancel on the Query (cancel_requested), under __lock and before sending, so the flag is set by the time the session can report the cancellation. An attached query this cursor cancelled itself (cursor.close(), or executing on it) still completes empty, and executed queries are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
kcheng486
marked this pull request as ready for review
October 16, 2026 19:26
![salty-hambot[bot]](https://avatars.githubusercontent.com/in/2952321?s=60&v=4){kind=small}
Contributor
|
iiuc the driver already has the session endpoint, pending execution IDs and result handlers, and https://github.com/wherobots/sql-session/pull/219 lets it reclaim an execution without resubmitting the SQL. exposing ConnectionLostError, reconnect() and attach() as the normal recovery path means each consumer has to implement its own reconnect/backoff/cancellation loop for a transient socket failure. could we instead bake that into the driver and all consumers get it for free? |
…DB-760) Review on #80: exposing ConnectionLostError, reconnect() and attach() as the normal recovery path leaves every consumer to write its own reconnect, backoff and cancellation loop for a transient socket failure. The driver already has what that needs (the session endpoint, the pending execution IDs, their handlers and stores), so it now recovers by itself. When the reader thread loses the transport, it reconnects to the same session WebSocket instead of failing pending queries: - connect_direct() hands the Connection a transport factory: one WebSocket attempt to the same URI with the headers snapshot, bounded by the time left. 401/403/404 and other non-retryable 4xx, or a set cancel_event, raise InterfaceError and end the recovery at once; any other failure is retried. - The reader retries with jittered exponential backoff (0.5s doubling to 5s) for up to reconnect_timeout_seconds, a new connect() and connect_direct() parameter (default 60s, None disables). The default fits, after up to ~40s of keepalive detection, in sql-session's 120s grace period for orphaned executions. A transport lost again within 10s of reconnecting continues the same recovery, deadline and backoff included, so a session that accepts and drops at once can't be reconnected to forever. - It reconnects eagerly, with or without pending queries. - On success it re-registers every pending query as Cursor.attach() registers one (fresh watchdog, _Attach(resumed=True), handler and store kept) and sends retrieve_results with attach: true for each, so the existing attach reply handling applies: owner again on a new session, polled on an old one, ExecutionCancelledError if the orphan lease ran out, ExecutionNotFoundError (with a reconnect-specific message) if the statement never reached the session. Cancels this connection already sent are sent again after the attaches. - The reader holds __send_lock for the whole recovery, so execute() and cancels wait for it: they then go out over the new transport after the re-attaches, or fail with ConnectionLostError once the connection is closed. close() cuts a backoff short. - A send that fails on a lost transport no longer closes the connection when reconnecting is on: it aborts the transport so the reader notices, and its query is re-attached with the rest. - Only when reconnecting fails, or is off, do pending queries get ConnectionLostError. A Connection built without a factory behaves as before. cancel_requested is now set in __send, once the cancel is admitted and just before it goes out, rather than before the cancel waits for the send lock. Otherwise a cancel waiting out a reconnect would take the blame for a cancellation the session made itself (the orphan lease) and report it as an empty success. This also resolves the review finding on #80 about that window. Connection.reconnect() and Cursor.attach() stay as the manual path for when automatic recovery is off or has given up. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
9 tasks
…ursor.attach() (DB-760) With the connection reconnecting and re-attaching by itself, the manual path is redundant: after automatic recovery gives up, the session's orphan lease has all but run out too. Keeping it would leave public API to support for little gain, and it is what the review on #80 objected to as the normal recovery path. - Remove Connection.reconnect() and the reconnect_factory connect_direct() passed for it; the transport factory stays. - Remove Cursor.attach() and its attach_fn plumbing; Cursor is back to main's apart from the execution_id property, kept for correlating an execution with its errors and logs. - Remove Connection.__attach_execution, and the "already tracking execution" check in __send, which only an attach could trigger. - Every attach is now a re-attach after reconnecting, so _Attach loses resumed and ExecutionNotFoundError has a single message. - ConnectionLostError, ExecutionNotFoundError and ExecutionCancelledError stay public: fetches still raise them. Their docs no longer point at the manual path. - The README drops the manual-recovery example. tests/test_attach.py drives the same reply handling through a reconnect instead of Cursor.attach(): reattach() executes query EID, loses the transport, and hands the connection the fake session to reconnect to. Tests of the removed API itself (argument checks, rejecting an in-flight or tracked ID, reconnect()'s factory) are gone. This also makes the review finding on #80 about attach(Y) cancelling the cursor's own query before rejecting Y moot. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
# Conflicts: # wherobots/db/connection.py # wherobots/db/constants.py # wherobots/db/driver.py
ClayMav
approved these changes
Oct 19, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
When the WebSocket to a SQL session drops mid-query (a network path failure, a keepalive timeout), the session keeps running the query, but the driver had no way to re-claim it. The pending fetch just waited, and a later cancel went nowhere over the dead socket, so the query ran on unread and couldn't be cancelled. In one recent case it ran for 52 minutes.
This adds a recovery path. The client reconnects to the same session and re-attaches to the query by its execution ID instead of losing it.
ConnectionLostError(OperationalError), with.session_idand.execution_id, is what pending queries now receive on transport loss. The message is unchanged, soexcept OperationalErrorkeeps working.ExecutionNotFoundError(OperationalError): the session doesn't know the execution, for example because it restarted or the connection reached a different session.Cursor.execution_id: the ID of the cursor's current or last execution.Cursor.attach(execution_id, store=None): resets likeexecute(), then sendsretrieve_resultswithattach: true.attached: true.ExecutionCancelledError(OperationalError): an attached execution that was cancelled, and not by this connection, for example by the session's orphan grace period running out before the re-attach. Without it, such a query would read as an empty successful result. A query this cursor cancelled itself still completes empty, as before.Connection.reconnect(): a new connection straight to the same session WebSocket, withconnect_direct()'s headers and settings. It skips the REST API, so a dead session fails fast instead of being replaced by a new one.Consumer: the Wherobots MCP server re-attaches its query jobs after a drop (internal PR).
Related Issues
Relates to Wherobots DB-760. Server side: the sql-session attach + orphan lease change (internal).
Requester Checklist
Complete these before marking Ready for Review
Visual Proof
tests/test_attach.py(52 new tests) drives a fake transport through:attached: true) and an old one (no flag, probes poll it);result_uri, and with noresult_uri, acknowledged and not;ExecutionCancelledError, versus this cursor's own cancel, which completes empty;NotSupportedErrorwith no factory;InterfaceErroroncecancel_eventis set.The new tests failed under mutation checks: removing the attach not-found branch, the eviction attribution, or the attach-probe rule.
Not tested against a real SQL session: the protocol is exercised only through the fake transport here. The studio-backend PR's tests exercise the MCP server's use of these APIs.
Size Justification (if L/XL)
About 470 lines of library code and docs, plus about 1,000 lines of tests. The changes can't be split usefully:
attachneedsConnectionLostErrorandexecution_idto be usable, andreconnectis the only way to get a connection to attach from. Walkthrough:errors.py: the three exceptions.cursor.py:execution_id,attach, and the shared__start().connection.py:__attachregistration, then the reply handling in__listen(Query.attachflags), then the probe rule in__probe_stale_queries.driver.py: the reconnect factory.Reviewer Checklist
If these are not met, close the tab — this PR is not ready for review
Model: Claude Opus 5.5 (claude-opus-5-5) · Harness: Claude Code
🤖 Generated with Claude Code