Finish 0.2.0

Author: chandr-andr

.github/workflows/test.yml

@@ -1,6 +1,11 @@
 name: Testing taskiq-redis
 
-on: push
+on:
+  push:
+    branches:
+      - main
+      - develop
+  pull_request:
 
 jobs:
   black:
@@ -78,3 +83,12 @@ jobs:
           POETRY_VIRTUALENVS_CREATE: false
       - name: Run pytest check
         run: poetry run pytest -vv -n auto --cov="taskiq_redis" .
+      - name: Generate report
+        run: poetry run coverage xml
+      - name: Upload coverage reports to Codecov with GitHub Action
+        uses: codecov/codecov-action@v3
+        if: matrix.py_version == '3.9'
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          fail_ci_if_error: true
+          verbose: true

README.md

@@ -16,18 +16,19 @@ pip install taskiq-redis
 # Usage
 
 Let's see the example with the redis broker and redis async result:
+
 ```python
 import asyncio
 
-from taskiq_redis.redis_broker import RedisBroker
+from taskiq_redis.redis_broker import ListQueueBroker
 from taskiq_redis.redis_backend import RedisAsyncResultBackend
 
-
 redis_async_result = RedisAsyncResultBackend(
     redis_url="redis://localhost:6379",
 )
 
-broker = RedisBroker(
+# Or you can use PubSubBroker if you need broadcasting
+broker = ListQueueBroker(
     url="redis://localhost:6379",
     result_backend=redis_async_result,
 )
@@ -48,9 +49,17 @@ async def main():
 asyncio.run(main())
 ```
 
-## RedisBroker configuration
+## PubSubBroker and ListQueueBroker configuration
+
+We have two brokers with similar interfaces, but with different logic.
+The PubSubBroker uses redis' pubsub mechanism and is very powerful,
+but it executes every task on all workers, because PUBSUB broadcasts message
+to all subscribers.
+
+If you want your messages to be processed only once, please use ListQueueBroker.
+It uses redis' [LPUSH](https://redis.io/commands/lpush/) and [BRPOP](https://redis.io/commands/brpop/) commands to deal with messages.
 
-RedisBroker parameters:
+Brokers parameters:
 * `url` - url to redis.
 * `task_id_generator` - custom task_id genertaor.
 * `result_backend` - custom result backend.

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "taskiq-redis"
-version = "0.1.0"
+version = "0.2.0"
 description = "Redis integration for taskiq"
 authors = ["taskiq-team <[email protected]>"]
 readme = "README.md"

taskiq_redis/__init__.py

@@ -1,5 +1,5 @@
 """Package for redis integration."""
 from taskiq_redis.redis_backend import RedisAsyncResultBackend
-from taskiq_redis.redis_broker import RedisBroker
+from taskiq_redis.redis_broker import ListQueueBroker
 
-__all__ = ["RedisAsyncResultBackend", "RedisBroker"]
+__all__ = ["RedisAsyncResultBackend", "ListQueueBroker"]

taskiq_redis/redis_broker.py

@@ -1,4 +1,5 @@
 import pickle
+from abc import abstractmethod
 from logging import getLogger
 from typing import Any, AsyncGenerator, Callable, Optional, TypeVar
 
@@ -12,8 +13,8 @@
 logger = getLogger("taskiq.redis_broker")
 
 
