dbos 0.23.0a13__py3-none-any.whl → 0.24.0__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]:
@@ -824,7 +942,9 @@ def decorate_step(
             stepOutcome = Outcome[R].make(functools.partial(func, *args, **kwargs))
             if retries_allowed:
                 stepOutcome = stepOutcome.retry(
-                    max_attempts, on_exception, lambda i: DBOSMaxStepRetriesExceeded()
+                    max_attempts,
+                    on_exception,
+                    lambda i: DBOSMaxStepRetriesExceeded(func.__name__, i),
                 )
 
             outcome = (

dbos/_db_wizard.py

@@ -28,7 +28,20 @@ class DatabaseConnection(TypedDict):
     local_suffix: Optional[bool]
 
 
-def db_wizard(config: "ConfigFile", config_file_path: str) -> "ConfigFile":
+def db_wizard(config: "ConfigFile") -> "ConfigFile":
+    """Checks database connectivity and helps the user start a database if needed
+
+    First, check connectivity to the database configured in the provided `config` object.
+    If it fails:
+     - Return an error if the connection failed due to incorrect credentials.
+     - Return an error if it detects a non-default configuration.
+     - Otherwise assume the configured database is not running and guide the user through setting it up.
+
+     The wizard will first attempt to start a local Postgres instance using Docker.
+     If Docker is not available, it will prompt the user to connect to a DBOS Cloud database.
+
+     Finally, if a database was configured, its connection details will be saved in the local `.dbos/db_connection` file.
+    """
     # 1. Check the connectivity to the database. Return if successful. If cannot connect, continue to the following steps.
     db_connection_error = _check_db_connectivity(config)
     if db_connection_error is None:
@@ -44,27 +57,18 @@ def db_wizard(config: "ConfigFile", config_file_path: str) -> "ConfigFile":
         raise DBOSInitializationError(
             f"Could not connect to Postgres: password authentication failed: {db_connection_error}"
         )
-    db_config = config["database"]
-
-    # Read the config file and check if the database hostname/port/username are set. If so, skip the wizard.
-    with open(config_file_path, "r") as file:
-        content = file.read()
-        local_config = yaml.safe_load(content)
-        if "database" not in local_config:
-            local_config["database"] = {}
-        local_config = cast("ConfigFile", local_config)
 
+    # If the database config is not the default one, surface the error and exit.
+    db_config = config["database"]  # FIXME: what if database is not in config?
     if (
-        local_config["database"].get("hostname")
-        or local_config["database"].get("port")
-        or local_config["database"].get("username")
-        or db_config["hostname"] != "localhost"
+        db_config["hostname"] != "localhost"
         or db_config["port"] != 5432
         or db_config["username"] != "postgres"
     ):
         raise DBOSInitializationError(
             f"Could not connect to the database. Exception: {db_connection_error}"
         )
+
     print("[yellow]Postgres not detected locally[/yellow]")
 
     # 3. If the database config is the default one, check if the user has Docker properly installed.

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
@@ -88,13 +89,23 @@ from ._context import (
     assert_current_dbos_context,
     get_local_dbos_context,
 )
-from ._dbos_config import ConfigFile, load_config, set_env_vars
+from ._dbos_config import (
+    ConfigFile,
+    DBOSConfig,
+    check_config_consistency,
+    is_dbos_configfile,
+    load_config,
+    overwrite_config,
+    process_config,
+    set_env_vars,
+    translate_dbos_config_to_config_file,
+)
 from ._error import (
     DBOSConflictingRegistrationError,
     DBOSException,
     DBOSNonExistentWorkflowError,
 )
-from ._logger import add_otlp_to_all_loggers, dbos_logger
+from ._logger import add_otlp_to_all_loggers, config_logger, dbos_logger, init_logger
 from ._sys_db import SystemDatabase
 
 # Most DBOS functions are just any callable F, so decorators / wrappers work on F
@@ -257,7 +268,7 @@ class DBOS:
     def __new__(
         cls: Type[DBOS],
         *,
-        config: Optional[ConfigFile] = None,
+        config: Optional[Union[ConfigFile, DBOSConfig]] = None,
         fastapi: Optional["FastAPI"] = None,
         flask: Optional["Flask"] = None,
         conductor_url: Optional[str] = None,
@@ -302,7 +313,7 @@ class DBOS:
     def __init__(
         self,
         *,
-        config: Optional[ConfigFile] = None,
+        config: Optional[Union[ConfigFile, DBOSConfig]] = None,
         fastapi: Optional["FastAPI"] = None,
         flask: Optional["Flask"] = None,
         conductor_url: Optional[str] = None,
@@ -312,12 +323,7 @@ class DBOS:
             return
 
         self._initialized: bool = True
-        if config is None:
-            config = load_config()
-        set_env_vars(config)
-        dbos_tracer.config(config)
-        dbos_logger.info("Initializing DBOS")
-        self.config: ConfigFile = config
+
         self._launched: bool = False
         self._debug_mode: bool = False
         self._sys_db_field: Optional[SystemDatabase] = None
@@ -334,6 +340,36 @@ class DBOS:
         self.conductor_key: Optional[str] = conductor_key
         self.conductor_websocket: Optional[ConductorWebsocket] = None
 
+        init_logger()
+
+        unvalidated_config: Optional[ConfigFile] = None
+
+        if config is None:
+            # If no config is provided, load it from dbos-config.yaml
+            unvalidated_config = load_config(run_process_config=False)
+        elif is_dbos_configfile(config):
+            unvalidated_config = cast(ConfigFile, config)
+            if os.environ.get("DBOS__CLOUD") == "true":
+                unvalidated_config = overwrite_config(unvalidated_config)
+            check_config_consistency(name=unvalidated_config["name"])
+        else:
+            unvalidated_config = translate_dbos_config_to_config_file(
+                cast(DBOSConfig, config)
+            )
+            if os.environ.get("DBOS__CLOUD") == "true":
+                unvalidated_config = overwrite_config(unvalidated_config)
+            check_config_consistency(name=unvalidated_config["name"])
+
+        if unvalidated_config is not None:
+            self.config: ConfigFile = process_config(data=unvalidated_config)
+        else:
+            raise ValueError("No valid configuration was loaded.")
+
+        set_env_vars(self.config)
+        config_logger(self.config)
+        dbos_tracer.config(self.config)
+        dbos_logger.info("Initializing DBOS")
+
         # If using FastAPI, set up middleware and lifecycle events
         if self.fastapi is not None:
             from ._fastapi import setup_fastapi_middleware
@@ -419,10 +455,17 @@ class DBOS:
             if debug_mode:
                 return
 
-            admin_port = self.config["runtimeConfig"].get("admin_port")
+            admin_port = self.config.get("runtimeConfig", {}).get("admin_port")
             if admin_port is None:
                 admin_port = 3001
-            self._admin_server_field = AdminServer(dbos=self, port=admin_port)
+            run_admin_server = self.config.get("runtimeConfig", {}).get(
+                "run_admin_server"
+            )
+            if run_admin_server:
+                try:
+                    self._admin_server_field = AdminServer(dbos=self, port=admin_port)
+                except Exception as e:
+                    dbos_logger.warning(f"Failed to start admin server: {e}")
 
             workflow_ids = self._sys_db.get_pending_workflows(
                 GlobalParams.executor_id, GlobalParams.app_version
@@ -660,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
@@ -923,7 +957,9 @@ class DBOS:
         reg = _get_or_create_dbos_registry()
         if reg.config is not None:
             return reg.config
-        config = load_config()
+        config = (
+            load_config()
+        )  # This will return the processed & validated config (with defaults)
         reg.config = config
         return config
 
@@ -1063,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.