Add Client.subscribe_with_manual_ack() (#105)

vrslev · web-flow · commit 0b665751019e · 2025-02-03T17:54:12.000+03:00
diff --git a/README.md b/README.md
@@ -108,6 +108,20 @@ You can pass custom headers to `client.subscribe()`:
 await client.subscribe("DLQ", handle_message_from_dlq, ack="client", headers={"selector": "location = 'Europe'"}, on_suppressed_exception=print)
 ```
 
+#### Handling ACK/NACKs yourself
+
+If you want to send ACK and NACK frames yourself, you can use `client.subscribe_with_manual_ack()`:
+
+```python
+async def handle_message_from_dlq(message_frame: stompman.AckableMessageFrame) -> None:
+    print(message_frame.body)
+    await message_frame.ack()
+
+await client.subscribe_with_manual_ack("DLQ", handle_message_from_dlq, ack="client")
+```
+
+Note that this way exceptions won't be suppressed automatically.
+
 ### Cleaning Up
 
 stompman takes care of cleaning up resources automatically. When you leave the context of async context managers `stompman.Client()`, or `client.begin()`, the necessary frames will be sent to the server.
diff --git a/packages/faststream-stomp/faststream_stomp/subscriber.py b/packages/faststream-stomp/faststream_stomp/subscriber.py
@@ -39,7 +39,7 @@ def __init__(
         self.headers = headers
         self.on_suppressed_exception = on_suppressed_exception
         self.suppressed_exception_classes = suppressed_exception_classes
-        self._subscription: stompman.Subscription | None = None
+        self._subscription: stompman.AutoAckSubscription | None = None
 
         super().__init__(
             no_ack=self.ack == "auto",
diff --git a/packages/stompman/stompman/__init__.py b/packages/stompman/stompman/__init__.py
@@ -31,16 +31,18 @@
     UnsubscribeFrame,
 )
 from stompman.serde import FrameParser, dump_frame
-from stompman.subscription import Subscription
+from stompman.subscription import AckableMessageFrame, AutoAckSubscription, ManualAckSubscription
 from stompman.transaction import Transaction
 
 __all__ = [
     "AbortFrame",
     "AckFrame",
     "AckMode",
+    "AckableMessageFrame",
     "AnyClientFrame",
     "AnyRealServerFrame",
     "AnyServerFrame",
+    "AutoAckSubscription",
     "BeginFrame",
     "Client",
     "CommitFrame",
@@ -57,13 +59,13 @@
     "FrameParser",
     "Heartbeat",
     "HeartbeatFrame",
+    "ManualAckSubscription",
     "MessageFrame",
     "NackFrame",
     "ReceiptFrame",
     "SendFrame",
     "StompProtocolConnectionIssue",
     "SubscribeFrame",
-    "Subscription",
     "Transaction",
     "UnsubscribeFrame",
     "UnsupportedProtocolVersion",
diff --git a/packages/stompman/stompman/client.py b/packages/stompman/stompman/client.py
@@ -1,6 +1,6 @@
 import asyncio
 import inspect
-from collections.abc import AsyncGenerator, Awaitable, Callable
+from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine
 from contextlib import AsyncExitStack, asynccontextmanager
 from dataclasses import dataclass, field
 from functools import partial
@@ -21,7 +21,7 @@
     ReceiptFrame,
     SendFrame,
 )
-from stompman.subscription import Subscription
+from stompman.subscription import AckableMessageFrame, ActiveSubscriptions, AutoAckSubscription, ManualAckSubscription
 from stompman.transaction import Transaction
 
 
@@ -47,7 +47,7 @@ class Client:
     connection_class: type[AbstractConnection] = Connection
 
     _connection_manager: ConnectionManager = field(init=False)
-    _active_subscriptions: dict[str, "Subscription"] = field(default_factory=dict, init=False)
+    _active_subscriptions: ActiveSubscriptions = field(default_factory=dict, init=False)
     _active_transactions: set[Transaction] = field(default_factory=set, init=False)
     _exit_stack: AsyncExitStack = field(default_factory=AsyncExitStack, init=False)
     _heartbeat_task: asyncio.Task[None] = field(init=False)
@@ -113,7 +113,15 @@ async def _listen_to_frames(self) -> None:
                 match frame:
                     case MessageFrame():
                         if subscription := self._active_subscriptions.get(frame.headers["subscription"]):
-                            task_group.create_task(subscription._run_handler(frame=frame))  # noqa: SLF001
+                            task_group.create_task(
+                                subscription._run_handler(frame=frame)  # noqa: SLF001
+                                if isinstance(subscription, AutoAckSubscription)
+                                else subscription.handler(
+                                    AckableMessageFrame(
+                                        headers=frame.headers, body=frame.body, _subscription=subscription
+                                    )
+                                )
+                            )
                     case ErrorFrame():
                         if self.on_error_frame:
                             self.on_error_frame(frame)
@@ -152,8 +160,8 @@ async def subscribe(
         headers: dict[str, str] | None = None,
         on_suppressed_exception: Callable[[Exception, MessageFrame], Any],
         suppressed_exception_classes: tuple[type[Exception], ...] = (Exception,),
-    ) -> "Subscription":
-        subscription = Subscription(
+    ) -> "AutoAckSubscription":
+        subscription = AutoAckSubscription(
             destination=destination,
             handler=handler,
             headers=headers,
@@ -165,3 +173,22 @@ async def subscribe(
         )
         await subscription._subscribe()  # noqa: SLF001
         return subscription
+
+    async def subscribe_with_manual_ack(
+        self,
+        destination: str,
+        handler: Callable[[AckableMessageFrame], Coroutine[Any, Any, Any]],
+        *,
+        ack: AckMode = "client-individual",
+        headers: dict[str, str] | None = None,
+    ) -> "ManualAckSubscription":
+        subscription = ManualAckSubscription(
+            destination=destination,
+            handler=handler,
+            headers=headers,
+            ack=ack,
+            _connection_manager=self._connection_manager,
+            _active_subscriptions=self._active_subscriptions,
+        )
+        await subscription._subscribe()  # noqa: SLF001
+        return subscription
diff --git a/packages/stompman/stompman/subscription.py b/packages/stompman/stompman/subscription.py
@@ -1,4 +1,4 @@
-from collections.abc import Awaitable, Callable
+from collections.abc import Awaitable, Callable, Coroutine
 from dataclasses import dataclass, field
 from typing import Any
 from uuid import uuid4
@@ -14,62 +14,79 @@
     UnsubscribeFrame,
 )
 
-ActiveSubscriptions = dict[str, "Subscription"]
+ActiveSubscriptions = dict[str, "AutoAckSubscription | ManualAckSubscription"]
 
 
 @dataclass(kw_only=True, slots=True)
-class Subscription:
+class BaseSubscription:
     id: str = field(default_factory=lambda: _make_subscription_id(), init=False)  # noqa: PLW0108
     destination: str
     headers: dict[str, str] | None
-    handler: Callable[[MessageFrame], Awaitable[Any]]
     ack: AckMode
-    on_suppressed_exception: Callable[[Exception, MessageFrame], Any]
-    suppressed_exception_classes: tuple[type[Exception], ...]
     _connection_manager: ConnectionManager
     _active_subscriptions: ActiveSubscriptions
 
-    _should_handle_ack_nack: bool = field(init=False)
-
-    def __post_init__(self) -> None:
-        self._should_handle_ack_nack = self.ack in {"client", "client-individual"}
-
     async def _subscribe(self) -> None:
         await self._connection_manager.write_frame_reconnecting(
             SubscribeFrame.build(
                 subscription_id=self.id, destination=self.destination, ack=self.ack, headers=self.headers
             )
         )
-        self._active_subscriptions[self.id] = self
+        self._active_subscriptions[self.id] = self  # type: ignore[assignment]
 
     async def unsubscribe(self) -> None:
         del self._active_subscriptions[self.id]
         await self._connection_manager.maybe_write_frame(UnsubscribeFrame(headers={"id": self.id}))
 
+    async def _nack(self, frame: MessageFrame) -> None:
+        if self.id in self._active_subscriptions and (ack_id := frame.headers.get("ack")):
+            await self._connection_manager.maybe_write_frame(
+                NackFrame(headers={"id": ack_id, "subscription": frame.headers["subscription"]})
+            )
+
+    async def _ack(self, frame: MessageFrame) -> None:
+        if self.id in self._active_subscriptions and (ack_id := frame.headers["ack"]):
+            await self._connection_manager.maybe_write_frame(
+                AckFrame(headers={"id": ack_id, "subscription": frame.headers["subscription"]})
+            )
+
+
+@dataclass(kw_only=True, slots=True)
+class AutoAckSubscription(BaseSubscription):
+    handler: Callable[[MessageFrame], Awaitable[Any]]
+    on_suppressed_exception: Callable[[Exception, MessageFrame], Any]
+    suppressed_exception_classes: tuple[type[Exception], ...]
+    _should_handle_ack_nack: bool = field(init=False)
+
+    def __post_init__(self) -> None:
+        self._should_handle_ack_nack = self.ack in {"client", "client-individual"}
+
     async def _run_handler(self, *, frame: MessageFrame) -> None:
         try:
             await self.handler(frame)
         except self.suppressed_exception_classes as exception:
-            if (
-                self._should_handle_ack_nack
-                and self.id in self._active_subscriptions
-                and (ack_id := frame.headers["ack"])
-            ):
-                await self._connection_manager.maybe_write_frame(
-                    NackFrame(headers={"id": ack_id, "subscription": frame.headers["subscription"]})
-                )
+            if self._should_handle_ack_nack:
+                await self._nack(frame)
             self.on_suppressed_exception(exception, frame)
         else:
-            if (
-                self._should_handle_ack_nack
-                and self.id in self._active_subscriptions
-                and (ack_id := frame.headers["ack"])
-            ):
-                await self._connection_manager.maybe_write_frame(
-                    AckFrame(
-                        headers={"id": ack_id, "subscription": frame.headers["subscription"]},
-                    )
-                )
+            if self._should_handle_ack_nack:
+                await self._ack(frame)
+
+
+@dataclass(kw_only=True, slots=True)
+class ManualAckSubscription(BaseSubscription):
+    handler: Callable[["AckableMessageFrame"], Coroutine[Any, Any, Any]]
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class AckableMessageFrame(MessageFrame):
+    _subscription: ManualAckSubscription
+
+    async def ack(self) -> None:
+        await self._subscription._ack(self)  # noqa: SLF001
+
+    async def nack(self) -> None:
+        await self._subscription._nack(self)  # noqa: SLF001
 
 
 def _make_subscription_id() -> str:
diff --git a/packages/stompman/test_stompman/test_subscription.py b/packages/stompman/test_stompman/test_subscription.py
@@ -278,6 +278,35 @@ async def test_client_listen_auto_ack_nack(monkeypatch: pytest.MonkeyPatch, fake
     )
 
 
+async def test_client_listen_manual_ack_nack_ok(monkeypatch: pytest.MonkeyPatch, faker: faker.Faker) -> None:
+    subscription_id, destination, message_id, ack_id = faker.pystr(), faker.pystr(), faker.pystr(), faker.pystr()
+    monkeypatch.setattr(stompman.subscription, "_make_subscription_id", mock.Mock(return_value=subscription_id))
+
+    message_frame = build_dataclass(
+        MessageFrame,
+        headers={"destination": destination, "message-id": message_id, "subscription": subscription_id, "ack": ack_id},
+    )
+    connection_class, collected_frames = create_spying_connection(*get_read_frames_with_lifespan([message_frame]))
+
+    async def handle_message(message: stompman.subscription.AckableMessageFrame) -> None:
+        await message.ack()
+        await message.nack()
+
+    async with EnrichedClient(connection_class=connection_class) as client:
+        subscription = await client.subscribe_with_manual_ack(destination, handle_message)
+        await asyncio.sleep(0)
+        await asyncio.sleep(0)
+        await subscription.unsubscribe()
+
+    assert collected_frames == enrich_expected_frames(
+        SubscribeFrame(headers={"ack": "client-individual", "destination": destination, "id": subscription_id}),
+        message_frame,
+        AckFrame(headers={"subscription": subscription_id, "id": ack_id}),
+        NackFrame(headers={"subscription": subscription_id, "id": ack_id}),
+        UnsubscribeFrame(headers={"id": subscription_id}),
+    )
+
+
 async def test_client_listen_raises_on_aexit(monkeypatch: pytest.MonkeyPatch, faker: faker.Faker) -> None:
     monkeypatch.setattr("asyncio.sleep", partial(asyncio.sleep, 0))
 
