dbos 0.26.0a19__py3-none-any.whl → 0.26.0a22__py3-none-any.whl

dbos/__init__.py

@@ -5,8 +5,7 @@ from ._dbos import DBOS, DBOSConfiguredInstance, WorkflowHandle
 from ._dbos_config import ConfigFile, DBOSConfig, get_dbos_database_url, load_config
 from ._kafka_message import KafkaMessage
 from ._queue import Queue
-from ._sys_db import GetWorkflowsInput, WorkflowStatusString
-from ._workflow_commands import WorkflowStatus
+from ._sys_db import GetWorkflowsInput, WorkflowStatus, WorkflowStatusString
 
 __all__ = [
     "ConfigFile",

dbos/_client.py

@@ -1,5 +1,6 @@
 import asyncio
 import sys
+import time
 import uuid
 from typing import Any, Generic, List, Optional, TypedDict, TypeVar
 
@@ -19,11 +20,11 @@ from dbos._serialization import WorkflowInputs
 from dbos._sys_db import (
     StepInfo,
     SystemDatabase,
+    WorkflowStatus,
     WorkflowStatusInternal,
     WorkflowStatusString,
 )
 from dbos._workflow_commands import (
-    WorkflowStatus,
     fork_workflow,
     get_workflow,
     list_queued_workflows,
@@ -39,6 +40,7 @@ class EnqueueOptions(TypedDict):
     queue_name: str
     workflow_id: NotRequired[str]
     app_version: NotRequired[str]
+    workflow_timeout: NotRequired[float]
 
 
 class WorkflowHandleClientPolling(Generic[R]):
@@ -54,7 +56,7 @@ class WorkflowHandleClientPolling(Generic[R]):
         res: R = self._sys_db.await_workflow_result(self.workflow_id)
         return res
 
-    def get_status(self) -> "WorkflowStatus":
+    def get_status(self) -> WorkflowStatus:
         status = get_workflow(self._sys_db, self.workflow_id, True)
         if status is None:
             raise DBOSNonExistentWorkflowError(self.workflow_id)
@@ -76,7 +78,7 @@ class WorkflowHandleClientAsyncPolling(Generic[R]):
         )
         return res
 
-    async def get_status(self) -> "WorkflowStatus":
+    async def get_status(self) -> WorkflowStatus:
         status = await asyncio.to_thread(
             get_workflow, self._sys_db, self.workflow_id, True
         )
@@ -107,6 +109,7 @@ class DBOSClient:
         workflow_id = options.get("workflow_id")
         if workflow_id is None:
             workflow_id = str(uuid.uuid4())
+        workflow_timeout = options.get("workflow_timeout", None)
 
         status: WorkflowStatusInternal = {
             "workflow_uuid": workflow_id,
@@ -127,6 +130,10 @@ class DBOSClient:
             "executor_id": None,
             "recovery_attempts": None,
             "app_id": None,
+            "workflow_timeout_ms": (
+                int(workflow_timeout * 1000) if workflow_timeout is not None else None
+            ),
+            "workflow_deadline_epoch_ms": None,
         }
 
         inputs: WorkflowInputs = {
@@ -134,7 +141,9 @@ class DBOSClient:
             "kwargs": kwargs,
         }
 
-        self._sys_db.init_workflow(status, _serialization.serialize_args(inputs))
+        self._sys_db.init_workflow(
+            status, _serialization.serialize_args(inputs), max_recovery_attempts=None
+        )
         return workflow_id
 
     def enqueue(
@@ -188,9 +197,13 @@ class DBOSClient:
             "recovery_attempts": None,
             "app_id": None,
             "app_version": None,
+            "workflow_timeout_ms": None,
+            "workflow_deadline_epoch_ms": None,
         }
         with self._sys_db.engine.begin() as conn:
-            self._sys_db.insert_workflow_status(status, conn)
+            self._sys_db.insert_workflow_status(
+                status, conn, max_recovery_attempts=None
+            )
         self._sys_db.send(status["workflow_uuid"], 0, destination_id, message, topic)
 
     async def send_async(

dbos/_conductor/protocol.py

@@ -3,8 +3,7 @@ from dataclasses import asdict, dataclass
 from enum import Enum
 from typing import List, Optional, Type, TypedDict, TypeVar
 
-from dbos._sys_db import StepInfo
-from dbos._workflow_commands import WorkflowStatus
+from dbos._sys_db import StepInfo, WorkflowStatus
 
 
 class MessageType(str, Enum):

dbos/_context.py

@@ -93,6 +93,11 @@ class DBOSContext:
         self.assumed_role: Optional[str] = None
         self.step_status: Optional[StepStatus] = None
 
+        # A user-specified workflow timeout. Takes priority over a propagated deadline.
+        self.workflow_timeout_ms: Optional[int] = None
+        # A propagated workflow deadline.
+        self.workflow_deadline_epoch_ms: Optional[int] = None
+
     def create_child(self) -> DBOSContext:
         rv = DBOSContext()
         rv.logger = self.logger
@@ -360,11 +365,60 @@ class SetWorkflowID:
         return False  # Did not handle
 
 
+class SetWorkflowTimeout:
+    """
+    Set the workflow timeout (in seconds) to be used for the enclosed workflow invocations.
+
+    Typical Usage
+        ```
+        with SetWorkflowTimeout(<timeout in seconds>):
+            result = workflow_function(...)
+        ```
+    """
+
+    def __init__(self, workflow_timeout_sec: Optional[float]) -> None:
+        if workflow_timeout_sec and not workflow_timeout_sec > 0:
+            raise Exception(
+                f"Invalid workflow timeout {workflow_timeout_sec}. Timeouts must be positive."
+            )
+        self.created_ctx = False
+        self.workflow_timeout_ms = (
+            int(workflow_timeout_sec * 1000)
+            if workflow_timeout_sec is not None
+            else None
+        )
+        self.saved_workflow_timeout: Optional[int] = None
+
+    def __enter__(self) -> SetWorkflowTimeout:
+        # Code to create a basic context
+        ctx = get_local_dbos_context()
+        if ctx is None:
+            self.created_ctx = True
+            _set_local_dbos_context(DBOSContext())
+        ctx = assert_current_dbos_context()
+        self.saved_workflow_timeout = ctx.workflow_timeout_ms
+        ctx.workflow_timeout_ms = self.workflow_timeout_ms
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> Literal[False]:
+        assert_current_dbos_context().workflow_timeout_ms = self.saved_workflow_timeout
+        # Code to clean up the basic context if we created it
+        if self.created_ctx:
+            _clear_local_dbos_context()
+        return False  # Did not handle
+
+
 class EnterDBOSWorkflow(AbstractContextManager[DBOSContext, Literal[False]]):
     def __init__(self, attributes: TracedAttributes) -> None:
         self.created_ctx = False
         self.attributes = attributes
         self.is_temp_workflow = attributes["name"] == "temp_wf"
+        self.saved_workflow_timeout: Optional[int] = None
 
     def __enter__(self) -> DBOSContext:
         # Code to create a basic context
@@ -374,6 +428,10 @@ class EnterDBOSWorkflow(AbstractContextManager[DBOSContext, Literal[False]]):
             ctx = DBOSContext()
             _set_local_dbos_context(ctx)
         assert not ctx.is_within_workflow()
+        # Unset the workflow_timeout_ms context var so it is not applied to this
+        # workflow's children (instead we propagate the deadline)
+        self.saved_workflow_timeout = ctx.workflow_timeout_ms
+        ctx.workflow_timeout_ms = None
         ctx.start_workflow(
             None, self.attributes, self.is_temp_workflow
         )  # Will get from the context's next workflow ID
@@ -388,6 +446,10 @@ class EnterDBOSWorkflow(AbstractContextManager[DBOSContext, Literal[False]]):
         ctx = assert_current_dbos_context()
         assert ctx.is_within_workflow()
         ctx.end_workflow(exc_value, self.is_temp_workflow)
+        # Restore the saved workflow timeout
+        ctx.workflow_timeout_ms = self.saved_workflow_timeout
+        # Clear any propagating timeout
+        ctx.workflow_deadline_epoch_ms = None
         # Code to clean up the basic context if we created it
         if self.created_ctx:
             _clear_local_dbos_context()

dbos/_core.py

@@ -3,6 +3,7 @@ import functools
 import inspect
 import json
 import sys
+import threading
 import time
 import traceback
 from concurrent.futures import Future
@@ -14,11 +15,9 @@ from typing import (
     Coroutine,
     Generic,
     Optional,
-    Tuple,
     TypeVar,
     Union,
     cast,
-    overload,
 )
 
 from dbos._outcome import Immediate, NoResult, Outcome, Pending
@@ -59,7 +58,6 @@ from ._error import (
 )
 from ._registrations import (
     DEFAULT_MAX_RECOVERY_ATTEMPTS,
-    DBOSFuncInfo,
     get_config_name,
     get_dbos_class_name,
     get_dbos_func_name,
@@ -75,6 +73,7 @@ from ._serialization import WorkflowInputs
 from ._sys_db import (
     GetEventWorkflowContext,
     OperationResultInternal,
+    WorkflowStatus,
     WorkflowStatusInternal,
     WorkflowStatusString,
 )
@@ -87,7 +86,6 @@ if TYPE_CHECKING:
         DBOSRegistry,
         IsolationLevel,
     )
-    from ._workflow_commands import WorkflowStatus
 
 from sqlalchemy.exc import DBAPIError, InvalidRequestError
 
@@ -119,7 +117,7 @@ class WorkflowHandleFuture(Generic[R]):
         self.dbos._sys_db.record_get_result(self.workflow_id, serialized_r, None)
         return r
 
-    def get_status(self) -> "WorkflowStatus":
+    def get_status(self) -> WorkflowStatus:
         stat = self.dbos.get_workflow_status(self.workflow_id)
         if stat is None:
             raise DBOSNonExistentWorkflowError(self.workflow_id)
@@ -146,7 +144,7 @@ class WorkflowHandlePolling(Generic[R]):
         self.dbos._sys_db.record_get_result(self.workflow_id, serialized_r, None)
         return r
 
-    def get_status(self) -> "WorkflowStatus":
+    def get_status(self) -> WorkflowStatus:
         stat = self.dbos.get_workflow_status(self.workflow_id)
         if stat is None:
             raise DBOSNonExistentWorkflowError(self.workflow_id)
@@ -181,7 +179,7 @@ class WorkflowHandleAsyncTask(Generic[R]):
         )
         return r
 
-    async def get_status(self) -> "WorkflowStatus":
+    async def get_status(self) -> WorkflowStatus:
         stat = await asyncio.to_thread(self.dbos.get_workflow_status, self.workflow_id)
         if stat is None:
             raise DBOSNonExistentWorkflowError(self.workflow_id)
@@ -217,7 +215,7 @@ class WorkflowHandleAsyncPolling(Generic[R]):
         )
         return r
 
