Add method for synchronous listening (#89)

* Sync publisher, using new psycopg features

* Update more type hints

Author: AlanCoding

dispatcher/brokers/base.py

@@ -1,14 +1,18 @@
-from typing import AsyncGenerator, Optional, Protocol
+from typing import Any, AsyncGenerator, Callable, Coroutine, Iterator, Optional, Protocol
 
 
 class BaseBroker(Protocol):
-    async def aprocess_notify(self, connected_callback=None) -> AsyncGenerator[tuple[str, str], None]:
+    async def aprocess_notify(
+        self, connected_callback: Optional[Optional[Callable[[], Coroutine[Any, Any, None]]]] = None
+    ) -> AsyncGenerator[tuple[str, str], None]:
         yield ('', '')  # yield affects CPython type https://github.com/python/mypy/pull/18422
 
     async def apublish_message(self, channel: Optional[str] = None, message: str = '') -> None: ...
 
     async def aclose(self) -> None: ...
 
+    def process_notify(self, connected_callback: Optional[Callable] = None, timeout: float = 5.0, max_messages: int = 1) -> Iterator[tuple[str, str]]: ...
+
     def publish_message(self, channel=None, message=None): ...
 
     def close(self): ...

dispatcher/brokers/pg_notify.py

@@ -1,5 +1,5 @@
 import logging
-from typing import AsyncGenerator, Callable, Optional, Union
+from typing import Any, AsyncGenerator, Callable, Coroutine, Iterator, Optional, Union
 
 import psycopg
 
@@ -112,11 +112,21 @@ async def aget_connection(self) -> psycopg.AsyncConnection:
             return connection  # slightly weird due to MyPY
         return self._async_connection
 
-    async def aprocess_notify(self, connected_callback: Optional[Callable] = None) -> AsyncGenerator[tuple[str, str], None]:  # public
+    def get_listen_query(self, channel: str) -> psycopg.sql.Composed:
+        """Returns SQL command for listening on pg_notify channel
+
+        This uses the psycopg utilities which ensure correct escaping so SQL injection is not possible.
+        Return value is a valid argument for cursor.execute()
+        """
+        return psycopg.sql.SQL("LISTEN {};").format(psycopg.sql.Identifier(channel))
+
+    async def aprocess_notify(
+        self, connected_callback: Optional[Callable[[], Coroutine[Any, Any, None]]] = None
+    ) -> AsyncGenerator[tuple[str, str], None]:  # public
         connection = await self.aget_connection()
         async with connection.cursor() as cur:
             for channel in self.channels:
-                await cur.execute(f"LISTEN {channel};")
+                await cur.execute(self.get_listen_query(channel))
                 logger.info(f"Set up pg_notify listening on channel '{channel}'")
 
             if connected_callback:
@@ -178,6 +188,27 @@ def get_connection(self) -> psycopg.Connection:
             return connection
         return self._sync_connection
 
+    def process_notify(self, connected_callback: Optional[Callable] = None, timeout: float = 5.0, max_messages: int = 1) -> Iterator[tuple[str, str]]:
+        """Blocking method that listens for messages on subscribed pg_notify channels until timeout
+
+        This has two different exit conditions:
+        - received max_messages number of messages or more
+        - taken longer than the specified timeout condition
+        """
+        connection = self.get_connection()
+
+        with connection.cursor() as cur:
+            for channel in self.channels:
+                cur.execute(self.get_listen_query(channel))
+                logger.info(f"Set up pg_notify listening on channel '{channel}'")
+
+            if connected_callback:
+                connected_callback()
+
+            logger.debug('Starting listening for pg_notify notifications')
+            for notify in connection.notifies(timeout=timeout, stop_after=max_messages):
+                yield (notify.channel, notify.payload)
+
     def publish_message(self, channel: Optional[str] = None, message: str = '') -> None:
         connection = self.get_connection()
         channel = self.get_publish_channel(channel)

dispatcher/control.py

@@ -116,22 +116,27 @@ async def acontrol(self, command, data=None):
         producer = self.make_producer(Control.generate_reply_queue_name())  # reply queue not used
         await control_callbacks.connected_callback(producer)
 
-    def control_with_reply(self, command, expected_replies=1, timeout=1, data=None):
+    def control_with_reply(self, command: str, expected_replies: int = 1, timeout: float = 1.0, data: Optional[dict] = None) -> list[dict]:
         logger.info('control-and-reply {} to {}'.format(command, self.queuename))
         start = time.time()
         reply_queue = Control.generate_reply_queue_name()
         send_data = {'control': command, 'reply_to': reply_queue}
         if data:
             send_data['control_data'] = data
 
-        producer = self.make_producer(reply_queue)
+        broker = get_broker(self.broker_name, self.broker_config, channels=[reply_queue])
 
-        loop = asyncio.new_event_loop()
-        try:
-            replies = loop.run_until_complete(self.acontrol_with_reply_internal(producer, send_data, expected_replies, timeout))
-        finally:
-            loop.close()
-            loop = None
+        def connected_callback():
+            payload = json.dumps(send_data)
+            if self.queuename:
+                broker.publish_message(channel=self.queuename, message=payload)
+            else:
+                broker.publish_message(message=payload)
+
+        replies = []
+        for channel, payload in broker.process_notify(connected_callback=connected_callback, max_messages=expected_replies, timeout=timeout):
+            reply_data = json.loads(payload)
+            replies.append(reply_data)
 
         logger.info(f'control-and-reply message returned in {time.time() - start} seconds')
         return replies

pyproject.toml

@@ -53,7 +53,7 @@ profile = "black"
 line_length = 160
 
 [project.optional-dependencies]
-pg_notify = ["psycopg[binary]"]
+pg_notify = ["psycopg[binary]>=3.2.0"]
 
 [tool.pytest.ini_options]
 log_cli_level = "DEBUG"

tests/integration/brokers/test_pg_notify.py

@@ -1,3 +1,6 @@
+import time
+import multiprocessing
+
 import pytest
 
 from dispatcher.brokers.pg_notify import Broker, create_connection, acreate_connection
@@ -16,6 +19,78 @@ def test_sync_connection_from_config_reuse(conn_config):
     assert conn is not create_connection(**conn_config)
 
 
+def test_sync_listen_timeout(conn_config):
+    broker = Broker(config=conn_config)
+    timeout_value = 0.05
+    start = time.monotonic()
+    assert list(broker.process_notify(timeout=timeout_value)) == []
+    delta = time.monotonic() - start
+    assert delta > timeout_value
+
+
+def _send_message(conn_config):
+    broker = Broker(config=conn_config)
+    if broker._sync_connection:
+        broker._sync_connection.close()
+
+    broker.publish_message('test_sync_listen_receive', 'test_message')
+
+
+def test_sync_listen_receive(conn_config):
+    messages = []
+    with multiprocessing.Pool(processes=1) as pool:
+        def send_from_subprocess():
+            pool.apply(_send_message, args=(conn_config,))
+
+        broker = Broker(config=conn_config, channels=('test_sync_listen_receive',))
+        timeout_value = 2.0
+        start = time.monotonic()
+        for channel, message in broker.process_notify(connected_callback=send_from_subprocess, timeout=timeout_value):
+            messages.append(message)
+        delta = time.monotonic() - start
+
+    assert messages == ['test_message']
+    assert delta < timeout_value
+
+
+def test_sync_listen_receive_multi_message(conn_config):
+    """Tests that the expected messages exit condition works, we get 3 messages, not just 1"""
+    messages = []
+    with multiprocessing.Pool(processes=1) as pool:
+        def send_from_subprocess():
+            pool.apply(_send_message, args=(conn_config,))
+            pool.apply(_send_message, args=(conn_config,))
+            pool.apply(_send_message, args=(conn_config,))
+
+        broker = Broker(config=conn_config, channels=('test_sync_listen_receive',))
+        timeout_value = 2.0
+        start = time.monotonic()
+        for channel, message in broker.process_notify(connected_callback=send_from_subprocess, max_messages=3):
+            messages.append(message)
+        delta = time.monotonic() - start
+
+    assert messages == ['test_message' for i in range(3)]
+    assert delta < timeout_value
+
+
+def test_get_message_then_timeout(conn_config):
+    """Tests that the expected messages exit condition works, we get 3 messages, not just 1"""
+    messages = []
+    with multiprocessing.Pool(processes=1) as pool:
+        def send_from_subprocess():
+            pool.apply(_send_message, args=(conn_config,))
+
+        broker = Broker(config=conn_config, channels=('test_sync_listen_receive',))
+        timeout_value = 0.5
+        start = time.monotonic()
+        for channel, message in broker.process_notify(connected_callback=send_from_subprocess, timeout=timeout_value, max_messages=2):
+            messages.append(message)
+        delta = time.monotonic() - start
+
+    assert messages == ['test_message']
+    assert delta > timeout_value  # goes until timeout
+
+
 @pytest.mark.asyncio
 async def test_async_connection_from_config_reuse(conn_config):
     broker = Broker(config=conn_config)