Better Exception Serialization (#323)

Author: kraftp

dbos/_error.py

@@ -134,12 +134,17 @@ class DBOSNotAuthorizedError(DBOSException):
     """Exception raised by DBOS role-based security when the user is not authorized to access a function."""
 
     def __init__(self, msg: str):
+        self.msg = msg
         super().__init__(
             msg,
             dbos_error_code=DBOSErrorCode.NotAuthorized.value,
         )
         self.status_code = 403
 
+    def __reduce__(self) -> Any:
+        # Tell jsonpickle how to reconstruct this object
+        return (self.__class__, (self.msg,))
+
 
 class DBOSMaxStepRetriesExceeded(DBOSException):
     """Exception raised when a step was retried the maximimum number of times without success."""
@@ -185,11 +190,21 @@ class DBOSQueueDeduplicatedError(DBOSException):
     def __init__(
         self, workflow_id: str, queue_name: str, deduplication_id: str
     ) -> None:
+        self.workflow_id = workflow_id
+        self.queue_name = queue_name
+        self.deduplication_id = deduplication_id
         super().__init__(
             f"Workflow {workflow_id} was deduplicated due to an existing workflow in queue {queue_name} with deduplication ID {deduplication_id}.",
             dbos_error_code=DBOSErrorCode.QueueDeduplicated.value,
         )
 
+    def __reduce__(self) -> Any:
+        # Tell jsonpickle how to reconstruct this object
+        return (
+            self.__class__,
+            (self.workflow_id, self.queue_name, self.deduplication_id),
+        )
+
 
 #######################################
 ## BaseException

dbos/_serialization.py

@@ -1,8 +1,10 @@
 import types
-from typing import Any, Dict, Tuple, TypedDict
+from typing import Any, Dict, Optional, Tuple, TypedDict
 
 import jsonpickle  # type: ignore
 
+from ._logger import dbos_logger
+
 
 class WorkflowInputs(TypedDict):
     args: Tuple[Any, ...]
@@ -51,5 +53,54 @@ def deserialize_args(serialized_data: str) -> WorkflowInputs:
 
 def deserialize_exception(serialized_data: str) -> Exception:
     """Deserialize JSON string back to a Python Exception using jsonpickle."""
-    upo: Exception = jsonpickle.decode(serialized_data)
-    return upo
+    exc: Exception = jsonpickle.decode(serialized_data)
+    return exc
+
+
+def safe_deserialize(
+    workflow_id: str,
+    *,
+    serialized_input: Optional[str],
+    serialized_output: Optional[str],
+    serialized_exception: Optional[str],
+) -> tuple[Optional[WorkflowInputs], Optional[Any], Optional[Exception]]:
+    """
+    This function safely deserializes a workflow's recorded input and output/exception.
+    If any of them is not deserializable, it logs a warning and returns a string instead of throwing an exception.
+
+    This function is used in workflow introspection methods (get_workflows and get_queued_workflow)
+    to ensure errors related to nondeserializable objects are observable.
+    """
+    input: Optional[WorkflowInputs]
+    try:
+        input = (
+            deserialize_args(serialized_input) if serialized_input is not None else None
+        )
+    except Exception as e:
+        dbos_logger.warning(
+            f"Warning: input object could not be deserialized for workflow {workflow_id}, returning as string: {e}"
+        )
+        input = serialized_input  # type: ignore
+    output: Optional[Any]
+    try:
+        output = (
+            deserialize(serialized_output) if serialized_output is not None else None
+        )
+    except Exception as e:
+        dbos_logger.warning(
+            f"Warning: output object could not be deserialized for workflow {workflow_id}, returning as string: {e}"
+        )
+        output = serialized_output
+    exception: Optional[Exception]
+    try:
+        exception = (
+            deserialize_exception(serialized_exception)
+            if serialized_exception is not None
+            else None
+        )
+    except Exception as e:
+        dbos_logger.warning(
+            f"Warning: exception object could not be deserialized for workflow {workflow_id}, returning as string: {e}"
+        )
+        exception = serialized_exception  # type: ignore
+    return input, output, exception

dbos/_sys_db.py

@@ -901,13 +901,15 @@ def get_workflows(
             info.app_version = row[14]
             info.app_id = row[15]
 
-            inputs = _serialization.deserialize_args(row[16])
-            if inputs is not None:
-                info.input = inputs
-            if info.status == WorkflowStatusString.SUCCESS.value:
-                info.output = _serialization.deserialize(row[17])
-            elif info.status == WorkflowStatusString.ERROR.value:
-                info.error = _serialization.deserialize_exception(row[18])
+            inputs, output, exception = _serialization.safe_deserialize(
+                info.workflow_id,
+                serialized_input=row[16],
+                serialized_output=row[17],
+                serialized_exception=row[18],
+            )
+            info.input = inputs
+            info.output = output
+            info.error = exception
 
             infos.append(info)
         return infos
@@ -1007,13 +1009,15 @@ def get_queued_workflows(
             info.app_version = row[14]
             info.app_id = row[15]
 
-            inputs = _serialization.deserialize_args(row[16])
-            if inputs is not None:
-                info.input = inputs
-            if info.status == WorkflowStatusString.SUCCESS.value:
-                info.output = _serialization.deserialize(row[17])
-            elif info.status == WorkflowStatusString.ERROR.value:
-                info.error = _serialization.deserialize_exception(row[18])
+            inputs, output, exception = _serialization.safe_deserialize(
+                info.workflow_id,
+                serialized_input=row[16],
+                serialized_output=row[17],
+                serialized_exception=row[18],
+            )
+            info.input = inputs
+            info.output = output
+            info.error = exception
 
             infos.append(info)
 

tests/test_failures.py

@@ -11,9 +11,16 @@
 from dbos._error import (
     DBOSDeadLetterQueueError,
     DBOSMaxStepRetriesExceeded,
+    DBOSNotAuthorizedError,
+    DBOSQueueDeduplicatedError,
     DBOSUnexpectedStepError,
 )
 from dbos._registrations import DEFAULT_MAX_RECOVERY_ATTEMPTS
+from dbos._serialization import (
+    deserialize_exception,
+    safe_deserialize,
+    serialize_exception,
+)
 from dbos._sys_db import WorkflowStatusString
 
 from .conftest import queue_entries_are_cleaned_up
@@ -434,3 +441,42 @@ def failing_workflow() -> str:
     recovery_handles = DBOS._recover_pending_workflows()
     assert len(recovery_handles) == 1
     assert recovery_handles[0].get_result() == recovery_handles[0].workflow_id
+
+
+def test_error_serialization() -> None:
+    # Verify that each exception that can be thrown in a workflow
+    # is serializable and deserializable
+    # DBOSMaxStepRetriesExceeded
+    e: Exception = DBOSMaxStepRetriesExceeded("step", 1)
+    d = deserialize_exception(serialize_exception(e))
+    assert isinstance(d, DBOSMaxStepRetriesExceeded)
+    assert str(d) == str(e)
+    # DBOSNotAuthorizedError
+    e = DBOSNotAuthorizedError("no")
+    d = deserialize_exception(serialize_exception(e))
+    assert isinstance(d, DBOSNotAuthorizedError)
+    assert str(d) == str(e)
+    # DBOSQueueDeduplicatedError
+    e = DBOSQueueDeduplicatedError("id", "queue", "dedup")
+    d = deserialize_exception(serialize_exception(e))
+    assert isinstance(d, DBOSQueueDeduplicatedError)
+    assert str(d) == str(e)
+
+    # Test safe_deserialize
+    class BadException(Exception):
+        def __init__(self, one: int, two: int) -> None:
+            super().__init__(f"Message: {one}, {two}")
+
+    bad_exception = BadException(1, 2)
+    with pytest.raises(TypeError):
+        deserialize_exception(serialize_exception(bad_exception))
+    input, output, exception = safe_deserialize(
+        "my_id",
+        serialized_input=None,
+        serialized_exception=serialize_exception(bad_exception),
+        serialized_output=None,
+    )
+    assert input is None
+    assert output is None
+    assert isinstance(exception, str)
+    assert "Message: 1, 2" in exception