✨ NEW: Expose aio_pika.Connection.add_close_callback

kiwipy/rmq/communicator.py

@@ -541,13 +541,15 @@ async def async_connect(
     task_prefetch_count=defaults.TASK_PREFETCH_COUNT,
     encoder=defaults.ENCODER,
     decoder=defaults.DECODER,
-    testing_mode=False
+    testing_mode=False,
+    connection_close_callback: Optional[aio_pika.types.CloseCallbackType] = None,
 ) -> RmqCommunicator:
     # pylint: disable=too-many-arguments
     """Convenience method that returns a connected communicator.
 
     :param connection_params: parameters that will be passed to the connection factory to create the connection
     :param connection_factory: the factory method to open the aio_pika connection with
+    :param connection_close_callback: This will be called after the connection is closed
     :param message_exchange: The name of the RMQ message exchange to use
     :param queue_expires: the expiry time for standard queues in milliseconds. This is the time after which, if there
         are no subscribers, a queue will automatically be deleted by RabbitMQ.
@@ -565,6 +567,9 @@ async def async_connect(
     else:
         connection = await connection_factory(connection_params)
 
+    if connection_close_callback is not None:
+        connection.add_close_callback(connection_close_callback)
+
     communicator = RmqCommunicator(
         connection=connection,
 

kiwipy/rmq/threadcomms.py

@@ -5,7 +5,7 @@
 from concurrent.futures import Future as ThreadFuture
 import functools
 import logging
-from typing import Union
+from typing import Optional, Union
 
 import aio_pika
 import deprecation
@@ -45,12 +45,14 @@ def connect(
         encoder=defaults.ENCODER,
         decoder=defaults.DECODER,
         testing_mode=False,
-        async_task_timeout=TASK_TIMEOUT
+        async_task_timeout=TASK_TIMEOUT,
+        connection_close_callback: Optional[aio_pika.types.CloseCallbackType] = None,
     ):
         # pylint: disable=too-many-arguments
         comm = cls(
             connection_params,
             connection_factory,
+            connection_close_callback=connection_close_callback,
             message_exchange=message_exchange,
             task_exchange=task_exchange,
             task_queue=task_queue,
@@ -78,12 +80,14 @@ def __init__(
         encoder=defaults.ENCODER,
         decoder=defaults.DECODER,
         testing_mode=False,
-        async_task_timeout=TASK_TIMEOUT
+        async_task_timeout=TASK_TIMEOUT,
+        connection_close_callback: Optional[aio_pika.types.CloseCallbackType] = None,
     ):
         # pylint: disable=too-many-arguments
         """
         :param connection_params: parameters that will be passed to the connection factory to create the connection
         :param connection_factory: the factory method to open the aio_pika connection with
+        :param connection_close_callback: This will be called after the connection is closed
         :param message_exchange: The name of the RMQ message exchange to use
         :param queue_expires: the expiry time for standard queues in milliseconds. This is the time after which, if
             there are no subscribers, a queue will automatically be deleted by RabbitMQ.
@@ -109,7 +113,7 @@ def __init__(
             communicator.async_connect(
                 connection_params=connection_params,
                 connection_factory=connection_factory,
-
+                connection_close_callback=connection_close_callback,
                 # Messages
                 message_exchange=message_exchange,
                 queue_expires=queue_expires,
@@ -339,7 +343,8 @@ def connect(
     task_prefetch_count=defaults.TASK_PREFETCH_COUNT,
     encoder=defaults.ENCODER,
     decoder=defaults.DECODER,
-    testing_mode=False
+    testing_mode=False,
+    connection_close_callback: Optional[aio_pika.types.CloseCallbackType] = None,
 ) -> RmqThreadCommunicator:
     """
     Establish a RabbitMQ communicator connection
@@ -348,6 +353,7 @@ def connect(
     return RmqThreadCommunicator.connect(
         connection_params=connection_params,
         connection_factory=connection_factory,
+        connection_close_callback=connection_close_callback,
         message_exchange=message_exchange,
         task_exchange=task_exchange,
         task_queue=task_queue,

test/rmq/test_rmq_thread_communicator.py

@@ -234,6 +234,25 @@ def test_task_processing_exception(thread_task_queue: rmq.RmqThreadTaskQueue):
             pass
 
 
+def test_connection_close_callback():
+    """Test that a callback set with `connection_close_callback` is correctly called."""
+    result = []
+
+    def close_callback(sender, exc):  # pylint: disable=unused-argument
+        result.append('called')
+
+    communicator = rmq.connect(
+        connection_params={'url': 'amqp://guest:guest@localhost:5672/'},
+        connection_close_callback=close_callback,
+        message_exchange=f'{__file__}.{shortuuid.uuid()}',
+        task_exchange=f'{__file__}.{shortuuid.uuid()}',
+        task_queue=f'{__file__}.{shortuuid.uuid()}',
+        testing_mode=True
+    )
+    communicator.close()
+    assert result == ['called']
+
+
 @pytest.mark.skipif(sys.version_info < (3, 6), reason='`pytest-notebook` plugin requires Python >= 3.6')
 def test_jupyter_notebook():
     """Test that the `RmqThreadCommunicator` can be used in a Jupyter notebook."""