From a2d27a85c041dad7cc6a2a93c50603a6eb226ed7 Mon Sep 17 00:00:00 2001 From: Ryan Avery Date: Fri, 4 Sep 2026 22:01:07 -0700 Subject: [PATCH 1/7] 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/7] 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/7] 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/7] 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/7] 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 33657063188ada6d685e4a656e6e2c316655792c Mon Sep 17 00:00:00 2001 From: Simon Fishel Date: Mon, 21 Sep 2026 18:51:40 -0700 Subject: [PATCH 6/7] fix(dbapi): always deliver pending failures when transport teardown fails __fail_pending latched __closed before calling abort_connection, but only OSError was caught and the try/finally that sets __shutdown_done started below it. Any other exception from the abort skipped the delivery loop, and since __closed is a one-way latch every later __fail_pending early-returned: pending queries were never failed, __shutdown_done never set, and fetchall() blocked forever -- the hang this PR exists to remove. Run everything after the latch under one try/finally, and make transport termination catch every exception so delivery is never skipped. --- tests/test_disconnect.py | 19 +++++++++++++++++++ wherobots/db/connection.py | 22 ++++++++++++++-------- 2 files changed, 33 insertions(+), 8 deletions(-) diff --git a/tests/test_disconnect.py b/tests/test_disconnect.py index ef700e5..15b156b 100644 --- a/tests/test_disconnect.py +++ b/tests/test_disconnect.py @@ -470,6 +470,25 @@ def test_abort_closes_socket_even_if_shutdown_errors(): ws.socket.close.assert_called_once() +def test_abort_failure_still_fails_pending_and_completes_shutdown(): + # A non-OSError from the adapter (e.g. a renamed attribute in a future + # websockets release) must not skip delivery: __closed is a one-way latch. + ws = Transport() + conn = Connection(ws) + cursor = conn.cursor() + cursor.execute("SELECT 1") + ws.socket.shutdown.side_effect = AttributeError("no attribute 'socket'") + conn.close() + with pytest.raises(OperationalError): + cursor.fetchall() + assert not conn._Connection__queries + assert conn._Connection__shutdown_done.is_set() + assert conn._Connection__shutdown_owner is None + with pytest.raises(OperationalError): + conn.cursor().execute("SELECT 2") + conn.close() + + def test_failure_is_not_delivered_before_transport_is_disabled(): ws = Transport() conn = Connection(ws) diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index eddcc1f..a6d7430 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -178,15 +178,13 @@ def __fail_pending(self) -> None: return self.__closed = True self.__shutdown_owner = threading.get_ident() + # __closed is a one-way latch: nothing below may be skipped, or pending + # queries are stranded with no path to recovery. 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() - try: + self.__terminate_transport() + with self.__lock: + pending = list(self.__queries.values()) + self.__queries.clear() for query in pending: try: query.handler( @@ -202,6 +200,14 @@ def __fail_pending(self) -> None: self.__shutdown_owner = None self.__shutdown_done.set() + def __terminate_transport(self) -> None: + """Disable the transport. Never raises: delivery must not be skipped.""" + try: + abort_connection(self.__ws) + except Exception: + # The adapter still closes the socket object in its finally block. + logging.exception("Could not abort SQL session transport") + def __listen(self) -> None: """Waits for the next message from the SQL session and processes it. From 66d0daa770199b9a586270e959817189e456d8d9 Mon Sep 17 00:00:00 2001 From: Simon Fishel Date: Mon, 21 Sep 2026 18:53:28 -0700 Subject: [PATCH 7/7] fix(dbapi): close the WebSocket cleanly when no send is in flight Every close() routed through abort_connection, which shuts the socket down without sending a close frame -- even on a healthy `with connect()` exit. The server could not distinguish a client exit from a crash and had to fall back to its idle timeout to reclaim the session. The unconditional abort exists because ws.close() takes the library's protocol mutex, which a stalled sendall() may hold. That is only true when a send is actually in flight, and __send_lock tells us: a non-blocking acquire succeeds only when no sender is active, in which case close() performs the handshake; otherwise, and on reader-side failures where the transport is already broken, we abort as before. The handshake is bounded by close_timeout on the connection. --- tests/test_disconnect.py | 145 +++++++++++++++++++++++++++++++++---- tests/test_driver.py | 15 ++++ wherobots/db/connection.py | 34 +++++++-- wherobots/db/constants.py | 2 + wherobots/db/driver.py | 2 + 5 files changed, 177 insertions(+), 21 deletions(-) diff --git a/tests/test_disconnect.py b/tests/test_disconnect.py index 15b156b..8121875 100644 --- a/tests/test_disconnect.py +++ b/tests/test_disconnect.py @@ -32,6 +32,8 @@ def __init__(self): self.incoming = queue.Queue() self.sent = [] self.aborted = threading.Event() + self.closed = threading.Event() + self.close_calls = 0 self.socket = MagicMock() self.socket.shutdown.side_effect = self.shutdown @@ -58,6 +60,13 @@ def shutdown(self, how): self.aborted.set() self.incoming.put(ConnectionClosedOK(None, None)) + def close(self): + # Graceful handshake: the library's close() completes and recv sees EOF. + self.close_calls += 1 + self.closed.set() + self.aborted.set() + self.incoming.put(ConnectionClosedOK(None, None)) + def deliver(ws, execution_id): ws.incoming.put( @@ -143,7 +152,67 @@ def test_close_fails_pending_and_joins_reader(): cursor.fetchall() assert not conn._Connection__thread.is_alive() conn.close() + assert ws.close_calls == 1 + ws.socket.shutdown.assert_not_called() + + +def test_idle_close_performs_close_handshake(): + ws = Transport() + conn = Connection(ws) + conn.close() + assert ws.close_calls == 1 + ws.socket.shutdown.assert_not_called() + assert not conn._Connection__thread.is_alive() + assert not conn._Connection__send_lock.locked() + + +def test_close_with_stalled_send_falls_back_to_abort(): + ws = Transport() + conn = Connection(ws) + cursor = conn.cursor() + sending, timed_out = threading.Event(), threading.Event() + original_send = ws.send + + def blocked_send(value): + sending.set() + # Only actual transport shutdown releases the writer. Record a timeout + # rather than asserting here: __send re-raises, but only the test body + # can report it. + if not ws.aborted.wait(timeout=3): + timed_out.set() + original_send(value) + + ws.send = blocked_send + sender = threading.Thread(target=cursor.execute, args=("SELECT 1",)) + sender.start() + try: + assert sending.wait(timeout=1) + conn.close() + ws.socket.shutdown.assert_called_once() + assert ws.close_calls == 0 + assert isinstance(cursor._Cursor__queue.get(timeout=2).error, OperationalError) + finally: + conn.close() + sender.join(timeout=3) + assert not sender.is_alive() + assert not timed_out.is_set() + + +@pytest.mark.parametrize( + "error", [ConnectionClosedError(None, None), OSError("transport lost")] +) +def test_reader_failure_aborts_without_close_handshake(error): + ws = Transport() + conn = Connection(ws) + cursor = conn.cursor() + cursor.execute("SELECT 1") + ws.incoming.put(error) + conn._Connection__thread.join(timeout=3) + assert not conn._Connection__thread.is_alive() ws.socket.shutdown.assert_called_once() + assert ws.close_calls == 0 + with pytest.raises(OperationalError): + cursor.fetchall() @pytest.mark.parametrize( @@ -459,7 +528,8 @@ def test_concurrent_close_delivers_once(): 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() + assert ws.close_calls == 1 + ws.socket.shutdown.assert_not_called() def test_abort_closes_socket_even_if_shutdown_errors(): @@ -477,6 +547,7 @@ def test_abort_failure_still_fails_pending_and_completes_shutdown(): conn = Connection(ws) cursor = conn.cursor() cursor.execute("SELECT 1") + ws.close = MagicMock(side_effect=AttributeError("no attribute 'close'")) ws.socket.shutdown.side_effect = AttributeError("no attribute 'socket'") conn.close() with pytest.raises(OperationalError): @@ -489,38 +560,80 @@ def test_abort_failure_still_fails_pending_and_completes_shutdown(): conn.close() -def test_failure_is_not_delivered_before_transport_is_disabled(): +def _delayed(hook, started, release, timed_out): + """Wrap a teardown hook so the test controls when it completes. + + Production code swallows exceptions from the transport, so an ``assert`` + inside the hook would be invisible; record a timeout instead and let the + test body assert on it. + """ + + def delayed(*args): + started.set() + if not release.wait(timeout=3): + timed_out.set() + hook(*args) + + return delayed + + +def test_failure_is_not_delivered_before_graceful_close_completes(): 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 + started, release, timed_out = (threading.Event() for _ in range(3)) + ws.close = _delayed(ws.close, started, release, timed_out) closer = threading.Thread(target=conn.close) closer.start() try: - assert abort_started.wait(timeout=1) + assert started.wait(timeout=1) assert cursor._Cursor__queue.empty() + # Senders fail fast on the latch; they don't wait for the handshake. + begun = time.monotonic() with pytest.raises(OperationalError): conn.cursor().execute("SELECT 2") - release_abort.set() + assert time.monotonic() - begun < 1 + assert cursor._Cursor__queue.empty() + release.set() closer.join(timeout=2) assert not closer.is_alive() - assert ws.aborted.is_set() + assert not timed_out.is_set() + assert ws.close_calls == 1 + ws.socket.shutdown.assert_not_called() assert isinstance(cursor._Cursor__queue.get(timeout=1).error, OperationalError) finally: - release_abort.set() + release.set() closer.join(timeout=3) +def test_failure_is_not_delivered_before_transport_is_aborted(): + ws = Transport() + conn = Connection(ws) + cursor = conn.cursor() + cursor.execute("SELECT 1") + started, release, timed_out = (threading.Event() for _ in range(3)) + ws.socket.shutdown.side_effect = _delayed(ws.shutdown, started, release, timed_out) + # Reader-side failure: the abort runs on the reader thread. + ws.incoming.put(ConnectionClosedError(None, None)) + try: + assert started.wait(timeout=1) + assert cursor._Cursor__queue.empty() + with pytest.raises(OperationalError): + conn.cursor().execute("SELECT 2") + assert cursor._Cursor__queue.empty() + release.set() + conn._Connection__thread.join(timeout=2) + assert not conn._Connection__thread.is_alive() + assert not timed_out.is_set() + ws.socket.shutdown.assert_called_once() + assert ws.close_calls == 0 + assert isinstance(cursor._Cursor__queue.get(timeout=1).error, OperationalError) + finally: + release.set() + conn.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. diff --git a/tests/test_driver.py b/tests/test_driver.py index eeb0169..6c9cb7d 100644 --- a/tests/test_driver.py +++ b/tests/test_driver.py @@ -8,6 +8,7 @@ import pytest import requests +from wherobots.db.constants import DEFAULT_CLOSE_TIMEOUT_SECONDS from wherobots.db.driver import ( DEFAULT_HTTP_TIMEOUT, _check_cancelled, @@ -234,6 +235,20 @@ def test_cancel_before_ws_connect(self, mock_ws): mock_ws.assert_not_called() +class TestConnectDirectWebSocket: + @patch("wherobots.db.driver.websockets.sync.client.connect") + def test_close_handshake_is_bounded(self, mock_ws): + from websockets.exceptions import ConnectionClosedOK + + ws = MagicMock() + ws.recv.side_effect = ConnectionClosedOK(None, None) + mock_ws.return_value = ws + conn = connect_direct(uri="wss://compute.example.com/sql/org/session-id") + conn.close() + _, kwargs = mock_ws.call_args + assert kwargs["close_timeout"] == DEFAULT_CLOSE_TIMEOUT_SECONDS + + class TestWherobotsClientHeader: """connect() emits/appends the shared X-Wherobots-Client hop. diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index a6d7430..e2f3be2 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -97,13 +97,20 @@ def __exit__(self, exc_type, exc_val, exc_tb): self.close() def close(self) -> None: - """Abort the transport, fail pending work, and wait up to 1s for the reader. + """Close the transport, fail pending work, and wait up to 1s for the reader. + + The WebSocket close handshake is attempted when no send is in flight, + so the server can distinguish a client exit from a crash; otherwise the + socket is aborted. The handshake is bounded by the ``close_timeout`` of + the underlying ``ClientConnection`` (1s via ``connect``/ + ``connect_direct``; the library default of 10s for a caller-built + socket), and shares the 1s budget below with the reader join. 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() + self.__fail_pending(graceful=True) # A handler may close its own connection during terminal delivery. if self.__shutdown_owner == threading.get_ident(): return @@ -170,7 +177,7 @@ def __connection_error(self, execution_id: str) -> OperationalError: ) return OperationalError(message) - def __fail_pending(self) -> None: + def __fail_pending(self, graceful: bool = False) -> 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: @@ -181,7 +188,7 @@ def __fail_pending(self) -> None: # __closed is a one-way latch: nothing below may be skipped, or pending # queries are stranded with no path to recovery. try: - self.__terminate_transport() + self.__terminate_transport(graceful) with self.__lock: pending = list(self.__queries.values()) self.__queries.clear() @@ -200,8 +207,25 @@ def __fail_pending(self) -> None: self.__shutdown_owner = None self.__shutdown_done.set() - def __terminate_transport(self) -> None: + def __terminate_transport(self, graceful: bool) -> None: """Disable the transport. Never raises: delivery must not be skipped.""" + # A non-blocking acquire tells us whether a sender is inside ws.send() + # right now, possibly stalled holding the library's protocol mutex, + # which close() would deadlock on. Releasing immediately is safe: the + # __closed latch is already set, and __send re-checks it under __lock + # after taking __send_lock, so no sender can reach ws.send() from here + # on. Holding the lock across the handshake would only park unrelated + # senders (including the reader's own __request_results) for up to + # close_timeout instead of letting them fail fast. On a reader-side + # failure the transport is already broken and a close frame is + # pointless, so callers abort directly. + if graceful and self.__send_lock.acquire(blocking=False): + self.__send_lock.release() + try: + self.__ws.close() + return + except Exception: + logging.debug("Graceful close failed; aborting", exc_info=True) try: abort_connection(self.__ws) except Exception: diff --git a/wherobots/db/constants.py b/wherobots/db/constants.py index 809b635..0b409cd 100644 --- a/wherobots/db/constants.py +++ b/wherobots/db/constants.py @@ -14,6 +14,8 @@ DEFAULT_SESSION_TYPE: SessionType = SessionType.MULTI DEFAULT_STORAGE_FORMAT: StorageFormat = StorageFormat.PARQUET DEFAULT_READ_TIMEOUT_SECONDS: float = 0.25 +# 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 MAX_MESSAGE_SIZE: int = 100 * 2**20 # 100MiB diff --git a/wherobots/db/driver.py b/wherobots/db/driver.py index 63d1b24..932cff6 100644 --- a/wherobots/db/driver.py +++ b/wherobots/db/driver.py @@ -20,6 +20,7 @@ from .connection import Connection from .constants import ( + DEFAULT_CLOSE_TIMEOUT_SECONDS, DEFAULT_ENDPOINT, DEFAULT_READ_TIMEOUT_SECONDS, DEFAULT_SESSION_TYPE, @@ -325,6 +326,7 @@ def ws_connect() -> websockets.sync.client.ClientConnection: additional_headers=headers, max_size=MAX_MESSAGE_SIZE, open_timeout=DEFAULT_HTTP_TIMEOUT, + close_timeout=DEFAULT_CLOSE_TIMEOUT_SECONDS, ssl=ssl_context, )