Adjusted AsyncPersistentSubscription to support task cancellation, and adjusted examples.

johnbywater · johnbywater · commit e4b1db0a1b09 · 2026-05-08T21:57:42.000+01:00
diff --git a/README.md b/README.md
@@ -2514,7 +2514,10 @@ giving `RecordedEvent` objects. It also has `ack()`, `nack()` and `stop()`
 methods.
 
 ```python
-subscription = client.read_subscription_to_all(group_name=group_name1)
+with client.read_subscription_to_all(
+    group_name=group_name1
+) as subscription:
+    ...
 ```
 
 The `ack()` method should be used by a consumer to "acknowledge" to the server that
@@ -2536,15 +2539,18 @@ examples below.
 ```python
 received_events = []
 
-for event in subscription:
-    received_events.append(event)
+with client.read_subscription_to_all(
+    group_name=group_name1
+) as subscription:
+    for event in subscription:
+        received_events.append(event)
 
-    # Acknowledge the received event.
-    subscription.ack(event)
+        # Acknowledge the received event.
+        subscription.ack(event)
 
-    # Stop when 'event9' has been received.
-    if event == event9:
-        subscription.stop()
+        # Stop when 'event9' has been received.
+        if event == event9:
+            subscription.stop()
 ```
 
 The `nack()` should be used by a consumer to "negatively acknowledge" to the server that
@@ -2802,28 +2808,25 @@ This method returns a `PersistentSubscription` object, which is an iterator
 giving `RecordedEvent` objects, that also has `ack()`, `nack()` and `stop()`
 methods.
 
-```python
-subscription = client.read_subscription_to_stream(
-    group_name=group_name2,
-    stream_name=stream_name2,
-)
-```
-
 The example below iterates over the subscription object, and calls `ack()`.
 The subscription's `stop()` method is called when we have received `event6`,
 stopping the iteration, so that we can continue with the examples below.
 
 ```python
 events = []
-for event in subscription:
-    events.append(event)
+with client.read_subscription_to_stream(
+    group_name=group_name2,
+    stream_name=stream_name2,
+) as subscription:
+    for event in subscription:
+        events.append(event)
 
-    # Acknowledge the received event.
-    subscription.ack(event)
+        # Acknowledge the received event.
+        subscription.ack(event)
 
-    # Stop when 'event6' has been received.
-    if event == event6:
-        subscription.stop()
+        # Stop when 'event6' has been received.
+        if event == event6:
+            break
 ```
 
 We can check we received all the events that were appended to `stream_name2`
diff --git a/kurrentdbclient/persistent.py b/kurrentdbclient/persistent.py
@@ -501,15 +501,51 @@ def __init__(
         AbstractAsyncPersistentSubscription.__init__(self)
         self._read_reqs = read_reqs
         self._stream_stream_call = stream_stream_call
-        self._stream_stream_call_iter = stream_stream_call.__aiter__()
         self._expected_group_name = expected_group_name
         self._stream_name = stream_name
         self._is_stopping = False
+        self._read_resp_iter = stream_stream_call.__aiter__()
+        self._read_resp_queue = asyncio.Queue[persistent_pb2.ReadResp | None](
+            maxsize=10
+        )
+        self._read_resp_stream_worker_co = self._read_resp_stream_worker()
+        self._read_resp_stream_worker_task = asyncio.create_task(
+            self._read_resp_stream_worker_co
+        )
+
+    async def _read_resp_stream_worker(self) -> None:
+        """Coroutine for isolated task that only handles the gRPC stream."""
+        try:
+            while True:
+                read_resp = await self._read_resp_iter.__anext__()
+                await self._read_resp_queue.put(read_resp)
+                if self._has_iter_error_for_testing():
+                    raise AioRpcError(
+                        grpc.StatusCode.INTERNAL,
+                        grpc.aio.Metadata(),
+                        grpc.aio.Metadata(),
+                        "",
+                        "",
+                    )
+        except BaseException:
+            # Drain the queue (to avoid blocking on put).
+            try:
+                while True:
+                    self._read_resp_queue.get_nowait()
+            except asyncio.queues.QueueEmpty:
+                pass
+            # Unblock waiting on get().
+            await self._read_resp_queue.put(None)
+            # Reraise the error (appears at 'await task').
+            raise
 
     async def init(self) -> None:
         try:
-            first_read_resp = await self._get_next_read_resp()
-            if first_read_resp.WhichOneof("content") == "subscription_confirmation":
+            first_read_resp = await self._read_resp_queue.get()
+            if (
+                first_read_resp is not None
+                and first_read_resp.WhichOneof("content") == "subscription_confirmation"
+            ):
                 expected_stream_name = (
                     self._stream_name if self._stream_name is not None else "$all"
                 )
@@ -522,14 +558,20 @@ async def init(self) -> None:
                 if (
                     confirmed_group_name != self._expected_group_name
                     or confirmed_stream_name != expected_stream_name
-                ):  # pragma: no cover
-                    raise SubscriptionConfirmationError
+                ):
+                    msg = (
+                        f"Expected group name: {self._expected_group_name}, "
+                        f"confirmed group name: {confirmed_group_name}, "
+                        f"expected stream name: {expected_stream_name}, "
+                        f"confirmed stream name: {expected_stream_name}"
+                    )
+                    raise SubscriptionConfirmationError(msg)
                 self._subscription_id = subscription_id
                 self._read_reqs.subscription_id = subscription_id.encode()
 
