dbos 0.25.1__py3-none-any.whl → 0.26.0__py3-none-any.whl

dbos/_sys_db.py

@@ -1,20 +1,23 @@
 import datetime
+import json
 import logging
 import os
 import re
 import threading
 import time
+import uuid
 from enum import Enum
 from typing import (
     TYPE_CHECKING,
     Any,
+    Callable,
     Dict,
     List,
     Literal,
     Optional,
     Sequence,
-    Set,
     TypedDict,
+    TypeVar,
 )
 
 import psycopg
@@ -25,7 +28,7 @@ from alembic.config import Config
 from sqlalchemy.exc import DBAPIError
 from sqlalchemy.sql import func
 
-from dbos._utils import GlobalParams
+from dbos._utils import INTERNAL_QUEUE_NAME, GlobalParams
 
 from . import _serialization
 from ._context import get_local_dbos_context
@@ -34,6 +37,8 @@ from ._error import (
     DBOSConflictingWorkflowError,
     DBOSDeadLetterQueueError,
     DBOSNonExistentWorkflowError,
+    DBOSUnexpectedStepError,
+    DBOSWorkflowCancelledError,
     DBOSWorkflowConflictIDError,
 )
 from ._logger import dbos_logger
@@ -60,6 +65,50 @@ WorkflowStatuses = Literal[
 ]
 
 
+class WorkflowStatus:
+    # The workflow ID
+    workflow_id: str
+    # The workflow status. Must be one of ENQUEUED, PENDING, SUCCESS, ERROR, CANCELLED, or RETRIES_EXCEEDED
+    status: str
+    # The name of the workflow function
+    name: str
+    # The name of the workflow's class, if any
+    class_name: Optional[str]
+    # The name with which the workflow's class instance was configured, if any
+    config_name: Optional[str]
+    # The user who ran the workflow, if specified
+    authenticated_user: Optional[str]
+    # The role with which the workflow ran, if specified
+    assumed_role: Optional[str]
+    # All roles which the authenticated user could assume
+    authenticated_roles: Optional[list[str]]
+    # The deserialized workflow input object
+    input: Optional[_serialization.WorkflowInputs]
+    # The workflow's output, if any
+    output: Optional[Any] = None
+    # The error the workflow threw, if any
+    error: Optional[Exception] = None
+    # Workflow start time, as a Unix epoch timestamp in ms
+    created_at: Optional[int]
+    # Last time the workflow status was updated, as a Unix epoch timestamp in ms
+    updated_at: Optional[int]
+    # If this workflow was enqueued, on which queue
+    queue_name: Optional[str]
+    # The executor to most recently executed this workflow
+    executor_id: Optional[str]
+    # The application version on which this workflow was started
+    app_version: Optional[str]
+
+    # INTERNAL FIELDS
+
+    # The ID of the application executing this workflow
+    app_id: Optional[str]
+    # The number of times this workflow's execution has been attempted
+    recovery_attempts: Optional[int]
+    # The HTTP request that triggered the workflow, if known
+    request: Optional[str]
+
+
 class WorkflowStatusInternal(TypedDict):
     workflow_uuid: str
     status: WorkflowStatuses
@@ -79,6 +128,11 @@ 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]
 
 
 class RecordedResult(TypedDict):
@@ -128,6 +182,9 @@ class GetWorkflowsInput:
         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):
@@ -141,11 +198,6 @@ class GetQueuedWorkflowsInput(TypedDict):
     sort_desc: Optional[bool]  # Sort by created_at in DESC or ASC order
 
 
-class GetWorkflowsOutput:
-    def __init__(self, workflow_uuids: List[str]):
-        self.workflow_uuids = workflow_uuids
-
-
 class GetPendingWorkflowsOutput:
     def __init__(self, *, workflow_uuid: str, queue_name: Optional[str] = None):
         self.workflow_uuid: str = workflow_uuid
@@ -278,12 +330,14 @@ class SystemDatabase:
     def insert_workflow_status(
         self,
         status: WorkflowStatusInternal,
+        conn: sa.Connection,
         *,
-        max_recovery_attempts: int = DEFAULT_MAX_RECOVERY_ATTEMPTS,
-    ) -> WorkflowStatuses:
+        max_recovery_attempts: Optional[int],
+    ) -> tuple[WorkflowStatuses, Optional[int]]:
         if self._debug_mode:
             raise Exception("called insert_workflow_status in debug mode")
         wf_status: WorkflowStatuses = status["status"]
