taskiq-redis 1.0.2__py3-none-any.whl → 1.0.3__py3-none-any.whl

taskiq_redis/__init__.py

@@ -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 @@ from taskiq_redis.schedule_source import (
 )
 
 __all__ = [
-    "RedisAsyncClusterResultBackend",
-    "RedisAsyncResultBackend",
-    "RedisAsyncSentinelResultBackend",
     "ListQueueBroker",
-    "PubSubBroker",
     "ListQueueClusterBroker",
     "ListQueueSentinelBroker",
+    "PubSubBroker",
     "PubSubSentinelBroker",
-    "RedisScheduleSource",
+    "RedisAsyncClusterResultBackend",
+    "RedisAsyncResultBackend",
+    "RedisAsyncSentinelResultBackend",
     "RedisClusterScheduleSource",
+    "RedisScheduleSource",
     "RedisSentinelScheduleSource",
+    "RedisStreamBroker",
+    "RedisStreamClusterBroker",
+    "RedisStreamSentinelBroker",
 ]

taskiq_redis/exceptions.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/redis_backend.py

@@ -16,10 +16,10 @@ from redis.asyncio import BlockingConnectionPool, Redis, Sentinel
 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 (
@@ -34,8 +34,8 @@ else:
     from typing_extensions import TypeAlias
 
 if TYPE_CHECKING:
-    _Redis: TypeAlias = Redis[bytes]
-    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool[Connection]
+    _Redis: TypeAlias = Redis[bytes]  # type: ignore
+    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool[Connection]  # type: ignore
 else:
     _Redis: TypeAlias = Redis
     _BlockingConnectionPool: TypeAlias = BlockingConnectionPool
@@ -56,6 +56,7 @@ class RedisAsyncResultBackend(AsyncResultBackend[_ReturnType]):
         result_px_time: Optional[int] = None,
         max_connection_pool_size: Optional[int] = None,
         serializer: Optional[TaskiqSerializer] = None,
+        prefix_str: Optional[str] = None,
         **connection_kwargs: Any,
     ) -> None:
         """
@@ -82,6 +83,7 @@ class RedisAsyncResultBackend(AsyncResultBackend[_ReturnType]):
         self.keep_results = keep_results
         self.result_ex_time = result_ex_time
         self.result_px_time = result_px_time
+        self.prefix_str = prefix_str
 
         unavailable_conditions = any(
             (
@@ -90,14 +92,15 @@ class RedisAsyncResultBackend(AsyncResultBackend[_ReturnType]):
             ),
         )
         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:
+            return task_id
+        return f"{self.prefix_str}:{task_id}"
 
     async def shutdown(self) -> None:
         """Closes redis connection."""
@@ -119,7 +122,7 @@ class RedisAsyncResultBackend(AsyncResultBackend[_ReturnType]):
         :param result: TaskiqResult instance.
         """
         redis_set_params: Dict[str, Union[str, int, bytes]] = {
-            "name": task_id,
+            "name": self._task_name(task_id),
             "value": self.serializer.dumpb(model_dump(result)),
         }
         if self.result_ex_time:
@@ -139,7 +142,7 @@ class RedisAsyncResultBackend(AsyncResultBackend[_ReturnType]):
         :returns: True if the result is ready else False.
         """
         async with Redis(connection_pool=self.redis_pool) as redis:
-            return bool(await redis.exists(task_id))
+            return bool(await redis.exists(self._task_name(task_id)))
 
     async def get_result(
         self,
@@ -154,14 +157,15 @@ class RedisAsyncResultBackend(AsyncResultBackend[_ReturnType]):
         :raises ResultIsMissingError: if there is no result when trying to get it.
         :return: task's return value.
         """
+        task_name = self._task_name(task_id)
         async with Redis(connection_pool=self.redis_pool) as redis:
             if self.keep_results:
                 result_value = await redis.get(
-                    name=task_id,
+                    name=task_name,
                 )
             else:
                 result_value = await redis.getdel(
-                    name=task_id,
+                    name=task_name,
                 )
 
         if result_value is None:
