dbos 0.25.0a16__py3-none-any.whl → 0.26.0__py3-none-any.whl

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
@@ -195,8 +200,10 @@ class DBOSContext:
     def end_handler(self, exc_value: Optional[BaseException]) -> None:
         self._end_span(exc_value)
 
-    def get_current_span(self) -> Span:
-        return self.spans[-1]
+    def get_current_span(self) -> Optional[Span]:
+        if len(self.spans):
+            return self.spans[-1]
+        return None
 
     def _start_span(self, attributes: TracedAttributes) -> None:
         attributes["operationUUID"] = (
@@ -358,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
@@ -372,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
@@ -386,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
@@ -52,13 +51,13 @@ from ._error import (
     DBOSMaxStepRetriesExceeded,
     DBOSNonExistentWorkflowError,
     DBOSRecoveryError,
+    DBOSUnexpectedStepError,
     DBOSWorkflowCancelledError,
     DBOSWorkflowConflictIDError,
     DBOSWorkflowFunctionNotFoundError,
 )
 from ._registrations import (
     DEFAULT_MAX_RECOVERY_ATTEMPTS,
-    DBOSFuncInfo,
     get_config_name,
     get_dbos_class_name,
     get_dbos_func_name,
@@ -74,6 +73,7 @@ from ._serialization import WorkflowInputs
 from ._sys_db import (
     GetEventWorkflowContext,
     OperationResultInternal,
+    WorkflowStatus,
     WorkflowStatusInternal,
     WorkflowStatusString,
 )
@@ -86,7 +86,6 @@ if TYPE_CHECKING:
         DBOSRegistry,
         IsolationLevel,
     )
-    from ._workflow_commands import WorkflowStatus
 
 from sqlalchemy.exc import DBAPIError, InvalidRequestError
 
@@ -118,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)
@@ -145,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)
@@ -180,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)
@@ -216,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)
@@ -226,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": (
@@ -266,31 +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:
-        # Synchronously record the status and inputs for workflows
-        # TODO: Make this transactional (and with the queue step below)
-        wf_status = dbos._sys_db.insert_workflow_status(
-            status, 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,
+    )
 
-        # TODO: Modify the inputs if they were changed by `update_workflow_inputs`
-        dbos._sys_db.update_workflow_inputs(wfid, _serialization.serialize_args(inputs))
+    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}"
+                )
 
-        if queue is not None and wf_status == WorkflowStatusString.ENQUEUED.value:
-            dbos._sys_db.enqueue(wfid, queue)
+        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
 
@@ -364,7 +390,9 @@ def _execute_workflow_wthread(
                 if isinstance(result, Immediate):
                     return cast(Immediate[R], result)()
                 else:
-                    return asyncio.run(cast(Pending[R], result)())
+                    return dbos._background_event_loop.submit_coroutine(
+                        cast(Pending[R], result)()
+                    )
             except Exception:
                 dbos.logger.error(
                     f"Exception encountered in asynchronous workflow: {traceback.format_exc()}"
@@ -398,9 +426,7 @@ async def _execute_workflow_async(
                 raise
 
 
-def execute_workflow_by_id(
-    dbos: "DBOS", workflow_id: str, startNew: bool = False
-) -> "WorkflowHandle[Any]":
+def execute_workflow_by_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[Any]":
     status = dbos._sys_db.get_workflow_status(workflow_id)
     if not status:
         raise DBOSRecoveryError(workflow_id, "Workflow status not found")
@@ -441,7 +467,7 @@ def execute_workflow_by_id(
             class_object = dbos._registry.class_info_map[class_name]
             inputs["args"] = (class_object,) + inputs["args"]
 
-        if startNew:
+        with SetWorkflowID(workflow_id):
             return start_workflow(
                 dbos,
                 wf_func,
@@ -450,16 +476,6 @@ def execute_workflow_by_id(
                 *inputs["args"],
                 **inputs["kwargs"],
             )
-        else:
-            with SetWorkflowID(workflow_id):
-                return start_workflow(
-                    dbos,
-                    wf_func,
-                    status["queue_name"],
-                    True,
-                    *inputs["args"],
-                    **inputs["kwargs"],
-                )
 
 
 def _get_new_wf() -> tuple[str, DBOSContext]:
@@ -516,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
@@ -534,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,
     )
 
@@ -598,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
@@ -619,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,
     )
 
@@ -668,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
 
@@ -695,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
         )
@@ -732,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,
             )
 
@@ -780,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)
@@ -794,19 +828,23 @@ def decorate_transaction(
     dbosreg: "DBOSRegistry", isolation_level: "IsolationLevel" = "SERIALIZABLE"
 ) -> Callable[[F], F]:
     def decorator(func: F) -> F:
+
+        transaction_name = func.__qualname__
+
         def invoke_tx(*args: Any, **kwargs: Any) -> Any:
             if dbosreg.dbos is None:
                 raise DBOSException(
                     f"Function {func.__name__} invoked before DBOS initialized"
                 )
 
+            dbos = dbosreg.dbos
             ctx = assert_current_dbos_context()
