Added post-save hook in middlewares. (#41)

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

Author: s3rius

taskiq/abc/middleware.py

@@ -81,6 +81,21 @@ def post_execute(
         :param result: result of execution for current task.
         """
 
+    def post_save(
+        self,
+        message: "TaskiqMessage",
+        result: "TaskiqResult[Any]",
+    ) -> "Union[None, Coroutine[Any, Any, None]]":
+        """
+        Post save hook.
+
+        This function is called after result of
+        the executions is saved in the result_backend.
+
+        :param message: processed message.
+        :param result: returned value.
+        """
+
     def on_error(
         self,
         message: "TaskiqMessage",

taskiq/cli/receiver.py

@@ -73,7 +73,7 @@ def __init__(self, broker: AsyncBroker, cli_args: TaskiqArgs) -> None:
             max_workers=cli_args.max_threadpool_threads,
         )
 
-    async def callback(  # noqa: C901
+    async def callback(  # noqa: C901, WPS213
         self,
         message: BrokerMessage,
         raise_err: bool = False,
@@ -142,6 +142,10 @@ async def callback(  # noqa: C901
             if raise_err:
                 raise exc
 
+        for middleware in self.broker.middlewares:
+            if middleware.__class__.post_save != TaskiqMiddleware.post_save:
+                await maybe_awaitable(middleware.post_save(taskiq_msg, result))
+
     async def run_task(  # noqa: C901, WPS210
         self,
         target: Callable[..., Any],