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

taskiq_redis/__init__.py

@@ -1,14 +1,20 @@
 """Package for redis integration."""
+
+from taskiq_redis.list_schedule_source import ListRedisScheduleSource
 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 +23,19 @@ from taskiq_redis.schedule_source import (
 )
 
 __all__ = [
-    "RedisAsyncClusterResultBackend",
-    "RedisAsyncResultBackend",
-    "RedisAsyncSentinelResultBackend",
     "ListQueueBroker",
-    "PubSubBroker",
     "ListQueueClusterBroker",
     "ListQueueSentinelBroker",
+    "ListRedisScheduleSource",
+    "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/list_schedule_source.py

@@ -0,0 +1,229 @@
+import datetime
+from logging import getLogger
+from typing import Any, List, Optional
+
+from redis.asyncio import BlockingConnectionPool, Redis
+from taskiq import ScheduledTask, ScheduleSource
+from taskiq.abc.serializer import TaskiqSerializer
+from taskiq.compat import model_dump, model_validate
+from taskiq.serializers import PickleSerializer
+from typing_extensions import Self
+
+logger = getLogger("taskiq.redis_schedule_source")
+
+
+class ListRedisScheduleSource(ScheduleSource):
+    """Schecule source based on arrays."""
+
+    def __init__(
+        self,
+        url: str,
+        prefix: str = "schedule",
+        max_connection_pool_size: Optional[int] = None,
+        serializer: Optional[TaskiqSerializer] = None,
+        bufffer_size: int = 50,
+        skip_past_schedules: bool = False,
+        **connection_kwargs: Any,
+    ) -> None:
+        super().__init__()
+        self._prefix = prefix
+        self._buffer_size = bufffer_size
+        self._connection_pool = BlockingConnectionPool.from_url(
+            url=url,
+            max_connections=max_connection_pool_size,
+            **connection_kwargs,
+        )
+        if serializer is None:
+            serializer = PickleSerializer()
+        self._serializer = serializer
+        self._is_first_run = True
+        self._previous_schedule_source: Optional[ScheduleSource] = None
+        self._delete_schedules_after_migration: bool = True
+        self._skip_past_schedules = skip_past_schedules
+
+    async def startup(self) -> None:
+        """
+        Startup the schedule source.
+
+        By default this function does nothing.
+        But if the previous schedule source is set,
+        it will try to migrate schedules from it.
+        """
+        if self._previous_schedule_source is not None:
+            logger.info("Migrating schedules from previous source")
+            await self._previous_schedule_source.startup()
+            schedules = await self._previous_schedule_source.get_schedules()
+            logger.info(f"Found {len(schedules)}")
+            for schedule in schedules:
+                await self.add_schedule(schedule)
+                if self._delete_schedules_after_migration:
+                    await self._previous_schedule_source.delete_schedule(
+                        schedule.schedule_id,
+                    )
+            await self._previous_schedule_source.shutdown()
+            logger.info("Migration complete")
+
+    def _get_time_key(self, time: datetime.datetime) -> str:
+        """Get the key for a time-based schedule."""
+        if time.tzinfo is None:
+            time = time.replace(tzinfo=datetime.timezone.utc)
+        iso_time = time.astimezone(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M")
+        return f"{self._prefix}:time:{iso_time}"
+
+    def _get_cron_key(self) -> str:
+        """Get the key for a cron-based schedule."""
+        return f"{self._prefix}:cron"
+
+    def _get_data_key(self, schedule_id: str) -> str:
+        """Get the key for a schedule data."""
+        return f"{self._prefix}:data:{schedule_id}"
+
+    def _parse_time_key(self, key: str) -> Optional[datetime.datetime]:
+        """Get time value from the timed-key."""
+        try:
+            dt_str = key.split(":", 2)[2]
+            return datetime.datetime.strptime(dt_str, "%Y-%m-%dT%H:%M").replace(
+                tzinfo=datetime.timezone.utc,
+            )
+        except ValueError:
+            logger.debug("Failed to parse time key %s", key)
+            return None
+
+    async def _get_previous_time_schedules(self) -> list[bytes]:
+        """
+        Function that gets all timed schedules that are in the past.
+
+        Since this source doesn't retrieve all the schedules at once,
+        we need to get all the schedules that are in the past and haven't
+        been sent yet.
+
+        We do this by getting all the time keys and checking if the time
+        is less than the current time.
+
+        This function is called only during the first run to minimize
+        the number of requests to the Redis server.
+        """
+        logger.info("Getting previous time schedules")
+        minute_before = datetime.datetime.now(
+            datetime.timezone.utc,
+        ).replace(second=0, microsecond=0) - datetime.timedelta(
+            minutes=1,
+        )
+        schedules = []
+        async with Redis(connection_pool=self._connection_pool) as redis:
+            time_keys: list[str] = []
+            # We need to get all the time keys and check if the time is less than
+            # the current time.
+            async for key in redis.scan_iter(f"{self._prefix}:time:*"):
+                key_time = self._parse_time_key(key.decode())
+                if key_time and key_time <= minute_before:
+                    time_keys.append(key.decode())
+            for key in time_keys:
+                schedules.extend(await redis.lrange(key, 0, -1))  # type: ignore
+
+        return schedules
+
+    async def delete_schedule(self, schedule_id: str) -> None:
+        """Delete a schedule from the source."""
+        async with Redis(connection_pool=self._connection_pool) as redis:
+            schedule = await redis.getdel(self._get_data_key(schedule_id))
+            if schedule is not None:
+                logger.debug("Deleting schedule %s", schedule_id)
+                schedule = model_validate(
+                    ScheduledTask,
+                    self._serializer.loadb(schedule),
+                )
+                # We need to remove the schedule from the cron or time list.
+                if schedule.cron is not None:
+                    await redis.lrem(self._get_cron_key(), 0, schedule_id)  # type: ignore
+                elif schedule.time is not None:
+                    time_key = self._get_time_key(schedule.time)
+                    await redis.lrem(time_key, 0, schedule_id)  # type: ignore
+
+    async def add_schedule(self, schedule: "ScheduledTask") -> None:
+        """Add a schedule to the source."""
+        async with Redis(connection_pool=self._connection_pool) as redis:
+            # At first we set data key which contains the schedule data.
+            await redis.set(
+                f"{self._prefix}:data:{schedule.schedule_id}",
+                self._serializer.dumpb(model_dump(schedule)),
+            )
+            # Then we add the schedule to the cron or time list.
+            # This is an optimization, so we can get all the schedules
+            # for the current time much faster.
+            if schedule.cron is not None:
+                await redis.rpush(self._get_cron_key(), schedule.schedule_id)  # type: ignore
+            elif schedule.time is not None:
+                await redis.rpush(  # type: ignore
+                    self._get_time_key(schedule.time),
+                    schedule.schedule_id,
+                )
+
+    async def post_send(self, task: ScheduledTask) -> None:
+        """Delete a task after it's completed."""
+        if task.time is not None:
+            await self.delete_schedule(task.schedule_id)
+
+    async def get_schedules(self) -> List["ScheduledTask"]:
+        """
+        Get all schedules.
+
+        This function gets all the schedules from the schedule source.
+        What it does is get all the cron schedules and time schedules
+        for the current time and return them.
+
+        If it's the first run, it also gets all the time schedules
+        that are in the past and haven't been sent yet.
+        """
+        schedules = []
+        current_time = datetime.datetime.now(datetime.timezone.utc)
+        timed: list[bytes] = []
+        # Only during first run, we need to get previous time schedules
+        if self._is_first_run and not self._skip_past_schedules:
+            timed = await self._get_previous_time_schedules()
+            self._is_first_run = False
+        async with Redis(connection_pool=self._connection_pool) as redis:
+            buffer = []
+            crons = await redis.lrange(self._get_cron_key(), 0, -1)  # type: ignore
+            logger.debug("Got cron scheduleds: %s", crons)
+            if crons:
+                buffer.extend(crons)
+            timed.extend(await redis.lrange(self._get_time_key(current_time), 0, -1))  # type: ignore
+            logger.debug("Got timed scheduleds: %s", crons)
+            if timed:
+                buffer.extend(timed)
+            while buffer:
+                schedules.extend(
+                    await redis.mget(
+                        (
+                            self._get_data_key(x.decode())
+                            for x in buffer[: self._buffer_size]
+                        ),
+                    ),
+                )
+                buffer = buffer[self._buffer_size :]
+
+        return [
+            model_validate(ScheduledTask, self._serializer.loadb(schedule))
+            for schedule in schedules
+            if schedule
+        ]
+
+    def with_migrate_from(
+        self,
+        source: ScheduleSource,
+        delete_schedules: bool = True,
+    ) -> Self:
+        """
+        Enable migration from previous schedule source.
+
+        If this function is called during declaration,
+        the source will try to migrate schedules from the previous source.
+
+        :param source: previous schedule source
+        :param delete_schedules: delete schedules during migration process
+            from the previous source.
+        """
+        self._previous_schedule_source = source
+        self._delete_schedules_after_migration = delete_schedules
+        return self

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