+        workflow_deadline_epoch_ms: Optional[int] = status["workflow_deadline_epoch_ms"]
 
         cmd = (
             pg.insert(SystemSchema.workflow_status)
@@ -306,6 +360,8 @@ class SystemDatabase:
                 recovery_attempts=(
                     1 if wf_status != WorkflowStatusString.ENQUEUED.value else 0
                 ),
+                workflow_timeout_ms=status["workflow_timeout_ms"],
+                workflow_deadline_epoch_ms=status["workflow_deadline_epoch_ms"],
             )
             .on_conflict_do_update(
                 index_elements=["workflow_uuid"],
@@ -319,10 +375,9 @@ class SystemDatabase:
             )
         )
 
-        cmd = cmd.returning(SystemSchema.workflow_status.c.recovery_attempts, SystemSchema.workflow_status.c.status, SystemSchema.workflow_status.c.name, SystemSchema.workflow_status.c.class_name, SystemSchema.workflow_status.c.config_name, SystemSchema.workflow_status.c.queue_name)  # type: ignore
+        cmd = cmd.returning(SystemSchema.workflow_status.c.recovery_attempts, SystemSchema.workflow_status.c.status, SystemSchema.workflow_status.c.workflow_deadline_epoch_ms, SystemSchema.workflow_status.c.name, SystemSchema.workflow_status.c.class_name, SystemSchema.workflow_status.c.config_name, SystemSchema.workflow_status.c.queue_name)  # type: ignore
 
-        with self.engine.begin() as c:
-            results = c.execute(cmd)
+        results = conn.execute(cmd)
 
         row = results.fetchone()
         if row is not None:
@@ -330,51 +385,58 @@ class SystemDatabase:
             # A mismatch indicates a workflow starting with the same UUID but different functions, which would throw an exception.
             recovery_attempts: int = row[0]
             wf_status = row[1]
+            workflow_deadline_epoch_ms = row[2]
             err_msg: Optional[str] = None
-            if row[2] != status["name"]:
-                err_msg = f"Workflow already exists with a different function name: {row[2]}, but the provided function name is: {status['name']}"
-            elif row[3] != status["class_name"]:
-                err_msg = f"Workflow already exists with a different class name: {row[3]}, but the provided class name is: {status['class_name']}"
-            elif row[4] != status["config_name"]:
-                err_msg = f"Workflow already exists with a different config name: {row[4]}, but the provided config name is: {status['config_name']}"
-            elif row[5] != status["queue_name"]:
+            if row[3] != status["name"]:
+                err_msg = f"Workflow already exists with a different function name: {row[3]}, but the provided function name is: {status['name']}"
+            elif row[4] != status["class_name"]:
+                err_msg = f"Workflow already exists with a different class name: {row[4]}, but the provided class name is: {status['class_name']}"
+            elif row[5] != status["config_name"]:
+                err_msg = f"Workflow already exists with a different config name: {row[5]}, but the provided config name is: {status['config_name']}"
+            elif row[6] != status["queue_name"]:
                 # This is a warning because a different queue name is not necessarily an error.
                 dbos_logger.warning(
-                    f"Workflow already exists in queue: {row[5]}, but the provided queue name is: {status['queue_name']}. The queue is not updated."
+                    f"Workflow already exists in queue: {row[6]}, but the provided queue name is: {status['queue_name']}. The queue is not updated."
                 )
             if err_msg is not None:
                 raise DBOSConflictingWorkflowError(status["workflow_uuid"], err_msg)
 
             # Every time we start executing a workflow (and thus attempt to insert its status), we increment `recovery_attempts` by 1.
             # When this number becomes equal to `maxRetries + 1`, we mark the workflow as `RETRIES_EXCEEDED`.
