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

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()