Merge tag '0.0.5' into develop

0.0.5

Author: s3rius

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "taskiq"
-version = "0.0.4"
+version = "0.0.5"
 description = "Distributed task queue with full async support"
 authors = ["Pavel Kirilin <[email protected]>"]
 maintainers = ["Pavel Kirilin <[email protected]>"]

taskiq/cli/async_task_runner.py

@@ -167,13 +167,14 @@ async def run_task(  # noqa: C901, WPS210, WPS211
     )
     if found_exception is not None:
         for middleware in middlewares:
-            await maybe_awaitable(
-                middleware.on_error(
-                    message,
-                    result,
-                    found_exception,
-                ),
-            )
+            if middleware.__class__.on_error != TaskiqMiddleware.on_error:
+                await maybe_awaitable(
+                    middleware.on_error(
+                        message,
+                        result,
+                        found_exception,
+                    ),
+                )
 
     return result
 
@@ -223,12 +224,18 @@ async def async_listen_messages(  # noqa: C901, WPS210, WPS213
                 exc_info=True,
             )
             continue
+        logger.info(
+            "Executing task %s with ID: %s",
+            taskiq_msg.task_name,
+            taskiq_msg.task_id,
+        )
         for middleware in broker.middlewares:
-            taskiq_msg = await maybe_awaitable(
-                middleware.pre_execute(
-                    taskiq_msg,
-                ),
-            )
+            if middleware.__class__.pre_execute != TaskiqMiddleware.pre_execute:
+                taskiq_msg = await maybe_awaitable(
+                    middleware.pre_execute(
+                        taskiq_msg,
+                    ),
+                )
 
         result = await run_task(
             target=broker.available_tasks[message.task_name].original_func,
@@ -239,7 +246,8 @@ async def async_listen_messages(  # noqa: C901, WPS210, WPS213
             middlewares=broker.middlewares,
         )
         for middleware in broker.middlewares:
-            await maybe_awaitable(middleware.post_execute(taskiq_msg, result))
+            if middleware.__class__.post_execute != TaskiqMiddleware.post_execute:
+                await maybe_awaitable(middleware.post_execute(taskiq_msg, result))
         try:
             await broker.result_backend.set_result(message.task_id, result)
         except Exception as exc:

taskiq/kicker.py

@@ -1,11 +1,12 @@
 from dataclasses import asdict, is_dataclass
 from logging import getLogger
-from typing import (
+from typing import (  # noqa: WPS235
     TYPE_CHECKING,
     Any,
     Coroutine,
     Dict,
     Generic,
+    Optional,
     TypeVar,
     Union,
     overload,
@@ -14,6 +15,7 @@
 from pydantic import BaseModel
 from typing_extensions import ParamSpec
 
+from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.exceptions import SendTaskError
 from taskiq.message import TaskiqMessage
 from taskiq.task import AsyncTaskiqTask, SyncTaskiqTask
@@ -41,6 +43,7 @@ def __init__(
         self.task_name = task_name
         self.broker = broker
         self.labels = labels
+        self.custom_task_id: Optional[str] = None
 
     def with_labels(
         self,
@@ -55,6 +58,19 @@ def with_labels(
         self.labels.update(labels)
         return self
 
+    def with_task_id(self, task_id: str) -> "AsyncKicker[_FuncParams, _ReturnType]":
+        """
+        Set task_id for current execution.
+
+        Please use this method with caution,
+        because it may brake the logic of getting results.
+
+        :param task_id: custom task id.
+        :return: kicker with custom task id.
+        """
+        self.custom_task_id = task_id
+        return self
+
     def with_broker(
         self,
         broker: "AsyncBroker",
@@ -87,7 +103,7 @@ async def kiq(  # noqa: D102
     ) -> AsyncTaskiqTask[_ReturnType]:
         ...
 
-    async def kiq(
+    async def kiq(  # noqa: C901
         self,
         *args: _FuncParams.args,
         **kwargs: _FuncParams.kwargs,
@@ -110,15 +126,16 @@ async def kiq(
         )
         message = self._prepare_message(*args, **kwargs)
         for middleware in self.broker.middlewares:
-            message = await maybe_awaitable(middleware.pre_send(message))
-
+            if middleware.__class__.pre_send != TaskiqMiddleware.pre_send:
+                message = await maybe_awaitable(middleware.pre_send(message))
         try:
             await self.broker.kick(self.broker.formatter.dumps(message))
         except Exception as exc:
             raise SendTaskError() from exc
 
         for middleware in self.broker.middlewares:
-            await maybe_awaitable(middleware.post_send(message))
+            if middleware.__class__.post_send != TaskiqMiddleware.post_send:
+                await maybe_awaitable(middleware.post_send(message))
 
         return AsyncTaskiqTask(
             task_id=message.task_id,
@@ -198,8 +215,12 @@ def _prepare_message(  # noqa: WPS210
         for label, label_val in self.labels.items():
             labels[label] = str(label_val)
 
+        task_id = self.custom_task_id
+        if task_id is None:
+            task_id = self.broker.id_generator()
+
         return TaskiqMessage(
-            task_id=self.broker.id_generator(),
+            task_id=task_id,
             task_name=self.task_name,
             labels=labels,
             args=formatted_args,

taskiq/middlewares/retry_middleware.py

@@ -1,3 +1,4 @@
+from copy import deepcopy
 from logging import getLogger
 from typing import Any
 
@@ -40,19 +41,19 @@ async def on_error(
         # Check if retrying is enabled for the task.
         if retry_on_error != "True":
             return
-
+        new_msg = deepcopy(message)
         # Getting number of previous retries.
-        retries = int(message.labels.get("_retries", 0)) + 1
-        message.labels["_retries"] = str(retries)
-        max_retries = int(message.labels.get("max_retries", self.default_retry_count))
+        retries = int(new_msg.labels.get("_retries", 0)) + 1
+        new_msg.labels["_retries"] = str(retries)
+        max_retries = int(new_msg.labels.get("max_retries", self.default_retry_count))
         if retries < max_retries:
             logger.info(
                 "Task '%s' invocation failed. Retrying.",
                 message.task_name,
             )
-            message.labels["_parent"] = message.task_id
-            message.task_id = self.broker.id_generator()
-            broker_message = self.broker.formatter.dumps(message=message)
+            new_msg.labels["_parent"] = message.task_id
+            new_msg.task_id = self.broker.id_generator()
+            broker_message = self.broker.formatter.dumps(message=new_msg)
             await self.broker.kick(broker_message)
         else:
             logger.warning(