-    async def get_status(self) -> "WorkflowStatus":
+    async def get_status(self) -> WorkflowStatus:
         stat = await asyncio.to_thread(self.dbos.get_workflow_status, self.workflow_id)
         if stat is None:
             raise DBOSNonExistentWorkflowError(self.workflow_id)
@@ -227,19 +225,30 @@ class WorkflowHandleAsyncPolling(Generic[R]):
 def _init_workflow(
     dbos: "DBOS",
     ctx: DBOSContext,
+    *,
     inputs: WorkflowInputs,
     wf_name: str,
     class_name: Optional[str],
     config_name: Optional[str],
-    temp_wf_type: Optional[str],
-    queue: Optional[str] = None,
-    max_recovery_attempts: int = DEFAULT_MAX_RECOVERY_ATTEMPTS,
+    queue: Optional[str],
+    workflow_timeout_ms: Optional[int],
+    workflow_deadline_epoch_ms: Optional[int],
+    max_recovery_attempts: Optional[int],
 ) -> WorkflowStatusInternal:
     wfid = (
         ctx.workflow_id
         if len(ctx.workflow_id) > 0
         else ctx.id_assigned_for_next_workflow
     )
+
+    # In debug mode, just return the existing status
+    if dbos.debug_mode:
+        get_status_result = dbos._sys_db.get_workflow_status(wfid)
+        if get_status_result is None:
+            raise DBOSNonExistentWorkflowError(wfid)
+        return get_status_result
+
+    # Initialize a workflow status object from the context
     status: WorkflowStatusInternal = {
         "workflow_uuid": wfid,
         "status": (
@@ -267,25 +276,47 @@ def _init_workflow(
         "queue_name": queue,
         "created_at": None,
         "updated_at": None,
+        "workflow_timeout_ms": workflow_timeout_ms,
+        "workflow_deadline_epoch_ms": workflow_deadline_epoch_ms,
     }
 
     # If we have a class name, the first arg is the instance and do not serialize
     if class_name is not None:
         inputs = {"args": inputs["args"][1:], "kwargs": inputs["kwargs"]}
 
-    wf_status = status["status"]
-    if dbos.debug_mode:
-        get_status_result = dbos._sys_db.get_workflow_status(wfid)
-        if get_status_result is None:
-            raise DBOSNonExistentWorkflowError(wfid)
-        wf_status = get_status_result["status"]
-    else:
-        wf_status = dbos._sys_db.init_workflow(
-            status,
-            _serialization.serialize_args(inputs),
-            max_recovery_attempts=max_recovery_attempts,
-        )
+    # Synchronously record the status and inputs for workflows
+    wf_status, workflow_deadline_epoch_ms = dbos._sys_db.init_workflow(
+        status,
+        _serialization.serialize_args(inputs),
+        max_recovery_attempts=max_recovery_attempts,
+    )
 
+    if workflow_deadline_epoch_ms is not None:
+        evt = threading.Event()
+        dbos.stop_events.append(evt)
+
+        def timeout_func() -> None:
+            try:
+                assert workflow_deadline_epoch_ms is not None
+                time_to_wait_sec = (
+                    workflow_deadline_epoch_ms - (time.time() * 1000)
+                ) / 1000
+                if time_to_wait_sec > 0:
+                    was_stopped = evt.wait(time_to_wait_sec)
+                    if was_stopped:
+                        return
+                dbos._sys_db.cancel_workflow(wfid)
+            except Exception as e:
+                dbos.logger.warning(
+                    f"Exception in timeout thread for workflow {wfid}: {e}"
+                )
+
+        timeout_thread = threading.Thread(target=timeout_func, daemon=True)
+        timeout_thread.start()
+        dbos._background_threads.append(timeout_thread)
+
+    ctx.workflow_deadline_epoch_ms = workflow_deadline_epoch_ms
+    status["workflow_deadline_epoch_ms"] = workflow_deadline_epoch_ms
     status["status"] = wf_status
     return status
 
@@ -501,6 +532,13 @@ def start_workflow(
         "kwargs": kwargs,
     }
 
+    local_ctx = get_local_dbos_context()
+    workflow_timeout_ms, workflow_deadline_epoch_ms = _get_timeout_deadline(
+        local_ctx, queue_name
+    )
+    workflow_timeout_ms = (
+        local_ctx.workflow_timeout_ms if local_ctx is not None else None
+    )
     new_wf_id, new_wf_ctx = _get_new_wf()
 
     ctx = new_wf_ctx
@@ -519,8 +557,9 @@ def start_workflow(
         wf_name=get_dbos_func_name(func),
         class_name=get_dbos_class_name(fi, func, args),
         config_name=get_config_name(fi, func, args),
-        temp_wf_type=get_temp_workflow_type(func),
         queue=queue_name,
+        workflow_timeout_ms=workflow_timeout_ms,
+        workflow_deadline_epoch_ms=workflow_deadline_epoch_ms,
         max_recovery_attempts=fi.max_recovery_attempts,
     )
 
@@ -583,6 +622,10 @@ async def start_workflow_async(
         "kwargs": kwargs,
     }
 
+    local_ctx = get_local_dbos_context()
+    workflow_timeout_ms, workflow_deadline_epoch_ms = _get_timeout_deadline(
+        local_ctx, queue_name
+    )
     new_wf_id, new_wf_ctx = _get_new_wf()
 
     ctx = new_wf_ctx
@@ -604,8 +647,9 @@ async def start_workflow_async(
         wf_name=get_dbos_func_name(func),
         class_name=get_dbos_class_name(fi, func, args),
         config_name=get_config_name(fi, func, args),
-        temp_wf_type=get_temp_workflow_type(func),
         queue=queue_name,
+        workflow_timeout_ms=workflow_timeout_ms,
+        workflow_deadline_epoch_ms=workflow_deadline_epoch_ms,
         max_recovery_attempts=fi.max_recovery_attempts,
     )
 
@@ -653,7 +697,7 @@ else:
 def workflow_wrapper(
     dbosreg: "DBOSRegistry",
     func: Callable[P, R],
-    max_recovery_attempts: int = DEFAULT_MAX_RECOVERY_ATTEMPTS,
+    max_recovery_attempts: Optional[int] = DEFAULT_MAX_RECOVERY_ATTEMPTS,
 ) -> Callable[P, R]:
     func.__orig_func = func  # type: ignore
 
@@ -680,6 +724,9 @@ def workflow_wrapper(
             "kwargs": kwargs,
         }
         ctx = get_local_dbos_context()
+        workflow_timeout_ms, workflow_deadline_epoch_ms = _get_timeout_deadline(
+            ctx, queue=None
+        )
         enterWorkflowCtxMgr = (
             EnterDBOSChildWorkflow if ctx and ctx.is_workflow() else EnterDBOSWorkflow
         )
@@ -717,7 +764,9 @@ def workflow_wrapper(
                 wf_name=get_dbos_func_name(func),
                 class_name=get_dbos_class_name(fi, func, args),
                 config_name=get_config_name(fi, func, args),
-                temp_wf_type=get_temp_workflow_type(func),
+                queue=None,
+                workflow_timeout_ms=workflow_timeout_ms,
+                workflow_deadline_epoch_ms=workflow_deadline_epoch_ms,
                 max_recovery_attempts=max_recovery_attempts,
             )
 
@@ -765,7 +814,7 @@ def workflow_wrapper(
 
 
 def decorate_workflow(
-    reg: "DBOSRegistry", max_recovery_attempts: int
+    reg: "DBOSRegistry", max_recovery_attempts: Optional[int]
 ) -> Callable[[Callable[P, R]], Callable[P, R]]:
     def _workflow_decorator(func: Callable[P, R]) -> Callable[P, R]:
         wrapped_func = workflow_wrapper(reg, func, max_recovery_attempts)
@@ -1211,3 +1260,24 @@ def get_event(
     else:
         # Directly call it outside of a workflow
         return dbos._sys_db.get_event(workflow_id, key, timeout_seconds)
+
+
+def _get_timeout_deadline(
+    ctx: Optional[DBOSContext], queue: Optional[str]
+) -> tuple[Optional[int], Optional[int]]:
+    if ctx is None:
+        return None, None
+    # If a timeout is explicitly specified, use it over any propagated deadline
+    if ctx.workflow_timeout_ms:
+        if queue:
+            # Queued workflows are assigned a deadline on dequeue
+            return ctx.workflow_timeout_ms, None
+        else:
+            # Otherwise, compute the deadline immediately
+            return (
+                ctx.workflow_timeout_ms,
+                int(time.time() * 1000) + ctx.workflow_timeout_ms,
+            )
+    # Otherwise, return the propagated deadline, if any
+    else:
+        return None, ctx.workflow_deadline_epoch_ms

dbos/_dbos.py

@@ -31,9 +31,9 @@ from typing import (
 from opentelemetry.trace import Span
 
 from dbos._conductor.conductor import ConductorWebsocket
+from dbos._sys_db import WorkflowStatus
 from dbos._utils import INTERNAL_QUEUE_NAME, GlobalParams
 from dbos._workflow_commands import (
-    WorkflowStatus,
     fork_workflow,
     list_queued_workflows,
     list_workflows,
@@ -68,7 +68,7 @@ from ._registrations import (
 )
 from ._roles import default_required_roles, required_roles
 from ._scheduler import ScheduledWorkflow, scheduled
-from ._sys_db import StepInfo, reset_system_database
+from ._sys_db import StepInfo, WorkflowStatus, reset_system_database
 from ._tracer import dbos_tracer
 
 if TYPE_CHECKING:
@@ -114,7 +114,7 @@ from ._error import (
 from ._event_loop import BackgroundEventLoop
 from ._logger import add_otlp_to_all_loggers, config_logger, dbos_logger, init_logger
 from ._sys_db import SystemDatabase
-from ._workflow_commands import WorkflowStatus, get_workflow, list_workflow_steps
+from ._workflow_commands import get_workflow, list_workflow_steps
 
 # Most DBOS functions are just any callable F, so decorators / wrappers work on F
 # There are cases where the parameters P and return value R should be separate
@@ -600,7 +600,7 @@ class DBOS:
     # Decorators for DBOS functionality
     @classmethod
     def workflow(
-        cls, *, max_recovery_attempts: int = DEFAULT_MAX_RECOVERY_ATTEMPTS
+        cls, *, max_recovery_attempts: Optional[int] = DEFAULT_MAX_RECOVERY_ATTEMPTS
     ) -> Callable[[Callable[P, R]], Callable[P, R]]:
         """Decorate a function for use as a DBOS workflow."""
         return decorate_workflow(_get_or_create_dbos_registry(), max_recovery_attempts)

dbos/_migrations/versions/83f3732ae8e7_workflow_timeout.py

@@ -0,0 +1,44 @@
+"""workflow_timeout
+
+Revision ID: 83f3732ae8e7
+Revises: f4b9b32ba814
+Create Date: 2025-04-16 17:05:36.642395
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "83f3732ae8e7"
+down_revision: Union[str, None] = "f4b9b32ba814"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.add_column(
+        "workflow_status",
+        sa.Column(
+            "workflow_timeout_ms",
+            sa.BigInteger(),
+            nullable=True,
+        ),
+        schema="dbos",
+    )
+    op.add_column(
+        "workflow_status",
+        sa.Column(
+            "workflow_deadline_epoch_ms",
+            sa.BigInteger(),
+            nullable=True,
+        ),
+        schema="dbos",
+    )
+
+
+def downgrade() -> None:
+    op.drop_column("workflow_status", "workflow_deadline_epoch_ms", schema="dbos")
+    op.drop_column("workflow_status", "workflow_timeout_ms", schema="dbos")

dbos/_registrations.py

@@ -51,7 +51,7 @@ class DBOSFuncInfo:
     class_info: Optional[DBOSClassInfo] = None
     func_type: DBOSFuncType = DBOSFuncType.Unknown
     required_roles: Optional[List[str]] = None
-    max_recovery_attempts: int = DEFAULT_MAX_RECOVERY_ATTEMPTS
+    max_recovery_attempts: Optional[int] = DEFAULT_MAX_RECOVERY_ATTEMPTS
 
 
 def get_or_create_class_info(

dbos/_schemas/system_database.py

@@ -54,7 +54,9 @@ class SystemSchema:
             nullable=True,
             server_default=text("'0'::bigint"),
         ),
-        Column("queue_name", Text),
+        Column("queue_name", Text, nullable=True),
+        Column("workflow_timeout_ms", BigInteger, nullable=True),
+        Column("workflow_deadline_epoch_ms", BigInteger, nullable=True),
         Index("workflow_status_created_at_index", "created_at"),
         Index("workflow_status_executor_id_index", "executor_id"),
     )