-class RedisBroker(AsyncBroker):
-    """Broker that works with Redis."""
+class BaseRedisBroker(AsyncBroker):
+    """Base broker that works with Redis."""
 
     def __init__(
         self,
@@ -44,31 +45,12 @@ def __init__(
             max_connections=max_connection_pool_size,
             **connection_kwargs,
         )
-
-        self.redis_pubsub_channel = queue_name
+        self.queue_name = queue_name
 
     async def shutdown(self) -> None:
         """Closes redis connection pool."""
         await self.connection_pool.disconnect()
 
-    async def kick(self, message: BrokerMessage) -> None:
-        """
-        Sends a message to the redis broker list.
-
-        This function constructs message for redis
-        and sends it.
-
-        The message is pickled dict object with message,
-        task_id, task_name and labels.
-
-        :param message: message to send.
-        """
-        async with Redis(connection_pool=self.connection_pool) as redis_conn:
-            await redis_conn.publish(
-                self.redis_pubsub_channel,
-                pickle.dumps(message),
-            )
-
     async def listen(self) -> AsyncGenerator[BrokerMessage, None]:
         """
         Listen redis queue for new messages.
@@ -78,24 +60,60 @@ async def listen(self) -> AsyncGenerator[BrokerMessage, None]:
 
         :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]:
+        """
+        Generator for reading raw data from Redis.
+
+        :yields: raw data.
+        """
+        yield  # type: ignore
+
+
+class PubSubBroker(BaseRedisBroker):
+    """Broker that works with Redis and broadcasts tasks to all workers."""
+
+    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))
+
+    async def _listen_to_raw_messages(self) -> AsyncGenerator[bytes, None]:
         async with Redis(connection_pool=self.connection_pool) as redis_conn:
             redis_pubsub_channel = redis_conn.pubsub()
-            await redis_pubsub_channel.subscribe(self.redis_pubsub_channel)
+            await redis_pubsub_channel.subscribe(self.queue_name)
             async for message in redis_pubsub_channel.listen():
-                if message:
-                    try:
-                        redis_message = pickle.loads(
-                            message["data"],
-                        )
-                        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,
-                        )
+                if not 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 with Redis(connection_pool=self.connection_pool) as redis_conn:
+            await redis_conn.lpush(self.queue_name, pickle.dumps(message))
+
+    async def _listen_to_raw_messages(self) -> AsyncGenerator[bytes, None]:
+        redis_brpop_data_position = 1
+        async with Redis(connection_pool=self.connection_pool) as redis_conn:
+            while True:  # noqa: WPS457
+                yield (await redis_conn.brpop(self.queue_name))[
+                    redis_brpop_data_position
+                ]

tests/test_broker.py

@@ -0,0 +1,84 @@
+import asyncio
+import uuid
+
+import pytest
+from taskiq import AsyncBroker, BrokerMessage
+
+from taskiq_redis.redis_broker import ListQueueBroker, PubSubBroker
+
+
+async def get_message(broker: AsyncBroker) -> BrokerMessage:  # type: ignore
+    """
+    Get a message from the broker.
+
+    :param broker: async message broker.
+    :return: first message from listen method.
+    """
+    async for message in broker.listen():  # noqa: WPS328
+        return message
+
+
+@pytest.fixture
+def valid_broker_message() -> BrokerMessage:
+    """
+    Generate valid broker message for tests.
+
+    :returns: broker message.
+    """
+    return BrokerMessage(
+        task_id=uuid.uuid4().hex,
+        task_name=uuid.uuid4().hex,
+        message="my_msg",
+        labels={
+            "label1": "val1",
+        },
+    )
+
+
+@pytest.mark.anyio
+async def test_pub_sub_broker(
+    valid_broker_message: BrokerMessage,
+    redis_url: str,
+) -> None:
+    """
+    Test that messages are published and read correctly by PubSubBroker.
+
+    We create two workers that listen and send a message to them.
+    Expect both workers to receive the same message we sent.
+    """
+    broker = PubSubBroker(url=redis_url, queue_name=uuid.uuid4().hex)
+    worker1_task = asyncio.create_task(get_message(broker))
+    worker2_task = asyncio.create_task(get_message(broker))
+    await asyncio.sleep(0.3)
+
+    await broker.kick(valid_broker_message)
+    await asyncio.sleep(0.3)
+
+    message1 = worker1_task.result()
+    message2 = worker2_task.result()
+    assert message1 == valid_broker_message
+    assert message1 == message2
+
+
+@pytest.mark.anyio
+async def test_list_queue_broker(
+    valid_broker_message: BrokerMessage,
+    redis_url: str,
+) -> None:
+    """
+    Test that messages are published and read correctly by ListQueueBroker.
+
+    We create two workers that listen and send a message to them.
+    Expect only one worker to receive the same message we sent.
+    """
+    broker = ListQueueBroker(url=redis_url, queue_name=uuid.uuid4().hex)
+    worker1_task = asyncio.create_task(get_message(broker))
+    worker2_task = asyncio.create_task(get_message(broker))
+    await asyncio.sleep(0.3)
+
+    await broker.kick(valid_broker_message)
+    await asyncio.sleep(0.3)
+
+    assert worker1_task.done() != worker2_task.done()
+    message = worker1_task.result() if worker1_task.done() else worker2_task.result()
+    assert message == valid_broker_message