added custom serializers

Author: Andrei Neagu

packages/service-library/src/servicelib/aiohttp/long_running_tasks/_server.py

@@ -5,19 +5,23 @@
 from typing import Any
 
 from aiohttp import web
+from aiohttp.web import HTTPException
 from common_library.json_serialization import json_dumps
 from pydantic import AnyHttpUrl, TypeAdapter
-from servicelib.long_running_tasks.task import Namespace
 from settings_library.redis import RedisSettings
 
 from ...aiohttp import status
 from ...long_running_tasks import lrt_api
+from ...long_running_tasks._error_serialization import (
+    BaseObjectSerializer,
+    register_custom_serialization,
+)
 from ...long_running_tasks.constants import (
     DEFAULT_STALE_TASK_CHECK_INTERVAL,
     DEFAULT_STALE_TASK_DETECT_TIMEOUT,
 )
 from ...long_running_tasks.models import TaskContext, TaskGet
-from ...long_running_tasks.task import RegisteredTaskName
+from ...long_running_tasks.task import Namespace, RegisteredTaskName
 from ..typing_extension import Handler
 from . import _routes
 from ._constants import (
@@ -118,6 +122,22 @@ def _wrap_and_add_routes(
         )
 
 
+class AiohttpHTTPExceptionSerializer(BaseObjectSerializer[HTTPException]):
+    @classmethod
+    def get_init_kwargs_from_object(cls, obj: HTTPException) -> dict:
+        return {
+            "status_code": obj.status_code,
+            "reason": obj.reason,
+            "text": obj.text,
+            "headers": dict(obj.headers) if obj.headers else None,
+        }
+
+    @classmethod
+    def prepare_object_init_kwargs(cls, data: dict) -> dict:
+        data.pop("status_code")
+        return data
+
+
 def setup(
     app: web.Application,
     *,
@@ -140,6 +160,8 @@ def setup(
     """
 
     async def on_cleanup_ctx(app: web.Application) -> AsyncGenerator[None, None]:
+        register_custom_serialization(HTTPException, AiohttpHTTPExceptionSerializer)
+
         # add error handlers
         app.middlewares.append(base_long_running_error_handler)
 

packages/service-library/src/servicelib/long_running_tasks/_error_serialization.py

@@ -0,0 +1,85 @@
+import base64
+import logging
+import pickle
+from abc import ABC, abstractmethod
+from typing import Final, Generic, TypeVar
+
+_logger = logging.getLogger(__name__)
+
+
+T = TypeVar("T")
+
+
+class BaseObjectSerializer(ABC, Generic[T]):
+
+    @classmethod
+    @abstractmethod
+    def get_init_kwargs_from_object(cls, obj: T) -> dict:
+        """dictionary reppreseting the kwargs passed to the __init__ method"""
+
+    @classmethod
+    @abstractmethod
+    def prepare_object_init_kwargs(cls, data: dict) -> dict:
+        """cleanup data to be used as kwargs for the __init__ method if required"""
+
+
+_SERIALIZERS: Final[dict[type, type[BaseObjectSerializer]]] = {}
+
+
+def register_custom_serialization(
+    object_type: type, object_serializer: type[BaseObjectSerializer]
+) -> None:
+    """Register a custom serializer for a specific object type.
+
+    Arguments:
+        object_type -- the type or parent class of the object to be serialized
+        object_serializer -- custom implementation of BaseObjectSerializer for the object type
+    """
+    _SERIALIZERS[object_type] = object_serializer
+
+
+_TYPE_FIELD: Final[str] = "__pickle__type__field__"
+_MODULE_FIELD: Final[str] = "__pickle__module__field__"
+
+
+def error_to_string(e: Exception) -> str:
+    """Serialize exception to base64-encoded string."""
+    to_serialize: Exception | dict = e
+    object_class = type(e)
+
+    for registered_class, object_serializer in _SERIALIZERS.items():
+        if issubclass(object_class, registered_class):
+            to_serialize = {
+                _TYPE_FIELD: type(e).__name__,
+                _MODULE_FIELD: type(e).__module__,
+                **object_serializer.get_init_kwargs_from_object(e),
+            }
+            break
+
+    return base64.b85encode(pickle.dumps(to_serialize)).decode("utf-8")
+
+
+def error_from_string(error_str: str) -> Exception:
+    """Deserialize exception from base64-encoded string."""
+    data = pickle.loads(base64.b85decode(error_str))  # noqa: S301
+
+    if isinstance(data, dict) and _TYPE_FIELD in data and _MODULE_FIELD in data:
+        try:
+            # Import the module and get the exception class
+            module = __import__(data[_MODULE_FIELD], fromlist=[data[_TYPE_FIELD]])
+            exception_class = getattr(module, data[_TYPE_FIELD])
+
+            for registered_class, object_serializer in _SERIALIZERS.items():
+                if issubclass(exception_class, registered_class):
+                    # remove unrequired
+                    data.pop(_TYPE_FIELD)
+                    data.pop(_MODULE_FIELD)
+
+                    return exception_class(  # type: ignore[no-any-return]
+                        **object_serializer.prepare_object_init_kwargs(data)
+                    )
+        except (ImportError, AttributeError, TypeError) as e:
+            msg = f"Could not reconstruct object from data: {data}"
+            raise ValueError(msg) from e
+
+    return data  # type: ignore[no-any-return]

packages/service-library/src/servicelib/long_running_tasks/task.py

@@ -1,10 +1,8 @@
 import asyncio
-import base64
 import datetime
 import functools
 import inspect
 import logging
-import pickle
 import traceback
 import urllib.parse
 from contextlib import suppress
@@ -19,6 +17,7 @@
 from settings_library.redis import RedisDatabase, RedisSettings
 
 from ..redis import RedisClientSDK, exclusive
+from ._error_serialization import error_from_string, error_to_string
 from ._store.base import BaseStore
 from ._store.redis import RedisStore
 from .errors import (
@@ -101,14 +100,6 @@ async def _get_tasks_to_remove(
     return tasks_to_remove
 
 
-def _error_to_string(e: Exception) -> str:
-    return base64.b85encode(pickle.dumps(e)).decode("utf-8")
-
-
-def _error_from_string(error_str: str) -> Exception:
-    return pickle.loads(base64.b85decode(error_str))  # type: ignore[no-any-return]  # noqa: S301
-
-
 class TasksManager:  # pylint:disable=too-many-instance-attributes
     """
     Monitors execution and results retrieval of a collection of asyncio.Tasks
@@ -278,10 +269,10 @@ async def _status_update_worker(self) -> None:
                     continue
                 except asyncio.CancelledError:
                     task_data.result_field = ResultField(
-                        error=_error_to_string(TaskCancelledError(task_id=task_id))
+                        error=error_to_string(TaskCancelledError(task_id=task_id))
                     )
                 except Exception as e:  # pylint:disable=broad-except
-                    task_data.result_field = ResultField(error=_error_to_string(e))
+                    task_data.result_field = ResultField(error=error_to_string(e))
 
                 await self._tasks_data.set_task_data(task_id, task_data)
 
@@ -369,7 +360,7 @@ async def get_task_result(
             raise TaskNotCompletedError(task_id=task_id)
 
         if tracked_task.result_field.error is not None:
-            raise _error_from_string(tracked_task.result_field.error)
+            raise error_from_string(tracked_task.result_field.error)
 
         return tracked_task.result_field.result
 

packages/service-library/tests/long_running_tasks/test_long_running_tasks__error_serialization.py

@@ -0,0 +1,47 @@
+from typing import Any
+
+import pytest
+from aiohttp.web import HTTPException, HTTPInternalServerError
+from servicelib.aiohttp.long_running_tasks._server import AiohttpHTTPExceptionSerializer
+from servicelib.long_running_tasks._error_serialization import (
+    error_from_string,
+    error_to_string,
+    register_custom_serialization,
+)
+
+register_custom_serialization(HTTPException, AiohttpHTTPExceptionSerializer)
+
+
+class PositionalArguments:
+    def __init__(self, arg1, arg2, *args):
+        self.arg1 = arg1
+        self.arg2 = arg2
+        self.args = args
+
+
+class MixedArguments:
+    def __init__(self, arg1, arg2, kwarg1=None, kwarg2=None):
+        self.arg1 = arg1
+        self.arg2 = arg2
+        self.kwarg1 = kwarg1
+        self.kwarg2 = kwarg2
+
+
+@pytest.mark.parametrize(
+    "obj",
+    [
+        HTTPInternalServerError(reason="Uh-oh!", text="Failure!"),
+        PositionalArguments("arg1", "arg2", "arg3", "arg4"),
+        MixedArguments("arg1", "arg2", kwarg1="kwarg1", kwarg2="kwarg2"),
+        "a_string",
+        1,
+    ],
+)
+def test_serialization(obj: Any):
+    str_data = error_to_string(obj)
+
+    reconstructed_obj = error_from_string(str_data)
+
+    assert type(reconstructed_obj) is type(obj)
+    if hasattr(obj, "__dict__"):
+        assert reconstructed_obj.__dict__ == obj.__dict__