Merge pull request #17 from corruptmane/develop

Author: s3rius

taskiq_nats/__init__.py

@@ -11,10 +11,12 @@
     PushBasedJetStreamBroker,
 )
 from taskiq_nats.result_backend import NATSObjectStoreResultBackend
+from taskiq_nats.schedule_source import NATSKeyValueScheduleSource
 
 __all__ = [
     "NatsBroker",
     "PushBasedJetStreamBroker",
     "PullBasedJetStreamBroker",
     "NATSObjectStoreResultBackend",
+    "NATSKeyValueScheduleSource",
 ]

taskiq_nats/schedule_source.py

@@ -0,0 +1,113 @@
+import logging
+from typing import Any, Final, List, Optional, Union
+
+import nats
+from nats import NATS
+from nats.js import JetStreamContext
+from nats.js.errors import BucketNotFoundError, NoKeysError
+from nats.js.kv import KeyValue
+from taskiq import ScheduledTask, ScheduleSource
+from taskiq.abc.serializer import TaskiqSerializer
+from taskiq.compat import model_dump, model_validate
+from taskiq.serializers import PickleSerializer
+
+log = logging.getLogger(__name__)
+
+
+class NATSKeyValueScheduleSource(ScheduleSource):
+    """
+    Source of schedules for NATS Key-Value storage.
+
+    This class allows you to store schedules in NATS Key-Value storage.
+    Also it supports dynamic schedules.
+    """
+
+    def __init__(
+        self,
+        servers: Union[str, List[str]],
+        bucket_name: str = "taskiq_schedules",
+        prefix: str = "schedule",
+        serializer: Optional[TaskiqSerializer] = None,
+        **connect_options: Any,
+    ) -> None:
+        """Construct new result backend.
+
+        :param servers: NATS servers.
+        :param bucket_name: name of the bucket where schedules would be stored.
+        :param prefix: prefix for nats kv storage schedule keys.
+        :param serializer: serializer for data.
+        :param connect_kwargs: additional arguments for nats `connect()` method.
+        """
+        self.servers: Final = servers
+        self.bucket_name: Final = bucket_name
+        self.prefix: Final = prefix
+        self.serializer = serializer or PickleSerializer()
+        self.connect_options: Final = connect_options
+
+        self.nats_client: NATS
+        self.nats_jetstream: JetStreamContext
+        self.kv: KeyValue
+
+    async def startup(self) -> None:
+        """Create new connection to NATS.
+
+        Initialize JetStream context and new KeyValue instance.
+        """
+        self.nats_client = await nats.connect(
+            servers=self.servers,
+            **self.connect_options,
+        )
+        self.nats_jetstream = self.nats_client.jetstream()
+
+        try:
+            self.kv = await self.nats_jetstream.key_value(self.bucket_name)
+        except BucketNotFoundError:
+            self.kv = await self.nats_jetstream.create_key_value(
+                bucket=self.bucket_name,
+            )
+
+    async def shutdown(self) -> None:
+        """Close nats connection."""
+        if self.nats_client.is_closed:
+            return
+        await self.nats_client.close()
+
+    async def delete_schedule(self, schedule_id: str) -> None:
+        """Remove schedule by id."""
+        await self.kv.delete(f"{self.prefix}.{schedule_id}")
+
+    async def add_schedule(self, schedule: ScheduledTask) -> None:
+        """
+        Add schedule to NATS Key-Value storage.
+
+        :param schedule: schedule to add.
+        :param schedule_id: schedule id.
+        """
+        await self.kv.put(
+            f"{self.prefix}.{schedule.schedule_id}",
+            self.serializer.dumpb(model_dump(schedule)),
+        )
+
+    async def get_schedules(self) -> List[ScheduledTask]:
+        """
+        Get all schedules from NATS Key-Value storage.
+
+        This method is used by scheduler to get all schedules.
+
+        :return: list of schedules.
+        """
+        try:
+            schedules = await self.kv.history(f"{self.prefix}.*")
+        except NoKeysError:
+            return []
+
+        return [
+            model_validate(ScheduledTask, self.serializer.loadb(schedule.value))
+            for schedule in schedules
+            if schedule and schedule.value
+        ]
+
+    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)

tests/test_schedule_source.py

@@ -0,0 +1,86 @@
+import datetime as dt
+import uuid
+from typing import List
+
+import pytest
+from taskiq import ScheduledTask
+
+from taskiq_nats import NATSKeyValueScheduleSource
+
+
+@pytest.mark.anyio
+async def test_set_schedule(nats_urls: List[str]) -> None:
+    prefix = uuid.uuid4().hex
+    source = NATSKeyValueScheduleSource(servers=nats_urls, prefix=prefix)
+    await source.startup()
+    schedule = ScheduledTask(
+        task_name="test_task",
+        labels={},
+        args=[],
+        kwargs={},
+        cron="* * * * *",
+    )
+    await source.add_schedule(schedule)
+    schedules = await source.get_schedules()
+    assert schedules == [schedule]
+    await source.shutdown()
+
+
+@pytest.mark.anyio
+async def test_delete_schedule(nats_urls: List[str]) -> None:
+    prefix = uuid.uuid4().hex
+    source = NATSKeyValueScheduleSource(servers=nats_urls, prefix=prefix)
+    await source.startup()
+    schedule = ScheduledTask(
+        task_name="test_task",
+        labels={},
+        args=[],
+        kwargs={},
+        cron="* * * * *",
+    )
+    await source.add_schedule(schedule)
+    schedules = await source.get_schedules()
+    assert schedules == [schedule]
+    await source.delete_schedule(schedule.schedule_id)
+    schedules = await source.get_schedules()
+    # Schedules are empty.
+    assert not schedules
+    await source.shutdown()
+
+
+@pytest.mark.anyio
+async def test_post_run_cron(nats_urls: List[str]) -> None:
+    prefix = uuid.uuid4().hex
+    source = NATSKeyValueScheduleSource(servers=nats_urls, prefix=prefix)
+    await source.startup()
+    schedule = ScheduledTask(
+        task_name="test_task",
+        labels={},
+        args=[],
+        kwargs={},
+        cron="* * * * *",
+    )
+    await source.add_schedule(schedule)
+    assert await source.get_schedules() == [schedule]
+    await source.post_send(schedule)
+    assert await source.get_schedules() == [schedule]
+    await source.shutdown()
+
+
+@pytest.mark.anyio
+async def test_post_run_time(nats_urls: List[str]) -> None:
+    prefix = uuid.uuid4().hex
+    source = NATSKeyValueScheduleSource(servers=nats_urls, prefix=prefix)
+    await source.startup()
+    schedule = ScheduledTask(
+        task_name="test_task",
+        labels={},
+        args=[],
+        kwargs={},
+        time=dt.datetime(2000, 1, 1),
+    )
+    await source.add_schedule(schedule)
+    assert await source.get_schedules() == [schedule]
+    await source.post_send(schedule)
+    assert await source.get_schedules() == []
+    await source.shutdown()

tests/utils.py

@@ -5,10 +5,10 @@
 
 async def read_message(broker: AsyncBroker) -> Union[bytes, AckableMessage]:
     """
-    Read signle message from the broker's listen method.
+    Read single message from the broker's listen method.
 
     :param broker: current broker.
-    :return: firs message.
+    :return: first message.
     """
     msg: Union[bytes, AckableMessage] = b"error"
     async for message in broker.listen():