Skip to content

Commit

Permalink
add AsyncKicker.with_schedule_id() (#338)
Browse files Browse the repository at this point in the history
  • Loading branch information
MehdiRtal authored Sep 6, 2024
1 parent ee4e5e6 commit 437687a
Showing 1 changed file with 20 additions and 2 deletions.
22 changes: 20 additions & 2 deletions taskiq/kicker.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ def __init__(
self.broker = broker
self.labels = labels
self.custom_task_id: Optional[str] = None
self.custom_schedule_id: Optional[str] = None

def with_labels(
self,
Expand Down Expand Up @@ -77,6 +78,19 @@ def with_task_id(self, task_id: str) -> "AsyncKicker[_FuncParams, _ReturnType]":
self.custom_task_id = task_id
return self

def with_schedule_id(
self,
schedule_id: str,
) -> "AsyncKicker[_FuncParams, _ReturnType]":
"""
Set schedule_id for current execution.
:param schedule_id: custom schedule id.
:return: kicker with custom schedule id.
"""
self.custom_schedule_id = schedule_id
return self

def with_broker(
self,
broker: "AsyncBroker",
Expand Down Expand Up @@ -166,7 +180,9 @@ async def schedule_by_cron(
:return: schedule id.
"""
schedule_id = self.broker.id_generator()
schedule_id = self.custom_schedule_id
if schedule_id is None:
schedule_id = self.broker.id_generator()
message = self._prepare_message(*args, **kwargs)
cron_offset = None
if isinstance(cron, CronSpec):
Expand Down Expand Up @@ -201,7 +217,9 @@ async def schedule_by_time(
:param args: function's args.
:param kwargs: function's kwargs.
"""
schedule_id = self.broker.id_generator()
schedule_id = self.custom_schedule_id
if schedule_id is None:
schedule_id = self.broker.id_generator()
message = self._prepare_message(*args, **kwargs)
scheduled = ScheduledTask(
schedule_id=schedule_id,
Expand Down

0 comments on commit 437687a

Please sign in to comment.