Added shutdown and startup events.

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

Author: s3rius

taskiq/abc/broker.py

@@ -147,7 +147,10 @@ def __init__(
         self.is_worker_process = False
         self._related_tasks: Set[AsyncTaskiqDecoratedTask[..., Any]] = set()
 
-    def close(self) -> None:
+    async def startup(self) -> None:
+        """Do something when starting broker."""
+
+    async def shutdown(self) -> None:
         """
         Close the broker.
 

taskiq/abc/result_backend.py

@@ -44,6 +44,12 @@ def value(self) -> Any:
 class AsyncResultBackend(ABC, Generic[_ReturnType]):
     """Async result backend."""
 
+    async def startup(self) -> None:
+        """Do something when starting broker."""
+
+    async def shutdown(self) -> None:
+        """Do something on shutdown."""
+
     def generate_task(self, task_id: str) -> "AsyncTaskiqTask[_ReturnType]":
         """
         Generates new task.

taskiq/cli/task_runner.py

@@ -169,6 +169,8 @@ async def async_listen_messages(  # noqa: C901, WPS210
     :param broker: broker to listen to.
     :param cli_args: CLI arguments for worker.
     """
+    logger.info("Runing startup event.")
+    await broker.startup()
     logger.info("Listening started.")
     task_registry: Dict[str, Callable[..., Any]] = {}
     task_signatures: Dict[str, inspect.Signature] = {}