Support Async WF and Step functions (#168)

This PR introduces `Outcome[T]` - a type to allow us to polymorphically
compose decorator behavior for sync and async methods. Updates WF and
Step decorators to use `Outcome[T]` in order to support sync and async
WF/step functions. Adds async versions of DBOS APIs
`send/recv/get_event/set_event/sleep` for use in async workflows.

Note, this PR does *NOT* add support for async transactions. Due to
design of SQLAlchemy's async support, adding async tx function support
is more involved and will be addressed in a future PR.

fixes #112

Author: devhawk

dbos/_core.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import atexit
 import json
 import os
@@ -13,6 +14,7 @@
     TYPE_CHECKING,
     Any,
     Callable,
+    Coroutine,
     Generic,
     List,
     Literal,
@@ -21,6 +23,9 @@
     Tuple,
     Type,
     TypeVar,
+    Union,
+    cast,
+    overload,
 )
 
 from opentelemetry.trace import Span
@@ -71,6 +76,7 @@
 from ._admin_server import AdminServer
 from ._app_db import ApplicationDatabase
 from ._context import (
+    DBOSContext,
     EnterDBOSStep,
     TracedAttributes,
     assert_current_dbos_context,
@@ -432,7 +438,7 @@ def register_instance(cls, inst: object) -> None:
     @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 @@ def step(
         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 @@ def kafka_consumer(
                 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 @@ def send(
         """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 @@ def recv(cls, topic: Optional[str] = None, timeout_seconds: float = 60) -> Any:
         """
         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,25 +677,34 @@ def sleep(cls, seconds: float) -> None:
             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
                 )
         else:
             # 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 @@ def set_event(cls, key: str, value: Any) -> None:
         """
         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 @@ def get_event(cls, workflow_id: str, key: str, timeout_seconds: float = 60) -> A
         """
         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)."""