dbos 2.1.0a2__tar.gz → 2.2.0__tar.gz

{dbos-2.1.0a2 → dbos-2.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dbos
-Version: 2.1.0a2
+Version: 2.2.0
 Summary: Ultra-lightweight durable execution in Python
 Author-Email: "DBOS, Inc." <contact@dbos.dev>
 License: MIT

{dbos-2.1.0a2 → dbos-2.2.0}/dbos/__init__.py

@@ -12,6 +12,7 @@ from ._dbos_config import DBOSConfig
 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 @@ __all__ = [
     "Queue",
     "Debouncer",
     "DebouncerClient",
+    "Serializer",
 ]

{dbos-2.1.0a2 → dbos-2.2.0}/dbos/_app_db.py

@@ -8,8 +8,8 @@ from sqlalchemy.exc import DBAPIError
 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 @@ class ApplicationDatabase(ABC):
                 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 @@ class ApplicationDatabase(ABC):
         """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-2.1.0a2 → dbos-2.2.0}/dbos/_client.py

@@ -16,7 +16,6 @@ from typing import (
 
 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 @@ if TYPE_CHECKING:
 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,
@@ -63,6 +62,7 @@ class EnqueueOptions(_EnqueueOptionsRequired, total=False):
     deduplication_id: str
     priority: int
     max_recovery_attempts: int
+    queue_partition_key: str
 
 
 def validate_enqueue_options(options: EnqueueOptions) -> None:
@@ -127,7 +127,9 @@ class DBOSClient:
         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 +152,7 @@ class DBOSClient:
             },
             engine=system_database_engine,
             schema=dbos_system_schema,
+            serializer=serializer,
         )
         self._sys_db.check_connection()
         if application_database_url:
@@ -161,6 +164,7 @@ class DBOSClient:
                     "pool_size": 2,
                 },
                 schema=dbos_system_schema,
+                serializer=serializer,
             )
 
     def destroy(self) -> None:
@@ -182,6 +186,7 @@ class DBOSClient:
             "deduplication_id": options.get("deduplication_id"),
             "priority": options.get("priority"),
             "app_version": options.get("app_version"),
+            "queue_partition_key": options.get("queue_partition_key"),
         }
 
         inputs: WorkflowInputs = {
@@ -217,7 +222,8 @@ class DBOSClient:
                 if enqueue_options_internal["priority"] is not None
                 else 0
             ),
-            "inputs": _serialization.serialize_args(inputs),
+            "inputs": self._serializer.serialize(inputs),
+            "queue_partition_key": enqueue_options_internal["queue_partition_key"],
         }
 
         self._sys_db.init_workflow(
@@ -282,7 +288,8 @@ class DBOSClient:
             "workflow_deadline_epoch_ms": None,
             "deduplication_id": None,
             "priority": 0,
-            "inputs": _serialization.serialize_args({"args": (), "kwargs": {}}),
+            "inputs": self._serializer.serialize({"args": (), "kwargs": {}}),
+            "queue_partition_key": None,
         }
         with self._sys_db.engine.begin() as conn:
             self._sys_db._insert_workflow_status(

{dbos-2.1.0a2 → dbos-2.2.0}/dbos/_context.py

@@ -120,6 +120,8 @@ class DBOSContext:
         self.deduplication_id: Optional[str] = None
         # A user-specified priority for the enqueuing workflow.
         self.priority: Optional[int] = None
+        # If the workflow is enqueued on a partitioned queue, its partition key
+        self.queue_partition_key: Optional[str] = None
 
     def create_child(self) -> DBOSContext:
         rv = DBOSContext()
@@ -479,6 +481,7 @@ class SetEnqueueOptions:
         deduplication_id: Optional[str] = None,
         priority: Optional[int] = None,
         app_version: Optional[str] = None,
+        queue_partition_key: Optional[str] = None,
     ) -> None:
         self.created_ctx = False
         self.deduplication_id: Optional[str] = deduplication_id
@@ -491,6 +494,8 @@ class SetEnqueueOptions:
         self.saved_priority: Optional[int] = None
         self.app_version: Optional[str] = app_version
         self.saved_app_version: Optional[str] = None
+        self.queue_partition_key = queue_partition_key
+        self.saved_queue_partition_key: Optional[str] = None
 
     def __enter__(self) -> SetEnqueueOptions:
         # Code to create a basic context
@@ -505,6 +510,8 @@ class SetEnqueueOptions:
         ctx.priority = self.priority
         self.saved_app_version = ctx.app_version
         ctx.app_version = self.app_version
+        self.saved_queue_partition_key = ctx.queue_partition_key
+        ctx.queue_partition_key = self.queue_partition_key
         return self
 
     def __exit__(
@@ -517,6 +524,7 @@ class SetEnqueueOptions:
         curr_ctx.deduplication_id = self.saved_deduplication_id
         curr_ctx.priority = self.saved_priority
         curr_ctx.app_version = self.saved_app_version
+        curr_ctx.queue_partition_key = self.saved_queue_partition_key
         # Code to clean up the basic context if we created it
         if self.created_ctx:
             _clear_local_dbos_context()

{dbos-2.1.0a2 → dbos-2.2.0}/dbos/_core.py

@@ -23,7 +23,6 @@ from typing import (
 from dbos._outcome import Immediate, NoResult, Outcome, Pending
 from dbos._utils import GlobalParams, retriable_postgres_exception
 
-from . import _serialization
 from ._app_db import ApplicationDatabase, TransactionResultInternal
 from ._context import (
     DBOSAssumeRole,
@@ -116,10 +115,10 @@ class WorkflowHandleFuture(Generic[R]):
         try:
             r = self.future.result()
         except Exception as e:
-            serialized_e = _serialization.serialize_exception(e)
+            serialized_e = self.dbos._serializer.serialize(e)
             self.dbos._sys_db.record_get_result(self.workflow_id, None, serialized_e)
             raise
-        serialized_r = _serialization.serialize(r)
+        serialized_r = self.dbos._serializer.serialize(r)
         self.dbos._sys_db.record_get_result(self.workflow_id, serialized_r, None)
         return r
 
@@ -143,10 +142,10 @@ class WorkflowHandlePolling(Generic[R]):
         try:
             r: R = self.dbos._sys_db.await_workflow_result(self.workflow_id)
         except Exception as e:
-            serialized_e = _serialization.serialize_exception(e)
+            serialized_e = self.dbos._serializer.serialize(e)
             self.dbos._sys_db.record_get_result(self.workflow_id, None, serialized_e)
             raise
-        serialized_r = _serialization.serialize(r)
+        serialized_r = self.dbos._serializer.serialize(r)
         self.dbos._sys_db.record_get_result(self.workflow_id, serialized_r, None)
         return r
 
@@ -171,7 +170,7 @@ class WorkflowHandleAsyncTask(Generic[R]):
         try:
             r = await self.task
         except Exception as e:
-            serialized_e = _serialization.serialize_exception(e)
+            serialized_e = self.dbos._serializer.serialize(e)
             await asyncio.to_thread(
                 self.dbos._sys_db.record_get_result,
                 self.workflow_id,
@@ -179,7 +178,7 @@ class WorkflowHandleAsyncTask(Generic[R]):
                 serialized_e,
             )
             raise
-        serialized_r = _serialization.serialize(r)
+        serialized_r = self.dbos._serializer.serialize(r)
         await asyncio.to_thread(
             self.dbos._sys_db.record_get_result, self.workflow_id, serialized_r, None
         )
@@ -207,7 +206,7 @@ class WorkflowHandleAsyncPolling(Generic[R]):
                 self.dbos._sys_db.await_workflow_result, self.workflow_id
             )
         except Exception as e:
-            serialized_e = _serialization.serialize_exception(e)
+            serialized_e = self.dbos._serializer.serialize(e)
             await asyncio.to_thread(
                 self.dbos._sys_db.record_get_result,
                 self.workflow_id,
@@ -215,7 +214,7 @@ class WorkflowHandleAsyncPolling(Generic[R]):
                 serialized_e,
             )
             raise
-        serialized_r = _serialization.serialize(r)
+        serialized_r = self.dbos._serializer.serialize(r)
         await asyncio.to_thread(
             self.dbos._sys_db.record_get_result, self.workflow_id, serialized_r, None
         )
@@ -303,7 +302,12 @@ def _init_workflow(
             if enqueue_options is not None
             else 0
         ),
-        "inputs": _serialization.serialize_args(inputs),
+        "inputs": dbos._serializer.serialize(inputs),
+        "queue_partition_key": (
+            enqueue_options["queue_partition_key"]
+            if enqueue_options is not None
+            else None
+        ),
     }
 
     # Synchronously record the status and inputs for workflows
@@ -319,7 +323,7 @@ def _init_workflow(
                 "function_id": ctx.parent_workflow_fid,
                 "function_name": wf_name,
                 "output": None,
-                "error": _serialization.serialize_exception(e),
+                "error": dbos._serializer.serialize(e),
             }
             dbos._sys_db.record_operation_result(result)
         raise
@@ -378,7 +382,7 @@ def _get_wf_invoke_func(
                 dbos._sys_db.update_workflow_outcome(
                     status["workflow_uuid"],
                     "SUCCESS",
-                    output=_serialization.serialize(output),
+                    output=dbos._serializer.serialize(output),
                 )
             return output
         except DBOSWorkflowConflictIDError:
@@ -392,7 +396,7 @@ def _get_wf_invoke_func(
                 dbos._sys_db.update_workflow_outcome(
                     status["workflow_uuid"],
                     "ERROR",
-                    error=_serialization.serialize_exception(error),
+                    error=dbos._serializer.serialize(error),
                 )
             raise
         finally:
@@ -464,7 +468,7 @@ def execute_workflow_by_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[An
     status = dbos._sys_db.get_workflow_status(workflow_id)
     if not status:
         raise DBOSRecoveryError(workflow_id, "Workflow status not found")
-    inputs = _serialization.deserialize_args(status["inputs"])
+    inputs: WorkflowInputs = dbos._serializer.deserialize(status["inputs"])
     wf_func = dbos._registry.workflow_info_map.get(status["name"], None)
     if not wf_func:
         raise DBOSWorkflowFunctionNotFoundError(
@@ -572,6 +576,9 @@ def start_workflow(
         deduplication_id=local_ctx.deduplication_id if local_ctx is not None else None,
         priority=local_ctx.priority if local_ctx is not None else None,
         app_version=local_ctx.app_version if local_ctx is not None else None,
+        queue_partition_key=(
+            local_ctx.queue_partition_key if local_ctx is not None else None
+        ),
     )
     new_wf_id, new_wf_ctx = _get_new_wf()
 
@@ -665,6 +672,9 @@ async def start_workflow_async(
         deduplication_id=local_ctx.deduplication_id if local_ctx is not None else None,
         priority=local_ctx.priority if local_ctx is not None else None,
         app_version=local_ctx.app_version if local_ctx is not None else None,
+        queue_partition_key=(
+            local_ctx.queue_partition_key if local_ctx is not None else None
+        ),
     )
     new_wf_id, new_wf_ctx = _get_new_wf()
 
@@ -837,11 +847,11 @@ def workflow_wrapper(
             try:
                 r = func()
             except Exception as e:
-                serialized_e = _serialization.serialize_exception(e)
+                serialized_e = dbos._serializer.serialize(e)
                 assert workflow_id is not None
                 dbos._sys_db.record_get_result(workflow_id, None, serialized_e)
                 raise
-            serialized_r = _serialization.serialize(r)
+            serialized_r = dbos._serializer.serialize(r)
             assert workflow_id is not None
             dbos._sys_db.record_get_result(workflow_id, serialized_r, None)
             return r
@@ -948,15 +958,15 @@ def decorate_transaction(
                                         f"Replaying transaction, id: {ctx.function_id}, name: {attributes['name']}"
                                     )
                                     if recorded_output["error"]:
-                                        deserialized_error = (
-                                            _serialization.deserialize_exception(
+                                        deserialized_error: Exception = (
+                                            dbos._serializer.deserialize(
                                                 recorded_output["error"]
                                             )
                                         )
                                         has_recorded_error = True
                                         raise deserialized_error
                                     elif recorded_output["output"]:
-                                        return _serialization.deserialize(
+                                        return dbos._serializer.deserialize(
                                             recorded_output["output"]
                                         )
                                     else:
@@ -969,7 +979,9 @@ def decorate_transaction(
                                     )
 
                                 output = func(*args, **kwargs)
-                                txn_output["output"] = _serialization.serialize(output)
+                                txn_output["output"] = dbos._serializer.serialize(
+                                    output
+                                )
                                 assert (
                                     ctx.sql_session is not None
                                 ), "Cannot find a database connection"
@@ -1010,8 +1022,8 @@ def decorate_transaction(
                         finally:
                             # Don't record the error if it was already recorded
                             if txn_error and not has_recorded_error:
-                                txn_output["error"] = (
-                                    _serialization.serialize_exception(txn_error)
+                                txn_output["error"] = dbos._serializer.serialize(
+                                    txn_error
                                 )
                                 dbos._app_db.record_transaction_error(txn_output)
             return output
@@ -1128,10 +1140,10 @@ def decorate_step(
                 try:
                     output = func()
                 except Exception as error:
-                    step_output["error"] = _serialization.serialize_exception(error)
+                    step_output["error"] = dbos._serializer.serialize(error)
                     dbos._sys_db.record_operation_result(step_output)
                     raise
-                step_output["output"] = _serialization.serialize(output)
+                step_output["output"] = dbos._serializer.serialize(output)
                 dbos._sys_db.record_operation_result(step_output)
                 return output
 
@@ -1147,13 +1159,13 @@ def decorate_step(
                         f"Replaying step, id: {ctx.function_id}, name: {attributes['name']}"
                     )
                     if recorded_output["error"] is not None:
-                        deserialized_error = _serialization.deserialize_exception(
+                        deserialized_error: Exception = dbos._serializer.deserialize(
                             recorded_output["error"]
                         )
                         raise deserialized_error
                     elif recorded_output["output"] is not None:
                         return cast(
-                            R, _serialization.deserialize(recorded_output["output"])
+                            R, dbos._serializer.deserialize(recorded_output["output"])
                         )
                     else:
                         raise Exception("Output and error are both None")

{dbos-2.1.0a2 → dbos-2.2.0}/dbos/_dbos.py

@@ -31,6 +31,7 @@ from typing import (
 
 from dbos._conductor.conductor import ConductorWebsocket
 from dbos._debouncer import debouncer_workflow
+from dbos._serialization import DefaultSerializer, Serializer
 from dbos._sys_db import SystemDatabase, WorkflowStatus
 from dbos._utils import INTERNAL_QUEUE_NAME, GlobalParams
 from dbos._workflow_commands import fork_workflow, list_queued_workflows, list_workflows
@@ -341,6 +342,8 @@ class DBOS:
         self.conductor_websocket: Optional[ConductorWebsocket] = None
         self._background_event_loop: BackgroundEventLoop = BackgroundEventLoop()
         self._active_workflows_set: set[str] = set()
+        serializer = config.get("serializer")
+        self._serializer: Serializer = serializer if serializer else DefaultSerializer()
 
         # Globally set the application version and executor ID.
         # In DBOS Cloud, instead use the values supplied through environment variables.
@@ -449,28 +452,34 @@ class DBOS:
             assert self._config["database"]["sys_db_engine_kwargs"] is not None
             # Get the schema configuration, use "dbos" as default
             schema = self._config.get("dbos_system_schema", "dbos")
+            dbos_logger.debug("Creating system database")
             self._sys_db_field = SystemDatabase.create(
                 system_database_url=get_system_database_url(self._config),
                 engine_kwargs=self._config["database"]["sys_db_engine_kwargs"],
                 engine=self._config["system_database_engine"],
                 debug_mode=debug_mode,
                 schema=schema,
+                serializer=self._serializer,
             )
             assert self._config["database"]["db_engine_kwargs"] is not None
             if self._config["database_url"]:
+                dbos_logger.debug("Creating application database")
                 self._app_db_field = ApplicationDatabase.create(
                     database_url=self._config["database_url"],
                     engine_kwargs=self._config["database"]["db_engine_kwargs"],
                     debug_mode=debug_mode,
                     schema=schema,
+                    serializer=self._serializer,
                 )
 
             if debug_mode:
                 return
 
             # Run migrations for the system and application databases
+            dbos_logger.debug("Running system database migrations")
             self._sys_db.run_migrations()
             if self._app_db:
+                dbos_logger.debug("Running application database migrations")
                 self._app_db.run_migrations()
 
             admin_port = self._config.get("runtimeConfig", {}).get("admin_port")
@@ -481,10 +490,12 @@ class DBOS:
             )
             if run_admin_server:
                 try:
+                    dbos_logger.debug("Starting admin server")
                     self._admin_server_field = AdminServer(dbos=self, port=admin_port)
                 except Exception as e:
                     dbos_logger.warning(f"Failed to start admin server: {e}")
 
+            dbos_logger.debug("Retrieving local pending workflows for recovery")
             workflow_ids = self._sys_db.get_pending_workflows(
                 GlobalParams.executor_id, GlobalParams.app_version
             )
@@ -500,6 +511,7 @@ class DBOS:
             self._executor.submit(startup_recovery_thread, self, workflow_ids)
 
             # Listen to notifications
+            dbos_logger.debug("Starting notifications listener thread")
             notification_listener_thread = threading.Thread(
                 target=self._sys_db._notification_listener,
                 daemon=True,
@@ -511,6 +523,7 @@ class DBOS:
             self._registry.get_internal_queue()
 
             # Start the queue thread
+            dbos_logger.debug("Starting queue thread")
             evt = threading.Event()
             self.background_thread_stop_events.append(evt)
             bg_queue_thread = threading.Thread(
@@ -526,6 +539,7 @@ class DBOS:
                     self.conductor_url = f"wss://{dbos_domain}/conductor/v1alpha1"
                 evt = threading.Event()
                 self.background_thread_stop_events.append(evt)
+                dbos_logger.debug("Starting Conductor thread")
                 self.conductor_websocket = ConductorWebsocket(
                     self,
                     conductor_url=self.conductor_url,
@@ -536,6 +550,7 @@ class DBOS:
                 self._background_threads.append(self.conductor_websocket)
 
             # Grab any pollers that were deferred and start them
+            dbos_logger.debug("Starting event receivers")
             for evt, func, args, kwargs in self._registry.pollers:
                 self.poller_stop_events.append(evt)
                 poller_thread = threading.Thread(

{dbos-2.1.0a2 → dbos-2.2.0}/dbos/_dbos_config.py

@@ -7,6 +7,8 @@ import sqlalchemy as sa
 import yaml
 from sqlalchemy import make_url
 
+from dbos._serialization import Serializer
+
 from ._error import DBOSInitializationError
 from ._logger import dbos_logger
 from ._schemas.system_database import SystemSchema
@@ -37,6 +39,7 @@ class DBOSConfig(TypedDict, total=False):
         enable_otlp (bool): If True, enable built-in DBOS OTLP tracing and logging.
         system_database_engine (sa.Engine): A custom system database engine. If provided, DBOS will not create an engine but use this instead.
         conductor_key (str): An API key for DBOS Conductor. Pass this in to connect your process to Conductor.
+        serializer (Serializer): A custom serializer and deserializer DBOS uses when storing program data in the system database
     """
 
     name: str
@@ -57,6 +60,7 @@ class DBOSConfig(TypedDict, total=False):
     enable_otlp: Optional[bool]
     system_database_engine: Optional[sa.Engine]
     conductor_key: Optional[str]
+    serializer: Optional[Serializer]
 
 
 class RuntimeConfig(TypedDict, total=False):
@@ -67,16 +71,6 @@ class RuntimeConfig(TypedDict, total=False):
 
 
 class DatabaseConfig(TypedDict, total=False):
-    """
-    Internal data structure containing the DBOS database configuration.
-    Attributes:
-        sys_db_name (str): System database name
-        sys_db_pool_size (int): System database pool size
-        db_engine_kwargs (Dict[str, Any]): SQLAlchemy engine kwargs
-        migrate (List[str]): Migration commands to run on startup
-        dbos_system_schema (str): Schema name for DBOS system tables. Defaults to "dbos".
-    """
-
     sys_db_pool_size: Optional[int]
     db_engine_kwargs: Optional[Dict[str, Any]]
     sys_db_engine_kwargs: Optional[Dict[str, Any]]

{dbos-2.1.0a2 → dbos-2.2.0}/dbos/_migration.py

@@ -203,8 +203,14 @@ CREATE TABLE \"{schema}\".event_dispatch_kv (
 """
 
 
+def get_dbos_migration_two(schema: str) -> str:
+    return f"""
+ALTER TABLE \"{schema}\".workflow_status ADD COLUMN queue_partition_key TEXT;
+"""
+
+
 def get_dbos_migrations(schema: str) -> list[str]:
-    return [get_dbos_migration_one(schema)]
+    return [get_dbos_migration_one(schema), get_dbos_migration_two(schema)]
 
 
 def get_sqlite_timestamp_expr() -> str:
@@ -293,4 +299,8 @@ CREATE TABLE streams (
 );
 """
 
-sqlite_migrations = [sqlite_migration_one]
+sqlite_migration_two = """
+ALTER TABLE workflow_status ADD COLUMN queue_partition_key TEXT;
+"""
+
+sqlite_migrations = [sqlite_migration_one, sqlite_migration_two]