-            else:  # pragma: no cover
+            else:
                 msg = f"Expected subscription confirmation, got: {first_read_resp}"
-                raise KurrentDBClientError(msg)
+                raise SubscriptionConfirmationError(msg)
         except BaseException:
             await self.stop(wait_until_stopped=False)
             raise
@@ -539,11 +581,13 @@ def subscription_id(self) -> str:
         return self._subscription_id
 
     async def __anext__(self) -> RecordedEvent:
-        if self._is_context_manager_active and self._is_stopping:
-            raise StopAsyncIteration
         try:
             while True:
-                read_resp = await self._get_next_read_resp()
+                if self._is_context_manager_active and self._is_stopping:
+                    raise StopAsyncIteration
+                read_resp = await self._read_resp_queue.get()
+                if read_resp is None:
+                    raise StopAsyncIteration
                 content_oneof = read_resp.WhichOneof("content")
                 if content_oneof == "event":
                     recorded_event = construct_recorded_event(read_resp.event)
@@ -560,37 +604,23 @@ async def __anext__(self) -> RecordedEvent:
             await self.stop(wait_until_stopped=False)
             raise
 
-    async def _get_next_read_resp(self) -> persistent_pb2.ReadResp:
-        try:
-            response = await self._stream_stream_call_iter.__anext__()
-            if self._has_iter_error_for_testing():
-                raise AioRpcError(
-                    grpc.StatusCode.INTERNAL,
-                    grpc.aio.Metadata(),
-                    grpc.aio.Metadata(),
-                    "",
-                    "",
-                )
-
-        except asyncio.CancelledError as e:
-            await self.stop(wait_until_stopped=False)
-            if self._read_reqs.errored:
-                raise ExceptionIteratingRequestsError from self._read_reqs.errored
-            raise StopAsyncIteration from e
-        except grpc.aio.AioRpcError as e:
-            await self.stop(wait_until_stopped=False)
-            raise handle_rpc_error(e) from None
-        else:
-            return response
-
     async def stop(self, *, wait_until_stopped: bool = True) -> None:
         if self._is_context_manager_active:
             self._is_stopping = True
         elif not await self._set_is_stopped():
             await self._read_reqs.stop(wait_until_stopped=wait_until_stopped)
+            self._read_resp_stream_worker_task.cancel()
             self._stream_stream_call.cancel()
             await asyncio.sleep(0.05)
             self._grpc_streamers.remove(self)
+            try:
+                await self._read_resp_stream_worker_task
+            except asyncio.CancelledError:
+                pass
+            except grpc.RpcError as e:
+                raise handle_rpc_error(e) from e
+            if self._read_reqs.errored:
+                raise ExceptionIteratingRequestsError from self._read_reqs.errored
 
     async def ack(self, item: UUID | RecordedEvent) -> None:
         await self._read_reqs.ack(event_id=self._get_event_id(item))