-            if recovery_attempts > max_recovery_attempts + 1:
-                with self.engine.begin() as c:
-                    c.execute(
-                        sa.delete(SystemSchema.workflow_queue).where(
-                            SystemSchema.workflow_queue.c.workflow_uuid
-                            == status["workflow_uuid"]
-                        )
+            if (
+                (wf_status != "SUCCESS" and wf_status != "ERROR")
+                and max_recovery_attempts is not None
+                and recovery_attempts > max_recovery_attempts + 1
+            ):
+                delete_cmd = sa.delete(SystemSchema.workflow_queue).where(
+                    SystemSchema.workflow_queue.c.workflow_uuid
+                    == status["workflow_uuid"]
+                )
+                conn.execute(delete_cmd)
+
+                dlq_cmd = (
+                    sa.update(SystemSchema.workflow_status)
+                    .where(
+                        SystemSchema.workflow_status.c.workflow_uuid
+                        == status["workflow_uuid"]
                     )
-                    c.execute(
-                        sa.update(SystemSchema.workflow_status)
-                        .where(
-                            SystemSchema.workflow_status.c.workflow_uuid
-                            == status["workflow_uuid"]
-                        )
-                        .where(
-                            SystemSchema.workflow_status.c.status
-                            == WorkflowStatusString.PENDING.value
-                        )
-                        .values(
-                            status=WorkflowStatusString.RETRIES_EXCEEDED.value,
-                            queue_name=None,
-                        )
+                    .where(
+                        SystemSchema.workflow_status.c.status
+                        == WorkflowStatusString.PENDING.value
+                    )
+                    .values(
+                        status=WorkflowStatusString.RETRIES_EXCEEDED.value,
+                        queue_name=None,
                     )
+                )
+                conn.execute(dlq_cmd)
+                # Need to commit here because we're throwing an exception
+                conn.commit()
                 raise DBOSDeadLetterQueueError(
                     status["workflow_uuid"], max_recovery_attempts
                 )
 
-        return wf_status
+        return wf_status, workflow_deadline_epoch_ms
 
     def update_workflow_status(
         self,
@@ -432,6 +494,18 @@ class SystemDatabase:
         if self._debug_mode:
             raise Exception("called cancel_workflow in debug mode")
         with self.engine.begin() as c:
+            # Check the status of the workflow. If it is complete, do nothing.
+            row = c.execute(
+                sa.select(
+                    SystemSchema.workflow_status.c.status,
+                ).where(SystemSchema.workflow_status.c.workflow_uuid == workflow_id)
+            ).fetchone()
+            if (
+                row is None
+                or row[0] == WorkflowStatusString.SUCCESS.value
+                or row[0] == WorkflowStatusString.ERROR.value
+            ):
+                return
             # Remove the workflow from the queues table so it does not block the table
             c.execute(
                 sa.delete(SystemSchema.workflow_queue).where(
@@ -447,13 +521,12 @@ class SystemDatabase:
                 )
             )
 
-    def resume_workflow(
-        self,
-        workflow_id: str,
-    ) -> None:
+    def resume_workflow(self, workflow_id: str) -> None:
         if self._debug_mode:
             raise Exception("called resume_workflow in debug mode")
         with self.engine.begin() as c:
+            # Execute with snapshot isolation in case of concurrent calls on the same workflow
+            c.execute(sa.text("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ"))
             # Check the status of the workflow. If it is complete, do nothing.
             row = c.execute(
                 sa.select(
@@ -472,13 +545,113 @@ class SystemDatabase:
                     SystemSchema.workflow_queue.c.workflow_uuid == workflow_id
                 )
             )
-            # Set the workflow's status to PENDING and clear its recovery attempts.
+            # Enqueue the workflow on the internal queue
+            c.execute(
+                pg.insert(SystemSchema.workflow_queue).values(
+                    workflow_uuid=workflow_id,
+                    queue_name=INTERNAL_QUEUE_NAME,
+                )
+            )
+            # Set the workflow's status to ENQUEUED and clear its recovery attempts and deadline.
             c.execute(
                 sa.update(SystemSchema.workflow_status)
                 .where(SystemSchema.workflow_status.c.workflow_uuid == workflow_id)
-                .values(status=WorkflowStatusString.PENDING.value, recovery_attempts=0)
+                .values(
+                    status=WorkflowStatusString.ENQUEUED.value,
+                    recovery_attempts=0,
+                    workflow_deadline_epoch_ms=None,
+                )
             )
 
+    def get_max_function_id(self, workflow_uuid: str) -> Optional[int]:
+        with self.engine.begin() as conn:
+            max_function_id_row = conn.execute(
+                sa.select(
+                    sa.func.max(SystemSchema.operation_outputs.c.function_id)
+                ).where(SystemSchema.operation_outputs.c.workflow_uuid == workflow_uuid)
+            ).fetchone()
+
+            max_function_id = max_function_id_row[0] if max_function_id_row else None
+
+            return max_function_id
+
+    def fork_workflow(
+        self, original_workflow_id: str, forked_workflow_id: str, start_step: int = 1
+    ) -> str:
+
+        status = self.get_workflow_status(original_workflow_id)
+        if status is None:
+            raise Exception(f"Workflow {original_workflow_id} not found")
+        inputs = self.get_workflow_inputs(original_workflow_id)
+        if inputs is None:
+            raise Exception(f"Workflow {original_workflow_id} not found")
+
+        with self.engine.begin() as c:
+            # Create an entry for the forked workflow with the same
+            # initial values as the original.
+            c.execute(
+                pg.insert(SystemSchema.workflow_status).values(
+                    workflow_uuid=forked_workflow_id,
+                    status=WorkflowStatusString.ENQUEUED.value,
+                    name=status["name"],
+                    class_name=status["class_name"],
+                    config_name=status["config_name"],
+                    application_version=status["app_version"],
+                    application_id=status["app_id"],
+                    request=status["request"],
+                    authenticated_user=status["authenticated_user"],
+                    authenticated_roles=status["authenticated_roles"],
+                    assumed_role=status["assumed_role"],
+                    queue_name=INTERNAL_QUEUE_NAME,
+                )
+            )
+            # Copy the original workflow's inputs into the forked workflow
+            c.execute(
+                pg.insert(SystemSchema.workflow_inputs).values(
+                    workflow_uuid=forked_workflow_id,
+                    inputs=_serialization.serialize_args(inputs),
+                )
+            )
+
+            if start_step > 1:
+
+                # Copy the original workflow's outputs into the forked workflow
+                insert_stmt = sa.insert(SystemSchema.operation_outputs).from_select(
+                    [
+                        "workflow_uuid",
+                        "function_id",
+                        "output",
+                        "error",
+                        "function_name",
+                        "child_workflow_id",
+                    ],
+                    sa.select(
+                        sa.literal(forked_workflow_id).label("workflow_uuid"),
+                        SystemSchema.operation_outputs.c.function_id,
+                        SystemSchema.operation_outputs.c.output,
+                        SystemSchema.operation_outputs.c.error,
+                        SystemSchema.operation_outputs.c.function_name,
+                        SystemSchema.operation_outputs.c.child_workflow_id,
+                    ).where(
+                        (
+                            SystemSchema.operation_outputs.c.workflow_uuid
+                            == original_workflow_id
+                        )
+                        & (SystemSchema.operation_outputs.c.function_id < start_step)
+                    ),
+                )
+
+                c.execute(insert_stmt)
+
+            # Enqueue the forked workflow on the internal queue
+            c.execute(
+                pg.insert(SystemSchema.workflow_queue).values(
+                    workflow_uuid=forked_workflow_id,
+                    queue_name=INTERNAL_QUEUE_NAME,
+                )
+            )
+        return forked_workflow_id
+
     def get_workflow_status(
         self, workflow_uuid: str
     ) -> Optional[WorkflowStatusInternal]:
@@ -500,6 +673,8 @@ class SystemDatabase:
                     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,
                 ).where(SystemSchema.workflow_status.c.workflow_uuid == workflow_uuid)
             ).fetchone()
             if row is None:
@@ -523,12 +698,12 @@ class SystemDatabase:
                 "updated_at": row[12],
                 "app_version": row[13],
                 "app_id": row[14],
+                "workflow_deadline_epoch_ms": row[15],
+                "workflow_timeout_ms": row[16],
             }
             return status
 
