dbos 0.16.1__py3-none-any.whl → 0.17.0__py3-none-any.whl

dbos/_core.py

@@ -1,10 +1,27 @@
+import asyncio
+import functools
+import inspect
 import json
 import sys
 import time
 import traceback
 from concurrent.futures import Future
 from functools import wraps
-from typing import TYPE_CHECKING, Any, Callable, Generic, Optional, Tuple, TypeVar, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Coroutine,
+    Generic,
+    Optional,
+    Tuple,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
+
+from dbos._outcome import Immediate, NoResult, Outcome, Pending
 
 from ._app_db import ApplicationDatabase, TransactionResultInternal
 
@@ -181,39 +198,38 @@ def _init_workflow(
     return status
 
 
-def _execute_workflow(
+def _get_wf_invoke_func(
     dbos: "DBOS",
     status: WorkflowStatusInternal,
-    func: "Workflow[P, R]",
-    *args: Any,
-    **kwargs: Any,
-) -> R:
-    try:
-        output = func(*args, **kwargs)
-        status["status"] = "SUCCESS"
-        status["output"] = _serialization.serialize(output)
-        if status["queue_name"] is not None:
-            queue = dbos._registry.queue_info_map[status["queue_name"]]
-            dbos._sys_db.remove_from_queue(status["workflow_uuid"], queue)
-        dbos._sys_db.buffer_workflow_status(status)
-    except DBOSWorkflowConflictIDError:
-        # Retrieve the workflow handle and wait for the result.
-        # Must use existing_workflow=False because workflow status might not be set yet for single transaction workflows.
-        wf_handle: "WorkflowHandle[R]" = dbos.retrieve_workflow(
-            status["workflow_uuid"], existing_workflow=False
-        )
-        output = wf_handle.get_result()
-        return output
-    except Exception as error:
-        status["status"] = "ERROR"
-        status["error"] = _serialization.serialize_exception(error)
-        if status["queue_name"] is not None:
-            queue = dbos._registry.queue_info_map[status["queue_name"]]
-            dbos._sys_db.remove_from_queue(status["workflow_uuid"], queue)
-        dbos._sys_db.update_workflow_status(status)
-        raise
+) -> Callable[[Callable[[], R]], R]:
+    def persist(func: Callable[[], R]) -> R:
+        try:
+            output = func()
+            status["status"] = "SUCCESS"
+            status["output"] = _serialization.serialize(output)
+            if status["queue_name"] is not None:
+                queue = dbos._registry.queue_info_map[status["queue_name"]]
+                dbos._sys_db.remove_from_queue(status["workflow_uuid"], queue)
+            dbos._sys_db.buffer_workflow_status(status)
+            return output
+        except DBOSWorkflowConflictIDError:
+            # Retrieve the workflow handle and wait for the result.
+            # Must use existing_workflow=False because workflow status might not be set yet for single transaction workflows.
+            wf_handle: "WorkflowHandle[R]" = dbos.retrieve_workflow(
+                status["workflow_uuid"], existing_workflow=False
+            )
+            output = wf_handle.get_result()
+            return output
+        except Exception as error:
+            status["status"] = "ERROR"
+            status["error"] = _serialization.serialize_exception(error)
+            if status["queue_name"] is not None:
+                queue = dbos._registry.queue_info_map[status["queue_name"]]
+                dbos._sys_db.remove_from_queue(status["workflow_uuid"], queue)
+            dbos._sys_db.update_workflow_status(status)
+            raise
 
-    return output
+    return persist
 
 
 def _execute_workflow_wthread(
@@ -231,7 +247,15 @@ def _execute_workflow_wthread(
     with DBOSContextSwap(ctx):
         with EnterDBOSWorkflow(attributes):
             try:
-                return _execute_workflow(dbos, status, func, *args, **kwargs)
+                result = (
+                    Outcome[R]
+                    .make(functools.partial(func, *args, **kwargs))
+                    .then(_get_wf_invoke_func(dbos, status))
+                )
+                if isinstance(result, Immediate):
+                    return cast(Immediate[R], result)()
+                else:
+                    return asyncio.run(cast(Pending[R], result)())
             except Exception:
                 dbos.logger.error(
                     f"Exception encountered in asynchronous workflow: {traceback.format_exc()}"
@@ -305,6 +329,18 @@ def execute_workflow_by_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[An
                 )
 
 
+@overload
+def start_workflow(
+    dbos: "DBOS",
+    func: "Workflow[P, Coroutine[Any, Any, R]]",
+    queue_name: Optional[str],
+    execute_workflow: bool,
+    *args: P.args,
+    **kwargs: P.kwargs,
+) -> "WorkflowHandle[R]": ...
+
+
+@overload
 def start_workflow(
     dbos: "DBOS",
     func: "Workflow[P, R]",
@@ -312,6 +348,16 @@ def start_workflow(
     execute_workflow: bool,
     *args: P.args,
     **kwargs: P.kwargs,
+) -> "WorkflowHandle[R]": ...
+
+
+def start_workflow(
+    dbos: "DBOS",
+    func: "Workflow[P, Union[R, Coroutine[Any, Any, R]]]",
+    queue_name: Optional[str],
+    execute_workflow: bool,
+    *args: P.args,
+    **kwargs: P.kwargs,
 ) -> "WorkflowHandle[R]":
     fself: Optional[object] = None
     if hasattr(func, "__self__"):
@@ -396,16 +442,16 @@ def start_workflow(
 
 def workflow_wrapper(
     dbosreg: "DBOSRegistry",
-    func: F,
+    func: Callable[P, R],
     max_recovery_attempts: int = DEFAULT_MAX_RECOVERY_ATTEMPTS,
-) -> F:
+) -> Callable[P, R]:
     func.__orig_func = func  # type: ignore
 
     fi = get_or_create_func_info(func)
     fi.max_recovery_attempts = max_recovery_attempts
 
     @wraps(func)
-    def wrapper(*args: Any, **kwargs: Any) -> Any:
+    def wrapper(*args: Any, **kwargs: Any) -> R:
         if dbosreg.dbos is None:
             raise DBOSException(
                 f"Function {func.__name__} invoked before DBOS initialized"
@@ -425,7 +471,10 @@ def workflow_wrapper(
         enterWorkflowCtxMgr = (
             EnterDBOSChildWorkflow if ctx and ctx.is_workflow() else EnterDBOSWorkflow
         )
-        with enterWorkflowCtxMgr(attributes), DBOSAssumeRole(rr):
+
+        wfOutcome = Outcome[R].make(functools.partial(func, *args, **kwargs))
+
+        def init_wf() -> Callable[[Callable[[], R]], R]:
             ctx = assert_current_dbos_context()  # Now the child ctx
             status = _init_workflow(
                 dbos,
@@ -441,16 +490,23 @@ def workflow_wrapper(
             dbos.logger.debug(
                 f"Running workflow, id: {ctx.workflow_id}, name: {get_dbos_func_name(func)}"
             )
-            return _execute_workflow(dbos, status, func, *args, **kwargs)
 
-    wrapped_func = cast(F, wrapper)
-    return wrapped_func
+            return _get_wf_invoke_func(dbos, status)
+
+        outcome = (
+            wfOutcome.wrap(init_wf)
+            .also(DBOSAssumeRole(rr))
+            .also(enterWorkflowCtxMgr(attributes))
+        )
+        return outcome()  # type: ignore
+
+    return wrapper
 
 
 def decorate_workflow(
     reg: "DBOSRegistry", max_recovery_attempts: int
-) -> Callable[[F], F]:
-    def _workflow_decorator(func: F) -> F:
+) -> 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)
         reg.register_wf_function(func.__qualname__, wrapped_func)
         return wrapped_func
@@ -473,7 +529,8 @@ def decorate_transaction(
                     "name": func.__name__,
                     "operationType": OperationType.TRANSACTION.value,
                 }
-                with EnterDBOSTransaction(session, attributes=attributes) as ctx:
+                with EnterDBOSTransaction(session, attributes=attributes):
+                    ctx = assert_current_dbos_context()
                     txn_output: TransactionResultInternal = {
                         "workflow_uuid": ctx.workflow_id,
                         "function_id": ctx.function_id,
@@ -562,6 +619,11 @@ def decorate_transaction(
                             raise
             return output
 
+        if inspect.iscoroutinefunction(func):
+            raise DBOSException(
+                f"Function {func.__name__} is a coroutine function, but DBOS.transaction does not support coroutine functions"
+            )
+
         fi = get_or_create_func_info(func)
 
         @wraps(func)
@@ -603,8 +665,8 @@ def decorate_step(
     interval_seconds: float = 1.0,
     max_attempts: int = 3,
     backoff_rate: float = 2.0,
-) -> Callable[[F], F]:
-    def decorator(func: F) -> F:
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    def decorator(func: Callable[P, R]) -> Callable[P, R]:
 
         def invoke_step(*args: Any, **kwargs: Any) -> Any:
             if dbosreg.dbos is None:
@@ -617,13 +679,48 @@ def decorate_step(
                 "name": func.__name__,
                 "operationType": OperationType.STEP.value,
             }
-            with EnterDBOSStep(attributes) as ctx:
+
+            attempts = max_attempts if retries_allowed else 1
+            max_retry_interval_seconds: float = 3600  # 1 Hour
+
+            def on_exception(attempt: int, error: BaseException) -> float:
+                dbos.logger.warning(
+                    f"Step being automatically retried. (attempt {attempt} 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,
+                    },
+                )
+                return min(
+                    interval_seconds * (backoff_rate**attempt),
+                    max_retry_interval_seconds,
+                )
+
+            def record_step_result(func: Callable[[], R]) -> R:
+                ctx = assert_current_dbos_context()
                 step_output: OperationResultInternal = {
                     "workflow_uuid": ctx.workflow_id,
                     "function_id": ctx.function_id,
                     "output": None,
                     "error": None,
                 }
+
+                try:
+                    output = func()
+                    step_output["output"] = _serialization.serialize(output)
+                    return output
+                except Exception as error:
+                    step_output["error"] = _serialization.serialize_exception(error)
+                    raise
+                finally:
+                    dbos._sys_db.record_operation_result(step_output)
+
+            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
                 )
@@ -637,57 +734,29 @@ def decorate_step(
                         )
                         raise deserialized_error
                     elif recorded_output["output"] is not None:
-                        return _serialization.deserialize(recorded_output["output"])
+                        return cast(
+                            R, _serialization.deserialize(recorded_output["output"])
+                        )
                     else:
                         raise Exception("Output and error are both None")
                 else:
                     dbos.logger.debug(
                         f"Running step, id: {ctx.function_id}, name: {attributes['name']}"
                     )
+                    return NoResult()
 
-                output = None
-                error = None
-                local_max_attempts = max_attempts if retries_allowed else 1
-                max_retry_interval_seconds: float = 3600  # 1 Hour
-                local_interval_seconds = interval_seconds
-                for attempt in range(1, local_max_attempts + 1):
-                    try:
-                        output = func(*args, **kwargs)
-                        step_output["output"] = _serialization.serialize(output)
-                        error = None
-                        break
-                    except Exception as err:
-                        error = err
-                        if retries_allowed:
-                            dbos.logger.warning(
-                                f"Step being automatically retried. (attempt {attempt} of {local_max_attempts}). {traceback.format_exc()}"
-                            )
-                            ctx.get_current_span().add_event(
-                                f"Step attempt {attempt} failed",
-                                {
-                                    "error": str(error),
-                                    "retryIntervalSeconds": local_interval_seconds,
-                                },
-                            )
-                            if attempt == local_max_attempts:
-                                error = DBOSMaxStepRetriesExceeded()
-                            else:
-                                time.sleep(local_interval_seconds)
-                                local_interval_seconds = min(
-                                    local_interval_seconds * backoff_rate,
-                                    max_retry_interval_seconds,
-                                )
-
-                step_output["error"] = (
-                    _serialization.serialize_exception(error)
-                    if error is not None
-                    else None
+            stepOutcome = Outcome[R].make(functools.partial(func, *args, **kwargs))
+            if retries_allowed:
+                stepOutcome = stepOutcome.retry(
+                    max_attempts, on_exception, lambda i: DBOSMaxStepRetriesExceeded()
                 )
-                dbos._sys_db.record_operation_result(step_output)
 
-                if error is not None:
-                    raise error
-                return output
+            outcome = (
+                stepOutcome.then(record_step_result)
+                .intercept(check_existing_result)
+                .also(EnterDBOSStep(attributes))
+            )
+            return outcome()
 
         fi = get_or_create_func_info(func)
 
@@ -711,16 +780,25 @@ def decorate_step(
                 assert tempwf
                 return tempwf(*args, **kwargs)
 
-        def temp_wf(*args: Any, **kwargs: Any) -> Any:
+        def temp_wf_sync(*args: Any, **kwargs: Any) -> Any:
             return wrapper(*args, **kwargs)
 
+        async def temp_wf_async(*args: Any, **kwargs: Any) -> Any:
+            return await wrapper(*args, **kwargs)
+
+        # Other code in transact-py depends on the name of temporary workflow functions to be "temp_wf"
+        # so set the name of both sync and async temporary workflow functions explicitly
+        temp_wf_sync.__name__ = "temp_wf"
+        temp_wf_async.__name__ = "temp_wf"
+
+        temp_wf = temp_wf_async if inspect.iscoroutinefunction(func) else temp_wf_sync
         wrapped_wf = workflow_wrapper(dbosreg, temp_wf)
         set_dbos_func_name(temp_wf, "<temp>." + func.__qualname__)
         set_temp_workflow_type(temp_wf, "step")
         dbosreg.register_wf_function(get_dbos_func_name(temp_wf), wrapped_wf)
         wrapper.__orig_func = temp_wf  # type: ignore
 
-        return cast(F, wrapper)
+        return cast(Callable[P, R], wrapper)
 
     return decorator
 
@@ -732,7 +810,8 @@ def send(
         attributes: TracedAttributes = {
             "name": "send",
         }
-        with EnterDBOSStep(attributes) as ctx:
+        with EnterDBOSStep(attributes):
+            ctx = assert_current_dbos_context()
             dbos._sys_db.send(
                 ctx.workflow_id,
                 ctx.curr_step_function_id,
@@ -759,7 +838,8 @@ def recv(dbos: "DBOS", topic: Optional[str] = None, timeout_seconds: float = 60)
         attributes: TracedAttributes = {
             "name": "recv",
         }
-        with EnterDBOSStep(attributes) as ctx:
+        with EnterDBOSStep(attributes):
+            ctx = assert_current_dbos_context()
             ctx.function_id += 1  # Reserve for the sleep
             timeout_function_id = ctx.function_id
             return dbos._sys_db.recv(
@@ -784,7 +864,8 @@ def set_event(dbos: "DBOS", key: str, value: Any) -> None:
         attributes: TracedAttributes = {
             "name": "set_event",
         }
-        with EnterDBOSStep(attributes) as ctx:
+        with EnterDBOSStep(attributes):
+            ctx = assert_current_dbos_context()
             dbos._sys_db.set_event(
                 ctx.workflow_id, ctx.curr_step_function_id, key, value
             )
@@ -805,7 +886,8 @@ def get_event(
         attributes: TracedAttributes = {
             "name": "get_event",
         }
-        with EnterDBOSStep(attributes) as ctx:
+        with EnterDBOSStep(attributes):
+            ctx = assert_current_dbos_context()
             ctx.function_id += 1
             timeout_function_id = ctx.function_id
             caller_ctx: GetEventWorkflowContext = {

dbos/_dbos.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import atexit
 import json
 import os
@@ -13,6 +14,7 @@ from typing import (
     TYPE_CHECKING,
     Any,
     Callable,
+    Coroutine,
     Generic,
     List,
     Literal,
@@ -21,6 +23,9 @@ from typing import (
     Tuple,
     Type,
     TypeVar,
+    Union,
+    cast,
+    overload,
 )
 
 from opentelemetry.trace import Span
@@ -40,7 +45,7 @@ from ._core import (
     start_workflow,
     workflow_wrapper,
 )
-from ._queue import Queue, _queue_thread
+from ._queue import Queue, queue_thread
 from ._recovery import recover_pending_workflows, startup_recovery_thread
 from ._registrations import (
     DEFAULT_MAX_RECOVERY_ATTEMPTS,
@@ -71,6 +76,7 @@ else:
 from ._admin_server import AdminServer
 from ._app_db import ApplicationDatabase
 from ._context import (
+    DBOSContext,
     EnterDBOSStep,
     TracedAttributes,
     assert_current_dbos_context,
@@ -277,6 +283,7 @@ class DBOS:
         self.flask: Optional["Flask"] = flask
         self._executor_field: Optional[ThreadPoolExecutor] = None
         self._background_threads: List[threading.Thread] = []
+        self._executor_id: str = os.environ.get("DBOS__VMID", "local")
 
         # If using FastAPI, set up middleware and lifecycle events
         if self.fastapi is not None:
@@ -377,7 +384,7 @@ class DBOS:
             evt = threading.Event()
             self.stop_events.append(evt)
             bg_queue_thread = threading.Thread(
-                target=_queue_thread, args=(evt, self), daemon=True
+                target=queue_thread, args=(evt, self), daemon=True
             )
             bg_queue_thread.start()
             self._background_threads.append(bg_queue_thread)
@@ -432,7 +439,7 @@ class DBOS:
     @classmethod
     def workflow(
         cls, *, max_recovery_attempts: int = DEFAULT_MAX_RECOVERY_ATTEMPTS
-    ) -> Callable[[F], F]:
+    ) -> 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)
 
@@ -457,7 +464,7 @@ class DBOS:
         interval_seconds: float = 1.0,
         max_attempts: int = 3,
         backoff_rate: float = 2.0,
-    ) -> Callable[[F], F]:
+    ) -> Callable[[Callable[P, R]], Callable[P, R]]:
         """
         Decorate and configure a function for use as a DBOS step.
 
@@ -542,15 +549,36 @@ class DBOS:
                 f"{e.name} dependency not found. Please install {e.name} via your package manager."
             ) from e
 
+    @overload
+    @classmethod
+    def start_workflow(
+        cls,
+        func: Workflow[P, Coroutine[Any, Any, R]],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> WorkflowHandle[R]: ...
+
+    @overload
     @classmethod
     def start_workflow(
         cls,
         func: Workflow[P, R],
         *args: P.args,
         **kwargs: P.kwargs,
+    ) -> WorkflowHandle[R]: ...
+
+    @classmethod
+    def start_workflow(
+        cls,
+        func: Workflow[P, Union[R, Coroutine[Any, Any, R]]],
+        *args: P.args,
+        **kwargs: P.kwargs,
     ) -> WorkflowHandle[R]:
         """Invoke a workflow function in the background, returning a handle to the ongoing execution."""
-        return start_workflow(_get_dbos_instance(), func, None, True, *args, **kwargs)
+        return cast(
+            WorkflowHandle[R],
+            start_workflow(_get_dbos_instance(), func, None, True, *args, **kwargs),
+        )
 
     @classmethod
     def get_workflow_status(cls, workflow_id: str) -> Optional[WorkflowStatus]:
@@ -602,6 +630,13 @@ class DBOS:
         """Send a message to a workflow execution."""
         return send(_get_dbos_instance(), destination_id, message, topic)
 
+    @classmethod
+    async def send_async(
+        cls, destination_id: str, message: Any, topic: Optional[str] = None
+    ) -> None:
+        """Send a message to a workflow execution."""
+        await asyncio.to_thread(lambda: DBOS.send(destination_id, message, topic))
+
     @classmethod
     def recv(cls, topic: Optional[str] = None, timeout_seconds: float = 60) -> Any:
         """
@@ -612,13 +647,25 @@ class DBOS:
         """
         return recv(_get_dbos_instance(), topic, timeout_seconds)
 
+    @classmethod
+    async def recv_async(
+        cls, topic: Optional[str] = None, timeout_seconds: float = 60
+    ) -> Any:
+        """
+        Receive a workflow message.
+
+        This function is to be called from within a workflow.
+        `recv_async` will return the message sent on `topic`, asyncronously waiting if necessary.
+        """
+        return await asyncio.to_thread(lambda: DBOS.recv(topic, timeout_seconds))
+
     @classmethod
     def sleep(cls, seconds: float) -> None:
         """
         Sleep for the specified time (in seconds).
 
-        It is important to use `DBOS.sleep` (as opposed to any other sleep) within workflows,
-        as the `DBOS.sleep`s are durable and completed sleeps will be skipped during recovery.
+        It is important to use `DBOS.sleep` or `DBOS.sleep_async` (as opposed to any other sleep) within workflows,
+        as the DBOS sleep methods are durable and completed sleeps will be skipped during recovery.
         """
         if seconds <= 0:
             return
@@ -631,7 +678,8 @@ class DBOS:
             attributes: TracedAttributes = {
                 "name": "sleep",
             }
-            with EnterDBOSStep(attributes) as ctx:
+            with EnterDBOSStep(attributes):
+                ctx = assert_current_dbos_context()
                 _get_dbos_instance()._sys_db.sleep(
                     ctx.workflow_id, ctx.curr_step_function_id, seconds
                 )
@@ -639,17 +687,25 @@ class DBOS:
             # Cannot call it from outside of a workflow
             raise DBOSException("sleep() must be called from within a workflow")
 
+    @classmethod
+    async def sleep_async(cls, seconds: float) -> None:
+        """
+        Sleep for the specified time (in seconds).
+
+        It is important to use `DBOS.sleep` or `DBOS.sleep_async` (as opposed to any other sleep) within workflows,
+        as the DBOS sleep methods are durable and completed sleeps will be skipped during recovery.
+        """
+        await asyncio.to_thread(lambda: DBOS.sleep(seconds))
+
     @classmethod
     def set_event(cls, key: str, value: Any) -> None:
         """
         Set a workflow event.
 
-        This function is to be called from within a workflow.
-
         `set_event` sets the `value` of `key` for the current workflow instance ID.
         This `value` can then be retrieved by other functions, using `get_event` below.
-
-        Each workflow invocation should only call set_event once per `key`.
+        If the event `key` already exists, its `value` is updated.
+        This function can only be called from within a workflow.
 
         Args:
             key(str): The event key / name within the workflow
@@ -658,6 +714,23 @@ class DBOS:
         """
         return set_event(_get_dbos_instance(), key, value)
 
+    @classmethod
+    async def set_event_async(cls, key: str, value: Any) -> None:
+        """
+        Set a workflow event.
+
+        `set_event_async` sets the `value` of `key` for the current workflow instance ID.
+        This `value` can then be retrieved by other functions, using `get_event` below.
+        If the event `key` already exists, its `value` is updated.
+        This function can only be called from within a workflow.
+
+        Args:
+            key(str): The event key / name within the workflow
+            value(Any): A serializable value to associate with the key
+
+        """
+        await asyncio.to_thread(lambda: DBOS.set_event(key, value))
+
     @classmethod
     def get_event(cls, workflow_id: str, key: str, timeout_seconds: float = 60) -> Any:
         """
@@ -673,6 +746,25 @@ class DBOS:
         """
         return get_event(_get_dbos_instance(), workflow_id, key, timeout_seconds)
 
+    @classmethod
+    async def get_event_async(
+        cls, workflow_id: str, key: str, timeout_seconds: float = 60
+    ) -> Any:
+        """
+        Return the `value` of a workflow event, waiting for it to occur if necessary.
+
+        `get_event_async` waits for a corresponding `set_event` by the workflow with ID `workflow_id` with the same `key`.
+
+        Args:
+            workflow_id(str): The workflow instance ID that is expected to call `set_event` on `key`
+            key(str): The event key / name within the workflow
+            timeout_seconds(float): The amount of time to wait, in case `set_event` has not yet been called byt the workflow
+
+        """
+        return await asyncio.to_thread(
+            lambda: DBOS.get_event(workflow_id, key, timeout_seconds)
+        )
+
     @classmethod
     def execute_workflow_id(cls, workflow_id: str) -> WorkflowHandle[Any]:
         """Execute a workflow by ID (for recovery)."""

dbos/_outcome.py

@@ -0,0 +1,206 @@
+import asyncio
+import contextlib
+import inspect
+import time
+from typing import Any, Callable, Coroutine, Optional, Protocol, TypeVar, Union, cast
+
+T = TypeVar("T")
+R = TypeVar("R")
+
+
+class NoResult:
+    _instance: Optional["NoResult"] = None
+    __slots__ = ()
+
+    def __new__(cls, *args: Any, **kwargs: Any) -> "NoResult":
+        if not cls._instance:
+            cls._instance = super(NoResult, cls).__new__(cls, *args, **kwargs)
+        return cls._instance
+
+
+# define Outcome protocol w/ common composition methods
+class Outcome(Protocol[T]):
+
+    def wrap(
+        self, before: Callable[[], Callable[[Callable[[], T]], R]]
+    ) -> "Outcome[R]": ...
+
+    def then(self, next: Callable[[Callable[[], T]], R]) -> "Outcome[R]": ...
+
+    def also(
+        self, cm: contextlib.AbstractContextManager[Any, bool]
+    ) -> "Outcome[T]": ...
+
+    def retry(
+        self,
+        attempts: int,
+        on_exception: Callable[[int, BaseException], float],
+        exceeded_retries: Callable[[int], BaseException],
+    ) -> "Outcome[T]": ...
+
+    def intercept(
+        self, interceptor: Callable[[], Union[NoResult, T]]
+    ) -> "Outcome[T]": ...
+
+    def __call__(self) -> Union[T, Coroutine[Any, Any, T]]: ...
+
+    # Helper function to create an Immediate or Pending Result, depending on if func is a coroutine function or not
+    @staticmethod
+    def make(func: Callable[[], Union[T, Coroutine[Any, Any, T]]]) -> "Outcome[T]":
+        return (
+            Pending(cast(Callable[[], Coroutine[Any, Any, T]], func))
+            if inspect.iscoroutinefunction(func)
+            else Immediate(cast(Callable[[], T], func))
+        )
+
+
+# Immediate Outcome - for composing non-async functions
+class Immediate(Outcome[T]):
+    __slots__ = "_func"
+
+    def __init__(self, func: Callable[[], T]):
+        self._func = func
+
+    def then(self, next: Callable[[Callable[[], T]], R]) -> "Immediate[R]":
+        return Immediate(lambda: next(self._func))
+
+    def wrap(
+        self, before: Callable[[], Callable[[Callable[[], T]], R]]
+    ) -> "Immediate[R]":
+        return Immediate(lambda: before()(self._func))
+
+    @staticmethod
+    def _intercept(
+        func: Callable[[], T], interceptor: Callable[[], Union[NoResult, T]]
+    ) -> T:
+        intercepted = interceptor()
+        return intercepted if not isinstance(intercepted, NoResult) else func()
+
+    def intercept(
+        self, interceptor: Callable[[], Union[NoResult, T]]
+    ) -> "Immediate[T]":
+        return Immediate[T](lambda: Immediate._intercept(self._func, interceptor))
+
+    @staticmethod
+    def _also(func: Callable[[], T], cm: contextlib.AbstractContextManager[Any, bool]) -> T:  # type: ignore
+        with cm:
+            return func()
+
+    def also(self, cm: contextlib.AbstractContextManager[Any, bool]) -> "Immediate[T]":
+        return Immediate[T](lambda: Immediate._also(self._func, cm))
+
+    @staticmethod
+    def _retry(
+        func: Callable[[], T],
+        attempts: int,
+        on_exception: Callable[[int, BaseException], float],
+        exceeded_retries: Callable[[int], BaseException],
+    ) -> T:
+        for i in range(attempts):
+            try:
+                return func()
+            except Exception as exp:
+                wait_time = on_exception(i, exp)
+                time.sleep(wait_time)
+
+        raise exceeded_retries(attempts)
+
+    def retry(
+        self,
+        attempts: int,
+        on_exception: Callable[[int, BaseException], float],
+        exceeded_retries: Callable[[int], BaseException],
+    ) -> "Immediate[T]":
+        assert attempts > 0
+        return Immediate[T](
+            lambda: Immediate._retry(
+                self._func, attempts, on_exception, exceeded_retries
+            )
+        )
+
+    def __call__(self) -> T:
+        return self._func()
+
+
+# Pending Outcome - for composing async functions
+class Pending(Outcome[T]):
+    __slots__ = "_func"
+
+    def __init__(self, func: Callable[[], Coroutine[Any, Any, T]]):
+        self._func = func
+
+    # Helper method in order to raise an exception in a lambda
+    @staticmethod
+    def _raise(ex: BaseException) -> T:
+        raise ex
+
+    async def _wrap(
+        func: Callable[[], Coroutine[Any, Any, T]],
+        before: Callable[[], Callable[[Callable[[], T]], R]],
+    ) -> R:
+        after = await asyncio.to_thread(before)
+        try:
+            value = await func()
+            return await asyncio.to_thread(after, lambda: value)
+        except BaseException as exp:
+            return await asyncio.to_thread(after, lambda: Pending._raise(exp))
+
+    def wrap(
+        self, before: Callable[[], Callable[[Callable[[], T]], R]]
+    ) -> "Pending[R]":
+        return Pending[R](lambda: Pending._wrap(self._func, before))
+
+    def then(self, next: Callable[[Callable[[], T]], R]) -> "Pending[R]":
+        return Pending[R](lambda: Pending._wrap(self._func, lambda: next))
+
+    @staticmethod
+    async def _also(  # type: ignore
+        func: Callable[[], Coroutine[Any, Any, T]],
+        cm: contextlib.AbstractContextManager[Any, bool],
+    ) -> T:
+        with cm:
+            return await func()
+
+    def also(self, cm: contextlib.AbstractContextManager[Any, bool]) -> "Pending[T]":
+        return Pending[T](lambda: Pending._also(self._func, cm))
+
+    @staticmethod
+    async def _intercept(
+        func: Callable[[], Coroutine[Any, Any, T]],
+        interceptor: Callable[[], Union[NoResult, T]],
+    ) -> T:
+        intercepted = await asyncio.to_thread(interceptor)
+        return intercepted if not isinstance(intercepted, NoResult) else await func()
+
+    def intercept(self, interceptor: Callable[[], Union[NoResult, T]]) -> "Pending[T]":
+        return Pending[T](lambda: Pending._intercept(self._func, interceptor))
+
+    @staticmethod
+    async def _retry(
+        func: Callable[[], Coroutine[Any, Any, T]],
+        attempts: int,
+        on_exception: Callable[[int, BaseException], float],
+        exceeded_retries: Callable[[int], BaseException],
+    ) -> T:
+        for i in range(attempts):
+            try:
+                return await func()
+            except Exception as exp:
+                wait_time = on_exception(i, exp)
+                await asyncio.sleep(wait_time)
+
+        raise exceeded_retries(attempts)
+
+    def retry(
+        self,
+        attempts: int,
+        on_exception: Callable[[int, BaseException], float],
+        exceeded_retries: Callable[[int], BaseException],
+    ) -> "Pending[T]":
+        assert attempts > 0
+        return Pending[T](
+            lambda: Pending._retry(self._func, attempts, on_exception, exceeded_retries)
+        )
+
+    async def __call__(self) -> T:
+        return await self._func()

dbos/_queue.py

@@ -51,13 +51,13 @@ class Queue:
         return start_workflow(dbos, func, self.name, False, *args, **kwargs)
 
 
-def _queue_thread(stop_event: threading.Event, dbos: "DBOS") -> None:
+def queue_thread(stop_event: threading.Event, dbos: "DBOS") -> None:
     while not stop_event.is_set():
         if stop_event.wait(timeout=1):
             return
         for _, queue in dbos._registry.queue_info_map.items():
             try:
-                wf_ids = dbos._sys_db.start_queued_workflows(queue)
+                wf_ids = dbos._sys_db.start_queued_workflows(queue, dbos._executor_id)
                 for id in wf_ids:
                     execute_workflow_by_id(dbos, id)
             except Exception:

dbos/_sys_db.py

@@ -1104,7 +1104,7 @@ class SystemDatabase:
                 .on_conflict_do_nothing()
             )
 
-    def start_queued_workflows(self, queue: "Queue") -> List[str]:
+    def start_queued_workflows(self, queue: "Queue", executor_id: str) -> List[str]:
         start_time_ms = int(time.time() * 1000)
         if queue.limiter is not None:
             limiter_period_ms = int(queue.limiter["period"] * 1000)
@@ -1159,7 +1159,7 @@ class SystemDatabase:
                     if len(ret_ids) + num_recent_queries >= queue.limiter["limit"]:
                         break
 
-                # To start a function, first set its status to PENDING
+                # To start a function, first set its status to PENDING and update its executor ID
                 c.execute(
                     SystemSchema.workflow_status.update()
                     .where(SystemSchema.workflow_status.c.workflow_uuid == id)
@@ -1167,7 +1167,10 @@ class SystemDatabase:
                         SystemSchema.workflow_status.c.status
                         == WorkflowStatusString.ENQUEUED.value
                     )
-                    .values(status=WorkflowStatusString.PENDING.value)
+                    .values(
+                        status=WorkflowStatusString.PENDING.value,
+                        executor_id=executor_id,
+                    )
                 )
 
                 # Then give it a start time

{dbos-0.16.1.dist-info → dbos-0.17.0.dist-info}/METADATA

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

{dbos-0.16.1.dist-info → dbos-0.17.0.dist-info}/RECORD

@@ -1,15 +1,15 @@
-dbos-0.16.1.dist-info/METADATA,sha256=DykTz31jHMRMUpYD1WXMx-LKRgiunj6uMjwLo7pcqsY,5020
-dbos-0.16.1.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
-dbos-0.16.1.dist-info/entry_points.txt,sha256=z6GcVANQV7Uw_82H9Ob2axJX6V3imftyZsljdh-M1HU,54
-dbos-0.16.1.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
+dbos-0.17.0.dist-info/METADATA,sha256=jFKo8TAjYuyNqaNqT6YpSBvRzFVuTW8uws5cODC_25Q,5020
+dbos-0.17.0.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
+dbos-0.17.0.dist-info/entry_points.txt,sha256=z6GcVANQV7Uw_82H9Ob2axJX6V3imftyZsljdh-M1HU,54
+dbos-0.17.0.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
 dbos/__init__.py,sha256=CxRHBHEthPL4PZoLbZhp3rdm44-KkRTT2-7DkK9d4QQ,724
 dbos/_admin_server.py,sha256=DOgzVp9kmwiebQqmJB1LcrZnGTxSMbZiGXdenc1wZDg,3163
 dbos/_app_db.py,sha256=_tv2vmPjjiaikwgxH3mqxgJ4nUUcG2-0uMXKWCqVu1c,5509
 dbos/_classproperty.py,sha256=f0X-_BySzn3yFDRKB2JpCbLYQ9tLwt1XftfshvY7CBs,626
 dbos/_context.py,sha256=KV3fd3-Rv6EWrYDUdHARxltSlNZGNtQtNSqeQ-gkXE8,18049
-dbos/_core.py,sha256=HUteo2HP9C4UvB-y41xwUphLZyYRTdb0sW5XqZ6QPAY,31167
+dbos/_core.py,sha256=NWJFQX5bECBvKlYH9pVmNJgmqFGYPnkHnOGjOlOQ3Ag,33504
 dbos/_croniter.py,sha256=hbhgfsHBqclUS8VeLnJ9PSE9Z54z6mi4nnrr1aUXn0k,47561
-dbos/_dbos.py,sha256=iYI-iY3kp_N0lNYYVKwssea7dgD7vsYHn2uBIhmZ2lM,31478
+dbos/_dbos.py,sha256=riYx_dkYFzqeVDYpmcA5ABdAYQFhwyDi4AwxIihDNKA,34809
 dbos/_dbos_config.py,sha256=f37eccN3JpCA32kRdQ4UsERjhYGcdLWv-N21ijnDZmY,6406
 dbos/_error.py,sha256=UETk8CoZL-TO2Utn1-E7OSWelhShWmKM-fOlODMR9PE,3893
 dbos/_fastapi.py,sha256=iyefCZq-ZDKRUjN_rgYQmFmyvWf4gPrSlC6CLbfq4a8,3419
@@ -25,7 +25,8 @@ dbos/_migrations/versions/a3b18ad34abe_added_triggers.py,sha256=Rv0ZsZYZ_WdgGEUL
 dbos/_migrations/versions/d76646551a6b_job_queue_limiter.py,sha256=8PyFi8rd6CN-mUro43wGhsg5wcQWKZPRHD6jw8R5pVc,986
 dbos/_migrations/versions/d76646551a6c_workflow_queue.py,sha256=G942nophZ2uC2vc4hGBC02Ptng1715roTjY3xiyzZU4,729
 dbos/_migrations/versions/eab0cc1d9a14_job_queue.py,sha256=uvhFOtqbBreCePhAxZfIT0qCAI7BiZTou9wt6QnbY7c,1412
-dbos/_queue.py,sha256=hAXwrfBmtv6BGrlmFq-Ol6b_ED-HDaYqSSxumMJC6Xo,1938
+dbos/_outcome.py,sha256=FDMgWVjZ06vm9xO-38H17mTqBImUYQxgKs_bDCSIAhE,6648
+dbos/_queue.py,sha256=5NZ6RfKQd8LQD8EeUXgrwu86r0AadKEqPIMmL_1ORuw,1956
 dbos/_recovery.py,sha256=jbzGYxICA2drzyzlBSy2UiXhKV_16tBVacKQdTkqf-w,2008
 dbos/_registrations.py,sha256=mei6q6_3R5uei8i_Wo_TqGZs85s10shOekDX41sFYD0,6642
 dbos/_request.py,sha256=cX1B3Atlh160phgS35gF1VEEV4pD126c9F3BDgBmxZU,929
@@ -35,7 +36,7 @@ dbos/_schemas/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dbos/_schemas/application_database.py,sha256=KeyoPrF7hy_ODXV7QNike_VFSD74QBRfQ76D7QyE9HI,966
 dbos/_schemas/system_database.py,sha256=7iw7eHJzEvkatHMOaHORoSvtfisF73wW5j8hRt_Ph14,5126
 dbos/_serialization.py,sha256=YCYv0qKAwAZ1djZisBC7khvKqG-5OcIv9t9EC5PFIog,1743
-dbos/_sys_db.py,sha256=NyakKa8lGMJ4WHydwukB8QubyGV9A6feK_bVepknZr4,48887
+dbos/_sys_db.py,sha256=uZKeCnGc2MgvEd0ID3nReBBZj21HzClP56TFkXTvIZE,49028
 dbos/_templates/hello/README.md,sha256=GhxhBj42wjTt1fWEtwNriHbJuKb66Vzu89G4pxNHw2g,930
 dbos/_templates/hello/__package/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dbos/_templates/hello/__package/main.py,sha256=eI0SS9Nwj-fldtiuSzIlIG6dC91GXXwdRsoHxv6S_WI,2719
@@ -51,4 +52,4 @@ dbos/cli.py,sha256=em1uAxrp5yyg53V7ZpmHFtqD6OJp2cMJkG9vGJPoFTA,10904
 dbos/dbos-config.schema.json,sha256=tS7x-bdFbFvpobcs3pIOhwun3yr_ndvTEYOn4BJjTzs,5889
 dbos/py.typed,sha256=QfzXT1Ktfk3Rj84akygc7_42z0lRpCq0Ilh8OXI6Zas,44
 version/__init__.py,sha256=L4sNxecRuqdtSFdpUGX3TtBi9KL3k7YsZVIvv-fv9-A,1678
-dbos-0.16.1.dist-info/RECORD,,
+dbos-0.17.0.dist-info/RECORD,,