BUG FIX: issue:493 Loss of task_id when using retry middlewares (#496)

* issue:493 add: test case add: task_id for retry scheduled task --------- Co-authored-by: SARomanchuk

Author: GefMar

taskiq/kicker.py

@@ -69,7 +69,10 @@ def with_labels(
  self.labels.update(labels)
  return self
 
- def with_task_id(self, task_id: str) -> "AsyncKicker[_FuncParams, _ReturnType]":
+ def with_task_id(
+ self,
+ task_id: Optional[str],
+ ) -> "AsyncKicker[_FuncParams, _ReturnType]":
  """
  Set task_id for current execution.
 
@@ -208,6 +211,7 @@ async def schedule_by_cron(
  labels=message.labels,
  args=message.args,
  kwargs=message.kwargs,
+ task_id=self.custom_task_id,
  cron=cron_str,
  cron_offset=cron_offset,
  )
@@ -239,6 +243,7 @@ async def schedule_by_time(
  labels=message.labels,
  args=message.args,
  kwargs=message.kwargs,
+ task_id=self.custom_task_id,
  time=time,
  )
  await source.add_schedule(scheduled)

taskiq/scheduler/scheduled_task/v1.py

@@ -12,6 +12,7 @@ class ScheduledTask(BaseModel):
  labels: Dict[str, Any]
  args: List[Any]
  kwargs: Dict[str, Any]
+ task_id: Optional[str] = None
  schedule_id: str = Field(default_factory=lambda: uuid.uuid4().hex)
  cron: Optional[str] = None
  cron_offset: Optional[Union[str, timedelta]] = None

taskiq/scheduler/scheduled_task/v2.py

@@ -13,6 +13,7 @@ class ScheduledTask(BaseModel):
  labels: Dict[str, Any]
  args: List[Any]
  kwargs: Dict[str, Any]
+ task_id: Optional[str] = None
  schedule_id: str = Field(default_factory=lambda: uuid.uuid4().hex)
  cron: Optional[str] = None
  cron_offset: Optional[Union[str, timedelta]] = None

taskiq/scheduler/scheduler.py

@@ -51,6 +51,7 @@ async def on_ready(self, source: "ScheduleSource", task: ScheduledTask) -> None:
  .with_labels(
  schedule_id=task.schedule_id,
  )
+ .with_task_id(task_id=task.task_id)
  .kiq(
  *task.args,
  **task.kwargs,

tests/test_retry_task.py

@@ -0,0 +1,43 @@
+import pytest
+
+from taskiq import (
+ Context,
+ InMemoryBroker,
+ SmartRetryMiddleware,
+ TaskiqDepends,
+ TaskiqScheduler,
+)
+from taskiq.schedule_sources import LabelScheduleSource
+
+
+@pytest.mark.parametrize(
+ "retry_count",
+ range(5),
+)
+@pytest.mark.anyio
+async def test_save_task_id_for_retry(retry_count: int) -> None:
+ broker = InMemoryBroker().with_middlewares(
+ SmartRetryMiddleware(
+ default_retry_count=retry_count + 1,
+ default_delay=0.1,
+ ),
+ )
+ scheduler = TaskiqScheduler(broker, [LabelScheduleSource(broker)])
+
+ check_interval = 0.5
+
+ @broker.task("exc_task", retry_on_error=True)
+ async def exc_task(count: int = 0, context: "Context" = TaskiqDepends()) -> int:
+ retry = int(context.message.labels.get("_retries", 0))
+ if retry < count:
+ raise Exception("test")
+ return retry
+
+ await broker.startup()
+ await scheduler.startup()
+
+ task_with_retry = await exc_task.kiq(retry_count)
+ task_with_retry_result = await task_with_retry.wait_result(
+ check_interval=check_interval,
+ )
+ assert task_with_retry_result.return_value == retry_count