Custom Serialization (#486)

Add support for using custom serializers and deserializers in DBOS
Python. You can now supply a custom serializer in DBOS configuration,
which completely replaces the default `pickle`-based serializer for
storing objects in the system database.

For example, here's how you would use a JSON serializer:

```python
from dbos import DBOS, DBOSConfig, Serializer

class JsonSerializer(Serializer):
    def serialize(self, data: Any) -> str:
        return json.dumps(data)

    def deserialize(cls, serialized_data: str) -> Any:
        return json.loads(serialized_data)

serializer = JsonSerializer()
config: DBOSConfig = {
    "name": "dbos-starter",
    "system_database_url": os.environ.get("DBOS_SYSTEM_DATABASE_URL"),
    "serializer": serializer
}
DBOS(config=config)
DBOS.launch()
```

Addresses #485

Author: kraftp

dbos/__init__.py

@@ -12,6 +12,7 @@
 from ._debouncer import Debouncer, DebouncerClient
 from ._kafka_message import KafkaMessage
 from ._queue import Queue
+from ._serialization import Serializer
 from ._sys_db import GetWorkflowsInput, WorkflowStatus, WorkflowStatusString
 
 __all__ = [
@@ -35,4 +36,5 @@
     "Queue",
     "Debouncer",
     "DebouncerClient",
+    "Serializer",
 ]

dbos/_app_db.py

@@ -8,8 +8,8 @@
 from sqlalchemy.orm import Session, sessionmaker
 
 from dbos._migration import get_sqlite_timestamp_expr
+from dbos._serialization import Serializer
 
-from . import _serialization
 from ._error import DBOSUnexpectedStepError, DBOSWorkflowConflictIDError
 from ._logger import dbos_logger
 from ._schemas.application_database import ApplicationSchema
@@ -34,17 +34,52 @@ class RecordedResult(TypedDict):
 
 class ApplicationDatabase(ABC):
 
+    @staticmethod
+    def create(
+        database_url: str,
+        engine_kwargs: Dict[str, Any],
+        schema: Optional[str],
+        serializer: Serializer,
+        debug_mode: bool = False,
+    ) -> "ApplicationDatabase":
+        """Factory method to create the appropriate ApplicationDatabase implementation based on URL."""
+        if database_url.startswith("sqlite"):
+            return SQLiteApplicationDatabase(
+                database_url=database_url,
+                engine_kwargs=engine_kwargs,
+                schema=schema,
+                serializer=serializer,
+                debug_mode=debug_mode,
+            )
+        else:
+            # Default to PostgreSQL for postgresql://, postgres://, or other URLs
+            return PostgresApplicationDatabase(
+                database_url=database_url,
+                engine_kwargs=engine_kwargs,
+                schema=schema,
+                serializer=serializer,
+                debug_mode=debug_mode,
+            )
+
     def __init__(
         self,
         *,
         database_url: str,
         engine_kwargs: Dict[str, Any],
+        serializer: Serializer,
+        schema: Optional[str],
         debug_mode: bool = False,
     ):
+        if database_url.startswith("sqlite"):
+            self.schema = None
+        else:
+            self.schema = schema if schema else "dbos"
+        ApplicationSchema.transaction_outputs.schema = schema
         self.engine = self._create_engine(database_url, engine_kwargs)
         self._engine_kwargs = engine_kwargs
         self.sessionmaker = sessionmaker(bind=self.engine)
         self.debug_mode = debug_mode
+        self.serializer = serializer
 
     @abstractmethod
     def _create_engine(
@@ -156,10 +191,12 @@ def get_transactions(self, workflow_uuid: str) -> List[StepInfo]:
                 function_id=row[0],
                 function_name=row[1],
                 output=(
-                    _serialization.deserialize(row[2]) if row[2] is not None else row[2]
+                    self.serializer.deserialize(row[2])
+                    if row[2] is not None
+                    else row[2]
                 ),
                 error=(
-                    _serialization.deserialize_exception(row[3])
+                    self.serializer.deserialize(row[3])
                     if row[3] is not None
                     else row[3]
                 ),
@@ -237,52 +274,10 @@ def _is_serialization_error(self, dbapi_error: DBAPIError) -> bool:
         """Check if the error is a serialization/concurrency error."""
         pass
 
-    @staticmethod
-    def create(
-        database_url: str,
-        engine_kwargs: Dict[str, Any],
-        schema: Optional[str],
-        debug_mode: bool = False,
-    ) -> "ApplicationDatabase":
-        """Factory method to create the appropriate ApplicationDatabase implementation based on URL."""
-        if database_url.startswith("sqlite"):
-            return SQLiteApplicationDatabase(
-                database_url=database_url,
-                engine_kwargs=engine_kwargs,
-                debug_mode=debug_mode,
-            )
-        else:
-            # Default to PostgreSQL for postgresql://, postgres://, or other URLs
-            return PostgresApplicationDatabase(
-                database_url=database_url,
-                engine_kwargs=engine_kwargs,
-                debug_mode=debug_mode,
-                schema=schema,
-            )
-
 
 class PostgresApplicationDatabase(ApplicationDatabase):
     """PostgreSQL-specific implementation of ApplicationDatabase."""
 
-    def __init__(
-        self,
-        *,
-        database_url: str,
-        engine_kwargs: Dict[str, Any],
-        schema: Optional[str],
-        debug_mode: bool = False,
-    ):
-        super().__init__(
-            database_url=database_url,
-            engine_kwargs=engine_kwargs,
-            debug_mode=debug_mode,
-        )
-        if schema is None:
-            self.schema = "dbos"
-        else:
-            self.schema = schema
-        ApplicationSchema.transaction_outputs.schema = schema
-
     def _create_engine(
         self, database_url: str, engine_kwargs: Dict[str, Any]
     ) -> sa.Engine:

dbos/_client.py

@@ -16,7 +16,6 @@
 
 import sqlalchemy as sa
 
-from dbos import _serialization
 from dbos._app_db import ApplicationDatabase
 from dbos._context import MaxPriority, MinPriority
 from dbos._sys_db import SystemDatabase
@@ -27,7 +26,7 @@
 from dbos._dbos_config import get_system_database_url, is_valid_database_url
 from dbos._error import DBOSException, DBOSNonExistentWorkflowError
 from dbos._registrations import DEFAULT_MAX_RECOVERY_ATTEMPTS
-from dbos._serialization import WorkflowInputs
+from dbos._serialization import DefaultSerializer, Serializer, WorkflowInputs
 from dbos._sys_db import (
     EnqueueOptionsInternal,
     StepInfo,
@@ -127,7 +126,9 @@ def __init__(
         system_database_engine: Optional[sa.Engine] = None,
         application_database_url: Optional[str] = None,
         dbos_system_schema: Optional[str] = "dbos",
+        serializer: Serializer = DefaultSerializer(),
     ):
+        self._serializer = serializer
         application_database_url = (
             database_url if database_url else application_database_url
         )
@@ -150,6 +151,7 @@ def __init__(
             },
             engine=system_database_engine,
             schema=dbos_system_schema,
+            serializer=serializer,
         )
         self._sys_db.check_connection()
         if application_database_url:
@@ -161,6 +163,7 @@ def __init__(
                     "pool_size": 2,
                 },
                 schema=dbos_system_schema,
+                serializer=serializer,
             )
 
     def destroy(self) -> None:
@@ -217,7 +220,7 @@ def _enqueue(self, options: EnqueueOptions, *args: Any, **kwargs: Any) -> str:
                 if enqueue_options_internal["priority"] is not None
                 else 0
             ),
-            "inputs": _serialization.serialize_args(inputs),
+            "inputs": self._serializer.serialize(inputs),
         }
 
         self._sys_db.init_workflow(
@@ -282,7 +285,7 @@ def send(
             "workflow_deadline_epoch_ms": None,
             "deduplication_id": None,
             "priority": 0,
-            "inputs": _serialization.serialize_args({"args": (), "kwargs": {}}),
+            "inputs": self._serializer.serialize({"args": (), "kwargs": {}}),
         }
         with self._sys_db.engine.begin() as conn:
             self._sys_db._insert_workflow_status(