fix(client): begin on real grpc.aio stubs, settle interrupts, confirm unknown-id cleanup

Four fixes, each with a regression test seen failing first. The
shielded begin RPC is now wrapped in a coroutine: a real grpc.aio stub
returns a UnaryUnaryCall awaitable, not a coroutine object, and
create_task() rejected it with TypeError on every real
begin_transaction() — AsyncMock had hidden this, so the new test wires
a call-like awaitable the way the transport does. A KeyboardInterrupt
or SystemExit raised from inside an awaited statement or commit now
settles the handle (aborted with detached cleanup, or indeterminate)
instead of parking it in-flight forever. And a cleanup rollback the
server answers with NOT_FOUND (unknown transaction id) now reads as
confirmed — nothing is held under that id — so no redundant retries
follow and an explicit rollback() can settle the handle, while lost
requests stay retriable.

Author: polaz

coordinode/coordinode/client.py

@@ -399,7 +399,7 @@ async def _best_effort_rollback(self) -> None:
         )
 
         self._cleanup_confirmed = False
-        with suppress(Exception):
+        try:
             await self._client._cypher_stub.RollbackTransaction(
                 RollbackTransactionRequest(transaction_id=self._id),
                 # Capped by the client's own timeout too: a caller who
@@ -408,6 +408,16 @@ async def _best_effort_rollback(self) -> None:
                 timeout=min(_CLEANUP_TIMEOUT_SECS, self._client._timeout),
             )
             self._cleanup_confirmed = True
