feat: reconnect and re-attach automatically after a lost connection (DB-760)

[kcheng486]

Summary

Stacked on #80, in response to Simon's review: consumers shouldn't each need their own reconnect, backoff and cancellation loop around reconnect() and attach(). The driver already has the session endpoint, the pending execution IDs, and their handlers and stores. So when the WebSocket drops, the connection now reconnects to the same session by itself and re-attaches every pending query. Callers just get their results. The second commit then removes the manual path, Connection.reconnect() and Cursor.attach(), so none of it becomes public API.

Merging: #80's own diff still adds reconnect() and attach(). To keep main from ever shipping them, merge this PR into #80's branch first, then merge #80.

• Reconnect. When the reader thread loses the transport, it reconnects straight to the same session WebSocket, using a snapshot of the original connection's headers and skipping the REST API. It never starts a new session.
  • Retries use jittered exponential backoff (0.5s, doubling up to 5s) for up to reconnect_timeout_seconds. That's a new connect() / connect_direct() parameter: default 60s, None turns reconnecting off.
  • A 401, 403, 404 or other non-retryable 4xx, or a set cancel_event, ends recovery at once.
  • Reconnecting is eager: it happens whether or not queries are pending.
• Re-attach. Each pending query gets retrieve_results with attach: true, keeping its handler and store. So feat: re-attach to a running execution after a lost connection (DB-760) #80's reply handling applies unchanged:
  • On a session with the attach change, this connection becomes the owner again.
  • An older session is polled for the outcome.
  • If the orphan lease ran out, the query fails with ExecutionCancelledError.
  • If the statement never reached the session, it fails with ExecutionNotFoundError, with a message specific to reconnecting.
  • Cancels already sent are sent again after the attaches.
• Sends during recovery. The reader holds the send lock while it recovers.
  • execute() and cancels wait. They then go out over the new transport, after the re-attaches, or fail with ConnectionLostError once the connection closes.
  • close() cuts a backoff short.
  • A send that fails on a dead transport hands its query to the reconnect instead of closing the connection.
• Giving up. Only if reconnecting fails, or is off, do pending queries get ConnectionLostError. A Connection built without a transport factory behaves exactly as before.
• Flapping. If the transport is lost again within 10s of reconnecting, the same recovery continues, with its deadline and backoff. A session that accepts the connection and drops it at once therefore can't be reconnected to forever.
• cancel_requested is now set when the cancel is actually sent, not before it waits for the send lock. Otherwise a cancel waiting out a reconnect would be blamed for the session's own orphan-lease cancellation, which would read as an empty success. This also fixes Hambot's second finding on feat: re-attach to a running execution after a lost connection (DB-760) #80.
• Manual path removed (second commit). Connection.reconnect(), Cursor.attach(), Connection.__attach_execution and the reconnect factory are gone. Once automatic recovery gives up, the session's orphan lease has nearly run out anyway. Cursor is back to main's version apart from the execution_id property.
  • Still public: ConnectionLostError, ExecutionNotFoundError and ExecutionCancelledError, since fetches still raise them.
  • Hambot's first finding on feat: re-attach to a running execution after a lost connection (DB-760) #80 (attach(Y) cancels the cursor's own query before rejecting Y) no longer applies.

Why 60s: keepalive takes up to about 40s (20s ping interval + 20s ping timeout) to notice a dead path, and sql-session cancels orphaned executions after 120s by default (wherobots/sql-session#219). 60s fits inside that.

Related Issues

Relates to DB-760. Base: #80. Server side: wherobots/sql-session#219. With this, the MCP server's reconnect loop in studio-backend#2894 should be unnecessary.

---

Requester Checklist

Complete these before marking Ready for Review

• I have self-reviewed my own code
• I have added/updated tests that prove my fix/feature works
• I have included visual proof (screenshot, video, or test output) if applicable
• All CI checks are passing
• PR size is S/M, OR I have justified the size and added a walkthrough
• I have updated documentation if needed

Visual Proof

tests/test_attach.py now runs #80's reply-handling tests through a reconnect instead of Cursor.attach(): reattach() executes query EID, drops the transport, and hands the connection the fake session to reconnect to. Tests of the removed API itself are gone.

tests/test_reconnect.py (31 new tests) drives fake transports and a fake factory through:

• a lost query completing over the new transport; a result lost with the connection arriving in the attach reply; all pending queries re-attached; an older session polled after re-attaching;
• not found after reconnecting; cancelled by the session while disconnected;
• a failed send handing its query to the reconnect;
• eager reconnect with nothing pending; execute() and a cancel waiting for the reconnect and following the re-attaches; a cancel sent before the loss being sent again;
• giving up when time runs out, when the session refuses, and when the transport flaps; reconnecting turned off, or given an invalid timeout;
• close() during backoff, during a connection attempt, and with an execute() waiting;
• connect_direct()'s factory: same URI, a snapshot of the headers, open_timeout bounded by the time left, 401/403/404 refused, 429/502/503 retried, cancel_event honoured.

Mutation checks: each of these regressions makes a test fail:

• starting a fresh window on every loss;
• setting the cancel flag before the send;
• not re-sending cancels;
• closing the connection on a failed send;
• not discarding a transport that connects after close();
• dropping the attach bookkeeping on re-attach;
• retrying an InterfaceError;
• retrying a 4xx;
• reconnecting with the caller's live headers dict.

$ uv run pytest tests/
307 passed        # 286 on #80 + 31 new + 3 in test_driver.py − 13 tests of the removed API; also run on Python 3.10
$ uv run pre-commit run --all-files
all hooks Passed

Not tested against a real SQL session or proxy:

• Status codes: treating 4xx as "session gone, stop" assumes that's what the proxy answers for a dead session.
• Bearer tokens: headers are snapshotted when the connection opens, so an expired token makes the reconnect fail with a 401.
• Blocking: during a reconnect, execute() or cursor.close() can block for up to reconnect_timeout_seconds. After close(), a waiting sender can stay blocked until the connection attempt in progress times out (at most 30s).
• No re-send: a statement that never reached the session gets ExecutionNotFoundError; the driver doesn't send it again. That's deliberate: it can't tell "never arrived" from "finished and evicted".

Size Justification (if L/XL)

Against #80: +385/−243 in library code and docs, and +823/−328 in tests. The two commits are separate steps: automatic recovery first, then removing the manual path. Walkthrough:

1. driver.py: open_transport, the single-attempt factory, and the new parameter.
2. connection.py:
   1. __main_loop → __reconnect: the recovery window and abort.
   2. __open_transport: backoff.
   3. __resume: the swap, re-attaches and cancels.
   4. __send: the failed-send hand-off and where cancel_requested is set.
3. README.md: "Recovering from a lost connection".
4. Second commit: what's removed, and tests/test_attach.py's reattach() helper.

---

Reviewer Checklist

If these are not met, close the tab — this PR is not ready for review

• Requester checklist above is complete
• All CI checks are passing
• Tests adequately cover the changes

Model: Claude Opus 5.5 (claude-opus-5-5) · Harness: Claude Code

🤖 Generated with Claude Code