Initial stream support.

Author: s3rius

docker-compose.yml

@@ -1,8 +1,6 @@
-version: '3.2'
-
 services:
   redis:
-    image: bitnami/redis:6.2.5
+    image: bitnami/redis:7.4.2
     environment:
       ALLOW_EMPTY_PASSWORD: "yes"
     healthcheck:
@@ -14,7 +12,7 @@ services:
     ports:
       - 7000:6379
   redis-node-0: &redis-node
-    image: docker.io/bitnami/redis-cluster:7.2
+    image: docker.io/bitnami/redis-cluster:7.4.2
     environment:
       ALLOW_EMPTY_PASSWORD: "yes"
       REDIS_NODES: "redis-node-0 redis-node-1 redis-node-2 redis-node-3 redis-node-4 redis-node-5"
@@ -38,7 +36,7 @@ services:
     <<: *redis-node
 
   redis-node-5:
-    image: docker.io/bitnami/redis-cluster:7.2
+    image: docker.io/bitnami/redis-cluster:7.4.2
     depends_on:
       - redis-node-0
       - redis-node-1
@@ -60,7 +58,7 @@ services:
       - 7001:6379
 
   redis-master:
-    image: bitnami/redis:6.2.5
+    image: bitnami/redis:7.4.2
     environment:
       ALLOW_EMPTY_PASSWORD: "yes"
     healthcheck:
@@ -71,7 +69,7 @@ services:
       start_period: 10s
 
   redis-sentinel:
-    image: bitnami/redis-sentinel:latest
+    image: bitnami/redis-sentinel:7.4.2
     depends_on:
       - redis-master
     environment:

poetry.lock

@@ -25,22 +25,21 @@ keywords = [
 ]
 
 [tool.poetry.dependencies]
-python = "^3.8.1"
-taskiq = ">=0.11.1,<1"
+python = "^3.9"
+taskiq = ">=0.11.12,<1"
 redis = "^5"
 
 [tool.poetry.group.dev.dependencies]
-pytest = "^7.0"
+pytest = "^8"
 mypy = "^1"
-black = "^22.3.0"
-pytest-cov = "^3.0.0"
-anyio = "^3.6.1"
-pytest-env = "^0.6.2"
+black = "^25"
+pytest-cov = "^6"
+anyio = "^4"
+pytest-env = "^1"
 fakeredis = "^2"
-pre-commit = "^2.20.0"
-pytest-xdist = { version = "^2.5.0", extras = ["psutil"] }
+pre-commit = "^4"
+pytest-xdist = { version = "^3", extras = ["psutil"] }
 ruff = "^0"
-types-redis = "^4.6.0.20240425"
 
 [tool.mypy]
 strict = true
@@ -65,7 +64,7 @@ build-backend = "poetry.core.masonry.api"
 [tool.ruff]
 # List of enabled rulsets.
 # See https://docs.astral.sh/ruff/rules/ for more information.
