add AsyncKicker.with_schedule_id() (#338)

Author: MehdiRtal

taskiq/kicker.py

@@ -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,
@@ -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",
@@ -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):
@@ -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,