changed syntax of subscribe

Author: sanderegg

packages/service-library/src/servicelib/rabbitmq/__init__.py

@@ -8,13 +8,16 @@
     RPCNotInitializedError,
     RPCServerError,
 )
-from ._models import ConsumerTag, QueueName
+from ._models import ConsumerTag, ExchangeName, QueueName
 from ._rpc_router import RPCRouter
 from ._utils import is_rabbitmq_responsive, wait_till_rabbitmq_responsive
 
 __all__: tuple[str, ...] = (
     "BIND_TO_ALL_TOPICS",
+    "ConsumerTag",
+    "ExchangeName",
     "is_rabbitmq_responsive",
+    "QueueName",
     "RabbitMQClient",
     "RabbitMQRPCClient",
     "RemoteMethodNotRegisteredError",
@@ -24,8 +27,6 @@
     "RPCRouter",
     "RPCServerError",
     "wait_till_rabbitmq_responsive",
-    "QueueName",
-    "ConsumerTag",
 )
 
 # nopycln: file

packages/service-library/src/servicelib/rabbitmq/_client.py

@@ -10,7 +10,14 @@
 
 from ..logging_utils import log_catch, log_context
 from ._client_base import RabbitMQClientBase
-from ._models import ConsumerTag, MessageHandler, QueueName, RabbitMessage
+from ._models import (
+    ConsumerTag,
+    ExchangeName,
+    MessageHandler,
+    QueueName,
+    RabbitMessage,
+    TopicName,
+)
 from ._utils import (
     RABBIT_QUEUE_MESSAGE_DEFAULT_TTL_MS,
     declare_queue,
@@ -27,7 +34,7 @@
 _DEFAULT_UNEXPECTED_ERROR_RETRY_DELAY_S: Final[float] = 1
 _DEFAULT_UNEXPECTED_ERROR_MAX_ATTEMPTS: Final[NonNegativeInt] = 15
 
-_DELAYED_EXCHANGE_NAME: Final[str] = "delayed_{exchange_name}"
+_DELAYED_EXCHANGE_NAME: Final[ExchangeName] = ExchangeName("delayed_{exchange_name}")
 
 
 def _get_x_death_count(message: aio_pika.abc.AbstractIncomingMessage) -> int:
@@ -139,12 +146,14 @@ async def _get_channel(self) -> aio_pika.abc.AbstractChannel:
             channel.close_callbacks.add(self._channel_close_callback)
             return channel
 
-    async def _get_consumer_tag(self, exchange_name) -> str:
-        return f"{get_rabbitmq_client_unique_name(self.client_name)}_{exchange_name}_{uuid4()}"
+    async def _create_consumer_tag(self, exchange_name) -> ConsumerTag:
+        return ConsumerTag(
+            f"{get_rabbitmq_client_unique_name(self.client_name)}_{exchange_name}_{uuid4()}"
+        )
 
     async def subscribe(
         self,
-        exchange_name: str,
+        exchange_name: ExchangeName,
         message_handler: MessageHandler,
         *,
         exclusive_queue: bool = True,
@@ -239,20 +248,19 @@ async def subscribe(
             )
             await delayed_queue.bind(delayed_exchange)
 
-            consumer_tag = await self._get_consumer_tag(exchange_name)
+            consumer_tag = await self._create_consumer_tag(exchange_name)
             await queue.consume(
                 partial(_on_message, message_handler, unexpected_error_max_attempts),
                 exclusive=exclusive_queue,
                 consumer_tag=consumer_tag,
             )
-            output: str = queue.name
-            return output, consumer_tag
+            return queue.name, consumer_tag
 
     async def add_topics(
         self,
-        exchange_name: str,
+        exchange_name: ExchangeName,
         *,
-        topics: list[str],
+        topics: list[TopicName],
     ) -> None:
         assert self._channel_pool  # nosec
 
@@ -276,9 +284,9 @@ async def add_topics(
 
     async def remove_topics(
         self,
-        exchange_name: str,
+        exchange_name: ExchangeName,
         *,
-        topics: list[str],
+        topics: list[TopicName],
     ) -> None:
         assert self._channel_pool  # nosec
         async with self._channel_pool.acquire() as channel:
@@ -305,14 +313,20 @@ async def unsubscribe(
     ) -> None:
         """This will delete the queue if there are no consumers left"""
         assert self._connection_pool  # nosec
-        if not self._connection_pool.is_closed:
-            assert self._channel_pool  # nosec
-            async with self._channel_pool.acquire() as channel:
-                queue = await channel.get_queue(queue_name)
-                # NOTE: we force delete here
-                await queue.delete(if_unused=False, if_empty=False)
+        if self._connection_pool.is_closed:
+            _logger.warning(
+                "Connection to RabbitMQ is already closed, skipping unsubscribe from queue..."
+            )
+            return
+        assert self._channel_pool  # nosec
+        async with self._channel_pool.acquire() as channel:
+            queue = await channel.get_queue(queue_name)
+            # NOTE: we force delete here
+            await queue.delete(if_unused=False, if_empty=False)
 
-    async def publish(self, exchange_name: str, message: RabbitMessage) -> None:
+    async def publish(
+        self, exchange_name: ExchangeName, message: RabbitMessage
+    ) -> None:
         """publish message in the exchange exchange_name.
         specifying a topic will use a TOPIC type of RabbitMQ Exchange instead of FANOUT
 
@@ -342,9 +356,13 @@ async def unsubscribe_consumer(
     ) -> None:
         """This will only remove the consumers without deleting the queue"""
         assert self._connection_pool  # nosec
-        if not self._connection_pool.is_closed:
-            assert self._channel_pool  # nosec
-            async with self._channel_pool.acquire() as channel:
-                assert isinstance(channel, aio_pika.RobustChannel)  # nosec
-                queue = await channel.get_queue(queue_name)
-                await queue.cancel(consumer_tag)
+        if self._connection_pool.is_closed:
+            _logger.warning(
+                "Connection to RabbitMQ is already closed, skipping unsubscribe consumers from queue..."
+            )
+            return
+        assert self._channel_pool  # nosec
+        async with self._channel_pool.acquire() as channel:
+            assert isinstance(channel, aio_pika.RobustChannel)  # nosec
+            queue = await channel.get_queue(queue_name)
+            await queue.cancel(consumer_tag)

packages/service-library/src/servicelib/rabbitmq/_models.py

@@ -11,8 +11,10 @@
 
 MessageHandler = Callable[[Any], Awaitable[bool]]
 
+ExchangeName: TypeAlias = str
 QueueName: TypeAlias = str
 ConsumerTag: TypeAlias = str
+TopicName: TypeAlias = str
 
 
 class RabbitMessage(Protocol):

packages/service-library/tests/rabbitmq/test_rabbitmq.py

@@ -15,7 +15,13 @@
 import pytest
 from faker import Faker
 from pytest_mock.plugin import MockerFixture
-from servicelib.rabbitmq import BIND_TO_ALL_TOPICS, RabbitMQClient, _client
+from servicelib.rabbitmq import (
+    BIND_TO_ALL_TOPICS,
+    ConsumerTag,
+    QueueName,
+    RabbitMQClient,
+    _client,
+)
 from servicelib.rabbitmq._client import _DEFAULT_UNEXPECTED_ERROR_MAX_ATTEMPTS
 from settings_library.rabbit import RabbitSettings
 from tenacity.asyncio import AsyncRetrying
@@ -325,7 +331,7 @@ async def test_publish_with_no_registered_subscriber(
     ttl_s: float = 0.1
     topics_count: int = 1 if topics is None else len(topics)
 
-    async def _publish_random_message():
+    async def _publish_random_message() -> None:
         if topics is None:
             message = random_rabbit_message()
             await publisher.publish(exchange_name, message)
@@ -335,8 +341,8 @@ async def _publish_random_message():
                 message = random_rabbit_message(topic=topic)
                 await publisher.publish(exchange_name, message)
 
-    async def _subscribe_consumer_to_queue():
-        await consumer.subscribe(
+    async def _subscribe_consumer_to_queue() -> tuple[QueueName, ConsumerTag]:
+        return await consumer.subscribe(
             exchange_name,
             mocked_message_parser,
             topics=topics,
@@ -346,27 +352,30 @@ async def _subscribe_consumer_to_queue():
             unexpected_error_retry_delay_s=ttl_s,
         )
 
-    async def _unsubscribe_consumer():
-        await consumer.unsubscribe_consumer(exchange_name)
+    async def _unsubscribe_consumer(
+        queue_name: QueueName, consumer_tag: ConsumerTag
+    ) -> None:
+        await consumer.unsubscribe_consumer(queue_name, consumer_tag)
 
     # CASE 1 (subscribe immediately after publishing message)
 
-    await _subscribe_consumer_to_queue()
-    await _unsubscribe_consumer()
+    consumer_1 = await _subscribe_consumer_to_queue()
+    await _unsubscribe_consumer(*consumer_1)
     await _publish_random_message()
     # reconnect immediately
-    await _subscribe_consumer_to_queue()
+    consumer_2 = await _subscribe_consumer_to_queue()
     # expected to receive a message (one per topic)
     await _assert_wait_for_messages(on_message_spy, 1 * topics_count)
 
     # CASE 2 (no subscriber attached when publishing)
     on_message_spy.reset_mock()
 
-    await _unsubscribe_consumer()
+    await _unsubscribe_consumer(*consumer_2)
     await _publish_random_message()
     # wait for message to expire (will be dropped)
     await asyncio.sleep(ttl_s * 2)
-    await _subscribe_consumer_to_queue()
+    _consumer_3 = await _subscribe_consumer_to_queue()
+
     # wait for a message to be possibly delivered
     await asyncio.sleep(ttl_s * 2)
     # nothing changed from before
@@ -604,7 +613,7 @@ async def test_rabbit_pub_sub_bind_and_unbind_topics(
     )
 
     # we should get no messages since no one was subscribed
-    queue_name = await consumer.subscribe(
+    queue_name, consumer_tag = await consumer.subscribe(
         exchange_name, mocked_message_parser, topics=[]
     )
     await _assert_message_received(mocked_message_parser, 0)
@@ -666,7 +675,7 @@ async def test_rabbit_adding_topics_to_a_fanout_exchange(
     message = random_rabbit_message()
     publisher = create_rabbitmq_client("publisher")
     consumer = create_rabbitmq_client("consumer")
-    queue_name = await consumer.subscribe(exchange_name, mocked_message_parser)
+    queue_name, _ = await consumer.subscribe(exchange_name, mocked_message_parser)
     await publisher.publish(exchange_name, message)
     await _assert_message_received(mocked_message_parser, 1, message)
     mocked_message_parser.reset_mock()
@@ -709,10 +718,12 @@ async def test_unsubscribe_consumer(
 ):
     exchange_name = f"{random_exchange_name()}"
     client = create_rabbitmq_client("consumer")
-    await client.subscribe(exchange_name, mocked_message_parser, exclusive_queue=False)
+    queue_name, consumer_tag = await client.subscribe(
+        exchange_name, mocked_message_parser, exclusive_queue=False
+    )
     # Unsubsribe just a consumer, the queue will be still there
-    await client.unsubscribe_consumer(exchange_name)
+    await client.unsubscribe_consumer(queue_name, consumer_tag)
     # Unsubsribe the queue
-    await client.unsubscribe(exchange_name)
+    await client.unsubscribe(queue_name)
     with pytest.raises(aio_pika.exceptions.ChannelNotFoundEntity):
-        await client.unsubscribe(exchange_name)
+        await client.unsubscribe(queue_name)