Updated redis broker to operate with raw bytes.

Signed-off-by: Pavel Kirilin <[email protected]>

Author: s3rius

taskiq_redis/redis_broker.py

@@ -1,5 +1,3 @@
-import pickle
-from abc import abstractmethod
 from logging import getLogger
 from typing import Any, AsyncGenerator, Callable, Optional, TypeVar
 
@@ -52,66 +50,63 @@ async def shutdown(self) -> None:
         await super().shutdown()
         await self.connection_pool.disconnect()
 
-    async def listen(self) -> AsyncGenerator[BrokerMessage, None]:
-        """
-        Listen redis queue for new messages.
 
-        This function listens to the queue
-        and yields new messages if they have BrokerMessage type.
+class PubSubBroker(BaseRedisBroker):
+    """Broker that works with Redis and broadcasts tasks to all workers."""
 
-        :yields: broker messages.
-        """
-        async for message in self._listen_to_raw_messages():
-            try:
-                redis_message = pickle.loads(message)
-                if isinstance(redis_message, BrokerMessage):
-                    yield redis_message
-            except (
-                TypeError,
-                AttributeError,
-                pickle.UnpicklingError,
-            ) as exc:
-                logger.debug(
-                    "Cannot read broker message %s",
-                    exc,
-                    exc_info=True,
-                )
-
-    @abstractmethod
-    async def _listen_to_raw_messages(self) -> AsyncGenerator[bytes, None]:
+    async def kick(self, message: BrokerMessage) -> None:
         """
-        Generator for reading raw data from Redis.
+        Publish message over PUBSUB channel.
 
-        :yields: raw data.
+        :param message: message to send.
         """
-        yield  # type: ignore
-
+        async with Redis(connection_pool=self.connection_pool) as redis_conn:
+            await redis_conn.publish(self.queue_name, message.message)
 
-class PubSubBroker(BaseRedisBroker):
-    """Broker that works with Redis and broadcasts tasks to all workers."""
+    async def listen(self) -> AsyncGenerator[bytes, None]:
+        """
+        Listen redis queue for new messages.
 
-    async def kick(self, message: BrokerMessage) -> None:  # noqa: D102
-        async with Redis(connection_pool=self.connection_pool) as redis_conn:
-            await redis_conn.publish(self.queue_name, pickle.dumps(message))
+        This function listens to the pubsub channel
+        and yields all messages with proper types.
 
-    async def _listen_to_raw_messages(self) -> AsyncGenerator[bytes, None]:
+        :yields: broker messages.
+        """
         async with Redis(connection_pool=self.connection_pool) as redis_conn:
             redis_pubsub_channel = redis_conn.pubsub()
             await redis_pubsub_channel.subscribe(self.queue_name)
             async for message in redis_pubsub_channel.listen():
                 if not message:
                     continue
+                if message["type"] != "message":
+                    logger.debug("Received non-message from redis: %s", message)
+                    continue
                 yield message["data"]
 
 
 class ListQueueBroker(BaseRedisBroker):
     """Broker that works with Redis and distributes tasks between workers."""
 
-    async def kick(self, message: BrokerMessage) -> None:  # noqa: D102
+    async def kick(self, message: BrokerMessage) -> None:
+        """
+        Put a message in a list.
+
+        This method appends a message to the list of all messages.
+
+        :param message: message to append.
+        """
         async with Redis(connection_pool=self.connection_pool) as redis_conn:
-            await redis_conn.lpush(self.queue_name, pickle.dumps(message))
+            await redis_conn.lpush(self.queue_name, message.message)
 
-    async def _listen_to_raw_messages(self) -> AsyncGenerator[bytes, None]:
+    async def listen(self) -> AsyncGenerator[bytes, None]:
+        """
+        Listen redis queue for new messages.
+
+        This function listens to the queue
+        and yields new messages if they have BrokerMessage type.
+
+        :yields: broker messages.
+        """
         redis_brpop_data_position = 1
         async with Redis(connection_pool=self.connection_pool) as redis_conn:
             while True:  # noqa: WPS457

tests/test_broker.py

@@ -7,7 +7,7 @@
 from taskiq_redis import ListQueueBroker, PubSubBroker
 
 
-async def get_message(broker: AsyncBroker) -> BrokerMessage:  # type: ignore
+async def get_message(broker: AsyncBroker) -> bytes:  # type: ignore
     """
     Get a message from the broker.
 
@@ -56,7 +56,7 @@ async def test_pub_sub_broker(
 
     message1 = worker1_task.result()
     message2 = worker2_task.result()
-    assert message1 == valid_broker_message
+    assert message1 == valid_broker_message.message
     assert message1 == message2
 
 
@@ -81,4 +81,4 @@ async def test_list_queue_broker(
 
     assert worker1_task.done() != worker2_task.done()
     message = worker1_task.result() if worker1_task.done() else worker2_task.result()
-    assert message == valid_broker_message
+    assert message == valid_broker_message.message