-    def await_workflow_result_internal(self, workflow_uuid: str) -> dict[str, Any]:
-        polling_interval_secs: float = 1.000
-
+    def await_workflow_result(self, workflow_id: str) -> Any:
         while True:
             with self.engine.begin() as c:
                 row = c.execute(
@@ -536,44 +711,26 @@ class SystemDatabase:
                         SystemSchema.workflow_status.c.status,
                         SystemSchema.workflow_status.c.output,
                         SystemSchema.workflow_status.c.error,
-                    ).where(
-                        SystemSchema.workflow_status.c.workflow_uuid == workflow_uuid
-                    )
+                    ).where(SystemSchema.workflow_status.c.workflow_uuid == workflow_id)
                 ).fetchone()
                 if row is not None:
                     status = row[0]
-                    if status == str(WorkflowStatusString.SUCCESS.value):
-                        return {
-                            "status": status,
-                            "output": row[1],
-                            "workflow_uuid": workflow_uuid,
-                        }
-
-                    elif status == str(WorkflowStatusString.ERROR.value):
-                        return {
-                            "status": status,
-                            "error": row[2],
-                            "workflow_uuid": workflow_uuid,
-                        }
-
+                    if status == WorkflowStatusString.SUCCESS.value:
+                        output = row[1]
+                        return _serialization.deserialize(output)
+                    elif status == WorkflowStatusString.ERROR.value:
+                        error = row[2]
+                        raise _serialization.deserialize_exception(error)
+                    elif status == WorkflowStatusString.CANCELLED.value:
+                        # Raise a normal exception here, not the cancellation exception
+                        # because the awaiting workflow is not being cancelled.
+                        raise Exception(f"Awaited workflow {workflow_id} was cancelled")
                 else:
                     pass  # CB: I guess we're assuming the WF will show up eventually.
