diff --git a/README.md b/README.md index ece9305..8f49122 100644 --- a/README.md +++ b/README.md @@ -200,6 +200,49 @@ 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 WebSocket to the SQL session drops while queries are running (a +network path failure, a keepalive timeout), the connection reconnects to +the same SQL session by itself and re-attaches every pending query, so +their results still arrive: there's nothing for your code to do. It keeps +trying, with backoff, for up to `reconnect_timeout_seconds` (60 seconds by +default; see below). It never starts a new session, and gives up at once +if the session refuses the connection, for example because it's gone. +While it reconnects, `execute()` and cancelling (closing a cursor, or +executing on it) wait for it to finish. + +A re-attached query can still fail: + +* If the SQL session doesn't know the execution (the statement didn't + reach it before the connection dropped, or the session restarted), + fetching raises `ExecutionNotFoundError`, a subclass of + `OperationalError`. +* If the SQL session cancelled the execution while the connection was + down (for example, its grace period for a lost connection ran out, or + another client cancelled it), fetching raises + `ExecutionCancelledError`, also an `OperationalError` subclass: its + outcome, including whether a write committed, is unknown. An execution + that you cancel yourself, by closing the cursor or executing on it, + completes with an empty result, as usual. +* An older SQL session that doesn't support attaching still answers with + the execution's state, and the driver polls it for the outcome (see + `stale_query_probe_seconds` below). If such a session reports no + location for a stored result, because it predates reporting it or + because the stored result is empty, the driver can't tell which, and + fails the query with an `OperationalError`. With probes disabled, a + running execution on such a session fails right away. + +If the connection can't reconnect in time, or reconnecting is turned off +(`reconnect_timeout_seconds=None`), fetching a pending query's outcome +raises `ConnectionLostError`, another `OperationalError` subclass. The +query may still have run, or be running, on the SQL session: its outcome, +including whether a write committed, is unknown. The error's +`execution_id` and `session_id` attributes identify it, and a cursor's +`execution_id` property holds the ID of its current or last execution. +(When `execute()` itself raises `ConnectionLostError`, the connection was +already lost and the statement was never sent.) + ### Runtime and region selection You can chose the Wherobots runtime you want to use using the `runtime` @@ -290,3 +333,8 @@ users may find useful: for its state again, in case the query's completion event was lost. Repeat probes back off to at most 8 times this interval. Defaults to 30 seconds; pass `None` to disable probing. +* `reconnect_timeout_seconds`: how long the connection keeps trying to + reconnect to the same SQL session after losing its WebSocket (see + "Recovering from a lost connection" above). Defaults to 60 seconds, + which fits in the SQL session's 120-second grace period for the + queries of a lost connection; pass `None` to turn reconnecting off. diff --git a/tests/test_attach.py b/tests/test_attach.py new file mode 100644 index 0000000..10df3fa --- /dev/null +++ b/tests/test_attach.py @@ -0,0 +1,830 @@ +"""Re-attaching pending queries after a reconnect: handling the replies.""" + +import json +import logging +import queue +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, + OperationalError, +) +from wherobots.db.connection import Connection +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-pending-when-lost" +"""The execution pending when the connection lost its transport.""" + + +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 reattach(session, store=None, **kwargs): + """A connection whose query EID was pending when its transport was lost. + + The connection reconnects to ``session``, its only transport to come, and + re-attaches the query there. Returns the connection and the query's + cursor. + """ + transports = iter([session]) + + def transport_factory(timeout): + transport = next(transports, None) + if transport is None: + raise OSError("connection refused") + return transport + + kwargs.setdefault("reconnect_timeout_seconds", 0.5) + lost = Session() + conn = open_connection(lost, transport_factory=transport_factory, **kwargs) + cursor = conn.cursor() + with patch("wherobots.db.connection.uuid.uuid4", return_value=EID): + cursor.execute("SELECT * FROM t", store=store) + lost.incoming.put(ConnectionClosedError(None, None)) + return conn, cursor + + +def tracked(conn, execution_id): + return conn._Connection__queries.get(execution_id) + + +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 + finally: + conn.close() + + +# Re-attaching on a session that supports it + + +def test_attach_request_carries_the_connection_settings(): + session = new_session() + conn, _ = reattach( + session, + store=STORE, + results_format=ResultsFormat.JSON, + data_compression=DataCompression.BROTLI, + geometry_representation=GeometryRepresentation.WKB, + ) + try: + assert wait_until(session.retrieves) + 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, cursor = reattach(session, store=STORE) + try: + 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, cursor = reattach(session, store=STORE) + try: + 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, cursor = reattach(session) + try: + 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, cursor = reattach(session, store=STORE) + try: + 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, cursor = reattach(session, store=STORE) + try: + 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, cursor = reattach(session, results_format=ResultsFormat.JSON) + try: + 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, cursor = reattach(session, store=STORE) + try: + 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, cursor = reattach(session, store=STORE, stale_query_probe_seconds=None) + try: + 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, cursor = reattach(session, store=STORE) + try: + 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, cursor = reattach(session) + try: + 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, cursor = reattach(session, store=STORE) + try: + 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, cursor = reattach(session, store=STORE) + try: + 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, cursor = reattach(session) + try: + error = outcome(cursor).error + assert type(error) is OperationalError + assert str(error) == "Table not found: t" + finally: + conn.close() + + +# Re-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, cursor = reattach(session, store=store) + try: + 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, cursor = reattach(session, store=STORE) + try: + 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, cursor = reattach(session, store=STORE) + try: + 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, cursor = reattach(session, store=STORE, stale_query_probe_seconds=None) + try: + 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() + + +# Re-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, cursor = reattach(session, store=STORE) + try: + 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} after reconnecting" 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, cursor = reattach(session) + try: + error = outcome(cursor).error + assert isinstance(error, ExecutionNotFoundError) + assert "after reconnecting" 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, cursor = reattach(session, store=STORE) + try: + 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, cursor = reattach(session, store=STORE) + try: + 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() + + +# The cursor's part + + +@pytest.mark.parametrize("then", ["close", "execute"]) +def test_the_cursor_still_owns_the_reattached_execution(then): + session = new_session("running") + conn, cursor = reattach(session, store=STORE) + try: + 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_transport_lost_again_without_reconnecting_raises_connection_lost(): + session = new_session("running") + conn, cursor = reattach(session, store=STORE, reconnect_timeout_seconds=0.2) + try: + 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() diff --git a/tests/test_driver.py b/tests/test_driver.py index f2909e7..5bf1907 100644 --- a/tests/test_driver.py +++ b/tests/test_driver.py @@ -192,6 +192,23 @@ def test_stale_query_probes_default_to_30_seconds(self, mock_post, mock_get): _, kwargs = _run_connect_full(mock_post, mock_get) assert kwargs["stale_query_probe_seconds"] == 30.0 + @pytest.mark.parametrize("reconnect_seconds", [None, 5.0]) + @patch("wherobots.db.driver.requests.get") + @patch("wherobots.db.driver.requests.post") + def test_reconnect_timeout_seconds_is_forwarded( + self, mock_post, mock_get, reconnect_seconds + ): + _, kwargs = _run_connect_full( + mock_post, mock_get, reconnect_timeout_seconds=reconnect_seconds + ) + assert kwargs["reconnect_timeout_seconds"] == reconnect_seconds + + @patch("wherobots.db.driver.requests.get") + @patch("wherobots.db.driver.requests.post") + def test_reconnecting_defaults_to_60_seconds(self, mock_post, mock_get): + _, kwargs = _run_connect_full(mock_post, mock_get) + assert kwargs["reconnect_timeout_seconds"] == 60.0 + @patch("wherobots.db.driver.requests.get") @patch("wherobots.db.driver.requests.post") def test_omitted_region_runtime_not_sent(self, mock_post, mock_get): diff --git a/tests/test_reconnect.py b/tests/test_reconnect.py new file mode 100644 index 0000000..daceadb --- /dev/null +++ b/tests/test_reconnect.py @@ -0,0 +1,744 @@ +"""Reconnecting by itself after a lost transport, and re-attaching queries.""" + +import itertools +import json +import logging +import queue +import threading +import time +from unittest.mock import MagicMock, patch + +import cbor2 +import pandas +import pyarrow +import pytest +from websockets.datastructures import Headers +from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK +from websockets.exceptions import InvalidStatus +from websockets.http11 import Response + +from wherobots.db import ( + ConnectionLostError, + ExecutionCancelledError, + ExecutionNotFoundError, +) +from wherobots.db.connection import Connection +from wherobots.db.driver import connect_direct +from wherobots.db.errors import InterfaceError +from wherobots.db.models import Store, StoreResult +from wherobots.db.types import StorageFormat + +PROBE = 0.05 +"""Probe interval used throughout: small, so the tests stay fast.""" + +READ_TIMEOUT = 0.01 + +NOT_FOUND = "Execution not found" + +STORE = Store.for_download(format=StorageFormat.PARQUET) + +RESULT_URI = "https://presigned.example.com/results.parquet" + +WS_CONNECT = "wherobots.db.driver.websockets.sync.client.connect" + + +@pytest.fixture(autouse=True) +def fast_backoff(monkeypatch): + monkeypatch.setattr("wherobots.db.connection._RECONNECT_FIRST_DELAY_SECONDS", 0.01) + monkeypatch.setattr("wherobots.db.connection._RECONNECT_MAX_DELAY_SECONDS", 0.05) + + +class Session: + """A fake SQL session transport. + + Records every request. ``on_retrieve`` answers ``retrieve_results`` + requests (attaches and probes). With ``fail_sends``, a send raises as on a + dead socket. + """ + + def __init__(self, on_retrieve=None): + self.incoming = queue.Queue() + self.sent = [] + self.on_retrieve = on_retrieve + self.fail_sends = False + self.close_timeout = 1.0 + self.socket = MagicMock() + self.socket.shutdown.side_effect = lambda how: self.close() + + def recv(self, timeout): + try: + value = self.incoming.get(timeout=timeout) + except queue.Empty: + raise TimeoutError from None + if isinstance(value, Exception): + raise value + # CBOR, as sql-session sends results: it carries their bytes. + return cbor2.dumps(value) + + def send(self, value): + if self.fail_sends: + raise OSError("broken pipe") + message = json.loads(value) + self.sent.append(message) + if message["kind"] == "retrieve_results" and self.on_retrieve is not None: + self.on_retrieve(self, message) + + def close(self): + self.incoming.put(ConnectionClosedOK(None, None)) + + def drop(self): + """The network path fails: the keepalive gives up on the connection.""" + self.incoming.put(ConnectionClosedError(None, None)) + + def put(self, **message): + self.incoming.put(message) + + def kinds(self): + return [message["kind"] for message in list(self.sent)] + + def requests(self, kind, execution_id=None): + return [ + message + for message in list(self.sent) + if message["kind"] == kind + and (execution_id is None or message["execution_id"] == execution_id) + ] + + +def answer(session, request, state="running", **fields): + """Answer a retrieve_results as sql-session does: execution_result.""" + fields.setdefault("results", None) + session.put( + kind="execution_result", + execution_id=request["execution_id"], + state=state, + **fields, + ) + + +def new_session(state="running", **fields): + """A session that supports attaching: it acknowledges each attach.""" + return Session( + on_retrieve=lambda s, m: answer( + s, m, state, attached=bool(m.get("attach")), **fields + ) + ) + + +class Factory: + """A transport factory handing out ``transports`` in turn. + + Each is a transport to return, an exception to raise, or a callable + producing either. Records the time limit of every call. + """ + + def __init__(self, *transports): + self.transports = iter(transports) + self.calls = [] + + def __call__(self, timeout): + self.calls.append(timeout) + transport = next(self.transports, OSError("connection refused")) + if callable(transport) and not isinstance(transport, Session): + transport = transport() + if isinstance(transport, Exception): + raise transport + return transport + + +def open_connection(session, factory, **kwargs): + kwargs.setdefault("stale_query_probe_seconds", PROBE) + kwargs.setdefault("session_id", "session-1") + kwargs.setdefault("reconnect_timeout_seconds", 3.0) + return Connection( + session, read_timeout=READ_TIMEOUT, transport_factory=factory, **kwargs + ) + + +def arrow_results(**columns): + """A ``results`` payload carrying ``columns`` as an Arrow IPC stream.""" + table = pyarrow.table(columns) + sink = pyarrow.BufferOutputStream() + with pyarrow.ipc.new_stream(sink, table.schema) as writer: + writer.write_table(table) + return {"result_bytes": sink.getvalue().to_pybytes(), "format": "arrow"} + + +def wait_until(predicate, timeout=3.0): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.005) + return predicate() + + +def tracked(conn, execution_id): + return conn._Connection__queries.get(execution_id) + + +def outcome(cursor, timeout=3.0): + """Wait for the cursor's terminal result, then hand it back to the cursor.""" + result = cursor._Cursor__queue.get(timeout=timeout) + cursor._Cursor__queue.put(result) + return result + + +def reader(conn): + return conn._Connection__thread + + +def in_thread(fn): + thread = threading.Thread(target=fn, daemon=True) + thread.start() + return thread + + +# Recovering pending queries + + +def test_lost_query_completes_over_the_new_transport(caplog): + caplog.set_level(logging.INFO) + lost, again = Session(), new_session("running") + factory = Factory(again) + conn = open_connection(lost, factory) + try: + cursor = conn.cursor() + cursor.execute("SELECT * FROM t", store=STORE) + execution_id = cursor.execution_id + lost.drop() + assert wait_until(lambda: again.requests("retrieve_results")) + assert again.sent[0] == { + "kind": "retrieve_results", + "execution_id": execution_id, + "attach": True, + } + again.put( + kind="state_updated", + execution_id=execution_id, + state="succeeded", + result_uri=RESULT_URI, + size=7, + ) + assert outcome(cursor).error is None + assert cursor.get_store_result() == StoreResult(RESULT_URI, 7) + assert len(factory.calls) == 1 + assert lost.socket.shutdown.called + finally: + conn.close() + messages = [r.getMessage() for r in caplog.records] + assert ( + "Lost the connection to SQL session session-1 with 1 pending queries; " + "reconnecting for up to 3s ..." + ) in messages + assert ( + "Reconnected to SQL session session-1; re-attaching 1 pending queries." + ) in messages + + +def test_result_lost_with_the_connection_arrives_in_the_attach_reply(): + # The query succeeded, but the terminal event went to the lost connection. + again = Session( + on_retrieve=lambda s, m: answer( + s, m, "succeeded", attached=True, results=arrow_results(x=[1, 2]) + ) + ) + lost = Session() + conn = open_connection(lost, Factory(again)) + try: + cursor = conn.cursor() + cursor.execute("SELECT x FROM t") + lost.drop() + assert outcome(cursor).error is None + assert list(cursor.fetchall()["x"]) == [1, 2] + finally: + conn.close() + + +def test_every_pending_query_is_reattached(): + lost, again = Session(), new_session("running") + conn = open_connection(lost, Factory(again)) + try: + cursors = [conn.cursor() for _ in range(3)] + for cursor in cursors: + cursor.execute("SELECT 1", store=STORE) + lost.drop() + ids = [cursor.execution_id for cursor in cursors] + assert wait_until(lambda: len(again.requests("retrieve_results")) == 3) + assert {m["execution_id"] for m in again.requests("retrieve_results")} == ( + set(ids) + ) + for execution_id in ids: + again.put( + kind="state_updated", + execution_id=execution_id, + state="succeeded", + result_uri=f"{RESULT_URI}?{execution_id}", + ) + for cursor in cursors: + assert outcome(cursor).error is None + assert cursor.get_store_result().result_uri.endswith(cursor.execution_id) + finally: + conn.close() + + +def test_reattached_query_is_polled_on_an_older_session(): + # An older session ignores the attach: its events go to the lost + # connection, so only polling tells the outcome. + replies = itertools.chain(["running"], itertools.repeat("succeeded")) + + def on_retrieve(s, m): + state = next(replies) + results = arrow_results(x=[3]) if state == "succeeded" else None + s.put( + kind="execution_result", + execution_id=m["execution_id"], + state=state, + results=results, + ) + + lost, again = Session(), Session(on_retrieve=on_retrieve) + conn = open_connection(lost, Factory(again)) + try: + cursor = conn.cursor() + cursor.execute("SELECT x FROM t") + lost.drop() + assert outcome(cursor).error is None + assert list(cursor.fetchall()["x"]) == [3] + assert len(again.requests("retrieve_results")) >= 2 + finally: + conn.close() + + +def test_reattached_query_unknown_to_the_session_raises_not_found(): + again = Session( + on_retrieve=lambda s, m: s.put( + kind="error", execution_id=m["execution_id"], message=NOT_FOUND + ) + ) + lost = Session() + conn = open_connection(lost, Factory(again)) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") + lost.drop() + error = outcome(cursor).error + assert isinstance(error, ExecutionNotFoundError) + assert error.execution_id == cursor.execution_id + assert str(error) == ( + "Execution not found: SQL session session-1 has no execution " + f"{cursor.execution_id} after reconnecting. The statement didn't " + "reach the session before the connection was lost, the session " + "restarted, or the session no longer tracks it." + ) + finally: + conn.close() + + +def test_query_cancelled_by_the_session_while_disconnected_raises(): + # E.g. the session's grace period for orphaned executions ran out. + lost, again = Session(), new_session("cancelled") + conn = open_connection(lost, Factory(again)) + try: + cursor = conn.cursor() + cursor.execute("UPDATE t SET x = 1") + lost.drop() + error = outcome(cursor).error + assert isinstance(error, ExecutionCancelledError) + assert error.execution_id == cursor.execution_id + finally: + conn.close() + + +def test_a_failed_send_hands_its_query_to_the_reconnect(): + again = Session( + on_retrieve=lambda s, m: s.put( + kind="error", execution_id=m["execution_id"], message=NOT_FOUND + ) + ) + lost = Session() + lost.fail_sends = True + conn = open_connection(lost, Factory(again)) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") # Doesn't raise: the query is re-attached. + assert wait_until(lambda: again.requests("retrieve_results")) + assert again.sent[0]["execution_id"] == cursor.execution_id + assert isinstance(outcome(cursor).error, ExecutionNotFoundError) + assert reader(conn).is_alive() + finally: + conn.close() + + +# Eagerness, and sends while reconnecting + + +def test_reconnects_eagerly_with_nothing_pending(): + lost, again = Session(), Session() + factory = Factory(again) + conn = open_connection(lost, factory) + try: + lost.drop() + assert wait_until(lambda: conn._Connection__ws is again) + assert factory.calls + cursor = conn.cursor() + cursor.execute("SELECT 1") + assert again.requests("execute_sql", cursor.execution_id) + assert not lost.requests("execute_sql") + finally: + conn.close() + + +def test_execute_waits_for_the_reconnect_and_follows_the_reattaches(): + release = threading.Event() + entered = threading.Event() + again = new_session("running") + + def slow_connect(): + entered.set() + release.wait(3) + return again + + lost = Session() + conn = open_connection(lost, Factory(slow_connect)) + try: + first = conn.cursor() + first.execute("SELECT 1", store=STORE) + lost.drop() + assert entered.wait(3) + second = conn.cursor() + executing = in_thread(lambda: second.execute("SELECT 2")) + executing.join(0.1) + assert executing.is_alive() # Waiting for the reconnect. + release.set() + executing.join(3) + assert not executing.is_alive() + assert again.kinds()[:2] == ["retrieve_results", "execute_sql"] + assert again.sent[0]["execution_id"] == first.execution_id + assert again.sent[1]["execution_id"] == second.execution_id + assert not lost.requests("execute_sql", second.execution_id) + finally: + release.set() + conn.close() + + +def test_cancel_waits_for_the_reconnect_and_follows_the_attach(): + release = threading.Event() + entered = threading.Event() + again = new_session("running") + + def slow_connect(): + entered.set() + release.wait(3) + return again + + lost = Session() + conn = open_connection(lost, Factory(slow_connect)) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") + execution_id = cursor.execution_id + lost.drop() + assert entered.wait(3) + closing = in_thread(cursor.close) + closing.join(0.1) + assert closing.is_alive() + # Not before the cancel goes out: a cancellation the session reports + # meanwhile isn't this connection's. + assert tracked(conn, execution_id).cancel_requested is False + release.set() + closing.join(3) + assert not closing.is_alive() + assert again.kinds()[:2] == ["retrieve_results", "cancel"] + assert tracked(conn, execution_id).cancel_requested is True + again.put(kind="state_updated", execution_id=execution_id, state="cancelled") + result = outcome(cursor) + assert result.error is None + assert isinstance(result.results, pandas.DataFrame) + assert result.results.empty + finally: + release.set() + conn.close() + + +def test_a_cancel_sent_before_the_loss_is_sent_again_after_the_attach(): + lost, again = Session(), new_session("running") + conn = open_connection(lost, Factory(again)) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") + execution_id = cursor.execution_id + cursor.close() + assert lost.requests("cancel", execution_id) + lost.drop() + assert wait_until(lambda: len(again.sent) >= 2) + assert again.kinds()[:2] == ["retrieve_results", "cancel"] + assert again.sent[1]["execution_id"] == execution_id + finally: + conn.close() + + +# Giving up + + +def test_pending_queries_fail_once_reconnecting_runs_out_of_time(): + lost = Session() + factory = Factory() # Every attempt fails. + conn = open_connection(lost, factory, reconnect_timeout_seconds=0.3) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") + started = time.monotonic() + lost.drop() + error = outcome(cursor).error + elapsed = time.monotonic() - started + assert isinstance(error, ConnectionLostError) + assert error.execution_id == cursor.execution_id + assert 0.25 <= elapsed < 2.0 + assert len(factory.calls) >= 2 + # Each attempt gets only the time left. + assert all(0 < limit <= 0.3 for limit in factory.calls) + assert wait_until(lambda: not reader(conn).is_alive()) + with pytest.raises(ConnectionLostError): + conn.cursor().execute("SELECT 2") + finally: + conn.close() + + +def test_a_refused_reconnect_fails_pending_queries_at_once(): + lost = Session() + factory = Factory(InterfaceError("SQL session refused the connection (HTTP 404)")) + conn = open_connection(lost, factory, reconnect_timeout_seconds=30) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") + lost.drop() + assert isinstance(outcome(cursor).error, ConnectionLostError) + assert len(factory.calls) == 1 + finally: + conn.close() + + +def test_a_transport_lost_again_at_once_continues_the_same_recovery(): + def dropping(): + session = Session() + session.drop() + return session + + lost = Session() + factory = Factory(*[dropping] * 1000) + conn = open_connection(lost, factory, reconnect_timeout_seconds=0.3) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") + lost.drop() + # Without a shared deadline, this would reconnect forever. + assert isinstance(outcome(cursor).error, ConnectionLostError) + assert 2 <= len(factory.calls) < 100 + finally: + conn.close() + + +def test_without_reconnecting_pending_queries_fail_at_once(): + lost = Session() + factory = Factory(Session()) + conn = open_connection(lost, factory, reconnect_timeout_seconds=None) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") + lost.drop() + assert isinstance(outcome(cursor).error, ConnectionLostError) + assert factory.calls == [] + finally: + conn.close() + + +@pytest.mark.parametrize("value", [0, -1, float("nan"), "soon", True]) +def test_an_invalid_reconnect_timeout_disables_reconnecting(value, caplog): + lost = Session() + factory = Factory(Session()) + conn = open_connection(lost, factory, reconnect_timeout_seconds=value) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") + lost.drop() + assert isinstance(outcome(cursor).error, ConnectionLostError) + assert factory.calls == [] + finally: + conn.close() + assert any( + "Invalid reconnect_timeout_seconds" in r.getMessage() for r in caplog.records + ) + + +# close() while reconnecting + + +def test_close_while_backing_off_fails_pending_queries_and_stops(): + lost = Session() + factory = Factory() # Every attempt fails. + conn = open_connection(lost, factory, reconnect_timeout_seconds=30) + cursor = conn.cursor() + cursor.execute("SELECT 1") + lost.drop() + assert wait_until(lambda: len(factory.calls) >= 2) + started = time.monotonic() + conn.close() + assert time.monotonic() - started < 2.0 + assert isinstance(outcome(cursor).error, ConnectionLostError) + assert wait_until(lambda: not reader(conn).is_alive()) + + +def test_close_while_connecting_discards_the_new_transport(): + again = new_session("running") + entered = threading.Event() + + def connect_during_close(): + entered.set() + assert wait_until(lambda: conn._Connection__closed) + return again + + lost = Session() + conn = open_connection(lost, Factory(connect_during_close)) + cursor = conn.cursor() + cursor.execute("SELECT 1") + lost.drop() + assert entered.wait(3) + conn.close() + assert isinstance(outcome(cursor).error, ConnectionLostError) + assert wait_until(lambda: not reader(conn).is_alive()) + assert again.socket.shutdown.called + assert again.sent == [] + + +def test_an_execute_waiting_for_the_reconnect_fails_when_closed(): + release = threading.Event() + entered = threading.Event() + + def slow_connect(): + entered.set() + release.wait(3) + return Session() + + lost = Session() + conn = open_connection(lost, Factory(slow_connect)) + try: + lost.drop() + assert entered.wait(3) + errors = [] + + def execute(): + try: + conn.cursor().execute("SELECT 1") + except Exception as error: + errors.append(error) + + executing = in_thread(execute) + executing.join(0.1) + assert executing.is_alive() + conn.close() + release.set() + executing.join(3) + assert len(errors) == 1 + assert isinstance(errors[0], ConnectionLostError) + finally: + release.set() + conn.close() + + +# The transport factory from connect_direct() + + +def test_connect_direct_reconnects_to_the_same_websocket(): + first, again = Session(), new_session("running") + headers = {"X-API-Key": "key"} + with patch(WS_CONNECT, side_effect=[first, again]) as ws_connect: + conn = connect_direct( + "wss://compute.example.com/sql/org/session-id", + headers=headers, + read_timeout=READ_TIMEOUT, + session_id="session-1", + stale_query_probe_seconds=PROBE, + reconnect_timeout_seconds=5, + ) + try: + headers["X-API-Key"] = "changed" + cursor = conn.cursor() + cursor.execute("SELECT 1", store=STORE) + first.drop() + assert wait_until(lambda: again.requests("retrieve_results")) + initial, reconnect = ws_connect.call_args_list + assert reconnect.kwargs["uri"] == initial.kwargs["uri"] + # What the connection was opened with, not the caller's later edit. + assert reconnect.kwargs["additional_headers"] == {"X-API-Key": "key"} + # A single attempt, bounded by the time left to reconnect in. + assert 0 < reconnect.kwargs["open_timeout"] <= 5 + for key in ("max_size", "close_timeout", "ssl"): + assert reconnect.kwargs[key] == initial.kwargs[key] + finally: + conn.close() + + +def invalid_status(status): + return InvalidStatus(Response(status, "", Headers())) + + +@pytest.mark.parametrize("status", [401, 403, 404]) +def test_connect_direct_gives_up_when_the_session_refuses(status): + first = Session() + with patch(WS_CONNECT, side_effect=[first, invalid_status(status)]) as ws_connect: + conn = connect_direct( + "wss://compute/sql", + read_timeout=READ_TIMEOUT, + reconnect_timeout_seconds=30, + ) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") + first.drop() + assert isinstance(outcome(cursor).error, ConnectionLostError) + assert ws_connect.call_count == 2 + finally: + conn.close() + + +@pytest.mark.parametrize("status", [429, 502, 503]) +def test_connect_direct_retries_a_transient_refusal(status): + first, again = Session(), new_session("running") + with patch( + WS_CONNECT, side_effect=[first, invalid_status(status), again] + ) as ws_connect: + conn = connect_direct( + "wss://compute/sql", + read_timeout=READ_TIMEOUT, + reconnect_timeout_seconds=30, + ) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") + first.drop() + assert wait_until(lambda: again.requests("retrieve_results")) + assert ws_connect.call_count == 3 + finally: + conn.close() + + +def test_connect_direct_does_not_reconnect_once_cancelled(): + cancel = threading.Event() + first = Session() + with patch(WS_CONNECT, side_effect=[first]) as ws_connect: + conn = connect_direct( + "wss://compute/sql", + read_timeout=READ_TIMEOUT, + cancel_event=cancel, + reconnect_timeout_seconds=30, + ) + try: + cursor = conn.cursor() + cursor.execute("SELECT 1") + cancel.set() + first.drop() + assert isinstance(outcome(cursor).error, ConnectionLostError) + assert ws_connect.call_count == 1 + finally: + conn.close() diff --git a/wherobots/db/__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 3d1f68a..5eff23a 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -1,6 +1,7 @@ import json import logging import math +import random import textwrap import threading import time @@ -16,10 +17,21 @@ import websockets.exceptions import websockets.sync.client -from .constants import DEFAULT_READ_TIMEOUT_SECONDS, DEFAULT_STALE_QUERY_PROBE_SECONDS +from .constants import ( + DEFAULT_READ_TIMEOUT_SECONDS, + DEFAULT_RECONNECT_TIMEOUT_SECONDS, + DEFAULT_STALE_QUERY_PROBE_SECONDS, +) from ._transport import abort_connection, socket_writable from .cursor import Cursor -from .errors import NotSupportedError, OperationalError +from .errors import ( + ConnectionLostError, + ExecutionCancelledError, + ExecutionNotFoundError, + InterfaceError, + NotSupportedError, + OperationalError, +) from .models import ExecutionResult, ProgressInfo, Store, StoreResult from .types import ( RequestKind, @@ -92,7 +104,7 @@ def _handshake_bound(ws: Any) -> float | None: Its execution cache is a small LRU that can evict a query that is still running, whose terminal state_updated still arrives (it's pushed from the query's own done-callback). So for a probe, this means "can't answer", not -"failed". +"failed". For a re-attach, it means the session doesn't know the execution. """ _PROBED_STATES = frozenset( @@ -105,9 +117,20 @@ def _handshake_bound(ws: Any) -> float | None: ) """States in which the driver is waiting on a session event it can re-ask for.""" +_RECONNECT_FIRST_DELAY_SECONDS = 0.5 +"""Wait after the first failed reconnect attempt; it doubles per attempt.""" + +_RECONNECT_MAX_DELAY_SECONDS = 5.0 +"""Upper bound on the wait between reconnect attempts.""" + +_RECONNECT_STABLE_SECONDS = 10.0 +"""A connection lost again sooner than this after reconnecting continues the +same recovery, with its deadline and backoff. Otherwise a session that accepts +the connection and drops it at once would be reconnected to forever.""" -def _probe_interval(value: Any) -> float | None: - """The validated stale-query probe interval, or ``None`` to disable it. + +def _positive_seconds(value: Any) -> float | None: + """A validated interval setting, or ``None`` to disable what it times. Like ``_handshake_bound``, returns ``None`` for a non-numeric (including ``bool``), non-finite or non-positive value. Never raises ``Exception``: @@ -180,6 +203,32 @@ class _Watch: """The one probe allowed after ``results_requested`` was sent.""" +@dataclass +class _Attach: + """Bookkeeping for a query re-attached after reconnecting. + + The session sent its events to the lost connection until then. Owned like + ``_Watch``. + """ + + reply_outstanding: bool = True + """An attach request was sent and no reply to a retrieve has arrived since. + + The attach is the query's first request on this transport, so the first + reply to a retrieve normally answers it. (After a pushed success, the + normal retrieve's reply can come first; it carries the results anyway.) + A probe sent while it's outstanding repeats the attach, so the order of + 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 +237,17 @@ class Query: handler: Callable[[Any], None] store: Store | None = None watch: _Watch = field(default_factory=_Watch) + attach: _Attach | None = None + """Set once the query is re-attached after reconnecting.""" + cancel_requested: bool = False + """This connection sent a cancel for the query. + + Unlike the reader-owned fields, a user thread writes it, in ``__send`` + under ``__lock``, once the cancel is admitted and just before it goes out: + so it's set by the time the session can report the cancellation, but not + while the cancel still waits (say, for a reconnect), when a cancellation + the session reports can't be this connection's. Read under ``__lock`` + too.""" class Connection: @@ -203,6 +263,12 @@ class Connection: A background thread listens for events from the SQL session, and handles update to the corresponding query state. Queries are tracked by their unique execution ID. + + Given a ``transport_factory`` (``connect()`` and ``connect_direct()`` pass one), the + background thread also recovers from a lost WebSocket: it reconnects to the same SQL + session, for up to ``reconnect_timeout_seconds``, and re-attaches every pending query, so + their results still arrive. Only when that fails do pending queries fail, with + ``ConnectionLostError``. """ def __init__( @@ -215,13 +281,50 @@ def __init__( session_id: str | None = None, stale_query_probe_seconds: float | None = DEFAULT_STALE_QUERY_PROBE_SECONDS, session_fallback: str | None = None, + *, + transport_factory: ( + Callable[[float], websockets.sync.client.ClientConnection] | None + ) = None, + reconnect_timeout_seconds: float | None = DEFAULT_RECONNECT_TIMEOUT_SECONDS, ): + """ + :param transport_factory: Opens a new WebSocket to the same SQL session, given + the seconds left to do so in, for automatic reconnection. It raises + ``InterfaceError`` if reconnecting can't succeed (the session refused the + connection, or the caller cancelled), and any other exception for a + failed attempt worth retrying. Without it, a lost transport fails the + pending queries at once. + :param reconnect_timeout_seconds: How long to keep trying to reconnect after + losing the transport. ``None``, or an invalid value (non-numeric, + non-finite, or not positive), disables automatic reconnection. + """ self.__ws = ws + self.__transport_factory = transport_factory + self.__reconnect_within = ( + _positive_seconds(reconnect_timeout_seconds) + if transport_factory is not None + else None + ) + if ( + transport_factory is not None + and self.__reconnect_within is None + and reconnect_timeout_seconds is not None + ): + logging.warning( + "Invalid reconnect_timeout_seconds (%s); automatic reconnection " + "disabled.", + _describe_invalid(reconnect_timeout_seconds), + ) + # The current recovery, which a loss soon after reconnecting continues. + # Only the reader thread touches these. + self.__recovery_deadline = 0.0 + self.__reconnect_delay = _RECONNECT_FIRST_DELAY_SECONDS + self.__reconnected_at: float | None = None self.__results_format = results_format self.__data_compression = data_compression self.__geometry_representation = geometry_representation self.__progress_handler: ProgressHandler | None = None - self.__probe_after = _probe_interval(stale_query_probe_seconds) + self.__probe_after = _positive_seconds(stale_query_probe_seconds) if self.__probe_after is None and stale_query_probe_seconds is not None: logging.warning( "Invalid stale_query_probe_seconds (%s); stale-query probes disabled.", @@ -334,7 +437,12 @@ def __main_loop(self) -> None: """Main background loop listening for messages from the SQL session.""" logging.info("Starting background connection handling loop...") try: - self.__receive_loop() + # The receive loop returns when the transport is gone: carry on over + # a new one, or stop. + while True: + self.__receive_loop() + if not self.__reconnect(): + return finally: self.__fail_pending() @@ -353,10 +461,10 @@ def __receive_loop(self) -> None: # Expected, retry next time continue except websockets.exceptions.ConnectionClosed: - logging.info("Connection closed; stopping main loop.") + logging.info("SQL session connection closed.") return except _TransportError: - logging.exception("SQL session transport failed; stopping main loop") + logging.exception("SQL session transport failed") return except Exception as e: logging.exception("Error handling message from SQL session", exc_info=e) @@ -370,6 +478,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 re-attached after a reconnect is probed the same way, which is + how the driver learns its outcome from an older session that ignores + the attach. + Runs on the reader thread, which owns every ``Query``'s watchdog fields. Never raises: a watchdog failure must not disturb the receive loop. """ @@ -391,17 +503,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; @@ -414,6 +542,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( @@ -428,13 +558,173 @@ def __probe_stale_queries(self) -> None: except Exception: logging.exception("Stale-query probe failed") - def __connection_error(self, execution_id: str) -> OperationalError: + def __reconnect(self) -> bool: + """Connect to the same SQL session again after losing the transport. + + Runs on the reader thread once ``__receive_loop`` has returned. Tries, + with backoff, until ``reconnect_timeout_seconds`` has passed, then + re-attaches every pending query over the new transport (see + ``__resume``). Returns ``False`` if the connection must close instead: + reconnecting is off, ``close()`` was called, the session refused the + connection, or time ran out. + + Holds ``__send_lock`` throughout, so user threads' sends wait for the + outcome: an ``execute()`` then goes out over the new transport, after + the re-attaches, or fails with ``ConnectionLostError`` once the + connection is closed. + """ + within = self.__reconnect_within + factory = self.__transport_factory + if within is None or factory is None: + return False + with self.__lock: + if self.__closed: + return False + lost = self.__ws + pending = len(self.__queries) + now = time.monotonic() + # A transport lost soon after reconnecting continues that recovery, + # backing off before its first attempt. + continued = ( + self.__reconnected_at is not None + and now - self.__reconnected_at < _RECONNECT_STABLE_SECONDS + ) + if not continued: + self.__recovery_deadline = now + within + self.__reconnect_delay = _RECONNECT_FIRST_DELAY_SECONDS + logging.warning( + "Lost the connection to SQL session %s with %d pending queries; " + "reconnecting for up to %.0fs ...", + self.__session_id or "(unknown)", + pending, + max(0.0, self.__recovery_deadline - now), + ) + # A sender stalled on the lost transport holds __send_lock: make its + # send fail now. + try: + abort_connection(lost) + except Exception: + logging.debug("Could not abort the lost transport", exc_info=True) + with self.__send_lock: + ws = self.__open_transport(factory, back_off_first=continued) + if ws is None: + return False + return self.__resume(ws) + + def __open_transport( + self, + factory: Callable[[float], websockets.sync.client.ClientConnection], + back_off_first: bool, + ) -> websockets.sync.client.ClientConnection | None: + """Open a new transport, retrying until the recovery's deadline. + + ``None`` if it gave up, or the connection is being closed: ``close()`` + cuts a backoff short, though not an attempt in progress. + """ + session = self.__session_id or "(unknown)" + attempts = 0 + while True: + if attempts or back_off_first: + remaining = self.__recovery_deadline - time.monotonic() + if remaining > 0: + # Jittered, so clients dropped together don't retry together. + delay = self.__reconnect_delay * random.uniform(0.5, 1.0) + self.__reconnect_delay = min( + self.__reconnect_delay * 2, _RECONNECT_MAX_DELAY_SECONDS + ) + if self.__shutdown_done.wait(min(delay, remaining)): + return None + with self.__lock: + if self.__closed: + return None + remaining = self.__recovery_deadline - time.monotonic() + if remaining <= 0: + logging.error( + "Could not reconnect to SQL session %s in time (%d attempts).", + session, + attempts, + ) + return None + attempts += 1 + try: + return factory(remaining) + except InterfaceError as error: + logging.error( + "Could not reconnect to SQL session %s: %s", session, error + ) + return None + except Exception as error: + logging.info( + "Reconnect attempt %d to SQL session %s failed: %s", + attempts, + session, + error, + ) + + def __resume(self, ws: websockets.sync.client.ClientConnection) -> bool: + """Carry on over a new transport, re-attaching every pending query. + + Called on the reader thread with ``__send_lock`` held. Each pending + query gets a ``retrieve_results`` with ``attach``, keeping its handler + and store: a session that supports attaching makes this connection the + owner again, pushing the execution's later events here, and answers + with ``attached``; an older one ignores the flag, and is polled. A query this connection cancelled is + cancelled again after its attach, as the cancel may have been lost + with the transport. + + ``False`` if ``close()`` was called while connecting. + """ + with self.__lock: + closed = self.__closed + if not closed: + self.__ws = ws + pending = list(self.__queries.values()) + for query in pending: + # The attach is a retrieve_results; its reply has the state. + query.state = ExecutionState.RESULTS_REQUESTED + query.watch = _Watch() + query.attach = _Attach() + cancelled = [q.execution_id for q in pending if q.cancel_requested] + if closed: + try: + abort_connection(ws) + except Exception: + logging.debug("Could not abort the new transport", exc_info=True) + return False + self.__reconnected_at = time.monotonic() + logging.info( + "Reconnected to SQL session %s; re-attaching %d pending queries.", + self.__session_id or "(unknown)", + len(pending), + ) + requests = [ + self.__retrieve_request(query.execution_id, attach=True) + for query in pending + ] + [ + {"kind": RequestKind.CANCEL.value, "execution_id": execution_id} + for execution_id in cancelled + ] + try: + for request in requests: + ws.send(json.dumps(request)) + except (websockets.exceptions.ConnectionClosed, OSError): + # Lost again already: the receive loop notices, and reconnects. + logging.warning("Lost the new SQL session connection while re-attaching.") + try: + abort_connection(ws) + except Exception: + logging.debug("Could not abort the new transport", exc_info=True) + return True + + def __connection_error(self, execution_id: str) -> ConnectionLostError: message = ( f"SQL connection lost (session={self.__session_id or 'unknown'}, " 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 @@ -517,6 +807,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 re-attached query. + 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. @@ -527,6 +820,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 @@ -606,6 +906,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: @@ -615,6 +922,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. @@ -653,6 +984,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 @@ -690,6 +1054,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, @@ -701,12 +1086,48 @@ 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} after reconnecting. The statement " + "didn't reach the session before the connection was " + "lost, the session restarted, or the session no longer " + "tracks it.", + session_id=self.__session_id, + execution_id=execution_id, + ) + ) + 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 @@ -757,7 +1178,8 @@ def __send( """Send a request; return ``False`` only if non-blocking and busy. Busy means another sender holds the send lock, or the socket's send - buffer is full. + buffer is full. While the reader reconnects, it holds the send lock, so + a blocking send waits for that to finish. """ # Serialization and redaction are local work. Fail before registration, # without poisoning unrelated cursors or misreporting a transport loss. @@ -780,8 +1202,15 @@ def __send( watch = query.watch watch.last_activity = watch.last_event = time.monotonic() self.__queries[query.execution_id] = query - elif message.get("execution_id") not in self.__queries: - return True + else: + tracked = self.__queries.get(message.get("execution_id")) + if tracked is None: + return True + if message.get("kind") == RequestKind.CANCEL.value: + # Only now that it's going out (see the field). + tracked.cancel_requested = True + # The reader replaces it only while holding the send lock. + ws = self.__ws # A non-blocking send runs on the reader thread (a probe). Not # waiting for the lock isn't enough: with the socket's send buffer # full (a peer that keeps sending but stops reading), ws.send() @@ -789,10 +1218,10 @@ def __send( # hold the library's protocol mutex, which keepalive pings need # too. Treat it like a busy lock: skip, and retry at the next # check. Checked outside __lock: it's a syscall on the transport. - if not blocking and not socket_writable(self.__ws): + if not blocking and not socket_writable(ws): return False try: - self.__ws.send(request) + ws.send(request) except (websockets.exceptions.ConnectionClosed, OSError): pass # Terminate outside the send gate, before delivering errors. except Exception: @@ -805,7 +1234,16 @@ def __send( return True finally: self.__send_lock.release() - self.__fail_pending() + if self.__reconnect_within is None: + self.__fail_pending() + return True + # Make sure the reader notices the loss, so it reconnects. A query + # registered above is then re-attached: if the request didn't reach the + # session, that fails it with ExecutionNotFoundError. + try: + abort_connection(ws) + except Exception: + logging.debug("Could not abort the lost transport", exc_info=True) return True @staticmethod @@ -878,17 +1316,10 @@ 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. - - A probe never waits for ``__send_lock`` or a full socket (see - ``__probe_stale_queries`` and ``__send``). - """ - query = self.__queries.get(execution_id) - if not query: - return True - - 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, } @@ -898,6 +1329,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. @@ -908,9 +1357,9 @@ 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: + if execution_id not in self.__queries: + return request = { "kind": RequestKind.CANCEL.value, diff --git a/wherobots/db/constants.py b/wherobots/db/constants.py index 5296ee9..961a12a 100644 --- a/wherobots/db/constants.py +++ b/wherobots/db/constants.py @@ -22,6 +22,11 @@ # The API sets this on its answer to a session request when it gave the request # a session under another key than its own, naming why ("quota"). SESSION_FALLBACK_HEADER: str = "X-Wherobots-Sql-Session-Fallback" +# How long the driver keeps trying to reconnect to the SQL session after losing +# the connection. It must fit, along with the up to ~40s that keepalive takes +# to notice a dead connection, in the session's grace period for orphaned +# executions (120s by default), after which the session cancels them. +DEFAULT_RECONNECT_TIMEOUT_SECONDS: float = 60.0 MAX_MESSAGE_SIZE: int = 100 * 2**20 # 100MiB PROTOCOL_VERSION: Version = Version("1.0.0") diff --git a/wherobots/db/cursor.py b/wherobots/db/cursor.py index 585c9d1..30b8e16 100644 --- a/wherobots/db/cursor.py +++ b/wherobots/db/cursor.py @@ -95,6 +95,16 @@ 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()``, and after one that was + rejected before submission. Errors about an execution, such as + :class:`~wherobots.db.errors.ConnectionLostError`, carry the same ID. + """ + return self.__current_execution_id + def __in_flight_execution_id(self) -> str | None: """The current execution's id if its result has not yet arrived. diff --git a/wherobots/db/driver.py b/wherobots/db/driver.py index 06fe601..622577c 100644 --- a/wherobots/db/driver.py +++ b/wherobots/db/driver.py @@ -23,6 +23,7 @@ DEFAULT_CLOSE_TIMEOUT_SECONDS, DEFAULT_ENDPOINT, DEFAULT_READ_TIMEOUT_SECONDS, + DEFAULT_RECONNECT_TIMEOUT_SECONDS, DEFAULT_SESSION_TYPE, DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS, DEFAULT_STALE_QUERY_PROBE_SECONDS, @@ -129,6 +130,7 @@ def connect( extra_headers: Union[Dict[str, str], None] = None, stale_query_probe_seconds: Union[float, None] = DEFAULT_STALE_QUERY_PROBE_SECONDS, session_key: Union[str, None] = None, + reconnect_timeout_seconds: Union[float, None] = DEFAULT_RECONNECT_TIMEOUT_SECONDS, ) -> Connection: """Create a connection to a Wherobots SQL session. @@ -161,6 +163,12 @@ def connect( and never one opened without a key. When a session of the key's own would exceed a quota, the API shares one under another key instead; the connection's ``session_fallback`` then says so. + :param reconnect_timeout_seconds: How long the connection keeps trying to + reconnect to the same SQL session after losing its WebSocket, then + re-attaches its pending queries, whose results still arrive. Only if + it can't reconnect in time do they fail, with ``ConnectionLostError``. + ``None`` disables reconnecting, as does an invalid value (non-numeric, + non-finite, or not positive): pending queries then fail at once. """ if not token and not api_key: raise ValueError("At least one of `token` or `api_key` is required") @@ -307,6 +315,7 @@ def get_session_uri() -> str: .rsplit("/", 1)[-1], stale_query_probe_seconds=stale_query_probe_seconds, session_fallback=session_fallback, + reconnect_timeout_seconds=reconnect_timeout_seconds, ) @@ -337,16 +346,36 @@ def connect_direct( session_id: str | None = None, stale_query_probe_seconds: Union[float, None] = DEFAULT_STALE_QUERY_PROBE_SECONDS, session_fallback: str | None = None, + reconnect_timeout_seconds: Union[float, None] = DEFAULT_RECONNECT_TIMEOUT_SECONDS, ) -> Connection: """Connect to a SQL session's WebSocket endpoint directly. + After losing its WebSocket, the returned connection reconnects to the same + endpoint by itself, with the same headers, honouring ``cancel_event`` (see + ``reconnect_timeout_seconds``). + :param stale_query_probe_seconds: See :func:`connect`. :param session_fallback: See :attr:`Connection.session_fallback`. + :param reconnect_timeout_seconds: See :func:`connect`. """ uri_with_protocol = f"{uri}/{protocol}" ssl_context = ssl.create_default_context() ssl_context.load_verify_locations(certifi.where()) + def open_ws( + headers: Union[Dict[str, str], None], open_timeout: float + ) -> websockets.sync.client.ClientConnection: + _check_cancelled(cancel_event) + logging.info("Connecting to SQL session at %s ...", uri_with_protocol) + return websockets.sync.client.connect( + uri=uri_with_protocol, + additional_headers=headers, + max_size=MAX_MESSAGE_SIZE, + open_timeout=open_timeout, + close_timeout=DEFAULT_CLOSE_TIMEOUT_SECONDS, + ssl=ssl_context, + ) + @tenacity.retry( stop=tenacity.stop_after_attempt(5), wait=tenacity.wait_exponential(multiplier=1, min=1, max=5), @@ -362,16 +391,7 @@ def connect_direct( reraise=True, ) def ws_connect() -> websockets.sync.client.ClientConnection: - _check_cancelled(cancel_event) - logging.info("Connecting to SQL session at %s ...", uri_with_protocol) - return websockets.sync.client.connect( - uri=uri_with_protocol, - additional_headers=headers, - max_size=MAX_MESSAGE_SIZE, - open_timeout=DEFAULT_HTTP_TIMEOUT, - close_timeout=DEFAULT_CLOSE_TIMEOUT_SECONDS, - ssl=ssl_context, - ) + return open_ws(headers, DEFAULT_HTTP_TIMEOUT) try: ws = ws_connect() @@ -380,6 +400,26 @@ 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 open_transport(timeout: float) -> websockets.sync.client.ClientConnection: + # A single attempt: the connection retries until its own deadline. + # Straight to the same session's WebSocket: connect() resolves the + # session through the REST API, which could start a new session when + # this one is gone. A dead session must fail fast instead. + try: + return open_ws(reconnect_headers, min(DEFAULT_HTTP_TIMEOUT, timeout)) + except websockets.exceptions.InvalidStatus as e: + status = e.response.status_code + if 400 <= status < 500 and status not in (408, 429): + # Gone, or no longer authorized: retrying can't help. + raise InterfaceError( + f"SQL session refused the connection (HTTP {status})" + ) from e + raise + return Connection( ws, read_timeout=read_timeout, @@ -389,4 +429,6 @@ def ws_connect() -> websockets.sync.client.ClientConnection: session_id=session_id, stale_query_probe_seconds=stale_query_probe_seconds, session_fallback=session_fallback, + transport_factory=open_transport, + reconnect_timeout_seconds=reconnect_timeout_seconds, ) diff --git a/wherobots/db/errors.py b/wherobots/db/errors.py index 4e18110..62eca0e 100644 --- a/wherobots/db/errors.py +++ b/wherobots/db/errors.py @@ -24,3 +24,80 @@ class ProgrammingError(DatabaseError): class NotSupportedError(DatabaseError): pass + + +class ConnectionLostError(OperationalError): + """The connection to the SQL session was lost while a query was pending. + + A connection from ``connect()`` or ``connect_direct()`` first reconnects + by itself and re-attaches its pending queries; this means that failed + within ``reconnect_timeout_seconds``, or that reconnecting is turned off. + The query may still have run, or be running, on the SQL session: its + outcome, for a write whether it committed, is unknown. Raised by + ``Cursor.execute()`` itself, rather than by a fetch, it means the + connection was already lost and the statement was never sent. + + Attributes: + session_id: The SQL session's ID, or None if the connection didn't + 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): + """A re-attached execution was cancelled, and not by this connection. + + The SQL session cancelled it before or while this connection re-attached + to it after reconnecting, for example because its grace period for a lost + connection ran out, or another client cancelled it. Its outcome is unknown: for a write, whether it + committed. + + Attributes: + 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 an execution to re-attach to. + + After reconnecting, the session didn't know the execution: the statement + didn't reach it before the connection was lost, the session restarted, or + it no longer tracks the execution. Its outcome can't be recovered. + + Attributes: + 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