dbos 0.20.0a9__tar.gz → 0.21.0a4__tar.gz

{dbos-0.20.0a9 → dbos-0.21.0a4}/PKG-INFO

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

{dbos-0.20.0a9 → dbos-0.21.0a4}/dbos/_context.py

@@ -63,7 +63,6 @@ class DBOSContext:
         self.parent_workflow_fid: int = -1
         self.workflow_id: str = ""
         self.function_id: int = -1
-        self.in_recovery: bool = False
 
         self.curr_step_function_id: int = -1
         self.curr_tx_function_id: int = -1
@@ -82,7 +81,6 @@ class DBOSContext:
         rv.is_within_set_workflow_id_block = self.is_within_set_workflow_id_block
         rv.parent_workflow_id = self.workflow_id
         rv.parent_workflow_fid = self.function_id
-        rv.in_recovery = self.in_recovery
         rv.authenticated_user = self.authenticated_user
         rv.authenticated_roles = (
             self.authenticated_roles[:]
@@ -335,34 +333,6 @@ class SetWorkflowID:
         return False  # Did not handle
 
 
-class SetWorkflowRecovery:
-    def __init__(self) -> None:
-        self.created_ctx = False
-
-    def __enter__(self) -> SetWorkflowRecovery:
-        # Code to create a basic context
-        ctx = get_local_dbos_context()
-        if ctx is None:
-            self.created_ctx = True
-            _set_local_dbos_context(DBOSContext())
-        assert_current_dbos_context().in_recovery = True
-
-        return self
-
-    def __exit__(
-        self,
-        exc_type: Optional[Type[BaseException]],
-        exc_value: Optional[BaseException],
-        traceback: Optional[TracebackType],
-    ) -> Literal[False]:
-        assert assert_current_dbos_context().in_recovery == True
-        assert_current_dbos_context().in_recovery = False
-        # Code to clean up the basic context if we created it
-        if self.created_ctx:
-            _clear_local_dbos_context()
-        return False  # Did not handle
-
-
 class EnterDBOSWorkflow(AbstractContextManager[DBOSContext, Literal[False]]):
     def __init__(self, attributes: TracedAttributes) -> None:
         self.created_ctx = False

{dbos-0.20.0a9 → dbos-0.21.0a4}/dbos/_core.py

@@ -185,8 +185,8 @@ def _init_workflow(
         # 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.update_workflow_status(
-            status, False, ctx.in_recovery, max_recovery_attempts=max_recovery_attempts
+        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))

{dbos-0.20.0a9 → dbos-0.21.0a4}/dbos/_dbos.py

@@ -800,14 +800,13 @@ class DBOS:
     @classmethod
     def cancel_workflow(cls, workflow_id: str) -> None:
         """Cancel a workflow by ID."""
-        _get_dbos_instance()._sys_db.set_workflow_status(
-            workflow_id, WorkflowStatusString.CANCELLED, False
-        )
+        _get_dbos_instance()._sys_db.cancel_workflow(workflow_id)
 
     @classmethod
-    def resume_workflow(cls, workflow_id: str) -> None:
+    def resume_workflow(cls, workflow_id: str) -> WorkflowHandle[Any]:
         """Resume a workflow by ID."""
-        execute_workflow_by_id(_get_dbos_instance(), workflow_id, False)
+        _get_dbos_instance()._sys_db.resume_workflow(workflow_id)
+        return execute_workflow_by_id(_get_dbos_instance(), workflow_id, False)
 
     @classproperty
     def logger(cls) -> Logger:

{dbos-0.20.0a9 → dbos-0.21.0a4}/dbos/_recovery.py

@@ -4,7 +4,6 @@ import time
 import traceback
 from typing import TYPE_CHECKING, Any, List
 
-from ._context import SetWorkflowRecovery
 from ._core import execute_workflow_by_id
 from ._error import DBOSWorkflowFunctionNotFoundError
 
