Change transaction API (#42)

Author: vrslev

README.md

@@ -46,9 +46,9 @@ await client.send(body=b"hi there!", destination="DLQ", headers={"persistent": "
 Or, to send messages in a transaction:
 
 ```python
-async with client.enter_transaction() as transaction:
+async with client.begin() as transaction:
     for _ in range(10):
-        await client.send(body=b"hi there!", destination="DLQ", transaction=transaction, headers={"persistent": "true"})
+        await transaction.send(body=b"hi there!", destination="DLQ", headers={"persistent": "true"})
         await asyncio.sleep(0.1)
 ```
 
@@ -108,7 +108,7 @@ async def handle_message(event: stompman.MessageEvent) -> None:
 
 ### Cleaning Up
 
-stompman takes care of cleaning up resources automatically. When you leave the context of async context managers `stompman.Client()`, `client.subscribe()`, or `client.enter_transaction()`, the necessary frames will be sent to the server.
+stompman takes care of cleaning up resources automatically. When you leave the context of async context managers `stompman.Client()`, `client.subscribe()`, or `client.begin()`, the necessary frames will be sent to the server.
 
 ### Handling Connectivity Issues
 

stompman/client.py

@@ -28,7 +28,6 @@
     NackFrame,
     ReceiptFrame,
     SendFrame,
-    SendHeaders,
     SubscribeFrame,
     UnsubscribeFrame,
 )
@@ -241,12 +240,12 @@ async def send_heartbeats_forever() -> None:
                     break
 
     @asynccontextmanager
-    async def enter_transaction(self) -> AsyncGenerator[str, None]:
+    async def begin(self) -> AsyncGenerator["Transaction", None]:
         transaction_id = str(uuid4())
         await self._connection.write_frame(BeginFrame(headers={"transaction": transaction_id}))
 
         try:
-            yield transaction_id
+            yield Transaction(id=transaction_id, _connection=self._connection)
         except Exception:
             if self._connection.active:
                 await self._connection.write_frame(AbortFrame(headers={"transaction": transaction_id}))
@@ -255,22 +254,18 @@ async def enter_transaction(self) -> AsyncGenerator[str, None]:
             if self._connection.active:
                 await self._connection.write_frame(CommitFrame(headers={"transaction": transaction_id}))
 
-    async def send(  # noqa: PLR0913
+    async def send(
         self,
         body: bytes,
         destination: str,
-        transaction: str | None = None,
         content_type: str | None = None,
         headers: dict[str, str] | None = None,
     ) -> None:
-        full_headers: SendHeaders = headers or {}  # type: ignore[assignment]
-        full_headers["destination"] = destination
-        full_headers["content-length"] = str(len(body))
-        if content_type is not None:
-            full_headers["content-type"] = content_type
-        if transaction is not None:
-            full_headers["transaction"] = transaction
-        await self._connection.write_frame(SendFrame(headers=full_headers, body=body))
+        await self._connection.write_frame(
+            SendFrame.build(
+                body=body, destination=destination, transaction=None, content_type=content_type, headers=headers
+            )
+        )
 
     @asynccontextmanager
     async def subscribe(self, destination: str) -> AsyncGenerator[None, None]:
@@ -300,6 +295,25 @@ async def listen(self) -> AsyncIterator["AnyListeningEvent"]:
                     raise AssertionError(msg, frame)
 
 
+@dataclass(kw_only=True, slots=True)
+class Transaction:
+    id: str
+    _connection: AbstractConnection
+
+    async def send(
+        self,
+        body: bytes,
+        destination: str,
+        content_type: str | None = None,
+        headers: dict[str, str] | None = None,
+    ) -> None:
+        await self._connection.write_frame(
+            SendFrame.build(
+                body=body, destination=destination, transaction=self.id, content_type=content_type, headers=headers
+            )
+        )
+
+
 @dataclass(kw_only=True, slots=True)
 class MessageEvent:
     body: bytes = field(init=False)

stompman/frames.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Literal, NotRequired, TypedDict
+from typing import Literal, NotRequired, Self, TypedDict
 
 ConnectHeaders = TypedDict(
     "ConnectHeaders",
@@ -140,6 +140,25 @@ class SendFrame:
     headers: SendHeaders
     body: bytes = b""
 
+    @classmethod
+    def build(  # noqa: PLR0913
+        cls,
+        *,
+        body: bytes,
+        destination: str,
+        transaction: str | None,
+        content_type: str | None,
+        headers: dict[str, str] | None,
+    ) -> Self:
+        full_headers: SendHeaders = headers or {}  # type: ignore[assignment]
+        full_headers["destination"] = destination
+        full_headers["content-length"] = str(len(body))
+        if content_type is not None:
+            full_headers["content-type"] = content_type
+        if transaction is not None:
+            full_headers["transaction"] = transaction
+        return cls(headers=full_headers, body=body)
+
 
 @dataclass(frozen=True, kw_only=True, slots=True)
 class SubscribeFrame:

testing/producer.py

@@ -5,14 +5,11 @@
 
 
 async def main() -> None:
-    async with (
-        stompman.Client(servers=[CONNECTION_PARAMETERS]) as client,
-        client.enter_transaction() as transaction,
-    ):
+    async with stompman.Client(servers=[CONNECTION_PARAMETERS]) as client, client.begin() as transaction:
         for _ in range(10):
-            await client.send(body=b"hi there!", destination="DLQ", transaction=transaction)
+            await transaction.send(body=b"hi there!", destination="DLQ")
         await asyncio.sleep(3)
-        await client.send(body=b"hi there!", destination="DLQ", transaction=transaction)
+        await transaction.send(body=b"hi there!", destination="DLQ")
 
 
 if __name__ == "__main__":

tests/integration.py

@@ -50,11 +50,9 @@ def destination() -> str:
 
 async def test_ok(destination: str) -> None:
     async def produce() -> None:
-        async with producer.enter_transaction() as transaction:
+        async with producer.begin() as transaction:
             for message in messages:
-                await producer.send(
-                    body=message, destination=destination, transaction=transaction, headers={"hello": "world"}
-                )
+                await transaction.send(body=message, destination=destination, headers={"hello": "world"})
 
     async def consume() -> None:
         received_messages = []
@@ -102,19 +100,19 @@ async def test_not_raises_connection_lost_error_in_subscription(client: stompman
 
 
 async def test_not_raises_connection_lost_error_in_transaction_without_send(client: stompman.Client) -> None:
-    async with client.enter_transaction():
+    async with client.begin():
         await client._connection.close()
 
 
 async def test_not_raises_connection_lost_error_in_transaction_with_send(
     client: stompman.Client, destination: str
 ) -> None:
-    async with client.enter_transaction() as transaction:
-        await client.send(b"first", destination=destination, transaction=transaction)
+    async with client.begin() as transaction:
+        await transaction.send(b"first", destination=destination)
         await client._connection.close()
 
         with pytest.raises(ConnectionLostError):
-            await client.send(b"second", destination=destination, transaction=transaction)
+            await transaction.send(b"second", destination=destination)
 
 
 async def test_raises_connection_lost_error_in_send(client: stompman.Client, destination: str) -> None:

tests/test_client.py

@@ -20,6 +20,7 @@
     ConnectedFrame,
     ConnectFrame,
     ConnectionConfirmationTimeoutError,
+    ConnectionLostError,
     ConnectionParameters,
     DisconnectFrame,
     ErrorEvent,
@@ -36,7 +37,6 @@
     UnsubscribeFrame,
     UnsupportedProtocolVersionError,
 )
-from stompman.errors import ConnectionLostError
 
 pytestmark = pytest.mark.anyio
 
@@ -470,53 +470,59 @@ async def test_send_message_and_enter_transaction_ok(monkeypatch: pytest.MonkeyP
     body = b"hello"
     destination = "/queue/test"
     expires = "whatever"
-    transaction = "myid"
+    transaction_id = "myid"
     content_type = "my-content-type"
-    monkeypatch.setattr(stompman.client, "uuid4", lambda: transaction)
+    monkeypatch.setattr(stompman.client, "uuid4", lambda: transaction_id)
 
     connection_class, collected_frames = create_spying_connection(get_read_frames_with_lifespan([]))
     async with (
         EnrichedClient(connection_class=connection_class) as client,
-        client.enter_transaction() as transaction,
+        client.begin() as transaction,
     ):
-        await client.send(
-            body=body,
-            destination=destination,
-            transaction=transaction,
-            content_type=content_type,
-            headers={"expires": expires},
+        await transaction.send(
+            body=body, destination=destination, content_type=content_type, headers={"expires": expires}
         )
+        await client.send(body=body, destination=destination, content_type=content_type, headers={"expires": expires})
 
     assert_frames_between_lifespan_match(
         collected_frames,
         [
-            BeginFrame(headers={"transaction": transaction}),
+            BeginFrame(headers={"transaction": transaction_id}),
             SendFrame(
                 headers={  # type: ignore[typeddict-unknown-key]
                     "content-length": str(len(body)),
                     "content-type": content_type,
                     "destination": destination,
-                    "transaction": transaction,
+                    "transaction": transaction_id,
                     "expires": expires,
                 },
                 body=b"hello",
             ),
-            CommitFrame(headers={"transaction": transaction}),
+            SendFrame(
+                headers={  # type: ignore[typeddict-unknown-key]
+                    "content-length": str(len(body)),
+                    "content-type": content_type,
+                    "destination": destination,
+                    "expires": expires,
+                },
+                body=b"hello",
+            ),
+            CommitFrame(headers={"transaction": transaction_id}),
         ],
     )
 
 
 async def test_send_message_and_enter_transaction_abort(monkeypatch: pytest.MonkeyPatch) -> None:
-    transaction = "myid"
-    monkeypatch.setattr(stompman.client, "uuid4", lambda: transaction)
+    transaction_id = "myid"
+    monkeypatch.setattr(stompman.client, "uuid4", lambda: transaction_id)
 
     connection_class, collected_frames = create_spying_connection(get_read_frames_with_lifespan([]))
     async with EnrichedClient(connection_class=connection_class) as client:
         with suppress(AssertionError):
-            async with client.enter_transaction() as transaction:
+            async with client.begin():
                 raise AssertionError
 
     assert_frames_between_lifespan_match(
         collected_frames,
-        [BeginFrame(headers={"transaction": transaction}), AbortFrame(headers={"transaction": transaction})],
+        [BeginFrame(headers={"transaction": transaction_id}), AbortFrame(headers={"transaction": transaction_id})],
     )