+        except grpc.RpcError as exc:
+            # NOT_FOUND is the server's "unknown transaction id": a definitive
+            # statement that nothing is held under this id any more — as
+            # settled as a successful rollback. Every other failure (a lost
+            # request most of all) leaves the cleanup unconfirmed, retriable.
+            with suppress(Exception):
+                if exc.code() == grpc.StatusCode.NOT_FOUND:
+                    self._cleanup_confirmed = True
+        except Exception:
+            pass
 
     async def cypher(
         self,
@@ -489,6 +499,16 @@ async def cypher(
             self._state = "aborted"
             self._spawn_cleanup()
             raise
+        except BaseException:
+            # KeyboardInterrupt / SystemExit raised from inside the awaited
+            # call (or any failure that is neither gRPC nor cancellation)
+            # matches none of the handlers above and would park the handle
+            # in "executing" forever. Settle conservatively: no commit was
+            # sent, so nothing can apply — close as aborted with a detached
+            # bounded cleanup, and let the exception propagate.
+            self._state = "aborted"
+            self._spawn_cleanup()
+            raise
         if self._abandoned:
             # The owning context exited while this statement was still in
             # flight: nobody is left to commit or roll back, so the buffered
@@ -560,6 +580,16 @@ async def commit(self) -> int:
                 # indeterminate cleanup, so it goes detached from here.
                 self._spawn_cleanup()
             raise
+        except BaseException:
+            # KeyboardInterrupt / SystemExit raised from inside the awaited
+            # call matches neither handler above and would park the handle in
+            # "committing" forever. The request may already have applied, so
+            # the only honest settlement is indeterminate — as the sync
+            # wrapper already does for interrupts at the loop boundary.
+            self._state = "indeterminate"
+            if self._abandoned:
+                self._spawn_cleanup()
+            raise
         self._state = "committed"
         return int(resp.applied_index)
 
@@ -924,9 +954,13 @@ async def begin_transaction(self) -> AsyncTransaction:
         # no handle, nothing for close() to drain, a pinned snapshot until
         # the idle sweep. On cancellation the task keeps running detached and
         # its late reply is handed straight to a rollback.
-        begin = asyncio.get_running_loop().create_task(
-            self._cypher_stub.BeginTransaction(BeginTransactionRequest(), timeout=self._timeout)
-        )
+        # Wrapped in a coroutine: a real grpc.aio stub returns a
+        # UnaryUnaryCall — an awaitable, NOT a coroutine object — and
+        # create_task() accepts only the latter.
+        async def _begin() -> Any:
+            return await self._cypher_stub.BeginTransaction(BeginTransactionRequest(), timeout=self._timeout)
+
+        begin = asyncio.get_running_loop().create_task(_begin())
         try:
             resp = await asyncio.shield(begin)
         except asyncio.CancelledError:

tests/unit/test_transactions.py

@@ -2455,3 +2455,104 @@ async def rollback(req, timeout=None):
             assert calls["n"] == 3, "a cancelled aborted-branch retry got no detached cleanup"
 
         asyncio.run(_inner())
+
+
+class TestBeginWorksWithGrpcStyleAwaitables:
+    """A real grpc.aio unary stub returns a UnaryUnaryCall — an AWAITABLE,
+    not a coroutine object — and loop.create_task() accepts only coroutines.
+    AsyncMock hides this by returning coroutines, so this test wires a
+    call-like awaitable the way the real transport does."""
+
+    def test_begin_accepts_a_non_coroutine_awaitable(self):
+        from unittest.mock import MagicMock
+
+        class _CallLike:
+            def __init__(self, resp):
+                self._resp = resp
+
+            def __await__(self):
+                async def _deliver():
+                    return self._resp
+
+                return _deliver().__await__()
+
+        async def _inner() -> None:
+            client = _async_client(
+                BeginTransaction=MagicMock(
+                    side_effect=lambda req, timeout=None: _CallLike(
+                        cypher_pb2.BeginTransactionResponse(transaction_id=42)
+                    )
+                )
+            )
+            tx = await client.begin_transaction()
+            assert tx.transaction_id == 42
+
+        asyncio.run(_inner())
+
+
+class TestAsyncStatementSettlesAfterProcessControlException:
+    """KeyboardInterrupt or SystemExit raised from inside the awaited
+    statement matches neither the gRPC nor the cancellation handler; the
+    handle must still close conservatively (aborted, cleanup unconfirmed)
+    instead of staying "executing" forever."""
+
+    def test_interrupted_statement_closes_the_handle(self):
+        from unittest.mock import AsyncMock
+
+        async def _inner() -> None:
+            client = _async_client(ExecuteCypher=AsyncMock(side_effect=KeyboardInterrupt))
+            tx = await client.begin_transaction()
+            with pytest.raises(KeyboardInterrupt):
+                await tx.cypher("CREATE (:A)")
+            assert tx._state == "aborted", "the interruption left the handle parked in-flight"
+            await tx.rollback()  # must be able to send the cleanup
+            await client.close()
+            assert client._cypher_stub.RollbackTransaction.await_count >= 1
+
+        asyncio.run(_inner())
+
+
+class TestAsyncCommitSettlesAfterProcessControlException:
+    """KeyboardInterrupt or SystemExit raised from inside the awaited commit
+    leaves the outcome unknowable — the request may already have applied —
+    so the handle must settle as indeterminate, not stay "committing"."""
+
+    def test_interrupted_commit_is_indeterminate(self):
+        from unittest.mock import AsyncMock
+
+        async def _inner() -> None:
+            client = _async_client(CommitTransaction=AsyncMock(side_effect=KeyboardInterrupt))
+            tx = await client.begin_transaction()
+            with pytest.raises(KeyboardInterrupt):
+                await tx.commit()
+            assert tx._state == "indeterminate"
+            with pytest.raises(RuntimeError, match="outcome is unknown"):
+                await tx.cypher("RETURN 1")
+
+        asyncio.run(_inner())
+
+
+class TestUnknownTransactionAnswerConfirmsCleanup:
+    """The server answering a cleanup rollback with NOT_FOUND ("unknown
+    transaction id") is a definitive statement that nothing is held: the
+    cleanup must read as confirmed, so no redundant retries follow and an
+    explicit rollback() can settle the handle."""
+
+    def test_not_found_confirms_the_cleanup(self):
+        from contextlib import suppress as ctx_suppress
+        from unittest.mock import AsyncMock
+
+        async def _inner() -> None:
+            client = _async_client(
+                ExecuteCypher=AsyncMock(side_effect=_ServerRejected()),
+                RollbackTransaction=AsyncMock(side_effect=_TransportError(grpc.StatusCode.NOT_FOUND)),
+            )
+            tx = await client.begin_transaction()
+            with ctx_suppress(grpc.RpcError):
+                await tx.cypher("CREATE (:A)")  # aborts; cleanup answered NOT_FOUND
+            assert tx._cleanup_confirmed is True, "a definitive unknown-id answer read as a lost request"
+            await tx.rollback()
+            assert tx._state == "rolled_back"
+            assert client._cypher_stub.RollbackTransaction.await_count == 1, "redundant cleanup retry"
+
+        asyncio.run(_inner())