-            if dbosreg.is_workflow_cancelled(ctx.workflow_id):
+            status = dbos._sys_db.get_workflow_status(ctx.workflow_id)
+            if status and status["status"] == WorkflowStatusString.CANCELLED.value:
                 raise DBOSWorkflowCancelledError(
                     f"Workflow {ctx.workflow_id} is cancelled. Aborting transaction {func.__name__}."
                 )
 
-            dbos = dbosreg.dbos
             with dbos._app_db.sessionmaker() as session:
                 attributes: TracedAttributes = {
                     "name": func.__name__,
@@ -822,17 +860,12 @@ def decorate_transaction(
                         "txn_snapshot": "",  # TODO: add actual snapshot
                         "executor_id": None,
                         "txn_id": None,
+                        "function_name": transaction_name,
                     }
                     retry_wait_seconds = 0.001
                     backoff_factor = 1.5
                     max_retry_wait_seconds = 2.0
                     while True:
-
-                        if dbosreg.is_workflow_cancelled(ctx.workflow_id):
-                            raise DBOSWorkflowCancelledError(
-                                f"Workflow {ctx.workflow_id} is cancelled. Aborting transaction {func.__name__}."
-                            )
-
                         has_recorded_error = False
                         txn_error: Optional[Exception] = None
                         try:
@@ -849,6 +882,7 @@ def decorate_transaction(
                                         session,
                                         ctx.workflow_id,
                                         ctx.function_id,
+                                        transaction_name,
                                     )
                                 )
                                 if dbos.debug_mode and recorded_output is None:
@@ -892,10 +926,12 @@ def decorate_transaction(
                         except DBAPIError as dbapi_error:
                             if dbapi_error.orig.sqlstate == "40001":  # type: ignore
                                 # Retry on serialization failure
-                                ctx.get_current_span().add_event(
-                                    "Transaction Serialization Failure",
-                                    {"retry_wait_seconds": retry_wait_seconds},
-                                )
+                                span = ctx.get_current_span()
+                                if span:
+                                    span.add_event(
+                                        "Transaction Serialization Failure",
+                                        {"retry_wait_seconds": retry_wait_seconds},
+                                    )
                                 time.sleep(retry_wait_seconds)
                                 retry_wait_seconds = min(
                                     retry_wait_seconds * backoff_factor,
@@ -910,6 +946,8 @@ def decorate_transaction(
                             )
                             txn_error = invalid_request_error
                             raise
+                        except DBOSUnexpectedStepError:
+                            raise
                         except Exception as error:
                             txn_error = error
                             raise
@@ -975,7 +1013,7 @@ def decorate_step(
 ) -> Callable[[Callable[P, R]], Callable[P, R]]:
     def decorator(func: Callable[P, R]) -> Callable[P, R]:
 
-        stepName = func.__qualname__
+        step_name = func.__qualname__
 
         def invoke_step(*args: Any, **kwargs: Any) -> Any:
             if dbosreg.dbos is None:
@@ -989,13 +1027,6 @@ def decorate_step(
                 "operationType": OperationType.STEP.value,
             }
 
-            # Check if the workflow is cancelled
-            ctx = assert_current_dbos_context()
-            if dbosreg.is_workflow_cancelled(ctx.workflow_id):
-                raise DBOSWorkflowCancelledError(
-                    f"Workflow {ctx.workflow_id} is cancelled. Aborting step {func.__name__}."
-                )
-
             attempts = max_attempts if retries_allowed else 1
             max_retry_interval_seconds: float = 3600  # 1 Hour
 
@@ -1004,13 +1035,15 @@ def decorate_step(
                     f"Step being automatically retried. (attempt {attempt + 1} of {attempts}). {traceback.format_exc()}"
                 )
                 ctx = assert_current_dbos_context()
-                ctx.get_current_span().add_event(
-                    f"Step attempt {attempt} failed",
-                    {
-                        "error": str(error),
-                        "retryIntervalSeconds": interval_seconds,
-                    },
-                )
+                span = ctx.get_current_span()
+                if span:
+                    span.add_event(
+                        f"Step attempt {attempt} failed",
+                        {
+                            "error": str(error),
+                            "retryIntervalSeconds": interval_seconds,
+                        },
+                    )
                 return min(
                     interval_seconds * (backoff_rate**attempt),
                     max_retry_interval_seconds,
@@ -1021,7 +1054,7 @@ def decorate_step(
                 step_output: OperationResultInternal = {
                     "workflow_uuid": ctx.workflow_id,
                     "function_id": ctx.function_id,
-                    "function_name": stepName,
+                    "function_name": step_name,
                     "output": None,
                     "error": None,
                 }
@@ -1039,7 +1072,7 @@ def decorate_step(
             def check_existing_result() -> Union[NoResult, R]:
                 ctx = assert_current_dbos_context()
                 recorded_output = dbos._sys_db.check_operation_execution(
-                    ctx.workflow_id, ctx.function_id
+                    ctx.workflow_id, ctx.function_id, step_name
                 )
                 if dbos.debug_mode and recorded_output is None:
                     raise DBOSException("Step output not found in debug mode")
@@ -1227,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