fix: isolate cursor results per execution and stop cancelling completed queries (WBC-922)

Author: james-willis

tests/test_cursor.py

@@ -6,15 +6,19 @@
 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.models import ExecutionResult
 
 
 def _make_cursor():
@@ -211,6 +215,114 @@ 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]
+
+
+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()
+
+
 # ---------------------------------------------------------------------------
 # _substitute_parameters unit tests
 # ---------------------------------------------------------------------------

wherobots/db/cursor.py

@@ -93,8 +93,19 @@ 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 result has been fetched (or is waiting in the queue), the
+        execution is complete and must not be cancelled.
+        """
+        if (
+            self.__current_execution_id is not None
+            and self.__results is None
+            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:
@@ -140,9 +151,14 @@ 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.__current_row = 0
@@ -151,7 +167,7 @@ def execute(
 
         self.__current_execution_id = self.__exec_fn(
             _substitute_parameters(operation, parameters),
-            self.__on_execution_result,
+            self.__queue.put,
             store,
         )
 
@@ -193,8 +209,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