From fb3a29173383527efa22157186f7b2f0d728da6f Mon Sep 17 00:00:00 2001 From: Kevin Date: Mon, 19 Oct 2026 15:41:37 -0700 Subject: [PATCH 1/2] feat: reconnect and re-attach automatically after a lost connection (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 --- README.md | 86 +++-- tests/test_attach.py | 2 + tests/test_driver.py | 17 + tests/test_reconnect.py | 744 +++++++++++++++++++++++++++++++++++++ wherobots/db/connection.py | 313 ++++++++++++++-- wherobots/db/constants.py | 5 + wherobots/db/driver.py | 59 ++- wherobots/db/errors.py | 3 + 8 files changed, 1153 insertions(+), 76 deletions(-) create mode 100644 tests/test_reconnect.py diff --git a/README.md b/README.md index a5c1945..62d8511 100644 --- a/README.md +++ b/README.md @@ -202,17 +202,51 @@ progress information is available. ### Recovering from a lost connection -If the connection to the SQL session is lost while a query is pending, -fetching its outcome raises `ConnectionLostError`, a subclass of -`OperationalError`. The query may still be running on the SQL session; -the error's `execution_id` and `session_id` attributes identify it. A -cursor's `execution_id` property also holds the ID of its current or -last execution. (When `execute()` itself raises `ConnectionLostError`, -the connection was already lost and the statement was never sent.) - -To recover the outcome, open a new connection to the same SQL session -with `Connection.reconnect()`, and re-attach a cursor to the execution -with `Cursor.attach()`: +If the WebSocket to the SQL session drops while queries are running (a +network path failure, a keepalive timeout), the connection reconnects to +the same SQL session by itself and re-attaches every pending query, so +their results still arrive: there's nothing for your code to do. It keeps +trying, with backoff, for up to `reconnect_timeout_seconds` (60 seconds by +default; see below). It never starts a new session, and gives up at once +if the session refuses the connection, for example because it's gone. +While it reconnects, `execute()` and cancelling (closing a cursor, or +executing on it) wait for it to finish. + +A re-attached query can still fail: + +* If the SQL session doesn't know the execution (the statement didn't + reach it before the connection dropped, or the session restarted), + fetching raises `ExecutionNotFoundError`, a subclass of + `OperationalError`. +* If the SQL session cancelled the execution while the connection was + down (for example, its grace period for a lost connection ran out, or + another client cancelled it), fetching raises + `ExecutionCancelledError`, also an `OperationalError` subclass: its + outcome, including whether a write committed, is unknown. An execution + that you cancel yourself, by closing the cursor or executing on it, + completes with an empty result, as usual. +* An older SQL session that doesn't support attaching still answers with + the execution's state, and the driver polls it for the outcome (see + `stale_query_probe_seconds` below). If such a session reports no + location for a stored result, because it predates reporting it or + because the stored result is empty, the driver can't tell which, and + fails the query with an `OperationalError`. With probes disabled, a + running execution on such a session fails right away. + +#### When reconnecting fails + +If the connection can't reconnect in time, or reconnecting is turned off +(`reconnect_timeout_seconds=None`), fetching a pending query's outcome +raises `ConnectionLostError`, another `OperationalError` subclass. The +query may still be running on the SQL session; the error's +`execution_id` and `session_id` attributes identify it. A cursor's +`execution_id` property also holds the ID of its current or last +execution. (When `execute()` itself raises `ConnectionLostError`, the +connection was already lost and the statement was never sent.) + +To recover the outcome yourself, open a new connection to the same SQL +session with `Connection.reconnect()`, and re-attach a cursor to the +execution with `Cursor.attach()`: ```python from wherobots.db import ConnectionLostError, Store, connect @@ -237,27 +271,10 @@ with connect(api_key='...') as conn: an `InterfaceError`. * `attach(execution_id, store=None)` takes the `Store` the execution was started with, if any. Afterwards, fetch the outcome as after - `execute()`. Like `execute()`, it first cancels the cursor's own - in-flight execution, and the cursor then owns the attached one: - closing the cursor, or executing on it, before the outcome arrives - cancels the execution. -* If the SQL session doesn't know the execution (it ran on another - session, or the session restarted), fetching raises - `ExecutionNotFoundError`, another `OperationalError` subclass. -* If the execution was cancelled on the SQL session before or while you - re-attached (for example, its grace period for a lost connection ran - out, or another client cancelled it), fetching raises - `ExecutionCancelledError`, also an `OperationalError` subclass: its - outcome, including whether a write committed, is unknown. An attached - execution that you cancel yourself, by closing the cursor or executing - on it, completes with an empty result, as an executed one does. -* An older SQL session that doesn't support attaching still answers with - the execution's state, and the driver polls it for the outcome (see - `stale_query_probe_seconds` below). If such a session reports no - location for a stored result, because it predates reporting it or - because the stored result is empty, the driver can't tell which, and - fails the query with an `OperationalError`. With probes disabled, - attaching to a running execution on such a session fails right away. + `execute()`, with the same failures as above. Like `execute()`, it + first cancels the cursor's own in-flight execution, and the cursor + then owns the attached one: closing the cursor, or executing on it, + before the outcome arrives cancels the execution. ### Runtime and region selection @@ -341,3 +358,8 @@ users may find useful: for its state again, in case the query's completion event was lost. Repeat probes back off to at most 8 times this interval. Defaults to 30 seconds; pass `None` to disable probing. +* `reconnect_timeout_seconds`: how long the connection keeps trying to + reconnect to the same SQL session after losing its WebSocket (see + "Recovering from a lost connection" above). Defaults to 60 seconds, + which fits in the SQL session's 120-second grace period for the + queries of a lost connection; pass `None` to turn reconnecting off. diff --git a/tests/test_attach.py b/tests/test_attach.py index a940a00..9eba4b4 100644 --- a/tests/test_attach.py +++ b/tests/test_attach.py @@ -1071,6 +1071,8 @@ def test_lost_query_is_recovered_by_reconnecting_and_attaching(): read_timeout=READ_TIMEOUT, session_id="session-1", stale_query_probe_seconds=PROBE, + # The manual path, for when the connection doesn't reconnect. + reconnect_timeout_seconds=None, ) cursor = conn.cursor() cursor.execute("SELECT * FROM t", store=STORE) diff --git a/tests/test_driver.py b/tests/test_driver.py index 3295cd0..2fcabfe 100644 --- a/tests/test_driver.py +++ b/tests/test_driver.py @@ -92,6 +92,23 @@ def test_stale_query_probes_default_to_30_seconds(self, mock_post, mock_get): _, kwargs = _run_connect_full(mock_post, mock_get) assert kwargs["stale_query_probe_seconds"] == 30.0 + @pytest.mark.parametrize("reconnect_seconds", [None, 5.0]) + @patch("wherobots.db.driver.requests.get") + @patch("wherobots.db.driver.requests.post") + def test_reconnect_timeout_seconds_is_forwarded( + self, mock_post, mock_get, reconnect_seconds + ): + _, kwargs = _run_connect_full( + mock_post, mock_get, reconnect_timeout_seconds=reconnect_seconds + ) + assert kwargs["reconnect_timeout_seconds"] == reconnect_seconds + + @patch("wherobots.db.driver.requests.get") + @patch("wherobots.db.driver.requests.post") + def test_reconnecting_defaults_to_60_seconds(self, mock_post, mock_get): + _, kwargs = _run_connect_full(mock_post, mock_get) + assert kwargs["reconnect_timeout_seconds"] == 60.0 + @patch("wherobots.db.driver.requests.get") @patch("wherobots.db.driver.requests.post") def test_omitted_region_runtime_not_sent(self, mock_post, mock_get): diff --git a/tests/test_reconnect.py b/tests/test_reconnect.py new file mode 100644 index 0000000..daceadb --- /dev/null +++ b/tests/test_reconnect.py @@ -0,0 +1,744 @@ +"""Reconnecting by itself after a lost transport, and re-attaching queries.""" + +import itertools +import json +import logging +import queue +import threading +import time +from unittest.mock import MagicMock, patch + +import cbor2 +import pandas +import pyarrow +import pytest +from websockets.datastructures import Headers +from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK +from websockets.exceptions import InvalidStatus +from websockets.http11 import Response + +from wherobots.db import ( + ConnectionLostError, + ExecutionCancelledError, + ExecutionNotFoundError, +) +from wherobots.db.connection import Connection +from wherobots.db.driver import connect_direct +from wherobots.db.errors import InterfaceError +from wherobots.db.models import Store, StoreResult +from wherobots.db.types import StorageFormat + +PROBE = 0.05 +"""Probe interval used throughout: small, so the tests stay fast.""" + +READ_TIMEOUT = 0.01 + +NOT_FOUND = "Execution not found" + +STORE = Store.for_download(format=StorageFormat.PARQUET) + +RESULT_URI = "https://presigned.example.com/results.parquet" + +WS_CONNECT = "wherobots.db.driver.websockets.sync.client.connect" + + +@pytest.fixture(autouse=True) +def fast_backoff(monkeypatch): + monkeypatch.setattr("wherobots.db.connection._RECONNECT_FIRST_DELAY_SECONDS", 0.01) + monkeypatch.setattr("wherobots.db.connection._RECONNECT_MAX_DELAY_SECONDS", 0.05) + + +class Session: + """A fake SQL session transport. + + Records every request. ``on_retrieve`` answers ``retrieve_results`` + requests (attaches and probes). With ``fail_sends``, a send raises as on a + dead socket. + """ + + def __init__(self, on_retrieve=None): + self.incoming = queue.Queue() + self.sent = [] + self.on_retrieve = on_retrieve + self.fail_sends = False + self.close_timeout = 1.0 + self.socket = MagicMock() + self.socket.shutdown.side_effect = lambda how: self.close() + + def recv(self, timeout): + try: + value = self.incoming.get(timeout=timeout) + except queue.Empty: + raise TimeoutError from None + if isinstance(value, Exception): + raise value + # CBOR, as sql-session sends results: it carries their bytes. + return cbor2.dumps(value) + + def send(self, value): + if self.fail_sends: + raise OSError("broken pipe") + message = json.loads(value) + self.sent.append(message) + if message["kind"] == "retrieve_results" and self.on_retrieve is not None: + self.on_retrieve(self, message) + + def close(self): + self.incoming.put(ConnectionClosedOK(None, None)) + + def drop(self): + """The network path fails: the keepalive gives up on the connection.""" + self.incoming.put(ConnectionClosedError(None, None)) + + def put(self, **message): + self.incoming.put(message) + + def kinds(self): + return [message["kind"] for message in list(self.sent)] + + def requests(self, kind, execution_id=None): + return [ + message + for message in list(self.sent) + if message["kind"] == kind + and (execution_id is None or message["execution_id"] == execution_id) + ] + + +def answer(session, request, state="running", **fields): + """Answer a retrieve_results as sql-session does: execution_result.""" + fields.setdefault("results", None) + session.put( + kind="execution_result", + execution_id=request["execution_id"], + state=state, + **fields, + ) + + +def new_session(state="running", **fields): + """A session that supports attaching: it acknowledges each attach.""" + return Session( + on_retrieve=lambda s, m: answer( + s, m, state, attached=bool(m.get("attach")), **fields + ) + ) + + +class Factory: + """A transport factory handing out ``transports`` in turn. + + Each is a transport to return, an exception to raise, or a callable + producing either. Records the time limit of every call. + """ + + def __init__(self, *transports): + self.transports = iter(transports) + self.calls = [] + + def __call__(self, timeout): + self.calls.append(timeout) + transport = next(self.transports, OSError("connection refused")) + if callable(transport) and not isinstance(transport, Session): + transport = transport() + if isinstance(transport, Exception): + raise transport + return transport + + +def open_connection(session, factory, **kwargs): + kwargs.setdefault("stale_query_probe_seconds", PROBE) + kwargs.setdefault("session_id", "session-1") + kwargs.setdefault("reconnect_timeout_seconds", 3.0) + return Connection( + session, read_timeout=READ_TIMEOUT, transport_factory=factory, **kwargs + ) + + +def arrow_results(**columns): + """A ``results`` payload carrying ``columns`` as an Arrow IPC stream.""" + table = pyarrow.table(columns) + sink = pyarrow.BufferOutputStream() + with pyarrow.ipc.new_stream(sink, table.schema) as writer: + writer.write_table(table) + return {"result_bytes": sink.getvalue().to_pybytes(), "format": "arrow"} + + +def wait_until(predicate, timeout=3.0): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.005) + return predicate() + + +def tracked(conn, execution_id): + return conn._Connection__queries.get(execution_id) + + +def outcome(cursor, timeout=3.0): + """Wait for the cursor's terminal result, then hand it back to the cursor.""" + result = cursor._Cursor__queue.get(timeout=timeout) + cursor._Cursor__queue.put(result) + return result + + +def reader(conn): + return conn._Connection__thread + + +def in_thread(fn): + thread = threading.Thread(target=fn, daemon=True) + thread.start() + return thread + + +# Recovering pending queries + + +def test_lost_query_completes_over_the_new_transport(caplog): + caplog.set_level(logging.INFO) + lost, again = Session(), new_session("running") + factory = Factory(again) + conn = open_connection(lost, factory) + try: + cursor = conn.cursor() + cursor.execute("SELECT * FROM t", store=STORE) + execution_id = cursor.execution_id + lost.drop() + assert wait_until(lambda: again.requests("retrieve_results")) + assert again.sent[0] == { + "kind": "retrieve_results", + "execution_id": execution_id, + "attach": True, + } + again.put( + kind="state_updated", + execution_id=execution_id, + state="succeeded", + result_uri=RESULT_URI, + size=7, + ) + assert outcome(cursor).error is None + assert cursor.get_store_result() == StoreResult(RESULT_URI, 7) + assert len(factory.calls) == 1 + assert lost.socket.shutdown.called + finally: + conn.close() + messages = [r.getMessage() for r in caplog.records] + assert ( + "Lost the connection to SQL session session-1 with 1 pending queries; " + "reconnecting for up to 3s ..." + ) in messages + assert ( + "Reconnected to SQL session session-1; re-attaching 1 pending queries." + ) in messages + + +def test_result_lost_with_the_connection_arrives_in_the_attach_reply(): + # The query succeeded, but the terminal event went to the lost connection. + again = Session( + on_retrieve=lambda s, m: answer( + s, m, "succeeded", attached=True, results=arrow_results(x=[1, 2]) + ) + ) + lost = Session() + conn = open_connection(lost, Factory(again)) + try: + cursor = conn.cursor() + cursor.execute("SELECT x FROM t") + lost.drop() + assert outcome(cursor).error is None + assert list(cursor.fetchall()["x"]) == [1, 2] + finally: + conn.close() + + +def test_every_pending_query_is_reattached(): + lost, again = Session(), new_session("running") + conn = open_connection(lost, Factory(again)) + try: + cursors = [conn.cursor() for _ in range(3)] + for cursor in cursors: + cursor.execute("SELECT 1", store=STORE) + lost.drop() + ids = [cursor.execution_id for cursor in cursors] + assert wait_until(lambda: len(again.requests("retrieve_results")) == 3) + assert {m["execution_id"] for m in again.requests("retrieve_results")} == ( + set(ids) + ) + for execution_id in ids: + again.put( + kind="state_updated", + execution_id=execution_id, + state="succeeded", + result_uri=f"{RESULT_URI}?{execution_id}", + ) + for cursor in cursors: + assert outcome(cursor).error is None + assert cursor.get_store_result().result_uri.endswith(cursor.execution_id) + finally: + conn.close() + + +def test_reattached_query_is_polled_on_an_older_session(): + # An older session ignores the attach: its events go to the lost + # connection, so only polling tells the outcome. + replies = itertools.chain(["running"], itertools.repeat("succeeded")) + + def on_retrieve(s, m): + state = next(replies) + results = arrow_results(x=[3]) if state == "succeeded" else None + s.put( + kind="execution_result", + execution_id=m["execution_id"], + state=state, + results=results, + ) + + lost, again = Session(), Session(on_retrieve=on_retrieve) + conn = open_connection(lost, Factory(again)) + try: + cursor = conn.cursor() + cursor.execute("SELECT x FROM t") + lost.drop() + assert outcome(cursor).error is None + assert list(cursor.fetchall()["x"]) == [3] + assert len(again.requests("retrieve_results")) >= 2 + finally: + conn.close() + + +def test_reattached_query_unknown_to_the_session_raises_not_found(): + again = Session( + on_retrieve=lambda s, m: s.put( + kind="error", execution_id=m["execution_id"], message=NOT_FOUND + ) + ) + lost = Session() + conn = open_connection(lost, Factory(again)) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") + lost.drop() + error = outcome(cursor).error + assert isinstance(error, ExecutionNotFoundError) + assert error.execution_id == cursor.execution_id + assert str(error) == ( + "Execution not found: SQL session session-1 has no execution " + f"{cursor.execution_id} after reconnecting. The statement didn't " + "reach the session before the connection was lost, the session " + "restarted, or the session no longer tracks it." + ) + finally: + conn.close() + + +def test_query_cancelled_by_the_session_while_disconnected_raises(): + # E.g. the session's grace period for orphaned executions ran out. + lost, again = Session(), new_session("cancelled") + conn = open_connection(lost, Factory(again)) + try: + cursor = conn.cursor() + cursor.execute("UPDATE t SET x = 1") + lost.drop() + error = outcome(cursor).error + assert isinstance(error, ExecutionCancelledError) + assert error.execution_id == cursor.execution_id + finally: + conn.close() + + +def test_a_failed_send_hands_its_query_to_the_reconnect(): + again = Session( + on_retrieve=lambda s, m: s.put( + kind="error", execution_id=m["execution_id"], message=NOT_FOUND + ) + ) + lost = Session() + lost.fail_sends = True + conn = open_connection(lost, Factory(again)) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") # Doesn't raise: the query is re-attached. + assert wait_until(lambda: again.requests("retrieve_results")) + assert again.sent[0]["execution_id"] == cursor.execution_id + assert isinstance(outcome(cursor).error, ExecutionNotFoundError) + assert reader(conn).is_alive() + finally: + conn.close() + + +# Eagerness, and sends while reconnecting + + +def test_reconnects_eagerly_with_nothing_pending(): + lost, again = Session(), Session() + factory = Factory(again) + conn = open_connection(lost, factory) + try: + lost.drop() + assert wait_until(lambda: conn._Connection__ws is again) + assert factory.calls + cursor = conn.cursor() + cursor.execute("SELECT 1") + assert again.requests("execute_sql", cursor.execution_id) + assert not lost.requests("execute_sql") + finally: + conn.close() + + +def test_execute_waits_for_the_reconnect_and_follows_the_reattaches(): + release = threading.Event() + entered = threading.Event() + again = new_session("running") + + def slow_connect(): + entered.set() + release.wait(3) + return again + + lost = Session() + conn = open_connection(lost, Factory(slow_connect)) + try: + first = conn.cursor() + first.execute("SELECT 1", store=STORE) + lost.drop() + assert entered.wait(3) + second = conn.cursor() + executing = in_thread(lambda: second.execute("SELECT 2")) + executing.join(0.1) + assert executing.is_alive() # Waiting for the reconnect. + release.set() + executing.join(3) + assert not executing.is_alive() + assert again.kinds()[:2] == ["retrieve_results", "execute_sql"] + assert again.sent[0]["execution_id"] == first.execution_id + assert again.sent[1]["execution_id"] == second.execution_id + assert not lost.requests("execute_sql", second.execution_id) + finally: + release.set() + conn.close() + + +def test_cancel_waits_for_the_reconnect_and_follows_the_attach(): + release = threading.Event() + entered = threading.Event() + again = new_session("running") + + def slow_connect(): + entered.set() + release.wait(3) + return again + + lost = Session() + conn = open_connection(lost, Factory(slow_connect)) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") + execution_id = cursor.execution_id + lost.drop() + assert entered.wait(3) + closing = in_thread(cursor.close) + closing.join(0.1) + assert closing.is_alive() + # Not before the cancel goes out: a cancellation the session reports + # meanwhile isn't this connection's. + assert tracked(conn, execution_id).cancel_requested is False + release.set() + closing.join(3) + assert not closing.is_alive() + assert again.kinds()[:2] == ["retrieve_results", "cancel"] + assert tracked(conn, execution_id).cancel_requested is True + again.put(kind="state_updated", execution_id=execution_id, state="cancelled") + result = outcome(cursor) + assert result.error is None + assert isinstance(result.results, pandas.DataFrame) + assert result.results.empty + finally: + release.set() + conn.close() + + +def test_a_cancel_sent_before_the_loss_is_sent_again_after_the_attach(): + lost, again = Session(), new_session("running") + conn = open_connection(lost, Factory(again)) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") + execution_id = cursor.execution_id + cursor.close() + assert lost.requests("cancel", execution_id) + lost.drop() + assert wait_until(lambda: len(again.sent) >= 2) + assert again.kinds()[:2] == ["retrieve_results", "cancel"] + assert again.sent[1]["execution_id"] == execution_id + finally: + conn.close() + + +# Giving up + + +def test_pending_queries_fail_once_reconnecting_runs_out_of_time(): + lost = Session() + factory = Factory() # Every attempt fails. + conn = open_connection(lost, factory, reconnect_timeout_seconds=0.3) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") + started = time.monotonic() + lost.drop() + error = outcome(cursor).error + elapsed = time.monotonic() - started + assert isinstance(error, ConnectionLostError) + assert error.execution_id == cursor.execution_id + assert 0.25 <= elapsed < 2.0 + assert len(factory.calls) >= 2 + # Each attempt gets only the time left. + assert all(0 < limit <= 0.3 for limit in factory.calls) + assert wait_until(lambda: not reader(conn).is_alive()) + with pytest.raises(ConnectionLostError): + conn.cursor().execute("SELECT 2") + finally: + conn.close() + + +def test_a_refused_reconnect_fails_pending_queries_at_once(): + lost = Session() + factory = Factory(InterfaceError("SQL session refused the connection (HTTP 404)")) + conn = open_connection(lost, factory, reconnect_timeout_seconds=30) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") + lost.drop() + assert isinstance(outcome(cursor).error, ConnectionLostError) + assert len(factory.calls) == 1 + finally: + conn.close() + + +def test_a_transport_lost_again_at_once_continues_the_same_recovery(): + def dropping(): + session = Session() + session.drop() + return session + + lost = Session() + factory = Factory(*[dropping] * 1000) + conn = open_connection(lost, factory, reconnect_timeout_seconds=0.3) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") + lost.drop() + # Without a shared deadline, this would reconnect forever. + assert isinstance(outcome(cursor).error, ConnectionLostError) + assert 2 <= len(factory.calls) < 100 + finally: + conn.close() + + +def test_without_reconnecting_pending_queries_fail_at_once(): + lost = Session() + factory = Factory(Session()) + conn = open_connection(lost, factory, reconnect_timeout_seconds=None) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") + lost.drop() + assert isinstance(outcome(cursor).error, ConnectionLostError) + assert factory.calls == [] + finally: + conn.close() + + +@pytest.mark.parametrize("value", [0, -1, float("nan"), "soon", True]) +def test_an_invalid_reconnect_timeout_disables_reconnecting(value, caplog): + lost = Session() + factory = Factory(Session()) + conn = open_connection(lost, factory, reconnect_timeout_seconds=value) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") + lost.drop() + assert isinstance(outcome(cursor).error, ConnectionLostError) + assert factory.calls == [] + finally: + conn.close() + assert any( + "Invalid reconnect_timeout_seconds" in r.getMessage() for r in caplog.records + ) + + +# close() while reconnecting + + +def test_close_while_backing_off_fails_pending_queries_and_stops(): + lost = Session() + factory = Factory() # Every attempt fails. + conn = open_connection(lost, factory, reconnect_timeout_seconds=30) + cursor = conn.cursor() + cursor.execute("SELECT 1") + lost.drop() + assert wait_until(lambda: len(factory.calls) >= 2) + started = time.monotonic() + conn.close() + assert time.monotonic() - started < 2.0 + assert isinstance(outcome(cursor).error, ConnectionLostError) + assert wait_until(lambda: not reader(conn).is_alive()) + + +def test_close_while_connecting_discards_the_new_transport(): + again = new_session("running") + entered = threading.Event() + + def connect_during_close(): + entered.set() + assert wait_until(lambda: conn._Connection__closed) + return again + + lost = Session() + conn = open_connection(lost, Factory(connect_during_close)) + cursor = conn.cursor() + cursor.execute("SELECT 1") + lost.drop() + assert entered.wait(3) + conn.close() + assert isinstance(outcome(cursor).error, ConnectionLostError) + assert wait_until(lambda: not reader(conn).is_alive()) + assert again.socket.shutdown.called + assert again.sent == [] + + +def test_an_execute_waiting_for_the_reconnect_fails_when_closed(): + release = threading.Event() + entered = threading.Event() + + def slow_connect(): + entered.set() + release.wait(3) + return Session() + + lost = Session() + conn = open_connection(lost, Factory(slow_connect)) + try: + lost.drop() + assert entered.wait(3) + errors = [] + + def execute(): + try: + conn.cursor().execute("SELECT 1") + except Exception as error: + errors.append(error) + + executing = in_thread(execute) + executing.join(0.1) + assert executing.is_alive() + conn.close() + release.set() + executing.join(3) + assert len(errors) == 1 + assert isinstance(errors[0], ConnectionLostError) + finally: + release.set() + conn.close() + + +# The transport factory from connect_direct() + + +def test_connect_direct_reconnects_to_the_same_websocket(): + first, again = Session(), new_session("running") + headers = {"X-API-Key": "key"} + with patch(WS_CONNECT, side_effect=[first, again]) as ws_connect: + conn = connect_direct( + "wss://compute.example.com/sql/org/session-id", + headers=headers, + read_timeout=READ_TIMEOUT, + session_id="session-1", + stale_query_probe_seconds=PROBE, + reconnect_timeout_seconds=5, + ) + try: + headers["X-API-Key"] = "changed" + cursor = conn.cursor() + cursor.execute("SELECT 1", store=STORE) + first.drop() + assert wait_until(lambda: again.requests("retrieve_results")) + initial, reconnect = ws_connect.call_args_list + assert reconnect.kwargs["uri"] == initial.kwargs["uri"] + # What the connection was opened with, not the caller's later edit. + assert reconnect.kwargs["additional_headers"] == {"X-API-Key": "key"} + # A single attempt, bounded by the time left to reconnect in. + assert 0 < reconnect.kwargs["open_timeout"] <= 5 + for key in ("max_size", "close_timeout", "ssl"): + assert reconnect.kwargs[key] == initial.kwargs[key] + finally: + conn.close() + + +def invalid_status(status): + return InvalidStatus(Response(status, "", Headers())) + + +@pytest.mark.parametrize("status", [401, 403, 404]) +def test_connect_direct_gives_up_when_the_session_refuses(status): + first = Session() + with patch(WS_CONNECT, side_effect=[first, invalid_status(status)]) as ws_connect: + conn = connect_direct( + "wss://compute/sql", + read_timeout=READ_TIMEOUT, + reconnect_timeout_seconds=30, + ) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") + first.drop() + assert isinstance(outcome(cursor).error, ConnectionLostError) + assert ws_connect.call_count == 2 + finally: + conn.close() + + +@pytest.mark.parametrize("status", [429, 502, 503]) +def test_connect_direct_retries_a_transient_refusal(status): + first, again = Session(), new_session("running") + with patch( + WS_CONNECT, side_effect=[first, invalid_status(status), again] + ) as ws_connect: + conn = connect_direct( + "wss://compute/sql", + read_timeout=READ_TIMEOUT, + reconnect_timeout_seconds=30, + ) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") + first.drop() + assert wait_until(lambda: again.requests("retrieve_results")) + assert ws_connect.call_count == 3 + finally: + conn.close() + + +def test_connect_direct_does_not_reconnect_once_cancelled(): + cancel = threading.Event() + first = Session() + with patch(WS_CONNECT, side_effect=[first]) as ws_connect: + conn = connect_direct( + "wss://compute/sql", + read_timeout=READ_TIMEOUT, + cancel_event=cancel, + reconnect_timeout_seconds=30, + ) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") + cancel.set() + first.drop() + assert isinstance(outcome(cursor).error, ConnectionLostError) + assert ws_connect.call_count == 1 + finally: + conn.close() diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index 1acc6e6..8bf9b6d 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -1,6 +1,7 @@ import json import logging import math +import random import textwrap import threading import time @@ -16,13 +17,18 @@ import websockets.exceptions import websockets.sync.client -from .constants import DEFAULT_READ_TIMEOUT_SECONDS, DEFAULT_STALE_QUERY_PROBE_SECONDS +from .constants import ( + DEFAULT_READ_TIMEOUT_SECONDS, + DEFAULT_RECONNECT_TIMEOUT_SECONDS, + DEFAULT_STALE_QUERY_PROBE_SECONDS, +) from ._transport import abort_connection, socket_writable from .cursor import Cursor from .errors import ( ConnectionLostError, ExecutionCancelledError, ExecutionNotFoundError, + InterfaceError, NotSupportedError, OperationalError, ProgrammingError, @@ -112,9 +118,20 @@ def _handshake_bound(ws: Any) -> float | None: ) """States in which the driver is waiting on a session event it can re-ask for.""" +_RECONNECT_FIRST_DELAY_SECONDS = 0.5 +"""Wait after the first failed reconnect attempt; it doubles per attempt.""" + +_RECONNECT_MAX_DELAY_SECONDS = 5.0 +"""Upper bound on the wait between reconnect attempts.""" + +_RECONNECT_STABLE_SECONDS = 10.0 +"""A connection lost again sooner than this after reconnecting continues the +same recovery, with its deadline and backoff. Otherwise a session that accepts +the connection and drops it at once would be reconnected to forever.""" + -def _probe_interval(value: Any) -> float | None: - """The validated stale-query probe interval, or ``None`` to disable it. +def _positive_seconds(value: Any) -> float | None: + """A validated interval setting, or ``None`` to disable what it times. Like ``_handshake_bound``, returns ``None`` for a non-numeric (including ``bool``), non-finite or non-positive value. Never raises ``Exception``: @@ -189,9 +206,11 @@ class _Watch: @dataclass class _Attach: - """Bookkeeping for a query re-claimed with ``Cursor.attach()``. + """Bookkeeping for a query re-claimed with an attach. - Such a query was started by another connection. Owned like ``_Watch``. + Such a query was started by another connection (``Cursor.attach()``), or + by this one before it lost its transport and reconnected. Owned like + ``_Watch``. """ reply_outstanding: bool = True @@ -210,6 +229,8 @@ class _Attach: lost connection; only polling it can tell the outcome.""" probed: bool = False """The one probe allowed while ``reply_outstanding`` was sent.""" + resumed: bool = False + """Re-attached by this connection after reconnecting, not by the user.""" @dataclass @@ -225,9 +246,12 @@ class Query: cancel_requested: bool = False """This connection sent a cancel for the query. - Unlike the reader-owned fields, a user thread writes it, under ``__lock`` - and before the cancel is sent, so it's set by the time the session can - report the cancellation. Read under ``__lock`` too.""" + Unlike the reader-owned fields, a user thread writes it, in ``__send`` + under ``__lock``, once the cancel is admitted and just before it goes out: + so it's set by the time the session can report the cancellation, but not + while the cancel still waits (say, for a reconnect), when a cancellation + the session reports can't be this connection's. Read under ``__lock`` + too.""" class Connection: @@ -243,6 +267,12 @@ class Connection: A background thread listens for events from the SQL session, and handles update to the corresponding query state. Queries are tracked by their unique execution ID. + + Given a ``transport_factory`` (``connect()`` and ``connect_direct()`` pass one), the + background thread also recovers from a lost WebSocket: it reconnects to the same SQL + session, for up to ``reconnect_timeout_seconds``, and re-attaches every pending query, so + their results still arrive. Only when that fails do pending queries fail, with + ``ConnectionLostError``. """ def __init__( @@ -256,14 +286,50 @@ def __init__( stale_query_probe_seconds: float | None = DEFAULT_STALE_QUERY_PROBE_SECONDS, *, reconnect_factory: Callable[[], "Connection"] | None = None, + transport_factory: ( + Callable[[float], websockets.sync.client.ClientConnection] | None + ) = None, + reconnect_timeout_seconds: float | None = DEFAULT_RECONNECT_TIMEOUT_SECONDS, ): + """ + :param transport_factory: Opens a new WebSocket to the same SQL session, given + the seconds left to do so in, for automatic reconnection. It raises + ``InterfaceError`` if reconnecting can't succeed (the session refused the + connection, or the caller cancelled), and any other exception for a + failed attempt worth retrying. Without it, a lost transport fails the + pending queries at once. + :param reconnect_timeout_seconds: How long to keep trying to reconnect after + losing the transport. ``None``, or an invalid value (non-numeric, + non-finite, or not positive), disables automatic reconnection. + """ self.__ws = ws self.__reconnect_factory = reconnect_factory + self.__transport_factory = transport_factory + self.__reconnect_within = ( + _positive_seconds(reconnect_timeout_seconds) + if transport_factory is not None + else None + ) + if ( + transport_factory is not None + and self.__reconnect_within is None + and reconnect_timeout_seconds is not None + ): + logging.warning( + "Invalid reconnect_timeout_seconds (%s); automatic reconnection " + "disabled.", + _describe_invalid(reconnect_timeout_seconds), + ) + # The current recovery, which a loss soon after reconnecting continues. + # Only the reader thread touches these. + self.__recovery_deadline = 0.0 + self.__reconnect_delay = _RECONNECT_FIRST_DELAY_SECONDS + self.__reconnected_at: float | None = None self.__results_format = results_format self.__data_compression = data_compression self.__geometry_representation = geometry_representation self.__progress_handler: ProgressHandler | None = None - self.__probe_after = _probe_interval(stale_query_probe_seconds) + self.__probe_after = _positive_seconds(stale_query_probe_seconds) if self.__probe_after is None and stale_query_probe_seconds is not None: logging.warning( "Invalid stale_query_probe_seconds (%s); stale-query probes disabled.", @@ -349,9 +415,10 @@ def reconnect(self) -> "Connection": Connects straight to the session's WebSocket, with the headers and settings this connection was opened with (``connect()``'s, or ``connect_direct()``'s), so a session that is gone fails fast instead - of being replaced by a new one. Use it after a - :class:`~wherobots.db.errors.ConnectionLostError` to re-attach to the - lost execution with ``Cursor.attach()``. + of being replaced by a new one. The connection already reconnects by + itself after a lost transport; use this when that's turned off or has + given up, after a :class:`~wherobots.db.errors.ConnectionLostError`, to + re-attach to the lost execution with ``Cursor.attach()``. The new connection doesn't inherit the progress handler. This connection is left as it is; close it as usual. @@ -387,7 +454,12 @@ def __main_loop(self) -> None: """Main background loop listening for messages from the SQL session.""" logging.info("Starting background connection handling loop...") try: - self.__receive_loop() + # The receive loop returns when the transport is gone: carry on over + # a new one, or stop. + while True: + self.__receive_loop() + if not self.__reconnect(): + return finally: self.__fail_pending() @@ -406,10 +478,10 @@ def __receive_loop(self) -> None: # Expected, retry next time continue except websockets.exceptions.ConnectionClosed: - logging.info("Connection closed; stopping main loop.") + logging.info("SQL session connection closed.") return except _TransportError: - logging.exception("SQL session transport failed; stopping main loop") + logging.exception("SQL session transport failed") return except Exception as e: logging.exception("Error handling message from SQL session", exc_info=e) @@ -423,9 +495,9 @@ def __probe_stale_queries(self) -> None: with the query's current state (and its results or result location once it has succeeded), or with an error for an unknown execution. - A query attached with ``Cursor.attach()`` is probed the same way, which - is how the driver learns its outcome from an older session that - ignores the attach. + A query attached with ``Cursor.attach()``, or re-attached after a + reconnect, is probed the same way, which is how the driver learns its + outcome from an older session that ignores the attach. Runs on the reader thread, which owns every ``Query``'s watchdog fields. Never raises: a watchdog failure must not disturb the receive loop. @@ -503,6 +575,164 @@ def __probe_stale_queries(self) -> None: except Exception: logging.exception("Stale-query probe failed") + def __reconnect(self) -> bool: + """Connect to the same SQL session again after losing the transport. + + Runs on the reader thread once ``__receive_loop`` has returned. Tries, + with backoff, until ``reconnect_timeout_seconds`` has passed, then + re-attaches every pending query over the new transport (see + ``__resume``). Returns ``False`` if the connection must close instead: + reconnecting is off, ``close()`` was called, the session refused the + connection, or time ran out. + + Holds ``__send_lock`` throughout, so user threads' sends wait for the + outcome: an ``execute()`` then goes out over the new transport, after + the re-attaches, or fails with ``ConnectionLostError`` once the + connection is closed. + """ + within = self.__reconnect_within + factory = self.__transport_factory + if within is None or factory is None: + return False + with self.__lock: + if self.__closed: + return False + lost = self.__ws + pending = len(self.__queries) + now = time.monotonic() + # A transport lost soon after reconnecting continues that recovery, + # backing off before its first attempt. + continued = ( + self.__reconnected_at is not None + and now - self.__reconnected_at < _RECONNECT_STABLE_SECONDS + ) + if not continued: + self.__recovery_deadline = now + within + self.__reconnect_delay = _RECONNECT_FIRST_DELAY_SECONDS + logging.warning( + "Lost the connection to SQL session %s with %d pending queries; " + "reconnecting for up to %.0fs ...", + self.__session_id or "(unknown)", + pending, + max(0.0, self.__recovery_deadline - now), + ) + # A sender stalled on the lost transport holds __send_lock: make its + # send fail now. + try: + abort_connection(lost) + except Exception: + logging.debug("Could not abort the lost transport", exc_info=True) + with self.__send_lock: + ws = self.__open_transport(factory, back_off_first=continued) + if ws is None: + return False + return self.__resume(ws) + + def __open_transport( + self, + factory: Callable[[float], websockets.sync.client.ClientConnection], + back_off_first: bool, + ) -> websockets.sync.client.ClientConnection | None: + """Open a new transport, retrying until the recovery's deadline. + + ``None`` if it gave up, or the connection is being closed: ``close()`` + cuts a backoff short, though not an attempt in progress. + """ + session = self.__session_id or "(unknown)" + attempts = 0 + while True: + if attempts or back_off_first: + remaining = self.__recovery_deadline - time.monotonic() + if remaining > 0: + # Jittered, so clients dropped together don't retry together. + delay = self.__reconnect_delay * random.uniform(0.5, 1.0) + self.__reconnect_delay = min( + self.__reconnect_delay * 2, _RECONNECT_MAX_DELAY_SECONDS + ) + if self.__shutdown_done.wait(min(delay, remaining)): + return None + with self.__lock: + if self.__closed: + return None + remaining = self.__recovery_deadline - time.monotonic() + if remaining <= 0: + logging.error( + "Could not reconnect to SQL session %s in time (%d attempts).", + session, + attempts, + ) + return None + attempts += 1 + try: + return factory(remaining) + except InterfaceError as error: + logging.error( + "Could not reconnect to SQL session %s: %s", session, error + ) + return None + except Exception as error: + logging.info( + "Reconnect attempt %d to SQL session %s failed: %s", + attempts, + session, + error, + ) + + def __resume(self, ws: websockets.sync.client.ClientConnection) -> bool: + """Carry on over a new transport, re-attaching every pending query. + + Called on the reader thread with ``__send_lock`` held. Each pending + query is set up as ``Cursor.attach()`` sets up an attached one, keeping + its handler and store, so the replies are handled the same way: a + session that supports attaching makes this connection the owner again, + and an older one is polled. A query this connection cancelled is + cancelled again after its attach, as the cancel may have been lost + with the transport. + + ``False`` if ``close()`` was called while connecting. + """ + with self.__lock: + closed = self.__closed + if not closed: + self.__ws = ws + pending = list(self.__queries.values()) + for query in pending: + # The attach is a retrieve_results; its reply has the state. + query.state = ExecutionState.RESULTS_REQUESTED + query.watch = _Watch() + query.attach = _Attach(resumed=True) + cancelled = [q.execution_id for q in pending if q.cancel_requested] + if closed: + try: + abort_connection(ws) + except Exception: + logging.debug("Could not abort the new transport", exc_info=True) + return False + self.__reconnected_at = time.monotonic() + logging.info( + "Reconnected to SQL session %s; re-attaching %d pending queries.", + self.__session_id or "(unknown)", + len(pending), + ) + requests = [ + self.__retrieve_request(query.execution_id, attach=True) + for query in pending + ] + [ + {"kind": RequestKind.CANCEL.value, "execution_id": execution_id} + for execution_id in cancelled + ] + try: + for request in requests: + ws.send(json.dumps(request)) + except (websockets.exceptions.ConnectionClosed, OSError): + # Lost again already: the receive loop notices, and reconnects. + logging.warning("Lost the new SQL session connection while re-attaching.") + try: + abort_connection(ws) + except Exception: + logging.debug("Could not abort the new transport", exc_info=True) + return True + def __connection_error(self, execution_id: str) -> ConnectionLostError: message = ( f"SQL connection lost (session={self.__session_id or 'unknown'}, " @@ -877,13 +1107,19 @@ def fail_attached(error: OperationalError) -> None: if attach_reply and not_found: # Unlike for a probe, there's no terminal event to wait for: # the session doesn't know the execution at all. + if attach is not None and attach.resumed: + why = ( + "after reconnecting. The statement didn't reach the " + "session before the connection was lost" + ) + else: + why = "to attach to. It ran on a different session" fail_attached( ExecutionNotFoundError( f"Execution not found: SQL session " f"{self.__session_id or '(unknown)'} has no execution " - f"{execution_id} to attach to. It ran on a different " - "session, the session restarted, or the session no " - "longer tracks it.", + f"{execution_id} {why}, the session restarted, or the " + "session no longer tracks it.", session_id=self.__session_id, execution_id=execution_id, ) @@ -964,7 +1200,8 @@ def __send( """Send a request; return ``False`` only if non-blocking and busy. Busy means another sender holds the send lock, or the socket's send - buffer is full. + buffer is full. While the reader reconnects, it holds the send lock, so + a blocking send waits for that to finish. """ # Serialization and redaction are local work. Fail before registration, # without poisoning unrelated cursors or misreporting a transport loss. @@ -995,8 +1232,15 @@ def __send( watch = query.watch watch.last_activity = watch.last_event = time.monotonic() self.__queries[query.execution_id] = query - elif message.get("execution_id") not in self.__queries: - return True + else: + tracked = self.__queries.get(message.get("execution_id")) + if tracked is None: + return True + if message.get("kind") == RequestKind.CANCEL.value: + # Only now that it's going out (see the field). + tracked.cancel_requested = True + # The reader replaces it only while holding the send lock. + ws = self.__ws # A non-blocking send runs on the reader thread (a probe). Not # waiting for the lock isn't enough: with the socket's send buffer # full (a peer that keeps sending but stops reading), ws.send() @@ -1004,10 +1248,10 @@ def __send( # hold the library's protocol mutex, which keepalive pings need # too. Treat it like a busy lock: skip, and retry at the next # check. Checked outside __lock: it's a syscall on the transport. - if not blocking and not socket_writable(self.__ws): + if not blocking and not socket_writable(ws): return False try: - self.__ws.send(request) + ws.send(request) except (websockets.exceptions.ConnectionClosed, OSError): pass # Terminate outside the send gate, before delivering errors. except Exception: @@ -1020,7 +1264,16 @@ def __send( return True finally: self.__send_lock.release() - self.__fail_pending() + if self.__reconnect_within is None: + self.__fail_pending() + return True + # Make sure the reader notices the loss, so it reconnects. A query + # registered above is then re-attached: if the request didn't reach the + # session, that fails it with ExecutionNotFoundError. + try: + abort_connection(ws) + except Exception: + logging.debug("Could not abort the lost transport", exc_info=True) return True @staticmethod @@ -1164,12 +1417,8 @@ def __request_results( def __cancel_query(self, execution_id: str) -> None: """Cancels the query with the given execution ID.""" with self.__lock: - query = self.__queries.get(execution_id) - if query is None: + if execution_id not in self.__queries: return - # Before the request goes out: the session's cancelled state can - # only follow it. - query.cancel_requested = True request = { "kind": RequestKind.CANCEL.value, diff --git a/wherobots/db/constants.py b/wherobots/db/constants.py index 965bd59..8049c75 100644 --- a/wherobots/db/constants.py +++ b/wherobots/db/constants.py @@ -19,6 +19,11 @@ DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS: float = 900 # Silence, per query, before the driver re-asks the session for its state. DEFAULT_STALE_QUERY_PROBE_SECONDS: float = 30.0 +# How long the driver keeps trying to reconnect to the SQL session after losing +# the connection. It must fit, along with the up to ~40s that keepalive takes +# to notice a dead connection, in the session's grace period for orphaned +# executions (120s by default), after which the session cancels them. +DEFAULT_RECONNECT_TIMEOUT_SECONDS: float = 60.0 MAX_MESSAGE_SIZE: int = 100 * 2**20 # 100MiB PROTOCOL_VERSION: Version = Version("1.0.0") diff --git a/wherobots/db/driver.py b/wherobots/db/driver.py index 26b97cd..abe3473 100644 --- a/wherobots/db/driver.py +++ b/wherobots/db/driver.py @@ -23,6 +23,7 @@ DEFAULT_CLOSE_TIMEOUT_SECONDS, DEFAULT_ENDPOINT, DEFAULT_READ_TIMEOUT_SECONDS, + DEFAULT_RECONNECT_TIMEOUT_SECONDS, DEFAULT_SESSION_TYPE, DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS, DEFAULT_STALE_QUERY_PROBE_SECONDS, @@ -127,6 +128,7 @@ def connect( cancel_event: Union[threading.Event, None] = None, extra_headers: Union[Dict[str, str], None] = None, stale_query_probe_seconds: Union[float, None] = DEFAULT_STALE_QUERY_PROBE_SECONDS, + reconnect_timeout_seconds: Union[float, None] = DEFAULT_RECONNECT_TIMEOUT_SECONDS, ) -> Connection: """Create a connection to a Wherobots SQL session. @@ -154,6 +156,12 @@ def connect( back off to at most 8 times this interval. ``None`` disables probing; an invalid value (non-numeric, non-finite, or not positive) also disables it. + :param reconnect_timeout_seconds: How long the connection keeps trying to + reconnect to the same SQL session after losing its WebSocket, then + re-attaches its pending queries, whose results still arrive. Only if + it can't reconnect in time do they fail, with ``ConnectionLostError``. + ``None`` disables reconnecting, as does an invalid value (non-numeric, + non-finite, or not positive): pending queries then fail at once. """ if not token and not api_key: raise ValueError("At least one of `token` or `api_key` is required") @@ -280,6 +288,7 @@ def get_session_uri() -> str: .path.rstrip("/") .rsplit("/", 1)[-1], stale_query_probe_seconds=stale_query_probe_seconds, + reconnect_timeout_seconds=reconnect_timeout_seconds, ) @@ -309,18 +318,36 @@ def connect_direct( cancel_event: Union[threading.Event, None] = None, session_id: str | None = None, stale_query_probe_seconds: Union[float, None] = DEFAULT_STALE_QUERY_PROBE_SECONDS, + reconnect_timeout_seconds: Union[float, None] = DEFAULT_RECONNECT_TIMEOUT_SECONDS, ) -> Connection: """Connect to a SQL session's WebSocket endpoint directly. - The returned connection's ``reconnect()`` connects to the same endpoint - again, with the same arguments (``cancel_event`` included). + After losing its WebSocket, the returned connection reconnects to the same + endpoint by itself (see ``reconnect_timeout_seconds``). Its ``reconnect()`` + connects to the same endpoint again, with the same arguments + (``cancel_event`` included). :param stale_query_probe_seconds: See :func:`connect`. + :param reconnect_timeout_seconds: See :func:`connect`. """ uri_with_protocol = f"{uri}/{protocol}" ssl_context = ssl.create_default_context() ssl_context.load_verify_locations(certifi.where()) + def open_ws( + headers: Union[Dict[str, str], None], open_timeout: float + ) -> websockets.sync.client.ClientConnection: + _check_cancelled(cancel_event) + logging.info("Connecting to SQL session at %s ...", uri_with_protocol) + return websockets.sync.client.connect( + uri=uri_with_protocol, + additional_headers=headers, + max_size=MAX_MESSAGE_SIZE, + open_timeout=open_timeout, + close_timeout=DEFAULT_CLOSE_TIMEOUT_SECONDS, + ssl=ssl_context, + ) + @tenacity.retry( stop=tenacity.stop_after_attempt(5), wait=tenacity.wait_exponential(multiplier=1, min=1, max=5), @@ -336,16 +363,7 @@ def connect_direct( reraise=True, ) def ws_connect() -> websockets.sync.client.ClientConnection: - _check_cancelled(cancel_event) - logging.info("Connecting to SQL session at %s ...", uri_with_protocol) - return websockets.sync.client.connect( - uri=uri_with_protocol, - additional_headers=headers, - max_size=MAX_MESSAGE_SIZE, - open_timeout=DEFAULT_HTTP_TIMEOUT, - close_timeout=DEFAULT_CLOSE_TIMEOUT_SECONDS, - ssl=ssl_context, - ) + return open_ws(headers, DEFAULT_HTTP_TIMEOUT) try: ws = ws_connect() @@ -358,6 +376,20 @@ def ws_connect() -> websockets.sync.client.ClientConnection: # sent even if the caller later mutates it. reconnect_headers = dict(headers) if isinstance(headers, dict) else headers + def open_transport(timeout: float) -> websockets.sync.client.ClientConnection: + # A single attempt: the connection retries until its own deadline. + # Straight to the session's WebSocket, as for reconnect() below. + try: + return open_ws(reconnect_headers, min(DEFAULT_HTTP_TIMEOUT, timeout)) + except websockets.exceptions.InvalidStatus as e: + status = e.response.status_code + if 400 <= status < 500 and status not in (408, 429): + # Gone, or no longer authorized: retrying can't help. + raise InterfaceError( + f"SQL session refused the connection (HTTP {status})" + ) from e + raise + def reconnect() -> Connection: # Straight to the same session's WebSocket: connect() resolves the # session through the REST API, which could start a new session when @@ -373,6 +405,7 @@ def reconnect() -> Connection: cancel_event=cancel_event, session_id=session_id, stale_query_probe_seconds=stale_query_probe_seconds, + reconnect_timeout_seconds=reconnect_timeout_seconds, ) return Connection( @@ -384,4 +417,6 @@ def reconnect() -> Connection: session_id=session_id, stale_query_probe_seconds=stale_query_probe_seconds, reconnect_factory=reconnect, + transport_factory=open_transport, + reconnect_timeout_seconds=reconnect_timeout_seconds, ) diff --git a/wherobots/db/errors.py b/wherobots/db/errors.py index 0926bb7..b6d57c2 100644 --- a/wherobots/db/errors.py +++ b/wherobots/db/errors.py @@ -29,6 +29,9 @@ class NotSupportedError(DatabaseError): class ConnectionLostError(OperationalError): """The connection to the SQL session was lost while a query was pending. + A connection from ``connect()`` or ``connect_direct()`` first reconnects + by itself and re-attaches its pending queries; this means that failed + within ``reconnect_timeout_seconds``, or that reconnecting is turned off. The query may still be running on the SQL session. To recover its outcome, open a new connection to the same session with ``Connection.reconnect()`` and re-attach to the execution with ``Cursor.attach(execution_id)``. From f2406787a30f5787f50c63dbba4942e04850c56b Mon Sep 17 00:00:00 2001 From: Kevin Date: Mon, 19 Oct 2026 15:48:46 -0700 Subject: [PATCH 2/2] refactor: drop the manual recovery path, Connection.reconnect() and Cursor.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 --- README.md | 45 +---- tests/test_attach.py | 392 ++++++------------------------------- wherobots/db/connection.py | 112 ++--------- wherobots/db/cursor.py | 69 ++----- wherobots/db/driver.py | 28 +-- wherobots/db/errors.py | 20 +- 6 files changed, 113 insertions(+), 553 deletions(-) diff --git a/README.md b/README.md index 62d8511..b788513 100644 --- a/README.md +++ b/README.md @@ -233,48 +233,15 @@ A re-attached query can still fail: fails the query with an `OperationalError`. With probes disabled, a running execution on such a session fails right away. -#### When reconnecting fails - If the connection can't reconnect in time, or reconnecting is turned off (`reconnect_timeout_seconds=None`), fetching a pending query's outcome raises `ConnectionLostError`, another `OperationalError` subclass. The -query may still be running on the SQL session; the error's -`execution_id` and `session_id` attributes identify it. A cursor's -`execution_id` property also holds the ID of its current or last -execution. (When `execute()` itself raises `ConnectionLostError`, the -connection was already lost and the statement was never sent.) - -To recover the outcome yourself, open a new connection to the same SQL -session with `Connection.reconnect()`, and re-attach a cursor to the -execution with `Cursor.attach()`: - -```python -from wherobots.db import ConnectionLostError, Store, connect - -store = Store.for_download() -with connect(api_key='...') as conn: - curr = conn.cursor() - curr.execute("SELECT ...", store=store) - try: - store_result = curr.get_store_result() - except ConnectionLostError as e: - with conn.reconnect() as again: - curr = again.cursor() - curr.attach(e.execution_id, store=store) - store_result = curr.get_store_result() -``` - -* `reconnect()` connects straight to the same SQL session's WebSocket, - with the headers and settings of the original connection, including - its `cancel_event` (but not its progress handler). Unlike `connect()`, - it never starts a new session: if the session is gone, it fails with - an `InterfaceError`. -* `attach(execution_id, store=None)` takes the `Store` the execution was - started with, if any. Afterwards, fetch the outcome as after - `execute()`, with the same failures as above. Like `execute()`, it - first cancels the cursor's own in-flight execution, and the cursor - then owns the attached one: closing the cursor, or executing on it, - before the outcome arrives cancels the execution. +query may still have run, or be running, on the SQL session: its outcome, +including whether a write committed, is unknown. The error's +`execution_id` and `session_id` attributes identify it, and a cursor's +`execution_id` property holds the ID of its current or last execution. +(When `execute()` itself raises `ConnectionLostError`, the connection was +already lost and the statement was never sent.) ### Runtime and region selection diff --git a/tests/test_attach.py b/tests/test_attach.py index 9eba4b4..10df3fa 100644 --- a/tests/test_attach.py +++ b/tests/test_attach.py @@ -1,9 +1,8 @@ -"""Recovering a query whose connection was lost: reconnect, then re-attach.""" +"""Re-attaching pending queries after a reconnect: handling the replies.""" import json import logging import queue -import threading import time from unittest.mock import MagicMock, patch @@ -16,14 +15,9 @@ ConnectionLostError, ExecutionCancelledError, ExecutionNotFoundError, - NotSupportedError, OperationalError, - ProgrammingError, ) from wherobots.db.connection import Connection -from wherobots.db.cursor import Cursor -from wherobots.db.driver import connect, connect_direct -from wherobots.db.errors import InterfaceError from wherobots.db.models import Store, StoreResult from wherobots.db.types import ( DataCompression, @@ -47,10 +41,8 @@ RESULT_URI = "https://presigned.example.com/results.parquet" -EID = "exec-started-elsewhere" -"""An execution started by another (lost) connection.""" - -WS_CONNECT = "wherobots.db.driver.websockets.sync.client.connect" +EID = "exec-pending-when-lost" +"""The execution pending when the connection lost its transport.""" class Session: @@ -155,6 +147,31 @@ def open_connection(session, **kwargs): return Connection(session, read_timeout=READ_TIMEOUT, **kwargs) +def reattach(session, store=None, **kwargs): + """A connection whose query EID was pending when its transport was lost. + + The connection reconnects to ``session``, its only transport to come, and + re-attaches the query there. Returns the connection and the query's + cursor. + """ + transports = iter([session]) + + def transport_factory(timeout): + transport = next(transports, None) + if transport is None: + raise OSError("connection refused") + return transport + + kwargs.setdefault("reconnect_timeout_seconds", 0.5) + lost = Session() + conn = open_connection(lost, transport_factory=transport_factory, **kwargs) + cursor = conn.cursor() + with patch("wherobots.db.connection.uuid.uuid4", return_value=EID): + cursor.execute("SELECT * FROM t", store=store) + lost.incoming.put(ConnectionClosedError(None, None)) + return conn, cursor + + def tracked(conn, execution_id): return conn._Connection__queries.get(execution_id) @@ -238,25 +255,24 @@ def test_execution_id_tracks_the_current_execution(): with pytest.raises(TypeError): cursor.execute("SELECT 2", store=bad) assert cursor.execution_id is None - cursor.attach(EID) - assert cursor.execution_id == EID finally: conn.close() -# Attaching on a session that supports it +# Re-attaching on a session that supports it def test_attach_request_carries_the_connection_settings(): session = new_session() - conn = open_connection( + conn, _ = reattach( session, + store=STORE, results_format=ResultsFormat.JSON, data_compression=DataCompression.BROTLI, geometry_representation=GeometryRepresentation.WKB, ) try: - conn.cursor().attach(EID, store=STORE) + assert wait_until(session.retrieves) assert session.retrieves()[0] == { "kind": "retrieve_results", "execution_id": EID, @@ -273,10 +289,8 @@ def test_attach_request_carries_the_connection_settings(): def test_attach_on_new_session_completes_from_pushed_event(caplog): caplog.set_level(logging.INFO) session = new_session("running") - conn = open_connection(session) + conn, cursor = reattach(session, store=STORE) try: - cursor = conn.cursor() - cursor.attach(EID, store=STORE) assert wait_until(lambda: tracked(conn, EID).state == ExecutionState.RUNNING) assert tracked(conn, EID).attach.acknowledged assert deliveries(cursor) == 0 @@ -305,10 +319,8 @@ def test_attach_reply_with_result_uri_completes(attached): s, m, "succeeded", result_uri=RESULT_URI, size=9, **fields ) ) - conn = open_connection(session) + conn, cursor = reattach(session, store=STORE) try: - cursor = conn.cursor() - cursor.attach(EID, store=STORE) assert outcome(cursor).error is None assert cursor.get_store_result() == StoreResult(RESULT_URI, 9) assert tracked(conn, EID) is None @@ -331,10 +343,8 @@ def on_retrieve(s, m): ) session = Session(on_retrieve=on_retrieve) - conn = open_connection(session) + conn, cursor = reattach(session) try: - cursor = conn.cursor() - cursor.attach(EID) assert outcome(cursor).error is None assert cursor.fetchall()["x"].tolist() == [1, 2] assert cursor.rowcount == 2 @@ -344,10 +354,8 @@ def on_retrieve(s, m): def test_acknowledged_empty_store_result_completes_empty(): session = new_session("succeeded") - conn = open_connection(session) + conn, cursor = reattach(session, store=STORE) try: - cursor = conn.cursor() - cursor.attach(EID, store=STORE) result = outcome(cursor) assert result.error is None assert result.store_result is None @@ -380,10 +388,8 @@ def on_retrieve(s, m): s.incoming.put(message) session = Session(on_retrieve=on_retrieve) - conn = open_connection(session) + conn, cursor = reattach(session, store=STORE) try: - cursor = conn.cursor() - cursor.attach(EID, store=STORE) assert outcome(cursor).store_result.result_uri == RESULT_URI assert wait_until(session.incoming.empty) time.sleep(PROBE * 4) @@ -414,10 +420,8 @@ def on_retrieve(s, m): s.incoming.put(rows) session = Session(on_retrieve=on_retrieve) - conn = open_connection(session, results_format=ResultsFormat.JSON) + conn, cursor = reattach(session, results_format=ResultsFormat.JSON) try: - cursor = conn.cursor() - cursor.attach(EID) assert outcome(cursor).results == [{"x": 1}] assert wait_until(lambda: len(session.retrieves()) == 2) assert wait_until(session.incoming.empty) @@ -438,10 +442,8 @@ def on_retrieve(s, m): s.put(kind="error", execution_id=EID, message=NOT_FOUND) session = Session(on_retrieve=on_retrieve) - conn = open_connection(session) + conn, cursor = reattach(session, store=STORE) try: - cursor = conn.cursor() - cursor.attach(EID, store=STORE) assert wait_until(lambda: len(session.retrieves()) == 2) assert wait_until(lambda: not tracked(conn, EID).watch.probeable) assert deliveries(cursor) == 0 @@ -458,10 +460,8 @@ def on_retrieve(s, m): def test_acknowledged_attach_waits_with_probes_disabled(): session = new_session("running") - conn = open_connection(session, stale_query_probe_seconds=None) + conn, cursor = reattach(session, store=STORE, stale_query_probe_seconds=None) try: - cursor = conn.cursor() - cursor.attach(EID, store=STORE) assert wait_until(lambda: tracked(conn, EID).state == ExecutionState.RUNNING) time.sleep(PROBE * 4) assert deliveries(cursor) == 0 @@ -493,10 +493,8 @@ def assert_cancelled_elsewhere(error): def test_attach_to_an_already_cancelled_execution_raises(make_session): # E.g. the session's orphan grace period ran out before the reconnect. session = make_session("cancelled") - conn = open_connection(session) + conn, cursor = reattach(session, store=STORE) try: - cursor = conn.cursor() - cursor.attach(EID, store=STORE) assert_cancelled_elsewhere(outcome(cursor).error) with pytest.raises(ExecutionCancelledError): cursor.get_store_result() @@ -507,10 +505,8 @@ def test_attach_to_an_already_cancelled_execution_raises(make_session): def test_cancelled_pushed_after_attach_raises(): session = new_session("running") - conn = open_connection(session) + conn, cursor = reattach(session) try: - cursor = conn.cursor() - cursor.attach(EID) assert wait_until(lambda: tracked(conn, EID).state == ExecutionState.RUNNING) session.put(kind="state_updated", execution_id=EID, state="cancelled") assert_cancelled_elsewhere(outcome(cursor).error) @@ -522,10 +518,8 @@ def test_cancelled_pushed_after_attach_raises(): def test_cancelled_probe_reply_after_attach_on_old_session_raises(): states = iter(["running"]) session = Session(on_retrieve=lambda s, m: answer(s, m, next(states, "cancelled"))) - conn = open_connection(session) + conn, cursor = reattach(session, store=STORE) try: - cursor = conn.cursor() - cursor.attach(EID, store=STORE) assert_cancelled_elsewhere(outcome(cursor).error) assert len(session.retrieves()) == 2 finally: @@ -540,10 +534,8 @@ def on_retrieve(s, m): answer(s, m, state, **fields) session = Session(on_retrieve=on_retrieve) - conn = open_connection(session) + conn, cursor = reattach(session, store=STORE) try: - cursor = conn.cursor() - cursor.attach(EID, store=STORE) assert wait_until(lambda: tracked(conn, EID).state == ExecutionState.RUNNING) cursor.close() assert [m["execution_id"] for _, m in session.requests("cancel")] == [EID] @@ -585,10 +577,8 @@ def on_retrieve(s, m): s.put(kind="error", execution_id=EID, message="Table not found: t") session = Session(on_retrieve=on_retrieve) - conn = open_connection(session) + conn, cursor = reattach(session) try: - cursor = conn.cursor() - cursor.attach(EID) error = outcome(cursor).error assert type(error) is OperationalError assert str(error) == "Table not found: t" @@ -596,7 +586,7 @@ def on_retrieve(s, m): conn.close() -# Attaching on an older session, which ignores the attach +# Re-attaching on an older session, which ignores the attach @pytest.mark.parametrize("store", [STORE, None]) @@ -624,10 +614,8 @@ def on_retrieve(s, m): ) session = Session(on_retrieve=on_retrieve) - conn = open_connection(session) + conn, cursor = reattach(session, store=store) try: - cursor = conn.cursor() - cursor.attach(EID, store=store) result = outcome(cursor, timeout=5) assert result.error is None if store is not None: @@ -657,10 +645,8 @@ def on_retrieve(s, m): answer(s, m, next(replies, "succeeded")) session = Session(on_retrieve=on_retrieve) - conn = open_connection(session) + conn, cursor = reattach(session, store=STORE) try: - cursor = conn.cursor() - cursor.attach(EID, store=STORE) error = outcome(cursor).error assert type(error) is OperationalError assert "succeeded" in str(error) @@ -683,10 +669,8 @@ def on_retrieve(s, m): s.put(kind="error", execution_id=EID, message=NOT_FOUND) session = Session(on_retrieve=on_retrieve) - conn = open_connection(session) + conn, cursor = reattach(session, store=STORE) try: - cursor = conn.cursor() - cursor.attach(EID, store=STORE) error = outcome(cursor).error assert isinstance(error, ExecutionNotFoundError) assert "no longer tracks" in str(error) @@ -699,10 +683,8 @@ def on_retrieve(s, m): @pytest.mark.parametrize("state", ["running", "pending"]) def test_old_session_running_attach_fails_when_probes_are_disabled(state): session = old_session(state) - conn = open_connection(session, stale_query_probe_seconds=None) + conn, cursor = reattach(session, store=STORE, stale_query_probe_seconds=None) try: - cursor = conn.cursor() - cursor.attach(EID, store=STORE) error = outcome(cursor).error assert type(error) is OperationalError assert "probes are disabled" in str(error) @@ -711,7 +693,7 @@ def test_old_session_running_attach_fails_when_probes_are_disabled(state): conn.close() -# Attaching to an execution the session doesn't know +# Re-attaching to an execution the session doesn't know def test_attach_not_found_raises_execution_not_found(): @@ -720,16 +702,14 @@ def test_attach_not_found_raises_execution_not_found(): kind="error", execution_id=m["execution_id"], message=NOT_FOUND ) ) - conn = open_connection(session) + conn, cursor = reattach(session, store=STORE) try: - cursor = conn.cursor() - cursor.attach(EID, store=STORE) error = outcome(cursor).error assert isinstance(error, ExecutionNotFoundError) assert isinstance(error, OperationalError) assert error.session_id == "session-1" assert error.execution_id == EID - assert f"no execution {EID} to attach to" in str(error) + assert f"no execution {EID} after reconnecting" in str(error) with pytest.raises(ExecutionNotFoundError): cursor.get_store_result() assert tracked(conn, EID) is None @@ -747,13 +727,11 @@ def on_retrieve(s, m): s.put(kind="error", execution_id=EID, message=NOT_FOUND) session = Session(on_retrieve=on_retrieve) - conn = open_connection(session) + conn, cursor = reattach(session) try: - cursor = conn.cursor() - cursor.attach(EID) error = outcome(cursor).error assert isinstance(error, ExecutionNotFoundError) - assert "to attach to" in str(error) + assert "after reconnecting" in str(error) assert [r.get("attach") for r in session.retrieves()] == [True, True] finally: conn.close() @@ -771,10 +749,8 @@ def on_retrieve(s, m): s.put(kind="error", execution_id=EID, message=NOT_FOUND) session = Session(on_retrieve=on_retrieve) - conn = open_connection(session) + conn, cursor = reattach(session, store=STORE) try: - cursor = conn.cursor() - cursor.attach(EID, store=STORE) if acknowledged: # Its terminal event is pushed here: keep waiting for it. assert wait_until(lambda: not tracked(conn, EID).watch.probeable) @@ -799,10 +775,8 @@ def on_retrieve(s, m): def test_unanswered_attach_gets_one_probe_after_4n_that_repeats_it(): session = Session() # Doesn't answer, at first. - conn = open_connection(session) + conn, cursor = reattach(session, store=STORE) try: - cursor = conn.cursor() - cursor.attach(EID, store=STORE) assert wait_until(lambda: len(session.retrieves()) >= 2) attached_at, probed_at = session.retrieve_times() assert probed_at - attached_at >= PROBE * 4 @@ -825,40 +799,14 @@ def test_unanswered_attach_gets_one_probe_after_4n_that_repeats_it(): conn.close() -# Cursor.attach semantics - - -def test_attach_cancels_the_cursors_own_in_flight_execution_first(): - session = new_session("running") - conn = open_connection(session) - try: - cursor = conn.cursor() - cursor.execute("SELECT 1") - previous = cursor.execution_id - cursor.attach(EID, store=STORE) - (_, cancel), *_ = session.requests("cancel") - assert cancel["execution_id"] == previous - # The superseded execution's late result can't leak into the attach. - session.put(kind="state_updated", execution_id=previous, state="cancelled") - session.put( - kind="state_updated", - execution_id=EID, - state="succeeded", - result_uri=RESULT_URI, - ) - assert outcome(cursor).store_result.result_uri == RESULT_URI - assert deliveries(cursor) == 1 - finally: - conn.close() +# The cursor's part @pytest.mark.parametrize("then", ["close", "execute"]) -def test_the_cursor_owns_the_attached_execution(then): +def test_the_cursor_still_owns_the_reattached_execution(then): session = new_session("running") - conn = open_connection(session) + conn, cursor = reattach(session, store=STORE) try: - cursor = conn.cursor() - cursor.attach(EID, store=STORE) assert wait_until(lambda: tracked(conn, EID).state == ExecutionState.RUNNING) if then == "close": cursor.close() @@ -869,77 +817,10 @@ def test_the_cursor_owns_the_attached_execution(then): conn.close() -def test_attach_to_the_cursors_own_in_flight_execution_is_rejected(): - session = Session() - conn = open_connection(session) - try: - cursor = conn.cursor() - cursor.execute("SELECT 1") - own = cursor.execution_id - with pytest.raises(ProgrammingError, match="already waiting"): - cursor.attach(own) - # Not cancelled, and still waiting on its result. - assert session.requests("cancel") == [] - assert cursor.execution_id == own - session.put(kind="state_updated", execution_id=own, state="cancelled") - assert outcome(cursor).error is None - finally: - conn.close() - - -def test_attach_to_an_execution_tracked_by_another_cursor_is_rejected(): - session = Session() - conn = open_connection(session) - try: - other = conn.cursor() - other.execute("SELECT 1") - cursor = conn.cursor() - with pytest.raises(ProgrammingError, match="already tracking"): - cursor.attach(other.execution_id) - assert cursor.execution_id is None - assert session.retrieves() == [] - # The other cursor still gets its result. - session.put( - kind="state_updated", execution_id=other.execution_id, state="cancelled" - ) - assert outcome(other).error is None - finally: - conn.close() - - -@pytest.mark.parametrize("execution_id", ["", None, 42]) -def test_attach_requires_an_execution_id(execution_id): - conn = open_connection(Session()) - try: - with pytest.raises(ProgrammingError): - conn.cursor().attach(execution_id) - finally: - conn.close() - - -def test_cursor_without_an_attach_function_cannot_attach(): - cursor = Cursor(MagicMock(), MagicMock()) - with pytest.raises(NotSupportedError): - cursor.attach(EID) - - -def test_attach_on_a_closed_connection_raises_connection_lost(): - session = Session() - conn = open_connection(session) - conn.close() - cursor = conn.cursor() - with pytest.raises(ConnectionLostError) as exc: - cursor.attach(EID) - assert exc.value.execution_id == EID - assert cursor.execution_id is None - - -def test_transport_loss_while_attached_raises_connection_lost(): +def test_transport_lost_again_without_reconnecting_raises_connection_lost(): session = new_session("running") - conn = open_connection(session) + conn, cursor = reattach(session, store=STORE, reconnect_timeout_seconds=0.2) try: - cursor = conn.cursor() - cursor.attach(EID, store=STORE) assert wait_until(lambda: tracked(conn, EID).state == ExecutionState.RUNNING) session.incoming.put(ConnectionClosedError(None, None)) error = outcome(cursor).error @@ -947,152 +828,3 @@ def test_transport_loss_while_attached_raises_connection_lost(): assert error.execution_id == EID finally: conn.close() - - -# Connection.reconnect - - -def test_reconnect_opens_the_same_session_websocket_with_the_same_settings(): - sessions = [Session(), new_session("running"), Session()] - headers = {"X-API-Key": "key", "X-Trace": "abc"} - with patch(WS_CONNECT, side_effect=sessions) as ws_connect: - conn = connect_direct( - "wss://compute.example.com/sql/org/session-id", - headers=headers, - read_timeout=READ_TIMEOUT, - results_format=ResultsFormat.JSON, - data_compression=DataCompression.BROTLI, - geometry_representation=GeometryRepresentation.WKB, - session_id="session-1", - stale_query_probe_seconds=SLOW_PROBE, - ) - # Not inherited, and a later change to the caller's dict isn't sent. - conn.set_progress_handler(lambda info: None) - headers["X-Trace"] = "changed" - again = conn.reconnect() - third = again.reconnect() - try: - assert again is not conn - calls = ws_connect.call_args_list - assert len(calls) == 3 - # The first connect still gets the caller's headers verbatim. - assert calls[0].kwargs["additional_headers"] is headers - for call in calls: - assert call.kwargs["uri"] == ( - "wss://compute.example.com/sql/org/session-id/1.0.0" - ) - for call in calls[1:]: - assert call.kwargs["additional_headers"] == { - "X-API-Key": "key", - "X-Trace": "abc", - } - ignored = {"ssl", "additional_headers"} - first_kwargs = {k: v for k, v in calls[0].kwargs.items() if k not in ignored} - for call in calls[1:]: - assert { - k: v for k, v in call.kwargs.items() if k not in ignored - } == first_kwargs - assert again._Connection__progress_handler is None - assert again._Connection__probe_after == SLOW_PROBE - assert again._Connection__session_id == "session-1" - # The connection's settings shape the attach request. - cursor = again.cursor() - cursor.attach(EID, store=STORE) - assert sessions[1].retrieves()[0] == { - "kind": "retrieve_results", - "execution_id": EID, - "format": "json", - "compression": "brotli", - "geometry": "wkb", - "attach": True, - } - finally: - for connection in (conn, again, third): - connection.close() - - -@patch("wherobots.db.driver.requests.get") -@patch("wherobots.db.driver.requests.post") -def test_reconnect_after_connect_goes_straight_to_the_websocket(mock_post, mock_get): - post_resp = MagicMock() - post_resp.url = "https://api.example.com/sql/session/test-id" - mock_post.return_value = post_resp - get_resp = MagicMock() - get_resp.json.return_value = { - "status": "READY", - "appMeta": {"url": "https://compute.example.com/sql/org/session-id"}, - } - mock_get.return_value = get_resp - with patch(WS_CONNECT, side_effect=[Session(), Session()]) as ws_connect: - conn = connect(api_key="key", stale_query_probe_seconds=PROBE) - again = conn.reconnect() - try: - # No new session lookup: a session that's gone must fail fast. - assert mock_post.call_count == 1 - assert mock_get.call_count == 1 - first, second = ws_connect.call_args_list - assert second.kwargs["uri"] == first.kwargs["uri"] - assert first.kwargs["uri"] == ( - "wss://compute.example.com/sql/org/session-id/1.0.0" - ) - assert second.kwargs["additional_headers"] == first.kwargs["additional_headers"] - assert second.kwargs["additional_headers"]["X-API-Key"] == "key" - assert again._Connection__session_id == "test-id" - finally: - conn.close() - again.close() - - -def test_reconnect_honours_the_cancel_event(): - cancel = threading.Event() - with patch(WS_CONNECT, side_effect=[Session()]) as ws_connect: - conn = connect_direct("wss://compute/sql", cancel_event=cancel) - cancel.set() - with pytest.raises(InterfaceError, match="cancelled by caller"): - conn.reconnect() - conn.close() - assert ws_connect.call_count == 1 - - -def test_reconnect_without_a_factory_is_not_supported(): - conn = Connection(Session(), read_timeout=READ_TIMEOUT) - try: - with pytest.raises(NotSupportedError): - conn.reconnect() - finally: - conn.close() - - -def test_lost_query_is_recovered_by_reconnecting_and_attaching(): - lost, recovered = Session(), new_session("running") - with patch(WS_CONNECT, side_effect=[lost, recovered]): - conn = connect_direct( - "wss://compute/sql", - read_timeout=READ_TIMEOUT, - session_id="session-1", - stale_query_probe_seconds=PROBE, - # The manual path, for when the connection doesn't reconnect. - reconnect_timeout_seconds=None, - ) - cursor = conn.cursor() - cursor.execute("SELECT * FROM t", store=STORE) - lost.incoming.put(ConnectionClosedError(None, None)) - error = outcome(cursor).error - assert isinstance(error, ConnectionLostError) - conn.close() - again = conn.reconnect() - try: - assert error.execution_id == cursor.execution_id - recovered_cursor = again.cursor() - recovered_cursor.attach(error.execution_id, store=STORE) - recovered.put( - kind="state_updated", - execution_id=error.execution_id, - state="succeeded", - result_uri=RESULT_URI, - size=7, - ) - assert outcome(recovered_cursor).error is None - assert recovered_cursor.get_store_result() == StoreResult(RESULT_URI, 7) - finally: - again.close() diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index 8bf9b6d..ae303c0 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -31,7 +31,6 @@ InterfaceError, NotSupportedError, OperationalError, - ProgrammingError, ) from .models import ExecutionResult, ProgressInfo, Store, StoreResult from .types import ( @@ -105,7 +104,7 @@ def _handshake_bound(ws: Any) -> float | None: Its execution cache is a small LRU that can evict a query that is still running, whose terminal state_updated still arrives (it's pushed from the query's own done-callback). So for a probe, this means "can't answer", not -"failed". For an attach, it means the session doesn't know the execution. +"failed". For a re-attach, it means the session doesn't know the execution. """ _PROBED_STATES = frozenset( @@ -206,17 +205,16 @@ class _Watch: @dataclass class _Attach: - """Bookkeeping for a query re-claimed with an attach. + """Bookkeeping for a query re-attached after reconnecting. - Such a query was started by another connection (``Cursor.attach()``), or - by this one before it lost its transport and reconnected. Owned like + The session sent its events to the lost connection until then. Owned like ``_Watch``. """ reply_outstanding: bool = True """An attach request was sent and no reply to a retrieve has arrived since. - The attach is the query's first request on this connection, so the first + The attach is the query's first request on this transport, so the first reply to a retrieve normally answers it. (After a pushed success, the normal retrieve's reply can come first; it carries the results anyway.) A probe sent while it's outstanding repeats the attach, so the order of @@ -229,8 +227,6 @@ class _Attach: lost connection; only polling it can tell the outcome.""" probed: bool = False """The one probe allowed while ``reply_outstanding`` was sent.""" - resumed: bool = False - """Re-attached by this connection after reconnecting, not by the user.""" @dataclass @@ -242,7 +238,7 @@ class Query: store: Store | None = None watch: _Watch = field(default_factory=_Watch) attach: _Attach | None = None - """Set for a query attached here rather than executed here.""" + """Set once the query is re-attached after reconnecting.""" cancel_requested: bool = False """This connection sent a cancel for the query. @@ -285,7 +281,6 @@ def __init__( session_id: str | None = None, stale_query_probe_seconds: float | None = DEFAULT_STALE_QUERY_PROBE_SECONDS, *, - reconnect_factory: Callable[[], "Connection"] | None = None, transport_factory: ( Callable[[float], websockets.sync.client.ClientConnection] | None ) = None, @@ -303,7 +298,6 @@ def __init__( non-finite, or not positive), disables automatic reconnection. """ self.__ws = ws - self.__reconnect_factory = reconnect_factory self.__transport_factory = transport_factory self.__reconnect_within = ( _positive_seconds(reconnect_timeout_seconds) @@ -407,33 +401,7 @@ def rollback(self) -> None: raise NotSupportedError def cursor(self) -> Cursor: - return Cursor(self.__execute_sql, self.__cancel_query, self.__attach_execution) - - def reconnect(self) -> "Connection": - """Open a new connection to the same SQL session. - - Connects straight to the session's WebSocket, with the headers and - settings this connection was opened with (``connect()``'s, or - ``connect_direct()``'s), so a session that is gone fails fast instead - of being replaced by a new one. The connection already reconnects by - itself after a lost transport; use this when that's turned off or has - given up, after a :class:`~wherobots.db.errors.ConnectionLostError`, to - re-attach to the lost execution with ``Cursor.attach()``. - - The new connection doesn't inherit the progress handler. This - connection is left as it is; close it as usual. - - Raises ``NotSupportedError`` for a connection constructed without a - reconnect factory (not by ``connect()`` or ``connect_direct()``), and - ``InterfaceError`` if the new connection can't be established. - """ - factory = self.__reconnect_factory - if factory is None: - raise NotSupportedError( - "This connection wasn't opened by connect() or connect_direct(); " - "it can't reconnect" - ) - return factory() + return Cursor(self.__execute_sql, self.__cancel_query) def set_progress_handler(self, handler: ProgressHandler | None) -> None: """Register a callback invoked for execution progress events. @@ -495,9 +463,9 @@ def __probe_stale_queries(self) -> None: with the query's current state (and its results or result location once it has succeeded), or with an error for an unknown execution. - A query attached with ``Cursor.attach()``, or re-attached after a - reconnect, is probed the same way, which is how the driver learns its - outcome from an older session that ignores the attach. + A query re-attached after a reconnect is probed the same way, which is + how the driver learns its outcome from an older session that ignores + the attach. Runs on the reader thread, which owns every ``Query``'s watchdog fields. Never raises: a watchdog failure must not disturb the receive loop. @@ -682,10 +650,10 @@ def __resume(self, ws: websockets.sync.client.ClientConnection) -> bool: """Carry on over a new transport, re-attaching every pending query. Called on the reader thread with ``__send_lock`` held. Each pending - query is set up as ``Cursor.attach()`` sets up an attached one, keeping - its handler and store, so the replies are handled the same way: a - session that supports attaching makes this connection the owner again, - and an older one is polled. A query this connection cancelled is + query gets a ``retrieve_results`` with ``attach``, keeping its handler + and store: a session that supports attaching makes this connection the + owner again, pushing the execution's later events here, and answers + with ``attached``; an older one ignores the flag, and is polled. A query this connection cancelled is cancelled again after its attach, as the cancel may have been lost with the transport. @@ -700,7 +668,7 @@ def __resume(self, ws: websockets.sync.client.ClientConnection) -> bool: # The attach is a retrieve_results; its reply has the state. query.state = ExecutionState.RESULTS_REQUESTED query.watch = _Watch() - query.attach = _Attach(resumed=True) + query.attach = _Attach() cancelled = [q.execution_id for q in pending if q.cancel_requested] if closed: try: @@ -824,7 +792,7 @@ def __listen(self) -> None: # Whether this event answers a probe (possibly along with the normal # retrieve). Read before the bookkeeping below resets it. probe_reply = False - # Likewise, whether it answers the attach of a query attached here. + # Likewise, whether it answers the attach of a re-attached query. attach_reply = False attach = query.attach if query is not None else None if query is not None: @@ -1107,19 +1075,14 @@ def fail_attached(error: OperationalError) -> None: if attach_reply and not_found: # Unlike for a probe, there's no terminal event to wait for: # the session doesn't know the execution at all. - if attach is not None and attach.resumed: - why = ( - "after reconnecting. The statement didn't reach the " - "session before the connection was lost" - ) - else: - why = "to attach to. It ran on a different session" fail_attached( ExecutionNotFoundError( f"Execution not found: SQL session " f"{self.__session_id or '(unknown)'} has no execution " - f"{execution_id} {why}, the session restarted, or the " - "session no longer tracks it.", + f"{execution_id} after reconnecting. The statement " + "didn't reach the session before the connection was " + "lost, the session restarted, or the session no longer " + "tracks it.", session_id=self.__session_id, execution_id=execution_id, ) @@ -1220,14 +1183,6 @@ def __send( raise self.__connection_error(query.execution_id) return True if query is not None: - if query.execution_id in self.__queries: - # Only an attach can repeat an ID (executions get fresh - # UUIDs). Replacing the registration would strand the - # first query's handler. - raise ProgrammingError( - "This connection is already tracking execution " - f"{query.execution_id}" - ) # The watchdog's clock starts at registration. watch = query.watch watch.last_activity = watch.last_event = time.monotonic() @@ -1346,35 +1301,6 @@ def __execute_sql( ) return execution_id - def __attach_execution( - self, - execution_id: str, - handler: Callable[[Any], None], - store: Store | None = None, - ) -> str: - """Re-claims an execution started by another connection, by its ID. - - Sends retrieve_results with ``attach``: a session that supports it - makes this connection the execution's owner, which receives its later - events, then answers as to any retrieve_results, with ``attached``. An - older session ignores the flag and only answers. - """ - logging.info("Attaching to execution %s ...", execution_id) - self.__send( - self.__retrieve_request(execution_id, attach=True), - Query( - # Another connection submitted the statement; it's unknown here. - sql="", - execution_id=execution_id, - # The attach is a retrieve_results; its reply has the state. - state=ExecutionState.RESULTS_REQUESTED, - handler=handler, - store=store, - attach=_Attach(), - ), - ) - return execution_id - def __retrieve_request( self, execution_id: str, attach: bool = False ) -> Dict[str, Any]: diff --git a/wherobots/db/cursor.py b/wherobots/db/cursor.py index c26d3a9..30b8e16 100644 --- a/wherobots/db/cursor.py +++ b/wherobots/db/cursor.py @@ -1,9 +1,9 @@ import math import queue import re -from typing import Any, Callable, List, Tuple, Dict +from typing import Any, List, Tuple, Dict -from .errors import NotSupportedError, ProgrammingError +from .errors import ProgrammingError from .models import ExecutionResult, Store, StoreResult # Matches pyformat parameter markers: %(name)s @@ -67,10 +67,9 @@ def replacer(match: re.Match) -> str: class Cursor: - def __init__(self, exec_fn, cancel_fn, attach_fn=None) -> None: + def __init__(self, exec_fn, cancel_fn) -> None: self.__exec_fn = exec_fn self.__cancel_fn = cancel_fn - self.__attach_fn = attach_fn self.__queue: queue.Queue = queue.Queue() self.__results: list[Any] | None = None @@ -100,10 +99,9 @@ def rowcount(self) -> int: def execution_id(self) -> str | None: """The ID of this cursor's current or last execution. - ``None`` before the first ``execute()`` or ``attach()``, and after one - that was rejected before submission. After a - :class:`~wherobots.db.errors.ConnectionLostError`, pass it to - ``attach()`` on a cursor of a reconnected connection. + ``None`` before the first ``execute()``, and after one that was + rejected before submission. Errors about an execution, such as + :class:`~wherobots.db.errors.ConnectionLostError`, carry the same ID. """ return self.__current_execution_id @@ -182,55 +180,6 @@ def execute( parameters: Dict[str, Any] | None = None, store: Store | None = None, ) -> None: - self.__start( - lambda handler: self.__exec_fn( - _substitute_parameters(operation, parameters), handler, store - ) - ) - - def attach(self, execution_id: str, store: Store | None = None) -> None: - """Re-attach to an execution already on the SQL session, by its ID. - - Use it to recover a query whose connection was lost: open a new - connection to the same session with ``Connection.reconnect()``, then - attach one of its cursors to ``ConnectionLostError.execution_id``. - Afterwards, fetch the outcome as after ``execute()``. - - ``store`` must be the ``Store`` the execution was started with, if - any: it tells the driver to expect a stored result. Like ``execute()``, - this first cancels the cursor's own in-flight execution, and from then - on the cursor owns the attached one: closing the cursor, or executing - on it, before the outcome arrives cancels it. - - Fetching raises :class:`~wherobots.db.errors.ExecutionNotFoundError` - if the session doesn't know the execution, for example because it ran - on a different session or the session restarted, and - :class:`~wherobots.db.errors.ExecutionCancelledError` if the session - cancelled it before or while re-attaching (unless this cursor - cancelled it, which completes empty, as for ``execute()``). An older - SQL session - that doesn't support attaching still answers with the execution's - state; the driver then polls it for the outcome (see - ``stale_query_probe_seconds``). If such a session reports no location - for a stored result (it predates reporting it, or the result is - empty), the driver can't tell which, and fails the query with an - ``OperationalError``. - """ - if self.__attach_fn is None: - raise NotSupportedError("This cursor can't attach to executions") - if not isinstance(execution_id, str) or not execution_id: - raise ProgrammingError("execution_id must be a non-empty string") - if execution_id == self.__in_flight_execution_id(): - # Cancelling it first, as a new execution would, would cancel the - # very execution being attached to. - raise ProgrammingError( - f"This cursor is already waiting for execution {execution_id}" - ) - attach_fn = self.__attach_fn - self.__start(lambda handler: attach_fn(execution_id, handler, store)) - - def __start(self, submit: Callable[[Callable[[Any], None]], str]) -> None: - """Reset for a new execution, then submit it with its result handler.""" in_flight = self.__in_flight_execution_id() if in_flight: self.__cancel_fn(in_flight) @@ -250,7 +199,11 @@ def __start(self, submit: Callable[[Callable[[Any], None]], str]) -> None: # A rejected submission must not leave the previous execution ID paired # with this new empty queue (which would make a later fetch wait forever). self.__current_execution_id = None - self.__current_execution_id = submit(self.__queue.put) + self.__current_execution_id = self.__exec_fn( + _substitute_parameters(operation, parameters), + self.__queue.put, + store, + ) def get_store_result(self) -> StoreResult | None: """Get the store result for the last executed query. diff --git a/wherobots/db/driver.py b/wherobots/db/driver.py index abe3473..6c2d234 100644 --- a/wherobots/db/driver.py +++ b/wherobots/db/driver.py @@ -323,9 +323,8 @@ def connect_direct( """Connect to a SQL session's WebSocket endpoint directly. After losing its WebSocket, the returned connection reconnects to the same - endpoint by itself (see ``reconnect_timeout_seconds``). Its ``reconnect()`` - connects to the same endpoint again, with the same arguments - (``cancel_event`` included). + endpoint by itself, with the same headers, honouring ``cancel_event`` (see + ``reconnect_timeout_seconds``). :param stale_query_probe_seconds: See :func:`connect`. :param reconnect_timeout_seconds: See :func:`connect`. @@ -378,7 +377,9 @@ def ws_connect() -> websockets.sync.client.ClientConnection: def open_transport(timeout: float) -> websockets.sync.client.ClientConnection: # A single attempt: the connection retries until its own deadline. - # Straight to the session's WebSocket, as for reconnect() below. + # Straight to the same session's WebSocket: connect() resolves the + # session through the REST API, which could start a new session when + # this one is gone. A dead session must fail fast instead. try: return open_ws(reconnect_headers, min(DEFAULT_HTTP_TIMEOUT, timeout)) except websockets.exceptions.InvalidStatus as e: @@ -390,24 +391,6 @@ def open_transport(timeout: float) -> websockets.sync.client.ClientConnection: ) from e raise - def reconnect() -> Connection: - # Straight to the same session's WebSocket: connect() resolves the - # session through the REST API, which could start a new session when - # this one is gone. A dead session must fail fast instead. - return connect_direct( - uri=uri, - protocol=protocol, - headers=reconnect_headers, - read_timeout=read_timeout, - results_format=results_format, - data_compression=data_compression, - geometry_representation=geometry_representation, - cancel_event=cancel_event, - session_id=session_id, - stale_query_probe_seconds=stale_query_probe_seconds, - reconnect_timeout_seconds=reconnect_timeout_seconds, - ) - return Connection( ws, read_timeout=read_timeout, @@ -416,7 +399,6 @@ def reconnect() -> Connection: geometry_representation=geometry_representation, session_id=session_id, stale_query_probe_seconds=stale_query_probe_seconds, - reconnect_factory=reconnect, transport_factory=open_transport, reconnect_timeout_seconds=reconnect_timeout_seconds, ) diff --git a/wherobots/db/errors.py b/wherobots/db/errors.py index b6d57c2..62eca0e 100644 --- a/wherobots/db/errors.py +++ b/wherobots/db/errors.py @@ -32,11 +32,10 @@ class ConnectionLostError(OperationalError): A connection from ``connect()`` or ``connect_direct()`` first reconnects by itself and re-attaches its pending queries; this means that failed within ``reconnect_timeout_seconds``, or that reconnecting is turned off. - The query may still be running on the SQL session. To recover its outcome, - open a new connection to the same session with ``Connection.reconnect()`` - and re-attach to the execution with ``Cursor.attach(execution_id)``. - Raised by ``Cursor.execute()`` itself, rather than by a fetch, it means - the connection was already lost and the statement was never sent. + The query may still have run, or be running, on the SQL session: its + outcome, for a write whether it committed, is unknown. Raised by + ``Cursor.execute()`` itself, rather than by a fetch, it means the + connection was already lost and the statement was never sent. Attributes: session_id: The SQL session's ID, or None if the connection didn't @@ -56,11 +55,11 @@ def __init__( class ExecutionCancelledError(OperationalError): - """An attached execution was cancelled, and not by this connection. + """A re-attached execution was cancelled, and not by this connection. The SQL session cancelled it before or while this connection re-attached - to it, for example because its orphan grace period ran out, or another - client cancelled it. Its outcome is unknown: for a write, whether it + to it after reconnecting, for example because its grace period for a lost + connection ran out, or another client cancelled it. Its outcome is unknown: for a write, whether it committed. Attributes: @@ -81,9 +80,10 @@ def __init__( class ExecutionNotFoundError(OperationalError): - """The SQL session has no record of the execution to attach to. + """The SQL session has no record of an execution to re-attach to. - The execution ran on a different SQL session, the session restarted, or + After reconnecting, the session didn't know the execution: the statement + didn't reach it before the connection was lost, the session restarted, or it no longer tracks the execution. Its outcome can't be recovered. Attributes: