From a2d27a85c041dad7cc6a2a93c50603a6eb226ed7 Mon Sep 17 00:00:00 2001 From: Ryan Avery Date: Fri, 4 Sep 2026 22:01:07 -0700 Subject: [PATCH 1/6] fix(dbapi): fail pending queries on SQL connection loss --- tests/test_connection_query_tracking.py | 7 +- tests/test_disconnect.py | 247 ++++++++++++++++++++++++ tests/test_empty_store_results.py | 9 +- wherobots/db/connection.py | 102 ++++++++-- wherobots/db/driver.py | 22 +++ 5 files changed, 365 insertions(+), 22 deletions(-) create mode 100644 tests/test_disconnect.py diff --git a/tests/test_connection_query_tracking.py b/tests/test_connection_query_tracking.py index e69089c..1c7108a 100644 --- a/tests/test_connection_query_tracking.py +++ b/tests/test_connection_query_tracking.py @@ -9,7 +9,7 @@ import json import queue -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import cbor2 import pyarrow @@ -22,9 +22,8 @@ def _make_connection(): """Create a Connection with a mocked WebSocket.""" mock_ws = MagicMock() - # Prevent the background thread from running the main loop - mock_ws.protocol.state = 4 # CLOSED state, so __main_loop exits immediately - return Connection(mock_ws) + with patch("wherobots.db.connection.threading.Thread.start"): + return Connection(mock_ws) def _track_query(conn, execution_id="exec-1", state=ExecutionState.RUNNING, store=None): diff --git a/tests/test_disconnect.py b/tests/test_disconnect.py new file mode 100644 index 0000000..7fdb282 --- /dev/null +++ b/tests/test_disconnect.py @@ -0,0 +1,247 @@ +"""Connection loss must complete each pending cursor exactly once.""" +import json +import queue +import threading +import time +from unittest.mock import MagicMock, patch + +import pandas +import pytest +from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK +from websockets.protocol import State + +from wherobots.db.connection import Connection +from wherobots.db.driver import connect_direct +from wherobots.db.errors import OperationalError + + +class Transport: + def __init__(self): + self.protocol = MagicMock(state=State.OPEN) + self.incoming = queue.Queue() + self.sent = [] + + def recv(self, timeout): + value = self.incoming.get(timeout=3) + if isinstance(value, Exception): + self.protocol.state = State.CLOSED + raise value + return json.dumps(value) + + def send(self, value): + self.sent.append(json.loads(value)) + + def close(self): + self.incoming.put(ConnectionClosedOK(None, None)) + + +@pytest.mark.parametrize( + "error", + [ + ConnectionClosedError(None, None), + ConnectionClosedOK(None, None), + OSError("transport lost"), + ], +) +def test_disconnect_unblocks_all_cursors_and_rejects_new_queries(error): + ws = Transport() + conn = Connection(ws, session_id="session-1") + cursors = [conn.cursor() for _ in range(3)] + for cursor in cursors: + cursor.execute("MERGE INTO secret VALUES ('private')") + ws.incoming.put(error) + conn._Connection__thread.join(timeout=3) + assert not conn._Connection__thread.is_alive() + for cursor, request in zip(cursors, ws.sent): + with pytest.raises(OperationalError) as exc: + cursor.fetchall() + assert "session-1" in str(exc.value) + assert request["execution_id"] in str(exc.value) + assert "Commit outcome is unknown" in str(exc.value) + assert "private" not in str(exc.value) + with pytest.raises(OperationalError): + cursor.fetchall() + assert not conn._Connection__queries + with pytest.raises(OperationalError): + conn.cursor().execute("INSERT INTO t VALUES (1)") + assert len(ws.sent) == 3 + + +def test_delivered_result_wins_close_and_is_not_overwritten(): + ws = Transport() + conn = Connection(ws) + cursor = conn.cursor() + cursor.execute("SELECT 1") + with patch.object( + conn, "_handle_results", return_value=pandas.DataFrame({"x": [1]}) + ): + ws.incoming.put( + { + "kind": "execution_result", + "execution_id": ws.sent[0]["execution_id"], + "state": "succeeded", + "results": {"ignored": True}, + } + ) + ws.incoming.put(ConnectionClosedOK(None, None)) + conn._Connection__thread.join(timeout=3) + assert cursor.fetchall()["x"].tolist() == [1] + assert cursor._Cursor__queue.empty() + + +def test_close_fails_pending_without_waiting_for_status(): + ws = Transport() + details = MagicMock() + conn = Connection(ws, failure_details=details) + cursor = conn.cursor() + cursor.execute("SELECT 1") + conn.close() + with pytest.raises(OperationalError): + cursor.fetchall() + details.assert_not_called() + conn._Connection__thread.join(timeout=3) + + +def test_stalled_enrichment_is_bounded_once_for_all_cursors(): + release = threading.Event() + ws = Transport() + + def lookup(): + release.wait(timeout=10) + return "late" + + conn = Connection(ws, failure_details=lookup) + cursors = [conn.cursor() for _ in range(3)] + for cursor in cursors: + cursor.execute("SELECT 1") + started = time.monotonic() + ws.incoming.put(ConnectionClosedError(None, None)) + conn._Connection__thread.join(timeout=3) + try: + assert not conn._Connection__thread.is_alive() + assert time.monotonic() - started < 3 + for cursor in cursors: + with pytest.raises(OperationalError, match="Commit outcome is unknown"): + cursor.fetchall() + finally: + release.set() + + +@pytest.mark.parametrize( + "status,payload", + [ + (200, {"firstFailure": {"message": "Evicted: ephemeral-storage"}}), + (404, {}), + (503, {}), + (200, {}), + (200, None), + ], +) +def test_http_enrichment_best_effort(status, payload): + ws = Transport() + response = MagicMock(status_code=status) + response.json.return_value = payload + response.__enter__.return_value = response + with patch( + "wherobots.db.driver.websockets.sync.client.connect", return_value=ws + ), patch("wherobots.db.driver.requests.get", return_value=response) as get: + conn = connect_direct( + "wss://compute/sql", + headers={"Authorization": "Bearer test"}, + session_status_url="https://api/sql/session/session-1", + ) + cursor = conn.cursor() + cursor.execute("SELECT 1") + ws.incoming.put(ConnectionClosedError(None, None)) + conn._Connection__thread.join(timeout=3) + assert not conn._Connection__thread.is_alive() + with pytest.raises(OperationalError) as exc: + cursor.fetchall() + assert ("ephemeral-storage" in str(exc.value)) == ( + status == 200 and bool(payload) + ) + assert get.call_args.kwargs["timeout"] == 1.0 + assert get.call_args.kwargs["allow_redirects"] is False + + +def test_send_failure_does_not_leave_pending_query(): + ws = Transport() + conn = Connection(ws) + ws.send = MagicMock(side_effect=ConnectionClosedError(None, None)) + cursor = conn.cursor() + cursor.execute("INSERT INTO t VALUES (1)") + with pytest.raises(OperationalError): + cursor.fetchall() + assert not conn._Connection__queries + conn.close() + conn._Connection__thread.join(timeout=3) + + +def test_buffered_result_is_drained_even_when_transport_is_already_closed(): + ws = Transport() + with patch("wherobots.db.connection.threading.Thread.start"): + conn = Connection(ws) + cursor = conn.cursor() + cursor.execute("SELECT 1") + ws.incoming.put( + { + "kind": "execution_result", + "execution_id": ws.sent[0]["execution_id"], + "state": "succeeded", + "results": {"ignored": True}, + } + ) + ws.incoming.put(ConnectionClosedOK(None, None)) + ws.protocol.state = State.CLOSED + with patch.object( + conn, "_handle_results", return_value=pandas.DataFrame({"x": [1]}) + ): + conn._Connection__main_loop() + assert cursor.fetchall()["x"].tolist() == [1] + + +@pytest.mark.parametrize( + "error", [OSError("HTTP unavailable"), ValueError("invalid JSON")] +) +def test_enrichment_errors_preserve_connection_failure(error): + ws = Transport() + conn = Connection(ws, failure_details=MagicMock(side_effect=error)) + cursor = conn.cursor() + cursor.execute("SELECT 1") + ws.incoming.put(ConnectionClosedError(None, None)) + conn._Connection__thread.join(timeout=3) + assert not conn._Connection__thread.is_alive() + with pytest.raises(OperationalError, match="Commit outcome is unknown"): + cursor.fetchall() + + +def test_close_racing_result_decode_delivers_only_one_terminal_outcome(): + decoding = threading.Event() + release = threading.Event() + ws = Transport() + conn = Connection(ws) + cursor = conn.cursor() + cursor.execute("SELECT 1") + + def decode(*args): + decoding.set() + assert release.wait(timeout=3) + return pandas.DataFrame({"x": [1]}) + + with patch.object(conn, "_handle_results", side_effect=decode): + ws.incoming.put( + { + "kind": "execution_result", + "execution_id": ws.sent[0]["execution_id"], + "state": "succeeded", + "results": {"ignored": True}, + } + ) + assert decoding.wait(timeout=3) + conn.close() + release.set() + conn._Connection__thread.join(timeout=3) + assert not conn._Connection__thread.is_alive() + with pytest.raises(OperationalError): + cursor.fetchall() + assert cursor._Cursor__queue.empty() diff --git a/tests/test_empty_store_results.py b/tests/test_empty_store_results.py index 5af05a0..daa02f1 100644 --- a/tests/test_empty_store_results.py +++ b/tests/test_empty_store_results.py @@ -22,9 +22,8 @@ class TestEmptyStoreResults: def _make_connection_and_cursor(self): """Create a Connection with a mocked WebSocket and return (connection, cursor).""" mock_ws = MagicMock() - # Prevent the background thread from running the main loop - mock_ws.protocol.state = 4 # CLOSED state, so __main_loop exits immediately - conn = Connection(mock_ws) + with patch("wherobots.db.connection.threading.Thread.start"): + conn = Connection(mock_ws) cursor = conn.cursor() return conn, cursor @@ -150,8 +149,8 @@ class TestDefensiveNullResults: def _make_connection_and_cursor(self): mock_ws = MagicMock() - mock_ws.protocol.state = 4 - conn = Connection(mock_ws) + with patch("wherobots.db.connection.threading.Thread.start"): + conn = Connection(mock_ws) cursor = conn.cursor() return conn, cursor diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index 7be2f88..cdb30f7 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -1,5 +1,6 @@ import json import logging +import queue import textwrap import threading import uuid @@ -12,7 +13,6 @@ import pyarrow import cbor2 import websockets.exceptions -import websockets.protocol import websockets.sync.client from .constants import DEFAULT_READ_TIMEOUT_SECONDS @@ -64,6 +64,8 @@ def __init__( results_format: ResultsFormat | None = None, data_compression: DataCompression | None = None, geometry_representation: GeometryRepresentation | None = None, + session_id: str | None = None, + failure_details: Callable[[], str | None] | None = None, ): self.__ws = ws self.__read_timeout = read_timeout @@ -72,6 +74,10 @@ def __init__( self.__geometry_representation = geometry_representation self.__progress_handler: ProgressHandler | None = None + self.__session_id = session_id + self.__failure_details = failure_details + self.__lock = threading.Lock() + self.__closed = False self.__queries: dict[str, Query] = {} self.__thread = threading.Thread( target=self.__main_loop, daemon=True, name="wherobots-connection" @@ -85,6 +91,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): self.close() def close(self) -> None: + self.__fail_pending(enrich=False) self.__ws.close() def commit(self) -> None: @@ -114,17 +121,78 @@ def set_progress_handler(self, handler: ProgressHandler | None) -> None: def __main_loop(self) -> None: """Main background loop listening for messages from the SQL session.""" logging.info("Starting background connection handling loop...") - while self.__ws.protocol.state < websockets.protocol.State.CLOSING: + try: + self.__receive_loop() + finally: + self.__fail_pending() + + def __receive_loop(self) -> None: + # recv drains buffered results before raising ConnectionClosed. + while True: try: self.__listen() except TimeoutError: # Expected, retry next time continue - except websockets.exceptions.ConnectionClosedOK: + except websockets.exceptions.ConnectionClosed: logging.info("Connection closed; stopping main loop.") return except Exception as e: logging.exception("Error handling message from SQL session", exc_info=e) + return + + def __connection_error( + self, execution_id: str, details: str | None = None + ) -> OperationalError: + message = ( + f"SQL connection lost (session={self.__session_id or 'unknown'}, " + f"execution={execution_id}). Commit outcome is unknown; " + "verify the operation before retrying writes." + ) + if details: + message += f" Session failure: {details}" + return OperationalError(message) + + def __fail_pending(self, enrich: bool = True) -> None: + # Claim terminal delivery atomically with query registration/result delivery. + with self.__lock: + if self.__closed: + return + self.__closed = True + pending = list(self.__queries.values()) + self.__queries.clear() + details = None + if pending and enrich and self.__failure_details is not None: + # requests' socket timeouts don't bound DNS or a trickling response. + # One daemon lookup per connection bounds the callers' total wait too. + result_queue: queue.Queue = queue.Queue(maxsize=1) + + def lookup() -> None: + try: + result_queue.put(self.__failure_details()) + except Exception: + result_queue.put(None) + + try: + threading.Thread( + target=lookup, daemon=True, name="wherobots-failure-details" + ).start() + details = result_queue.get(timeout=2.0) + except (queue.Empty, RuntimeError): + # Enrichment must not prevent failure delivery, even if the + # process cannot start another thread. + pass + for query in pending: + try: + query.handler( + ExecutionResult( + error=self.__connection_error(query.execution_id, details) + ) + ) + except Exception: + logging.exception( + "Could not deliver connection failure to query handler" + ) def __listen(self) -> None: """Waits for the next message from the SQL session and processes it. @@ -168,8 +236,10 @@ def complete_query(result: ExecutionResult) -> None: # Terminal delivery: stop tracking the query first. Keeping it in # __queries would retain its handler — and the results the handler # references — for the connection's lifetime (WBC-922). - self.__queries.pop(execution_id, None) - query.handler(result) + with self.__lock: + claimed = self.__queries.pop(execution_id, None) + if claimed is not None: + claimed.handler(result) # Incoming state transitions are handled here. if kind == EventKind.STATE_UPDATED or kind == EventKind.EXECUTION_RESULT: @@ -318,13 +388,16 @@ def __execute_sql( if store: request["store"] = store.to_dict() - self.__queries[execution_id] = Query( - sql=sql, - execution_id=execution_id, - state=ExecutionState.EXECUTION_REQUESTED, - handler=handler, - store=store, - ) + with self.__lock: + if self.__closed: + raise self.__connection_error(execution_id) + self.__queries[execution_id] = Query( + sql=sql, + execution_id=execution_id, + state=ExecutionState.EXECUTION_REQUESTED, + handler=handler, + store=store, + ) # Redact literal values before logging: this driver is embedded by other # services, so raw SQL here would leak into their log streams (WBC-139). @@ -334,7 +407,10 @@ def __execute_sql( get_statement_type(sql), textwrap.shorten(redact_sql(sql), width=200), ) - self.__send(request) + try: + self.__send(request) + except Exception: + self.__fail_pending() return execution_id def __request_results(self, execution_id: str) -> None: diff --git a/wherobots/db/driver.py b/wherobots/db/driver.py index 2b9f40c..242daf5 100644 --- a/wherobots/db/driver.py +++ b/wherobots/db/driver.py @@ -267,6 +267,7 @@ def get_session_uri() -> str: data_compression=data_compression, geometry_representation=geometry_representation, cancel_event=cancel_event, + session_status_url=session_id_url, ) @@ -294,6 +295,7 @@ def connect_direct( data_compression: Union[DataCompression, None] = None, geometry_representation: Union[GeometryRepresentation, None] = None, cancel_event: Union[threading.Event, None] = None, + session_status_url: str | None = None, ) -> Connection: uri_with_protocol = f"{uri}/{protocol}" ssl_context = ssl.create_default_context() @@ -331,10 +333,30 @@ def ws_connect() -> websockets.sync.client.ClientConnection: except Exception as e: raise InterfaceError("Failed to connect to SQL session!") from e + def failure_details() -> str | None: + if session_status_url is None: + return None + # Never follow a status redirect with the caller's credentials. + with requests.get( + session_status_url, headers=headers, timeout=1.0, allow_redirects=False + ) as response: + if response.status_code != 200: + return None + payload = response.json() + failure = payload.get("firstFailure") if isinstance(payload, dict) else None + if not isinstance(failure, dict): + return None + message = failure.get("message") + return message[:4096] if isinstance(message, str) else None + return Connection( ws, read_timeout=read_timeout, results_format=results_format, data_compression=data_compression, geometry_representation=geometry_representation, + session_id=urllib.parse.urlparse(session_status_url).path.rsplit("/", 1)[-1] + if session_status_url + else None, + failure_details=failure_details if session_status_url else None, ) From 48999f38866c6b768f1570bf5df76a352f01ce12 Mon Sep 17 00:00:00 2001 From: Ryan Avery Date: Mon, 21 Sep 2026 10:11:28 -0700 Subject: [PATCH 2/6] fix(dbapi): synchronize query send with shutdown --- tests/test_disconnect.py | 110 +++++++++++++++++++++++++++++++++++++ wherobots/db/connection.py | 51 +++++++++++------ 2 files changed, 145 insertions(+), 16 deletions(-) diff --git a/tests/test_disconnect.py b/tests/test_disconnect.py index 7fdb282..4f4e753 100644 --- a/tests/test_disconnect.py +++ b/tests/test_disconnect.py @@ -215,6 +215,116 @@ def test_enrichment_errors_preserve_connection_failure(error): cursor.fetchall() +def test_enrichment_error_is_logged_at_debug(caplog): + ws = Transport() + conn = Connection( + ws, failure_details=MagicMock(side_effect=ValueError("invalid JSON")) + ) + cursor = conn.cursor() + cursor.execute("SELECT 1") + with caplog.at_level("DEBUG"): + ws.incoming.put(ConnectionClosedError(None, None)) + conn._Connection__thread.join(timeout=3) + assert "Failure-details lookup failed: invalid JSON" in caplog.text + + +def test_enrichment_thread_start_error_is_logged_at_debug(caplog): + ws = Transport() + conn = Connection(ws, failure_details=MagicMock(return_value="details")) + cursor = conn.cursor() + cursor.execute("SELECT 1") + with caplog.at_level("DEBUG"), patch( + "wherobots.db.connection.threading.Thread.start", + side_effect=RuntimeError("thread unavailable"), + ): + ws.incoming.put(ConnectionClosedError(None, None)) + conn._Connection__thread.join(timeout=3) + assert "Could not start failure-details lookup: thread unavailable" in caplog.text + + +@pytest.mark.parametrize( + "decode_error", [ValueError("malformed payload"), OSError("decoder I/O error")] +) +def test_result_decode_error_does_not_fail_other_queries(decode_error): + decoded = threading.Event() + ws = Transport() + conn = Connection(ws) + bad_cursor = conn.cursor() + good_cursor = conn.cursor() + bad_cursor.execute("SELECT bad") + good_cursor.execute("SELECT good") + + def decode(execution_id, results): + if execution_id == ws.sent[0]["execution_id"]: + decoded.set() + raise decode_error + return pandas.DataFrame({"x": [1]}) + + with patch.object(conn, "_handle_results", side_effect=decode): + for request in ws.sent: + ws.incoming.put( + { + "kind": "execution_result", + "execution_id": request["execution_id"], + "state": "succeeded", + "results": {"ignored": True}, + } + ) + assert decoded.wait(timeout=3) + assert good_cursor.fetchall()["x"].tolist() == [1] + assert conn._Connection__thread.is_alive() + conn.close() + with pytest.raises(OperationalError): + bad_cursor.fetchall() + conn._Connection__thread.join(timeout=3) + + +def test_close_waits_for_registered_query_to_be_sent(): + send_started = threading.Event() + release_send = threading.Event() + close_started = threading.Event() + close_finished = threading.Event() + events = [] + ws = Transport() + + def send(value): + send_started.set() + assert release_send.wait(timeout=3) + ws.sent.append(json.loads(value)) + events.append("send") + + def close(): + events.append("close") + ws.incoming.put(ConnectionClosedOK(None, None)) + + ws.send = send + ws.close = close + conn = Connection(ws) + cursor = conn.cursor() + execute_thread = threading.Thread(target=cursor.execute, args=("SELECT 1",)) + execute_thread.start() + assert send_started.wait(timeout=3) + + def close_connection(): + close_started.set() + conn.close() + close_finished.set() + + close_thread = threading.Thread(target=close_connection) + close_thread.start() + assert close_started.wait(timeout=3) + assert not close_finished.is_set() + release_send.set() + execute_thread.join(timeout=3) + close_thread.join(timeout=3) + conn._Connection__thread.join(timeout=3) + assert not execute_thread.is_alive() + assert not close_thread.is_alive() + assert events == ["send", "close"] + with pytest.raises(OperationalError): + cursor.fetchall() + + def test_close_racing_result_decode_delivers_only_one_terminal_outcome(): decoding = threading.Event() release = threading.Event() diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index cdb30f7..80365eb 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -33,6 +33,10 @@ """A callable invoked with a :class:`ProgressInfo` on every progress event.""" +class _TransportError(Exception): + """An I/O failure raised while receiving from the WebSocket.""" + + @dataclass class Query: sql: str @@ -137,9 +141,11 @@ def __receive_loop(self) -> None: except websockets.exceptions.ConnectionClosed: logging.info("Connection closed; stopping main loop.") return + except _TransportError: + logging.exception("SQL session transport failed; stopping main loop") + return except Exception as e: logging.exception("Error handling message from SQL session", exc_info=e) - return def __connection_error( self, execution_id: str, details: str | None = None @@ -170,7 +176,8 @@ def __fail_pending(self, enrich: bool = True) -> None: def lookup() -> None: try: result_queue.put(self.__failure_details()) - except Exception: + except Exception as e: + logging.debug("Failure-details lookup failed: %s", e) result_queue.put(None) try: @@ -178,10 +185,12 @@ def lookup() -> None: target=lookup, daemon=True, name="wherobots-failure-details" ).start() details = result_queue.get(timeout=2.0) - except (queue.Empty, RuntimeError): + except queue.Empty: # Enrichment must not prevent failure delivery, even if the # process cannot start another thread. pass + except RuntimeError as e: + logging.debug("Could not start failure-details lookup: %s", e) for query in pending: try: query.handler( @@ -359,7 +368,12 @@ def __redacted_request(message: Dict[str, Any]) -> str: return json.dumps(message) def __recv(self) -> Dict[str, Any]: - frame = self.__ws.recv(timeout=self.__read_timeout) + try: + frame = self.__ws.recv(timeout=self.__read_timeout) + except OSError as e: + # Distinguish transport I/O failures from OSErrors raised later by + # protocol parsing or result decoding; only the former are terminal. + raise _TransportError from e if isinstance(frame, str): message = json.loads(frame) elif isinstance(frame, bytes): @@ -388,6 +402,15 @@ def __execute_sql( if store: request["store"] = store.to_dict() + # Redact literal values before logging: this driver is embedded by other + # services, so raw SQL here would leak into their log streams (WBC-139). + logging.info( + "Executing SQL query %s (%s): %s", + execution_id, + get_statement_type(sql), + textwrap.shorten(redact_sql(sql), width=200), + ) + send_failed = False with self.__lock: if self.__closed: raise self.__connection_error(execution_id) @@ -398,18 +421,14 @@ def __execute_sql( handler=handler, store=store, ) - - # Redact literal values before logging: this driver is embedded by other - # services, so raw SQL here would leak into their log streams (WBC-139). - logging.info( - "Executing SQL query %s (%s): %s", - execution_id, - get_statement_type(sql), - textwrap.shorten(redact_sql(sql), width=200), - ) - try: - self.__send(request) - except Exception: + try: + # Keep registration and transmission atomic with respect to + # shutdown: close() must not claim this query before its SQL is + # sent, then report failure while the request is still emitted. + self.__send(request) + except Exception: + send_failed = True + if send_failed: self.__fail_pending() return execution_id From bf9b1e0e62f87234aaeb4d276a19d7c249dd789d Mon Sep 17 00:00:00 2001 From: Ryan Avery Date: Mon, 21 Sep 2026 13:26:21 -0700 Subject: [PATCH 3/6] fix(dbapi): make shutdown independent of stalled sends --- tests/test_disconnect.py | 486 +++++++++++++++++++++---------------- tests/test_driver.py | 15 +- wherobots/db/_transport.py | 24 ++ wherobots/db/connection.py | 139 ++++++----- wherobots/db/driver.py | 27 +-- 5 files changed, 395 insertions(+), 296 deletions(-) create mode 100644 wherobots/db/_transport.py diff --git a/tests/test_disconnect.py b/tests/test_disconnect.py index 4f4e753..84772cc 100644 --- a/tests/test_disconnect.py +++ b/tests/test_disconnect.py @@ -1,6 +1,8 @@ -"""Connection loss must complete each pending cursor exactly once.""" +"""Transport loss must complete pending queries without waiting for a sender.""" +import errno import json import queue +import socket import threading import time from unittest.mock import MagicMock, patch @@ -9,10 +11,15 @@ import pytest from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK from websockets.protocol import State +from websockets.sync.client import ClientConnection +from websockets.client import ClientProtocol +from websockets.uri import parse_uri -from wherobots.db.connection import Connection +from wherobots.db._transport import abort_connection +from wherobots.db.connection import Connection, Query from wherobots.db.driver import connect_direct from wherobots.db.errors import OperationalError +from wherobots.db.types import ExecutionState class Transport: @@ -20,21 +27,43 @@ def __init__(self): self.protocol = MagicMock(state=State.OPEN) self.incoming = queue.Queue() self.sent = [] + self.aborted = threading.Event() + self.socket = MagicMock() + self.socket.shutdown.side_effect = self.shutdown def recv(self, timeout): - value = self.incoming.get(timeout=3) + try: + value = self.incoming.get(timeout=timeout) + except queue.Empty: + raise TimeoutError from None if isinstance(value, Exception): - self.protocol.state = State.CLOSED + if not isinstance(value, TimeoutError): + self.protocol.state = State.CLOSED raise value return json.dumps(value) def send(self, value): + if self.aborted.is_set(): + raise ConnectionClosedError(None, None) self.sent.append(json.loads(value)) - def close(self): + def shutdown(self, how): + assert how == socket.SHUT_RDWR + self.aborted.set() self.incoming.put(ConnectionClosedOK(None, None)) +def deliver(ws, execution_id): + ws.incoming.put( + { + "kind": "execution_result", + "execution_id": execution_id, + "state": "succeeded", + "results": None, + } + ) + + @pytest.mark.parametrize( "error", [ @@ -61,10 +90,29 @@ def test_disconnect_unblocks_all_cursors_and_rejects_new_queries(error): assert "private" not in str(exc.value) with pytest.raises(OperationalError): cursor.fetchall() + assert cursor._Cursor__queue.empty() assert not conn._Connection__queries with pytest.raises(OperationalError): conn.cursor().execute("INSERT INTO t VALUES (1)") assert len(ws.sent) == 3 + assert ws.aborted.is_set() + + +def test_idle_timeouts_then_result_leave_connection_usable(): + ws = Transport() + conn = Connection(ws) + cursor = conn.cursor() + cursor.execute("SELECT 1") + ws.incoming.put(TimeoutError()) + ws.incoming.put(TimeoutError()) + deliver(ws, ws.sent[0]["execution_id"]) + assert cursor._Cursor__queue.get(timeout=3).error is None + assert conn._Connection__thread.is_alive() + assert not conn._Connection__closed + conn.cursor().execute("SELECT 2") + assert len(ws.sent) == 2 + conn.close() + assert not conn._Connection__thread.is_alive() def test_delivered_result_wins_close_and_is_not_overwritten(): @@ -72,109 +120,39 @@ def test_delivered_result_wins_close_and_is_not_overwritten(): conn = Connection(ws) cursor = conn.cursor() cursor.execute("SELECT 1") - with patch.object( - conn, "_handle_results", return_value=pandas.DataFrame({"x": [1]}) - ): - ws.incoming.put( - { - "kind": "execution_result", - "execution_id": ws.sent[0]["execution_id"], - "state": "succeeded", - "results": {"ignored": True}, - } - ) - ws.incoming.put(ConnectionClosedOK(None, None)) - conn._Connection__thread.join(timeout=3) - assert cursor.fetchall()["x"].tolist() == [1] + deliver(ws, ws.sent[0]["execution_id"]) + ws.incoming.put(ConnectionClosedOK(None, None)) + conn._Connection__thread.join(timeout=3) + assert cursor._Cursor__queue.get(timeout=1).error is None assert cursor._Cursor__queue.empty() -def test_close_fails_pending_without_waiting_for_status(): +def test_close_fails_pending_and_joins_reader(): ws = Transport() - details = MagicMock() - conn = Connection(ws, failure_details=details) + conn = Connection(ws) cursor = conn.cursor() cursor.execute("SELECT 1") conn.close() with pytest.raises(OperationalError): cursor.fetchall() - details.assert_not_called() - conn._Connection__thread.join(timeout=3) - - -def test_stalled_enrichment_is_bounded_once_for_all_cursors(): - release = threading.Event() - ws = Transport() - - def lookup(): - release.wait(timeout=10) - return "late" - - conn = Connection(ws, failure_details=lookup) - cursors = [conn.cursor() for _ in range(3)] - for cursor in cursors: - cursor.execute("SELECT 1") - started = time.monotonic() - ws.incoming.put(ConnectionClosedError(None, None)) - conn._Connection__thread.join(timeout=3) - try: - assert not conn._Connection__thread.is_alive() - assert time.monotonic() - started < 3 - for cursor in cursors: - with pytest.raises(OperationalError, match="Commit outcome is unknown"): - cursor.fetchall() - finally: - release.set() + assert not conn._Connection__thread.is_alive() + conn.close() + ws.socket.shutdown.assert_called_once() @pytest.mark.parametrize( - "status,payload", - [ - (200, {"firstFailure": {"message": "Evicted: ephemeral-storage"}}), - (404, {}), - (503, {}), - (200, {}), - (200, None), - ], + "error", [ConnectionClosedError(None, None), OSError("send failed")] ) -def test_http_enrichment_best_effort(status, payload): - ws = Transport() - response = MagicMock(status_code=status) - response.json.return_value = payload - response.__enter__.return_value = response - with patch( - "wherobots.db.driver.websockets.sync.client.connect", return_value=ws - ), patch("wherobots.db.driver.requests.get", return_value=response) as get: - conn = connect_direct( - "wss://compute/sql", - headers={"Authorization": "Bearer test"}, - session_status_url="https://api/sql/session/session-1", - ) - cursor = conn.cursor() - cursor.execute("SELECT 1") - ws.incoming.put(ConnectionClosedError(None, None)) - conn._Connection__thread.join(timeout=3) - assert not conn._Connection__thread.is_alive() - with pytest.raises(OperationalError) as exc: - cursor.fetchall() - assert ("ephemeral-storage" in str(exc.value)) == ( - status == 200 and bool(payload) - ) - assert get.call_args.kwargs["timeout"] == 1.0 - assert get.call_args.kwargs["allow_redirects"] is False - - -def test_send_failure_does_not_leave_pending_query(): +def test_send_failure_does_not_leave_pending_query(error): ws = Transport() conn = Connection(ws) - ws.send = MagicMock(side_effect=ConnectionClosedError(None, None)) + ws.send = MagicMock(side_effect=error) cursor = conn.cursor() cursor.execute("INSERT INTO t VALUES (1)") with pytest.raises(OperationalError): cursor.fetchall() assert not conn._Connection__queries conn.close() - conn._Connection__thread.join(timeout=3) def test_buffered_result_is_drained_even_when_transport_is_already_closed(): @@ -183,146 +161,118 @@ def test_buffered_result_is_drained_even_when_transport_is_already_closed(): conn = Connection(ws) cursor = conn.cursor() cursor.execute("SELECT 1") - ws.incoming.put( - { - "kind": "execution_result", - "execution_id": ws.sent[0]["execution_id"], - "state": "succeeded", - "results": {"ignored": True}, - } - ) + deliver(ws, ws.sent[0]["execution_id"]) ws.incoming.put(ConnectionClosedOK(None, None)) ws.protocol.state = State.CLOSED - with patch.object( - conn, "_handle_results", return_value=pandas.DataFrame({"x": [1]}) - ): - conn._Connection__main_loop() - assert cursor.fetchall()["x"].tolist() == [1] + conn._Connection__main_loop() + assert cursor._Cursor__queue.get(timeout=1).error is None @pytest.mark.parametrize( - "error", [OSError("HTTP unavailable"), ValueError("invalid JSON")] + "decode_error", [ValueError("bad payload"), OSError("decoder error")] ) -def test_enrichment_errors_preserve_connection_failure(error): - ws = Transport() - conn = Connection(ws, failure_details=MagicMock(side_effect=error)) - cursor = conn.cursor() - cursor.execute("SELECT 1") - ws.incoming.put(ConnectionClosedError(None, None)) - conn._Connection__thread.join(timeout=3) - assert not conn._Connection__thread.is_alive() - with pytest.raises(OperationalError, match="Commit outcome is unknown"): - cursor.fetchall() - - -def test_enrichment_error_is_logged_at_debug(caplog): +def test_result_decode_error_does_not_fail_other_queries(decode_error): ws = Transport() - conn = Connection( - ws, failure_details=MagicMock(side_effect=ValueError("invalid JSON")) - ) - cursor = conn.cursor() - cursor.execute("SELECT 1") - with caplog.at_level("DEBUG"): - ws.incoming.put(ConnectionClosedError(None, None)) - conn._Connection__thread.join(timeout=3) - assert "Failure-details lookup failed: invalid JSON" in caplog.text + conn = Connection(ws) + bad = conn.cursor() + good = conn.cursor() + bad.execute("SELECT bad") + good.execute("SELECT good") + with patch.object(conn, "_handle_results", side_effect=decode_error): + ws.incoming.put( + { + "kind": "execution_result", + "execution_id": ws.sent[0]["execution_id"], + "state": "succeeded", + "results": {"ignored": True}, + } + ) + deliver(ws, ws.sent[1]["execution_id"]) + assert good._Cursor__queue.get(timeout=3).error is None + assert not conn._Connection__closed + conn.close() -def test_enrichment_thread_start_error_is_logged_at_debug(caplog): +def test_serialization_error_does_not_register_or_fail_other_queries(): ws = Transport() - conn = Connection(ws, failure_details=MagicMock(return_value="details")) - cursor = conn.cursor() - cursor.execute("SELECT 1") - with caplog.at_level("DEBUG"), patch( - "wherobots.db.connection.threading.Thread.start", - side_effect=RuntimeError("thread unavailable"), - ): - ws.incoming.put(ConnectionClosedError(None, None)) - conn._Connection__thread.join(timeout=3) - assert "Could not start failure-details lookup: thread unavailable" in caplog.text + conn = Connection(ws) + good = conn.cursor() + good.execute("SELECT 1") + store = MagicMock() + store.to_dict.return_value = {"invalid": object()} + with pytest.raises(TypeError): + conn.cursor().execute("SELECT 2", store=store) + assert len(ws.sent) == len(conn._Connection__queries) == 1 + deliver(ws, ws.sent[0]["execution_id"]) + assert good._Cursor__queue.get(timeout=3).error is None + assert not conn._Connection__closed + conn.close() -@pytest.mark.parametrize( - "decode_error", [ValueError("malformed payload"), OSError("decoder I/O error")] -) -def test_result_decode_error_does_not_fail_other_queries(decode_error): - decoded = threading.Event() +def test_nontransport_send_error_is_propagated_and_query_is_untracked(): ws = Transport() conn = Connection(ws) - bad_cursor = conn.cursor() - good_cursor = conn.cursor() - bad_cursor.execute("SELECT bad") - good_cursor.execute("SELECT good") - - def decode(execution_id, results): - if execution_id == ws.sent[0]["execution_id"]: - decoded.set() - raise decode_error - return pandas.DataFrame({"x": [1]}) - - with patch.object(conn, "_handle_results", side_effect=decode): - for request in ws.sent: - ws.incoming.put( - { - "kind": "execution_result", - "execution_id": request["execution_id"], - "state": "succeeded", - "results": {"ignored": True}, - } - ) - assert decoded.wait(timeout=3) - assert good_cursor.fetchall()["x"].tolist() == [1] - assert conn._Connection__thread.is_alive() + good = conn.cursor() + good.execute("SELECT 1") + with patch.object(ws, "send", side_effect=ValueError("API misuse")): + with pytest.raises(ValueError, match="API misuse"): + conn.cursor().execute("SELECT 2") + assert len(conn._Connection__queries) == 1 + assert not conn._Connection__closed conn.close() - with pytest.raises(OperationalError): - bad_cursor.fetchall() - conn._Connection__thread.join(timeout=3) -def test_close_waits_for_registered_query_to_be_sent(): - send_started = threading.Event() - release_send = threading.Event() - close_started = threading.Event() - close_finished = threading.Event() - events = [] +@pytest.mark.parametrize("shutdown", ["close", "reader"]) +def test_stalled_send_does_not_block_result_delivery_or_shutdown(shutdown): ws = Transport() + conn = Connection(ws) + a = conn.cursor() + b = conn.cursor() + a.execute("SELECT 1") + sending = threading.Event() + original_send = ws.send + + def blocked_send(value): + sending.set() + # Only actual transport shutdown releases the writer, not the test. + assert ws.aborted.wait(timeout=3) + original_send(value) + + ws.send = blocked_send + sender = threading.Thread(target=b.execute, args=("INSERT INTO t VALUES (1)",)) + sender.start() + try: + assert sending.wait(timeout=1) + deliver(ws, ws.sent[0]["execution_id"]) + assert a._Cursor__queue.get(timeout=1).error is None + if shutdown == "close": + conn.close() + else: + ws.incoming.put(ConnectionClosedError(None, None)) + assert isinstance(b._Cursor__queue.get(timeout=2).error, OperationalError) + sender.join(timeout=2) + assert not sender.is_alive() + assert len(ws.sent) == 1 + assert b._Cursor__queue.empty() + finally: + conn.close() + sender.join(timeout=3) - def send(value): - send_started.set() - assert release_send.wait(timeout=3) - ws.sent.append(json.loads(value)) - events.append("send") - - def close(): - events.append("close") - ws.incoming.put(ConnectionClosedOK(None, None)) - ws.send = send - ws.close = close +def test_close_from_reader_callback_does_not_join_itself(): + ws = Transport() conn = Connection(ws) - cursor = conn.cursor() - execute_thread = threading.Thread(target=cursor.execute, args=("SELECT 1",)) - execute_thread.start() - assert send_started.wait(timeout=3) + finished = threading.Event() - def close_connection(): - close_started.set() + def progress(_): conn.close() - close_finished.set() - - close_thread = threading.Thread(target=close_connection) - close_thread.start() - assert close_started.wait(timeout=3) - assert not close_finished.is_set() - release_send.set() - execute_thread.join(timeout=3) - close_thread.join(timeout=3) - conn._Connection__thread.join(timeout=3) - assert not execute_thread.is_alive() - assert not close_thread.is_alive() - assert events == ["send", "close"] - with pytest.raises(OperationalError): - cursor.fetchall() + finished.set() + + conn.set_progress_handler(progress) + ws.incoming.put({"kind": "execution_progress", "execution_id": "progress"}) + assert finished.wait(timeout=2) + conn._Connection__thread.join(timeout=2) + assert not conn._Connection__thread.is_alive() def test_close_racing_result_decode_delivers_only_one_terminal_outcome(): @@ -335,7 +285,7 @@ def test_close_racing_result_decode_delivers_only_one_terminal_outcome(): def decode(*args): decoding.set() - assert release.wait(timeout=3) + assert release.wait(timeout=5) return pandas.DataFrame({"x": [1]}) with patch.object(conn, "_handle_results", side_effect=decode): @@ -347,11 +297,127 @@ def decode(*args): "results": {"ignored": True}, } ) - assert decoding.wait(timeout=3) + assert decoding.wait(timeout=2) + started = time.monotonic() conn.close() + assert time.monotonic() - started < 2 + assert conn._Connection__thread.is_alive() + with pytest.raises(OperationalError): + cursor.fetchall() release.set() - conn._Connection__thread.join(timeout=3) - assert not conn._Connection__thread.is_alive() - with pytest.raises(OperationalError): - cursor.fetchall() + conn._Connection__thread.join(timeout=2) + assert cursor._Cursor__queue.empty() + + +def test_concurrent_close_delivers_once(): + ws = Transport() + conn = Connection(ws) + cursor = conn.cursor() + cursor.execute("SELECT 1") + closers = [threading.Thread(target=conn.close) for _ in range(4)] + for closer in closers: + closer.start() + for closer in closers: + closer.join(timeout=2) + assert not closer.is_alive() + assert isinstance(cursor._Cursor__queue.get(timeout=1).error, OperationalError) assert cursor._Cursor__queue.empty() + ws.socket.shutdown.assert_called_once() + + +def test_abort_closes_socket_even_if_shutdown_errors(): + ws = MagicMock() + ws.socket.shutdown.side_effect = OSError(errno.EIO, "shutdown failure") + with pytest.raises(OSError): + abort_connection(ws) + ws.socket.close.assert_called_once() + + +def test_failure_is_not_delivered_before_transport_is_disabled(): + ws = Transport() + conn = Connection(ws) + cursor = conn.cursor() + cursor.execute("SELECT 1") + abort_started = threading.Event() + release_abort = threading.Event() + shutdown = ws.shutdown + + def delayed_abort(how): + abort_started.set() + assert release_abort.wait(timeout=3) + shutdown(how) + + ws.socket.shutdown.side_effect = delayed_abort + closer = threading.Thread(target=conn.close) + closer.start() + try: + assert abort_started.wait(timeout=1) + assert cursor._Cursor__queue.empty() + with pytest.raises(OperationalError): + conn.cursor().execute("SELECT 2") + release_abort.set() + closer.join(timeout=2) + assert not closer.is_alive() + assert ws.aborted.is_set() + assert isinstance(cursor._Cursor__queue.get(timeout=1).error, OperationalError) + finally: + release_abort.set() + closer.join(timeout=3) + + +def test_real_websocket_stalled_send_is_interrupted_by_close(): + # Real library protocol mutex + socket.sendall; the peer never reads. + local, peer = socket.socketpair() + local.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 4096) + protocol = ClientProtocol(parse_uri("ws://localhost"), state=State.OPEN) + ws = ClientConnection(local, protocol) + conn = Connection(ws) + entered = threading.Event() + send_data = ws.send_data + + def observe_send(): + entered.set() + send_data() + + outcomes = queue.Queue() + query = Query( + "SELECT 1", "blocked", ExecutionState.EXECUTION_REQUESTED, outcomes.put + ) + with patch.object(ws, "send_data", side_effect=observe_send): + sender = threading.Thread( + target=conn._Connection__send, + args=( + { + "kind": "execute_sql", + "execution_id": "blocked", + "statement": "x" * (8 * 1024 * 1024), + }, + query, + ), + ) + sender.start() + try: + assert entered.wait(timeout=3) + assert sender.is_alive() + conn.close() + assert isinstance(outcomes.get(timeout=2).error, OperationalError) + sender.join(timeout=3) + assert not sender.is_alive() + assert not conn._Connection__thread.is_alive() + ws.recv_events_thread.join(timeout=2) + assert not ws.recv_events_thread.is_alive() + finally: + peer.close() + conn.close() + sender.join(timeout=3) + + +def test_direct_connection_uses_explicit_session_id(): + ws = Transport() + with patch("wherobots.db.driver.websockets.sync.client.connect", return_value=ws): + conn = connect_direct("wss://compute/sql", session_id="session-1") + cursor = conn.cursor() + cursor.execute("SELECT 1") + conn.close() + with pytest.raises(OperationalError, match="session=session-1"): + cursor.fetchall() diff --git a/tests/test_driver.py b/tests/test_driver.py index 7232f35..eeb0169 100644 --- a/tests/test_driver.py +++ b/tests/test_driver.py @@ -29,7 +29,7 @@ def _run_connect(mock_post, mock_get, **connect_kwargs): return kwargs -def _run_connect_full(mock_post, mock_get, **connect_kwargs): +def _run_connect_full(mock_post, mock_get, session_url=None, **connect_kwargs): """Drive a successful connect(). Returns a tuple of (kwargs passed to requests.post, kwargs passed to the @@ -38,7 +38,7 @@ def _run_connect_full(mock_post, mock_get, **connect_kwargs): """ post_resp = MagicMock() post_resp.status_code = 200 - post_resp.url = "https://api.example.com/sql/session/test-id" + post_resp.url = session_url or "https://api.example.com/sql/session/test-id" post_resp.raise_for_status = MagicMock() mock_post.return_value = post_resp @@ -63,6 +63,17 @@ def _run_connect_full(mock_post, mock_get, **connect_kwargs): class TestConnectRegionRuntime: """region/runtime accept enum|str and are omitted when not provided.""" + @pytest.mark.parametrize("suffix", ["", "/", "/?ignored=value"]) + @patch("wherobots.db.driver.requests.get") + @patch("wherobots.db.driver.requests.post") + def test_session_id_is_derived_from_status_url(self, mock_post, mock_get, suffix): + _, kwargs = _run_connect_full( + mock_post, + mock_get, + session_url="https://api.example.com/sql/session/test-id" + suffix, + ) + assert kwargs["session_id"] == "test-id" + @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/wherobots/db/_transport.py b/wherobots/db/_transport.py new file mode 100644 index 0000000..8f7989c --- /dev/null +++ b/wherobots/db/_transport.py @@ -0,0 +1,24 @@ +"""Transport termination independent of the WebSocket protocol's send lock.""" + +import errno +import socket + +from websockets.sync.client import ClientConnection + + +def abort_connection(ws: ClientConnection) -> None: + """Disable further writes and wake blocked I/O before publishing failures. + + ClientConnection.close() takes the library's protocol mutex, which a + stalled sendall() may hold. Shutdown the owned socket instead; the library's + receive thread will observe EOF/error and finish protocol cleanup itself. + Already transmitted bytes may still execute remotely. + """ + try: + ws.socket.shutdown(socket.SHUT_RDWR) + except OSError as exc: + if exc.errno not in (errno.ENOTCONN, errno.EBADF, errno.ECONNRESET): + raise + finally: + # Even if shutdown fails, prevent any later send on this socket object. + ws.socket.close() diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index 80365eb..e4fb00c 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -1,8 +1,8 @@ import json import logging -import queue import textwrap import threading +import time import uuid from dataclasses import dataclass from typing import Any, Callable, Dict @@ -16,6 +16,7 @@ import websockets.sync.client from .constants import DEFAULT_READ_TIMEOUT_SECONDS +from ._transport import abort_connection from .cursor import Cursor from .errors import NotSupportedError, OperationalError from .models import ExecutionResult, ProgressInfo, Store, StoreResult @@ -69,7 +70,6 @@ def __init__( data_compression: DataCompression | None = None, geometry_representation: GeometryRepresentation | None = None, session_id: str | None = None, - failure_details: Callable[[], str | None] | None = None, ): self.__ws = ws self.__read_timeout = read_timeout @@ -79,8 +79,10 @@ def __init__( self.__progress_handler: ProgressHandler | None = None self.__session_id = session_id - self.__failure_details = failure_details self.__lock = threading.Lock() + self.__send_lock = threading.Lock() + self.__shutdown_done = threading.Event() + self.__shutdown_owner: int | None = None self.__closed = False self.__queries: dict[str, Query] = {} self.__thread = threading.Thread( @@ -95,8 +97,19 @@ def __exit__(self, exc_type, exc_val, exc_tb): self.close() def close(self) -> None: - self.__fail_pending(enrich=False) - self.__ws.close() + """Abort the transport, fail pending work, and wait up to 1s for the reader. + + Closing doesn't imply that server-side writes were rolled back. A + decoder or callback can outlive the bounded reader join. + """ + deadline = time.monotonic() + 1.0 + self.__fail_pending() + # A handler may close its own connection during terminal delivery. + if self.__shutdown_owner == threading.get_ident(): + return + self.__shutdown_done.wait(max(0.0, deadline - time.monotonic())) + if self.__thread is not threading.current_thread(): + self.__thread.join(timeout=max(0.0, deadline - time.monotonic())) def commit(self) -> None: raise NotSupportedError @@ -133,6 +146,8 @@ def __main_loop(self) -> None: def __receive_loop(self) -> None: # recv drains buffered results before raising ConnectionClosed. while True: + if self.__shutdown_done.is_set(): + return try: self.__listen() except TimeoutError: @@ -147,61 +162,45 @@ def __receive_loop(self) -> None: except Exception as e: logging.exception("Error handling message from SQL session", exc_info=e) - def __connection_error( - self, execution_id: str, details: str | None = None - ) -> OperationalError: + def __connection_error(self, execution_id: str) -> OperationalError: message = ( f"SQL connection lost (session={self.__session_id or 'unknown'}, " f"execution={execution_id}). Commit outcome is unknown; " "verify the operation before retrying writes." ) - if details: - message += f" Session failure: {details}" return OperationalError(message) - def __fail_pending(self, enrich: bool = True) -> None: - # Claim terminal delivery atomically with query registration/result delivery. + def __fail_pending(self) -> None: + # Stop admission first. Do not wait for __send_lock: its owner may be + # blocked in network I/O. __closed means closing until shutdown_done. with self.__lock: if self.__closed: return self.__closed = True + self.__shutdown_owner = threading.get_ident() + try: + abort_connection(self.__ws) + except OSError: + # The adapter still closes the socket object in its finally block. + logging.exception("Socket shutdown failed; socket was closed") + with self.__lock: pending = list(self.__queries.values()) self.__queries.clear() - details = None - if pending and enrich and self.__failure_details is not None: - # requests' socket timeouts don't bound DNS or a trickling response. - # One daemon lookup per connection bounds the callers' total wait too. - result_queue: queue.Queue = queue.Queue(maxsize=1) - - def lookup() -> None: + try: + for query in pending: try: - result_queue.put(self.__failure_details()) - except Exception as e: - logging.debug("Failure-details lookup failed: %s", e) - result_queue.put(None) - - try: - threading.Thread( - target=lookup, daemon=True, name="wherobots-failure-details" - ).start() - details = result_queue.get(timeout=2.0) - except queue.Empty: - # Enrichment must not prevent failure delivery, even if the - # process cannot start another thread. - pass - except RuntimeError as e: - logging.debug("Could not start failure-details lookup: %s", e) - for query in pending: - try: - query.handler( - ExecutionResult( - error=self.__connection_error(query.execution_id, details) + query.handler( + ExecutionResult( + error=self.__connection_error(query.execution_id) + ) ) - ) - except Exception: - logging.exception( - "Could not deliver connection failure to query handler" - ) + except Exception: + logging.exception( + "Could not deliver connection failure to query handler" + ) + finally: + self.__shutdown_owner = None + self.__shutdown_done.set() def __listen(self) -> None: """Waits for the next message from the SQL session and processes it. @@ -345,14 +344,38 @@ def _handle_results(self, execution_id: str, results: Dict[str, Any]) -> Any: else: return OperationalError(f"Unsupported results format {result_format}") - def __send(self, message: Dict[str, Any]) -> None: + def __send(self, message: Dict[str, Any], query: Query | None = None) -> None: + # Serialization and redaction are local work. Fail before registration, + # without poisoning unrelated cursors or misreporting a transport loss. request = json.dumps(message) # Only compute the redacted request (json.dumps + sqlparse parse) when # DEBUG is actually enabled; the log argument is evaluated eagerly, so an # 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)) - self.__ws.send(request) + with self.__send_lock: + with self.__lock: + if self.__closed: + if query is not None: + raise self.__connection_error(query.execution_id) + return + if query is not None: + self.__queries[query.execution_id] = query + elif message.get("execution_id") not in self.__queries: + return + try: + self.__ws.send(request) + except (websockets.exceptions.ConnectionClosed, OSError): + pass # Terminate outside the send gate, before delivering errors. + except Exception: + # API/programming errors aren't evidence of connection loss. + if query is not None: + with self.__lock: + self.__queries.pop(query.execution_id, None) + raise + else: + return + self.__fail_pending() @staticmethod def __redacted_request(message: Dict[str, Any]) -> str: @@ -370,6 +393,8 @@ def __redacted_request(message: Dict[str, Any]) -> str: def __recv(self) -> Dict[str, Any]: try: frame = self.__ws.recv(timeout=self.__read_timeout) + except TimeoutError: + raise # Idle polls are expected, not terminal I/O failures. except OSError as e: # Distinguish transport I/O failures from OSErrors raised later by # protocol parsing or result decoding; only the former are terminal. @@ -410,26 +435,16 @@ def __execute_sql( get_statement_type(sql), textwrap.shorten(redact_sql(sql), width=200), ) - send_failed = False - with self.__lock: - if self.__closed: - raise self.__connection_error(execution_id) - self.__queries[execution_id] = Query( + self.__send( + request, + Query( sql=sql, execution_id=execution_id, state=ExecutionState.EXECUTION_REQUESTED, handler=handler, store=store, - ) - try: - # Keep registration and transmission atomic with respect to - # shutdown: close() must not claim this query before its SQL is - # sent, then report failure while the request is still emitted. - self.__send(request) - except Exception: - send_failed = True - if send_failed: - self.__fail_pending() + ), + ) return execution_id def __request_results(self, execution_id: str) -> None: diff --git a/wherobots/db/driver.py b/wherobots/db/driver.py index 242daf5..63d1b24 100644 --- a/wherobots/db/driver.py +++ b/wherobots/db/driver.py @@ -267,7 +267,9 @@ def get_session_uri() -> str: data_compression=data_compression, geometry_representation=geometry_representation, cancel_event=cancel_event, - session_status_url=session_id_url, + session_id=urllib.parse.urlparse(session_id_url) + .path.rstrip("/") + .rsplit("/", 1)[-1], ) @@ -295,7 +297,7 @@ def connect_direct( data_compression: Union[DataCompression, None] = None, geometry_representation: Union[GeometryRepresentation, None] = None, cancel_event: Union[threading.Event, None] = None, - session_status_url: str | None = None, + session_id: str | None = None, ) -> Connection: uri_with_protocol = f"{uri}/{protocol}" ssl_context = ssl.create_default_context() @@ -333,30 +335,11 @@ def ws_connect() -> websockets.sync.client.ClientConnection: except Exception as e: raise InterfaceError("Failed to connect to SQL session!") from e - def failure_details() -> str | None: - if session_status_url is None: - return None - # Never follow a status redirect with the caller's credentials. - with requests.get( - session_status_url, headers=headers, timeout=1.0, allow_redirects=False - ) as response: - if response.status_code != 200: - return None - payload = response.json() - failure = payload.get("firstFailure") if isinstance(payload, dict) else None - if not isinstance(failure, dict): - return None - message = failure.get("message") - return message[:4096] if isinstance(message, str) else None - return Connection( ws, read_timeout=read_timeout, results_format=results_format, data_compression=data_compression, geometry_representation=geometry_representation, - session_id=urllib.parse.urlparse(session_status_url).path.rsplit("/", 1)[-1] - if session_status_url - else None, - failure_details=failure_details if session_status_url else None, + session_id=session_id, ) From d9b1b602b240bc90243074587ff40942d2ed21b4 Mon Sep 17 00:00:00 2001 From: Ryan Avery Date: Mon, 21 Sep 2026 13:37:23 -0700 Subject: [PATCH 4/6] test(dbapi): cover TLS shutdown and rejected cursor reuse --- tests/test_disconnect.py | 70 ++++++++++++++++++++++++++++++++++++++-- wherobots/db/cursor.py | 3 ++ 2 files changed, 71 insertions(+), 2 deletions(-) diff --git a/tests/test_disconnect.py b/tests/test_disconnect.py index 84772cc..04f4bb0 100644 --- a/tests/test_disconnect.py +++ b/tests/test_disconnect.py @@ -3,6 +3,9 @@ import json import queue import socket +import ssl +import subprocess +import shutil import threading import time from unittest.mock import MagicMock, patch @@ -18,7 +21,7 @@ from wherobots.db._transport import abort_connection from wherobots.db.connection import Connection, Query from wherobots.db.driver import connect_direct -from wherobots.db.errors import OperationalError +from wherobots.db.errors import OperationalError, ProgrammingError from wherobots.db.types import ExecutionState @@ -209,6 +212,23 @@ def test_serialization_error_does_not_register_or_fail_other_queries(): conn.close() +def test_rejected_reexecution_does_not_leave_a_stale_execution_id(): + ws = Transport() + conn = Connection(ws) + cursor = conn.cursor() + cursor.execute("SELECT 1") + deliver(ws, ws.sent[0]["execution_id"]) + # Consume the terminal outcome without depending on empty-result slicing. + assert cursor._Cursor__get_results() is None + store = MagicMock() + store.to_dict.return_value = {"invalid": object()} + with pytest.raises(TypeError): + cursor.execute("SELECT 2", store=store) + with pytest.raises(ProgrammingError, match="No query"): + cursor.fetchall() + conn.close() + + def test_nontransport_send_error_is_propagated_and_query_is_untracked(): ws = Transport() conn = Connection(ws) @@ -365,10 +385,56 @@ def delayed_abort(how): closer.join(timeout=3) -def test_real_websocket_stalled_send_is_interrupted_by_close(): +@pytest.mark.parametrize("tls", [False, True]) +def test_real_websocket_stalled_send_is_interrupted_by_close(tls, tmp_path): # Real library protocol mutex + socket.sendall; the peer never reads. local, peer = socket.socketpair() local.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 4096) + if tls: + openssl = shutil.which("openssl") + if openssl is None: + local.close() + peer.close() + pytest.skip("TLS fixture requires openssl") + key, cert = tmp_path / "key.pem", tmp_path / "cert.pem" + subprocess.run( + [ + openssl, + "req", + "-x509", + "-newkey", + "rsa:2048", + "-nodes", + "-keyout", + str(key), + "-out", + str(cert), + "-days", + "1", + "-subj", + "/CN=localhost", + ], + check=True, + capture_output=True, + ) + server_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_context.load_cert_chain(cert, key) + client_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + client_context.check_hostname = False + client_context.verify_mode = ( + ssl.CERT_NONE + ) # Local generated test certificate only. + peers = queue.Queue() + server_socket = peer + handshake = threading.Thread( + target=lambda: peers.put( + server_context.wrap_socket(server_socket, server_side=True) + ) + ) + handshake.start() + local = client_context.wrap_socket(local, server_hostname="localhost") + peer = peers.get(timeout=3) + handshake.join(timeout=3) protocol = ClientProtocol(parse_uri("ws://localhost"), state=State.OPEN) ws = ClientConnection(local, protocol) conn = Connection(ws) diff --git a/wherobots/db/cursor.py b/wherobots/db/cursor.py index af0a4e6..585c9d1 100644 --- a/wherobots/db/cursor.py +++ b/wherobots/db/cursor.py @@ -186,6 +186,9 @@ def execute( self.__rowcount = -1 self.__description = 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 = self.__exec_fn( _substitute_parameters(operation, parameters), self.__queue.put, From e211002d5fc55652152d7d9ca2b8ed7556e254b1 Mon Sep 17 00:00:00 2001 From: Ryan Avery Date: Mon, 21 Sep 2026 14:39:15 -0700 Subject: [PATCH 5/6] fix: complete query-local result failures without closing connection --- tests/test_disconnect.py | 125 +++++++++++++++++++++++++++++++++++-- wherobots/db/connection.py | 45 ++++++++++--- 2 files changed, 157 insertions(+), 13 deletions(-) diff --git a/tests/test_disconnect.py b/tests/test_disconnect.py index 04f4bb0..ef700e5 100644 --- a/tests/test_disconnect.py +++ b/tests/test_disconnect.py @@ -11,6 +11,7 @@ from unittest.mock import MagicMock, patch import pandas +import cbor2 import pytest from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK from websockets.protocol import State @@ -43,6 +44,8 @@ def recv(self, timeout): if not isinstance(value, TimeoutError): self.protocol.state = State.CLOSED raise value + if isinstance(value, bytes): + return value return json.dumps(value) def send(self, value): @@ -172,11 +175,17 @@ def test_buffered_result_is_drained_even_when_transport_is_already_closed(): @pytest.mark.parametrize( - "decode_error", [ValueError("bad payload"), OSError("decoder error")] + "decode_error", + [ + ValueError("private payload"), + OSError("private payload"), + TimeoutError("private payload"), + ConnectionClosedError(None, None), + ], ) -def test_result_decode_error_does_not_fail_other_queries(decode_error): +def test_result_decode_error_completes_only_affected_query(decode_error, caplog): ws = Transport() - conn = Connection(ws) + conn = Connection(ws, session_id="session-1") bad = conn.cursor() good = conn.cursor() bad.execute("SELECT bad") @@ -192,8 +201,113 @@ def test_result_decode_error_does_not_fail_other_queries(decode_error): ) deliver(ws, ws.sent[1]["execution_id"]) assert good._Cursor__queue.get(timeout=3).error is None + outcome = bad._Cursor__queue.get(timeout=3) + assert isinstance(outcome.error, OperationalError) + assert "Could not decode" in str(outcome.error) + assert "session-1" in str(outcome.error) + assert ws.sent[0]["execution_id"] in str(outcome.error) + assert "private payload" not in str(outcome.error) + caplog.text + assert "connection lost" not in str(outcome.error) + assert not conn._Connection__queries + bad._Cursor__queue.put(outcome) + for fetch in (bad.fetchall, bad.fetchall, bad.get_store_result): + with pytest.raises(OperationalError) as exc: + fetch() + assert exc.value is outcome.error + assert not conn._Connection__closed + good.execute("SELECT next") + deliver(ws, ws.sent[-1]["execution_id"]) + assert good._Cursor__queue.get(timeout=3).error is None + conn.close() + assert bad._Cursor__queue.empty() + + +@pytest.mark.parametrize( + "results", + [ + {"format": "json", "result_bytes": b"private malformed JSON"}, + {"format": "arrow", "result_bytes": b"private malformed Arrow"}, + {"format": "unsupported", "result_bytes": b"private"}, + {"format": "arrow"}, + ["private malformed object"], + "", + False, + ], +) +def test_malformed_result_payload_delivers_one_error(results, caplog): + ws = Transport() + conn = Connection(ws) + cursor = conn.cursor() + cursor.execute("SELECT 1") + message = cbor2.dumps( + { + "kind": "execution_result", + "execution_id": ws.sent[0]["execution_id"], + "state": "succeeded", + "results": results, + } + ) + ws.incoming.put(message) + outcome = cursor._Cursor__queue.get(timeout=3) + assert isinstance(outcome.error, OperationalError) + assert "Could not decode" in str(outcome.error) + assert "private" not in str(outcome.error) + caplog.text + assert not conn._Connection__queries + # Duplicate server messages and later shutdown cannot deliver again. + ws.incoming.put(message) + assert not conn._Connection__closed + conn.close() + assert cursor._Cursor__queue.empty() + + +@pytest.mark.parametrize("state", [None, "unknown", 123, {}]) +def test_invalid_state_completes_only_identified_query(state): + ws = Transport() + conn = Connection(ws) + bad, good = conn.cursor(), conn.cursor() + bad.execute("SELECT bad") + good.execute("SELECT good") + ws.incoming.put( + { + "kind": "execution_result", + "execution_id": ws.sent[0]["execution_id"], + "state": state, + } + ) + outcome = bad._Cursor__queue.get(timeout=3) + assert isinstance(outcome.error, OperationalError) + assert "Could not interpret" in str(outcome.error) + assert ws.sent[0]["execution_id"] not in conn._Connection__queries + deliver(ws, ws.sent[1]["execution_id"]) + assert good._Cursor__queue.get(timeout=3).error is None + assert not conn._Connection__closed + conn.close() + + +def test_local_retrieve_send_error_completes_only_affected_query(): + ws = Transport() + conn = Connection(ws) + bad, good = conn.cursor(), conn.cursor() + bad.execute("SELECT bad") + good.execute("SELECT good") + with patch.object(ws, "send", side_effect=ValueError("private API error")): + ws.incoming.put( + { + "kind": "state_updated", + "execution_id": ws.sent[0]["execution_id"], + "state": "succeeded", + } + ) + outcome = bad._Cursor__queue.get(timeout=3) + assert isinstance(outcome.error, OperationalError) + assert "Could not request" in str(outcome.error) + assert "private" not in str(outcome.error) + assert ws.sent[0]["execution_id"] not in conn._Connection__queries + deliver(ws, ws.sent[1]["execution_id"]) + assert good._Cursor__queue.get(timeout=3).error is None assert not conn._Connection__closed conn.close() + assert bad._Cursor__queue.empty() def test_serialization_error_does_not_register_or_fail_other_queries(): @@ -295,7 +409,8 @@ def progress(_): assert not conn._Connection__thread.is_alive() -def test_close_racing_result_decode_delivers_only_one_terminal_outcome(): +@pytest.mark.parametrize("decode_fails", [False, True]) +def test_close_racing_result_decode_delivers_only_one_terminal_outcome(decode_fails): decoding = threading.Event() release = threading.Event() ws = Transport() @@ -306,6 +421,8 @@ def test_close_racing_result_decode_delivers_only_one_terminal_outcome(): def decode(*args): decoding.set() assert release.wait(timeout=5) + if decode_fails: + raise ValueError("private payload") return pandas.DataFrame({"x": [1]}) with patch.object(conn, "_handle_results", side_effect=decode): diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index e4fb00c..eddcc1f 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -249,13 +249,27 @@ def complete_query(result: ExecutionResult) -> None: if claimed is not None: claimed.handler(result) + def fail_query(action: str, error: Exception) -> None: + # This is a query-local failure, not evidence of transport loss. + # Exception text may contain result data; report only its type. + query.state = ExecutionState.FAILED + message = ( + f"Could not {action} SQL results " + f"(session={self.__session_id or 'unknown'}, " + f"execution={execution_id}; {type(error).__name__}). " + "The statement may have completed; verify the operation before " + "retrying writes." + ) + logging.error("%s", message) + complete_query(ExecutionResult(error=OperationalError(message))) + # Incoming state transitions are handled here. if kind == EventKind.STATE_UPDATED or kind == EventKind.EXECUTION_RESULT: try: query.state = ExecutionState[message["state"].upper()] logging.info("Query %s is now %s.", execution_id, query.state) - except KeyError: - logging.warning("Invalid state update message for %s", execution_id) + except (KeyError, AttributeError, TypeError) as error: + fail_query("interpret", error) return if query.state == ExecutionState.SUCCEEDED: @@ -290,21 +304,34 @@ def complete_query(result: ExecutionResult) -> None: return # No store configured, request results normally - self.__request_results(execution_id) + try: + self.__request_results(execution_id) + except Exception as error: + # Transport failures are handled by __send; a local + # retrieval failure must not orphan this execution. + fail_query("request", error) return # Otherwise, process the results from the execution_result event. results = message.get("results") - if not results or not isinstance(results, dict): + if results is None or results == {}: logging.warning("Got no results back from %s.", execution_id) query.state = ExecutionState.COMPLETED complete_query(ExecutionResult()) return - query.state = ExecutionState.COMPLETED - complete_query( - ExecutionResult(results=self._handle_results(execution_id, results)) - ) + try: + if not isinstance(results, dict): + raise TypeError("Expected a result object") + decoded = self._handle_results(execution_id, results) + except Exception as error: + # Even OSError/TimeoutError here belong to decoding, not + # recv. Claim and deliver exactly once, including if close + # concurrently claims this execution. + fail_query("decode", error) + else: + query.state = ExecutionState.COMPLETED + complete_query(ExecutionResult(results=decoded)) elif query.state == ExecutionState.CANCELLED: logging.info( "Query %s has been cancelled; returning empty results.", @@ -342,7 +369,7 @@ def _handle_results(self, execution_id: str, results: Dict[str, Any]) -> Any: with pyarrow.ipc.open_stream(stream) as reader: return reader.read_pandas() else: - return OperationalError(f"Unsupported results format {result_format}") + raise NotSupportedError("Unsupported results format") def __send(self, message: Dict[str, Any], query: Query | None = None) -> None: # Serialization and redaction are local work. Fail before registration, From 4576820291ab2e975eacd315561b1aac7a32c3cb Mon Sep 17 00:00:00 2001 From: Ryan Avery Date: Mon, 21 Sep 2026 13:37:25 -0700 Subject: [PATCH 6/6] feat(dbapi): log bounded session diagnostics after connection loss --- tests/test_session_diagnostics.py | 278 ++++++++++++++++++++++++++++++ wherobots/db/_diagnostics.py | 125 ++++++++++++++ wherobots/db/connection.py | 13 +- wherobots/db/driver.py | 20 +++ 4 files changed, 434 insertions(+), 2 deletions(-) create mode 100644 tests/test_session_diagnostics.py create mode 100644 wherobots/db/_diagnostics.py diff --git a/tests/test_session_diagnostics.py b/tests/test_session_diagnostics.py new file mode 100644 index 0000000..041a419 --- /dev/null +++ b/tests/test_session_diagnostics.py @@ -0,0 +1,278 @@ +"""Optional diagnostics are bounded and never delay a cursor's terminal result.""" +import json +import threading +from unittest.mock import MagicMock, patch + +import pytest +from websockets.exceptions import ConnectionClosedError + +from wherobots.db import _diagnostics +from wherobots.db._diagnostics import SessionDiagnostics +from wherobots.db.connection import Connection +from wherobots.db.driver import connect_direct +from wherobots.db.errors import OperationalError +from test_disconnect import Transport + + +def response(payload, status=200): + result = MagicMock(status_code=status) + result.__enter__.return_value = result + result.iter_content.return_value = [json.dumps(payload).encode()] + return result + + +@pytest.fixture +def workers(monkeypatch): + # Keep real concurrency; remove only jitter and record workers for cleanup. + threads = [] + thread_type = threading.Thread + + def thread(*args, **kwargs): + result = thread_type(*args, **kwargs) + threads.append(result) + return result + + monkeypatch.setattr(_diagnostics.threading, "Thread", thread) + monkeypatch.setattr(_diagnostics.time, "sleep", lambda _: None) + yield threads + for worker in threads: + worker.join(timeout=3) + assert not worker.is_alive() + + +@pytest.mark.parametrize( + "payload,expected", + [ + ( + {"status": "FAILED", "firstFailure": {"message": "Evicted"}}, + ("FAILED", "Evicted"), + ), + ({"status": "READY"}, ("READY", None)), + ({"firstFailure": {"message": "Evicted"}}, (None, "Evicted")), + ({}, None), + (None, None), + ([], None), + ({"status": 1, "firstFailure": {"message": 5}}, None), + ({"firstFailure": "invalid"}, None), + ], +) +def test_status_and_failure_parsing(payload, expected): + with patch.object( + _diagnostics.requests, "get", return_value=response(payload) + ) as get: + assert ( + SessionDiagnostics._fetch( + "https://api/session/1", {"Authorization": "Bearer test"} + ) + == expected + ) + assert get.call_args.kwargs == { + "headers": {"Authorization": "Bearer test"}, + "timeout": 1.0, + "allow_redirects": False, + "stream": True, + } + + +def test_message_and_status_lengths_are_bounded(): + payload = {"status": "s" * 200, "firstFailure": {"message": "m" * 5000}} + with patch.object(_diagnostics.requests, "get", return_value=response(payload)): + status, message = SessionDiagnostics._fetch("https://api/session/1", {}) + assert status == "s" * 128 + assert message == "m" * 4096 + + +def test_streaming_response_size_is_bounded(): + result = response({}) + chunks_consumed = [] + + def chunks(chunk_size): + for i in range(1000): + chunks_consumed.append(i) + yield b"x" * chunk_size + + result.iter_content.side_effect = chunks + with patch.object(_diagnostics.requests, "get", return_value=result): + assert SessionDiagnostics._fetch("https://api/session/1", {}) is None + assert len(chunks_consumed) == 9 + result.__exit__.assert_called_once() + + +@pytest.mark.parametrize("status", [302, 404, 503]) +def test_http_errors_are_logged_without_parsing(status, caplog): + result = response({}, status) + with caplog.at_level("DEBUG"), patch.object( + _diagnostics.requests, "get", return_value=result + ): + assert SessionDiagnostics._fetch("https://api/session/1", {}) is None + result.iter_content.assert_not_called() + assert f"HTTP status: {status}" in caplog.text + + +def test_inflight_and_recent_requests_are_shared(workers, caplog): + service = SessionDiagnostics() + entered = threading.Event() + release = threading.Event() + + def fetch(*args): + entered.set() + assert release.wait(timeout=3) + return "FAILED", "Evicted" + + with caplog.at_level("WARNING"), patch.object( + service, "_fetch", side_effect=fetch + ) as get: + try: + service.request("https://api/session/1", {"Authorization": "secret"}, "1") + assert entered.wait(timeout=1) + for _ in range(50): + service.request( + "https://api/session/1", {"Authorization": "secret"}, "1" + ) + assert len(workers) == 1 + release.set() + workers[0].join(timeout=2) + service.request("https://api/session/1", {"Authorization": "secret"}, "1") + get.assert_called_once() + finally: + release.set() + assert "status='FAILED'" in caplog.text + assert "firstFailure='Evicted'" in caplog.text + assert "secret" not in caplog.text + + +def test_authentication_and_endpoint_contexts_are_not_shared(workers): + service = SessionDiagnostics() + with patch.object(service, "_fetch", return_value=None) as get: + for url, token in [ + ("https://one/session/1", "a"), + ("https://one/session/1", "b"), + ("https://two/session/1", "a"), + ]: + service.request(url, {"Authorization": token}, "1") + workers[-1].join(timeout=1) + assert get.call_count == 3 + assert all( + isinstance(key[1], bytes) and len(key[1]) == 32 for key in service._entries + ) + + +def test_worker_budget_bounds_stalled_lookups(workers): + service = SessionDiagnostics() + release = threading.Event() + + def fetch(*args): + assert release.wait(timeout=3) + return None + + with patch.object(service, "_fetch", side_effect=fetch): + try: + for i in range(100): + service.request(f"https://api/session/{i}", {}, str(i)) + assert len(workers) == service._active == 4 + assert len(service._entries) == 4 + finally: + release.set() + for worker in workers: + worker.join(timeout=2) + assert service._active == 0 + + +def test_cache_expires_and_has_a_fixed_size(workers, monkeypatch): + service = SessionDiagnostics() + monkeypatch.setattr(_diagnostics, "_MAX_ENTRIES", 2) + with patch.object(service, "_fetch", return_value=None) as get: + for i in range(3): + service.request(f"https://api/session/{i}", {}, str(i)) + workers[-1].join(timeout=1) + assert len(service._entries) == 2 + assert get.call_count == 3 + for key in service._entries: + service._entries[key] = 0.0 + service.request("https://api/session/2", {}, "2") + workers[-1].join(timeout=1) + assert get.call_count == 4 + + +def test_worker_exception_is_debug_logged_without_credentials(workers, caplog): + service = SessionDiagnostics() + with caplog.at_level("DEBUG"), patch.object( + service, "_fetch", side_effect=ValueError("secret token") + ): + service.request("https://api/session/1", {}, "1") + workers[0].join(timeout=1) + assert "lookup failed (ValueError)" in caplog.text + assert "secret token" not in caplog.text + assert service._active == 0 + + +def test_thread_start_failure_is_debug_logged_and_releases_slot(caplog): + service = SessionDiagnostics() + with caplog.at_level("DEBUG"), patch.object( + threading.Thread, "start", side_effect=RuntimeError + ): + service.request("https://api/session/1", {}, "1") + assert "Could not start session diagnostic thread" in caplog.text + assert service._active == 0 + + +@pytest.mark.parametrize("failure", ["send", "receive"]) +def test_stalled_lookup_does_not_delay_failure_delivery(failure, workers): + service = SessionDiagnostics() + entered = threading.Event() + release = threading.Event() + + def fetch(*args): + entered.set() + assert release.wait(timeout=3) + return None + + ws = Transport() + conn = Connection( + ws, on_connection_lost=lambda: service.request("https://api/session/1", {}, "1") + ) + with patch.object(service, "_fetch", side_effect=fetch): + try: + if failure == "send": + ws.send = MagicMock(side_effect=ConnectionClosedError(None, None)) + cursor = conn.cursor() + cursor.execute("SELECT 1") + if failure == "receive": + ws.incoming.put(ConnectionClosedError(None, None)) + assert isinstance( + cursor._Cursor__queue.get(timeout=1).error, OperationalError + ) + assert entered.wait(timeout=1) + assert not release.is_set() # Error preceded lookup completion. + finally: + release.set() + conn.close() + + +def test_explicit_close_does_not_request_diagnostics(): + callback = MagicMock() + conn = Connection(Transport(), on_connection_lost=callback) + cursor = conn.cursor() + cursor.execute("SELECT 1") + conn.close() + callback.assert_not_called() + + +@pytest.mark.parametrize("suffix", ["", "/", "/?ignored=1"]) +def test_direct_status_url_derivation_and_header_snapshot(suffix): + headers = {"Authorization": "original"} + url = "https://api/session/session-1" + suffix + with patch( + "wherobots.db.driver.websockets.sync.client.connect", return_value=Transport() + ), patch.object(_diagnostics.session_diagnostics, "request") as request: + conn = connect_direct( + "wss://compute/sql", session_status_url=url, headers=headers + ) + headers["Authorization"] = "changed" + cursor = conn.cursor() + cursor.execute("SELECT 1") + conn._Connection__ws.incoming.put(ConnectionClosedError(None, None)) + conn._Connection__thread.join(timeout=2) + request.assert_called_once_with(url, {"Authorization": "original"}, "session-1") + with pytest.raises(OperationalError, match="session=session-1"): + cursor.fetchall() diff --git a/wherobots/db/_diagnostics.py b/wherobots/db/_diagnostics.py new file mode 100644 index 0000000..28ba023 --- /dev/null +++ b/wherobots/db/_diagnostics.py @@ -0,0 +1,125 @@ +"""Bounded, best-effort session diagnostics after cursor failure delivery.""" + +import hashlib +import json +import logging +import random +import threading +import time +from collections import OrderedDict + +import requests + + +_logger = logging.getLogger(__name__) +_MAX_WORKERS = 4 +_MAX_ENTRIES = 128 +_CACHE_SECONDS = 30.0 +_MAX_RESPONSE_BYTES = 65536 + + +class SessionDiagnostics: + """Share in-flight and recent lookups by URL and authentication context. + + A hung DNS/TLS/HTTP operation retains its slot, so even failures that ignore + socket timeouts cannot cause unbounded threads. At capacity we skip optional + diagnostics. There are no retries, and no caller waits for a lookup. + """ + + def __init__(self) -> None: + self._lock = threading.Lock() + self._entries: OrderedDict[tuple[str, bytes], float | None] = OrderedDict() + self._active = 0 + + def request( + self, url: str, headers: dict[str, str] | None, session_id: str + ) -> None: + headers = dict(headers or {}) + # Distinct credentials must never share diagnostic responses. Keep only + # a digest in the short-lived registry, never credentials in log output. + context = hashlib.sha256(json.dumps(sorted(headers.items())).encode()).digest() + key = (url, context) + now = time.monotonic() + with self._lock: + for expired, deadline in list(self._entries.items()): + if deadline is not None and deadline <= now: + del self._entries[expired] + if key in self._entries: + return + if self._active >= _MAX_WORKERS: + _logger.debug("Session diagnostics at capacity; skipping lookup") + return + if len(self._entries) >= _MAX_ENTRIES: + # Never evict a running request: it must remain deduplicated. + oldest = next( + k for k, deadline in self._entries.items() if deadline is not None + ) + del self._entries[oldest] + self._entries[key] = None + self._active += 1 + + def lookup() -> None: + try: + # Spread correlated failures across processes without delaying + # execute(), fetch(), or connection shutdown. + time.sleep(random.uniform(0.0, 0.5)) + details = self._fetch(url, headers) + if details is not None: + status, message = details + _logger.warning( + "SQL session diagnostic (session=%s): status=%r; firstFailure=%r", + session_id, + status, + message, + ) + except Exception as exc: + # Request exceptions may contain URLs/credentials. Log the type, + # not arbitrary exception text or the authenticated request. + _logger.debug( + "Session diagnostic lookup failed (%s)", type(exc).__name__ + ) + finally: + with self._lock: + self._active -= 1 + self._entries[key] = time.monotonic() + _CACHE_SECONDS + + try: + threading.Thread( + target=lookup, daemon=True, name="wherobots-session-diagnostic" + ).start() + except RuntimeError: + with self._lock: + self._active -= 1 + self._entries[key] = time.monotonic() + _CACHE_SECONDS + _logger.debug("Could not start session diagnostic thread") + + @staticmethod + def _fetch( + url: str, headers: dict[str, str] + ) -> tuple[str | None, str | None] | None: + with requests.get( + url, headers=headers, timeout=1.0, allow_redirects=False, stream=True + ) as response: + if response.status_code != 200: + _logger.debug( + "Session diagnostic HTTP status: %s", response.status_code + ) + return None + data = bytearray() + for chunk in response.iter_content(chunk_size=8192): + if len(data) + len(chunk) > _MAX_RESPONSE_BYTES: + _logger.debug("Session diagnostic response exceeds byte limit") + return None + data.extend(chunk) + payload = json.loads(data) + if not isinstance(payload, dict): + return None + status = payload.get("status") + status = status[:128] if isinstance(status, str) else None + failure = payload.get("firstFailure") + message = failure.get("message") if isinstance(failure, dict) else None + message = message[:4096] if isinstance(message, str) else None + return (status, message) if status is not None or message is not None else None + + +session_diagnostics = SessionDiagnostics() diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index eddcc1f..96ab8e2 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -70,6 +70,7 @@ def __init__( data_compression: DataCompression | None = None, geometry_representation: GeometryRepresentation | None = None, session_id: str | None = None, + on_connection_lost: Callable[[], None] | None = None, ): self.__ws = ws self.__read_timeout = read_timeout @@ -79,6 +80,9 @@ def __init__( self.__progress_handler: ProgressHandler | None = None self.__session_id = session_id + # Internal notification hook: schedule optional diagnostics, never wait + # for remote results here. All cursor failures are delivered first. + self.__on_connection_lost = on_connection_lost self.__lock = threading.Lock() self.__send_lock = threading.Lock() self.__shutdown_done = threading.Event() @@ -103,7 +107,7 @@ def close(self) -> None: decoder or callback can outlive the bounded reader join. """ deadline = time.monotonic() + 1.0 - self.__fail_pending() + self.__fail_pending(notify=False) # A handler may close its own connection during terminal delivery. if self.__shutdown_owner == threading.get_ident(): return @@ -170,7 +174,7 @@ def __connection_error(self, execution_id: str) -> OperationalError: ) return OperationalError(message) - def __fail_pending(self) -> None: + def __fail_pending(self, notify: bool = True) -> None: # Stop admission first. Do not wait for __send_lock: its owner may be # blocked in network I/O. __closed means closing until shutdown_done. with self.__lock: @@ -201,6 +205,11 @@ def __fail_pending(self) -> None: finally: self.__shutdown_owner = None self.__shutdown_done.set() + if notify and pending and self.__on_connection_lost is not None: + try: + self.__on_connection_lost() + except Exception: + logging.debug("Could not schedule session diagnostics") def __listen(self) -> None: """Waits for the next message from the SQL session and processes it. diff --git a/wherobots/db/driver.py b/wherobots/db/driver.py index 63d1b24..3ef119d 100644 --- a/wherobots/db/driver.py +++ b/wherobots/db/driver.py @@ -267,6 +267,7 @@ def get_session_uri() -> str: data_compression=data_compression, geometry_representation=geometry_representation, cancel_event=cancel_event, + session_status_url=session_id_url, session_id=urllib.parse.urlparse(session_id_url) .path.rstrip("/") .rsplit("/", 1)[-1], @@ -298,6 +299,7 @@ def connect_direct( geometry_representation: Union[GeometryRepresentation, None] = None, cancel_event: Union[threading.Event, None] = None, session_id: str | None = None, + session_status_url: str | None = None, ) -> Connection: uri_with_protocol = f"{uri}/{protocol}" ssl_context = ssl.create_default_context() @@ -335,6 +337,23 @@ def ws_connect() -> websockets.sync.client.ClientConnection: except Exception as e: raise InterfaceError("Failed to connect to SQL session!") from e + on_connection_lost = None + if session_status_url is not None: + from ._diagnostics import session_diagnostics + + if session_id is None: + session_id = ( + urllib.parse.urlparse(session_status_url) + .path.rstrip("/") + .rsplit("/", 1)[-1] + ) + diagnostic_headers = dict(headers or {}) + + def on_connection_lost() -> None: + session_diagnostics.request( + session_status_url, diagnostic_headers, session_id or "unknown" + ) + return Connection( ws, read_timeout=read_timeout, @@ -342,4 +361,5 @@ def ws_connect() -> websockets.sync.client.ClientConnection: data_compression=data_compression, geometry_representation=geometry_representation, session_id=session_id, + on_connection_lost=on_connection_lost, )