diff --git a/tests/test_connection_query_tracking.py b/tests/test_connection_query_tracking.py new file mode 100644 index 0000000..e69089c --- /dev/null +++ b/tests/test_connection_query_tracking.py @@ -0,0 +1,259 @@ +"""Tests for connection-level query lifecycle tracking (WBC-922). + +A query must be removed from the connection's tracking dict as soon as its +terminal result is delivered, on every terminal path — success with a store +result, success with an empty result, cancellation, and error. Retaining +completed queries pins their handlers (and any results those reference) for +the connection's lifetime. +""" + +import json +import queue +from unittest.mock import MagicMock + +import cbor2 +import pyarrow + +from wherobots.db.connection import Connection, Query +from wherobots.db.models import ExecutionResult, Store +from wherobots.db.types import ExecutionState, StorageFormat + + +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) + + +def _track_query(conn, execution_id="exec-1", state=ExecutionState.RUNNING, store=None): + """Register a query on the connection and return its result queue.""" + result_queue = queue.Queue() + query = Query( + sql="SELECT 1", + execution_id=execution_id, + state=state, + handler=result_queue.put, + store=store, + ) + conn._Connection__queries[execution_id] = query + return result_queue + + +def _deliver(conn, message): + """Feed one message through the connection's listener.""" + conn._Connection__ws.recv.return_value = json.dumps(message) + conn._Connection__listen() + + +def _deliver_binary(conn, message): + """Feed one CBOR-encoded message (used for binary result payloads).""" + conn._Connection__ws.recv.return_value = cbor2.dumps(message) + conn._Connection__listen() + + +class TestTerminalDeliveryStopsTracking: + """Each terminal path must pop the query from __queries.""" + + def test_store_result_success_is_untracked(self): + conn = _make_connection() + result_queue = _track_query( + conn, store=Store.for_download(format=StorageFormat.GEOJSON) + ) + + _deliver( + conn, + { + "kind": "state_updated", + "execution_id": "exec-1", + "state": "succeeded", + "result_uri": "s3://results/exec-1", + "size": 42, + }, + ) + + result = result_queue.get(timeout=1) + assert result.store_result.result_uri == "s3://results/exec-1" + assert "exec-1" not in conn._Connection__queries + + def test_empty_store_success_is_untracked(self): + conn = _make_connection() + result_queue = _track_query( + conn, store=Store.for_download(format=StorageFormat.GEOJSON) + ) + + _deliver( + conn, + { + "kind": "state_updated", + "execution_id": "exec-1", + "state": "succeeded", + "result_uri": None, + "size": None, + }, + ) + + result = result_queue.get(timeout=1) + assert isinstance(result, ExecutionResult) + assert "exec-1" not in conn._Connection__queries + + def test_empty_execution_result_is_untracked(self): + conn = _make_connection() + result_queue = _track_query(conn) + + _deliver( + conn, + { + "kind": "execution_result", + "execution_id": "exec-1", + "state": "succeeded", + "results": None, + }, + ) + + result = result_queue.get(timeout=1) + assert result.results is None + assert "exec-1" not in conn._Connection__queries + + def test_json_results_success_is_untracked(self): + """The succeeded path with an actual JSON payload delivers decoded + rows and stops tracking the query.""" + conn = _make_connection() + result_queue = _track_query(conn) + + _deliver_binary( + conn, + { + "kind": "execution_result", + "execution_id": "exec-1", + "state": "succeeded", + "results": { + "result_bytes": b'[{"x": 1}, {"x": 2}]', + "format": "json", + }, + }, + ) + + result = result_queue.get(timeout=1) + assert result.results == [{"x": 1}, {"x": 2}] + assert "exec-1" not in conn._Connection__queries + + def test_arrow_results_success_is_untracked(self): + """The succeeded path with an actual Arrow IPC payload delivers a + DataFrame and stops tracking the query.""" + conn = _make_connection() + result_queue = _track_query(conn) + + table = pyarrow.table({"x": [1, 2, 3]}) + sink = pyarrow.BufferOutputStream() + with pyarrow.ipc.new_stream(sink, table.schema) as writer: + writer.write_table(table) + + _deliver_binary( + conn, + { + "kind": "execution_result", + "execution_id": "exec-1", + "state": "succeeded", + "results": { + "result_bytes": sink.getvalue().to_pybytes(), + "format": "arrow", + }, + }, + ) + + result = result_queue.get(timeout=1) + assert result.results["x"].tolist() == [1, 2, 3] + assert "exec-1" not in conn._Connection__queries + + def test_cancelled_query_is_untracked(self): + conn = _make_connection() + result_queue = _track_query(conn) + + _deliver( + conn, + { + "kind": "state_updated", + "execution_id": "exec-1", + "state": "cancelled", + }, + ) + + result = result_queue.get(timeout=1) + assert result.results.empty + assert "exec-1" not in conn._Connection__queries + + def test_errored_query_is_untracked(self): + conn = _make_connection() + result_queue = _track_query(conn) + + _deliver( + conn, + { + "kind": "error", + "execution_id": "exec-1", + "message": "boom", + }, + ) + + result = result_queue.get(timeout=1) + assert result.error is not None + assert "exec-1" not in conn._Connection__queries + + def test_cancel_of_untracked_query_is_noop(self): + """Cancelling an execution that already completed (and was popped) + must not send anything over the wire.""" + conn = _make_connection() + conn._Connection__ws.send.reset_mock() + + conn._Connection__cancel_query("exec-gone") + + conn._Connection__ws.send.assert_not_called() + + def test_non_terminal_state_update_keeps_tracking(self): + """A running-state update is not terminal; the query stays tracked.""" + conn = _make_connection() + _track_query(conn, state=ExecutionState.EXECUTION_REQUESTED) + + _deliver( + conn, + { + "kind": "state_updated", + "execution_id": "exec-1", + "state": "running", + }, + ) + + assert "exec-1" in conn._Connection__queries + + def test_failed_state_keeps_tracking_until_error_event(self): + """A failed-state update is not terminal by itself — the query must + stay tracked so the follow-up error event can deliver the message.""" + conn = _make_connection() + result_queue = _track_query(conn) + + _deliver( + conn, + { + "kind": "state_updated", + "execution_id": "exec-1", + "state": "failed", + }, + ) + + assert "exec-1" in conn._Connection__queries + assert result_queue.empty() + + _deliver( + conn, + { + "kind": "error", + "execution_id": "exec-1", + "message": "boom", + }, + ) + + result = result_queue.get(timeout=1) + assert "boom" in str(result.error) + assert "exec-1" not in conn._Connection__queries diff --git a/tests/test_cursor.py b/tests/test_cursor.py index 7f6c585..2cfe7a9 100644 --- a/tests/test_cursor.py +++ b/tests/test_cursor.py @@ -6,15 +6,24 @@ 2. Pyformat parameter substitution (%(name)s) works correctly with type-aware SQL quoting. 3. Unknown parameter keys raise ProgrammingError. +4. Fetches only ever observe the most recent execution's result set, and + re-executing never cancels an already-completed statement (WBC-922). """ from datetime import date +import pandas import pytest from unittest.mock import MagicMock from wherobots.db.cursor import Cursor, _substitute_parameters, _quote_value -from wherobots.db.errors import ProgrammingError +from wherobots.db.errors import OperationalError, ProgrammingError +from wherobots.db.models import ( + ExecutionResult, + StorageFormat, + Store, + StoreResult, +) def _make_cursor(): @@ -211,6 +220,194 @@ def test_unknown_parameter_raises(self): cursor.execute(sql, parameters={"id": 42}) +# --------------------------------------------------------------------------- +# Result isolation and cancellation tests (WBC-922) +# --------------------------------------------------------------------------- + + +def _make_async_cursor(): + """Create a Cursor whose exec_fn records each execution's handler. + + Tests deliver results by invoking a recorded handler, mimicking the + connection's asynchronous result callbacks. + """ + handlers = [] + + def exec_fn(sql, handler, store): + handlers.append(handler) + return f"exec-{len(handlers)}" + + cancel_fn = MagicMock() + return Cursor(exec_fn, cancel_fn), handlers, cancel_fn + + +class TestCursorResultIsolation: + """Fetches must only observe the most recent execution's result set.""" + + def test_unfetched_result_does_not_leak_into_next_execute(self): + cursor, handlers, _ = _make_async_cursor() + + cursor.execute("SELECT 1") + handlers[0](ExecutionResult(results=pandas.DataFrame({"x": [1]}))) + + # Re-execute without fetching the first result. + cursor.execute("SELECT 2") + handlers[1](ExecutionResult(results=pandas.DataFrame({"x": [2]}))) + + assert cursor.fetchall()["x"].tolist() == [2] + + def test_late_result_from_superseded_execution_is_ignored(self): + cursor, handlers, _ = _make_async_cursor() + + cursor.execute("SELECT 1") + # First query still in flight when the second is executed. + cursor.execute("SELECT 2") + + # The first query's result arrives late (e.g. the empty result the + # connection delivers for a cancelled query), then the second's. + handlers[0](ExecutionResult(results=pandas.DataFrame())) + handlers[1](ExecutionResult(results=pandas.DataFrame({"x": [2]}))) + + assert cursor.fetchall()["x"].tolist() == [2] + + def test_fetch_after_fetch_returns_same_results(self): + cursor, handlers, _ = _make_async_cursor() + + cursor.execute("SELECT 1") + handlers[0](ExecutionResult(results=pandas.DataFrame({"x": [1]}))) + + assert cursor.fetchall()["x"].tolist() == [1] + assert cursor.fetchall()["x"].tolist() == [1] + + def test_get_store_result_is_idempotent(self): + """A second get_store_result() must not block on the drained queue.""" + cursor, handlers, _ = _make_async_cursor() + + cursor.execute("SELECT 1", store=Store(format=StorageFormat.PARQUET)) + handlers[0]( + ExecutionResult(store_result=StoreResult(result_uri="s3://r/1", size=42)) + ) + + assert cursor.get_store_result().result_uri == "s3://r/1" + assert cursor.get_store_result().result_uri == "s3://r/1" + + def test_fetch_after_error_reraises(self): + """Repeated fetches of a failed execution re-raise its error instead + of blocking on the drained queue.""" + cursor, handlers, _ = _make_async_cursor() + + cursor.execute("SELECT broken") + handlers[0](ExecutionResult(error=OperationalError("boom"))) + + with pytest.raises(OperationalError, match="boom"): + cursor.fetchall() + with pytest.raises(OperationalError, match="boom"): + cursor.fetchall() + + +class TestCursorCancellation: + """Only genuinely in-flight executions may be cancelled.""" + + def test_execute_cancels_in_flight_previous_query(self): + cursor, _, cancel_fn = _make_async_cursor() + + cursor.execute("SELECT 1") + cursor.execute("SELECT 2") + + cancel_fn.assert_called_once_with("exec-1") + + def test_execute_does_not_cancel_completed_previous_query(self): + cursor, handlers, cancel_fn = _make_async_cursor() + + cursor.execute("MERGE INTO t USING s ON t.id = s.id ...") + handlers[0](ExecutionResult(results=pandas.DataFrame())) + + # The DML completed (result queued, not fetched); executing another + # statement must not attempt to cancel it. + cursor.execute("SELECT 1") + + cancel_fn.assert_not_called() + + def test_close_cancels_in_flight_query(self): + cursor, _, cancel_fn = _make_async_cursor() + + cursor.execute("SELECT 1") + cursor.close() + + cancel_fn.assert_called_once_with("exec-1") + + def test_close_does_not_cancel_completed_query(self): + cursor, handlers, cancel_fn = _make_async_cursor() + + cursor.execute("SELECT 1") + handlers[0](ExecutionResult(results=pandas.DataFrame({"x": [1]}))) + cursor.close() + + cancel_fn.assert_not_called() + + def test_close_without_execute_does_not_cancel(self): + cursor, _, cancel_fn = _make_async_cursor() + + cursor.close() + + cancel_fn.assert_not_called() + + def test_execute_does_not_cancel_completed_store_query(self): + """Store-backed executions never populate __results; completion must + still be recognized so they aren't cancelled (WBC-922 review).""" + cursor, handlers, cancel_fn = _make_async_cursor() + + cursor.execute("SELECT * FROM t", store=Store(format=StorageFormat.PARQUET)) + handlers[0]( + ExecutionResult(store_result=StoreResult(result_uri="s3://r/1", size=42)) + ) + assert cursor.get_store_result().result_uri == "s3://r/1" + + cursor.execute("SELECT 1") + + cancel_fn.assert_not_called() + + def test_close_does_not_cancel_completed_store_query(self): + cursor, handlers, cancel_fn = _make_async_cursor() + + cursor.execute("SELECT * FROM t", store=Store(format=StorageFormat.PARQUET)) + handlers[0]( + ExecutionResult(store_result=StoreResult(result_uri="s3://r/1", size=42)) + ) + cursor.get_store_result() + cursor.close() + + cancel_fn.assert_not_called() + + def test_execute_does_not_cancel_completed_empty_result_query(self): + """An execution that completes with neither rows nor a store result + (e.g. store configured but empty result set) must not be cancelled.""" + cursor, handlers, cancel_fn = _make_async_cursor() + + cursor.execute( + "SELECT 1 WHERE 1 = 0", store=Store(format=StorageFormat.PARQUET) + ) + handlers[0](ExecutionResult()) + assert cursor.get_store_result() is None + + cursor.execute("SELECT 2") + + cancel_fn.assert_not_called() + + def test_execute_does_not_cancel_failed_query(self): + """A failed execution is terminal; re-executing must not cancel it.""" + cursor, handlers, cancel_fn = _make_async_cursor() + + cursor.execute("SELECT broken") + handlers[0](ExecutionResult(error=OperationalError("boom"))) + with pytest.raises(OperationalError): + cursor.fetchall() + + cursor.execute("SELECT 1") + + cancel_fn.assert_not_called() + + # --------------------------------------------------------------------------- # _substitute_parameters unit tests # --------------------------------------------------------------------------- diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index cdb420b..7be2f88 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -164,6 +164,13 @@ def __listen(self) -> None: ) return + 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) + # Incoming state transitions are handled here. if kind == EventKind.STATE_UPDATED or kind == EventKind.EXECUTION_RESULT: try: @@ -191,7 +198,7 @@ def __listen(self) -> None: store_result.size, ) query.state = ExecutionState.COMPLETED - query.handler(ExecutionResult(store_result=store_result)) + complete_query(ExecutionResult(store_result=store_result)) return if query.store is not None: @@ -201,7 +208,7 @@ def __listen(self) -> None: execution_id, ) query.state = ExecutionState.COMPLETED - query.handler(ExecutionResult()) + complete_query(ExecutionResult()) return # No store configured, request results normally @@ -213,11 +220,11 @@ def __listen(self) -> None: if not results or not isinstance(results, dict): logging.warning("Got no results back from %s.", execution_id) query.state = ExecutionState.COMPLETED - query.handler(ExecutionResult()) + complete_query(ExecutionResult()) return query.state = ExecutionState.COMPLETED - query.handler( + complete_query( ExecutionResult(results=self._handle_results(execution_id, results)) ) elif query.state == ExecutionState.CANCELLED: @@ -225,8 +232,7 @@ def __listen(self) -> None: "Query %s has been cancelled; returning empty results.", execution_id, ) - query.handler(ExecutionResult(results=pandas.DataFrame())) - self.__queries.pop(execution_id) + complete_query(ExecutionResult(results=pandas.DataFrame())) elif query.state == ExecutionState.FAILED: # Don't do anything here; the ERROR event is coming with more # details. @@ -234,7 +240,7 @@ def __listen(self) -> None: elif kind == EventKind.ERROR: query.state = ExecutionState.FAILED error = message.get("message") - query.handler(ExecutionResult(error=OperationalError(error))) + complete_query(ExecutionResult(error=OperationalError(error))) else: logging.warning("Received unknown %s event!", kind) diff --git a/wherobots/db/cursor.py b/wherobots/db/cursor.py index 51fea0f..af0a4e6 100644 --- a/wherobots/db/cursor.py +++ b/wherobots/db/cursor.py @@ -74,6 +74,8 @@ def __init__(self, exec_fn, cancel_fn) -> None: self.__queue: queue.Queue = queue.Queue() self.__results: list[Any] | None = None self.__store_result: StoreResult | None = None + self.__complete: bool = False + self.__error: Exception | None = None self.__current_execution_id: str | None = None self.__current_row: int = 0 @@ -93,21 +95,49 @@ def description(self) -> List[Tuple] | None: def rowcount(self) -> int: return self.__rowcount - def __on_execution_result(self, result) -> None: - self.__queue.put(result) + def __in_flight_execution_id(self) -> str | None: + """The current execution's id if its result has not yet arrived. + + Once a terminal result has been fetched (``__complete``) or is + waiting in the queue, the execution is finished and must not be + cancelled. + + Results are delivered from the connection's reader thread, so one + may arrive between the ``empty()`` check and a cancellation that + follows it. That race is benign: the cancel targets an execution + that already completed — a server-side no-op, same as when a query + finishes while a legitimate cancel is in flight — and the delivered + result sits in this execution's own queue, so it can never be + observed by a later execution's fetches. + """ + if ( + self.__current_execution_id is not None + and not self.__complete + and self.__queue.empty() + ): + return self.__current_execution_id + return None def __get_results(self) -> List[Tuple[Any, ...]] | None: if not self.__current_execution_id: raise ProgrammingError("No query has been executed yet") - if self.__results is not None: + if self.__complete: + if self.__error: + raise self.__error return self.__results execution_result = self.__queue.get() if not isinstance(execution_result, ExecutionResult): raise ProgrammingError("Unexpected result type") + # Whatever the outcome — rows, store export, empty result, or error — + # the execution has reached a terminal state. ``__results`` alone + # cannot signal this: store-backed and empty executions never set it. + self.__complete = True + if execution_result.error: - raise execution_result.error + self.__error = execution_result.error + raise self.__error self.__store_result = execution_result.store_result results = execution_result.results @@ -140,18 +170,25 @@ def execute( parameters: Dict[str, Any] | None = None, store: Store | None = None, ) -> None: - if self.__current_execution_id: - self.__cancel_fn(self.__current_execution_id) - + in_flight = self.__in_flight_execution_id() + if in_flight: + self.__cancel_fn(in_flight) + + # Each execution gets its own queue, and the handler closes over it: + # a late result from a superseded execution lands in the orphaned + # queue and can never be observed by fetches of the current one. + self.__queue = queue.Queue() self.__results = None self.__store_result = None + self.__complete = False + self.__error = None self.__current_row = 0 self.__rowcount = -1 self.__description = None self.__current_execution_id = self.__exec_fn( _substitute_parameters(operation, parameters), - self.__on_execution_result, + self.__queue.put, store, ) @@ -193,8 +230,9 @@ def fetchall(self) -> List[Any]: def close(self) -> None: """Close the cursor.""" - if self.__results is None and self.__current_execution_id: - self.__cancel_fn(self.__current_execution_id) + in_flight = self.__in_flight_execution_id() + if in_flight: + self.__cancel_fn(in_flight) def __iter__(self): return self