From 4c034e3e4e5961da977ae0aa16419476dde55e83 Mon Sep 17 00:00:00 2001 From: Kevin Date: Fri, 16 Oct 2026 10:48:24 -0700 Subject: [PATCH 1/2] feat: re-attach to a running execution after a lost connection (DB-760) When the WebSocket to a SQL session drops mid-query, the session keeps running the execution, but the driver had no way to re-claim it (2026-10-14: an MCP query job's socket hit a keepalive timeout and its Spark job ran on for 52 minutes, uncancellable). Add a recovery path: reconnect to the same session, then re-attach to the execution by its ID. - ConnectionLostError(OperationalError), with session_id and execution_id, replaces the plain OperationalError delivered to pending queries on transport loss. Same message, so `except OperationalError` keeps working. - Connection.reconnect() opens a new connection straight to the same session WebSocket, with connect_direct()'s headers and settings, via a factory connect_direct() hands the Connection. Going around the REST API means a dead session fails fast instead of being replaced. A Connection built without a factory raises NotSupportedError. - Cursor.execution_id exposes the current or last execution's ID. - Cursor.attach(execution_id, store=None) resets like execute(), then sends retrieve_results with attach: true. A session that supports it makes this connection the owner and answers with attached: true; an older one ignores the flag and only answers. Reply handling for an attached query: - running/pending: keep waiting. Acknowledged, events are pushed here; otherwise the stale-query watchdog polls it with its usual backoff. - succeeded with results or result_uri: complete as usual. - succeeded, store, no result_uri: complete empty if acknowledged, otherwise fail with an OperationalError, since the location went to the lost connection (or the result is empty) and nothing else will arrive. - "Execution not found" for the attach: ExecutionNotFoundError, without the watchdog's evicted-keep-waiting guard. A later not-found is read by acknowledgement: keep waiting if acknowledged, else ExecutionNotFoundError. - Unacknowledged and still running with probes disabled: fail now, rather than wait forever. The attach is itself a retrieve whose reply may carry the full result, so while it's unanswered the watchdog allows one probe after 4N, as after the normal retrieve. That probe repeats the attach: the session dispatches each request as its own task, so the two replies can arrive in either order, and either then reports the acknowledgement. Attaching an ID this connection already tracks raises ProgrammingError (replacing the registration would strand the first handler), as does attaching the cursor's own in-flight execution, which the reset would otherwise cancel. Co-Authored-By: Claude Opus 5.5 --- README.md | 52 ++ tests/test_attach.py | 1010 ++++++++++++++++++++++++++++++++++++ wherobots/db/__init__.py | 4 + wherobots/db/connection.py | 274 +++++++++- wherobots/db/cursor.py | 69 ++- wherobots/db/driver.py | 25 + wherobots/db/errors.py | 49 ++ 7 files changed, 1459 insertions(+), 24 deletions(-) create mode 100644 tests/test_attach.py diff --git a/README.md b/README.md index d99cfb0..d22e294 100644 --- a/README.md +++ b/README.md @@ -200,6 +200,58 @@ Progress events are best-effort and may not be available for all query types or server versions. The handler is simply not invoked when no 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()`: + +```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()`. 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. +* 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. + ### Runtime and region selection You can chose the Wherobots runtime you want to use using the `runtime` diff --git a/tests/test_attach.py b/tests/test_attach.py new file mode 100644 index 0000000..6e8a6cf --- /dev/null +++ b/tests/test_attach.py @@ -0,0 +1,1010 @@ +"""Recovering a query whose connection was lost: reconnect, then re-attach.""" + +import json +import logging +import queue +import threading +import time +from unittest.mock import MagicMock, patch + +import cbor2 +import pyarrow +import pytest +from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK + +from wherobots.db import ( + ConnectionLostError, + 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, + ExecutionState, + GeometryRepresentation, + ResultsFormat, + StorageFormat, +) + +PROBE = 0.05 +"""Probe interval used throughout: small, so the tests stay fast.""" + +SLOW_PROBE = 0.2 +"""For assertions that something does *not* happen within N.""" + +READ_TIMEOUT = 0.01 + +NOT_FOUND = "Execution not found" + +STORE = Store.for_download(format=StorageFormat.PARQUET) + +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" + + +class Session: + """A fake SQL session transport. + + Records every request with its send time. ``on_retrieve`` answers + ``retrieve_results`` requests: attaches and probes. + """ + + def __init__(self, on_retrieve=None): + self.incoming = queue.Queue() + self.sent = [] + self.on_retrieve = on_retrieve + 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 + if isinstance(value, bytes): + return value + return json.dumps(value) + + def send(self, value): + message = json.loads(value) + self.sent.append((time.monotonic(), 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 put(self, **message): + self.incoming.put(message) + + def requests(self, kind, execution_id=None): + return [ + (sent_at, message) + for sent_at, message in list(self.sent) + if message["kind"] == kind + and (execution_id is None or message["execution_id"] == execution_id) + ] + + def retrieves(self, execution_id=None): + return [m for _, m in self.requests("retrieve_results", execution_id)] + + def retrieve_times(self, execution_id=None): + return [t for t, _ in self.requests("retrieve_results", execution_id)] + + +def answer(session, request, state="running", **fields): + """Answer a retrieve_results as sql-session does: execution_result.""" + session.put( + kind="execution_result", + execution_id=request["execution_id"], + state=state, + results=None, + **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 + ) + ) + + +def old_session(state="running", **fields): + """A session that predates attaching: it ignores the attach flag.""" + return Session(on_retrieve=lambda s, m: answer(s, m, state, **fields)) + + +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 open_connection(session, **kwargs): + kwargs.setdefault("stale_query_probe_seconds", PROBE) + kwargs.setdefault("session_id", "session-1") + return Connection(session, read_timeout=READ_TIMEOUT, **kwargs) + + +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. + + Fetching through the cursor would block forever on a regression; this + fails the test instead, and leaves the cursor's own fetches working. + """ + result = cursor._Cursor__queue.get(timeout=timeout) + cursor._Cursor__queue.put(result) + return result + + +def deliveries(cursor): + return cursor._Cursor__queue.qsize() + + +# ConnectionLostError and Cursor.execution_id + + +@pytest.mark.parametrize( + "loss", [ConnectionClosedError(None, None), OSError("transport lost")] +) +def test_transport_loss_raises_connection_lost_error_with_ids(loss): + session = Session() + conn = open_connection(session) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1", store=STORE) + session.incoming.put(loss) + error = outcome(cursor).error + assert isinstance(error, ConnectionLostError) + assert isinstance(error, OperationalError) + assert error.session_id == "session-1" + assert error.execution_id == cursor.execution_id + assert str(error) == ( + f"SQL connection lost (session=session-1, " + f"execution={cursor.execution_id}). Commit outcome is unknown; " + "verify the operation before retrying writes." + ) + with pytest.raises(ConnectionLostError): + cursor.get_store_result() + # Submitting on the lost connection raises it too. + with pytest.raises(ConnectionLostError) as exc: + conn.cursor().execute("SELECT 2") + assert exc.value.session_id == "session-1" + finally: + conn.close() + + +def test_connection_lost_error_without_a_session_id(): + session = Session() + conn = Connection(session, read_timeout=READ_TIMEOUT) + cursor = conn.cursor() + cursor.execute("SELECT 1") + conn.close() + error = outcome(cursor).error + assert isinstance(error, ConnectionLostError) + assert error.session_id is None + assert error.execution_id == cursor.execution_id + assert "session=unknown" in str(error) + + +def test_execution_id_tracks_the_current_execution(): + session = Session() + conn = open_connection(session) + try: + cursor = conn.cursor() + assert cursor.execution_id is None + cursor.execute("SELECT 1") + first = session.requests("execute_sql")[-1][1]["execution_id"] + assert cursor.execution_id == first + session.put(kind="state_updated", execution_id=first, state="cancelled") + outcome(cursor) + # Still the last execution's once it has completed. + assert cursor.execution_id == first + bad = MagicMock() + bad.to_dict.return_value = {"invalid": object()} + 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 + + +def test_attach_request_carries_the_connection_settings(): + session = new_session() + conn = open_connection( + session, + results_format=ResultsFormat.JSON, + data_compression=DataCompression.BROTLI, + geometry_representation=GeometryRepresentation.WKB, + ) + try: + conn.cursor().attach(EID, store=STORE) + assert session.retrieves()[0] == { + "kind": "retrieve_results", + "execution_id": EID, + "format": "json", + "compression": "brotli", + "geometry": "wkb", + "attach": True, + } + assert session.requests("execute_sql") == [] + finally: + conn.close() + + +def test_attach_on_new_session_completes_from_pushed_event(caplog): + caplog.set_level(logging.INFO) + session = new_session("running") + conn = open_connection(session) + 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 + assert f"Attached to execution {EID}" in caplog.text + assert "doesn't support attaching" not in caplog.text + # The session pushes the terminal event to the new owner. + session.put( + kind="state_updated", + execution_id=EID, + state="succeeded", + result_uri=RESULT_URI, + size=5, + ) + assert outcome(cursor).error is None + assert cursor.get_store_result() == StoreResult(RESULT_URI, 5) + assert tracked(conn, EID) is None + finally: + conn.close() + + +@pytest.mark.parametrize("attached", [True, None]) +def test_attach_reply_with_result_uri_completes(attached): + fields = {"attached": attached} if attached else {} + session = Session( + on_retrieve=lambda s, m: answer( + s, m, "succeeded", result_uri=RESULT_URI, size=9, **fields + ) + ) + conn = open_connection(session) + 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 + finally: + conn.close() + + +def test_attach_reply_with_rows_completes(): + def on_retrieve(s, m): + s.incoming.put( + cbor2.dumps( + { + "kind": "execution_result", + "execution_id": m["execution_id"], + "state": "succeeded", + "results": arrow_results(x=[1, 2]), + "attached": True, + } + ) + ) + + session = Session(on_retrieve=on_retrieve) + conn = open_connection(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 + finally: + conn.close() + + +def test_acknowledged_empty_store_result_completes_empty(): + session = new_session("succeeded") + conn = open_connection(session) + try: + cursor = conn.cursor() + cursor.attach(EID, store=STORE) + result = outcome(cursor) + assert result.error is None + assert result.store_result is None + assert cursor.get_store_result() is None + assert tracked(conn, EID) is None + finally: + conn.close() + + +@pytest.mark.parametrize("pushed_first", [True, False]) +def test_pushed_terminal_event_racing_the_attach_reply_delivers_once(pushed_first): + def on_retrieve(s, m): + pushed = dict( + kind="state_updated", + execution_id=EID, + state="succeeded", + result_uri=RESULT_URI, + size=1, + ) + reply = dict( + kind="execution_result", + execution_id=EID, + state="succeeded", + results=None, + result_uri=RESULT_URI, + size=1, + attached=True, + ) + for message in (pushed, reply) if pushed_first else (reply, pushed): + s.incoming.put(message) + + session = Session(on_retrieve=on_retrieve) + conn = open_connection(session) + 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) + assert deliveries(cursor) == 1 + assert tracked(conn, EID) is None + assert len(session.retrieves()) == 1 + assert not conn._Connection__closed + finally: + conn.close() + + +def test_pushed_success_racing_the_attach_reply_of_a_non_store_query(): + # The pushed state_updated makes the driver request the results normally; + # the attach reply then carries them, and the normal reply is a duplicate. + rows = cbor2.dumps( + { + "kind": "execution_result", + "execution_id": EID, + "state": "succeeded", + "results": {"result_bytes": b'[{"x": 1}]', "format": "json"}, + "attached": True, + } + ) + + def on_retrieve(s, m): + if m.get("attach"): + s.put(kind="state_updated", execution_id=EID, state="succeeded") + s.incoming.put(rows) + + session = Session(on_retrieve=on_retrieve) + conn = open_connection(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) + time.sleep(PROBE * 4) + assert deliveries(cursor) == 1 + assert tracked(conn, EID) is None + finally: + conn.close() + + +def test_acknowledged_attach_survives_a_probe_not_found(): + # Evicted from the cache of a session that supports attaching: it still + # pushes the terminal event here, so keep waiting, as for an execution. + def on_retrieve(s, m): + if m.get("attach"): + answer(s, m, "running", attached=True) + else: + s.put(kind="error", execution_id=EID, message=NOT_FOUND) + + session = Session(on_retrieve=on_retrieve) + conn = open_connection(session) + 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 + session.put( + kind="state_updated", + execution_id=EID, + state="succeeded", + result_uri=RESULT_URI, + ) + assert outcome(cursor).store_result.result_uri == RESULT_URI + finally: + conn.close() + + +def test_acknowledged_attach_waits_with_probes_disabled(): + session = new_session("running") + conn = open_connection(session, 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 + session.put( + kind="state_updated", + execution_id=EID, + state="succeeded", + result_uri=RESULT_URI, + ) + assert outcome(cursor).store_result.result_uri == RESULT_URI + finally: + conn.close() + + +def test_attach_to_a_cancelled_execution_completes_empty(): + session = new_session("cancelled") + conn = open_connection(session) + try: + cursor = conn.cursor() + cursor.attach(EID, store=STORE) + result = outcome(cursor) + assert result.error is None + assert result.results.empty + assert tracked(conn, EID) is None + finally: + conn.close() + + +@pytest.mark.parametrize("failed_reply", [True, False]) +def test_attach_to_a_failed_execution_reports_its_error(failed_reply): + # sql-session answers a retrieve for a failed execution with a failed + # execution_result, then an error event with the details. + def on_retrieve(s, m): + if failed_reply: + answer(s, m, "failed", attached=True) + s.put(kind="error", execution_id=EID, message="Table not found: t") + + session = Session(on_retrieve=on_retrieve) + conn = open_connection(session) + try: + cursor = conn.cursor() + cursor.attach(EID) + error = outcome(cursor).error + assert type(error) is OperationalError + assert str(error) == "Table not found: t" + finally: + conn.close() + + +# Attaching on an older session, which ignores the attach + + +@pytest.mark.parametrize("store", [STORE, None]) +def test_attach_on_old_session_polls_until_succeeded(store, caplog): + caplog.set_level(logging.INFO) + rows = arrow_results(x=[2]) + states = iter(["running", "pending", "running"]) + + def on_retrieve(s, m): + state = next(states, "succeeded") + if state != "succeeded": + answer(s, m, state) + elif store is not None: + answer(s, m, state, result_uri=RESULT_URI, size=3) + else: + s.incoming.put( + cbor2.dumps( + { + "kind": "execution_result", + "execution_id": EID, + "state": state, + "results": rows, + } + ) + ) + + session = Session(on_retrieve=on_retrieve) + conn = open_connection(session) + try: + cursor = conn.cursor() + cursor.attach(EID, store=store) + result = outcome(cursor, timeout=5) + assert result.error is None + if store is not None: + assert cursor.get_store_result() == StoreResult(RESULT_URI, 3) + else: + assert cursor.fetchall()["x"].tolist() == [2] + attach_request, *probes = session.retrieves() + assert attach_request["attach"] is True + assert len(probes) == 3 + assert all("attach" not in probe for probe in probes) + # The replies aren't events of their own: probes back off N, 2N, 4N. + times = session.retrieve_times() + gaps = [later - earlier for earlier, later in zip(times, times[1:])] + for gap, expected in zip(gaps, [1, 2, 4]): + assert gap >= PROBE * expected, gaps + assert "doesn't support attaching; polling" in caplog.text + assert tracked(conn, EID) is None + finally: + conn.close() + + +@pytest.mark.parametrize("after_probe", [False, True]) +def test_old_session_store_result_without_location_fails_clearly(after_probe): + replies = iter(["running"] if after_probe else []) + + def on_retrieve(s, m): + answer(s, m, next(replies, "succeeded")) + + session = Session(on_retrieve=on_retrieve) + conn = open_connection(session) + try: + cursor = conn.cursor() + cursor.attach(EID, store=STORE) + error = outcome(cursor).error + assert type(error) is OperationalError + assert "succeeded" in str(error) + assert "can't be recovered" in str(error) + assert "session=session-1" in str(error) + with pytest.raises(OperationalError): + cursor.get_store_result() + assert len(session.retrieves()) == (2 if after_probe else 1) + assert tracked(conn, EID) is None + finally: + conn.close() + + +def test_old_session_probe_not_found_fails_instead_of_waiting_forever(): + # The terminal event would go to the lost connection, never here. + def on_retrieve(s, m): + if m.get("attach"): + answer(s, m, "running") + else: + s.put(kind="error", execution_id=EID, message=NOT_FOUND) + + session = Session(on_retrieve=on_retrieve) + conn = open_connection(session) + try: + cursor = conn.cursor() + cursor.attach(EID, store=STORE) + error = outcome(cursor).error + assert isinstance(error, ExecutionNotFoundError) + assert "no longer tracks" in str(error) + assert error.execution_id == EID + assert tracked(conn, EID) is None + finally: + conn.close() + + +@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) + 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) + assert tracked(conn, EID) is None + finally: + conn.close() + + +# Attaching to an execution the session doesn't know + + +def test_attach_not_found_raises_execution_not_found(): + session = Session( + on_retrieve=lambda s, m: s.put( + kind="error", execution_id=m["execution_id"], message=NOT_FOUND + ) + ) + conn = open_connection(session) + 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) + with pytest.raises(ExecutionNotFoundError): + cursor.get_store_result() + assert tracked(conn, EID) is None + time.sleep(PROBE * 4) + assert len(session.retrieves()) == 1 + finally: + conn.close() + + +def test_not_found_for_a_repeated_attach_is_not_mistaken_for_eviction(): + # The first attach goes unanswered; the probe that repeats it gets + # not-found. It's an attach reply: fail, don't keep waiting. + def on_retrieve(s, m): + if len(s.retrieves()) > 1: + s.put(kind="error", execution_id=EID, message=NOT_FOUND) + + session = Session(on_retrieve=on_retrieve) + conn = open_connection(session) + try: + cursor = conn.cursor() + cursor.attach(EID) + error = outcome(cursor).error + assert isinstance(error, ExecutionNotFoundError) + assert "to attach to" in str(error) + assert [r.get("attach") for r in session.retrieves()] == [True, True] + finally: + conn.close() + + +@pytest.mark.parametrize("acknowledged", [True, False]) +def test_not_found_for_the_second_of_two_attach_replies(acknowledged): + # The attach and the probe repeating it are both outstanding. The first + # reply says running; the session then evicts the execution, so the + # second says not-found. Neither flag marks the second as a reply. + def on_retrieve(s, m): + if len(s.retrieves()) > 1: + fields = {"attached": True} if acknowledged else {} + answer(s, m, "running", **fields) + s.put(kind="error", execution_id=EID, message=NOT_FOUND) + + session = Session(on_retrieve=on_retrieve) + conn = open_connection(session) + 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) + assert deliveries(cursor) == 0 + session.put( + kind="state_updated", + execution_id=EID, + state="succeeded", + result_uri=RESULT_URI, + ) + assert outcome(cursor).store_result.result_uri == RESULT_URI + else: + error = outcome(cursor).error + assert isinstance(error, ExecutionNotFoundError) + assert "no longer tracks" in str(error) + finally: + conn.close() + + +# Watchdog: an unanswered attach + + +def test_unanswered_attach_gets_one_probe_after_4n_that_repeats_it(): + session = Session() # Doesn't answer, at first. + conn = open_connection(session) + 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 + assert session.retrieves()[1]["attach"] is True + # Never a second one while unanswered: the reply may be a full result. + time.sleep(PROBE * 16) + assert len(session.retrieves()) == 2 + # Once answered, probes resume with backoff, as plain retrieves. + session.put( + kind="execution_result", + execution_id=EID, + state="running", + results=None, + attached=True, + ) + assert wait_until(lambda: len(session.retrieves()) >= 3) + assert "attach" not in session.retrieves()[2] + assert deliveries(cursor) == 0 + finally: + 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() + + +@pytest.mark.parametrize("then", ["close", "execute"]) +def test_the_cursor_owns_the_attached_execution(then): + session = new_session("running") + conn = open_connection(session) + try: + cursor = conn.cursor() + cursor.attach(EID, store=STORE) + assert wait_until(lambda: tracked(conn, EID).state == ExecutionState.RUNNING) + if then == "close": + cursor.close() + else: + cursor.execute("SELECT 2") + assert [m["execution_id"] for _, m in session.requests("cancel")] == [EID] + finally: + 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(): + session = new_session("running") + conn = open_connection(session) + 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 + assert isinstance(error, ConnectionLostError) + 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, + ) + 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/__init__.py b/wherobots/db/__init__.py index 667daf5..f8831d6 100644 --- a/wherobots/db/__init__.py +++ b/wherobots/db/__init__.py @@ -3,7 +3,9 @@ from .driver import connect, connect_direct from .errors import ( Error, + ConnectionLostError, DatabaseError, + ExecutionNotFoundError, InternalError, InterfaceError, OperationalError, @@ -26,6 +28,8 @@ "InternalError", "InterfaceError", "OperationalError", + "ConnectionLostError", + "ExecutionNotFoundError", "ProgrammingError", "NotSupportedError", "Region", diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index 5d53298..e15d7b9 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -19,7 +19,13 @@ from .constants import DEFAULT_READ_TIMEOUT_SECONDS, DEFAULT_STALE_QUERY_PROBE_SECONDS from ._transport import abort_connection, socket_writable from .cursor import Cursor -from .errors import NotSupportedError, OperationalError +from .errors import ( + ConnectionLostError, + ExecutionNotFoundError, + NotSupportedError, + OperationalError, + ProgrammingError, +) from .models import ExecutionResult, ProgressInfo, Store, StoreResult from .types import ( RequestKind, @@ -92,7 +98,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". +"failed". For an attach, it means the session doesn't know the execution. """ _PROBED_STATES = frozenset( @@ -180,6 +186,31 @@ class _Watch: """The one probe allowed after ``results_requested`` was sent.""" +@dataclass +class _Attach: + """Bookkeeping for a query re-claimed with ``Cursor.attach()``. + + Such a query was started by another connection. 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 + 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 + their replies doesn't matter.""" + acknowledged: bool = False + """The session answered ``attached: true``: it pushes the execution's + later events to this connection, and echoes stored result locations. + + An older session ignores the attach and keeps sending the events to the + lost connection; only polling it can tell the outcome.""" + probed: bool = False + """The one probe allowed while ``reply_outstanding`` was sent.""" + + @dataclass class Query: sql: str @@ -188,6 +219,8 @@ class Query: handler: Callable[[Any], None] store: Store | None = None watch: _Watch = field(default_factory=_Watch) + attach: _Attach | None = None + """Set for a query attached here rather than executed here.""" class Connection: @@ -214,8 +247,11 @@ def __init__( geometry_representation: GeometryRepresentation | None = None, session_id: str | None = None, stale_query_probe_seconds: float | None = DEFAULT_STALE_QUERY_PROBE_SECONDS, + *, + reconnect_factory: Callable[[], "Connection"] | None = None, ): self.__ws = ws + self.__reconnect_factory = reconnect_factory self.__results_format = results_format self.__data_compression = data_compression self.__geometry_representation = geometry_representation @@ -298,7 +334,32 @@ def rollback(self) -> None: raise NotSupportedError def cursor(self) -> Cursor: - return Cursor(self.__execute_sql, self.__cancel_query) + 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. Use it 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() def set_progress_handler(self, handler: ProgressHandler | None) -> None: """Register a callback invoked for execution progress events. @@ -355,6 +416,10 @@ 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. + Runs on the reader thread, which owns every ``Query``'s watchdog fields. Never raises: a watchdog failure must not disturb the receive loop. """ @@ -376,17 +441,33 @@ def __probe_stale_queries(self) -> None: watch = query.watch if not watch.probeable or query.state not in _PROBED_STATES: continue + # The query's attach, while it's unanswered. + attach = query.attach + if attach is not None and not attach.reply_outstanding: + attach = None if watch.results_requested: if watch.results_probed: continue wait = probe_after * _RESULTS_PROBE_AFTER_FACTOR + elif attach is not None: + # The attach is a retrieve_results too, whose reply carries + # the whole result if the query has succeeded: like the + # normal retrieve, it gets one probe, after 4N. + if attach.probed: + continue + wait = probe_after * _RESULTS_PROBE_AFTER_FACTOR else: doublings = min(watch.probes_sent, _MAX_PROBE_BACKOFF_DOUBLINGS) wait = probe_after * 2**doublings silence = now - watch.last_activity if silence < wait: continue - if not self.__request_results(query.execution_id, probe=True): + # Until the attach is answered, the probe repeats it: the + # session may answer the two in either order, and either + # reply then says whether the session understood the attach. + if not self.__request_results( + query.execution_id, probe=True, attach=attach is not None + ): # A user thread is sending, possibly stalled, or the # socket can't take a write. Don't wait: that would stop # recv() for every query. Either holds for the rest too; @@ -399,6 +480,8 @@ def __probe_stale_queries(self) -> None: watch.probes_sent += 1 watch.probe_outstanding = True watch.results_probed = watch.results_requested + if attach is not None: + attach.probed = True # Space the next probe from this one, not from the last event. watch.last_activity = now logging.info( @@ -413,13 +496,15 @@ def __probe_stale_queries(self) -> None: except Exception: logging.exception("Stale-query probe failed") - def __connection_error(self, execution_id: str) -> OperationalError: + def __connection_error(self, execution_id: str) -> ConnectionLostError: message = ( f"SQL connection lost (session={self.__session_id or 'unknown'}, " f"execution={execution_id}). Commit outcome is unknown; " "verify the operation before retrying writes." ) - return OperationalError(message) + return ConnectionLostError( + message, session_id=self.__session_id, execution_id=execution_id + ) def __fail_pending(self, graceful: bool = False) -> None: # Stop admission first. Do not wait for __send_lock: its owner may be @@ -502,6 +587,9 @@ 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. + attach_reply = False + attach = query.attach if query is not None else None if query is not None: watch = query.watch # Any inbound event, progress included, shows the query is alive. @@ -512,6 +600,13 @@ def __listen(self) -> None: # silent query isn't probed every N seconds forever. probe_reply = watch.probe_outstanding watch.probe_outstanding = False + if attach is not None: + attach_reply = attach.reply_outstanding + attach.reply_outstanding = False + # Only a reply to an attach says so, and only from a + # session that supports it. Once said, it holds. + if message.get("attached") is True: + attach.acknowledged = True else: # A genuine event resets the backoff, but leaves an outstanding # probe outstanding: a non-terminal state_updated can arrive @@ -591,6 +686,13 @@ def complete_with_store_result(result_uri: str) -> None: query.state = ExecutionState.COMPLETED complete_query(ExecutionResult(store_result=store_result)) + def fail_attached(error: OperationalError) -> None: + # The session's answer to an attached query that leaves nothing + # to wait for on this connection. + query.state = ExecutionState.FAILED + logging.warning("%s", error) + complete_query(ExecutionResult(error=error)) + # Incoming state transitions are handled here. if kind == EventKind.STATE_UPDATED or kind == EventKind.EXECUTION_RESULT: try: @@ -600,6 +702,30 @@ def complete_with_store_result(result_uri: str) -> None: fail_query("interpret", error) return + if attach_reply and attach is not None: + if attach.acknowledged: + logging.info("Attached to execution %s.", execution_id) + elif query.state in (ExecutionState.PENDING, ExecutionState.RUNNING): + # An older session ignored the attach: the execution's + # later events go to the lost connection. Only probes can + # tell its outcome. + if self.__probe_after is None: + fail_attached( + OperationalError( + f"Execution {execution_id} is {query.state} on " + "a SQL session that doesn't support attaching " + f"(session={self.__session_id or 'unknown'}), and " + "stale-query probes are disabled, so this " + "connection can't learn its outcome." + ) + ) + return + logging.warning( + "SQL session doesn't support attaching; polling " + "execution %s for its outcome.", + execution_id, + ) + if query.state == ExecutionState.SUCCEEDED: # On a state_updated event telling us the query succeeded, # check if results are stored in cloud storage or need to be fetched. @@ -638,6 +764,39 @@ def complete_with_store_result(result_uri: str) -> None: return results = message.get("results") + if ( + query.store is not None + and (results is None or results == {}) + and attach is not None + ): + if attach.acknowledged: + # A session that supports attaching also echoes the + # result location, so this is an empty store result. + logging.info( + "Query %s completed with store configured but no " + "results to store.", + execution_id, + ) + query.state = ExecutionState.COMPLETED + complete_query(ExecutionResult()) + return + # An older session: either it predates echoing the + # location, which went to the lost connection with the + # terminal state_updated, or the stored result is empty. + # Nothing more will arrive here. Fail rather than wait + # forever, or report a real result as empty. + fail_attached( + OperationalError( + f"Execution {execution_id} succeeded, but the SQL " + f"session (session={self.__session_id or 'unknown'}) " + "doesn't support attaching and reported no stored " + "result location: the location went to the lost " + "connection, or the stored result is empty. It " + "can't be recovered." + ) + ) + return + if query.store is not None and (results is None or results == {}): # Succeeded without a result location: either an older # session that doesn't echo it, or an empty store result @@ -686,12 +845,47 @@ def complete_with_store_result(result_uri: str) -> None: pass elif kind == EventKind.ERROR: error = message.get("message") + not_found = error == _EXECUTION_NOT_FOUND + 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. + 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.", + session_id=self.__session_id, + execution_id=execution_id, + ) + ) + return + # For an attached query, any not-found while waiting is read as an + # eviction, not just a probe's: with the attach and the probe that + # repeats it both outstanding, the second reply has neither flag. if ( - probe_reply - and error == _EXECUTION_NOT_FOUND + (probe_reply or attach is not None) + and not_found and not query.watch.results_requested and query.state != ExecutionState.FAILED ): + if attach is not None and not attach.acknowledged: + # Evicted from an older session's cache, which sends the + # terminal event to the lost connection: waiting for it + # would wait forever. + fail_attached( + ExecutionNotFoundError( + f"Execution not found: SQL session " + f"{self.__session_id or '(unknown)'} no longer " + f"tracks execution {execution_id}, and doesn't " + "send its events to an attached connection (older " + "SQL session). Its outcome can't be recovered.", + session_id=self.__session_id, + execution_id=execution_id, + ) + ) + return # The session evicted the query from its cache. If the query is # still running, its terminal event will still arrive; if that # event was already lost, nothing more can be done (still @@ -761,6 +955,14 @@ 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() @@ -863,17 +1065,39 @@ def __execute_sql( ) return execution_id - def __request_results(self, execution_id: str, probe: bool = False) -> bool: - """Send retrieve_results; ``False`` if a probe found the sender busy. + 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. - A probe never waits for ``__send_lock`` or a full socket (see - ``__probe_stale_queries`` and ``__send``). + 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. """ - query = self.__queries.get(execution_id) - if not query: - return True + 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 - request = { + def __retrieve_request( + self, execution_id: str, attach: bool = False + ) -> Dict[str, Any]: + request: Dict[str, Any] = { "kind": RequestKind.RETRIEVE_RESULTS.value, "execution_id": execution_id, } @@ -883,6 +1107,24 @@ def __request_results(self, execution_id: str, probe: bool = False) -> bool: request["compression"] = self.__data_compression.value if self.__geometry_representation: request["geometry"] = self.__geometry_representation.value + if attach: + request["attach"] = True + return request + + def __request_results( + self, execution_id: str, probe: bool = False, attach: bool = False + ) -> bool: + """Send retrieve_results; ``False`` if a probe found the sender busy. + + A probe never waits for ``__send_lock`` or a full socket (see + ``__probe_stale_queries`` and ``__send``). ``attach`` repeats an + unanswered attach. + """ + query = self.__queries.get(execution_id) + if not query: + return True + + request = self.__retrieve_request(execution_id, attach=attach) if not probe: # A probe only asks; the reply carries the query's actual state. diff --git a/wherobots/db/cursor.py b/wherobots/db/cursor.py index 585c9d1..2258a19 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, List, Tuple, Dict +from typing import Any, Callable, List, Tuple, Dict -from .errors import ProgrammingError +from .errors import NotSupportedError, ProgrammingError from .models import ExecutionResult, Store, StoreResult # Matches pyformat parameter markers: %(name)s @@ -67,9 +67,10 @@ def replacer(match: re.Match) -> str: class Cursor: - def __init__(self, exec_fn, cancel_fn) -> None: + def __init__(self, exec_fn, cancel_fn, attach_fn=None) -> 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 @@ -95,6 +96,17 @@ def description(self) -> List[Tuple] | None: def rowcount(self) -> int: return self.__rowcount + @property + 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. + """ + return self.__current_execution_id + def __in_flight_execution_id(self) -> str | None: """The current execution's id if its result has not yet arrived. @@ -170,6 +182,51 @@ 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. 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) @@ -189,11 +246,7 @@ def execute( # 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 = self.__exec_fn( - _substitute_parameters(operation, parameters), - self.__queue.put, - store, - ) + self.__current_execution_id = submit(self.__queue.put) 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 399bced..26b97cd 100644 --- a/wherobots/db/driver.py +++ b/wherobots/db/driver.py @@ -312,6 +312,9 @@ def connect_direct( ) -> 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). + :param stale_query_probe_seconds: See :func:`connect`. """ uri_with_protocol = f"{uri}/{protocol}" @@ -351,6 +354,27 @@ def ws_connect() -> websockets.sync.client.ClientConnection: except Exception as e: raise InterfaceError("Failed to connect to SQL session!") from e + # Snapshot a dict of headers, so a reconnect sends what this connection + # sent even if the caller later mutates it. + reconnect_headers = dict(headers) if isinstance(headers, dict) else headers + + 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, + ) + return Connection( ws, read_timeout=read_timeout, @@ -359,4 +383,5 @@ def ws_connect() -> websockets.sync.client.ClientConnection: geometry_representation=geometry_representation, session_id=session_id, stale_query_probe_seconds=stale_query_probe_seconds, + reconnect_factory=reconnect, ) diff --git a/wherobots/db/errors.py b/wherobots/db/errors.py index 4e18110..e1f7c18 100644 --- a/wherobots/db/errors.py +++ b/wherobots/db/errors.py @@ -24,3 +24,52 @@ class ProgrammingError(DatabaseError): class NotSupportedError(DatabaseError): pass + + +class ConnectionLostError(OperationalError): + """The connection to the SQL session was lost while a query was pending. + + 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. + + Attributes: + session_id: The SQL session's ID, or None if the connection didn't + know it. + execution_id: The ID of the pending execution. + """ + + def __init__( + self, + message: str, + session_id: str | None = None, + execution_id: str | None = None, + ) -> None: + super().__init__(message) + self.session_id = session_id + self.execution_id = execution_id + + +class ExecutionNotFoundError(OperationalError): + """The SQL session has no record of the execution to attach to. + + The execution ran on a different SQL session, the session restarted, or + it no longer tracks the execution. Its outcome can't be recovered. + + Attributes: + session_id: The SQL session's ID, or None if the connection didn't + know it. + execution_id: The ID of the execution. + """ + + def __init__( + self, + message: str, + session_id: str | None = None, + execution_id: str | None = None, + ) -> None: + super().__init__(message) + self.session_id = session_id + self.execution_id = execution_id From 059a67c7199d3194556e84db6ae8aa86cde4644f Mon Sep 17 00:00:00 2001 From: Kevin Date: Fri, 16 Oct 2026 11:07:29 -0700 Subject: [PATCH 2/2] fix: report an attached execution cancelled elsewhere as an error (DB-760) An attached query whose outcome was cancelled completed like an executed one: an empty result, no error. With sql-session's orphan lease, a re-attach that arrives after the grace period has run out (reconnecting can take longer than 120 s, and a grace of 0 is a supported setting) then reads as a successful 0-row result. For a write cancelled midway, that reports it as succeeded. Add ExecutionCancelledError(OperationalError), with session_id and execution_id, exported like the other two. An attached query that reaches cancelled (from the attach reply, a probe reply or a pushed state_updated) without this connection having cancelled it now fails with it, saying the execution was cancelled before or while re-attaching and that its outcome, for a write whether it committed, is unknown. The connection records that it sent a cancel on the Query (cancel_requested), under __lock and before sending, so the flag is set by the time the session can report the cancellation. An attached query this cursor cancelled itself (cursor.close(), or executing on it) still completes empty, and executed queries are unchanged. Co-Authored-By: Claude Opus 5.5 --- README.md | 7 +++ tests/test_attach.py | 90 +++++++++++++++++++++++++++++++++++++- wherobots/db/__init__.py | 2 + wherobots/db/connection.py | 38 ++++++++++++++-- wherobots/db/cursor.py | 6 ++- wherobots/db/errors.py | 25 +++++++++++ 6 files changed, 162 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index d22e294..a5c1945 100644 --- a/README.md +++ b/README.md @@ -244,6 +244,13 @@ with connect(api_key='...') as conn: * 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 diff --git a/tests/test_attach.py b/tests/test_attach.py index 6e8a6cf..a940a00 100644 --- a/tests/test_attach.py +++ b/tests/test_attach.py @@ -14,6 +14,7 @@ from wherobots.db import ( ConnectionLostError, + ExecutionCancelledError, ExecutionNotFoundError, NotSupportedError, OperationalError, @@ -475,12 +476,81 @@ def test_acknowledged_attach_waits_with_probes_disabled(): conn.close() -def test_attach_to_a_cancelled_execution_completes_empty(): - session = new_session("cancelled") +def assert_cancelled_elsewhere(error): + assert isinstance(error, ExecutionCancelledError) + assert isinstance(error, OperationalError) + assert error.session_id == "session-1" + assert error.execution_id == EID + assert str(error) == ( + f"Execution {EID} was cancelled on SQL session session-1 before or " + "while this connection re-attached to it (for example, its orphan " + "grace period ran out, or another client cancelled it). Its outcome " + "is unknown: for a write, whether it committed." + ) + + +@pytest.mark.parametrize("make_session", [new_session, old_session]) +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) + try: + cursor = conn.cursor() + cursor.attach(EID, store=STORE) + assert_cancelled_elsewhere(outcome(cursor).error) + with pytest.raises(ExecutionCancelledError): + cursor.get_store_result() + assert tracked(conn, EID) is None + finally: + conn.close() + + +def test_cancelled_pushed_after_attach_raises(): + session = new_session("running") + conn = open_connection(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) + assert session.requests("cancel") == [] + finally: + conn.close() + + +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) + try: + cursor = conn.cursor() + cursor.attach(EID, store=STORE) + assert_cancelled_elsewhere(outcome(cursor).error) + assert len(session.retrieves()) == 2 + finally: + conn.close() + + +@pytest.mark.parametrize("acknowledged", [True, False]) +def test_own_cancel_of_an_attached_execution_completes_empty(acknowledged): + def on_retrieve(s, m): + state = "cancelled" if s.requests("cancel") else "running" + fields = {"attached": True} if acknowledged and m.get("attach") else {} + answer(s, m, state, **fields) + + session = Session(on_retrieve=on_retrieve) conn = open_connection(session) 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] + if acknowledged: + # Pushed to the new owner. An older session's next probe reply + # says cancelled instead. + session.put(kind="state_updated", execution_id=EID, state="cancelled") result = outcome(cursor) assert result.error is None assert result.results.empty @@ -489,6 +559,22 @@ def test_attach_to_a_cancelled_execution_completes_empty(): conn.close() +def test_executed_query_cancelled_elsewhere_still_completes_empty(): + session = Session() + conn = open_connection(session) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") + session.put( + kind="state_updated", execution_id=cursor.execution_id, state="cancelled" + ) + result = outcome(cursor) + assert result.error is None + assert result.results.empty + finally: + conn.close() + + @pytest.mark.parametrize("failed_reply", [True, False]) def test_attach_to_a_failed_execution_reports_its_error(failed_reply): # sql-session answers a retrieve for a failed execution with a failed diff --git a/wherobots/db/__init__.py b/wherobots/db/__init__.py index f8831d6..7324d31 100644 --- a/wherobots/db/__init__.py +++ b/wherobots/db/__init__.py @@ -5,6 +5,7 @@ Error, ConnectionLostError, DatabaseError, + ExecutionCancelledError, ExecutionNotFoundError, InternalError, InterfaceError, @@ -29,6 +30,7 @@ "InterfaceError", "OperationalError", "ConnectionLostError", + "ExecutionCancelledError", "ExecutionNotFoundError", "ProgrammingError", "NotSupportedError", diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index e15d7b9..1acc6e6 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -21,6 +21,7 @@ from .cursor import Cursor from .errors import ( ConnectionLostError, + ExecutionCancelledError, ExecutionNotFoundError, NotSupportedError, OperationalError, @@ -221,6 +222,12 @@ class Query: watch: _Watch = field(default_factory=_Watch) attach: _Attach | None = None """Set for a query attached here rather than executed here.""" + 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.""" class Connection: @@ -834,6 +841,27 @@ def fail_attached(error: OperationalError) -> None: query.state = ExecutionState.COMPLETED complete_query(ExecutionResult(results=decoded)) elif query.state == ExecutionState.CANCELLED: + if attach is not None: + with self.__lock: + requested = query.cancel_requested + if not requested: + # Not this connection's cancel: the session cancelled + # it before or while this connection re-attached. + # Completing empty would report, say, a write + # cancelled midway as a success. + cancelled = ExecutionCancelledError( + f"Execution {execution_id} was cancelled on SQL " + f"session {self.__session_id or '(unknown)'} before " + "or while this connection re-attached to it (for " + "example, its orphan grace period ran out, or " + "another client cancelled it). Its outcome is " + "unknown: for a write, whether it committed.", + session_id=self.__session_id, + execution_id=execution_id, + ) + logging.warning("%s", cancelled) + complete_query(ExecutionResult(error=cancelled)) + return logging.info( "Query %s has been cancelled; returning empty results.", execution_id, @@ -1135,9 +1163,13 @@ def __request_results( def __cancel_query(self, execution_id: str) -> None: """Cancels the query with the given execution ID.""" - query = self.__queries.get(execution_id) - if not query: - return + with self.__lock: + query = self.__queries.get(execution_id) + if query is None: + 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/cursor.py b/wherobots/db/cursor.py index 2258a19..c26d3a9 100644 --- a/wherobots/db/cursor.py +++ b/wherobots/db/cursor.py @@ -204,7 +204,11 @@ def attach(self, execution_id: str, store: Store | None = None) -> None: 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. An older SQL session + 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 diff --git a/wherobots/db/errors.py b/wherobots/db/errors.py index e1f7c18..0926bb7 100644 --- a/wherobots/db/errors.py +++ b/wherobots/db/errors.py @@ -52,6 +52,31 @@ def __init__( self.execution_id = execution_id +class ExecutionCancelledError(OperationalError): + """An 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 + committed. + + Attributes: + session_id: The SQL session's ID, or None if the connection didn't + know it. + execution_id: The ID of the execution. + """ + + def __init__( + self, + message: str, + session_id: str | None = None, + execution_id: str | None = None, + ) -> None: + super().__init__(message) + self.session_id = session_id + self.execution_id = execution_id + + class ExecutionNotFoundError(OperationalError): """The SQL session has no record of the execution to attach to.