diff --git a/README.md b/README.md index 8e89858..7d0b4f3 100644 --- a/README.md +++ b/README.md @@ -243,6 +243,13 @@ including whether a write committed, is unknown. The error's (When `execute()` itself raises `ConnectionLostError`, the connection was already lost and the statement was never sent.) +To restore a known execution on a new connection, connect to its original SQL +session and call `cursor.attach(execution_id, store=original_store)`. Fetch or +cancel it through that cursor as usual. Attaching sends no SQL, never creates +an execution, and does not extend server result retention. It raises the same +missing/cancelled execution errors described above. A second cursor cannot +attach an execution already tracked by that connection. + ### Stopping the SQL session Closing a connection leaves its SQL session running until diff --git a/tests/test_attach.py b/tests/test_attach.py index 10df3fa..a7f38fa 100644 --- a/tests/test_attach.py +++ b/tests/test_attach.py @@ -16,6 +16,7 @@ ExecutionCancelledError, ExecutionNotFoundError, OperationalError, + ProgrammingError, ) from wherobots.db.connection import Connection from wherobots.db.models import Store, StoreResult @@ -828,3 +829,261 @@ def test_transport_lost_again_without_reconnecting_raises_connection_lost(): assert error.execution_id == EID finally: conn.close() + + +# Public reattachment from a new client process. + + +def test_new_cursor_attaches_completed_rows_without_executing_sql(): + session = Session( + on_retrieve=lambda s, m: s.incoming.put( + cbor2.dumps( + { + "kind": "execution_result", + "execution_id": m["execution_id"], + "state": "succeeded", + "results": arrow_results(x=[1, 2]), + "attached": True, + } + ) + ) + ) + conn = open_connection(session, results_format=ResultsFormat.ARROW) + try: + cursor = conn.cursor() + cursor.attach(EID) + assert cursor.execution_id == EID + assert outcome(cursor).error is None + assert cursor.fetchall()["x"].tolist() == [1, 2] + assert session.retrieves() == [ + { + "kind": "retrieve_results", + "execution_id": EID, + "attach": True, + "format": "arrow", + } + ] + assert session.requests("execute_sql") == [] + assert tracked(conn, EID) is None + finally: + conn.close() + + +def test_new_cursor_attaches_running_stored_query(): + session = new_session("running") + conn = open_connection(session) + try: + cursor = conn.cursor() + cursor.attach(EID, store=STORE) + assert wait_until(lambda: len(session.retrieves()) == 1) + session.put( + kind="state_updated", + execution_id=EID, + state="succeeded", + result_uri=RESULT_URI, + size=42, + ) + assert outcome(cursor).store_result == StoreResult(RESULT_URI, size=42) + assert cursor.get_store_result().result_uri == RESULT_URI + assert session.requests("execute_sql") == [] + finally: + conn.close() + + +def test_new_cursor_missing_execution_is_not_resubmitted(): + session = Session( + on_retrieve=lambda s, m: s.put( + kind="error", + execution_id=m["execution_id"], + message=NOT_FOUND, + ) + ) + conn = open_connection(session) + try: + cursor = conn.cursor() + cursor.attach(EID) + assert isinstance(outcome(cursor).error, ExecutionNotFoundError) + with pytest.raises(ExecutionNotFoundError): + cursor.fetchall() + assert session.requests("execute_sql") == [] + finally: + conn.close() + + +def test_new_cursor_can_cancel_the_attached_execution(): + session = new_session("running") + conn = open_connection(session) + try: + cursor = conn.cursor() + cursor.attach(EID, store=STORE) + assert wait_until(lambda: tracked(conn, EID).state == ExecutionState.RUNNING) + cursor.close() + assert [m["execution_id"] for _, m in session.requests("cancel")] == [EID] + session.put(kind="state_updated", execution_id=EID, state="cancelled") + assert outcome(cursor).error is None + assert session.requests("execute_sql") == [] + finally: + conn.close() + + +def test_new_cursor_does_not_report_previously_cancelled_write_as_success(): + conn = open_connection(new_session("cancelled")) + try: + cursor = conn.cursor() + cursor.attach(EID, store=STORE) + assert_cancelled_elsewhere(outcome(cursor).error) + finally: + conn.close() + + +def test_duplicate_attach_preserves_the_original_cursor(): + session = new_session("running") + conn = open_connection(session) + try: + original = conn.cursor() + original.attach(EID) + assert wait_until(lambda: tracked(conn, EID).state == ExecutionState.RUNNING) + with pytest.raises(ProgrammingError, match="already tracked"): + conn.cursor().attach(EID) + with pytest.raises(ProgrammingError, match="already tracks"): + original.attach(EID) + assert session.requests("cancel") == [] + session.incoming.put( + cbor2.dumps( + { + "kind": "execution_result", + "execution_id": EID, + "state": "succeeded", + "results": arrow_results(x=[7]), + } + ) + ) + assert outcome(original).error is None + assert original.fetchall()["x"].tolist() == [7] + assert len(session.retrieves()) == 1 + finally: + conn.close() + + +@pytest.mark.parametrize("completed", [False, True]) +def test_duplicate_attach_preserves_the_replaced_execution(completed): + session = new_session("running") + conn = open_connection(session) + try: + original = conn.cursor() + original.attach("execution-a") + other = conn.cursor() + other.attach("execution-b") + assert wait_until(lambda: len(session.retrieves()) == 2) + + def finish_original(): + session.incoming.put( + cbor2.dumps( + { + "kind": "execution_result", + "execution_id": "execution-a", + "state": "succeeded", + "results": arrow_results(x=[7]), + } + ) + ) + + if completed: + finish_original() + assert outcome(original).error is None + assert original.fetchall()["x"].tolist() == [7] + + with pytest.raises(ProgrammingError, match="already tracked"): + original.attach("execution-b") + + assert original.execution_id == "execution-a" + assert session.requests("cancel") == [] + assert len(session.retrieves()) == 2 + if not completed: + finish_original() + assert outcome(original).error is None + assert original.fetchall()["x"].tolist() == [7] + finally: + conn.close() + + +def test_successful_attach_replaces_the_pending_execution(): + def on_retrieve(session, message): + if message["execution_id"] == "execution-a": + answer(session, message, "running", attached=True) + else: + session.incoming.put( + cbor2.dumps( + { + "kind": "execution_result", + "execution_id": "execution-b", + "state": "succeeded", + "results": arrow_results(x=[8]), + "attached": True, + } + ) + ) + + session = Session(on_retrieve=on_retrieve) + conn = open_connection(session) + try: + cursor = conn.cursor() + cursor.attach("execution-a") + cursor.attach("execution-b") + assert outcome(cursor).error is None + assert cursor.execution_id == "execution-b" + assert cursor.fetchall()["x"].tolist() == [8] + assert [m["execution_id"] for _, m in session.requests("cancel")] == [ + "execution-a" + ] + finally: + conn.close() + + +def test_new_cursor_attach_on_old_session_polls_without_replaying(): + def on_retrieve(s, m): + if len(s.retrieves()) == 1: + answer(s, m, "running") + else: + answer(s, m, "succeeded", result_uri=RESULT_URI, size=3) + + session = Session(on_retrieve=on_retrieve) + conn = open_connection(session) + try: + cursor = conn.cursor() + cursor.attach(EID, store=STORE) + assert outcome(cursor).store_result.result_uri == RESULT_URI + assert len(session.retrieves()) >= 2 + assert session.requests("execute_sql") == [] + finally: + conn.close() + + +def test_new_cursor_attach_survives_another_transport_loss(): + first = new_session("running") + second = new_session("succeeded", result_uri=RESULT_URI, size=3) + conn = open_connection(first, transport_factory=lambda timeout: second) + try: + cursor = conn.cursor() + cursor.attach(EID, store=STORE) + assert wait_until(lambda: tracked(conn, EID).state == ExecutionState.RUNNING) + first.incoming.put(ConnectionClosedError(None, None)) + assert outcome(cursor).store_result.result_uri == RESULT_URI + assert first.requests("execute_sql") == second.requests("execute_sql") == [] + finally: + conn.close() + + +@pytest.mark.parametrize("execution_id", [None, "", " ", 12]) +def test_attach_rejects_invalid_ids_without_touching_an_existing_cursor(execution_id): + session = new_session("running") + conn = open_connection(session) + try: + cursor = conn.cursor() + cursor.attach(EID) + with pytest.raises(ProgrammingError, match="non-empty string"): + cursor.attach(execution_id) + assert cursor.execution_id == EID + assert session.requests("cancel") == [] + finally: + conn.close() diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index d82c86b..ae78bc4 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -31,6 +31,7 @@ InterfaceError, NotSupportedError, OperationalError, + ProgrammingError, ) from .models import ExecutionResult, ProgressInfo, Store, StoreResult from .types import ( @@ -473,7 +474,7 @@ def rollback(self) -> None: raise NotSupportedError def cursor(self) -> Cursor: - return Cursor(self.__execute_sql, self.__cancel_query) + return Cursor(self.__execute_sql, self.__cancel_query, self.__attach_query) def set_progress_handler(self, handler: ProgressHandler | None) -> None: """Register a callback invoked for execution progress events. @@ -1255,6 +1256,11 @@ def __send( raise self.__connection_error(query.execution_id) return True if query is not None: + if query.execution_id in self.__queries: + raise ProgrammingError( + f"Execution {query.execution_id} is already tracked " + "by this connection" + ) # The watchdog's clock starts at registration. watch = query.watch watch.last_activity = watch.last_event = time.monotonic() @@ -1373,6 +1379,25 @@ def __execute_sql( ) return execution_id + def __attach_query( + self, + execution_id: str, + handler: Callable[[Any], None], + store: Store | None = None, + ) -> str: + self.__send( + self.__retrieve_request(execution_id, attach=True), + Query( + sql="", + execution_id=execution_id, + state=ExecutionState.RESULTS_REQUESTED, + handler=handler, + store=store, + attach=_Attach(), + ), + ) + return execution_id + def __retrieve_request( self, execution_id: str, attach: bool = False ) -> Dict[str, Any]: diff --git a/wherobots/db/cursor.py b/wherobots/db/cursor.py index 30b8e16..93d80cf 100644 --- a/wherobots/db/cursor.py +++ b/wherobots/db/cursor.py @@ -1,9 +1,9 @@ import math import queue import re -from typing import Any, List, Tuple, Dict +from typing import Any, Callable, List, Tuple, Dict -from .errors import ProgrammingError +from .errors import NotSupportedError, ProgrammingError from .models import ExecutionResult, Store, StoreResult # Matches pyformat parameter markers: %(name)s @@ -67,9 +67,16 @@ def replacer(match: re.Match) -> str: class Cursor: - def __init__(self, exec_fn, cancel_fn) -> None: + def __init__( + self, + exec_fn, + cancel_fn, + attach_fn: Callable[[str, Callable[[ExecutionResult], None], Store | None], str] + | None = None, + ) -> None: self.__exec_fn = exec_fn self.__cancel_fn = cancel_fn + self.__attach_fn = attach_fn self.__queue: queue.Queue = queue.Queue() self.__results: list[Any] | None = None @@ -180,6 +187,34 @@ def execute( parameters: Dict[str, Any] | None = None, store: Store | None = None, ) -> None: + self.__reset_execution() + self.__current_execution_id = self.__exec_fn( + _substitute_parameters(operation, parameters), + self.__queue.put, + store, + ) + + def attach(self, execution_id: str, store: Store | None = None) -> None: + """Retrieve an existing execution on this connection's SQL session. + + Connect to the original SQL session first. Pass the original Store for + stored results. This never submits SQL or extends server result retention; + a missing execution raises ExecutionNotFoundError when fetching. + """ + if not isinstance(execution_id, str) or not execution_id.strip(): + raise ProgrammingError("execution_id must be a non-empty string") + if self.__attach_fn is None: + raise NotSupportedError("This cursor does not support attaching executions") + if self.__in_flight_execution_id() == execution_id: + raise ProgrammingError("This cursor already tracks the execution") + # Preserve the current execution until the new attachment is accepted. + results: queue.Queue[ExecutionResult] = queue.Queue() + attached_id = self.__attach_fn(execution_id, results.put, store) + self.__reset_execution() + self.__queue = results + self.__current_execution_id = attached_id + + def __reset_execution(self) -> None: in_flight = self.__in_flight_execution_id() if in_flight: self.__cancel_fn(in_flight) @@ -199,11 +234,6 @@ def execute( # 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, - store, - ) def get_store_result(self) -> StoreResult | None: """Get the store result for the last executed query.