@@ -192,7 +196,7 @@ class RedisAsyncResultBackend(AsyncResultBackend[_ReturnType]):
         :param result: task's TaskProgress instance.
         """
         redis_set_params: Dict[str, Union[str, int, bytes]] = {
-            "name": task_id + PROGRESS_KEY_SUFFIX,
+            "name": self._task_name(task_id) + PROGRESS_KEY_SUFFIX,
             "value": self.serializer.dumpb(model_dump(progress)),
         }
         if self.result_ex_time:
@@ -215,7 +219,7 @@ class RedisAsyncResultBackend(AsyncResultBackend[_ReturnType]):
         """
         async with Redis(connection_pool=self.redis_pool) as redis:
             result_value = await redis.get(
-                name=task_id + PROGRESS_KEY_SUFFIX,
+                name=self._task_name(task_id) + PROGRESS_KEY_SUFFIX,
             )
 
         if result_value is None:
@@ -237,6 +241,7 @@ class RedisAsyncClusterResultBackend(AsyncResultBackend[_ReturnType]):
         result_ex_time: Optional[int] = None,
         result_px_time: Optional[int] = None,
         serializer: Optional[TaskiqSerializer] = None,
+        prefix_str: Optional[str] = None,
         **connection_kwargs: Any,
     ) -> None:
         """
@@ -253,7 +258,7 @@ class RedisAsyncClusterResultBackend(AsyncResultBackend[_ReturnType]):
         :raises ExpireTimeMustBeMoreThanZeroError: if result_ex_time
             and result_px_time are equal zero.
         """
