Merge pull request #12 from taskiq-python/develop

feat: add custom scheduler to support get_message callbacks

Author: Lancetnik

README.md

@@ -61,9 +61,8 @@ async def handler(msg: str):
     print(msg)
 
 # taskiq-faststream scheduling
-from taskiq import TaskiqScheduler
 from taskiq.schedule_sources import LabelScheduleSource
-from taskiq_faststream import BrokerWrapper
+from taskiq_faststream import BrokerWrapper, StreamScheduler
 
 # wrap FastStream object
 taskiq_broker = BrokerWrapper(broker)
@@ -80,7 +79,7 @@ taskiq_broker.task(
 )
 
 # create scheduler object
-scheduler = TaskiqScheduler(
+scheduler = StreamScheduler(
     broker=taskiq_broker,
     sources=[LabelScheduleSource(taskiq_broker)],
 )

pyproject.toml

@@ -71,10 +71,10 @@ test = [
 dev = [
     "taskiq-faststream[test]",
 
-    "mypy==1.7.0",
+    "mypy==1.7.1",
     "black==23.11.0",
     "isort==5.12.0",
-    "ruff==0.1.5",
+    "ruff==0.1.6",
     "pyupgrade-directories==0.3.0",
     "pre-commit==3.5.0",
 ]

taskiq_faststream/__about__.py

@@ -1,2 +1,2 @@
 """FastStream - taskiq integration to schedule FastStream tasks."""
-__version__ = "0.1.4"
+__version__ = "0.1.5"

taskiq_faststream/__init__.py

@@ -1,6 +1,8 @@
 from taskiq_faststream.broker import AppWrapper, BrokerWrapper
+from taskiq_faststream.scheduler import StreamScheduler
 
 __all__ = (
     "BrokerWrapper",
+    "StreamScheduler",
     "AppWrapper",
 )

taskiq_faststream/broker.py

@@ -10,6 +10,7 @@
 from taskiq.acks import AckableMessage
 from taskiq.decor import AsyncTaskiqDecoratedTask
 
+from taskiq_faststream.serializer import PatchedSerializer
 from taskiq_faststream.types import ScheduledTask
 from taskiq_faststream.utils import resolve_msg
 
@@ -32,6 +33,7 @@ class BrokerWrapper(AsyncBroker):
 
     def __init__(self, broker: BrokerAsyncUsecase[typing.Any, typing.Any]) -> None:
         super().__init__()
+        self.serializer = PatchedSerializer()
         self.broker = broker
 
     async def startup(self) -> None:
@@ -107,6 +109,7 @@ class AppWrapper(BrokerWrapper):
 
     def __init__(self, app: FastStream) -> None:
         super(BrokerWrapper, self).__init__()
+        self.serializer = PatchedSerializer()
         self.app = app
 
     async def startup(self) -> None:

taskiq_faststream/kicker.py

@@ -0,0 +1,13 @@
+from typing import Any
+
+from taskiq.kicker import AsyncKicker, _FuncParams, _ReturnType
+from taskiq.message import TaskiqMessage
+
+
+class LabelRespectKicker(AsyncKicker[_FuncParams, _ReturnType]):
+    """Patched kicker doesn't cast labels to str."""
+
+    def _prepare_message(self, *args: Any, **kwargs: Any) -> TaskiqMessage:
+        msg = super()._prepare_message(*args, **kwargs)
+        msg.labels = self.labels
+        return msg

taskiq_faststream/scheduler.py

@@ -0,0 +1,28 @@
+from typing import TYPE_CHECKING
+
+from taskiq.scheduler.scheduled_task import ScheduledTask
+from taskiq.scheduler.scheduler import TaskiqScheduler as Scheduler
+from taskiq.utils import maybe_awaitable
+
+from taskiq_faststream.kicker import LabelRespectKicker
+
+if TYPE_CHECKING:  # pragma: no cover
+    from taskiq.abc.schedule_source import ScheduleSource
+
+
+class StreamScheduler(Scheduler):
+    """Patched scheduler with custom kicker."""
+
+    async def on_ready(self, source: "ScheduleSource", task: ScheduledTask) -> None:
+        """
+        This method is called when task is ready to be enqueued.
+
+        It's triggered on proper time depending on `task.cron` or `task.time` attribute.
+        :param task: task to send
+        """
+        await maybe_awaitable(source.pre_send(task))
+        await LabelRespectKicker(task.task_name, self.broker, task.labels).kiq(
+            *task.args,
+            **task.kwargs,
+        )
+        await maybe_awaitable(source.post_send(task))

taskiq_faststream/serializer.py

@@ -0,0 +1,17 @@
+from typing import Any
+
+from taskiq.serializers.json_serializer import JSONSerializer
+
+
+class PatchedSerializer(JSONSerializer):
+    """Patched serializer removes labels."""
+
+    def dumpb(self, value: Any) -> bytes:
+        """
+        Dumps taskiq message to some broker message format.
+
+        :param message: message to send.
+        :return: Dumped message.
+        """
+        del value["labels"]
+        return super().dumpb(value)