Refactor internals (#24)

Author: vrslev

README.md

@@ -50,7 +50,7 @@ Or, to send messages in a transaction:
 ```python
 async with client.enter_transaction() as transaction:
     for _ in range(10):
-        await client.send(body=b"hi there!", destination="DLQ", transaction=transaction)
+        await client.send(body=b"hi there!", destination="DLQ", transaction=transaction, headers={"persistent": "true"})
         await asyncio.sleep(0.1)
 ```
 

stompman/__init__.py

@@ -1,5 +1,13 @@
-from stompman.client import Client, Heartbeat
-from stompman.connection import AbstractConnection, Connection, ConnectionParameters
+from stompman.client import (
+    AnyListeningEvent,
+    Client,
+    ConnectionParameters,
+    ErrorEvent,
+    Heartbeat,
+    HeartbeatEvent,
+    MessageEvent,
+)
+from stompman.connection import AbstractConnection, Connection
 from stompman.errors import (
     ConnectionConfirmationTimeoutError,
     ConnectionLostError,
@@ -26,7 +34,6 @@
     SubscribeFrame,
     UnsubscribeFrame,
 )
-from stompman.listening_events import AnyListeningEvent, ErrorEvent, HeartbeatEvent, MessageEvent
 
 __all__ = [
     "AbortFrame",

stompman/client.py

@@ -1,19 +1,20 @@
 import asyncio
-from collections.abc import AsyncGenerator, AsyncIterator
+from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable
 from contextlib import AsyncExitStack, asynccontextmanager
 from dataclasses import dataclass, field
 from types import TracebackType
-from typing import NamedTuple, Self
+from typing import Any, ClassVar, NamedTuple, Self, TypedDict
 from uuid import uuid4
 
-from stompman.connection import AbstractConnection, Connection, ConnectionParameters
+from stompman.connection import AbstractConnection, Connection
 from stompman.errors import (
     ConnectionConfirmationTimeoutError,
     FailedAllConnectAttemptsError,
     UnsupportedProtocolVersionError,
 )
 from stompman.frames import (
     AbortFrame,
+    AckFrame,
     BeginFrame,
     CommitFrame,
     ConnectedFrame,
@@ -22,14 +23,13 @@
     ErrorFrame,
     HeartbeatFrame,
     MessageFrame,
+    NackFrame,
     ReceiptFrame,
     SendFrame,
     SendHeaders,
     SubscribeFrame,
     UnsubscribeFrame,
 )
-from stompman.listening_events import AnyListeningEvent, ErrorEvent, HeartbeatEvent, MessageEvent
-from stompman.protocol import PROTOCOL_VERSION
 
 
 class Heartbeat(NamedTuple):
@@ -45,6 +45,56 @@ def from_header(cls, header: str) -> Self:
         return cls(int(first), int(second))
 
 
+class MultiHostHostLike(TypedDict):
+    username: str | None
+    password: str | None
+    host: str | None
+    port: int | None
+
+
+@dataclass
+class ConnectionParameters:
+    host: str
+    port: int
+    login: str
+    passcode: str = field(repr=False)
+
+    @classmethod
+    def from_pydantic_multihost_hosts(cls, hosts: list[MultiHostHostLike]) -> list[Self]:
+        """Create connection parameters from a list of `MultiHostUrl` objects.
+
+        .. code-block:: python
+        import stompman.
+
+        ArtemisDsn = typing.Annotated[
+            pydantic_core.MultiHostUrl,
+            pydantic.UrlConstraints(
+                host_required=True,
+                allowed_schemes=["tcp"],
+            ),
+        ]
+
+        async with stompman.Client(
+            servers=stompman.ConnectionParameters.from_pydantic_multihost_hosts(
+                ArtemisDsn("tcp://lev:pass@host1:61616,lev:pass@host1:61617,lev:pass@host2:61616").hosts()
+            ),
+        ):
+            ...
+        """
+        servers: list[Self] = []
+        for host in hosts:
+            if host["host"] is None:
+                raise ValueError("host must be set")
+            if host["port"] is None:
+                raise ValueError("port must be set")
+            if host["username"] is None:
+                raise ValueError("username must be set")
+            if host["password"] is None:
+                raise ValueError("password must be set")
+            servers.append(cls(host=host["host"], port=host["port"], login=host["username"], passcode=host["password"]))
+        return servers
+
+
 @dataclass
 class Client:
     servers: list[ConnectionParameters]
@@ -57,11 +107,14 @@ class Client:
     read_max_chunk_size: int = 1024 * 1024
     connection_class: type[AbstractConnection] = Connection
 
+    PROTOCOL_VERSION: ClassVar = "1.2"  # https://stomp.github.io/stomp-specification-1.2.html
+
     _connection: AbstractConnection = field(init=False)
+    _connection_parameters: ConnectionParameters = field(init=False)
     _exit_stack: AsyncExitStack = field(default_factory=AsyncExitStack, init=False)
 
     async def __aenter__(self) -> Self:
-        self._connection = await self._connect_to_any_server()
+        await self._connect_to_any_server()
         await self._exit_stack.enter_async_context(self._connection_lifespan())
         return self
 
@@ -71,26 +124,25 @@ async def __aexit__(
         await self._exit_stack.aclose()
         await self._connection.close()
 
-    async def _connect_to_one_server(self, server: ConnectionParameters) -> AbstractConnection | None:
+    async def _connect_to_one_server(
+        self, server: ConnectionParameters
+    ) -> tuple[AbstractConnection, ConnectionParameters] | None:
         for attempt in range(self.connect_retry_attempts):
-            connection = self.connection_class(
-                connection_parameters=server,
-                connect_timeout=self.connect_timeout,
-                read_timeout=self.read_timeout,
-                read_max_chunk_size=self.read_max_chunk_size,
-            )
-            if await connection.connect():
-                return connection
+            if connection := await self.connection_class.connect(
+                host=server.host, port=server.port, timeout=self.connect_timeout
+            ):
+                return connection, server
             await asyncio.sleep(self.connect_retry_interval * (attempt + 1))
         return None
 
-    async def _connect_to_any_server(self) -> AbstractConnection:
+    async def _connect_to_any_server(self) -> None:
         for maybe_connection_future in asyncio.as_completed(
             [self._connect_to_one_server(server) for server in self.servers]
         ):
-            maybe_connection = await maybe_connection_future
-            if maybe_connection:
-                return maybe_connection
+            maybe_result = await maybe_connection_future
+            if maybe_result:
+                self._connection, self._connection_parameters = maybe_result
+                return
         raise FailedAllConnectAttemptsError(
             servers=self.servers,
             retry_attempts=self.connect_retry_attempts,
@@ -112,24 +164,27 @@ async def _connection_lifespan(self) -> AsyncGenerator[None, None]:
         await self._connection.write_frame(
             ConnectFrame(
                 headers={
-                    "accept-version": PROTOCOL_VERSION,
+                    "accept-version": self.PROTOCOL_VERSION,
                     "heart-beat": self.heartbeat.to_header(),
-                    "host": self._connection.connection_parameters.host,
-                    "login": self._connection.connection_parameters.login,
-                    "passcode": self._connection.connection_parameters.passcode,
+                    "host": self._connection_parameters.host,
+                    "login": self._connection_parameters.login,
+                    "passcode": self._connection_parameters.passcode,
                 },
             )
         )
         try:
             connected_frame = await asyncio.wait_for(
-                self._connection.read_frame_of_type(ConnectedFrame), timeout=self.connection_confirmation_timeout
+                self._connection.read_frame_of_type(
+                    ConnectedFrame, max_chunk_size=self.read_max_chunk_size, timeout=self.read_timeout
+                ),
+                timeout=self.connection_confirmation_timeout,
             )
         except TimeoutError as exception:
             raise ConnectionConfirmationTimeoutError(self.connection_confirmation_timeout) from exception
 
-        if connected_frame.headers["version"] != PROTOCOL_VERSION:
+        if connected_frame.headers["version"] != self.PROTOCOL_VERSION:
             raise UnsupportedProtocolVersionError(
-                given_version=connected_frame.headers["version"], supported_version=PROTOCOL_VERSION
+                given_version=connected_frame.headers["version"], supported_version=self.PROTOCOL_VERSION
             )
 
         server_heartbeat = Heartbeat.from_header(connected_frame.headers["heart-beat"])
@@ -150,30 +205,9 @@ async def send_heartbeats_forever() -> None:
                 task.cancel()
 
         await self._connection.write_frame(DisconnectFrame(headers={"receipt": str(uuid4())}))
-        await self._connection.read_frame_of_type(ReceiptFrame)
-
-    @asynccontextmanager
-    async def subscribe(self, destination: str) -> AsyncGenerator[None, None]:
-        subscription_id = str(uuid4())
-        await self._connection.write_frame(
-            SubscribeFrame(headers={"id": subscription_id, "destination": destination, "ack": "client-individual"})
+        await self._connection.read_frame_of_type(
+            ReceiptFrame, max_chunk_size=self.read_max_chunk_size, timeout=self.read_timeout
         )
-        try:
-            yield
-        finally:
-            await self._connection.write_frame(UnsubscribeFrame(headers={"id": subscription_id}))
-
-    async def listen(self) -> AsyncIterator[AnyListeningEvent]:
-        async for frame in self._connection.read_frames():
-            match frame:
-                case MessageFrame():
-                    yield MessageEvent(_client=self, _frame=frame)
-                case ErrorFrame():
-                    yield ErrorEvent(_client=self, _frame=frame)
-                case HeartbeatFrame():
-                    yield HeartbeatEvent(_client=self, _frame=frame)
-                case ConnectedFrame() | ReceiptFrame():
-                    raise AssertionError("Should be unreachable! Report the issue.", frame)
 
     @asynccontextmanager
     async def enter_transaction(self) -> AsyncGenerator[str, None]:
@@ -204,3 +238,93 @@ async def send(  # noqa: PLR0913
         if transaction is not None:
             full_headers["transaction"] = transaction
         await self._connection.write_frame(SendFrame(headers=full_headers, body=body))
+
+    @asynccontextmanager
+    async def subscribe(self, destination: str) -> AsyncGenerator[None, None]:
+        subscription_id = str(uuid4())
+        await self._connection.write_frame(
+            SubscribeFrame(headers={"id": subscription_id, "destination": destination, "ack": "client-individual"})
+        )
+        try:
+            yield
+        finally:
+            await self._connection.write_frame(UnsubscribeFrame(headers={"id": subscription_id}))
+
+    async def listen(self) -> AsyncIterator["AnyListeningEvent"]:
+        async for frame in self._connection.read_frames(
+            max_chunk_size=self.read_max_chunk_size, timeout=self.read_timeout
+        ):
+            match frame:
+                case MessageFrame():
+                    yield MessageEvent(_client=self, _frame=frame)
+                case ErrorFrame():
+                    yield ErrorEvent(_client=self, _frame=frame)
+                case HeartbeatFrame():
+                    yield HeartbeatEvent(_client=self, _frame=frame)
+                case ConnectedFrame() | ReceiptFrame():
+                    raise AssertionError("Should be unreachable! Report the issue.", frame)
+
+
+@dataclass
+class MessageEvent:
+    body: bytes = field(init=False)
+    _frame: MessageFrame
+    _client: "Client" = field(repr=False)
+
+    def __post_init__(self) -> None:
+        self.body = self._frame.body
+
+    async def ack(self) -> None:
+        await self._client._connection.write_frame(
+            AckFrame(
+                headers={"id": self._frame.headers["message-id"], "subscription": self._frame.headers["subscription"]},
+            )
+        )
+
+    async def nack(self) -> None:
+        await self._client._connection.write_frame(
+            NackFrame(
+                headers={"id": self._frame.headers["message-id"], "subscription": self._frame.headers["subscription"]}
+            )
+        )
+
+    async def with_auto_ack(
+        self,
+        awaitable: Awaitable[None],
+        *,
+        on_suppressed_exception: Callable[[Exception, Self], Any],
+        supressed_exception_classes: tuple[type[Exception], ...] = (Exception,),
+    ) -> None:
+        called_nack = False
+        try:
+            await awaitable
+        except supressed_exception_classes as exception:
+            await self.nack()
+            called_nack = True
+            on_suppressed_exception(exception, self)
+        finally:
+            if not called_nack:
+                await self.ack()
+
+
+@dataclass
+class ErrorEvent:
+    message_header: str = field(init=False)
+    """Short description of the error."""
+    body: bytes = field(init=False)
+    """Long description of the error."""
+    _frame: ErrorFrame
+    _client: "Client" = field(repr=False)
+
+    def __post_init__(self) -> None:
+        self.message_header = self._frame.headers["message"]
+        self.body = self._frame.body
+
+
+@dataclass
+class HeartbeatEvent:
+    _frame: HeartbeatFrame
+    _client: "Client" = field(repr=False)
+
+
+AnyListeningEvent = MessageEvent | ErrorEvent | HeartbeatEvent