-
-            time.sleep(polling_interval_secs)
-
-    def await_workflow_result(self, workflow_uuid: str) -> Any:
-        stat = self.await_workflow_result_internal(workflow_uuid)
-        if not stat:
-            return None
-        status: str = stat["status"]
-        if status == str(WorkflowStatusString.SUCCESS.value):
-            return _serialization.deserialize(stat["output"])
-        elif status == str(WorkflowStatusString.ERROR.value):
-            raise _serialization.deserialize_exception(stat["error"])
-        return None
+            time.sleep(1)
 
     def update_workflow_inputs(
-        self, workflow_uuid: str, inputs: str, conn: Optional[sa.Connection] = None
+        self, workflow_uuid: str, inputs: str, conn: sa.Connection
     ) -> None:
         if self._debug_mode:
             raise Exception("called update_workflow_inputs in debug mode")
@@ -590,11 +747,8 @@ class SystemDatabase:
             )
             .returning(SystemSchema.workflow_inputs.c.inputs)
         )
-        if conn is not None:
-            row = conn.execute(cmd).fetchone()
-        else:
-            with self.engine.begin() as c:
-                row = c.execute(cmd).fetchone()
+
+        row = conn.execute(cmd).fetchone()
         if row is not None and row[0] != inputs:
             # In a distributed environment, scheduled workflows are enqueued multiple times with slightly different timestamps
             if not workflow_uuid.startswith("sched-"):
@@ -621,8 +775,37 @@ class SystemDatabase:
             )
             return inputs
 
-    def get_workflows(self, input: GetWorkflowsInput) -> GetWorkflowsOutput:
-        query = sa.select(SystemSchema.workflow_status.c.workflow_uuid)
+    def get_workflows(
+        self, input: GetWorkflowsInput, get_request: bool = False
+    ) -> List[WorkflowStatus]:
+        """
+        Retrieve a list of workflows result and inputs based on the input criteria. The result is a list of external-facing workflow status objects.
+        """
+        query = sa.select(
+            SystemSchema.workflow_status.c.workflow_uuid,
+            SystemSchema.workflow_status.c.status,
+            SystemSchema.workflow_status.c.name,
+            SystemSchema.workflow_status.c.request,
+            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_inputs.c.inputs,
+            SystemSchema.workflow_status.c.output,
+            SystemSchema.workflow_status.c.error,
+        ).join(
+            SystemSchema.workflow_inputs,
+            SystemSchema.workflow_status.c.workflow_uuid
+            == SystemSchema.workflow_inputs.c.workflow_uuid,
+        )
         if input.sort_desc:
             query = query.order_by(SystemSchema.workflow_status.c.created_at.desc())
         else:
@@ -655,6 +838,12 @@ class SystemDatabase:
             query = query.where(
                 SystemSchema.workflow_status.c.workflow_uuid.in_(input.workflow_ids)
             )
+        if input.workflow_id_prefix:
+            query = query.where(
+                SystemSchema.workflow_status.c.workflow_uuid.startswith(
+                    input.workflow_id_prefix
+                )
+            )
         if input.limit:
             query = query.limit(input.limit)
         if input.offset:
@@ -662,18 +851,76 @@ class SystemDatabase:
 
         with self.engine.begin() as c:
             rows = c.execute(query)
-        workflow_ids = [row[0] for row in rows]
 