diff --git a/kurrentdbclient/streams.py b/kurrentdbclient/streams.py
@@ -174,10 +174,12 @@ def __init__(
         self._unary_stream_call = unary_stream_call
         self._read_resp_iter = unary_stream_call.__aiter__()
         self._read_resp_queue = asyncio.Queue[streams_pb2.ReadResp | None](maxsize=10)
-        self._stream_worker_co = self._stream_worker()
-        self._stream_worker_task = asyncio.create_task(self._stream_worker_co)
+        self._read_resp_stream_worker_co = self._read_resp_stream_worker()
+        self._read_resp_stream_worker_task = asyncio.create_task(
+            self._read_resp_stream_worker_co
+        )
 
-    async def _stream_worker(self) -> None:
+    async def _read_resp_stream_worker(self) -> None:
         """Coroutine for isolated task that only handles the gRPC stream."""
         try:
             while True:
@@ -192,21 +194,25 @@ async def _stream_worker(self) -> None:
                         "",
                     )
         except StopAsyncIteration:
+            # End of the stream, signal to end the iteration.
             await self._read_resp_queue.put(None)
         except BaseException:
+            # Drain the queue (to avoid blocking on put).
             try:
                 while True:
                     self._read_resp_queue.get_nowait()
             except asyncio.queues.QueueEmpty:
                 pass
+            # Unblock waiting on get().
             await self._read_resp_queue.put(None)
+            # Reraise the error (appears at 'await task').
             raise
 
     async def __anext__(self) -> RecordedEvent:
         try:
             while True:
-                # if self._is_context_manager_active and self._is_stopping:
-                #     raise StopAsyncIteration
+                if self._is_context_manager_active and self._is_stopping:
+                    raise StopAsyncIteration
                 read_resp = await self._read_resp_queue.get()
                 if read_resp is None:
                     raise StopAsyncIteration
@@ -223,12 +229,11 @@ async def stop(self, *, wait_until_stopped: bool = True) -> None:
         if self._is_context_manager_active:
             self._is_stopping = True
         elif not await self._set_is_stopped():
-            self._stream_worker_task.cancel()
+            self._read_resp_stream_worker_task.cancel()
             self._unary_stream_call.cancel()
             self._grpc_streamers.remove(self)
-
             try:
-                await self._stream_worker_task
+                await self._read_resp_stream_worker_task
             except asyncio.CancelledError:
                 pass
             except grpc.RpcError as e:
diff --git a/tests/test_client.py b/tests/test_client.py
@@ -21,6 +21,7 @@
 from kurrentdbclient import (
     DEFAULT_EXCLUDE_FILTER,
     KDB_SYSTEM_EVENTS_REGEX,
+    CatchupSubscription,
     RecordedEvent,
     StreamState,
 )
@@ -3057,6 +3058,19 @@ def test_subscribe_to_all_can_be_stopped(self) -> None:
         # Iterating should stop.
         list(subscription)
 
+        # Exiting the context manager should stop the subscription.
+        subscription = self.client.subscribe_to_all()
+        with subscription:
+            pass
+        self.assertTrue(cast(CatchupSubscription, subscription)._is_stopped)
+
+        # Calling stop inside the context manager should terminate the iteration.
+        subscription = self.client.subscribe_to_all()
+        with subscription:
+            subscription.stop()
+            list(subscription)
+        self.assertTrue(cast(CatchupSubscription, subscription)._is_stopped)
+
     def _test_subscribe_to_all_raises_consumer_too_slow(self) -> None:
         # Todo: The server behaviour is too unreliable to run this test.
         self.construct_client()
@@ -3720,7 +3734,7 @@ def test_subscription_to_all_read_with_nack_stop(self) -> None:
         assert events[-2].data == event2.data
         assert events[-1].data == event3.data
 
-    def test_persistent_subscription_ack_after_stop(self) -> None:
+    def test_subscription_to_all_ack_after_stop(self) -> None:
         self.construct_client()
 
         group_name = str(uuid4())
diff --git a/tests/test_client_async.py b/tests/test_client_async.py
diff --git a/tests/test_docs.py b/tests/test_docs.py
