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

taskiq_redis/__init__.py

@@ -1,5 +1,6 @@
 """Package for redis integration."""
 
+from taskiq_redis.list_schedule_source import ListRedisScheduleSource
 from taskiq_redis.redis_backend import (
     RedisAsyncClusterResultBackend,
     RedisAsyncResultBackend,
@@ -25,6 +26,7 @@ __all__ = [
     "ListQueueBroker",
     "ListQueueClusterBroker",
     "ListQueueSentinelBroker",
+    "ListRedisScheduleSource",
     "PubSubBroker",
     "PubSubSentinelBroker",
     "RedisAsyncClusterResultBackend",

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):
+    """Schedule source based on arrays."""
+
+    def __init__(
+        self,
+        url: str,
+        prefix: str = "schedule",
+        max_connection_pool_size: Optional[int] = None,
+        serializer: Optional[TaskiqSerializer] = None,
+        buffer_size: int = 50,
+        skip_past_schedules: bool = False,
+        **connection_kwargs: Any,
+    ) -> None:
+        super().__init__()
+        self._prefix = prefix
+        self._buffer_size = buffer_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 %d cron schedules", len(crons))
+            if crons:
+                buffer.extend(crons)
+            timed.extend(await redis.lrange(self._get_time_key(current_time), 0, -1))  # type: ignore
+            logger.debug("Got %d timed schedules", len(timed))
+            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_broker.py

@@ -165,6 +165,7 @@ class RedisStreamBroker(BaseRedisBroker):
         consumer_id: str = "$",
         mkstream: bool = True,
         xread_block: int = 10000,
+        maxlen: Optional[int] = None,
         additional_streams: Optional[Dict[str, str]] = None,
         **connection_kwargs: Any,
     ) -> None:
@@ -184,6 +185,8 @@ class RedisStreamBroker(BaseRedisBroker):
         :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 maxlen: sets the maximum length of the stream
+            trims (the old values of) the stream each time a new element is added
         :param additional_streams: additional streams to read from.
             Each key is a stream name, value is a consumer id.
         """
@@ -200,6 +203,7 @@ class RedisStreamBroker(BaseRedisBroker):
         self.consumer_id = consumer_id
         self.mkstream = mkstream
         self.block = xread_block
+        self.maxlen = maxlen
         self.additional_streams = additional_streams or {}
 
     async def _declare_consumer_group(self) -> None:
@@ -235,7 +239,11 @@ class RedisStreamBroker(BaseRedisBroker):
         :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})
+            await redis_conn.xadd(
+                self.queue_name,
+                {b"data": message.message},
+                maxlen=self.maxlen,
+            )
 
     def _ack_generator(self, id: str) -> Callable[[], Awaitable[None]]:
         async def _ack() -> None:

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: taskiq-redis
-Version: 1.0.3
+Version: 1.0.5
 Summary: Redis integration for taskiq
 Keywords: taskiq,tasks,distributed,async,redis,result_backend
 Author: taskiq-team
@@ -149,3 +149,67 @@ RedisAsyncResultBackend parameters:
 > )
 > ```
 
+
+## Schedule sources
+
+
+You can use this package to add dynamic schedule sources. They are used to store
+schedules for taskiq scheduler.
+
+The advantage of using schedule sources from this package over default `LabelBased` source is that you can
+dynamically add schedules in it.
+
+We have two types of schedules:
+
+* `RedisScheduleSource`
+* `ListRedisScheduleSource`
+
+
+### RedisScheduleSource
+
+This source is super simple. It stores all schedules by key `{prefix}:{schedule_id}`. When scheduler requests
+schedules, it retrieves all values from redis that start with a given `prefix`.
+
+This is very ineficent and should not be used for high-volume schedules. Because if you have `1000` schedules, this scheduler will make at least `20` requests to retrieve them (we use `scan` and `mget` to minimize number of calls).
+
+### ListRedisScheduleSource
+
+This source holds values in lists.
+
+* For cron tasks it uses key `{prefix}:cron`.
+* For timed schedules it uses key `{prefix}:time:{time}` where `{time}` is actually time where schedules should run.
+
+The main advantage of this approach is that we only fetch tasks we need to run at a given time and do not perform any excesive calls to redis.
+
+
+### Migration from one source to another
+
+To migrate from `RedisScheduleSource` to `ListRedisScheduleSource` you can define the latter as this:
+
+```python
+# broker.py
+import asyncio
+import datetime
+
+from taskiq import TaskiqScheduler
+
+from taskiq_redis import ListRedisScheduleSource, RedisStreamBroker
+from taskiq_redis.schedule_source import RedisScheduleSource
+
+broker = RedisStreamBroker(url="redis://localhost:6379")
+
+old_source = RedisScheduleSource("redis://localhost/1", prefix="prefix1")
+array_source = ListRedisScheduleSource(
+    "redis://localhost/1",
+    prefix="prefix2",
+    # To migrate schedules from an old source.
+).with_migrate_from(
+    old_source,
+    # To delete schedules from an old source.
+    delete_schedules=True,
+)
+scheduler = TaskiqScheduler(broker, [array_source])
+```
+
+During startup the scheduler will try to migrate schedules from an old source to a new one. Please be sure to specify different prefixe just to avoid any kind of collision between these two.
+

{taskiq_redis-1.0.3.dist-info → taskiq_redis-1.0.5.dist-info}/RECORD

@@ -1,12 +1,13 @@
-taskiq_redis/__init__.py,sha256=nb2Lx4lVj9m20duzRFQk3nNVswQGylmwLsQN6Qc1lGI,1091
+taskiq_redis/__init__.py,sha256=Sl4m9rKxweU1t0m289Qtf0qm4xSSkkFHoOfKq6qaz6g,1192
 taskiq_redis/exceptions.py,sha256=7buBJ7CRVWd5WqVqSjtHO8cVL7QzZg-DOM3nB87t-Sk,738
+taskiq_redis/list_schedule_source.py,sha256=w7lSJ1-n889PXOwXiA_Jp6Wj4RsZpuXwFCVhOIc5KAw,9462
 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_broker.py,sha256=ep31kUxXx4XhGKjrzVjdNBmA6wLbTQoy7-DlKuwtLz4,10068
 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.5.dist-info/LICENSE,sha256=lEHEEE-ZxmuItxYgUMPiFWdRcAITxE8DFMNyAg4eOYE,1075
+taskiq_redis-1.0.5.dist-info/METADATA,sha256=uV2AJOyyE-f1_5A1LsvGpbf74zSueiJaNNRJ6kIqpvw,6573
+taskiq_redis-1.0.5.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+taskiq_redis-1.0.5.dist-info/RECORD,,

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

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