dbos 0.15.0a2__py3-none-any.whl → 0.16.0a2__py3-none-any.whl

dbos/_context.py

@@ -98,18 +98,27 @@ class DBOSContext:
             wfid = str(uuid.uuid4())
         return wfid
 
-    def start_workflow(self, wfid: Optional[str], attributes: TracedAttributes) -> None:
+    def start_workflow(
+        self,
+        wfid: Optional[str],
+        attributes: TracedAttributes,
+        is_temp_workflow: bool = False,
+    ) -> None:
         if wfid is None or len(wfid) == 0:
             wfid = self.assign_workflow_id()
             self.id_assigned_for_next_workflow = ""
         self.workflow_id = wfid
         self.function_id = 0
-        self._start_span(attributes)
+        if not is_temp_workflow:
+            self._start_span(attributes)
 
-    def end_workflow(self, exc_value: Optional[BaseException]) -> None:
+    def end_workflow(
+        self, exc_value: Optional[BaseException], is_temp_workflow: bool = False
+    ) -> None:
         self.workflow_id = ""
         self.function_id = -1
-        self._end_span(exc_value)
+        if not is_temp_workflow:
+            self._end_span(exc_value)
 
     def is_within_workflow(self) -> bool:
         return len(self.workflow_id) > 0
@@ -349,6 +358,7 @@ 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"
 
     def __enter__(self) -> DBOSContext:
         # Code to create a basic context
@@ -359,7 +369,7 @@ class EnterDBOSWorkflow(AbstractContextManager[DBOSContext, Literal[False]]):
             _set_local_dbos_context(ctx)
         assert not ctx.is_within_workflow()
         ctx.start_workflow(
-            None, self.attributes
+            None, self.attributes, self.is_temp_workflow
         )  # Will get from the context's next workflow ID
         return ctx
 
@@ -371,7 +381,7 @@ class EnterDBOSWorkflow(AbstractContextManager[DBOSContext, Literal[False]]):
     ) -> Literal[False]:
         ctx = assert_current_dbos_context()
         assert ctx.is_within_workflow()
-        ctx.end_workflow(exc_value)
+        ctx.end_workflow(exc_value, self.is_temp_workflow)
         # Code to clean up the basic context if we created it
         if self.created_ctx:
             _clear_local_dbos_context()

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, 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() -> Optional[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 None
 
-                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
@@ -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,
@@ -432,7 +438,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 +463,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 +548,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 +629,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 +646,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 +677,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 +686,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 +713,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 +745,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,190 @@
+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")
+
+
+# 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[[], Optional[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[[], Optional[T]]) -> T:
+        intercepted = interceptor()
+        return intercepted if intercepted else func()
+
+    def intercept(self, interceptor: Callable[[], Optional[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[[], Optional[T]],
+    ) -> T:
+        intercepted = await asyncio.to_thread(interceptor)
+        return intercepted if intercepted else await func()
+
+    def intercept(self, interceptor: Callable[[], Optional[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-0.15.0a2.dist-info → dbos-0.16.0a2.dist-info}/METADATA

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

{dbos-0.15.0a2.dist-info → dbos-0.16.0a2.dist-info}/RECORD

@@ -1,15 +1,15 @@
-dbos-0.15.0a2.dist-info/METADATA,sha256=8nu7jI2g326jK7Yx4cPZ1Z9yCPugQGlAqG-Aj8C6ZPw,5022
-dbos-0.15.0a2.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
-dbos-0.15.0a2.dist-info/entry_points.txt,sha256=z6GcVANQV7Uw_82H9Ob2axJX6V3imftyZsljdh-M1HU,54
-dbos-0.15.0a2.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
+dbos-0.16.0a2.dist-info/METADATA,sha256=kH8V4rGJRqSp0StPXPo-Ftmg-O1ITgB_Vyjn638AABU,5022
+dbos-0.16.0a2.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
+dbos-0.16.0a2.dist-info/entry_points.txt,sha256=z6GcVANQV7Uw_82H9Ob2axJX6V3imftyZsljdh-M1HU,54
+dbos-0.16.0a2.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=SJZTOB-saxCNNGT2LLop9LasqR-X2MQLTay3rjXRUCw,17748
-dbos/_core.py,sha256=HUteo2HP9C4UvB-y41xwUphLZyYRTdb0sW5XqZ6QPAY,31167
+dbos/_context.py,sha256=KV3fd3-Rv6EWrYDUdHARxltSlNZGNtQtNSqeQ-gkXE8,18049
+dbos/_core.py,sha256=F7ep-KA6c0tyrTxV7iwHjKRbxy9RyzCcxl-M8XsvVfE,33481
 dbos/_croniter.py,sha256=hbhgfsHBqclUS8VeLnJ9PSE9Z54z6mi4nnrr1aUXn0k,47561
-dbos/_dbos.py,sha256=iYI-iY3kp_N0lNYYVKwssea7dgD7vsYHn2uBIhmZ2lM,31478
+dbos/_dbos.py,sha256=lBDofWMgK-vVgrU6xmBn1miVv1cf7IamGa9QvPn-BxM,34740
 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,6 +25,7 @@ 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/_outcome.py,sha256=4BxjNfVhjI-M_m8eCWimpvkO8m-knAEyoDscFYq6SQ4,6237
 dbos/_queue.py,sha256=hAXwrfBmtv6BGrlmFq-Ol6b_ED-HDaYqSSxumMJC6Xo,1938
 dbos/_recovery.py,sha256=jbzGYxICA2drzyzlBSy2UiXhKV_16tBVacKQdTkqf-w,2008
 dbos/_registrations.py,sha256=mei6q6_3R5uei8i_Wo_TqGZs85s10shOekDX41sFYD0,6642
@@ -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.15.0a2.dist-info/RECORD,,
+dbos-0.16.0a2.dist-info/RECORD,,