-select = [
+lint.select = [
     "E",   # Error
     "F",   # Pyflakes
     "W",   # Pycodestyle
@@ -92,24 +91,22 @@ select = [
     "PL",  # PyLint checks
     "RUF", # Specific to Ruff checks
 ]
-ignore = [
+lint.ignore = [
     "D105",    # Missing docstring in magic method
     "D107",    # Missing docstring in __init__
     "D212",    # Multi-line docstring summary should start at the first line
     "D401",    # First line should be in imperative mood
     "D104",    # Missing docstring in public package
     "D100",    # Missing docstring in public module
-    "ANN102",  # Missing type annotation for self in method
-    "ANN101",  # Missing type annotation for argument
     "ANN401",  # typing.Any are disallowed in `**kwargs
     "PLR0913", # Too many arguments for function call
     "D106",    # Missing docstring in public nested class
 ]
 exclude = [".venv/"]
-mccabe = { max-complexity = 10 }
+lint.mccabe = { max-complexity = 10 }
 line-length = 88
 
-[tool.ruff.per-file-ignores]
+[tool.ruff.lint.per-file-ignores]
 "tests/*" = [
     "S101",   # Use of assert detected
     "S301",   # Use of pickle detected
@@ -119,12 +116,12 @@ line-length = 88
     "D101",   # Missing docstring in public class
 ]
 
-[tool.ruff.pydocstyle]
+[tool.ruff.lint.pydocstyle]
 convention = "pep257"
 ignore-decorators = ["typing.overload"]
 
-[tool.ruff.pylint]
+[tool.ruff.lint.pylint]
 allow-magic-value-types = ["int", "str", "float"]
 
-[tool.ruff.flake8-bugbear]
+[tool.ruff.lint.flake8-bugbear]
 extend-immutable-calls = ["taskiq_dependencies.Depends", "taskiq.TaskiqDepends"]

pyproject.toml

@@ -1,14 +1,19 @@
 """Package for redis integration."""
+
 from taskiq_redis.redis_backend import (
     RedisAsyncClusterResultBackend,
     RedisAsyncResultBackend,
     RedisAsyncSentinelResultBackend,
 )
-from taskiq_redis.redis_broker import ListQueueBroker, PubSubBroker
-from taskiq_redis.redis_cluster_broker import ListQueueClusterBroker
+from taskiq_redis.redis_broker import ListQueueBroker, PubSubBroker, RedisStreamBroker
+from taskiq_redis.redis_cluster_broker import (
+    ListQueueClusterBroker,
+    RedisStreamClusterBroker,
+)
 from taskiq_redis.redis_sentinel_broker import (
     ListQueueSentinelBroker,
     PubSubSentinelBroker,
+    RedisStreamSentinelBroker,
 )
 from taskiq_redis.schedule_source import (
     RedisClusterScheduleSource,
@@ -17,15 +22,18 @@
 )
 
 __all__ = [
-    "RedisAsyncClusterResultBackend",
-    "RedisAsyncResultBackend",
-    "RedisAsyncSentinelResultBackend",
     "ListQueueBroker",
-    "PubSubBroker",
     "ListQueueClusterBroker",
     "ListQueueSentinelBroker",
+    "PubSubBroker",
     "PubSubSentinelBroker",
-    "RedisScheduleSource",
+    "RedisAsyncClusterResultBackend",
+    "RedisAsyncResultBackend",
+    "RedisAsyncSentinelResultBackend",
     "RedisClusterScheduleSource",
+    "RedisScheduleSource",
     "RedisSentinelScheduleSource",
+    "RedisStreamBroker",
+    "RedisStreamClusterBroker",
+    "RedisStreamSentinelBroker",
 ]

taskiq_redis/__init__.py

@@ -8,10 +8,16 @@ class TaskIQRedisError(TaskiqError):
 class DuplicateExpireTimeSelectedError(ResultBackendError, TaskIQRedisError):
     """Error if two lifetimes are selected."""
 
+    __template__ = "Choose either result_ex_time or result_px_time."
+
 
 class ExpireTimeMustBeMoreThanZeroError(ResultBackendError, TaskIQRedisError):
     """Error if two lifetimes are less or equal zero."""
 
+    __template__ = (
+        "You must select one expire time param and it must be more than zero."
+    )
+
 
 class ResultIsMissingError(TaskIQRedisError, ResultGetError):
     """Error if there is no result when trying to get it."""

taskiq_redis/exceptions.py

@@ -16,10 +16,10 @@
 from redis.asyncio.cluster import RedisCluster
 from redis.asyncio.connection import Connection
 from taskiq import AsyncResultBackend
-from taskiq.abc.result_backend import TaskiqResult
 from taskiq.abc.serializer import TaskiqSerializer
 from taskiq.compat import model_dump, model_validate
 from taskiq.depends.progress_tracker import TaskProgress
+from taskiq.result import TaskiqResult
 from taskiq.serializers import PickleSerializer
 
 from taskiq_redis.exceptions import (
@@ -92,14 +92,10 @@ def __init__(
             ),
         )
         if unavailable_conditions:
-            raise ExpireTimeMustBeMoreThanZeroError(
-                "You must select one expire time param and it must be more than zero.",
-            )
+            raise ExpireTimeMustBeMoreThanZeroError
 
         if self.result_ex_time and self.result_px_time:
-            raise DuplicateExpireTimeSelectedError(
-                "Choose either result_ex_time or result_px_time.",
-            )
+            raise DuplicateExpireTimeSelectedError
 
     def _task_name(self, task_id: str) -> str:
         if self.prefix_str is None:

taskiq_redis/redis_backend.py

@@ -1,8 +1,19 @@
 import sys
+import uuid
 from logging import getLogger
-from typing import TYPE_CHECKING, Any, AsyncGenerator, Callable, Optional, TypeVar
-
-from redis.asyncio import BlockingConnectionPool, Connection, Redis
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncGenerator,
+    Awaitable,
+    Callable,
+    Dict,
+    Optional,
+    TypeVar,
+)
+
+from redis.asyncio import BlockingConnectionPool, Connection, Redis, ResponseError
+from taskiq import AckableMessage
 from taskiq.abc.broker import AsyncBroker
 from taskiq.abc.result_backend import AsyncResultBackend
 from taskiq.message import BrokerMessage
@@ -132,3 +143,110 @@ async def listen(self) -> AsyncGenerator[bytes, None]:
             except ConnectionError as exc:
                 logger.warning("Redis connection error: %s", exc)
                 continue
+
+
+class RedisStreamBroker(BaseRedisBroker):
+    """
+    Redis broker that uses streams for task distribution.
+
+    You can read more about streams here:
+    https://redis.io/docs/latest/develop/data-types/streams
+
+    This broker supports acknowledgment of messages.
+    """
+
+    def __init__(
+        self,
+        url: str,
+        queue_name: str = "taskiq",
+        max_connection_pool_size: Optional[int] = None,
+        consumer_group_name: str = "taskiq",
+        consumer_name: Optional[str] = None,
+        consumer_id: str = "$",
+        mkstream: bool = True,
+        xread_block: int = 10000,
+        additional_streams: Optional[Dict[str, str]] = None,
+        **connection_kwargs: Any,
+    ) -> None:
+        super().__init__(
+            url,
+            task_id_generator=None,
+            result_backend=None,
+            queue_name=queue_name,
+            max_connection_pool_size=max_connection_pool_size,
+            **connection_kwargs,
+        )
+        self.consumer_group_name = consumer_group_name
+        self.consumer_name = consumer_name or str(uuid.uuid4())
+        self.consumer_id = consumer_id
+        self.mkstream = mkstream
+        self.block = xread_block
+        self.additional_streams = additional_streams or {}
+
+    async def _declare_consumer_group(self) -> None:
+        """
+        Declare consumber group.
+
+        Required for proper work of the broker.
+        """
+        streams = {self.queue_name, *self.additional_streams.keys()}
+        async with Redis(connection_pool=self.connection_pool) as redis_conn:
+            for stream_name in streams:
+                try:
+                    await redis_conn.xgroup_create(
+                        stream_name,
+                        self.consumer_group_name,
+                        id=self.consumer_id,
+                        mkstream=self.mkstream,
+                    )
+                except ResponseError as err:
+                    logger.debug(err)
+
+    async def startup(self) -> None:
+        """Declare consumer group on startup."""
+        await super().startup()
+        await self._declare_consumer_group()
+
+    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.xadd(self.queue_name, {b"data": message.message})
+
+    def _ack_generator(self, id: str) -> Callable[[], Awaitable[None]]:
+        async def _ack() -> None:
+            async with Redis(connection_pool=self.connection_pool) as redis_conn:
+                await redis_conn.xack(
+                    self.queue_name,
+                    self.consumer_group_name,
+                    id,
+                )
+
+        return _ack
+
+    async def listen(self) -> AsyncGenerator[AckableMessage, None]:
+        """Listen to incoming messages."""
+        async with Redis(connection_pool=self.connection_pool) as redis_conn:
+            while True:
+                fetched = await redis_conn.xreadgroup(
+                    self.consumer_group_name,
+                    self.consumer_name,
+                    {
+                        self.queue_name: ">",
+                        **self.additional_streams,
+                    },
+                    block=self.block,
+                    noack=False,
+                )
+                for _, msg_list in fetched:
+                    for msg_id, msg in msg_list:
+                        logger.debug("Received message: %s", msg)
+                        yield AckableMessage(
+                            data=msg[b"data"],
+                            ack=self._ack_generator(msg_id),
+                        )