diff --git a/README.md b/README.md index d99cfb0..a5c1945 100644 --- a/README.md +++ b/README.md @@ -200,6 +200,65 @@ 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. +* 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. + ### 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..a940a00 --- /dev/null +++ b/tests/test_attach.py @@ -0,0 +1,1096 @@ +"""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, + 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, + 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 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 + assert tracked(conn, EID) is None + finally: + 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 + # 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..7324d31 100644 --- a/wherobots/db/__init__.py +++ b/wherobots/db/__init__.py @@ -3,7 +3,10 @@ from .driver import connect, connect_direct from .errors import ( Error, + ConnectionLostError, DatabaseError, + ExecutionCancelledError, + ExecutionNotFoundError, InternalError, InterfaceError, OperationalError, @@ -26,6 +29,9 @@ "InternalError", "InterfaceError", "OperationalError", + "ConnectionLostError", + "ExecutionCancelledError", + "ExecutionNotFoundError", "ProgrammingError", "NotSupportedError", "Region", diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index 5d53298..1acc6e6 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -19,7 +19,14 @@ 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, + ExecutionCancelledError, + ExecutionNotFoundError, + NotSupportedError, + OperationalError, + ProgrammingError, +) from .models import ExecutionResult, ProgressInfo, Store, StoreResult from .types import ( RequestKind, @@ -92,7 +99,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 +187,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 +220,14 @@ 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.""" + 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: @@ -214,8 +254,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 +341,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 +423,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 +448,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 +487,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 +503,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 +594,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 +607,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 +693,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 +709,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 +771,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 @@ -675,6 +841,27 @@ def complete_with_store_result(result_uri: str) -> 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, @@ -686,12 +873,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 +983,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 +1093,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 +1135,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. @@ -893,9 +1163,13 @@ def __request_results(self, execution_id: str, probe: bool = False) -> bool: 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 585c9d1..c26d3a9 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,55 @@ 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) @@ -189,11 +250,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..0926bb7 100644 --- a/wherobots/db/errors.py +++ b/wherobots/db/errors.py @@ -24,3 +24,77 @@ 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 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. + + 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