Repository navigation
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>
![salty-hambot[bot]](https://avatars.githubusercontent.com/in/2952321?s=60&v=4){kind=small}
| ) | ||
| attach_fn = self.__attach_fn | ||
| self.__start(lambda handler: attach_fn(execution_id, handler, store)) | ||
|
|
There was a problem hiding this comment.
medium: Reject duplicate attachments before cancelling the cursor's query
If this cursor has an in-flight execution X and another cursor on the same connection tracks execution Y, calling attach(Y) enters __start() and cancels X, then replaces its queue and clears its ID. Only afterward does Connection.__send() reject Y as already tracked. The call raises ProgrammingError, but X has been cancelled and its outcome is no longer retrievable through the cursor. Check whether Y is already tracked before resetting or cancelling X.
| # Before the request goes out: the session's cancelled state can | ||
| # only follow it. | ||
| query.cancel_requested = True | ||
|
|
There was a problem hiding this comment.
medium: Do not attribute cancellation to an unsent cancel request
__cancel_query() sets cancel_requested before it acquires the send lock and sends the cancel. If the session reports an independently cancelled attached execution during that interval, __listen() sees the flag and returns an empty successful result instead of ExecutionCancelledError, even though this connection's cancel has not reached the session. This can conceal the unknown outcome of an attached write.
|
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? |
This branch has not been deployed
2 participants
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