-        self.redis: RedisCluster[bytes] = RedisCluster.from_url(
+        self.redis: "RedisCluster" = RedisCluster.from_url(
             redis_url,
             **connection_kwargs,
         )
@@ -261,6 +266,7 @@ class RedisAsyncClusterResultBackend(AsyncResultBackend[_ReturnType]):
         self.keep_results = keep_results
         self.result_ex_time = result_ex_time
         self.result_px_time = result_px_time
+        self.prefix_str = prefix_str
 
         unavailable_conditions = any(
             (
@@ -269,18 +275,19 @@ class RedisAsyncClusterResultBackend(AsyncResultBackend[_ReturnType]):
             ),
         )
         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:
+            return task_id
+        return f"{self.prefix_str}:{task_id}"
 
     async def shutdown(self) -> None:
         """Closes redis connection."""
-        await self.redis.aclose()  # type: ignore[attr-defined]
+        await self.redis.aclose()
         await super().shutdown()
 
     async def set_result(
@@ -298,7 +305,7 @@ class RedisAsyncClusterResultBackend(AsyncResultBackend[_ReturnType]):
         :param result: TaskiqResult instance.
         """
         redis_set_params: Dict[str, Union[str, bytes, int]] = {
-            "name": task_id,
+            "name": self._task_name(task_id),
             "value": self.serializer.dumpb(model_dump(result)),
         }
         if self.result_ex_time:
@@ -316,7 +323,7 @@ class RedisAsyncClusterResultBackend(AsyncResultBackend[_ReturnType]):
 
         :returns: True if the result is ready else False.
         """
-        return bool(await self.redis.exists(task_id))  # type: ignore[attr-defined]
+        return bool(await self.redis.exists(self._task_name(task_id)))
 
     async def get_result(
         self,
@@ -331,13 +338,14 @@ class RedisAsyncClusterResultBackend(AsyncResultBackend[_ReturnType]):
         :raises ResultIsMissingError: if there is no result when trying to get it.
         :return: task's return value.
         """
+        task_name = self._task_name(task_id)
         if self.keep_results:
-            result_value = await self.redis.get(  # type: ignore[attr-defined]
-                name=task_id,
+            result_value = await self.redis.get(
+                name=task_name,
             )
         else:
-            result_value = await self.redis.getdel(  # type: ignore[attr-defined]
-                name=task_id,
+            result_value = await self.redis.getdel(
+                name=task_name,
             )
 
         if result_value is None:
@@ -368,7 +376,7 @@ class RedisAsyncClusterResultBackend(AsyncResultBackend[_ReturnType]):
         :param result: task's TaskProgress instance.
         """
         redis_set_params: Dict[str, Union[str, int, bytes]] = {
-            "name": task_id + PROGRESS_KEY_SUFFIX,
+            "name": self._task_name(task_id) + PROGRESS_KEY_SUFFIX,
             "value": self.serializer.dumpb(model_dump(progress)),
         }
         if self.result_ex_time:
@@ -388,8 +396,8 @@ class RedisAsyncClusterResultBackend(AsyncResultBackend[_ReturnType]):
         :param task_id: task's id.
         :return: task's TaskProgress instance.
         """
-        result_value = await self.redis.get(  # type: ignore[attr-defined]
-            name=task_id + PROGRESS_KEY_SUFFIX,
+        result_value = await self.redis.get(
+            name=self._task_name(task_id) + PROGRESS_KEY_SUFFIX,
         )
 
         if result_value is None:
@@ -414,6 +422,7 @@ class RedisAsyncSentinelResultBackend(AsyncResultBackend[_ReturnType]):
         min_other_sentinels: int = 0,
         sentinel_kwargs: Optional[Any] = None,
         serializer: Optional[TaskiqSerializer] = None,
+        prefix_str: Optional[str] = None,
         **connection_kwargs: Any,
     ) -> None:
         """
@@ -443,6 +452,7 @@ class RedisAsyncSentinelResultBackend(AsyncResultBackend[_ReturnType]):
         self.keep_results = keep_results
         self.result_ex_time = result_ex_time
         self.result_px_time = result_px_time
+        self.prefix_str = prefix_str
 
         unavailable_conditions = any(
             (
@@ -451,14 +461,15 @@ class RedisAsyncSentinelResultBackend(AsyncResultBackend[_ReturnType]):
             ),
         )
         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:
+            return task_id
+        return f"{self.prefix_str}:{task_id}"
 
     @asynccontextmanager
     async def _acquire_master_conn(self) -> AsyncIterator[_Redis]:
@@ -480,7 +491,7 @@ class RedisAsyncSentinelResultBackend(AsyncResultBackend[_ReturnType]):
         :param result: TaskiqResult instance.
         """
         redis_set_params: Dict[str, Union[str, bytes, int]] = {
-            "name": task_id,
+            "name": self._task_name(task_id),
             "value": self.serializer.dumpb(model_dump(result)),
         }
         if self.result_ex_time:
@@ -500,7 +511,7 @@ class RedisAsyncSentinelResultBackend(AsyncResultBackend[_ReturnType]):
         :returns: True if the result is ready else False.
         """
         async with self._acquire_master_conn() as redis:
-            return bool(await redis.exists(task_id))
+            return bool(await redis.exists(self._task_name(task_id)))
 
     async def get_result(
         self,
@@ -515,14 +526,15 @@ class RedisAsyncSentinelResultBackend(AsyncResultBackend[_ReturnType]):
         :raises ResultIsMissingError: if there is no result when trying to get it.
         :return: task's return value.
         """
+        task_name = self._task_name(task_id)
         async with self._acquire_master_conn() as redis:
             if self.keep_results:
                 result_value = await redis.get(
-                    name=task_id,
+                    name=task_name,
                 )
             else:
                 result_value = await redis.getdel(
-                    name=task_id,
+                    name=task_name,
                 )
 
         if result_value is None:
@@ -553,7 +565,7 @@ class RedisAsyncSentinelResultBackend(AsyncResultBackend[_ReturnType]):
         :param result: task's TaskProgress instance.
         """
         redis_set_params: Dict[str, Union[str, int, bytes]] = {
-            "name": task_id + PROGRESS_KEY_SUFFIX,
+            "name": self._task_name(task_id) + PROGRESS_KEY_SUFFIX,
             "value": self.serializer.dumpb(model_dump(progress)),
         }
         if self.result_ex_time:
@@ -576,7 +588,7 @@ class RedisAsyncSentinelResultBackend(AsyncResultBackend[_ReturnType]):
         """
         async with self._acquire_master_conn() as redis:
             result_value = await redis.get(
-                name=task_id + PROGRESS_KEY_SUFFIX,
+                name=self._task_name(task_id) + PROGRESS_KEY_SUFFIX,
             )
 
         if result_value is None:
@@ -590,4 +602,4 @@ class RedisAsyncSentinelResultBackend(AsyncResultBackend[_ReturnType]):
     async def shutdown(self) -> None:
         """Shutdown sentinel connections."""
         for sentinel in self.sentinel.sentinels:
-            await sentinel.aclose()  # type: ignore[attr-defined]
+            await sentinel.aclose()

taskiq_redis/redis_broker.py

@@ -1,8 +1,19 @@
 import sys
+import uuid
 from logging import getLogger
-from typing import TYPE_CHECKING, Any, AsyncGenerator, Callable, Optional, TypeVar
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncGenerator,
+    Awaitable,
+    Callable,
+    Dict,
+    Optional,
+    TypeVar,
+)
 
-from redis.asyncio import BlockingConnectionPool, Connection, Redis
+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
@@ -17,7 +28,7 @@ else:
     from typing_extensions import TypeAlias
 
 if TYPE_CHECKING:
-    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool[Connection]
+    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool[Connection]  # type: ignore
 else:
     _BlockingConnectionPool: TypeAlias = BlockingConnectionPool
 
@@ -111,7 +122,7 @@ class ListQueueBroker(BaseRedisBroker):
         """
         queue_name = message.labels.get("queue_name") or self.queue_name
         async with Redis(connection_pool=self.connection_pool) as redis_conn:
-            await redis_conn.lpush(queue_name, message.message)
+            await redis_conn.lpush(queue_name, message.message)  # type: ignore
 
     async def listen(self) -> AsyncGenerator[bytes, None]:
         """
@@ -126,9 +137,135 @@ class ListQueueBroker(BaseRedisBroker):
         while True:
             try:
                 async with Redis(connection_pool=self.connection_pool) as redis_conn:
-                    yield (await redis_conn.brpop(self.queue_name))[
+                    yield (await redis_conn.brpop(self.queue_name))[  # type: ignore
                         redis_brpop_data_position
                     ]
             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:
+        """
+        Constructs a new broker that uses streams.
+
+        :param url: url to redis.
+        :param queue_name: name for a key with stream in redis.
+        :param max_connection_pool_size: maximum number of connections in pool.
+            Each worker opens its own connection. Therefore this value has to be
+            at least number of workers + 1.
+        :param consumer_group_name: name for a consumer group.
+            Redis will keep track of acked messages for this group.
+        :param consumer_name: name for a consumer. By default it is a random uuid.
+        :param consumer_id: id for a consumer. ID of a message to start reading from.
+            $ means start from the latest message.
+        :param mkstream: create stream if it does not exist.
+        :param xread_block: block time in ms for xreadgroup.
+            Better to set it to a bigger value, to avoid unnecessary calls.
+        :param additional_streams: additional streams to read from.
+            Each key is a stream name, value is a consumer id.
+        """
+        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,  # type: ignore
+                    },
+                    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),
+                        )

taskiq_redis/redis_cluster_broker.py

@@ -1,9 +1,14 @@
-from typing import Any, AsyncGenerator
+import uuid
+from logging import getLogger
+from typing import Any, AsyncGenerator, Awaitable, Callable, Dict, Optional
 
-from redis.asyncio import RedisCluster
+from redis.asyncio import RedisCluster, ResponseError
+from taskiq import AckableMessage
 from taskiq.abc.broker import AsyncBroker
 from taskiq.message import BrokerMessage
 
+logger = getLogger(__name__)
+
 
 class BaseRedisClusterBroker(AsyncBroker):
     """Base broker that works with Redis Cluster."""
@@ -25,7 +30,7 @@ class BaseRedisClusterBroker(AsyncBroker):
         """
         super().__init__()
 
-        self.redis: RedisCluster[bytes] = RedisCluster.from_url(
+        self.redis: "RedisCluster[bytes]" = RedisCluster.from_url(  # type: ignore
             url=url,
             max_connections=max_connection_pool_size,
             **connection_kwargs,
@@ -35,7 +40,7 @@ class BaseRedisClusterBroker(AsyncBroker):
 
     async def shutdown(self) -> None:
         """Closes redis connection pool."""
-        await self.redis.aclose()  # type: ignore[attr-defined]
+        await self.redis.aclose()
         await super().shutdown()
 
 
@@ -50,7 +55,7 @@ class ListQueueClusterBroker(BaseRedisClusterBroker):
 
         :param message: message to append.
         """
-        await self.redis.lpush(self.queue_name, message.message)  # type: ignore[attr-defined]
+        await self.redis.lpush(self.queue_name, message.message)  # type: ignore
 
     async def listen(self) -> AsyncGenerator[bytes, None]:
         """
@@ -63,5 +68,121 @@ class ListQueueClusterBroker(BaseRedisClusterBroker):
         """
         redis_brpop_data_position = 1
         while True:
-            value = await self.redis.brpop([self.queue_name])  # type: ignore[attr-defined]
+            value = await self.redis.brpop([self.queue_name])  # type: ignore
             yield value[redis_brpop_data_position]
+
+
+class RedisStreamClusterBroker(BaseRedisClusterBroker):
+    """
+    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: int = 2**31,
+        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:
+        """
+        Constructs a new broker that uses streams.
+
+        :param url: url to redis.
+        :param queue_name: name for a key with stream in redis.
+        :param max_connection_pool_size: maximum number of connections in pool.
+            Each worker opens its own connection. Therefore this value has to be
+            at least number of workers + 1.
+        :param consumer_group_name: name for a consumer group.
+            Redis will keep track of acked messages for this group.
+        :param consumer_name: name for a consumer. By default it is a random uuid.
+        :param consumer_id: id for a consumer. ID of a message to start reading from.
+            $ means start from the latest message.
+        :param mkstream: create stream if it does not exist.
+        :param xread_block: block time in ms for xreadgroup.
+            Better to set it to a bigger value, to avoid unnecessary calls.
+        :param additional_streams: additional streams to read from.
+            Each key is a stream name, value is a consumer id.
+        """
+        super().__init__(
+            url,
+            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:
+        streams = {self.queue_name, *self.additional_streams.keys()}
+        async with self.redis 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.
+        """
+        await self.redis.xadd(self.queue_name, {b"data": message.message})
+
+    def _ack_generator(self, id: str) -> Callable[[], Awaitable[None]]:
+        async def _ack() -> None:
+            await self.redis.xack(
+                self.queue_name,
+                self.consumer_group_name,
+                id,
+            )
+
+        return _ack
+
+    async def listen(self) -> AsyncGenerator[AckableMessage, None]:
+        """Listen to the stream for new messages."""
+        while True:
+            fetched = await self.redis.xreadgroup(
+                self.consumer_group_name,
+                self.consumer_name,
+                {
+                    self.queue_name: ">",
+                    **self.additional_streams,  # type: ignore
+                },
+                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),
+                    )

taskiq_redis/redis_sentinel_broker.py

@@ -1,4 +1,5 @@
 import sys
+import uuid
 from contextlib import asynccontextmanager
 from logging import getLogger
 from typing import (
@@ -6,15 +7,18 @@ from typing import (
     Any,
     AsyncGenerator,
     AsyncIterator,
+    Awaitable,
     Callable,
+    Dict,
     List,
     Optional,
     Tuple,
     TypeVar,
 )
 
+from redis import ResponseError
 from redis.asyncio import Redis, Sentinel
-from taskiq import AsyncResultBackend, BrokerMessage
+from taskiq import AckableMessage, AsyncResultBackend, BrokerMessage
 from taskiq.abc.broker import AsyncBroker
 
 if sys.version_info >= (3, 10):
@@ -23,7 +27,7 @@ else:
     from typing_extensions import TypeAlias
 
 if TYPE_CHECKING:
-    _Redis: TypeAlias = Redis[bytes]
+    _Redis: TypeAlias = Redis[bytes]  # type: ignore
 else:
     _Redis: TypeAlias = Redis
 
@@ -113,7 +117,7 @@ class ListQueueSentinelBroker(BaseSentinelBroker):
         """
         queue_name = message.labels.get("queue_name") or self.queue_name
         async with self._acquire_master_conn() as redis_conn:
-            await redis_conn.lpush(queue_name, message.message)
+            await redis_conn.lpush(queue_name, message.message)  # type: ignore
 
     async def listen(self) -> AsyncGenerator[bytes, None]:
         """
@@ -127,6 +131,129 @@ class ListQueueSentinelBroker(BaseSentinelBroker):
         redis_brpop_data_position = 1
         async with self._acquire_master_conn() as redis_conn:
             while True:
-                yield (await redis_conn.brpop(self.queue_name))[
+                yield (await redis_conn.brpop(self.queue_name))[  # type: ignore
                     redis_brpop_data_position
                 ]
+
+
+class RedisStreamSentinelBroker(BaseSentinelBroker):
+    """
+    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,
+        sentinels: List[Tuple[str, int]],
+        master_name: str,
+        min_other_sentinels: int = 0,
+        queue_name: str = "taskiq",
+        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:
+        """
+        Constructs a new broker that uses streams.
+
+        :param sentinels: list of nodes to connect to.
+        :param queue_name: name for a key with stream in redis.
+        :param max_connection_pool_size: maximum number of connections in pool.
+            Each worker opens its own connection. Therefore this value has to be
+            at least number of workers + 1.
+        :param consumer_group_name: name for a consumer group.
+            Redis will keep track of acked messages for this group.
+        :param consumer_name: name for a consumer. By default it is a random uuid.
+        :param consumer_id: id for a consumer. ID of a message to start reading from.
+            $ means start from the latest message.
+        :param mkstream: create stream if it does not exist.
+        :param xread_block: block time in ms for xreadgroup.
+            Better to set it to a bigger value, to avoid unnecessary calls.
+        :param additional_streams: additional streams to read from.
+            Each key is a stream name, value is a consumer id.
+        """
+        super().__init__(
+            sentinels=sentinels,
+            master_name=master_name,
+            min_other_sentinels=min_other_sentinels,
+            task_id_generator=None,
+            result_backend=None,
+            queue_name=queue_name,
+            **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:
+        streams = {self.queue_name, *self.additional_streams.keys()}
+        async with self._acquire_master_conn() 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 self._acquire_master_conn() 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 self._acquire_master_conn() 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 the stream for new messages."""
+        async with self._acquire_master_conn() as redis_conn:
+            while True:
+                fetched = await redis_conn.xreadgroup(
+                    self.consumer_group_name,
+                    self.consumer_name,
+                    {
+                        self.queue_name: ">",
+                        **self.additional_streams,  # type: ignore
+                    },
+                    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),
+                        )

taskiq_redis/schedule_source.py

@@ -21,8 +21,8 @@ else:
     from typing_extensions import TypeAlias
 
 if TYPE_CHECKING:
-    _Redis: TypeAlias = Redis[bytes]
-    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool[Connection]
+    _Redis: TypeAlias = Redis[bytes]  # type: ignore
+    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool[Connection]  # type: ignore
 else:
     _Redis: TypeAlias = Redis
     _BlockingConnectionPool: TypeAlias = BlockingConnectionPool
@@ -140,7 +140,7 @@ class RedisClusterScheduleSource(ScheduleSource):
         **connection_kwargs: Any,
     ) -> None:
         self.prefix = prefix
-        self.redis: RedisCluster[bytes] = RedisCluster.from_url(
+        self.redis: "RedisCluster" = RedisCluster.from_url(
             url,
             **connection_kwargs,
         )
@@ -150,7 +150,7 @@ class RedisClusterScheduleSource(ScheduleSource):
 
     async def delete_schedule(self, schedule_id: str) -> None:
         """Remove schedule by id."""
-        await self.redis.delete(f"{self.prefix}:{schedule_id}")  # type: ignore[attr-defined]
+        await self.redis.delete(f"{self.prefix}:{schedule_id}")
 
     async def add_schedule(self, schedule: ScheduledTask) -> None:
         """
@@ -159,7 +159,7 @@ class RedisClusterScheduleSource(ScheduleSource):
         :param schedule: schedule to add.
         :param schedule_id: schedule id.
         """
-        await self.redis.set(  # type: ignore[attr-defined]
+        await self.redis.set(
             f"{self.prefix}:{schedule.schedule_id}",
             self.serializer.dumpb(model_dump(schedule)),
         )
@@ -173,8 +173,8 @@ class RedisClusterScheduleSource(ScheduleSource):
         :return: list of schedules.
         """
         schedules = []
-        async for key in self.redis.scan_iter(f"{self.prefix}:*"):  # type: ignore[attr-defined]
-            raw_schedule = await self.redis.get(key)  # type: ignore[attr-defined]
+        async for key in self.redis.scan_iter(f"{self.prefix}:*"):
+            raw_schedule = await self.redis.get(key)
             parsed_schedule = model_validate(
                 ScheduledTask,
                 self.serializer.loadb(raw_schedule),
@@ -189,7 +189,7 @@ class RedisClusterScheduleSource(ScheduleSource):
 
     async def shutdown(self) -> None:
         """Shut down the schedule source."""
-        await self.redis.aclose()  # type: ignore[attr-defined]
+        await self.redis.aclose()
 
 
 class RedisSentinelScheduleSource(ScheduleSource):
@@ -288,4 +288,4 @@ class RedisSentinelScheduleSource(ScheduleSource):
     async def shutdown(self) -> None:
         """Shut down the schedule source."""
         for sentinel in self.sentinel.sentinels:
-            await sentinel.aclose()  # type: ignore[attr-defined]
+            await sentinel.aclose()

{taskiq_redis-1.0.2.dist-info → taskiq_redis-1.0.3.dist-info}/METADATA

@@ -1,22 +1,23 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: taskiq-redis
-Version: 1.0.2
+Version: 1.0.3
 Summary: Redis integration for taskiq
-Home-page: https://github.com/taskiq-python/taskiq-redis
 Keywords: taskiq,tasks,distributed,async,redis,result_backend
 Author: taskiq-team
 Author-email: taskiq@norely.com
-Requires-Python: >=3.8.1,<4.0.0
+Requires-Python: >=3.9,<4.0
 Classifier: Programming Language :: Python
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3 :: Only
 Classifier: Programming Language :: Python :: 3.8
 Requires-Dist: redis (>=5,<6)
-Requires-Dist: taskiq (>=0.11.1,<1)
+Requires-Dist: taskiq (>=0.11.12,<1)
+Project-URL: Homepage, https://github.com/taskiq-python/taskiq-redis
 Project-URL: Repository, https://github.com/taskiq-python/taskiq-redis
 Description-Content-Type: text/markdown
 
@@ -43,17 +44,17 @@ Let's see the example with the redis broker and redis async result:
 # broker.py
 import asyncio
 
-from taskiq_redis import ListQueueBroker, RedisAsyncResultBackend
+from taskiq_redis import RedisAsyncResultBackend, RedisStreamBroker
 
-redis_async_result = RedisAsyncResultBackend(
+result_backend = RedisAsyncResultBackend(
     redis_url="redis://localhost:6379",
 )
 
 # Or you can use PubSubBroker if you need broadcasting
-broker = ListQueueBroker(
+# Or ListQueueBroker if you don't want acknowledges
+broker = RedisStreamBroker(
     url="redis://localhost:6379",
-    result_backend=redis_async_result,
-)
+).with_result_backend(result_backend)
 
 
 @broker.task
@@ -77,25 +78,48 @@ Launch the workers:
 Then run the main code:
 `python3 broker.py`
 
-## 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.
+## Brokers
 
-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.
+This package contains 6 broker implementations.
+3 broker types:
+* PubSub broker
+* ListQueue broker
+* Stream broker
 
-Brokers parameters:
-* `url` - url to redis.
-* `task_id_generator` - custom task_id genertaor.
-* `result_backend` - custom result backend.
-* `queue_name` - name of the pub/sub channel in redis.
-* `max_connection_pool_size` - maximum number of connections in pool.
-* Any other keyword arguments are passed to `redis.asyncio.BlockingConnectionPool`.
-  Notably, you can use `timeout` to set custom timeout in seconds for reconnects
-  (or set it to `None` to try reconnects indefinitely).
+Each of type is implemented for each redis architecture:
+* Single node
+* Cluster
+* Sentinel
+
+Here's a small breakdown of how they differ from eachother.
+
+
+### PubSub
+
+By default on old redis versions PUBSUB was the way of making redis into a queue.
+But using PUBSUB means that all messages delivered to all subscribed consumers.
+
+> [!WARNING]
+> This broker doesn't support acknowledgements. If during message processing
+> Worker was suddenly killed the message is going to be lost.
+
+### ListQueue
+
+This broker creates a list of messages at some key. Adding new tasks will be done
+by appending them from the left side using `lpush`, and taking them from the right side using `brpop`.
+
+> [!WARNING]
+> This broker doesn't support acknowledgements. If during message processing
+> Worker was suddenly killed the message is going to be lost.
+
+### Stream
+
+Stream brokers use redis [stream type](https://redis.io/docs/latest/develop/data-types/streams/) to store and fetch messages.
+
+> [!TIP]
+> This broker **supports** acknowledgements and therefore is fine to use in cases when data durability is
+> required.
 
 ## RedisAsyncResultBackend configuration
 
@@ -107,19 +131,21 @@ RedisAsyncResultBackend parameters:
 * Any other keyword arguments are passed to `redis.asyncio.BlockingConnectionPool`.
   Notably, you can use `timeout` to set custom timeout in seconds for reconnects
   (or set it to `None` to try reconnects indefinitely).
-> IMPORTANT: **It is highly recommended to use expire time ​​in RedisAsyncResultBackend**
+
+> [!WARNING]
+> **It is highly recommended to use expire time in RedisAsyncResultBackend**
 > If you want to add expiration, either `result_ex_time` or `result_px_time` must be set.
->```python
-># First variant
->redis_async_result = RedisAsyncResultBackend(
->    redis_url="redis://localhost:6379",
->    result_ex_time=1000,
->)
+> ```python
+> # First variant
+> redis_async_result = RedisAsyncResultBackend(
+>     redis_url="redis://localhost:6379",
+>     result_ex_time=1000,
+> )
 >
-># Second variant
->redis_async_result = RedisAsyncResultBackend(
->    redis_url="redis://localhost:6379",
->    result_px_time=1000000,
->)
->```
+> # Second variant
+> redis_async_result = RedisAsyncResultBackend(
+>     redis_url="redis://localhost:6379",
+>     result_px_time=1000000,
+> )
+> ```
 

taskiq_redis-1.0.3.dist-info/RECORD

@@ -0,0 +1,12 @@
+taskiq_redis/__init__.py,sha256=nb2Lx4lVj9m20duzRFQk3nNVswQGylmwLsQN6Qc1lGI,1091
+taskiq_redis/exceptions.py,sha256=7buBJ7CRVWd5WqVqSjtHO8cVL7QzZg-DOM3nB87t-Sk,738
+taskiq_redis/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+taskiq_redis/redis_backend.py,sha256=MLBaUN3Zx-DLvm1r-lgPU82_WZq9kc6oTxYI8LQjd6k,19882
+taskiq_redis/redis_broker.py,sha256=ZLn7LAHj8Sh_oyW5hMgD7PZPQfUdXNPKdqhBcr9Okmg,9775
+taskiq_redis/redis_cluster_broker.py,sha256=FuWl5fP7Fwr9FbytErmhcUGjRCdPexDK2Co2u6kpDlo,6591
+taskiq_redis/redis_sentinel_broker.py,sha256=wHnbG3xuD_ruhhwp4AXo91NNjq8v2iufUZ0i_HbBRVQ,9073
+taskiq_redis/schedule_source.py,sha256=hqpcs2D8W90KUDHREKblisnhGCE9dbVOtKtuJcOTGZw,9915
+taskiq_redis-1.0.3.dist-info/LICENSE,sha256=lEHEEE-ZxmuItxYgUMPiFWdRcAITxE8DFMNyAg4eOYE,1075
+taskiq_redis-1.0.3.dist-info/METADATA,sha256=whl7_U6GIIcNzVZqrLQkktES6Y91awk8-HtgaX5IQ8s,4391
+taskiq_redis-1.0.3.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
+taskiq_redis-1.0.3.dist-info/RECORD,,

{taskiq_redis-1.0.2.dist-info → taskiq_redis-1.0.3.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry-core 1.9.0
+Generator: poetry-core 2.1.1
 Root-Is-Purelib: true
 Tag: py3-none-any

taskiq_redis-1.0.2.dist-info/RECORD

@@ -1,12 +0,0 @@
-taskiq_redis/__init__.py,sha256=UEW3rQXt4jinMnAKJlpXQhyPDh6SU2in0bPgzfIo3y4,911
-taskiq_redis/exceptions.py,sha256=eS4bfZVAjyMsnFs3IF74uYwO1KZOlrYxhxgPqD49ztU,561
-taskiq_redis/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-taskiq_redis/redis_backend.py,sha256=mGU3rXJ727X-qYqSfaderTFVGz93NpCyHCf5PDFTjGk,19543
-taskiq_redis/redis_broker.py,sha256=JeoA3-quZYqa_wixJefMRYPkZe94x-qoQb6tQKkHLzg,4733
-taskiq_redis/redis_cluster_broker.py,sha256=CgPKkoEHZ1moNM-VNmzPQdjjNOrhiVUCNV-7FrUgqTo,2121
-taskiq_redis/redis_sentinel_broker.py,sha256=5MxUFIX7qRyDT7IHebLUhxAmmUwk1_b2sxjpSXRcjlo,4114
-taskiq_redis/schedule_source.py,sha256=bk96UBg8op-Xqg_PVETgyDb92cDaY69EAjpP8GvYSnY,10068
-taskiq_redis-1.0.2.dist-info/LICENSE,sha256=lEHEEE-ZxmuItxYgUMPiFWdRcAITxE8DFMNyAg4eOYE,1075
-taskiq_redis-1.0.2.dist-info/METADATA,sha256=6A_nDPLAmO92y_Db7vNUIGiOGqH7gTm_rCpb0KMIPOc,4030
-taskiq_redis-1.0.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-taskiq_redis-1.0.2.dist-info/RECORD,,