refactor

Author: Andrei Neagu

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

@@ -12,7 +12,7 @@
 
 from ...aiohttp import status
 from ...long_running_tasks import lrt_api
-from ...long_running_tasks._error_serialization import (
+from ...long_running_tasks._redis_serialization import (
     BaseObjectSerializer,
     register_custom_serialization,
 )

packages/service-library/src/servicelib/long_running_tasks/_error_serialization.py renamed to packages/service-library/src/servicelib/long_running_tasks/_redis_serialization.py

@@ -2,7 +2,7 @@
 import logging
 import pickle
 from abc import ABC, abstractmethod
-from typing import Final, Generic, TypeVar
+from typing import Any, Final, Generic, TypeVar
 
 _logger = logging.getLogger(__name__)
 
@@ -42,9 +42,9 @@ def register_custom_serialization(
 _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
+def object_to_string(e: Any) -> str:
+    """Serialize object to base64-encoded string."""
+    to_serialize: Any | dict = e
     object_class = type(e)
 
     for registered_class, object_serializer in _SERIALIZERS.items():
@@ -59,8 +59,8 @@ def error_to_string(e: Exception) -> str:
     return base64.b85encode(pickle.dumps(to_serialize)).decode("utf-8")
 
 
-def error_from_string(error_str: str) -> Exception:
-    """Deserialize exception from base64-encoded string."""
+def string_to_object(error_str: str) -> Any:
+    """Deserialize object 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:

packages/service-library/src/servicelib/long_running_tasks/_store/redis.py

@@ -1,3 +1,4 @@
+import logging
 from typing import Any, Final
 
 import redis.asyncio as aioredis
@@ -9,6 +10,8 @@
 from ..models import TaskContext, TaskData, TaskId
 from .base import BaseStore
 
+_logger = logging.getLogger(__name__)
+
 STORE_TYPE_TASK_DATA: Final[str] = "TD"
 STORE_TYPE_CANCELLED_TASKS: Final[str] = "CT"
 
@@ -49,6 +52,9 @@ async def get_task_data(self, task_id: TaskId) -> TaskData | None:
         return TypeAdapter(TaskData).validate_json(result) if result else None
 
     async def set_task_data(self, task_id: TaskId, value: TaskData) -> None:
+        _logger.debug(
+            "Setting task data for task_id=%s with data value=%s", task_id, value
+        )
         await self.redis.hset(
             self._get_redis_hash_key(STORE_TYPE_TASK_DATA),
             task_id,

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

@@ -17,7 +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 ._redis_serialization import object_to_string, string_to_object
 from ._store.base import BaseStore
 from ._store.redis import RedisStore
 from .errors import (
@@ -263,16 +263,18 @@ async def _status_update_worker(self) -> None:
 
                 # get task result
                 try:
-                    task_data.result_field = ResultField(result=task.result())
+                    task_data.result_field = ResultField(
+                        result=object_to_string(task.result())
+                    )
                 except asyncio.InvalidStateError:
                     # task was not completed try again next time and see if it is done
                     continue
                 except asyncio.CancelledError:
                     task_data.result_field = ResultField(
-                        error=error_to_string(TaskCancelledError(task_id=task_id))
+                        error=object_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=object_to_string(e))
 
                 await self._tasks_data.set_task_data(task_id, task_data)
 
@@ -360,9 +362,9 @@ 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 string_to_object(tracked_task.result_field.error)
 
-        return tracked_task.result_field.result
+        return string_to_object(tracked_task.result_field.result)
 
     async def cancel_task(
         self, task_id: TaskId, with_task_context: TaskContext

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

@@ -3,10 +3,10 @@
 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,
+from servicelib.long_running_tasks._redis_serialization import (
+    object_to_string,
     register_custom_serialization,
+    string_to_object,
 )
 
 register_custom_serialization(HTTPException, AiohttpHTTPExceptionSerializer)
@@ -38,9 +38,9 @@ def __init__(self, arg1, arg2, kwarg1=None, kwarg2=None):
     ],
 )
 def test_serialization(obj: Any):
-    str_data = error_to_string(obj)
+    str_data = object_to_string(obj)
 
-    reconstructed_obj = error_from_string(str_data)
+    reconstructed_obj = string_to_object(str_data)
 
     assert type(reconstructed_obj) is type(obj)
     if hasattr(obj, "__dict__"):

packages/service-library/tests/long_running_tasks/test_long_running_tasks__redis_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._redis_serialization import (
+    object_to_string,
+    register_custom_serialization,
+    string_to_object,
+)
+
+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 = object_to_string(obj)
+
+    reconstructed_obj = string_to_object(str_data)
+
+    assert type(reconstructed_obj) is type(obj)
+    if hasattr(obj, "__dict__"):
+        assert reconstructed_obj.__dict__ == obj.__dict__