-        return GetWorkflowsOutput(workflow_ids)
+        infos: List[WorkflowStatus] = []
+        for row in rows:
+            info = WorkflowStatus()
+            info.workflow_id = row[0]
+            info.status = row[1]
+            info.name = row[2]
+            info.request = row[3] if get_request else None
+            info.recovery_attempts = row[4]
+            info.config_name = row[5]
+            info.class_name = row[6]
+            info.authenticated_user = row[7]
+            info.authenticated_roles = (
+                json.loads(row[8]) if row[8] is not None else None
+            )
+            info.assumed_role = row[9]
+            info.queue_name = row[10]
+            info.executor_id = row[11]
+            info.created_at = row[12]
+            info.updated_at = row[13]
+            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])
+
+            infos.append(info)
+        return infos
 
     def get_queued_workflows(
-        self, input: GetQueuedWorkflowsInput
-    ) -> GetWorkflowsOutput:
-
-        query = sa.select(SystemSchema.workflow_queue.c.workflow_uuid).join(
-            SystemSchema.workflow_status,
-            SystemSchema.workflow_queue.c.workflow_uuid
-            == SystemSchema.workflow_status.c.workflow_uuid,
+        self, input: GetQueuedWorkflowsInput, get_request: bool = False
+    ) -> 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.
+        """
+        query = sa.select(
+            SystemSchema.workflow_status.c.workflow_uuid,
+            SystemSchema.workflow_status.c.status,
+            SystemSchema.workflow_status.c.name,
+            SystemSchema.workflow_status.c.request,
+            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_inputs.c.inputs,
+            SystemSchema.workflow_status.c.output,
+            SystemSchema.workflow_status.c.error,
+        ).select_from(
+            SystemSchema.workflow_queue.join(
+                SystemSchema.workflow_status,
+                SystemSchema.workflow_queue.c.workflow_uuid
+                == SystemSchema.workflow_status.c.workflow_uuid,
+            ).join(
+                SystemSchema.workflow_inputs,
+                SystemSchema.workflow_queue.c.workflow_uuid
+                == SystemSchema.workflow_inputs.c.workflow_uuid,
+            )
         )
         if input["sort_desc"]:
             query = query.order_by(SystemSchema.workflow_status.c.created_at.desc())
@@ -710,9 +957,40 @@ class SystemDatabase:
 
         with self.engine.begin() as c:
             rows = c.execute(query)
-        workflow_uuids = [row[0] for row in rows]
 
-        return GetWorkflowsOutput(workflow_uuids)
+        infos: List[WorkflowStatus] = []
+        for row in rows:
+            info = WorkflowStatus()
+            info.workflow_id = row[0]
+            info.status = row[1]
+            info.name = row[2]
+            info.request = row[3] if get_request else None
+            info.recovery_attempts = row[4]
+            info.config_name = row[5]
+            info.class_name = row[6]
+            info.authenticated_user = row[7]
+            info.authenticated_roles = (
+                json.loads(row[8]) if row[8] is not None else None
+            )
+            info.assumed_role = row[9]
+            info.queue_name = row[10]
+            info.executor_id = row[11]
+            info.created_at = row[12]
+            info.updated_at = row[13]
+            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])
+
+            infos.append(info)
+
+        return infos
 
     def get_pending_workflows(
         self, executor_id: str, app_version: str
@@ -844,28 +1122,74 @@ class SystemDatabase:
             raise
 
     def check_operation_execution(
-        self, workflow_uuid: str, function_id: int, conn: Optional[sa.Connection] = None
+        self,
+        workflow_id: str,
+        function_id: int,
+        function_name: str,
+        *,
+        conn: Optional[sa.Connection] = None,
     ) -> Optional[RecordedResult]:
-        sql = sa.select(
+        # First query: Retrieve the workflow status
+        workflow_status_sql = sa.select(
+            SystemSchema.workflow_status.c.status,
+        ).where(SystemSchema.workflow_status.c.workflow_uuid == workflow_id)
+
+        # Second query: Retrieve operation outputs if they exist
+        operation_output_sql = sa.select(
             SystemSchema.operation_outputs.c.output,
             SystemSchema.operation_outputs.c.error,
+            SystemSchema.operation_outputs.c.function_name,
         ).where(
-            SystemSchema.operation_outputs.c.workflow_uuid == workflow_uuid,
-            SystemSchema.operation_outputs.c.function_id == function_id,
+            (SystemSchema.operation_outputs.c.workflow_uuid == workflow_id)
+            & (SystemSchema.operation_outputs.c.function_id == function_id)
         )
 
-        # If in a transaction, use the provided connection
-        rows: Sequence[Any]
+        # Execute both queries
         if conn is not None:
-            rows = conn.execute(sql).all()
+            workflow_status_rows = conn.execute(workflow_status_sql).all()
+            operation_output_rows = conn.execute(operation_output_sql).all()
         else:
             with self.engine.begin() as c:
-                rows = c.execute(sql).all()
-        if len(rows) == 0:
+                workflow_status_rows = c.execute(workflow_status_sql).all()
+                operation_output_rows = c.execute(operation_output_sql).all()
+
+        # Check if the workflow exists
+        assert (
+            len(workflow_status_rows) > 0
+        ), f"Error: Workflow {workflow_id} does not exist"
+
+        # Get workflow status
+        workflow_status = workflow_status_rows[0][0]
+
+        # If the workflow is cancelled, raise the exception
+        if workflow_status == WorkflowStatusString.CANCELLED.value:
+            raise DBOSWorkflowCancelledError(
+                f"Workflow {workflow_id} is cancelled. Aborting function."
+            )
+
+        # If there are no operation outputs, return None
+        if not operation_output_rows:
             return None
+
+        # Extract operation output data
+        output, error, recorded_function_name = (
+            operation_output_rows[0][0],
+            operation_output_rows[0][1],
+            operation_output_rows[0][2],
+        )
+
+        # If the provided and recorded function name are different, throw an exception
+        if function_name != recorded_function_name:
+            raise DBOSUnexpectedStepError(
+                workflow_id=workflow_id,
+                step_id=function_id,
+                expected_name=function_name,
+                recorded_name=recorded_function_name,
+            )
+
         result: RecordedResult = {
-            "output": rows[0][0],
-            "error": rows[0][1],
+            "output": output,
+            "error": error,
         }
         return result
 
@@ -894,10 +1218,11 @@ class SystemDatabase:
         message: Any,
         topic: Optional[str] = None,
     ) -> None:
+        function_name = "DBOS.send"
         topic = topic if topic is not None else _dbos_null_topic
         with self.engine.begin() as c:
             recorded_output = self.check_operation_execution(
-                workflow_uuid, function_id, conn=c
+                workflow_uuid, function_id, function_name, conn=c
             )
             if self._debug_mode and recorded_output is None:
                 raise Exception(
@@ -930,7 +1255,7 @@ class SystemDatabase:
             output: OperationResultInternal = {
                 "workflow_uuid": workflow_uuid,
                 "function_id": function_id,
-                "function_name": "DBOS.send",
+                "function_name": function_name,
                 "output": None,
                 "error": None,
             }
@@ -944,10 +1269,13 @@ class SystemDatabase:
         topic: Optional[str],
         timeout_seconds: float = 60,
     ) -> Any:
+        function_name = "DBOS.recv"
         topic = topic if topic is not None else _dbos_null_topic
 
         # First, check for previous executions.
-        recorded_output = self.check_operation_execution(workflow_uuid, function_id)
+        recorded_output = self.check_operation_execution(
+            workflow_uuid, function_id, function_name
+        )
         if self._debug_mode and recorded_output is None:
             raise Exception("called recv in debug mode without a previous execution")
         if recorded_output is not None:
@@ -1024,7 +1352,7 @@ class SystemDatabase:
                 {
                     "workflow_uuid": workflow_uuid,
                     "function_id": function_id,
-                    "function_name": "DBOS.recv",
+                    "function_name": function_name,
                     "output": _serialization.serialize(
                         message
                     ),  # None will be serialized to 'null'
@@ -1098,7 +1426,10 @@ class SystemDatabase:
         seconds: float,
         skip_sleep: bool = False,
     ) -> float:
-        recorded_output = self.check_operation_execution(workflow_uuid, function_id)
+        function_name = "DBOS.sleep"
+        recorded_output = self.check_operation_execution(
+            workflow_uuid, function_id, function_name
+        )
         end_time: float
         if self._debug_mode and recorded_output is None:
             raise Exception("called sleep in debug mode without a previous execution")
@@ -1115,7 +1446,7 @@ class SystemDatabase:
                     {
                         "workflow_uuid": workflow_uuid,
                         "function_id": function_id,
-                        "function_name": "DBOS.sleep",
+                        "function_name": function_name,
                         "output": _serialization.serialize(end_time),
                         "error": None,
                     }
@@ -1134,9 +1465,10 @@ class SystemDatabase:
         key: str,
         message: Any,
     ) -> None:
+        function_name = "DBOS.setEvent"
         with self.engine.begin() as c:
             recorded_output = self.check_operation_execution(
-                workflow_uuid, function_id, conn=c
+                workflow_uuid, function_id, function_name, conn=c
             )
             if self._debug_mode and recorded_output is None:
                 raise Exception(
@@ -1163,7 +1495,7 @@ class SystemDatabase:
             output: OperationResultInternal = {
                 "workflow_uuid": workflow_uuid,
                 "function_id": function_id,
-                "function_name": "DBOS.setEvent",
+                "function_name": function_name,
                 "output": None,
                 "error": None,
             }
@@ -1176,6 +1508,7 @@ class SystemDatabase:
         timeout_seconds: float = 60,
         caller_ctx: Optional[GetEventWorkflowContext] = None,
     ) -> Any:
+        function_name = "DBOS.getEvent"
         get_sql = sa.select(
             SystemSchema.workflow_events.c.value,
         ).where(
@@ -1185,7 +1518,7 @@ class SystemDatabase:
         # Check for previous executions only if it's in a workflow
         if caller_ctx is not None:
             recorded_output = self.check_operation_execution(
-                caller_ctx["workflow_uuid"], caller_ctx["function_id"]
+                caller_ctx["workflow_uuid"], caller_ctx["function_id"], function_name
             )
             if self._debug_mode and recorded_output is None:
                 raise Exception(
@@ -1244,7 +1577,7 @@ class SystemDatabase:
                 {
                     "workflow_uuid": caller_ctx["workflow_uuid"],
                     "function_id": caller_ctx["function_id"],
-                    "function_name": "DBOS.getEvent",
+                    "function_name": function_name,
                     "output": _serialization.serialize(
                         value
                     ),  # None will be serialized to 'null'
@@ -1253,18 +1586,17 @@ class SystemDatabase:
             )
         return value
 
-    def enqueue(self, workflow_id: str, queue_name: str) -> None:
+    def enqueue(self, workflow_id: str, queue_name: str, conn: sa.Connection) -> None:
         if self._debug_mode:
             raise Exception("called enqueue in debug mode")
-        with self.engine.begin() as c:
-            c.execute(
-                pg.insert(SystemSchema.workflow_queue)
-                .values(
-                    workflow_uuid=workflow_id,
-                    queue_name=queue_name,
-                )
-                .on_conflict_do_nothing()
+        conn.execute(
+            pg.insert(SystemSchema.workflow_queue)
+            .values(
+                workflow_uuid=workflow_id,
+                queue_name=queue_name,
             )
+            .on_conflict_do_nothing()
+        )
 
     def start_queued_workflows(
         self, queue: "Queue", executor_id: str, app_version: str
@@ -1403,6 +1735,17 @@ class SystemDatabase:
                         status=WorkflowStatusString.PENDING.value,
                         application_version=app_version,
                         executor_id=executor_id,
+                        # If a timeout is set, set the deadline on dequeue
+                        workflow_deadline_epoch_ms=sa.case(
+                            (
+                                SystemSchema.workflow_status.c.workflow_timeout_ms.isnot(
+                                    None
+                                ),
+                                sa.func.extract("epoch", sa.func.now()) * 1000
+                                + SystemSchema.workflow_status.c.workflow_timeout_ms,
+                            ),
+                            else_=SystemSchema.workflow_status.c.workflow_deadline_epoch_ms,
+                        ),
                     )
                 )
                 if res.rowcount > 0:
@@ -1483,6 +1826,66 @@ class SystemDatabase:
                     )
                 return True
 
+    T = TypeVar("T")
+
+    def call_function_as_step(self, fn: Callable[[], T], function_name: str) -> T:
+        ctx = get_local_dbos_context()
+        if ctx and ctx.is_transaction():
+            raise Exception(f"Invalid call to `{function_name}` inside a transaction")
+        if ctx and ctx.is_workflow():
+            ctx.function_id += 1
+            res = self.check_operation_execution(
+                ctx.workflow_id, ctx.function_id, function_name
+            )
+            if res is not None:
+                if res["output"] is not None:
+                    resstat: SystemDatabase.T = _serialization.deserialize(
+                        res["output"]
+                    )
+                    return resstat
+                elif res["error"] is not None:
+                    raise _serialization.deserialize_exception(res["error"])
+                else:
+                    raise Exception(
+                        f"Recorded output and error are both None for {function_name}"
+                    )
+        result = fn()
+        if ctx and ctx.is_workflow():
+            self.record_operation_result(
+                {
+                    "workflow_uuid": ctx.workflow_id,
+                    "function_id": ctx.function_id,
+                    "function_name": function_name,
+                    "output": _serialization.serialize(result),
+                    "error": None,
+                }
+            )
+        return result
+
+    def init_workflow(
+        self,
+        status: WorkflowStatusInternal,
+        inputs: str,
+        *,
+        max_recovery_attempts: Optional[int],
+    ) -> tuple[WorkflowStatuses, Optional[int]]:
+        """
+        Synchronously record the status and inputs for workflows in a single transaction
+        """
+        with self.engine.begin() as conn:
+            wf_status, workflow_deadline_epoch_ms = self.insert_workflow_status(
+                status, conn, max_recovery_attempts=max_recovery_attempts
+            )
+            # TODO: Modify the inputs if they were changed by `update_workflow_inputs`
+            self.update_workflow_inputs(status["workflow_uuid"], inputs, conn)
+
+            if (
+                status["queue_name"] is not None
+                and wf_status == WorkflowStatusString.ENQUEUED.value
+            ):
+                self.enqueue(status["workflow_uuid"], status["queue_name"], conn)
+        return wf_status, workflow_deadline_epoch_ms
+
 
 def reset_system_database(config: ConfigFile) -> None:
     sysdb_name = (