dbos 0.24.1__py3-none-any.whl → 0.25.0__py3-none-any.whl

dbos/_core.py

@@ -81,13 +81,12 @@ from ._sys_db import (
 if TYPE_CHECKING:
     from ._dbos import (
         DBOS,
-        Workflow,
         WorkflowHandle,
         WorkflowHandleAsync,
-        WorkflowStatus,
         DBOSRegistry,
         IsolationLevel,
     )
+    from ._workflow_commands import WorkflowStatus
 
 from sqlalchemy.exc import DBAPIError, InvalidRequestError
 
@@ -109,7 +108,15 @@ class WorkflowHandleFuture(Generic[R]):
         return self.workflow_id
 
     def get_result(self) -> R:
-        return self.future.result()
+        try:
+            r = self.future.result()
+        except Exception as e:
+            serialized_e = _serialization.serialize_exception(e)
+            self.dbos._sys_db.record_get_result(self.workflow_id, None, serialized_e)
+            raise
+        serialized_r = _serialization.serialize(r)
+        self.dbos._sys_db.record_get_result(self.workflow_id, serialized_r, None)
+        return r
 
     def get_status(self) -> "WorkflowStatus":
         stat = self.dbos.get_workflow_status(self.workflow_id)
@@ -128,8 +135,15 @@ class WorkflowHandlePolling(Generic[R]):
         return self.workflow_id
 
     def get_result(self) -> R:
-        res: R = self.dbos._sys_db.await_workflow_result(self.workflow_id)
-        return res
+        try:
+            r: R = self.dbos._sys_db.await_workflow_result(self.workflow_id)
+        except Exception as e:
+            serialized_e = _serialization.serialize_exception(e)
+            self.dbos._sys_db.record_get_result(self.workflow_id, None, serialized_e)
+            raise
+        serialized_r = _serialization.serialize(r)
+        self.dbos._sys_db.record_get_result(self.workflow_id, serialized_r, None)
+        return r
 
     def get_status(self) -> "WorkflowStatus":
         stat = self.dbos.get_workflow_status(self.workflow_id)
@@ -149,7 +163,22 @@ class WorkflowHandleAsyncTask(Generic[R]):
         return self.workflow_id
 
     async def get_result(self) -> R:
-        return await self.task
+        try:
+            r = await self.task
+        except Exception as e:
+            serialized_e = _serialization.serialize_exception(e)
+            await asyncio.to_thread(
+                self.dbos._sys_db.record_get_result,
+                self.workflow_id,
+                None,
+                serialized_e,
+            )
+            raise
+        serialized_r = _serialization.serialize(r)
+        await asyncio.to_thread(
+            self.dbos._sys_db.record_get_result, self.workflow_id, serialized_r, None
+        )
+        return r
 
     async def get_status(self) -> "WorkflowStatus":
         stat = await asyncio.to_thread(self.dbos.get_workflow_status, self.workflow_id)
@@ -168,10 +197,24 @@ class WorkflowHandleAsyncPolling(Generic[R]):
         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
+        try:
+            r: R = await asyncio.to_thread(
+                self.dbos._sys_db.await_workflow_result, self.workflow_id
+            )
+        except Exception as e:
+            serialized_e = _serialization.serialize_exception(e)
+            await asyncio.to_thread(
+                self.dbos._sys_db.record_get_result,
+                self.workflow_id,
+                None,
+                serialized_e,
+            )
+            raise
+        serialized_r = _serialization.serialize(r)
+        await asyncio.to_thread(
+            self.dbos._sys_db.record_get_result, self.workflow_id, serialized_r, None
         )
-        return res
+        return r
 
     async def get_status(self) -> "WorkflowStatus":
         stat = await asyncio.to_thread(self.dbos.get_workflow_status, self.workflow_id)
@@ -236,22 +279,14 @@ def _init_workflow(
             raise DBOSNonExistentWorkflowError(wfid)
         wf_status = get_status_result["status"]
     else:
-        if temp_wf_type != "transaction" or queue is not None:
-            # Synchronously record the status and inputs for workflows and single-step workflows
-            # We also have to do this for single-step workflows because of the foreign key constraint on the operation outputs table
-            # TODO: Make this transactional (and with the queue step below)
-            wf_status = dbos._sys_db.insert_workflow_status(
-                status, max_recovery_attempts=max_recovery_attempts
-            )
-            # TODO: Modify the inputs if they were changed by `update_workflow_inputs`
-            dbos._sys_db.update_workflow_inputs(
-                wfid, _serialization.serialize_args(inputs)
-            )
-        else:
-            # Buffer the inputs for single-transaction workflows, but don't buffer the status
-            dbos._sys_db.buffer_workflow_inputs(
-                wfid, _serialization.serialize_args(inputs)
-            )
+        # Synchronously record the status and inputs for workflows
+        # TODO: Make this transactional (and with the queue step below)
+        wf_status = dbos._sys_db.insert_workflow_status(
+            status, max_recovery_attempts=max_recovery_attempts
+        )
+
+        # TODO: Modify the inputs if they were changed by `update_workflow_inputs`
+        dbos._sys_db.update_workflow_inputs(wfid, _serialization.serialize_args(inputs))
 
         if queue is not None and wf_status == WorkflowStatusString.ENQUEUED.value:
             dbos._sys_db.enqueue(wfid, queue)
@@ -265,6 +300,18 @@ def _get_wf_invoke_func(
     status: WorkflowStatusInternal,
 ) -> Callable[[Callable[[], R]], R]:
     def persist(func: Callable[[], R]) -> R:
+        if not dbos.debug_mode and (
+            status["status"] == WorkflowStatusString.ERROR.value
+            or status["status"] == WorkflowStatusString.SUCCESS.value
+        ):
+            dbos.logger.debug(
+                f"Workflow {status['workflow_uuid']} is already completed with status {status['status']}"
+            )
+            # Directly return the result if the workflow is already completed
+            recorded_result: R = dbos._sys_db.await_workflow_result(
+                status["workflow_uuid"]
+            )
+            return recorded_result
         try:
             output = func()
             status["status"] = "SUCCESS"
@@ -273,16 +320,12 @@ def _get_wf_invoke_func(
                 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)
+                dbos._sys_db.update_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
+            # Await the workflow result
+            r: R = dbos._sys_db.await_workflow_result(status["workflow_uuid"])
+            return r
         except DBOSWorkflowCancelledError as error:
             raise
         except Exception as error:
@@ -301,7 +344,7 @@ def _get_wf_invoke_func(
 def _execute_workflow_wthread(
     dbos: "DBOS",
     status: WorkflowStatusInternal,
-    func: "Workflow[P, R]",
+    func: "Callable[P, R]",
     ctx: DBOSContext,
     *args: Any,
     **kwargs: Any,
@@ -332,7 +375,7 @@ def _execute_workflow_wthread(
 async def _execute_workflow_async(
     dbos: "DBOS",
     status: WorkflowStatusInternal,
-    func: "Workflow[P, Coroutine[Any, Any, R]]",
+    func: "Callable[P, Coroutine[Any, Any, R]]",
     ctx: DBOSContext,
     *args: Any,
     **kwargs: Any,
@@ -446,7 +489,7 @@ def _get_new_wf() -> tuple[str, DBOSContext]:
 
 def start_workflow(
     dbos: "DBOS",
-    func: "Workflow[P, Union[R, Coroutine[Any, Any, R]]]",
+    func: "Callable[P, Union[R, Coroutine[Any, Any, R]]]",
     queue_name: Optional[str],
     execute_workflow: bool,
     *args: P.args,
@@ -475,6 +518,15 @@ def start_workflow(
 
     new_wf_id, new_wf_ctx = _get_new_wf()
 
+    ctx = new_wf_ctx
+    new_child_workflow_id = ctx.id_assigned_for_next_workflow
+    if ctx.has_parent():
+        child_workflow_id = dbos._sys_db.check_child_workflow(
+            ctx.parent_workflow_id, ctx.parent_workflow_fid
+        )
+        if child_workflow_id is not None:
+            return WorkflowHandlePolling(child_workflow_id, dbos)
+
     status = _init_workflow(
         dbos,
         new_wf_ctx,
@@ -488,6 +540,13 @@ def start_workflow(
     )
 
     wf_status = status["status"]
+    if ctx.has_parent():
+        dbos._sys_db.record_child_workflow(
+            ctx.parent_workflow_id,
+            new_child_workflow_id,
+            ctx.parent_workflow_fid,
+            get_dbos_func_name(func),
+        )
 
     if not execute_workflow or (
         not dbos.debug_mode
@@ -496,9 +555,6 @@ def start_workflow(
             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 WorkflowHandlePolling(new_wf_id, dbos)
 
     future = dbos._executor.submit(
@@ -515,7 +571,7 @@ def start_workflow(
 
 async def start_workflow_async(
     dbos: "DBOS",
-    func: "Workflow[P, Coroutine[Any, Any, R]]",
+    func: "Callable[P, Coroutine[Any, Any, R]]",
     queue_name: Optional[str],
     execute_workflow: bool,
     *args: P.args,
@@ -544,6 +600,17 @@ async def start_workflow_async(
 
     new_wf_id, new_wf_ctx = _get_new_wf()
 
+    ctx = new_wf_ctx
+    new_child_workflow_id = ctx.id_assigned_for_next_workflow
+    if ctx.has_parent():
+        child_workflow_id = await asyncio.to_thread(
+            dbos._sys_db.check_child_workflow,
+            ctx.parent_workflow_id,
+            ctx.parent_workflow_fid,
+        )
+        if child_workflow_id is not None:
+            return WorkflowHandleAsyncPolling(child_workflow_id, dbos)
+
     status = await asyncio.to_thread(
         _init_workflow,
         dbos,
@@ -557,6 +624,15 @@ async def start_workflow_async(
         max_recovery_attempts=fi.max_recovery_attempts,
     )
 
+    if ctx.has_parent():
+        await asyncio.to_thread(
+            dbos._sys_db.record_child_workflow,
+            ctx.parent_workflow_id,
+            new_child_workflow_id,
+            ctx.parent_workflow_fid,
+            get_dbos_func_name(func),
+        )
+
     wf_status = status["status"]
 
     if not execute_workflow or (
@@ -566,9 +642,6 @@ async def start_workflow_async(
             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)
@@ -628,8 +701,30 @@ def workflow_wrapper(
 
         wfOutcome = Outcome[R].make(functools.partial(func, *args, **kwargs))
 
+        workflow_id = None
+
         def init_wf() -> Callable[[Callable[[], R]], R]:
+
+            def recorded_result(
+                c_wfid: str, dbos: "DBOS"
+            ) -> Callable[[Callable[[], R]], R]:
+                def recorded_result_inner(func: Callable[[], R]) -> R:
+                    r: R = dbos._sys_db.await_workflow_result(c_wfid)
+                    return r
+
+                return recorded_result_inner
+
             ctx = assert_current_dbos_context()  # Now the child ctx
+            nonlocal workflow_id
+            workflow_id = ctx.workflow_id
+
+            if ctx.has_parent():
+                child_workflow_id = dbos._sys_db.check_child_workflow(
+                    ctx.parent_workflow_id, ctx.parent_workflow_fid
+                )
+                if child_workflow_id is not None:
+                    return recorded_result(child_workflow_id, dbos)
+
             status = _init_workflow(
                 dbos,
                 ctx,
@@ -640,17 +735,44 @@ def workflow_wrapper(
                 temp_wf_type=get_temp_workflow_type(func),
                 max_recovery_attempts=max_recovery_attempts,
             )
+
             # TODO: maybe modify the parameters if they've been changed by `_init_workflow`
             dbos.logger.debug(
                 f"Running workflow, id: {ctx.workflow_id}, name: {get_dbos_func_name(func)}"
             )
 
+            if ctx.has_parent():
+                dbos._sys_db.record_child_workflow(
+                    ctx.parent_workflow_id,
+                    ctx.workflow_id,
+                    ctx.parent_workflow_fid,
+                    get_dbos_func_name(func),
+                )
+
             return _get_wf_invoke_func(dbos, status)
 
+        def record_get_result(func: Callable[[], R]) -> R:
+            """
+            If a child workflow is invoked synchronously, this records the implicit "getResult" where the
+            parent retrieves the child's output. It executes in the CALLER'S context, not the workflow's.
+            """
+            try:
+                r = func()
+            except Exception as e:
+                serialized_e = _serialization.serialize_exception(e)
+                assert workflow_id is not None
+                dbos._sys_db.record_get_result(workflow_id, None, serialized_e)
+                raise
+            serialized_r = _serialization.serialize(r)
+            assert workflow_id is not None
+            dbos._sys_db.record_get_result(workflow_id, serialized_r, None)
+            return r
+
         outcome = (
             wfOutcome.wrap(init_wf)
             .also(DBOSAssumeRole(rr))
             .also(enterWorkflowCtxMgr(attributes))
+            .then(record_get_result)
         )
         return outcome()  # type: ignore
 
@@ -853,6 +975,8 @@ def decorate_step(
 ) -> Callable[[Callable[P, R]], Callable[P, R]]:
     def decorator(func: Callable[P, R]) -> Callable[P, R]:
 
+        stepName = func.__qualname__
+
         def invoke_step(*args: Any, **kwargs: Any) -> Any:
             if dbosreg.dbos is None:
                 raise DBOSException(
@@ -877,7 +1001,7 @@ def decorate_step(
 
             def on_exception(attempt: int, error: BaseException) -> float:
                 dbos.logger.warning(
-                    f"Step being automatically retried. (attempt {attempt} of {attempts}). {traceback.format_exc()}"
+                    f"Step being automatically retried. (attempt {attempt + 1} of {attempts}). {traceback.format_exc()}"
                 )
                 ctx = assert_current_dbos_context()
                 ctx.get_current_span().add_event(
@@ -897,19 +1021,20 @@ def decorate_step(
                 step_output: OperationResultInternal = {
                     "workflow_uuid": ctx.workflow_id,
                     "function_id": ctx.function_id,
+                    "function_name": stepName,
                     "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)
+                    raise
+                step_output["output"] = _serialization.serialize(output)
+                dbos._sys_db.record_operation_result(step_output)
+                return output
 
             def check_existing_result() -> Union[NoResult, R]:
                 ctx = assert_current_dbos_context()

dbos/_db_wizard.py

@@ -49,6 +49,7 @@ def db_wizard(config: "ConfigFile") -> "ConfigFile":
 
     # 2. If the error is due to password authentication or the configuration is non-default, surface the error and exit.
     error_str = str(db_connection_error)
+    dbos_logger.debug(f"Error connecting to Postgres: {error_str}")
     if (
         "password authentication failed" in error_str
         or "28P01" in error_str
@@ -182,17 +183,12 @@ def _check_db_connectivity(config: "ConfigFile") -> Optional[Exception]:
         host=config["database"]["hostname"],
         port=config["database"]["port"],
         database="postgres",
-        query={"connect_timeout": "1"},
+        query={"connect_timeout": "2"},
     )
     postgres_db_engine = create_engine(postgres_db_url)
     try:
         with postgres_db_engine.connect() as conn:
-            val = conn.execute(text("SELECT 1")).scalar()
-            if val != 1:
-                dbos_logger.error(
-                    f"Unexpected value returned from database: expected 1, received {val}"
-                )
-                return Exception()
+            conn.execute(text("SELECT 1")).scalar()
     except Exception as e:
         return e
     finally: