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, )