Merge pull request #16 from taskiq-python/issue-13

Add errors

Author: s3rius

taskiq/cli/async_task_runner.py

@@ -275,4 +275,9 @@ async def async_listen_messages(  # noqa: C901, WPS210, WPS213
         try:
             await broker.result_backend.set_result(message.task_id, result)
         except Exception as exc:
-            logger.exception(exc)
+            logger.exception(
+                "Can't set result in %s result backend: \n%s",
+                broker.result_backend.__class__.__name__,
+                exc,
+                exc_info=True,
+            )

taskiq/exceptions.py

@@ -4,3 +4,27 @@ class TaskiqError(Exception):
 
 class TaskiqResultTimeoutError(TaskiqError):
     """Waiting for task results has timed out."""
+
+
+class BrokerError(TaskiqError):
+    """Base class for all broker errors."""
+
+
+class SendTaskError(BrokerError):
+    """Error if the broker was unable to send the task to the queue."""
+
+
+class ResultBackendError(TaskiqError):
+    """Base class for all ResultBackend errors."""
+
+
+class ResultGetError(ResultBackendError):
+    """Error if ResultBackend was unable to get result."""
+
+
+class ResultSetError(ResultBackendError):
+    """Error if ResultBackend was unable to set result."""
+
+
+class ResultIsReadyError(ResultBackendError):
+    """Error if ResultBackend was unable to find out if the task is ready."""

taskiq/kicker.py

@@ -5,6 +5,7 @@
 
 from pydantic import BaseModel
 
+from taskiq.exceptions import SendTaskError
 from taskiq.message import TaskiqMessage
 from taskiq.task import AsyncTaskiqTask
 from taskiq.types_helpers import T_, FuncParams_, ReturnType_
@@ -105,13 +106,18 @@ async def kiq(
         :param args: function's arguments.
         :param kwargs: function's key word arguments.
 
+        :raises SendTaskError: if we can't send task to the broker.
+
         :returns: taskiq task.
         """
         logger.debug(
             f"Kicking {self.task_name} with args={args} and kwargs={kwargs}.",
         )
         message = self._prepare_message(*args, **kwargs)
-        await self.broker.kick(message)
+        try:
+            await self.broker.kick(message)
+        except Exception as exc:
+            raise SendTaskError() from exc
         return self.broker.result_backend.generate_task(message.task_id)
 
     @classmethod

taskiq/task.py

@@ -2,7 +2,11 @@
 from time import time
 from typing import TYPE_CHECKING, Generic
 
-from taskiq.exceptions import TaskiqResultTimeoutError
+from taskiq.exceptions import (
+    ResultGetError,
+    ResultIsReadyError,
+    TaskiqResultTimeoutError,
+)
 from taskiq.types_helpers import ReturnType_
 
 if TYPE_CHECKING:
@@ -25,18 +29,32 @@ async def is_ready(self) -> bool:
         """
         Checks if task is completed.
 
+        :raises ResultIsReadyError: if we can't get info about task readyness.
+
         :return: True if task is completed.
         """
-        return await self.result_backend.is_result_ready(self.task_id)
+        try:
+            return await self.result_backend.is_result_ready(self.task_id)
+        except Exception as exc:
+            raise ResultIsReadyError() from exc
 
     async def get_result(self, with_logs: bool = False) -> "TaskiqResult[ReturnType_]":
         """
         Get result of a task from result backend.
 
         :param with_logs: whether you want to fetch logs from worker.
+
+        :raises ResultGetError: if we can't get result from ResultBackend.
+
         :return: task's return value.
         """
-        return await self.result_backend.get_result(self.task_id, with_logs=with_logs)
+        try:
+            return await self.result_backend.get_result(
+                self.task_id,
+                with_logs=with_logs,
+            )
+        except Exception as exc:
+            raise ResultGetError() from exc
 
     async def wait_result(
         self,