Added middlewares.

Signed-off-by: Pavel Kirilin <[email protected]>

Author: s3rius

taskiq/abc/broker.py

@@ -1,11 +1,13 @@
 from abc import ABC, abstractmethod
 from functools import wraps
 from logging import getLogger
-from typing import (
+from typing import (  # noqa: WPS235
+    TYPE_CHECKING,
     Any,
     AsyncGenerator,
     Callable,
     Dict,
+    List,
     Optional,
     TypeVar,
     Union,
@@ -14,13 +16,16 @@
 
 from typing_extensions import ParamSpec
 
-from taskiq.abc.plugins.formatter import TaskiqFormatter
-from taskiq.abc.result_backend import AsyncResultBackend
 from taskiq.decor import AsyncTaskiqDecoratedTask
 from taskiq.message import BrokerMessage
 from taskiq.plugins.json_formatter import JSONFormatter
 from taskiq.result_backends.dummy import DummyResultBackend
 
+if TYPE_CHECKING:
+    from taskiq.abc.formatter import TaskiqFormatter
+    from taskiq.abc.middleware import TaskiqMiddleware
+    from taskiq.abc.result_backend import AsyncResultBackend
+
 _T = TypeVar("_T")  # noqa: WPS111
 _FuncParams = ParamSpec("_FuncParams")
 _ReturnType = TypeVar("_ReturnType")
@@ -41,14 +46,28 @@ class AsyncBroker(ABC):
 
     def __init__(
         self,
-        result_backend: Optional[AsyncResultBackend[_T]] = None,
+        result_backend: "Optional[AsyncResultBackend[_T]]" = None,
     ) -> None:
         if result_backend is None:
             result_backend = DummyResultBackend()
+        self.middlewares: "List[TaskiqMiddleware]" = []
         self.result_backend = result_backend
         self.is_worker_process = False
         self.decorator_class = AsyncTaskiqDecoratedTask
-        self.formatter: TaskiqFormatter = JSONFormatter()
+        self.formatter: "TaskiqFormatter" = JSONFormatter()
+
+    def add_middlewares(self, middlewares: "List[TaskiqMiddleware]") -> None:
+        """
+        Add a list of middlewares.
+
+        You should call this method to set middlewares,
+        since it saves current broker in all middlewares.
+
+        :param middlewares: list of middlewares.
+        """
+        for middleware in middlewares:
+            middleware.set_broker(self)
+            self.middlewares.append(middleware)
 
     async def startup(self) -> None:
         """Do something when starting broker."""

taskiq/abc/plugins/formatter.py renamed to taskiq/abc/formatter.py

@@ -0,0 +1,85 @@
+from typing import TYPE_CHECKING, Any, Coroutine, Dict, Union
+
+if TYPE_CHECKING:
+    from taskiq.abc.broker import AsyncBroker
+    from taskiq.message import TaskiqMessage
+    from taskiq.result import TaskiqResult
+
+
+class TaskiqMiddleware:
+    """Base class for middlewares."""
+
+    def __init__(self) -> None:
+        self.broker: "AsyncBroker" = None  # type: ignore
+
+    def set_broker(self, broker: "AsyncBroker") -> None:
+        """
+        Sets broker to middleware.
+
+        :param broker: broker to set.
+        """
+        self.broker = broker
+
+    def pre_send(
+        self,
+        message: "TaskiqMessage",
+        labels: Dict[str, Any],
+    ) -> "Union[TaskiqMessage, Coroutine[Any, Any, TaskiqMessage]]":
+        """
+        Hook that executes before sending the task to worker.
+
+        This is a client-side hook, that executes right before
+        the message is sent to broker.
+
+        :param message: message to send.
+        :param labels: task's labels.
+        :return: modified message.
+        """
+        return message
+
+    def post_send(
+        self,
+        message: "TaskiqMessage",
+        labels: Dict[str, Any],
+    ) -> "Union[None, Coroutine[Any, Any, None]]":
+        """
+        This hook is executed right after the task is sent.
+
+        This is a client-side hook. It executes right
+        after the messages is kicked in broker.
+
+        :param message: kicked message.
+        :param labels: labels for a message.
+        """
+
+    def pre_execute(
+        self,
+        message: "TaskiqMessage",
+        labels: Dict[str, Any],
+    ) -> "Union[TaskiqMessage, Coroutine[Any, Any, TaskiqMessage]]":
+        """
+        This hook is called before executing task.
+
+        This is a worker-side hook, wich means it
+        executes in the worker process.
+
+        :param message: incoming parsed taskiq message.
+        :param labels: task's labels without user-supplied lables.
+        :return: modified message.
+        """
+        return message
+
+    def post_execute(
+        self,
+        result: "TaskiqResult[Any]",
+        labels: Dict[str, Any],
+    ) -> "Union[None, Coroutine[Any, Any, None]]":
+        """
+        This hook executes after task is complete.
+
+        This is a worker-side hook. It's called
+        in worker process.
+
+        :param result: result of execution for current task.
+        :param labels: task's labels. Without user-supplied labels.
+        """

taskiq/abc/middleware.py

@@ -2,7 +2,6 @@
 from typing import Generic, TypeVar
 
 from taskiq.result import TaskiqResult
-from taskiq.task import AsyncTaskiqTask
 
 _ReturnType = TypeVar("_ReturnType")
 
@@ -16,19 +15,6 @@ async def startup(self) -> None:
     async def shutdown(self) -> None:
         """Do something on shutdown."""
 
-    def generate_task(self, task_id: str) -> "AsyncTaskiqTask[_ReturnType]":
-        """
-        Generates new task.
-
-        This function creates new AsyncTaskiqTask
-        that returned to client after calling kiq
-        method.
-
-        :param task_id: id of a task to save.
-        :return: task object.
-        """
-        return AsyncTaskiqTask(task_id=task_id, result_backend=self)
-
     @abstractmethod
     async def set_result(self, task_id: str, result: TaskiqResult[_ReturnType]) -> None:
         """

taskiq/abc/result_backend.py

@@ -21,15 +21,9 @@ def kicker(self) -> AsyncKicker[_Params, _ReturnType]:
         our shared broker and send task to it, instead
         of shared_broker.
 
-        :raises TaskiqError: if _default_broker is not set.
         :return: new kicker.
         """
-        broker = getattr(self.broker, "_default_broker", None)
-        if broker is None:
-            raise TaskiqError(
-                "You cannot use kiq directly on shared task "
-                "without setting the default_broker.",
-            )
+        broker = getattr(self.broker, "_default_broker", None) or self.broker
         return AsyncKicker(
             task_name=self.task_name,
             broker=broker,
@@ -45,15 +39,6 @@ def __init__(self) -> None:
         self._default_broker: Optional[AsyncBroker] = None
         self.decorator_class = SharedDecoratedTask
 
-    async def kick(self, message: BrokerMessage) -> None:
-        """
-        Shared broker cannot kick tasks.
-
-        :param message: message to send.
-        :raises TaskiqError: if called.
-        """
-        raise TaskiqError("Shared broker cannot kick tasks.")
-
     def default_broker(self, new_broker: AsyncBroker) -> None:
         """
         Updates default broker.
@@ -62,6 +47,18 @@ def default_broker(self, new_broker: AsyncBroker) -> None:
         """
         self._default_broker = new_broker
 
+    async def kick(self, message: BrokerMessage) -> None:
+        """
+        Shared broker cannot kick tasks.
+
+        :param message: message to send.
+        :raises TaskiqError: if called.
+        """
+        raise TaskiqError(
+            "You cannot use kiq directly on shared task "
+            "without setting the default_broker.",
+        )
+
     async def listen(self) -> AsyncGenerator[BrokerMessage, None]:  # type: ignore
         """
         Shared broker cannot listen to tasks.

taskiq/brokers/shared_broker.py

@@ -276,13 +276,29 @@ async def async_listen_messages(  # noqa: C901, WPS210, WPS213
                 exc_info=True,
             )
             continue
+        for middleware in broker.middlewares:
+            pre_ex_res = middleware.pre_execute(
+                taskiq_msg,
+                broker.available_tasks[message.task_name].labels,
+            )
+            if inspect.isawaitable(pre_ex_res):
+                taskiq_msg = await pre_ex_res
+            else:
+                taskiq_msg = pre_ex_res  # type: ignore
         result = await run_task(
             broker.available_tasks[message.task_name].original_func,
             task_signatures.get(message.task_name),
             taskiq_msg,
             cli_args.log_collector_format,
             executor,
         )
+        for middleware in broker.middlewares:
+            post_ex_res = middleware.post_execute(
+                result,
+                broker.available_tasks[message.task_name].labels,
+            )
+            if inspect.isawaitable(post_ex_res):
+                await post_ex_res
         try:
             await broker.result_backend.set_result(message.task_id, result)
         except Exception as exc:

taskiq/cli/async_task_runner.py

@@ -1,6 +1,16 @@
 from dataclasses import asdict, is_dataclass
+from inspect import isawaitable
 from logging import getLogger
-from typing import TYPE_CHECKING, Any, Coroutine, Dict, Generic, TypeVar, overload
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Coroutine,
+    Dict,
+    Generic,
+    TypeVar,
+    Union,
+    overload,
+)
 from uuid import uuid4
 
 from pydantic import BaseModel
@@ -27,38 +37,27 @@ def __init__(
         self,
         task_name: str,
         broker: "AsyncBroker",
-        labels: Dict[str, Any],
+        labels: Dict[
+            str,
+            Union[
+                str,
+                int,
+                float,
+            ],
+        ],
     ) -> None:
         self.task_name = task_name
         self.broker = broker
         self.labels = labels
 
-    def with_label(
-        self,
-        label_name: str,
-        value: Any,
-    ) -> "AsyncKicker[_FuncParams, _ReturnType]":
-        """
-        Update one single label.
-
-        This method is used to update
-        task's labels before sending.
-
-        :param label_name: name of the label to update.
-        :param value: label's value.
-        :return: kicker object with new labels.
-        """
-        self.labels[label_name] = value
-        return self
-
     def with_labels(
         self,
-        labels: Dict[str, Any],
+        **labels: Union[str, int, float],
     ) -> "AsyncKicker[_FuncParams, _ReturnType]":
         """
         Update function's labels before sending.
 
-        :param labels: dict with new labels.
+        :param labels: new labels.
         :return: kicker with new labels.
         """
         self.labels.update(labels)
@@ -96,7 +95,7 @@ async def kiq(  # noqa: D102
     ) -> AsyncTaskiqTask[_ReturnType]:
         ...
 
-    async def kiq(
+    async def kiq(  # noqa: C901
         self,
         *args: _FuncParams.args,
         **kwargs: _FuncParams.kwargs,
@@ -118,11 +117,24 @@ async def kiq(
             f"Kicking {self.task_name} with args={args} and kwargs={kwargs}.",
         )
         message = self._prepare_message(*args, **kwargs)
+        for middleware in self.broker.middlewares:
+            pre_send_res = middleware.pre_send(message, self.labels)
+            if isawaitable(pre_send_res):
+                message = await pre_send_res
+            else:
+                message = pre_send_res  # type: ignore
         try:
             await self.broker.kick(self.broker.formatter.dumps(message, self.labels))
         except Exception as exc:
             raise SendTaskError() from exc
-        return self.broker.result_backend.generate_task(message.task_id)
+        for middleware in self.broker.middlewares:
+            post_send_res = middleware.post_send(message, self.labels)
+            if isawaitable(post_send_res):
+                await post_send_res
+        return AsyncTaskiqTask(
+            task_id=message.task_id,
+            result_backend=self.broker.result_backend,
+        )
 
     @classmethod
     def _prepare_arg(cls, arg: Any) -> Any:

taskiq/kicker.py

@@ -1,6 +1,6 @@
 from typing import Any, Dict
 
-from taskiq.abc.plugins.formatter import TaskiqFormatter
+from taskiq.abc.formatter import TaskiqFormatter
 from taskiq.message import BrokerMessage, TaskiqMessage