diff --git a/README.md b/README.md index 07edb56..d99cfb0 100644 --- a/README.md +++ b/README.md @@ -277,3 +277,8 @@ users may find useful: caller needs to interrupt it (e.g. on client disconnect or timeout). The event is checked before each HTTP request, between retry attempts, and before the WebSocket handshake. +* `stale_query_probe_seconds`: how long a running query may go without + any message from the SQL session before the driver asks the session + 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. diff --git a/tests/test_disconnect.py b/tests/test_disconnect.py index 5f420a0..da70d8c 100644 --- a/tests/test_disconnect.py +++ b/tests/test_disconnect.py @@ -582,9 +582,10 @@ def slow_close(): def test_slow_handshake_does_not_consume_reader_join(): # The handshake uses most of close_timeout. A long read_timeout keeps the # reader parked in recv(), so only the end-of-stream delivered 0.3s after - # the handshake wakes it; close() must still be waiting then. + # the handshake wakes it; close() must still be waiting then. The watchdog + # would cap recv() at its 1s check interval, so it's off here. ws = Transport() - conn = Connection(ws, read_timeout=10) + conn = Connection(ws, read_timeout=10, stale_query_probe_seconds=None) def slow_close(): time.sleep(0.9) diff --git a/tests/test_driver.py b/tests/test_driver.py index 6c9cb7d..3295cd0 100644 --- a/tests/test_driver.py +++ b/tests/test_driver.py @@ -75,6 +75,23 @@ def test_session_id_is_derived_from_status_url(self, mock_post, mock_get, suffix ) assert kwargs["session_id"] == "test-id" + @pytest.mark.parametrize("probe_seconds", [None, 5.0]) + @patch("wherobots.db.driver.requests.get") + @patch("wherobots.db.driver.requests.post") + def test_stale_query_probe_seconds_is_forwarded( + self, mock_post, mock_get, probe_seconds + ): + _, kwargs = _run_connect_full( + mock_post, mock_get, stale_query_probe_seconds=probe_seconds + ) + assert kwargs["stale_query_probe_seconds"] == probe_seconds + + @patch("wherobots.db.driver.requests.get") + @patch("wherobots.db.driver.requests.post") + 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 + @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_watchdog.py b/tests/test_watchdog.py new file mode 100644 index 0000000..43d58b9 --- /dev/null +++ b/tests/test_watchdog.py @@ -0,0 +1,808 @@ +"""The stale-query watchdog re-asks the session when a terminal event is lost.""" + +import json +import logging +import os +import queue +import re +import socket +import threading +import time +from unittest.mock import MagicMock, patch + +import cbor2 +import pytest +from websockets.exceptions import ConnectionClosedOK + +from wherobots.db._transport import socket_writable +from wherobots.db.connection import Connection +from wherobots.db.driver import connect_direct +from wherobots.db.errors import OperationalError +from wherobots.db.models import Store +from wherobots.db.types import ExecutionState, 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: a larger N, so a +briefly stalled test thread on a loaded runner can't produce a false probe.""" + +NOT_FOUND = "Execution not found" + +READ_TIMEOUT = 0.01 + +STORE = Store.for_download(format=StorageFormat.PARQUET) + +RESULT_URI = "https://presigned.example.com/results.parquet" + + +class Session: + """A fake SQL session transport. + + Records every request with its send time, and lets a test answer probes + (``retrieve_results`` requests) through ``on_probe``. ``before_send`` runs + before a request is recorded, so a test can stall a send. + """ + + def __init__(self, on_probe=None): + self.incoming = queue.Queue() + self.sent = [] + self.on_probe = on_probe + self.before_send = None + 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) + if self.before_send is not None: + self.before_send(message) + self.sent.append((time.monotonic(), message)) + if message["kind"] == "retrieve_results" and self.on_probe is not None: + self.on_probe(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 probe_times(self, execution_id=None): + return [t for t, _ in self.requests("retrieve_results", execution_id)] + + +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 connect(session, **kwargs): + kwargs.setdefault("stale_query_probe_seconds", PROBE) + return Connection(session, read_timeout=READ_TIMEOUT, **kwargs) + + +def execute(conn, session, sql="SELECT 1", store=None): + """Start a query that the session reports as running, then goes quiet.""" + cursor = conn.cursor() + cursor.execute(sql, store=store) + execution_id = session.requests("execute_sql")[-1][1]["execution_id"] + session.put(kind="state_updated", execution_id=execution_id, state="running") + return cursor, execution_id + + +def results_of(cursor): + return cursor._Cursor__queue + + +def tracked(conn, execution_id): + return conn._Connection__queries.get(execution_id) + + +def reply(session, execution_id, state="succeeded", **fields): + session.put( + kind="execution_result", + execution_id=execution_id, + state=state, + results=None, + **fields, + ) + + +def test_probe_reply_with_result_uri_completes_store_query(caplog): + caplog.set_level(logging.INFO) + session = Session( + on_probe=lambda s, m: reply( + s, m["execution_id"], result_uri=RESULT_URI, size=42 + ) + ) + conn = connect(session) + try: + # The terminal state_updated is never delivered. + cursor, execution_id = execute(conn, session, store=STORE) + result = results_of(cursor).get(timeout=3) + assert result.error is None + assert result.store_result.result_uri == RESULT_URI + assert result.store_result.size == 42 + assert tracked(conn, execution_id) is None + assert len(session.probe_times()) == 1 + assert f"No events for query {execution_id}" in caplog.text + assert f"Query {execution_id} recovered" in caplog.text + finally: + conn.close() + + +def test_store_query_succeeded_without_result_uri_keeps_waiting(caplog): + caplog.set_level(logging.INFO) + session = Session(on_probe=lambda s, m: reply(s, m["execution_id"])) + conn = connect(session) + try: + cursor, execution_id = execute(conn, session, store=STORE) + assert wait_until(lambda: "no result location" in caplog.text) + warnings = [r for r in caplog.records if "no result location" in r.message] + assert warnings[0].levelno == logging.WARNING + # Not completed empty, and not re-probed: the answer wouldn't change. + time.sleep(PROBE * 8) + assert results_of(cursor).empty() + assert tracked(conn, execution_id).state == ExecutionState.SUCCEEDED + assert len(session.probe_times()) == 1 + # The genuine terminal event still completes it. + session.put( + kind="state_updated", + execution_id=execution_id, + state="succeeded", + result_uri=RESULT_URI, + size=7, + ) + assert results_of(cursor).get(timeout=3).store_result.result_uri == RESULT_URI + finally: + conn.close() + + +@pytest.mark.parametrize("late_event_first", [True, False]) +def test_late_state_updated_racing_probe_reply_delivers_once(late_event_first): + def on_probe(s, m): + late = dict( + kind="state_updated", + execution_id=m["execution_id"], + state="succeeded", + result_uri=RESULT_URI, + size=1, + ) + probe_reply = dict( + kind="execution_result", + execution_id=m["execution_id"], + state="succeeded", + results=None, + result_uri=RESULT_URI, + size=1, + ) + for message in (late, probe_reply) if late_event_first else (probe_reply, late): + s.incoming.put(message) + + session = Session(on_probe=on_probe) + conn = connect(session) + try: + cursor, execution_id = execute(conn, session, store=STORE) + assert results_of(cursor).get(timeout=3).store_result.result_uri == RESULT_URI + assert wait_until(session.incoming.empty) + time.sleep(PROBE * 4) + assert results_of(cursor).empty() + assert tracked(conn, execution_id) is None + assert len(session.probe_times()) == 1 + finally: + conn.close() + + +def test_disabled_watchdog_never_probes(): + session = Session() + conn = connect(session, stale_query_probe_seconds=None) + try: + cursor, execution_id = execute(conn, session, store=STORE) + time.sleep(PROBE * 8) + assert session.requests("retrieve_results") == [] + assert tracked(conn, execution_id) is not None + finally: + conn.close() + + +def test_steady_traffic_for_one_query_does_not_starve_anothers_watchdog(): + # read_timeout is well above the progress cadence, so recv() never times + # out while B is streaming. + session = Session() + conn = Connection(session, read_timeout=0.5, stale_query_probe_seconds=SLOW_PROBE) + stop = threading.Event() + try: + _, quiet = execute(conn, session) + _, busy = execute(conn, session) + + def stream_progress(): + while not stop.is_set(): + session.put(kind="execution_progress", execution_id=busy) + time.sleep(0.005) + + streamer = threading.Thread(target=stream_progress, daemon=True) + streamer.start() + assert wait_until(lambda: session.probe_times(quiet), timeout=2) + # Progress events are activity: the busy query is never probed. + assert session.probe_times(busy) == [] + finally: + stop.set() + conn.close() + + +def test_unanswered_probe_log_reports_silence_since_the_last_event(caplog): + caplog.set_level(logging.INFO) + session = Session() # Never answers probes. + conn = connect(session, stale_query_probe_seconds=SLOW_PROBE) + try: + _, execution_id = execute(conn, session) + assert wait_until(lambda: len(session.probe_times()) >= 2, timeout=5) + last_event = tracked(conn, execution_id).watch.last_event + second_probe_at = session.probe_times()[1] + finally: + conn.close() + logs = [ + r.getMessage() for r in caplog.records if "No events for query" in r.message + ] + assert len(logs) >= 2, logs + first, second = logs[:2] + assert "(probe 1)" in first + assert "(probe 2, previous probe unanswered)" in second + logged = float(re.search(r" in ([0-9.]+)s while ", second).group(1)) + # Probe 1 goes out N after the last event, probe 2 a further 2N later: the + # silence is >= 3N, not the 2N since probe 1. + assert logged >= SLOW_PROBE * 3 - 0.05, logged + # Loose on purpose: the >= 3N check above is what catches a regression; + # this only confirms the figure tracks the measured silence. + assert abs(logged - (second_probe_at - last_event)) < SLOW_PROBE, logged + + +def test_successive_probes_back_off_to_a_cap(): + session = Session( + on_probe=lambda s, m: reply(s, m["execution_id"], state="running") + ) + conn = connect(session) + try: + _, execution_id = execute(conn, session) + registered = session.requests("execute_sql")[0][0] + assert wait_until(lambda: len(session.probe_times()) >= 5, timeout=10) + times = session.probe_times()[:5] + assert times[0] - registered >= PROBE + gaps = [later - earlier for earlier, later in zip(times, times[1:])] + # N, then 2N, 4N, 8N, capped at 8N; never a tight loop. The upper + # bound (next doubling) still tells each step, and the cap, apart. + for gap, expected in zip(gaps, [2, 4, 8, 8]): + assert PROBE * expected <= gap < PROBE * expected * 2, gaps + finally: + conn.close() + + +def test_probe_recovers_non_store_query_and_requests_configured_format(): + rows = b'[{"x": 1}, {"x": 2}]' + + def on_probe(s, m): + s.incoming.put( + cbor2.dumps( + { + "kind": "execution_result", + "execution_id": m["execution_id"], + "state": "succeeded", + "results": {"result_bytes": rows, "format": "json"}, + } + ) + ) + + session = Session(on_probe=on_probe) + conn = connect(session, results_format=ResultsFormat.JSON) + try: + cursor, execution_id = execute(conn, session) + result = results_of(cursor).get(timeout=3) + assert result.error is None + assert result.results == [{"x": 1}, {"x": 2}] + (_, probe), *_ = session.requests("retrieve_results") + assert probe == { + "kind": "retrieve_results", + "execution_id": execution_id, + "format": "json", + } + finally: + conn.close() + + +def test_lost_retrieve_results_reply_is_re_requested(): + asked = [] + + def on_probe(s, m): + # Lose the reply to the driver's own request; answer the probe. + asked.append(m) + if len(asked) > 1: + s.incoming.put( + cbor2.dumps( + { + "kind": "execution_result", + "execution_id": m["execution_id"], + "state": "succeeded", + "results": {"result_bytes": b"[]", "format": "json"}, + } + ) + ) + + session = Session(on_probe=on_probe) + conn = connect(session) + try: + cursor, execution_id = execute(conn, session) + session.put(kind="state_updated", execution_id=execution_id, state="succeeded") + result = results_of(cursor).get(timeout=3) + assert result.error is None + assert result.results == [] + assert len(asked) == 2 + finally: + conn.close() + + +@pytest.mark.parametrize("state", ["running", "pending"]) +def test_non_terminal_probe_reply_keeps_waiting_and_probing(state): + session = Session(on_probe=lambda s, m: reply(s, m["execution_id"], state=state)) + conn = connect(session) + try: + cursor, execution_id = execute(conn, session) + assert wait_until(lambda: len(session.probe_times()) >= 2) + query = tracked(conn, execution_id) + assert query is not None + assert query.state == ExecutionState[state.upper()] + assert results_of(cursor).empty() + # The probe replies aren't events of their own: backoff keeps growing. + first, second = session.probe_times()[:2] + assert second - first >= PROBE * 2 + finally: + conn.close() + + +def test_genuine_events_postpone_probes_and_reset_backoff(): + session = Session( + on_probe=lambda s, m: reply(s, m["execution_id"], state="running") + ) + conn = connect(session, stale_query_probe_seconds=SLOW_PROBE) + try: + _, execution_id = execute(conn, session) + assert wait_until(lambda: len(session.probe_times()) >= 2) + query = tracked(conn, execution_id) + assert wait_until(lambda: query.watch.probes_sent >= 2) + # Progress is activity: no probes while it flows, and backoff resets. + deadline = time.monotonic() + SLOW_PROBE * 4 + probes_before = len(session.probe_times()) + while time.monotonic() < deadline: + session.put(kind="execution_progress", execution_id=execution_id) + time.sleep(SLOW_PROBE / 5) + assert wait_until(lambda: query.watch.probes_sent == 0, timeout=1) + assert len(session.probe_times()) == probes_before + finally: + conn.close() + + +def test_probe_not_found_keeps_waiting_for_an_evicted_query(caplog): + # The session's small LRU cache can evict a query that is still running. + # Its terminal state_updated still arrives (from the future's callback), + # but a probe gets "Execution not found". That must not fail the query. + session = Session( + on_probe=lambda s, m: s.put( + kind="error", execution_id=m["execution_id"], message=NOT_FOUND + ) + ) + conn = connect(session) + try: + cursor, execution_id = execute(conn, session, store=STORE) + assert wait_until(lambda: "no longer tracks execution" in caplog.text) + time.sleep(PROBE * 8) + assert results_of(cursor).empty() + assert tracked(conn, execution_id) is not None + # Not probeable any more: no further probes. + assert len(session.probe_times()) == 1 + session.put( + kind="state_updated", + execution_id=execution_id, + state="succeeded", + result_uri=RESULT_URI, + size=3, + ) + result = results_of(cursor).get(timeout=3) + assert result.error is None + assert result.store_result.result_uri == RESULT_URI + finally: + conn.close() + + +def test_progress_while_probe_outstanding_does_not_unmask_not_found(caplog): + # A progress event between the probe and its reply must not make the + # not-found reply look like an answer to a normal request. + def on_probe(s, m): + s.put(kind="execution_progress", execution_id=m["execution_id"]) + s.put(kind="error", execution_id=m["execution_id"], message=NOT_FOUND) + + session = Session(on_probe=on_probe) + conn = connect(session) + try: + cursor, execution_id = execute(conn, session) + assert wait_until(lambda: "no longer tracks execution" in caplog.text) + assert results_of(cursor).empty() + assert tracked(conn, execution_id) is not None + finally: + conn.close() + + +def test_running_state_update_while_probe_outstanding_does_not_unmask_not_found( + caplog, +): + # An evicted query that is still running: a genuine state_updated: running + # can arrive between the probe and its not-found reply. The reply is still + # the probe's and must not fail the query. + def on_probe(s, m): + s.put(kind="state_updated", execution_id=m["execution_id"], state="running") + s.put(kind="error", execution_id=m["execution_id"], message=NOT_FOUND) + + session = Session(on_probe=on_probe) + conn = connect(session) + try: + cursor, execution_id = execute(conn, session, store=STORE) + assert wait_until(lambda: "no longer tracks execution" in caplog.text) + time.sleep(PROBE * 8) + assert results_of(cursor).empty() + assert tracked(conn, execution_id) is not None + # Not probeable any more: no further probes. + assert len(session.probe_times()) == 1 + session.put( + kind="state_updated", + execution_id=execution_id, + state="succeeded", + result_uri=RESULT_URI, + size=3, + ) + result = results_of(cursor).get(timeout=3) + assert result.error is None + assert result.store_result.result_uri == RESULT_URI + finally: + conn.close() + + +def test_not_found_for_the_normal_results_request_still_fails(): + # After the genuine state_updated: succeeded, there is no other terminal + # event to wait for: a not-found reply fails the query, as without the + # watchdog. + session = Session( + on_probe=lambda s, m: s.put( + kind="error", execution_id=m["execution_id"], message=NOT_FOUND + ) + ) + conn = connect(session) + try: + cursor, execution_id = execute(conn, session) + session.put(kind="state_updated", execution_id=execution_id, state="succeeded") + error = results_of(cursor).get(timeout=3).error + assert isinstance(error, OperationalError) + assert NOT_FOUND in str(error) + assert tracked(conn, execution_id) is None + finally: + conn.close() + + +def test_not_found_after_results_requested_probe_still_fails(): + # The normal retrieve's reply is lost, and the one probe that follows it + # gets not-found: the terminal event has already arrived, so fail. + asked = [] + + def on_probe(s, m): + asked.append(m) + if len(asked) > 1: + s.put(kind="error", execution_id=m["execution_id"], message=NOT_FOUND) + + session = Session(on_probe=on_probe) + conn = connect(session) + try: + cursor, execution_id = execute(conn, session) + session.put(kind="state_updated", execution_id=execution_id, state="succeeded") + error = results_of(cursor).get(timeout=3).error + assert isinstance(error, OperationalError) + assert NOT_FOUND in str(error) + assert len(asked) == 2 + finally: + conn.close() + + +def test_results_requested_gets_one_probe_no_earlier_than_4n(): + session = Session() + conn = connect(session) + try: + _, execution_id = execute(conn, session) + session.put(kind="state_updated", execution_id=execution_id, state="succeeded") + assert wait_until(lambda: len(session.probe_times()) >= 2) + requested, probed = session.probe_times() + assert probed - requested >= PROBE * 4 + # Never a second probe: duplicate retrieves re-send the full result. + time.sleep(PROBE * 16) + assert len(session.probe_times()) == 2 + finally: + conn.close() + + +def test_stalled_user_send_does_not_hold_up_other_results(): + session = Session() + conn = connect(session) + sending, release = threading.Event(), threading.Event() + + def stall_new_queries(message): + if message["kind"] == "execute_sql" and message["statement"] == "SLOW": + sending.set() + release.wait(timeout=5) + + try: + _, quiet = execute(conn, session) + cursor, ready = execute(conn, session, sql="SELECT 2") + session.before_send = stall_new_queries + sender = threading.Thread(target=conn.cursor().execute, args=("SLOW",)) + sender.start() + assert sending.wait(timeout=1) + # Let a probe for the quiet query come due while the send is stalled. + time.sleep(PROBE * 3) + started = time.monotonic() + session.put( + kind="state_updated", + execution_id=ready, + state="cancelled", + ) + result = results_of(cursor).get(timeout=3) + delay = time.monotonic() - started + release.set() + sender.join(timeout=3) + assert result.error is None + assert delay < 1.0, delay + # The skipped probe goes out once the send lock is free. + assert wait_until(lambda: session.probe_times(quiet)) + finally: + release.set() + conn.close() + + +def fill_send_buffer(sock): + """Write to ``sock`` until its send buffer is full; restore its timeout.""" + timeout = sock.gettimeout() + sock.setblocking(False) + try: + # Large writes first, then single bytes: some kernels accept a last + # partial write. + for chunk in (b"x" * 65536, b"x"): + try: + while True: + sock.send(chunk) + except BlockingIOError: + pass + finally: + sock.settimeout(timeout) + + +class SocketSession(Session): + """A fake session whose requests are written to a real socket. + + Like websockets' sync client, ``send`` blocks while the socket's send + buffer is full (bounded here, so a regression fails instead of hanging). + """ + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.socket, self.peer = socket.socketpair() + self.socket.settimeout(5) + self.before_send = lambda message: self.socket.sendall( + json.dumps(message).encode() + ) + + def fill(self): + """Fill the send buffer, as a peer that stops reading would.""" + fill_send_buffer(self.socket) + + def drain(self): + self.peer.setblocking(False) + try: + while self.peer.recv(1 << 20): + pass + except BlockingIOError: + pass + + def dispose(self): + self.socket.close() + self.peer.close() + + +def test_probe_skips_a_full_socket_instead_of_stalling_the_reader(): + session = SocketSession() + conn = connect(session) + try: + _, quiet = execute(conn, session) + cursor, ready = execute(conn, session, sql="SELECT 2") + session.fill() + assert not socket_writable(session) + # Let a probe for the quiet query come due while the socket is full. + time.sleep(PROBE * 3) + started = time.monotonic() + session.put(kind="state_updated", execution_id=ready, state="cancelled") + result = results_of(cursor).get(timeout=3) + delay = time.monotonic() - started + assert result.error is None + assert delay < 1.0, delay + assert session.probe_times(quiet) == [] + assert tracked(conn, quiet) is not None + # The skipped probe goes out once the peer reads again. + session.drain() + assert wait_until(lambda: session.probe_times(quiet)) + finally: + conn.close() + session.dispose() + + +def test_socket_writable_proceeds_without_a_real_socket(): + # A failure must not disable probing: anything unusable counts as writable. + assert socket_writable(Session()) # MagicMock socket + assert socket_writable(object()) # No socket at all + + class Broken: + @property + def socket(self): + raise RuntimeError("no socket") + + assert socket_writable(Broken()) + closed = MagicMock() + closed.socket, peer = socket.socketpair() + closed.socket.close() + peer.close() + assert socket_writable(closed) + + +def test_socket_writable_handles_fds_above_fd_setsize(): + # select.select() rejects fds >= FD_SETSIZE (1024); a server holding many + # connections reaches them. + resource = pytest.importorskip("resource") # POSIX-only + target = 2048 + soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE) + if soft <= target: + if hard != resource.RLIM_INFINITY and hard <= target: + pytest.skip("RLIMIT_NOFILE too low for a high fd") + resource.setrlimit(resource.RLIMIT_NOFILE, (target + 1, hard)) + a, b = socket.socketpair() + high = None + try: + os.dup2(a.fileno(), target) + high = socket.socket(fileno=target) + ws = MagicMock() + ws.socket = high + assert socket_writable(ws) + fill_send_buffer(high) + assert not socket_writable(ws) + finally: + if high is not None: + high.close() + a.close() + b.close() + resource.setrlimit(resource.RLIMIT_NOFILE, (soft, hard)) + + +def test_recovery_is_logged_only_for_a_successful_probe_reply(caplog): + caplog.set_level(logging.INFO) + session = Session( + on_probe=lambda s, m: s.incoming.put( + cbor2.dumps( + { + "kind": "execution_result", + "execution_id": m["execution_id"], + "state": "succeeded", + "results": {"result_bytes": b"not json", "format": "json"}, + } + ) + ) + ) + conn = connect(session) + try: + cursor, execution_id = execute(conn, session) + assert isinstance(results_of(cursor).get(timeout=3).error, OperationalError) + assert "recovered" not in caplog.text + finally: + conn.close() + + +def test_watchdog_runs_when_read_timeout_is_none(): + # recv() with no timeout would block forever on a quiet connection, which + # is exactly the lost-terminal-event shape. + session = Session( + on_probe=lambda s, m: reply(s, m["execution_id"], state="running") + ) + conn = Connection(session, read_timeout=None, stale_query_probe_seconds=PROBE) + try: + _, execution_id = execute(conn, session) + assert wait_until(lambda: session.probe_times(execution_id)) + finally: + conn.close() + + +@pytest.mark.parametrize( + "value", + [0, -1, -0.5, float("nan"), float("inf"), float("-inf"), True, "30", object()], +) +def test_invalid_probe_interval_disables_watchdog(value, caplog): + session = Session() + conn = connect(session, stale_query_probe_seconds=value) + try: + assert "stale-query probes disabled" in caplog.text + assert conn._Connection__probe_after is None + _, execution_id = execute(conn, session) + time.sleep(PROBE * 3) + assert session.requests("retrieve_results") == [] + assert tracked(conn, execution_id) is not None + finally: + conn.close() + + +class _BadRepr: + def __repr__(self): + raise RuntimeError("no repr") + + +@pytest.mark.parametrize( + "value, expected", + [(0, "(0)"), (-5, "(-5)"), (_BadRepr(), "()")], +) +def test_invalid_probe_interval_warning_names_the_value(value, expected, caplog): + conn = connect(Session(), stale_query_probe_seconds=value) + try: + assert conn._Connection__probe_after is None + warnings = [r.getMessage() for r in caplog.records if "disabled" in r.message] + assert warnings == [ + f"Invalid stale_query_probe_seconds {expected}; " + "stale-query probes disabled." + ] + finally: + conn.close() + + +def test_invalid_probe_interval_warning_truncates_a_huge_value(caplog): + conn = connect(Session(), stale_query_probe_seconds=10**400) + try: + assert conn._Connection__probe_after is None + (warning,) = [r.getMessage() for r in caplog.records if "disabled" in r.message] + quoted = re.search(r"\((.*)\);", warning).group(1) + assert quoted.startswith("1000") and quoted.endswith("...") + assert len(quoted) <= 80 + finally: + conn.close() + + +def test_connect_direct_forwards_probe_interval(): + session = Session() + target = "wherobots.db.driver.websockets.sync.client.connect" + with patch(target, return_value=session): + conn = connect_direct("wss://compute/sql", stale_query_probe_seconds=PROBE) + try: + _, execution_id = execute(conn, session) + assert wait_until(lambda: session.probe_times(execution_id)) + finally: + conn.close() diff --git a/wherobots/db/_transport.py b/wherobots/db/_transport.py index 8f7989c..e9017c5 100644 --- a/wherobots/db/_transport.py +++ b/wherobots/db/_transport.py @@ -1,6 +1,8 @@ """Transport termination independent of the WebSocket protocol's send lock.""" import errno +import select +import selectors import socket from websockets.sync.client import ClientConnection @@ -22,3 +24,34 @@ def abort_connection(ws: ClientConnection) -> None: finally: # Even if shutdown fails, prevent any later send on this socket object. ws.socket.close() + + +def socket_writable(ws: ClientConnection) -> bool: + """Whether a small write to the transport's socket won't block right now. + + For a send that must not stall the reader thread: a full send buffer + (a peer that keeps sending but stops reading) would park sendall(), which + also holds the library's protocol mutex, so even keepalive pings stall. + + Never raises. Returns ``True`` (proceed, as without this check) on any + failure or for a transport without a real socket, such as a test fake: + ``False`` would silently disable the caller's sends. Uses ``poll`` (or a + selector), not ``select.select``, which fails for fds above FD_SETSIZE. + """ + try: + sock = getattr(ws, "socket", None) + if not isinstance(sock, socket.socket): + return True + fd = sock.fileno() + if fd < 0: + return True # Closed: let the send itself report it. + if hasattr(select, "poll"): + poller = select.poll() + poller.register(fd, select.POLLOUT) + # POLLERR/POLLHUP count as writable: the send then fails fast. + return bool(poller.poll(0)) + with selectors.DefaultSelector() as selector: + selector.register(fd, selectors.EVENT_WRITE) + return bool(selector.select(0)) + except Exception: + return True diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index 5e28b7b..5d53298 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -3,8 +3,9 @@ import math import textwrap import threading +import time import uuid -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any, Callable, Dict from wherobots.db.redaction import get_statement_type, redact_sql @@ -15,8 +16,8 @@ import websockets.exceptions import websockets.sync.client -from .constants import DEFAULT_READ_TIMEOUT_SECONDS -from ._transport import abort_connection +from .constants import DEFAULT_READ_TIMEOUT_SECONDS, DEFAULT_STALE_QUERY_PROBE_SECONDS +from ._transport import abort_connection, socket_writable from .cursor import Cursor from .errors import NotSupportedError, OperationalError from .models import ExecutionResult, ProgressInfo, Store, StoreResult @@ -69,6 +70,116 @@ def _handshake_bound(ws: Any) -> float | None: return None +_STALE_CHECK_MAX_INTERVAL_SECONDS = 1.0 +"""Upper bound on how often the reader thread scans for stale queries.""" + +_MAX_PROBE_BACKOFF_DOUBLINGS = 3 +"""Probe spacing doubles from N up to 8N (2**3 * N), then stays there.""" + +_RESULTS_PROBE_AFTER_FACTOR = 4 +"""A query whose results were already requested gets one probe, after 4N. + +Its terminal event has arrived; only the retrieve_results reply is missing. +The session re-serializes and re-sends the whole result (up to 100 MiB) for +every retrieve, synchronously on its event loop, so duplicates are costly. A +lost reply is also the rarer case: retrieve replies are sent inline on the +session's main loop, not pushed from a callback like state_updated. +""" + +_EXECUTION_NOT_FOUND = "Execution not found" +"""sql-session's error for a retrieve_results it has no cache entry for. + +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". +""" + +_PROBED_STATES = frozenset( + { + ExecutionState.EXECUTION_REQUESTED, + ExecutionState.PENDING, + ExecutionState.RUNNING, + ExecutionState.RESULTS_REQUESTED, + } +) +"""States in which the driver is waiting on a session event it can re-ask for.""" + + +def _probe_interval(value: Any) -> float | None: + """The validated stale-query probe interval, or ``None`` to disable it. + + Like ``_handshake_bound``, returns ``None`` for a non-numeric (including + ``bool``), non-finite or non-positive value. Never raises ``Exception``: + ``connect()`` callers pass this through, and a bad value must not cost + them their connection. + """ + try: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + value = float(value) + if not math.isfinite(value) or value <= 0: + return None + return value + except Exception: + return None + + +_INVALID_VALUE_REPR_CHARS = 80 +"""Longest repr of an invalid setting that a warning quotes.""" + + +def _describe_invalid(value: Any) -> str: + """A bounded ``repr`` of ``value`` for a warning. Never raises ``Exception``. + + Truncated so a huge value (``10**400``) can't flood the log; falls back to + the type name if ``repr`` fails, since this runs in the constructor. + """ + try: + text = repr(value) + except Exception: + try: + return f"" + except Exception: + return "" + if len(text) > _INVALID_VALUE_REPR_CHARS: + text = text[: _INVALID_VALUE_REPR_CHARS - 3] + "..." + return text + + +@dataclass +class _Watch: + """Stale-query watchdog bookkeeping for one query. + + Only the reader thread touches it once the query is registered; ``__send`` + stamps ``last_activity`` and ``last_event`` under ``__lock`` just before + registration. + """ + + last_activity: float = field(default_factory=time.monotonic) + """``time.monotonic()`` of the last inbound event (or probe) for the query. + + Spaces the probes; see ``last_event`` for the true silence.""" + last_event: float = field(default_factory=time.monotonic) + """``time.monotonic()`` of the last inbound event (or registration). + + Unlike ``last_activity``, probes don't update it, so the probe log reports + how long the query has really been silent.""" + probes_sent: int = 0 + """Probes sent since the last event that wasn't a reply to a retrieve.""" + probe_outstanding: bool = False + """A probe was sent and no reply to a retrieve has arrived since. + + Other events, a non-terminal ``state_updated`` included, leave it set: they + don't answer the probe, whose reply may still be on its way.""" + probeable: bool = True + """False once the session has said it can't answer probes for the query.""" + results_requested: bool = False + """The normal (non-probe) retrieve_results was sent.""" + results_probed: bool = False + """The one probe allowed after ``results_requested`` was sent.""" + + @dataclass class Query: sql: str @@ -76,6 +187,7 @@ class Query: state: ExecutionState handler: Callable[[Any], None] store: Store | None = None + watch: _Watch = field(default_factory=_Watch) class Connection: @@ -101,13 +213,29 @@ def __init__( data_compression: DataCompression | None = None, geometry_representation: GeometryRepresentation | None = None, session_id: str | None = None, + stale_query_probe_seconds: float | None = DEFAULT_STALE_QUERY_PROBE_SECONDS, ): self.__ws = ws - self.__read_timeout = read_timeout 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) + 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.", + _describe_invalid(stale_query_probe_seconds), + ) + self.__next_stale_check = 0.0 + # recv() must wake up for the watchdog's checks, also when the caller + # asked for no read timeout (websockets accepts None) or a long one. + self.__recv_timeout = read_timeout + if self.__probe_after is not None: + check = min(_STALE_CHECK_MAX_INTERVAL_SECONDS, self.__probe_after) + if read_timeout is None or ( + isinstance(read_timeout, (int, float)) and read_timeout > check + ): + self.__recv_timeout = check self.__session_id = session_id self.__lock = threading.Lock() @@ -200,6 +328,10 @@ def __receive_loop(self) -> None: while True: if self.__shutdown_done.is_set(): return + # Checked on every iteration, not only when recv() times out: on a + # busy connection, steady traffic for one query would otherwise + # starve another query's watchdog. + self.__probe_stale_queries() try: self.__listen() except TimeoutError: @@ -214,6 +346,73 @@ def __receive_loop(self) -> None: except Exception as e: logging.exception("Error handling message from SQL session", exc_info=e) + def __probe_stale_queries(self) -> None: + """Re-ask the session about queries that have gone quiet. + + Defence in depth for a lost terminal event: the driver otherwise trusts + a single pushed ``state_updated`` and would wait forever if it never + arrives. A probe is a ``retrieve_results`` request; the session answers + with the query's current state (and its results or result location + once it has succeeded), or with an error for an unknown execution. + + Runs on the reader thread, which owns every ``Query``'s watchdog fields. + Never raises: a watchdog failure must not disturb the receive loop. + """ + probe_after = self.__probe_after + if probe_after is None: + return + try: + now = time.monotonic() + if now < self.__next_stale_check: + return + self.__next_stale_check = now + min( + _STALE_CHECK_MAX_INTERVAL_SECONDS, probe_after + ) + with self.__lock: + if self.__closed: + return + candidates = list(self.__queries.values()) + for query in candidates: + watch = query.watch + if not watch.probeable or query.state not in _PROBED_STATES: + continue + if watch.results_requested: + if watch.results_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): + # 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; + # retry at the next check. + break + # Only a reply to a retrieve clears the flag; other events, + # state_updated included, don't. Still set means the last + # probe went unanswered. + unanswered = watch.probe_outstanding + watch.probes_sent += 1 + watch.probe_outstanding = True + watch.results_probed = watch.results_requested + # Space the next probe from this one, not from the last event. + watch.last_activity = now + logging.info( + "No events for query %s in %.1fs while %s; probed the SQL " + "session for its state (probe %d%s).", + query.execution_id, + now - watch.last_event, + query.state, + watch.probes_sent, + ", previous probe unanswered" if unanswered else "", + ) + except Exception: + logging.exception("Stale-query probe failed") + def __connection_error(self, execution_id: str) -> OperationalError: message = ( f"SQL connection lost (session={self.__session_id or 'unknown'}, " @@ -299,6 +498,27 @@ def __listen(self) -> None: # Invalid event. return + query = self.__queries.get(execution_id) + # Whether this event answers a probe (possibly along with the normal + # retrieve). Read before the bookkeeping below resets it. + probe_reply = False + if query is not None: + watch = query.watch + # Any inbound event, progress included, shows the query is alive. + watch.last_activity = watch.last_event = time.monotonic() + if kind == EventKind.EXECUTION_RESULT or kind == EventKind.ERROR: + # A reply to a retrieve_results, which may be our own probe: it + # resets the timer but not the backoff, so a long-running + # silent query isn't probed every N seconds forever. + probe_reply = watch.probe_outstanding + watch.probe_outstanding = False + else: + # A genuine event resets the backoff, but leaves an outstanding + # probe outstanding: a non-terminal state_updated can arrive + # between a probe and its reply, which must still be + # recognised as the probe's (see the not-found guard below). + watch.probes_sent = 0 + # Progress events are independent of the query state machine and don't # require a tracked query — the handler is connection-level. if kind == EventKind.EXECUTION_PROGRESS: @@ -318,7 +538,6 @@ def __listen(self) -> None: logging.exception("Progress handler raised an exception") return - query = self.__queries.get(execution_id) if not query: logging.warning( "Received %s event for unknown execution ID %s", kind, execution_id @@ -332,6 +551,18 @@ def complete_query(result: ExecutionResult) -> None: with self.__lock: claimed = self.__queries.pop(execution_id, None) if claimed is not None: + if probe_reply and result.error is None: + if claimed.watch.results_requested: + # The normal retrieve's reply may have won the race. + logging.info( + "Query %s results arrived after a stale-query probe.", + execution_id, + ) + else: + logging.info( + "Query %s recovered from a stale-query probe reply.", + execution_id, + ) claimed.handler(result) def fail_query(action: str, error: Exception) -> None: @@ -348,6 +579,18 @@ def fail_query(action: str, error: Exception) -> None: logging.error("%s", message) complete_query(ExecutionResult(error=OperationalError(message))) + def complete_with_store_result(result_uri: str) -> None: + # Results are stored in cloud storage + store_result = StoreResult(result_uri=result_uri, size=message.get("size")) + logging.info( + "Query %s results stored at: %s (size: %s)", + execution_id, + result_uri, + store_result.size, + ) + query.state = ExecutionState.COMPLETED + complete_query(ExecutionResult(store_result=store_result)) + # Incoming state transitions are handled here. if kind == EventKind.STATE_UPDATED or kind == EventKind.EXECUTION_RESULT: try: @@ -363,19 +606,7 @@ def fail_query(action: str, error: Exception) -> None: if kind == EventKind.STATE_UPDATED: result_uri = message.get("result_uri") if result_uri: - # Results are stored in cloud storage - store_result = StoreResult( - result_uri=result_uri, - size=message.get("size"), - ) - logging.info( - "Query %s results stored at: %s (size: %s)", - execution_id, - result_uri, - store_result.size, - ) - query.state = ExecutionState.COMPLETED - complete_query(ExecutionResult(store_result=store_result)) + complete_with_store_result(result_uri) return if query.store is not None: @@ -397,8 +628,34 @@ def fail_query(action: str, error: Exception) -> None: fail_query("request", error) return - # Otherwise, process the results from the execution_result event. + # Otherwise, this execution_result answers a retrieve_results. + # For a store-configured query, that can only be a stale-query + # probe: the normal path never requests results for one. + # Sessions with the probe contract echo the store location. + result_uri = message.get("result_uri") + if result_uri: + complete_with_store_result(result_uri) + return + results = message.get("results") + 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 + # (no store_path). The two are indistinguishable, and + # completing empty would silently drop an old session's + # real result. Keep waiting for the genuine state_updated, + # exactly as without the watchdog, and don't probe again: + # the answer won't change. Known limitation: an empty store + # result whose terminal event was lost can't be recovered. + query.watch.probeable = False + logging.warning( + "Query %s succeeded, but the SQL session's reply has no " + "result location (older SQL session, or an empty store " + "result); cannot recover it, still waiting.", + execution_id, + ) + return + if results is None or results == {}: logging.warning("Got no results back from %s.", execution_id) query.state = ExecutionState.COMPLETED @@ -428,8 +685,28 @@ def fail_query(action: str, error: Exception) -> None: # details. pass elif kind == EventKind.ERROR: - query.state = ExecutionState.FAILED error = message.get("message") + if ( + probe_reply + and error == _EXECUTION_NOT_FOUND + and not query.watch.results_requested + and query.state != ExecutionState.FAILED + ): + # 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 + # running and finished-with-a-lost-event look the same from + # here), so keep waiting either way. After the normal + # retrieve, though, there's nothing left to wait for. + query.watch.probeable = False + logging.warning( + "SQL session no longer tracks execution %s (likely evicted " + "from its result cache); cannot probe it, waiting for its " + "terminal event.", + execution_id, + ) + return + query.state = ExecutionState.FAILED complete_query(ExecutionResult(error=OperationalError(error))) else: logging.warning("Received unknown %s event!", kind) @@ -456,7 +733,17 @@ def _handle_results(self, execution_id: str, results: Dict[str, Any]) -> Any: else: raise NotSupportedError("Unsupported results format") - def __send(self, message: Dict[str, Any], query: Query | None = None) -> None: + def __send( + self, + message: Dict[str, Any], + query: Query | None = None, + blocking: bool = True, + ) -> bool: + """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. + """ # Serialization and redaction are local work. Fail before registration, # without poisoning unrelated cursors or misreporting a transport loss. request = json.dumps(message) @@ -465,16 +752,30 @@ def __send(self, message: Dict[str, Any], query: Query | None = None) -> None: # unguarded call would redact on every request even with DEBUG off. if logging.getLogger().isEnabledFor(logging.DEBUG): logging.debug("Request: %s", self.__redacted_request(message)) - with self.__send_lock: + if not self.__send_lock.acquire(blocking=blocking): + return False + try: with self.__lock: if self.__closed: if query is not None: raise self.__connection_error(query.execution_id) - return + return True if query is not None: + # 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 + return True + # 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() + # itself would block, stopping recv() for every query, and would + # 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): + return False try: self.__ws.send(request) except (websockets.exceptions.ConnectionClosed, OSError): @@ -486,8 +787,11 @@ def __send(self, message: Dict[str, Any], query: Query | None = None) -> None: self.__queries.pop(query.execution_id, None) raise else: - return + return True + finally: + self.__send_lock.release() self.__fail_pending() + return True @staticmethod def __redacted_request(message: Dict[str, Any]) -> str: @@ -504,7 +808,7 @@ def __redacted_request(message: Dict[str, Any]) -> str: def __recv(self) -> Dict[str, Any]: try: - frame = self.__ws.recv(timeout=self.__read_timeout) + frame = self.__ws.recv(timeout=self.__recv_timeout) except TimeoutError: raise # Idle polls are expected, not terminal I/O failures. except OSError as e: @@ -559,10 +863,15 @@ def __execute_sql( ) return execution_id - def __request_results(self, execution_id: str) -> None: + 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 + return True request = { "kind": RequestKind.RETRIEVE_RESULTS.value, @@ -575,9 +884,12 @@ def __request_results(self, execution_id: str) -> None: if self.__geometry_representation: request["geometry"] = self.__geometry_representation.value - query.state = ExecutionState.RESULTS_REQUESTED - logging.info("Requesting results from %s ...", execution_id) - self.__send(request) + if not probe: + # A probe only asks; the reply carries the query's actual state. + query.state = ExecutionState.RESULTS_REQUESTED + query.watch.results_requested = True + logging.info("Requesting results from %s ...", execution_id) + return self.__send(request, blocking=not probe) def __cancel_query(self, execution_id: str) -> None: """Cancels the query with the given execution ID.""" diff --git a/wherobots/db/constants.py b/wherobots/db/constants.py index 0b409cd..965bd59 100644 --- a/wherobots/db/constants.py +++ b/wherobots/db/constants.py @@ -17,6 +17,8 @@ # Bound on the WebSocket close handshake performed by a graceful close(). DEFAULT_CLOSE_TIMEOUT_SECONDS: float = 1.0 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 MAX_MESSAGE_SIZE: int = 100 * 2**20 # 100MiB PROTOCOL_VERSION: Version = Version("1.0.0") diff --git a/wherobots/db/driver.py b/wherobots/db/driver.py index 932cff6..399bced 100644 --- a/wherobots/db/driver.py +++ b/wherobots/db/driver.py @@ -25,6 +25,7 @@ DEFAULT_READ_TIMEOUT_SECONDS, DEFAULT_SESSION_TYPE, DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS, + DEFAULT_STALE_QUERY_PROBE_SECONDS, MAX_MESSAGE_SIZE, PARAM_STYLE, PROTOCOL_VERSION, @@ -125,6 +126,7 @@ def connect( geometry_representation: Union[GeometryRepresentation, None] = None, 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, ) -> Connection: """Create a connection to a Wherobots SQL session. @@ -146,6 +148,12 @@ def connect( own ``client=dbapi;ver=`` hop to the right of it. This header is advisory only and never affects authentication; ``extra_headers`` cannot be used to override the ``Authorization``/``X-API-Key`` headers. + :param stale_query_probe_seconds: How long a running query may go without + any event from the SQL session before the driver re-asks the session + for its state, in case its completion event was lost. Later probes + 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. """ if not token and not api_key: raise ValueError("At least one of `token` or `api_key` is required") @@ -271,6 +279,7 @@ def get_session_uri() -> str: session_id=urllib.parse.urlparse(session_id_url) .path.rstrip("/") .rsplit("/", 1)[-1], + stale_query_probe_seconds=stale_query_probe_seconds, ) @@ -299,7 +308,12 @@ def connect_direct( geometry_representation: Union[GeometryRepresentation, None] = None, cancel_event: Union[threading.Event, None] = None, session_id: str | None = None, + stale_query_probe_seconds: Union[float, None] = DEFAULT_STALE_QUERY_PROBE_SECONDS, ) -> Connection: + """Connect to a SQL session's WebSocket endpoint directly. + + :param stale_query_probe_seconds: See :func:`connect`. + """ uri_with_protocol = f"{uri}/{protocol}" ssl_context = ssl.create_default_context() ssl_context.load_verify_locations(certifi.where()) @@ -344,4 +358,5 @@ def ws_connect() -> websockets.sync.client.ClientConnection: data_compression=data_compression, geometry_representation=geometry_representation, session_id=session_id, + stale_query_probe_seconds=stale_query_probe_seconds, ) diff --git a/wherobots/db/types.py b/wherobots/db/types.py index 4c0745e..a1db45f 100644 --- a/wherobots/db/types.py +++ b/wherobots/db/types.py @@ -9,6 +9,9 @@ class ExecutionState(LowercaseStrEnum): EXECUTION_REQUESTED = auto() "Execution of a query has been requested by the driver." + PENDING = auto() + "The SQL session has accepted the query but not started running it yet." + RUNNING = auto() "The SQL session has reported the query is running."