dbos 2.1.0a2__py3-none-any.whl → 2.4.0a7__py3-none-any.whl

dbos/_sys_db.py

@@ -30,7 +30,6 @@ from dbos._utils import (
     retriable_sqlite_exception,
 )
 
-from . import _serialization
 from ._context import get_local_dbos_context
 from ._error import (
     DBOSAwaitedWorkflowCancelledError,
@@ -44,6 +43,7 @@ from ._error import (
 )
 from ._logger import dbos_logger
 from ._schemas.system_database import SystemSchema
+from ._serialization import Serializer, WorkflowInputs, safe_deserialize
 
 if TYPE_CHECKING:
     from ._queue import Queue
@@ -95,7 +95,7 @@ class WorkflowStatus:
     # All roles which the authenticated user could assume
     authenticated_roles: Optional[list[str]]
     # The deserialized workflow input object
-    input: Optional[_serialization.WorkflowInputs]
+    input: Optional[WorkflowInputs]
     # The workflow's output, if any
     output: Optional[Any] = None
     # The error the workflow threw, if any
@@ -106,7 +106,7 @@ class WorkflowStatus:
     updated_at: Optional[int]
     # If this workflow was enqueued, on which queue
     queue_name: Optional[str]
-    # The executor to most recently executed this workflow
+    # The executor to most recently execute this workflow
     executor_id: Optional[str]
     # The application version on which this workflow was started
     app_version: Optional[str]
@@ -114,6 +114,14 @@ class WorkflowStatus:
     workflow_timeout_ms: Optional[int]
     # The deadline of a workflow, computed by adding its timeout to its start time.
     workflow_deadline_epoch_ms: Optional[int]
+    # Unique ID for deduplication on a queue
+    deduplication_id: Optional[str]
+    # Priority of the workflow on the queue, starting from 1 ~ 2,147,483,647. Default 0 (highest priority).
+    priority: Optional[int]
+    # If this workflow is enqueued on a partitioned queue, its partition key
+    queue_partition_key: Optional[str]
+    # If this workflow was forked from another, that workflow's ID.
+    forked_from: Optional[str]
 
     # INTERNAL FIELDS
 
@@ -141,17 +149,13 @@ class WorkflowStatusInternal(TypedDict):
     app_version: Optional[str]
     app_id: Optional[str]
     recovery_attempts: Optional[int]
-    # The start-to-close timeout of the workflow in ms
     workflow_timeout_ms: Optional[int]
-    # The deadline of a workflow, computed by adding its timeout to its start time.
-    # Deadlines propagate to children. When the deadline is reached, the workflow is cancelled.
     workflow_deadline_epoch_ms: Optional[int]
-    # Unique ID for deduplication on a queue
     deduplication_id: Optional[str]
-    # Priority of the workflow on the queue, starting from 1 ~ 2,147,483,647. Default 0 (highest priority).
     priority: int
-    # Serialized workflow inputs
     inputs: str
+    queue_partition_key: Optional[str]
+    forked_from: Optional[str]
 
 
 class EnqueueOptionsInternal(TypedDict):
@@ -161,6 +165,8 @@ class EnqueueOptionsInternal(TypedDict):
     priority: Optional[int]
     # On what version the workflow is enqueued. Current version if not specified.
     app_version: Optional[str]
+    # If the workflow is enqueued on a partitioned queue, its partition key
+    queue_partition_key: Optional[str]
 
 
 class RecordedResult(TypedDict):
@@ -174,6 +180,7 @@ class OperationResultInternal(TypedDict):
     function_name: str
     output: Optional[str]  # JSON (jsonpickle)
     error: Optional[str]  # JSON (jsonpickle)
+    started_at_epoch_ms: int
 
 
 class GetEventWorkflowContext(TypedDict):
@@ -190,42 +197,34 @@ class GetWorkflowsInput:
     """
 
     def __init__(self) -> None:
-        self.workflow_ids: Optional[List[str]] = (
-            None  # Search only in these workflow IDs
-        )
-        self.name: Optional[str] = None  # The name of the workflow function
-        self.authenticated_user: Optional[str] = None  # The user who ran the workflow.
-        self.start_time: Optional[str] = None  # Timestamp in ISO 8601 format
-        self.end_time: Optional[str] = None  # Timestamp in ISO 8601 format
-        self.status: Optional[List[str]] = (
-            None  # Get workflows with one of these statuses
-        )
-        self.application_version: Optional[str] = (
-            None  # The application version that ran this workflow. = None
-        )
-        self.limit: Optional[int] = (
-            None  # Return up to this many workflows IDs. IDs are ordered by workflow creation time.
-        )
-        self.offset: Optional[int] = (
-            None  # Offset into the matching records for pagination
-        )
-        self.sort_desc: bool = (
-            False  # If true, sort by created_at in DESC order. Default false (in ASC order).
-        )
-        self.workflow_id_prefix: Optional[str] = (
-            None  # If set, search for workflow IDs starting with this string
-        )
-
-
-class GetQueuedWorkflowsInput(TypedDict):
-    queue_name: Optional[str]  # Get workflows belonging to this queue
-    status: Optional[list[str]]  # Get workflows with one of these statuses
-    start_time: Optional[str]  # Timestamp in ISO 8601 format
-    end_time: Optional[str]  # Timestamp in ISO 8601 format
-    limit: Optional[int]  # Return up to this many workflows IDs.
-    offset: Optional[int]  # Offset into the matching records for pagination
-    name: Optional[str]  # The name of the workflow function
-    sort_desc: Optional[bool]  # Sort by created_at in DESC or ASC order
+        # Search only in these workflow IDs
+        self.workflow_ids: Optional[List[str]] = None
+        # The name of the workflow function
+        self.name: Optional[str] = None
+        # The user who ran the workflow.
+        self.authenticated_user: Optional[str] = None
+        # Timestamp in ISO 8601 format
+        self.start_time: Optional[str] = None
+        # Timestamp in ISO 8601 format
+        self.end_time: Optional[str] = None
+        # Get workflows with one of these statuses
+        self.status: Optional[List[str]] = None
+        # The application version that ran this workflow.
+        self.application_version: Optional[str] = None
+        # Get workflows forked from this workflow ID.
+        self.forked_from: Optional[str] = None
+        # Return up to this many workflows IDs. IDs are ordered by workflow creation time.
+        self.limit: Optional[int] = None
+        # Offset into the matching records for pagination
+        self.offset: Optional[int] = None
+        # If true, sort by created_at in DESC order. Default false (in ASC order).
+        self.sort_desc: bool = False
+        # Search only for workflow IDs starting with this string
+        self.workflow_id_prefix: Optional[str] = None
+        # Search only for workflows enqueued on this queue
+        self.queue_name: Optional[str] = None
+        # Search only currently enqueued workflows
+        self.queues_only: bool = False
 
 
 class GetPendingWorkflowsOutput:
@@ -245,6 +244,10 @@ class StepInfo(TypedDict):
     error: Optional[Exception]
     # If the step starts or retrieves the result of a workflow, its ID
     child_workflow_id: Optional[str]
+    # The Unix epoch timestamp at which this step started
+    started_at_epoch_ms: Optional[int]
+    # The Unix epoch timestamp at which this step completed
+    completed_at_epoch_ms: Optional[int]
 
 
 _dbos_null_topic = "__null__topic__"
@@ -341,6 +344,42 @@ def db_retry(
 
 class SystemDatabase(ABC):
 
+    @staticmethod
+    def create(
+        system_database_url: str,
+        engine_kwargs: Dict[str, Any],
+        engine: Optional[sa.Engine],
+        schema: Optional[str],
+        serializer: Serializer,
+        executor_id: Optional[str],
+        debug_mode: bool = False,
+    ) -> "SystemDatabase":
+        """Factory method to create the appropriate SystemDatabase implementation based on URL."""
+        if system_database_url.startswith("sqlite"):
+            from ._sys_db_sqlite import SQLiteSystemDatabase
+
+            return SQLiteSystemDatabase(
+                system_database_url=system_database_url,
+                engine_kwargs=engine_kwargs,
+                engine=engine,
+                schema=schema,
+                serializer=serializer,
+                executor_id=executor_id,
+                debug_mode=debug_mode,
+            )
+        else:
+            from ._sys_db_postgres import PostgresSystemDatabase
+
+            return PostgresSystemDatabase(
+                system_database_url=system_database_url,
+                engine_kwargs=engine_kwargs,
+                engine=engine,
+                schema=schema,
+                serializer=serializer,
+                executor_id=executor_id,
+                debug_mode=debug_mode,
+            )
+
     def __init__(
         self,
         *,
@@ -348,6 +387,8 @@ class SystemDatabase(ABC):
         engine_kwargs: Dict[str, Any],
         engine: Optional[sa.Engine],
         schema: Optional[str],
+        serializer: Serializer,
+        executor_id: Optional[str],
         debug_mode: bool = False,
     ):
         import sqlalchemy.dialects.postgresql as pg
@@ -355,6 +396,8 @@ class SystemDatabase(ABC):
 
         self.dialect = sq if system_database_url.startswith("sqlite") else pg
 
+        self.serializer = serializer
+
         if system_database_url.startswith("sqlite"):
             self.schema = None
         else:
@@ -371,6 +414,8 @@ class SystemDatabase(ABC):
 
         self.notifications_map = ThreadSafeConditionDict()
         self.workflow_events_map = ThreadSafeConditionDict()
+        self.executor_id = executor_id
+
         self._listener_thread_lock = threading.Lock()
 
         # Now we can run background processes
@@ -454,6 +499,7 @@ class SystemDatabase(ABC):
                 deduplication_id=status["deduplication_id"],
                 priority=status["priority"],
                 inputs=status["inputs"],
+                queue_partition_key=status["queue_partition_key"],
             )
             .on_conflict_do_update(
                 index_elements=["workflow_uuid"],
@@ -665,6 +711,7 @@ class SystemDatabase(ABC):
                     assumed_role=status["assumed_role"],
                     queue_name=INTERNAL_QUEUE_NAME,
                     inputs=status["inputs"],
+                    forked_from=original_workflow_id,
                 )
             )
 
@@ -725,6 +772,8 @@ class SystemDatabase(ABC):
                     SystemSchema.workflow_status.c.deduplication_id,
                     SystemSchema.workflow_status.c.priority,
                     SystemSchema.workflow_status.c.inputs,
+                    SystemSchema.workflow_status.c.queue_partition_key,
+                    SystemSchema.workflow_status.c.forked_from,
                 ).where(SystemSchema.workflow_status.c.workflow_uuid == workflow_uuid)
             ).fetchone()
             if row is None:
@@ -752,6 +801,8 @@ class SystemDatabase(ABC):
                 "deduplication_id": row[16],
                 "priority": row[17],
                 "inputs": row[18],
+                "queue_partition_key": row[19],
+                "forked_from": row[20],
             }
             return status
 
@@ -797,10 +848,11 @@ class SystemDatabase(ABC):
                     status = row[0]
                     if status == WorkflowStatusString.SUCCESS.value:
                         output = row[1]
-                        return _serialization.deserialize(output)
+                        return self.serializer.deserialize(output)
                     elif status == WorkflowStatusString.ERROR.value:
                         error = row[2]
-                        raise _serialization.deserialize_exception(error)
+                        e: Exception = self.serializer.deserialize(error)
+                        raise e
                     elif status == WorkflowStatusString.CANCELLED.value:
                         # Raise AwaitedWorkflowCancelledError here, not the cancellation exception
                         # because the awaiting workflow is not being cancelled.
@@ -837,6 +889,10 @@ class SystemDatabase(ABC):
             SystemSchema.workflow_status.c.application_id,
             SystemSchema.workflow_status.c.workflow_deadline_epoch_ms,
             SystemSchema.workflow_status.c.workflow_timeout_ms,
+            SystemSchema.workflow_status.c.deduplication_id,
+            SystemSchema.workflow_status.c.priority,
+            SystemSchema.workflow_status.c.queue_partition_key,
+            SystemSchema.workflow_status.c.forked_from,
         ]
         if load_input:
             load_columns.append(SystemSchema.workflow_status.c.inputs)
@@ -844,7 +900,15 @@ class SystemDatabase(ABC):
             load_columns.append(SystemSchema.workflow_status.c.output)
             load_columns.append(SystemSchema.workflow_status.c.error)
 
-        query = sa.select(*load_columns)
+        if input.queues_only:
+            query = sa.select(*load_columns).where(
+                sa.and_(
+                    SystemSchema.workflow_status.c.queue_name.isnot(None),
+                    SystemSchema.workflow_status.c.status.in_(["ENQUEUED", "PENDING"]),
+                )
+            )
+        else:
+            query = sa.select(*load_columns)
         if input.sort_desc:
             query = query.order_by(SystemSchema.workflow_status.c.created_at.desc())
         else:
@@ -873,6 +937,10 @@ class SystemDatabase(ABC):
                 SystemSchema.workflow_status.c.application_version
                 == input.application_version
             )
+        if input.forked_from:
+            query = query.where(
+                SystemSchema.workflow_status.c.forked_from == input.forked_from
+            )
         if input.workflow_ids:
             query = query.where(
                 SystemSchema.workflow_status.c.workflow_uuid.in_(input.workflow_ids)
@@ -883,6 +951,10 @@ class SystemDatabase(ABC):
                     input.workflow_id_prefix
                 )
             )
+        if input.queue_name:
+            query = query.where(
+                SystemSchema.workflow_status.c.queue_name == input.queue_name
+            )
         if input.limit:
             query = query.limit(input.limit)
         if input.offset:
@@ -892,6 +964,7 @@ class SystemDatabase(ABC):
             rows = c.execute(query).fetchall()
 
         infos: List[WorkflowStatus] = []
+        workflow_ids: List[str] = []
         for row in rows:
             info = WorkflowStatus()
             info.workflow_id = row[0]
@@ -913,11 +986,16 @@ class SystemDatabase(ABC):
             info.app_id = row[14]
             info.workflow_deadline_epoch_ms = row[15]
             info.workflow_timeout_ms = row[16]
-
-            raw_input = row[17] if load_input else None
-            raw_output = row[18] if load_output else None
-            raw_error = row[19] if load_output else None
-            inputs, output, exception = _serialization.safe_deserialize(
+            info.deduplication_id = row[17]
+            info.priority = row[18]
+            info.queue_partition_key = row[19]
+            info.forked_from = row[20]
+
+            raw_input = row[21] if load_input else None
+            raw_output = row[22] if load_output else None
+            raw_error = row[23] if load_output else None
+            inputs, output, exception = safe_deserialize(
+                self.serializer,
                 info.workflow_id,
                 serialized_input=raw_input,
                 serialized_output=raw_output,
@@ -927,121 +1005,10 @@ class SystemDatabase(ABC):
             info.output = output
             info.error = exception
 
+            workflow_ids.append(info.workflow_id)
             infos.append(info)
         return infos
 
-    def get_queued_workflows(
-        self,
-        input: GetQueuedWorkflowsInput,
-        *,
-        load_input: bool = True,
-    ) -> List[WorkflowStatus]:
-        """
-        Retrieve a list of queued workflows result and inputs based on the input criteria. The result is a list of external-facing workflow status objects.
-        """
-        load_columns = [
-            SystemSchema.workflow_status.c.workflow_uuid,
-            SystemSchema.workflow_status.c.status,
-            SystemSchema.workflow_status.c.name,
-            SystemSchema.workflow_status.c.recovery_attempts,
-            SystemSchema.workflow_status.c.config_name,
-            SystemSchema.workflow_status.c.class_name,
-            SystemSchema.workflow_status.c.authenticated_user,
-            SystemSchema.workflow_status.c.authenticated_roles,
-            SystemSchema.workflow_status.c.assumed_role,
-            SystemSchema.workflow_status.c.queue_name,
-            SystemSchema.workflow_status.c.executor_id,
-            SystemSchema.workflow_status.c.created_at,
-            SystemSchema.workflow_status.c.updated_at,
-            SystemSchema.workflow_status.c.application_version,
-            SystemSchema.workflow_status.c.application_id,
-            SystemSchema.workflow_status.c.workflow_deadline_epoch_ms,
-            SystemSchema.workflow_status.c.workflow_timeout_ms,
-        ]
-        if load_input:
-            load_columns.append(SystemSchema.workflow_status.c.inputs)
-
-        query = sa.select(*load_columns).where(
-            sa.and_(
-                SystemSchema.workflow_status.c.queue_name.isnot(None),
-                SystemSchema.workflow_status.c.status.in_(["ENQUEUED", "PENDING"]),
-            )
-        )
-        if input["sort_desc"]:
-            query = query.order_by(SystemSchema.workflow_status.c.created_at.desc())
-        else:
-            query = query.order_by(SystemSchema.workflow_status.c.created_at.asc())
-
-        if input.get("name"):
-            query = query.where(SystemSchema.workflow_status.c.name == input["name"])
-
-        if input.get("queue_name"):
-            query = query.where(
-                SystemSchema.workflow_status.c.queue_name == input["queue_name"]
-            )
-
-        status = input.get("status", None)
-        if status:
-            query = query.where(SystemSchema.workflow_status.c.status.in_(status))
-        if "start_time" in input and input["start_time"] is not None:
-            query = query.where(
-                SystemSchema.workflow_status.c.created_at
-                >= datetime.datetime.fromisoformat(input["start_time"]).timestamp()
-                * 1000
-            )
-        if "end_time" in input and input["end_time"] is not None:
-            query = query.where(
-                SystemSchema.workflow_status.c.created_at
-                <= datetime.datetime.fromisoformat(input["end_time"]).timestamp() * 1000
-            )
-        if input.get("limit"):
-            query = query.limit(input["limit"])
-        if input.get("offset"):
-            query = query.offset(input["offset"])
-
-        with self.engine.begin() as c:
-            rows = c.execute(query).fetchall()
-
-        infos: List[WorkflowStatus] = []
-        for row in rows:
-            info = WorkflowStatus()
-            info.workflow_id = row[0]
-            info.status = row[1]
-            info.name = row[2]
-            info.recovery_attempts = row[3]
-            info.config_name = row[4]
-            info.class_name = row[5]
-            info.authenticated_user = row[6]
-            info.authenticated_roles = (
-                json.loads(row[7]) if row[7] is not None else None
-            )
-            info.assumed_role = row[8]
-            info.queue_name = row[9]
-            info.executor_id = row[10]
-            info.created_at = row[11]
-            info.updated_at = row[12]
-            info.app_version = row[13]
-            info.app_id = row[14]
-            info.workflow_deadline_epoch_ms = row[15]
-            info.workflow_timeout_ms = row[16]
-
-            raw_input = row[17] if load_input else None
-
-            # Error and Output are not loaded because they should always be None for queued workflows.
-            inputs, output, exception = _serialization.safe_deserialize(
-                info.workflow_id,
-                serialized_input=raw_input,
-                serialized_output=None,
-                serialized_exception=None,
-            )
-            info.input = inputs
-            info.output = output
-            info.error = exception
-
-            infos.append(info)
-
-        return infos
-
     def get_pending_workflows(
         self, executor_id: str, app_version: str
     ) -> list[GetPendingWorkflowsOutput]:
@@ -1075,11 +1042,14 @@ class SystemDatabase(ABC):
                     SystemSchema.operation_outputs.c.output,
                     SystemSchema.operation_outputs.c.error,
                     SystemSchema.operation_outputs.c.child_workflow_id,
+                    SystemSchema.operation_outputs.c.started_at_epoch_ms,
+                    SystemSchema.operation_outputs.c.completed_at_epoch_ms,
                 ).where(SystemSchema.operation_outputs.c.workflow_uuid == workflow_id)
             ).fetchall()
             steps = []
             for row in rows:
-                _, output, exception = _serialization.safe_deserialize(
+                _, output, exception = safe_deserialize(
+                    self.serializer,
                     workflow_id,
                     serialized_input=None,
                     serialized_output=row[2],
@@ -1091,6 +1061,8 @@ class SystemDatabase(ABC):
                     output=output,
                     error=exception,
                     child_workflow_id=row[4],
+                    started_at_epoch_ms=row[5],
+                    completed_at_epoch_ms=row[6],
                 )
                 steps.append(step)
             return steps
@@ -1103,10 +1075,33 @@ class SystemDatabase(ABC):
         error = result["error"]
         output = result["output"]
         assert error is None or output is None, "Only one of error or output can be set"
+        wf_executor_id_row = conn.execute(
+            sa.select(
+                SystemSchema.workflow_status.c.executor_id,
+            ).where(
+                SystemSchema.workflow_status.c.workflow_uuid == result["workflow_uuid"]
+            )
+        ).fetchone()
+        assert wf_executor_id_row is not None
+        wf_executor_id = wf_executor_id_row[0]
+        if self.executor_id is not None and wf_executor_id != self.executor_id:
+            dbos_logger.debug(
+                f'Resetting executor_id from {wf_executor_id} to {self.executor_id} for workflow {result["workflow_uuid"]}'
+            )
+            conn.execute(
+                sa.update(SystemSchema.workflow_status)
+                .values(executor_id=self.executor_id)
+                .where(
+                    SystemSchema.workflow_status.c.workflow_uuid
+                    == result["workflow_uuid"]
+                )
+            )
         sql = sa.insert(SystemSchema.operation_outputs).values(
             workflow_uuid=result["workflow_uuid"],
             function_id=result["function_id"],
             function_name=result["function_name"],
+            started_at_epoch_ms=result["started_at_epoch_ms"],
+            completed_at_epoch_ms=int(time.time() * 1000),
             output=output,
             error=error,
         )
@@ -1278,7 +1273,8 @@ class SystemDatabase(ABC):
         if row is None:
             return None
         elif row[1]:
-            raise _serialization.deserialize_exception(row[1])
+            e: Exception = self.serializer.deserialize(row[1])
+            raise e
         else:
             return str(row[0])
 
@@ -1292,6 +1288,7 @@ class SystemDatabase(ABC):
         topic: Optional[str] = None,
     ) -> None:
         function_name = "DBOS.send"
+        start_time = int(time.time() * 1000)
         topic = topic if topic is not None else _dbos_null_topic
         with self.engine.begin() as c:
             recorded_output = self._check_operation_execution_txn(
@@ -1317,7 +1314,7 @@ class SystemDatabase(ABC):
                     sa.insert(SystemSchema.notifications).values(
                         destination_uuid=destination_uuid,
                         topic=topic,
-                        message=_serialization.serialize(message),
+                        message=self.serializer.serialize(message),
                     )
                 )
             except DBAPIError as dbapi_error:
@@ -1328,6 +1325,7 @@ class SystemDatabase(ABC):
                 "workflow_uuid": workflow_uuid,
                 "function_id": function_id,
                 "function_name": function_name,
+                "started_at_epoch_ms": start_time,
                 "output": None,
                 "error": None,
             }
@@ -1343,6 +1341,7 @@ class SystemDatabase(ABC):
         timeout_seconds: float = 60,
     ) -> Any:
         function_name = "DBOS.recv"
+        start_time = int(time.time() * 1000)
         topic = topic if topic is not None else _dbos_null_topic
 
         # First, check for previous executions.
@@ -1354,7 +1353,7 @@ class SystemDatabase(ABC):
         if recorded_output is not None:
             dbos_logger.debug(f"Replaying recv, id: {function_id}, topic: {topic}")
             if recorded_output["output"] is not None:
-                return _serialization.deserialize(recorded_output["output"])
+                return self.serializer.deserialize(recorded_output["output"])
             else:
                 raise Exception("No output recorded in the last recv")
         else:
@@ -1421,13 +1420,14 @@ class SystemDatabase(ABC):
             rows = c.execute(delete_stmt).fetchall()
             message: Any = None
             if len(rows) > 0:
-                message = _serialization.deserialize(rows[0][0])
+                message = self.serializer.deserialize(rows[0][0])
             self._record_operation_result_txn(
                 {
                     "workflow_uuid": workflow_uuid,
                     "function_id": function_id,
                     "function_name": function_name,
-                    "output": _serialization.serialize(
+                    "started_at_epoch_ms": start_time,
+                    "output": self.serializer.serialize(
                         message
                     ),  # None will be serialized to 'null'
                     "error": None,
@@ -1453,36 +1453,6 @@ class SystemDatabase(ABC):
 
             PostgresSystemDatabase._reset_system_database(database_url)
 
-    @staticmethod
-    def create(
-        system_database_url: str,
-        engine_kwargs: Dict[str, Any],
-        engine: Optional[sa.Engine],
-        schema: Optional[str],
-        debug_mode: bool = False,
-    ) -> "SystemDatabase":
-        """Factory method to create the appropriate SystemDatabase implementation based on URL."""
-        if system_database_url.startswith("sqlite"):
-            from ._sys_db_sqlite import SQLiteSystemDatabase
-
-            return SQLiteSystemDatabase(
-                system_database_url=system_database_url,
-                engine_kwargs=engine_kwargs,
-                engine=engine,
-                schema=schema,
-                debug_mode=debug_mode,
-            )
-        else:
-            from ._sys_db_postgres import PostgresSystemDatabase
-
-            return PostgresSystemDatabase(
-                system_database_url=system_database_url,
-                engine_kwargs=engine_kwargs,
-                engine=engine,
-                schema=schema,
-                debug_mode=debug_mode,
-            )
-
     @db_retry()
     def sleep(
         self,
@@ -1492,6 +1462,7 @@ class SystemDatabase(ABC):
         skip_sleep: bool = False,
     ) -> float:
         function_name = "DBOS.sleep"
+        start_time = int(time.time() * 1000)
         recorded_output = self.check_operation_execution(
             workflow_uuid, function_id, function_name
         )
@@ -1502,7 +1473,7 @@ class SystemDatabase(ABC):
         if recorded_output is not None:
             dbos_logger.debug(f"Replaying sleep, id: {function_id}, seconds: {seconds}")
             assert recorded_output["output"] is not None, "no recorded end time"
-            end_time = _serialization.deserialize(recorded_output["output"])
+            end_time = self.serializer.deserialize(recorded_output["output"])
         else:
             dbos_logger.debug(f"Running sleep, id: {function_id}, seconds: {seconds}")
             end_time = time.time() + seconds
@@ -1512,7 +1483,8 @@ class SystemDatabase(ABC):
                         "workflow_uuid": workflow_uuid,
                         "function_id": function_id,
                         "function_name": function_name,
-                        "output": _serialization.serialize(end_time),
+                        "started_at_epoch_ms": start_time,
+                        "output": self.serializer.serialize(end_time),
                         "error": None,
                     }
                 )
@@ -1532,6 +1504,7 @@ class SystemDatabase(ABC):
         message: Any,
     ) -> None:
         function_name = "DBOS.setEvent"
+        start_time = int(time.time() * 1000)
         with self.engine.begin() as c:
             recorded_output = self._check_operation_execution_txn(
                 workflow_uuid, function_id, function_name, conn=c
@@ -1550,17 +1523,18 @@ class SystemDatabase(ABC):
                 .values(
                     workflow_uuid=workflow_uuid,
                     key=key,
-                    value=_serialization.serialize(message),
+                    value=self.serializer.serialize(message),
                 )
                 .on_conflict_do_update(
                     index_elements=["workflow_uuid", "key"],
-                    set_={"value": _serialization.serialize(message)},
+                    set_={"value": self.serializer.serialize(message)},
                 )
             )
             output: OperationResultInternal = {
                 "workflow_uuid": workflow_uuid,
                 "function_id": function_id,
                 "function_name": function_name,
+                "started_at_epoch_ms": start_time,
                 "output": None,
                 "error": None,
             }
@@ -1578,11 +1552,11 @@ class SystemDatabase(ABC):
                 .values(
                     workflow_uuid=workflow_uuid,
                     key=key,
-                    value=_serialization.serialize(message),
+                    value=self.serializer.serialize(message),
                 )
                 .on_conflict_do_update(
                     index_elements=["workflow_uuid", "key"],
-                    set_={"value": _serialization.serialize(message)},
+                    set_={"value": self.serializer.serialize(message)},
                 )
             )
 
@@ -1607,7 +1581,7 @@ class SystemDatabase(ABC):
             events: Dict[str, Any] = {}
             for row in rows:
                 key = row[0]
-                value = _serialization.deserialize(row[1])
+                value = self.serializer.deserialize(row[1])
                 events[key] = value
 
             return events
@@ -1621,6 +1595,7 @@ class SystemDatabase(ABC):
         caller_ctx: Optional[GetEventWorkflowContext] = None,
     ) -> Any:
         function_name = "DBOS.getEvent"
+        start_time = int(time.time() * 1000)
         get_sql = sa.select(
             SystemSchema.workflow_events.c.value,
         ).where(
@@ -1641,7 +1616,7 @@ class SystemDatabase(ABC):
                     f"Replaying get_event, id: {caller_ctx['function_id']}, key: {key}"
                 )
                 if recorded_output["output"] is not None:
-                    return _serialization.deserialize(recorded_output["output"])
+                    return self.serializer.deserialize(recorded_output["output"])
                 else:
                     raise Exception("No output recorded in the last get_event")
             else:
@@ -1666,7 +1641,7 @@ class SystemDatabase(ABC):
 
         value: Any = None
         if len(init_recv) > 0:
-            value = _serialization.deserialize(init_recv[0][0])
+            value = self.serializer.deserialize(init_recv[0][0])
         else:
             # Wait for the notification
             actual_timeout = timeout_seconds
@@ -1684,7 +1659,7 @@ class SystemDatabase(ABC):
             with self.engine.begin() as c:
                 final_recv = c.execute(get_sql).fetchall()
                 if len(final_recv) > 0:
-                    value = _serialization.deserialize(final_recv[0][0])
+                    value = self.serializer.deserialize(final_recv[0][0])
         condition.release()
         self.workflow_events_map.pop(payload)
 
@@ -1695,7 +1670,8 @@ class SystemDatabase(ABC):
                     "workflow_uuid": caller_ctx["workflow_uuid"],
                     "function_id": caller_ctx["function_id"],
                     "function_name": function_name,
-                    "output": _serialization.serialize(
+                    "started_at_epoch_ms": start_time,
+                    "output": self.serializer.serialize(
                         value
                     ),  # None will be serialized to 'null'
                     "error": None,
@@ -1703,8 +1679,41 @@ class SystemDatabase(ABC):
             )
         return value
 
+    @db_retry()
+    def get_queue_partitions(self, queue_name: str) -> List[str]:
+        """
+        Get all unique partition names associated with a queue for ENQUEUED workflows.
+
+        Args:
+            queue_name: The name of the queue to get partitions for
+
+        Returns:
+            A list of unique partition names for the queue
+        """
+        with self.engine.begin() as c:
+            query = (
+                sa.select(SystemSchema.workflow_status.c.queue_partition_key)
+                .distinct()
+                .where(SystemSchema.workflow_status.c.queue_name == queue_name)
+                .where(
+                    SystemSchema.workflow_status.c.status.in_(
+                        [
+                            WorkflowStatusString.ENQUEUED.value,
+                        ]
+                    )
+                )
+                .where(SystemSchema.workflow_status.c.queue_partition_key.isnot(None))
+            )
+
+            rows = c.execute(query).fetchall()
+            return [row[0] for row in rows]
+
     def start_queued_workflows(
-        self, queue: "Queue", executor_id: str, app_version: str
+        self,
+        queue: "Queue",
+        executor_id: str,
+        app_version: str,
+        queue_partition_key: Optional[str],
     ) -> List[str]:
         if self._debug_mode:
             return []
@@ -1723,6 +1732,10 @@ class SystemDatabase(ABC):
                     sa.select(sa.func.count())
                     .select_from(SystemSchema.workflow_status)
                     .where(SystemSchema.workflow_status.c.queue_name == queue.name)
+                    .where(
+                        SystemSchema.workflow_status.c.queue_partition_key
+                        == queue_partition_key
+                    )
                     .where(
                         SystemSchema.workflow_status.c.status
                         != WorkflowStatusString.ENQUEUED.value
@@ -1747,6 +1760,10 @@ class SystemDatabase(ABC):
                     )
                     .select_from(SystemSchema.workflow_status)
                     .where(SystemSchema.workflow_status.c.queue_name == queue.name)
+                    .where(
+                        SystemSchema.workflow_status.c.queue_partition_key
+                        == queue_partition_key
+                    )
                     .where(
                         SystemSchema.workflow_status.c.status
                         == WorkflowStatusString.PENDING.value
@@ -1788,6 +1805,10 @@ class SystemDatabase(ABC):
                 )
                 .select_from(SystemSchema.workflow_status)
                 .where(SystemSchema.workflow_status.c.queue_name == queue.name)
+                .where(
+                    SystemSchema.workflow_status.c.queue_partition_key
+                    == queue_partition_key
+                )
                 .where(
                     SystemSchema.workflow_status.c.status
                     == WorkflowStatusString.ENQUEUED.value
@@ -1888,6 +1909,7 @@ class SystemDatabase(ABC):
 
     def call_function_as_step(self, fn: Callable[[], T], function_name: str) -> T:
         ctx = get_local_dbos_context()
+        start_time = int(time.time() * 1000)
         if ctx and ctx.is_transaction():
             raise Exception(f"Invalid call to `{function_name}` inside a transaction")
         if ctx and ctx.is_workflow():
@@ -1897,12 +1919,13 @@ class SystemDatabase(ABC):
             )
             if res is not None:
                 if res["output"] is not None:
-                    resstat: SystemDatabase.T = _serialization.deserialize(
+                    resstat: SystemDatabase.T = self.serializer.deserialize(
                         res["output"]
                     )
                     return resstat
                 elif res["error"] is not None:
-                    raise _serialization.deserialize_exception(res["error"])
+                    e: Exception = self.serializer.deserialize(res["error"])
+                    raise e
                 else:
                     raise Exception(
                         f"Recorded output and error are both None for {function_name}"
@@ -1914,7 +1937,8 @@ class SystemDatabase(ABC):
                     "workflow_uuid": ctx.workflow_id,
                     "function_id": ctx.function_id,
                     "function_name": function_name,
-                    "output": _serialization.serialize(result),
+                    "started_at_epoch_ms": start_time,
+                    "output": self.serializer.serialize(result),
                     "error": None,
                 }
             )
@@ -1968,7 +1992,7 @@ class SystemDatabase(ABC):
             )
 
             # Serialize the value before storing
-            serialized_value = _serialization.serialize(value)
+            serialized_value = self.serializer.serialize(value)
 
             # Insert the new stream entry
             c.execute(
@@ -1992,6 +2016,7 @@ class SystemDatabase(ABC):
             if value == _dbos_stream_closed_sentinel
             else "DBOS.writeStream"
         )
+        start_time = int(time.time() * 1000)
 
         with self.engine.begin() as c:
 
@@ -2023,7 +2048,7 @@ class SystemDatabase(ABC):
             )
 
             # Serialize the value before storing
-            serialized_value = _serialization.serialize(value)
+            serialized_value = self.serializer.serialize(value)
 
             # Insert the new stream entry
             c.execute(
@@ -2038,6 +2063,7 @@ class SystemDatabase(ABC):
                 "workflow_uuid": workflow_uuid,
                 "function_id": function_id,
                 "function_name": function_name,
+                "started_at_epoch_ms": start_time,
                 "output": None,
                 "error": None,
             }
@@ -2068,7 +2094,7 @@ class SystemDatabase(ABC):
                 )
 
             # Deserialize the value before returning
-            return _serialization.deserialize(result[0])
+            return self.serializer.deserialize(result[0])
 
     def garbage_collect(
         self, cutoff_epoch_timestamp_ms: Optional[int], rows_threshold: Optional[int]