dbos 0.24.0a11__py3-none-any.whl → 0.24.0a12__py3-none-any.whl

dbos/_core.py

@@ -58,6 +58,7 @@ from ._error import (
 )
 from ._registrations import (
     DEFAULT_MAX_RECOVERY_ATTEMPTS,
+    DBOSFuncInfo,
     get_config_name,
     get_dbos_class_name,
     get_dbos_func_name,
@@ -82,6 +83,7 @@ if TYPE_CHECKING:
         DBOS,
         Workflow,
         WorkflowHandle,
+        WorkflowHandleAsync,
         WorkflowStatus,
         DBOSRegistry,
         IsolationLevel,
@@ -136,6 +138,48 @@ class WorkflowHandlePolling(Generic[R]):
         return stat
 
 
+class WorkflowHandleAsyncTask(Generic[R]):
+
+    def __init__(self, workflow_id: str, task: asyncio.Task[R], dbos: "DBOS"):
+        self.workflow_id = workflow_id
+        self.task = task
+        self.dbos = dbos
+
+    def get_workflow_id(self) -> str:
+        return self.workflow_id
+
+    async def get_result(self) -> R:
+        return await self.task
+
+    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)
+        return stat
+
+
+class WorkflowHandleAsyncPolling(Generic[R]):
+
+    def __init__(self, workflow_id: str, dbos: "DBOS"):
+        self.workflow_id = workflow_id
+        self.dbos = dbos
+
+    def get_workflow_id(self) -> str:
+        return self.workflow_id
+
+    async def get_result(self) -> R:
+        res: R = await asyncio.to_thread(
+            self.dbos._sys_db.await_workflow_result, self.workflow_id
+        )
+        return res
+
+    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)
+        return stat
+
+
 def _init_workflow(
     dbos: "DBOS",
     ctx: DBOSContext,
@@ -285,6 +329,32 @@ def _execute_workflow_wthread(
                 raise
 
 
+async def _execute_workflow_async(
+    dbos: "DBOS",
+    status: WorkflowStatusInternal,
+    func: "Workflow[P, Coroutine[Any, Any, R]]",
+    ctx: DBOSContext,
+    *args: Any,
+    **kwargs: Any,
+) -> R:
+    attributes: TracedAttributes = {
+        "name": func.__name__,
+        "operationType": OperationType.WORKFLOW.value,
+    }
+    with DBOSContextSwap(ctx):
+        with EnterDBOSWorkflow(attributes):
+            try:
+                result = Pending[R](functools.partial(func, *args, **kwargs)).then(
+                    _get_wf_invoke_func(dbos, status)
+                )
+                return await result()
+            except Exception:
+                dbos.logger.error(
+                    f"Exception encountered in asynchronous workflow: {traceback.format_exc()}"
+                )
+                raise
+
+
 def execute_workflow_by_id(
     dbos: "DBOS", workflow_id: str, startNew: bool = False
 ) -> "WorkflowHandle[Any]":
@@ -349,26 +419,29 @@ def execute_workflow_by_id(
                 )
 
 
-@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]": ...
+def _get_new_wf() -> tuple[str, DBOSContext]:
+    # Sequence of events for starting a workflow:
+    #   First - is there a WF already running?
+    #      (and not in step as that is an error)
+    #   Assign an ID to the workflow, if it doesn't have an app-assigned one
+    #      If this is a root workflow, assign a new ID
+    #      If this is a child workflow, assign parent wf id with call# suffix
+    #   Make a (system) DB record for the workflow
+    #   Pass the new context to a worker thread that will run the wf function
+    cur_ctx = get_local_dbos_context()
+    if cur_ctx is not None and cur_ctx.is_within_workflow():
+        assert cur_ctx.is_workflow()  # Not in a step
+        cur_ctx.function_id += 1
+        if len(cur_ctx.id_assigned_for_next_workflow) == 0:
+            cur_ctx.id_assigned_for_next_workflow = (
+                cur_ctx.workflow_id + "-" + str(cur_ctx.function_id)
+            )
 
+    new_wf_ctx = DBOSContext() if cur_ctx is None else cur_ctx.create_child()
+    new_wf_ctx.id_assigned_for_next_workflow = new_wf_ctx.assign_workflow_id()
+    new_wf_id = new_wf_ctx.id_assigned_for_next_workflow
 
-@overload
-def start_workflow(
-    dbos: "DBOS",
-    func: "Workflow[P, R]",
-    queue_name: Optional[str],
-    execute_workflow: bool,
-    *args: P.args,
-    **kwargs: P.kwargs,
-) -> "WorkflowHandle[R]": ...
+    return (new_wf_id, new_wf_ctx)
 
 
 def start_workflow(
@@ -379,6 +452,7 @@ def start_workflow(
     *args: P.args,
     **kwargs: P.kwargs,
 ) -> "WorkflowHandle[R]":
+
     # If the function has a class, add the class object as its first argument
     fself: Optional[object] = None
     if hasattr(func, "__self__"):
@@ -399,26 +473,7 @@ def start_workflow(
         "kwargs": kwargs,
     }
 
-    # Sequence of events for starting a workflow:
-    #   First - is there a WF already running?
-    #      (and not in step as that is an error)
-    #   Assign an ID to the workflow, if it doesn't have an app-assigned one
-    #      If this is a root workflow, assign a new ID
-    #      If this is a child workflow, assign parent wf id with call# suffix
-    #   Make a (system) DB record for the workflow
-    #   Pass the new context to a worker thread that will run the wf function
-    cur_ctx = get_local_dbos_context()
-    if cur_ctx is not None and cur_ctx.is_within_workflow():
-        assert cur_ctx.is_workflow()  # Not in a step
-        cur_ctx.function_id += 1
-        if len(cur_ctx.id_assigned_for_next_workflow) == 0:
-            cur_ctx.id_assigned_for_next_workflow = (
-                cur_ctx.workflow_id + "-" + str(cur_ctx.function_id)
-            )
-
-    new_wf_ctx = DBOSContext() if cur_ctx is None else cur_ctx.create_child()
-    new_wf_ctx.id_assigned_for_next_workflow = new_wf_ctx.assign_workflow_id()
-    new_wf_id = new_wf_ctx.id_assigned_for_next_workflow
+    new_wf_id, new_wf_ctx = _get_new_wf()
 
     status = _init_workflow(
         dbos,
@@ -458,6 +513,69 @@ def start_workflow(
     return WorkflowHandleFuture(new_wf_id, future, dbos)
 
 
+async def start_workflow_async(
+    dbos: "DBOS",
+    func: "Workflow[P, Coroutine[Any, Any, R]]",
+    queue_name: Optional[str],
+    execute_workflow: bool,
+    *args: P.args,
+    **kwargs: P.kwargs,
+) -> "WorkflowHandleAsync[R]":
+
+    # If the function has a class, add the class object as its first argument
+    fself: Optional[object] = None
+    if hasattr(func, "__self__"):
+        fself = func.__self__
+    if fself is not None:
+        args = (fself,) + args  # type: ignore
+
+    fi = get_func_info(func)
+    if fi is None:
+        raise DBOSWorkflowFunctionNotFoundError(
+            "<NONE>", f"start_workflow: function {func.__name__} is not registered"
+        )
+
+    func = cast("Workflow[P, R]", func.__orig_func)  # type: ignore
+
+    inputs: WorkflowInputs = {
+        "args": args,
+        "kwargs": kwargs,
+    }
+
+    new_wf_id, new_wf_ctx = _get_new_wf()
+
+    status = await asyncio.to_thread(
+        _init_workflow,
+        dbos,
+        new_wf_ctx,
+        inputs=inputs,
+        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,
+        max_recovery_attempts=fi.max_recovery_attempts,
+    )
+
+    wf_status = status["status"]
+
+    if not execute_workflow or (
+        not dbos.debug_mode
+        and (
+            wf_status == WorkflowStatusString.ERROR.value
+            or wf_status == WorkflowStatusString.SUCCESS.value
+        )
+    ):
+        dbos.logger.debug(
+            f"Workflow {new_wf_id} already completed with status {wf_status}. Directly returning a workflow handle."
+        )
+        return WorkflowHandleAsyncPolling(new_wf_id, dbos)
+
+    coro = _execute_workflow_async(dbos, status, func, new_wf_ctx, *args, **kwargs)
+    task = asyncio.create_task(coro)
+    return WorkflowHandleAsyncTask(new_wf_id, task, dbos)
+
+
 if sys.version_info < (3, 12):
 
     def _mark_coroutine(func: Callable[P, R]) -> Callable[P, R]:

dbos/_dbos.py

@@ -49,6 +49,7 @@ from ._core import (
     send,
     set_event,
     start_workflow,
+    start_workflow_async,
     workflow_wrapper,
 )
 from ._queue import Queue, queue_thread
@@ -702,35 +703,26 @@ 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]: ...
+    ) -> 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)
 
     @classmethod
-    def start_workflow(
+    async def start_workflow_async(
         cls,
-        func: Workflow[P, Union[R, Coroutine[Any, Any, R]]],
+        func: Workflow[P, 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 cast(
-            WorkflowHandle[R],
-            start_workflow(_get_dbos_instance(), func, None, True, *args, **kwargs),
+    ) -> WorkflowHandleAsync[R]:
+        """Invoke a workflow function on the event loop, returning a handle to the ongoing execution."""
+        return await start_workflow_async(
+            _get_dbos_instance(), func, None, True, *args, **kwargs
         )
 
     @classmethod
@@ -1107,6 +1099,35 @@ class WorkflowHandle(Generic[R], Protocol):
         ...
 
 
+class WorkflowHandleAsync(Generic[R], Protocol):
+    """
+    Handle to a workflow function.
+
+    `WorkflowHandleAsync` represents a current or previous workflow function invocation,
+    allowing its status and result to be accessed.
+
+    Attributes:
+        workflow_id(str): Workflow ID of the function invocation
+
+    """
+
+    def __init__(self, workflow_id: str) -> None: ...
+
+    workflow_id: str
+
+    def get_workflow_id(self) -> str:
+        """Return the applicable workflow ID."""
+        ...
+
+    async def get_result(self) -> R:
+        """Return the result of the workflow function invocation, waiting if necessary."""
+        ...
+
+    async def get_status(self) -> WorkflowStatus:
+        """Return the current workflow function invocation status as `WorkflowStatus`."""
+        ...
+
+
 class DBOSConfiguredInstance:
     """
     Base class for classes containing DBOS member functions.

dbos/_queue.py

@@ -1,16 +1,16 @@
 import threading
 import traceback
-from typing import TYPE_CHECKING, Optional, TypedDict
+from typing import TYPE_CHECKING, Any, Coroutine, Optional, TypedDict
 
 from psycopg import errors
 from sqlalchemy.exc import OperationalError
 
 from dbos._utils import GlobalParams
 
-from ._core import P, R, execute_workflow_by_id, start_workflow
+from ._core import P, R, execute_workflow_by_id, start_workflow, start_workflow_async
 
 if TYPE_CHECKING:
-    from ._dbos import DBOS, Workflow, WorkflowHandle
+    from ._dbos import DBOS, Workflow, WorkflowHandle, WorkflowHandleAsync
 
 
 class QueueRateLimit(TypedDict):
@@ -66,6 +66,17 @@ class Queue:
         dbos = _get_dbos_instance()
         return start_workflow(dbos, func, self.name, False, *args, **kwargs)
 
+    async def enqueue_async(
+        self,
+        func: "Workflow[P, Coroutine[Any, Any, R]]",
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> "WorkflowHandleAsync[R]":
+        from ._dbos import _get_dbos_instance
+
+        dbos = _get_dbos_instance()
+        return await start_workflow_async(dbos, func, self.name, False, *args, **kwargs)
+
 
 def queue_thread(stop_event: threading.Event, dbos: "DBOS") -> None:
     while not stop_event.is_set():

{dbos-0.24.0a11.dist-info → dbos-0.24.0a12.dist-info}/METADATA

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

{dbos-0.24.0a11.dist-info → dbos-0.24.0a12.dist-info}/RECORD

@@ -1,7 +1,7 @@
-dbos-0.24.0a11.dist-info/METADATA,sha256=Rn228j2NlzClo8eqCmbdGO0Qx4RfKcxEgZ_47DNeekE,5522
-dbos-0.24.0a11.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
-dbos-0.24.0a11.dist-info/entry_points.txt,sha256=_QOQ3tVfEjtjBlr1jS4sHqHya9lI2aIEIWkz8dqYp14,58
-dbos-0.24.0a11.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
+dbos-0.24.0a12.dist-info/METADATA,sha256=bbX82S62nJ9ZMnyppiRzrgMB_NuRFhigHkLeTK7aa_I,5522
+dbos-0.24.0a12.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
+dbos-0.24.0a12.dist-info/entry_points.txt,sha256=_QOQ3tVfEjtjBlr1jS4sHqHya9lI2aIEIWkz8dqYp14,58
+dbos-0.24.0a12.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
 dbos/__init__.py,sha256=uq9LP5uY96kIS9N0yKqlvDwADmtg_Hl30uSUhyuUr-4,754
 dbos/__main__.py,sha256=P7jAr-7L9XE5mrsQ7i4b-bLr2ap1tCQfhMByLCRWDj0,568
 dbos/_admin_server.py,sha256=YiVn5lywz2Vg8_juyNHOYl0HVEy48--7b4phwK7r92o,5732
@@ -13,10 +13,10 @@ dbos/_cloudutils/databases.py,sha256=_shqaqSvhY4n2ScgQ8IP5PDZvzvcx3YBKV8fj-cxhSY
 dbos/_conductor/conductor.py,sha256=5mFrE6yEAHWcFjkYr5p2AoPkV6iwYHxsB3K06yaOmqM,15112
 dbos/_conductor/protocol.py,sha256=W-BnX-bD5folo96NGrUytuCHW72H-OQwPrkJ9q5Au5c,5534
 dbos/_context.py,sha256=Ue5qu3rzLfRmPkz-UUZi9ZS8iXpapRN0NTM4mbA2QmQ,17738
-dbos/_core.py,sha256=8TTIj5shcm5hpKjouMA4VzMWl7R7MJlU-mAHB6xQBxE,37585
+dbos/_core.py,sha256=_a_rSkAWNLoHqzQbkqez0mpctkjDs301123ti3wmKHk,41340
 dbos/_croniter.py,sha256=XHAyUyibs_59sJQfSNWkP7rqQY6_XrlfuuCxk4jYqek,47559
 dbos/_db_wizard.py,sha256=YEW2qoy6hfHQv2fZ_4nHiPUeHMFofPpNTolJ1Kvw7AQ,8394
-dbos/_dbos.py,sha256=AiTuQYoaq9H5qH2dMsIVwNsF645PjBWzj9rjpa9snwo,42824
+dbos/_dbos.py,sha256=ymQnOZ8RQehcPVAjjJipoW8StxM7bktTyT_4a_Zlse8,43599
 dbos/_dbos_config.py,sha256=-DnOQ5gk3I8RPxNkL3GsAkaYvL8jv5umS8Xv5V7UQkw,20217
 dbos/_debug.py,sha256=mmgvLkqlrljMBBow9wk01PPur9kUf2rI_11dTJXY4gw,1822
 dbos/_error.py,sha256=B6Y9XLS1f6yrawxB2uAEYFMxFwk9BHhdxPNddKco-Fw,5399
@@ -35,7 +35,7 @@ dbos/_migrations/versions/d76646551a6b_job_queue_limiter.py,sha256=8PyFi8rd6CN-m
 dbos/_migrations/versions/d76646551a6c_workflow_queue.py,sha256=G942nophZ2uC2vc4hGBC02Ptng1715roTjY3xiyzZU4,729
 dbos/_migrations/versions/eab0cc1d9a14_job_queue.py,sha256=uvhFOtqbBreCePhAxZfIT0qCAI7BiZTou9wt6QnbY7c,1412
 dbos/_outcome.py,sha256=FDMgWVjZ06vm9xO-38H17mTqBImUYQxgKs_bDCSIAhE,6648
-dbos/_queue.py,sha256=I2gBc7zQ4G0vyDDBnKwIFzxtqfD7DxHO2IZ41brFSOM,2927
+dbos/_queue.py,sha256=OWUtbBAqdkDAArFWkwlF8STxykV4iQmrZxrF-_lavh4,3341
 dbos/_recovery.py,sha256=4KyZb0XJEUGH7ekYT1kpx38i6y5vygPeH75Ta7RZjYo,2596
 dbos/_registrations.py,sha256=_zy6k944Ll8QwqU12Kr3OP23ukVtm8axPNN1TS_kJRc,6717
 dbos/_request.py,sha256=cX1B3Atlh160phgS35gF1VEEV4pD126c9F3BDgBmxZU,929
@@ -65,4 +65,4 @@ dbos/cli/cli.py,sha256=pet2vf4GLlSDfxfQbsplM9uewD6pJK2ZpLgZlwgBU5w,15627
 dbos/dbos-config.schema.json,sha256=HtF_njVTGHLdzBGZ4OrGQz3qbPPT0Go-iwd1PgFVTNg,5847
 dbos/py.typed,sha256=QfzXT1Ktfk3Rj84akygc7_42z0lRpCq0Ilh8OXI6Zas,44
 version/__init__.py,sha256=L4sNxecRuqdtSFdpUGX3TtBi9KL3k7YsZVIvv-fv9-A,1678
-dbos-0.24.0a11.dist-info/RECORD,,
+dbos-0.24.0a12.dist-info/RECORD,,