Merge branch 'develop'

Author: s3rius

.flake8

@@ -1,9 +1,9 @@
 name: Release python package
 
 on:
-  push:
-    tags:
-      - "*"
+  release:
+    types:
+      - released
 
 jobs:
   deploy:

.github/workflows/release.yaml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "taskiq-nats"
-version = "0.5.0"
+version = "0.0.0"
 description = "NATS integration for taskiq"
 authors = ["taskiq-team <taskiq@norely.com>"]
 readme = "README.md"
@@ -31,7 +31,7 @@ pytest = "^7.2.2"
 pytest-xdist = "^3.2.1"
 anyio = "^3.6.2"
 pytest-cov = "^4.0.0"
-wemake-python-styleguide = "^0.18.0"
+ruff = "^0.7.2"
 
 
 [tool.mypy]
@@ -53,7 +53,7 @@ build-backend = "poetry.core.masonry.api"
 [tool.ruff]
 # List of enabled rulsets.
 # See https://docs.astral.sh/ruff/rules/ for more information.
-select = [
+lint.select = [
     "E",   # Error
     "F",   # Pyflakes
     "W",   # Pycodestyle
@@ -80,7 +80,7 @@ select = [
     "PL",  # PyLint checks
     "RUF", # Specific to Ruff checks
 ]
-ignore = [
+lint.ignore = [
     "D105",    # Missing docstring in magic method
     "D107",    # Missing docstring in __init__
     "D212",    # Multi-line docstring summary should start at the first line
@@ -94,10 +94,10 @@ ignore = [
     "D106",    # Missing docstring in public nested class
 ]
 exclude = [".venv/"]
-mccabe = { max-complexity = 10 }
+lint.mccabe = { max-complexity = 10 }
 line-length = 88
 
-[tool.ruff.per-file-ignores]
+[tool.ruff.lint.per-file-ignores]
 "tests/*" = [
     "S101",   # Use of assert detected
     "S301",   # Use of pickle detected
@@ -107,12 +107,12 @@ line-length = 88
     "D101",   # Missing docstring in public class
 ]
 
-[tool.ruff.pydocstyle]
+[tool.ruff.lint.pydocstyle]
 convention = "pep257"
 ignore-decorators = ["typing.overload"]
 
-[tool.ruff.pylint]
+[tool.ruff.lint.pylint]
 allow-magic-value-types = ["int", "str", "float"]
 
-[tool.ruff.flake8-bugbear]
+[tool.ruff.lint.flake8-bugbear]
 extend-immutable-calls = ["taskiq_dependencies.Depends", "taskiq.TaskiqDepends"]

poetry.lock

@@ -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",
 ]

pyproject.toml

@@ -1,7 +1,7 @@
 from typing import Any, Final, List, Optional, TypeVar, Union
 
 import nats
-from nats import NATS
+from nats.aio.client import Client
 from nats.js import JetStreamContext
 from nats.js.errors import BucketNotFoundError, ObjectNotFoundError
 from nats.js.object_store import ObjectStore
@@ -37,7 +37,7 @@ def __init__(
         self.serializer = serializer or PickleSerializer()
         self.connect_options: Final = connect_options
 
-        self.nats_client: NATS
+        self.nats_client: Client
         self.nats_jetstream: JetStreamContext
         self.object_store: ObjectStore
 

taskiq_nats/__init__.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)

taskiq_nats/result_backend.py

@@ -111,43 +111,39 @@ async def test_success_backend_default_result(
     assert result == default_taskiq_result
 
 
-async def test_success_backend_custom_result(
+async def test_error_backend_custom_result(
     nats_result_backend: NATSObjectStoreResultBackend[_ReturnType],
     custom_taskiq_result: TaskiqResult[_ReturnType],
     task_id: str,
 ) -> None:
     """
-    Tests normal behavior with custom result in TaskiqResult.
+    Tests that error is thrown on non-serializable result.
 
     :param custom_taskiq_result: TaskiqResult with custom result.
     :param task_id: ID for task.
     :param redis_url: url to redis.
     """
-    await nats_result_backend.set_result(
-        task_id=task_id,
-        result=custom_taskiq_result,
-    )
-    result = await nats_result_backend.get_result(task_id=task_id)
-
-    assert (
-        result.return_value.test_arg  # type: ignore
-        == custom_taskiq_result.return_value.test_arg  # type: ignore
-    )
-    assert result.is_err == custom_taskiq_result.is_err
-    assert result.execution_time == custom_taskiq_result.execution_time
-    assert result.log == custom_taskiq_result.log
+    with pytest.raises(ValueError):
+        await nats_result_backend.set_result(
+            task_id=task_id,
+            result=custom_taskiq_result,
+        )
 
 
 async def test_success_backend_is_result_ready(
     nats_result_backend: NATSObjectStoreResultBackend[_ReturnType],
-    custom_taskiq_result: TaskiqResult[_ReturnType],
     task_id: str,
 ) -> None:
     """Tests `is_result_ready` method."""
     assert not await nats_result_backend.is_result_ready(task_id=task_id)
     await nats_result_backend.set_result(
         task_id=task_id,
-        result=custom_taskiq_result,
+        result=TaskiqResult(
+            is_err=False,
+            log=None,
+            return_value="one",
+            execution_time=1,
+        ),
     )
 
     assert await nats_result_backend.is_result_ready(task_id=task_id)