@@ -19,8 +18,7 @@ def startup_recovery_thread(dbos: "DBOS", workflow_ids: List[str]) -> None:
     while not stop_event.is_set() and len(workflow_ids) > 0:
         try:
             for workflowID in list(workflow_ids):
-                with SetWorkflowRecovery():
-                    execute_workflow_by_id(dbos, workflowID)
+                execute_workflow_by_id(dbos, workflowID)
                 workflow_ids.remove(workflowID)
         except DBOSWorkflowFunctionNotFoundError:
             time.sleep(1)
@@ -45,8 +43,7 @@ def recover_pending_workflows(
         dbos.logger.debug(f"Pending workflows: {workflow_ids}")
 
         for workflowID in workflow_ids:
-            with SetWorkflowRecovery():
-                handle = execute_workflow_by_id(dbos, workflowID)
+            handle = execute_workflow_by_id(dbos, workflowID)
             workflow_handles.append(handle)
 
     dbos.logger.info("Recovered pending workflows")

{dbos-0.20.0a9 → dbos-0.21.0a4}/dbos/_sys_db.py

@@ -243,66 +243,50 @@ class SystemDatabase:
             dbos_logger.debug("Waiting for system buffers to be exported")
             time.sleep(1)
 
-    def update_workflow_status(
+    def insert_workflow_status(
         self,
         status: WorkflowStatusInternal,
-        replace: bool = True,
-        in_recovery: bool = False,
         *,
-        conn: Optional[sa.Connection] = None,
         max_recovery_attempts: int = DEFAULT_MAX_RECOVERY_ATTEMPTS,
     ) -> WorkflowStatuses:
         wf_status: WorkflowStatuses = status["status"]
 
-        cmd = pg.insert(SystemSchema.workflow_status).values(
-            workflow_uuid=status["workflow_uuid"],
-            status=status["status"],
-            name=status["name"],
-            class_name=status["class_name"],
-            config_name=status["config_name"],
-            output=status["output"],
-            error=status["error"],
-            executor_id=status["executor_id"],
-            application_version=status["app_version"],
-            application_id=status["app_id"],
-            request=status["request"],
-            authenticated_user=status["authenticated_user"],
-            authenticated_roles=status["authenticated_roles"],
-            assumed_role=status["assumed_role"],
-            queue_name=status["queue_name"],
-        )
-        if replace:
-            cmd = cmd.on_conflict_do_update(
-                index_elements=["workflow_uuid"],
-                set_=dict(
-                    status=status["status"],
-                    output=status["output"],
-                    error=status["error"],
-                ),
-            )
-        elif in_recovery:
-            cmd = cmd.on_conflict_do_update(
-                index_elements=["workflow_uuid"],
-                set_=dict(
-                    recovery_attempts=SystemSchema.workflow_status.c.recovery_attempts
-                    + 1,
+        cmd = (
+            pg.insert(SystemSchema.workflow_status)
+            .values(
+                workflow_uuid=status["workflow_uuid"],
+                status=status["status"],
+                name=status["name"],
+                class_name=status["class_name"],
+                config_name=status["config_name"],
+                output=status["output"],
+                error=status["error"],
+                executor_id=status["executor_id"],
+                application_version=status["app_version"],
+                application_id=status["app_id"],
+                request=status["request"],
+                authenticated_user=status["authenticated_user"],
+                authenticated_roles=status["authenticated_roles"],
+                assumed_role=status["assumed_role"],
+                queue_name=status["queue_name"],
+                recovery_attempts=(
+                    1 if wf_status != WorkflowStatusString.ENQUEUED.value else 0
                 ),
             )
-        else:
-            # A blank update so that we can return the existing status
-            cmd = cmd.on_conflict_do_update(
+            .on_conflict_do_update(
                 index_elements=["workflow_uuid"],
                 set_=dict(
-                    recovery_attempts=SystemSchema.workflow_status.c.recovery_attempts
+                    recovery_attempts=(
+                        SystemSchema.workflow_status.c.recovery_attempts + 1
+                    ),
                 ),
             )
+        )
+
         cmd = cmd.returning(SystemSchema.workflow_status.c.recovery_attempts, SystemSchema.workflow_status.c.status, SystemSchema.workflow_status.c.name, SystemSchema.workflow_status.c.class_name, SystemSchema.workflow_status.c.config_name, SystemSchema.workflow_status.c.queue_name)  # type: ignore
 
-        if conn is not None:
-            results = conn.execute(cmd)
-        else:
-            with self.engine.begin() as c:
-                results = c.execute(cmd)
+        with self.engine.begin() as c:
+            results = c.execute(cmd)
 
         row = results.fetchone()
         if row is not None:
@@ -325,7 +309,9 @@ class SystemDatabase:
             if err_msg is not None:
                 raise DBOSConflictingWorkflowError(status["workflow_uuid"], err_msg)
 
-            if in_recovery and recovery_attempts > max_recovery_attempts:
+            # Every time we start executing a workflow (and thus attempt to insert its status), we increment `recovery_attempts` by 1.
+            # When this number becomes equal to `maxRetries + 1`, we mark the workflow as `RETRIES_EXCEEDED`.
+            if recovery_attempts > max_recovery_attempts + 1:
                 with self.engine.begin() as c:
                     c.execute(
                         sa.delete(SystemSchema.workflow_queue).where(
@@ -352,38 +338,107 @@ class SystemDatabase:
                     status["workflow_uuid"], max_recovery_attempts
                 )
 
-        # Record we have exported status for this single-transaction workflow
+        return wf_status
+
+    def update_workflow_status(
+        self,
+        status: WorkflowStatusInternal,
+        *,
+        conn: Optional[sa.Connection] = None,
+    ) -> None:
+        wf_status: WorkflowStatuses = status["status"]
+
+        cmd = (
+            pg.insert(SystemSchema.workflow_status)
+            .values(
+                workflow_uuid=status["workflow_uuid"],
+                status=status["status"],
+                name=status["name"],
+                class_name=status["class_name"],
+                config_name=status["config_name"],
+                output=status["output"],
+                error=status["error"],
+                executor_id=status["executor_id"],
+                application_version=status["app_version"],
+                application_id=status["app_id"],
+                request=status["request"],
+                authenticated_user=status["authenticated_user"],
+                authenticated_roles=status["authenticated_roles"],
+                assumed_role=status["assumed_role"],
+                queue_name=status["queue_name"],
+                recovery_attempts=(
+                    1 if wf_status != WorkflowStatusString.ENQUEUED.value else 0
+                ),
+            )
+            .on_conflict_do_update(
+                index_elements=["workflow_uuid"],
+                set_=dict(
+                    status=status["status"],
+                    output=status["output"],
+                    error=status["error"],
+                ),
+            )
+        )
+
+        if conn is not None:
+            conn.execute(cmd)
+        else:
+            with self.engine.begin() as c:
+                c.execute(cmd)
+
+        # If this is a single-transaction workflow, record that its status has been exported
         if status["workflow_uuid"] in self._temp_txn_wf_ids:
             self._exported_temp_txn_wf_status.add(status["workflow_uuid"])
 
-        return wf_status
-
-    def set_workflow_status(
+    def cancel_workflow(
         self,
-        workflow_uuid: str,
-        status: WorkflowStatusString,
-        reset_recovery_attempts: bool,
+        workflow_id: str,
     ) -> None:
         with self.engine.begin() as c:
-            stmt = (
+            # Remove the workflow from the queues table so it does not block the table
+            c.execute(
+                sa.delete(SystemSchema.workflow_queue).where(
+                    SystemSchema.workflow_queue.c.workflow_uuid == workflow_id
+                )
+            )
+            # Set the workflow's status to CANCELLED
+            c.execute(
                 sa.update(SystemSchema.workflow_status)
-                .where(SystemSchema.workflow_status.c.workflow_uuid == workflow_uuid)
+                .where(SystemSchema.workflow_status.c.workflow_uuid == workflow_id)
                 .values(
-                    status=status,
+                    status=WorkflowStatusString.CANCELLED.value,
                 )
             )
-            c.execute(stmt)
 
-        if reset_recovery_attempts:
-            with self.engine.begin() as c:
-                stmt = (
-                    sa.update(SystemSchema.workflow_status)
-                    .where(
-                        SystemSchema.workflow_status.c.workflow_uuid == workflow_uuid
-                    )
-                    .values(recovery_attempts=reset_recovery_attempts)
+    def resume_workflow(
+        self,
+        workflow_id: str,
+    ) -> None:
+        with self.engine.begin() as c:
+            # Check the status of the workflow. If it is complete, do nothing.
+            row = c.execute(
+                sa.select(
+                    SystemSchema.workflow_status.c.status,
+                ).where(SystemSchema.workflow_status.c.workflow_uuid == workflow_id)
+            ).fetchone()
+            if (
+                row is None
+                or row[0] == WorkflowStatusString.SUCCESS.value
+                or row[0] == WorkflowStatusString.ERROR.value
+            ):
+                return
+            # Remove the workflow from the queues table so resume can safely be called on an ENQUEUED workflow
+            c.execute(
+                sa.delete(SystemSchema.workflow_queue).where(
+                    SystemSchema.workflow_queue.c.workflow_uuid == workflow_id
                 )
-                c.execute(stmt)
+            )
+            # Set the workflow's status to PENDING and clear its recovery attempts.
+            c.execute(
+                sa.update(SystemSchema.workflow_status)
+                .where(SystemSchema.workflow_status.c.workflow_uuid == workflow_id)
+                .values(status=WorkflowStatusString.PENDING.value, recovery_attempts=0)
+            )
 
     def get_workflow_status(
         self, workflow_uuid: str

{dbos-0.20.0a9 → dbos-0.21.0a4}/dbos/_workflow_commands.py

@@ -1,23 +1,14 @@
-import importlib
-import os
-import sys
-from typing import Any, List, Optional, cast
+from typing import List, Optional, cast
 
 import typer
-from rich import print
 
-from dbos import DBOS
-
-from . import _serialization, load_config
-from ._core import execute_workflow_by_id
-from ._dbos_config import ConfigFile, _is_valid_app_name
+from . import _serialization
+from ._dbos_config import ConfigFile
 from ._sys_db import (
     GetWorkflowsInput,
     GetWorkflowsOutput,
     SystemDatabase,
     WorkflowStatuses,
-    WorkflowStatusInternal,
-    WorkflowStatusString,
 )
 
 
@@ -41,7 +32,7 @@ class WorkflowInformation:
     queue_name: Optional[str]
 
 
-def _list_workflows(
+def list_workflows(
     config: ConfigFile,
     li: int,
     user: Optional[str],
@@ -91,17 +82,13 @@ def _list_workflows(
             sys_db.destroy()
 
 
-def _get_workflow(
+def get_workflow(
     config: ConfigFile, uuid: str, request: bool
 ) -> Optional[WorkflowInformation]:
-    sys_db = None
-
     try:
         sys_db = SystemDatabase(config)
-
         info = _get_workflow_info(sys_db, uuid, request)
         return info
-
     except Exception as e:
         typer.echo(f"Error getting workflow: {e}")
         return None
@@ -110,18 +97,13 @@ def _get_workflow(
             sys_db.destroy()
 
 
-def _cancel_workflow(config: ConfigFile, uuid: str) -> None:
-    # config = load_config()
-    sys_db = None
-
+def cancel_workflow(config: ConfigFile, uuid: str) -> None:
     try:
         sys_db = SystemDatabase(config)
-        sys_db.set_workflow_status(uuid, WorkflowStatusString.CANCELLED, False)
-        return
-
+        sys_db.cancel_workflow(uuid)
     except Exception as e:
         typer.echo(f"Failed to connect to DBOS system database: {e}")
-        return None
+        raise e
     finally:
         if sys_db:
             sys_db.destroy()

{dbos-0.20.0a9 → dbos-0.21.0a4}/dbos/cli/cli.py

@@ -19,7 +19,7 @@ from .. import load_config
 from .._app_db import ApplicationDatabase
 from .._dbos_config import _is_valid_app_name
 from .._sys_db import SystemDatabase, reset_system_database
-from .._workflow_commands import _cancel_workflow, _get_workflow, _list_workflows
+from .._workflow_commands import cancel_workflow, get_workflow, list_workflows
 from ..cli._github_init import create_template_from_github
 from ._template_init import copy_template, get_project_name, get_templates_directory
 
@@ -282,7 +282,7 @@ def list(
     ] = None,
 ) -> None:
     config = load_config()
-    workflows = _list_workflows(
+    workflows = list_workflows(
         config, limit, user, starttime, endtime, status, request, appversion
     )
     print(jsonpickle.encode(workflows, unpicklable=False))
@@ -301,7 +301,7 @@ def get(
     ] = True,
 ) -> None:
     config = load_config()
-    print(jsonpickle.encode(_get_workflow(config, uuid, request), unpicklable=False))
+    print(jsonpickle.encode(get_workflow(config, uuid, request), unpicklable=False))
 
 
 @workflow.command(
@@ -315,7 +315,7 @@ def cancel(
     ] = None,
 ) -> None:
     config = load_config()
-    _cancel_workflow(config, uuid)
+    cancel_workflow(config, uuid)
     print(f"Workflow {uuid} has been cancelled")
 
 

{dbos-0.20.0a9 → dbos-0.21.0a4}/pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
 ]
 requires-python = ">=3.9"
 readme = "README.md"
-version = "0.20.0a9"
+version = "0.21.0a4"
 
 [project.license]
 text = "MIT"

{dbos-0.20.0a9 → dbos-0.21.0a4}/tests/test_admin_server.py

@@ -151,54 +151,59 @@ runtimeConfig:
 
 
 def test_admin_workflow_resume(dbos: DBOS, config: ConfigFile) -> None:
+    counter: int = 0
 
     @DBOS.workflow()
     def simple_workflow() -> None:
-        print("Executed Simple workflow")
-        return
+        nonlocal counter
+        counter += 1
 
-    # run the workflow
+    # Run the workflow and flush its results
     simple_workflow()
-    time.sleep(1)
+    assert counter == 1
+    dbos._sys_db.wait_for_buffer_flush()
 
-    # get the workflow list
-    output = _workflow_commands._list_workflows(
+    # Verify the workflow has succeeded
+    output = _workflow_commands.list_workflows(
         config, 10, None, None, None, None, False, None
     )
     assert len(output) == 1, f"Expected list length to be 1, but got {len(output)}"
-
     assert output[0] != None, "Expected output to be not None"
-
     wfUuid = output[0].workflowUUID
-
-    info = _workflow_commands._get_workflow(config, wfUuid, True)
+    info = _workflow_commands.get_workflow(config, wfUuid, True)
     assert info is not None, "Expected output to be not None"
-
     assert info.status == "SUCCESS", f"Expected status to be SUCCESS"
 
+    # Cancel the workflow. Verify it was cancelled
     response = requests.post(
         f"http://localhost:3001/workflows/{wfUuid}/cancel", json=[], timeout=5
     )
     assert response.status_code == 204
+    info = _workflow_commands.get_workflow(config, wfUuid, True)
+    assert info is not None
+    assert info.status == "CANCELLED", f"Expected status to be CANCELLED"
 
-    info = _workflow_commands._get_workflow(config, wfUuid, True)
-    if info is not None:
-        assert info.status == "CANCELLED", f"Expected status to be CANCELLED"
-    else:
-        assert False, "Expected info to be not None"
-
+    # Resume the workflow. Verify that it succeeds again.
     response = requests.post(
         f"http://localhost:3001/workflows/{wfUuid}/resume", json=[], timeout=5
     )
     assert response.status_code == 204
+    dbos._sys_db.wait_for_buffer_flush()
+    assert counter == 2
+    info = _workflow_commands.get_workflow(config, wfUuid, True)
+    assert info is not None
+    assert info.status == "SUCCESS", f"Expected status to be SUCCESS"
 
-    time.sleep(1)
-
-    info = _workflow_commands._get_workflow(config, wfUuid, True)
-    if info is not None:
-        assert info.status == "SUCCESS", f"Expected status to be SUCCESS"
-    else:
-        assert False, "Expected info to be not None"
+    # Resume the workflow. Verify it does not run and status remains SUCCESS
+    response = requests.post(
+        f"http://localhost:3001/workflows/{wfUuid}/resume", json=[], timeout=5
+    )
+    assert response.status_code == 204
+    dbos._sys_db.wait_for_buffer_flush()
+    info = _workflow_commands.get_workflow(config, wfUuid, True)
+    assert info is not None
+    assert info.status == "SUCCESS", f"Expected status to be SUCCESS"
+    assert counter == 2
 
 
 def test_admin_workflow_restart(dbos: DBOS, config: ConfigFile) -> None:
@@ -213,7 +218,7 @@ def test_admin_workflow_restart(dbos: DBOS, config: ConfigFile) -> None:
     time.sleep(1)
 
     # get the workflow list
-    output = _workflow_commands._list_workflows(
+    output = _workflow_commands.list_workflows(
         config, 10, None, None, None, None, False, None
     )
     assert len(output) == 1, f"Expected list length to be 1, but got {len(output)}"
@@ -222,7 +227,7 @@ def test_admin_workflow_restart(dbos: DBOS, config: ConfigFile) -> None:
 
     wfUuid = output[0].workflowUUID
 
-    info = _workflow_commands._get_workflow(config, wfUuid, True)
+    info = _workflow_commands.get_workflow(config, wfUuid, True)
     assert info is not None, "Expected output to be not None"
 
     assert info.status == "SUCCESS", f"Expected status to be SUCCESS"
@@ -232,7 +237,7 @@ def test_admin_workflow_restart(dbos: DBOS, config: ConfigFile) -> None:
     )
     assert response.status_code == 204
 
-    info = _workflow_commands._get_workflow(config, wfUuid, True)
+    info = _workflow_commands.get_workflow(config, wfUuid, True)
     if info is not None:
         assert info.status == "CANCELLED", f"Expected status to be CANCELLED"
     else:
@@ -245,13 +250,13 @@ def test_admin_workflow_restart(dbos: DBOS, config: ConfigFile) -> None:
 
     time.sleep(1)
 
-    info = _workflow_commands._get_workflow(config, wfUuid, True)
+    info = _workflow_commands.get_workflow(config, wfUuid, True)
     if info is not None:
         assert info.status == "CANCELLED", f"Expected status to be CANCELLED"
     else:
         assert False, "Expected info to be not None"
 
-    output = _workflow_commands._list_workflows(
+    output = _workflow_commands.list_workflows(
         config, 10, None, None, None, None, False, None
     )
     assert len(output) == 2, f"Expected list length to be 2, but got {len(output)}"
@@ -261,7 +266,7 @@ def test_admin_workflow_restart(dbos: DBOS, config: ConfigFile) -> None:
     else:
         new_wfUuid = output[0].workflowUUID
 
-    info = _workflow_commands._get_workflow(config, new_wfUuid, True)
+    info = _workflow_commands.get_workflow(config, new_wfUuid, True)
     if info is not None:
         assert info.status == "SUCCESS", f"Expected status to be SUCCESS"
     else: