feat: exceptions serialization (#135)

Co-authored-by: Anton <[email protected]>

Author: Sobes76rus

.flake8

@@ -103,6 +103,8 @@ ignore =
   F403,
   ; Found wrong metadata variable
   WPS410,
+  ; Found commented out cod
+  E800,
 
 per-file-ignores =
   ; all tests

taskiq/abc/middleware.py

@@ -97,7 +97,7 @@ def on_error(
         self,
         message: "TaskiqMessage",
         result: "TaskiqResult[Any]",
-        exception: Exception,
+        exception: BaseException,
     ) -> "Union[None, Coroutine[Any, Any, None]]":
         """
         This function is called when exception is found.

taskiq/exceptions.py

@@ -28,3 +28,7 @@ class ResultSetError(ResultBackendError):
 
 class ResultIsReadyError(ResultBackendError):
     """Error if ResultBackend was unable to find out if the task is ready."""
+
+
+class SecurityError(TaskiqError):
+    """Security related exception."""

taskiq/middlewares/retry_middleware.py

@@ -22,7 +22,7 @@ async def on_error(
         self,
         message: "TaskiqMessage",
         result: "TaskiqResult[Any]",
-        exception: Exception,
+        exception: BaseException,
     ) -> None:
         """
         Retry on error.

taskiq/receiver/receiver.py

@@ -201,7 +201,7 @@ async def run_task(  # noqa: C901, WPS210
                     target,
                     message,
                 )
-        except Exception as exc:
+        except BaseException as exc:  # noqa: WPS424
             found_exception = exc
             logger.error(
                 "Exception found while executing function: %s",
@@ -219,6 +219,7 @@ async def run_task(  # noqa: C901, WPS210
             log=None,
             return_value=returned,
             execution_time=round(execution_time, 2),
+            error=found_exception,
         )
         # If exception is found we execute middlewares.
         if found_exception is not None:

taskiq/result.py

@@ -1,10 +1,28 @@
-from typing import Generic, Optional, TypeVar
+import json
+import pickle  # noqa: S403
+from functools import partial
+from typing import Any, Callable, Dict, Generic, Optional, TypeVar
 
+from pydantic import validator
 from pydantic.generics import GenericModel
+from typing_extensions import Self
+
+from taskiq.serialization import exception_to_python, prepare_exception
 
 _ReturnType = TypeVar("_ReturnType")
 
 
+def _json_encoder(value: Any, default: Callable[[Any], Any]) -> Any:
+    if isinstance(value, BaseException):
+        return prepare_exception(value, json)
+
+    return default(value)
+
+
+def _json_dumps(value: Any, *, default: Callable[[Any], Any], **kwargs: Any) -> str:
+    return json.dumps(value, default=partial(_json_encoder, default=default), **kwargs)
+
+
 class TaskiqResult(GenericModel, Generic[_ReturnType]):
     """Result of a remote task invocation."""
 
@@ -15,3 +33,37 @@ class TaskiqResult(GenericModel, Generic[_ReturnType]):
     log: Optional[str] = None
     return_value: _ReturnType
     execution_time: float
+
+    error: Optional[BaseException] = None
+
+    class Config:
+        arbitrary_types_allowed = True
+        json_dumps = _json_dumps  # type: ignore
+        json_loads = json.loads
+
+    def raise_for_error(self) -> "Self":
+        """Raise exception if `error`.
+
+        :raises error: task execution exception
+        :returns: TaskiqResult
+        """
+        if self.error is not None:
+            raise self.error
+        return self
+
+    def __getstate__(self) -> Dict[Any, Any]:
+        dict = super().__getstate__()  # noqa: WPS125
+        vals: Dict[str, Any] = dict["__dict__"]
+
+        if "error" in vals and vals["error"] is not None:
+            vals["error"] = prepare_exception(
+                vals["error"],
+                pickle,
+            )
+
+        return dict
+
+    @validator("error", pre=True)
+    @classmethod
+    def _validate_error(cls, value: Any) -> Optional[BaseException]:
+        return exception_to_python(value)