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