dbos 0.26.0a1__tar.gz → 0.26.0a5__tar.gz

{dbos-0.26.0a1 → dbos-0.26.0a5}/PKG-INFO

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

{dbos-0.26.0a1 → dbos-0.26.0a5}/dbos/_admin_server.py

@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, List, TypedDict
 
 from ._logger import dbos_logger
 from ._recovery import recover_pending_workflows
+from ._utils import GlobalParams
 
 if TYPE_CHECKING:
     from ._dbos import DBOS
@@ -44,6 +45,7 @@ class AdminServer:
 class AdminRequestHandler(BaseHTTPRequestHandler):
     def __init__(self, dbos: DBOS, *args: Any, **kwargs: Any) -> None:
         self.dbos = dbos
+        self.is_deactivated = False
         super().__init__(*args, **kwargs)
 
     def _end_headers(self) -> None:
@@ -59,10 +61,14 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
             self._end_headers()
             self.wfile.write("healthy".encode("utf-8"))
         elif self.path == _deactivate_path:
+            if not self.is_deactivated:
+                dbos_logger.info(
+                    f"Deactivating DBOS executor {GlobalParams.executor_id} with version {GlobalParams.app_version}. This executor will complete existing workflows but will not start new workflows."
+                )
+                self.is_deactivated = True
             # Stop all scheduled workflows, queues, and kafka loops
             for event in self.dbos.stop_events:
                 event.set()
-
             self.send_response(200)
             self._end_headers()
             self.wfile.write("deactivated".encode("utf-8"))

{dbos-0.26.0a1 → dbos-0.26.0a5}/dbos/_client.py

@@ -22,10 +22,9 @@ R = TypeVar("R", covariant=True)  # A generic type for workflow return values
 
 class EnqueueOptions(TypedDict):
     workflow_name: str
-    workflow_class_name: NotRequired[str]
     queue_name: str
-    app_version: NotRequired[str]
     workflow_id: NotRequired[str]
+    app_version: NotRequired[str]
 
 
 class WorkflowHandleClientPolling(Generic[R]):
@@ -86,7 +85,6 @@ class DBOSClient:
         workflow_name = options["workflow_name"]
         queue_name = options["queue_name"]
 
-        workflow_class_name = options.get("workflow_class_name")
         app_version = options.get("app_version")
         max_recovery_attempts = options.get("max_recovery_attempts")
         if max_recovery_attempts is None:
@@ -99,7 +97,7 @@ class DBOSClient:
             "workflow_uuid": workflow_id,
             "status": WorkflowStatusString.ENQUEUED.value,
             "name": workflow_name,
-            "class_name": workflow_class_name,
+            "class_name": None,
             "queue_name": queue_name,
             "app_version": app_version,
             "config_name": None,

{dbos-0.26.0a1 → dbos-0.26.0a5}/dbos/_context.py

@@ -195,8 +195,10 @@ class DBOSContext:
     def end_handler(self, exc_value: Optional[BaseException]) -> None:
         self._end_span(exc_value)
 
-    def get_current_span(self) -> Span:
-        return self.spans[-1]
+    def get_current_span(self) -> Optional[Span]:
+        if len(self.spans):
+            return self.spans[-1]
+        return None
 
     def _start_span(self, attributes: TracedAttributes) -> None:
         attributes["operationUUID"] = (

{dbos-0.26.0a1 → dbos-0.26.0a5}/dbos/_core.py

@@ -398,9 +398,7 @@ async def _execute_workflow_async(
                 raise
 
 
-def execute_workflow_by_id(
-    dbos: "DBOS", workflow_id: str, startNew: bool = False
-) -> "WorkflowHandle[Any]":
+def execute_workflow_by_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[Any]":
     status = dbos._sys_db.get_workflow_status(workflow_id)
     if not status:
         raise DBOSRecoveryError(workflow_id, "Workflow status not found")
@@ -441,7 +439,7 @@ def execute_workflow_by_id(
             class_object = dbos._registry.class_info_map[class_name]
             inputs["args"] = (class_object,) + inputs["args"]
 
-        if startNew:
+        with SetWorkflowID(workflow_id):
             return start_workflow(
                 dbos,
                 wf_func,
@@ -450,16 +448,6 @@ def execute_workflow_by_id(
                 *inputs["args"],
                 **inputs["kwargs"],
             )
-        else:
-            with SetWorkflowID(workflow_id):
-                return start_workflow(
-                    dbos,
-                    wf_func,
-                    status["queue_name"],
-                    True,
-                    *inputs["args"],
-                    **inputs["kwargs"],
-                )
 
 
 def _get_new_wf() -> tuple[str, DBOSContext]:
@@ -892,10 +880,12 @@ def decorate_transaction(
                         except DBAPIError as dbapi_error:
                             if dbapi_error.orig.sqlstate == "40001":  # type: ignore
                                 # Retry on serialization failure
-                                ctx.get_current_span().add_event(
-                                    "Transaction Serialization Failure",
-                                    {"retry_wait_seconds": retry_wait_seconds},
-                                )
+                                span = ctx.get_current_span()
+                                if span:
+                                    span.add_event(
+                                        "Transaction Serialization Failure",
+                                        {"retry_wait_seconds": retry_wait_seconds},
+                                    )
                                 time.sleep(retry_wait_seconds)
                                 retry_wait_seconds = min(
                                     retry_wait_seconds * backoff_factor,
@@ -1004,13 +994,15 @@ def decorate_step(
                     f"Step being automatically retried. (attempt {attempt + 1} of {attempts}). {traceback.format_exc()}"
                 )
                 ctx = assert_current_dbos_context()
-                ctx.get_current_span().add_event(
-                    f"Step attempt {attempt} failed",
-                    {
-                        "error": str(error),
-                        "retryIntervalSeconds": interval_seconds,
-                    },
-                )
+                span = ctx.get_current_span()
+                if span:
+                    span.add_event(
+                        f"Step attempt {attempt} failed",
+                        {
+                            "error": str(error),
+                            "retryIntervalSeconds": interval_seconds,
+                        },
+                    )
                 return min(
                     interval_seconds * (backoff_rate**attempt),
                     max_retry_interval_seconds,

{dbos-0.26.0a1 → dbos-0.26.0a5}/dbos/_dbos.py

@@ -33,7 +33,7 @@ from opentelemetry.trace import Span
 
 from dbos import _serialization
 from dbos._conductor.conductor import ConductorWebsocket
-from dbos._utils import GlobalParams
+from dbos._utils import INTERNAL_QUEUE_NAME, GlobalParams
 from dbos._workflow_commands import (
     WorkflowStatus,
     list_queued_workflows,
@@ -234,6 +234,13 @@ class DBOSRegistry:
             hasher.update(source.encode("utf-8"))
         return hasher.hexdigest()
 
+    def get_internal_queue(self) -> Queue:
+        """
+        Get or create the internal queue used for the DBOS scheduler, for Kafka, and for
+        programmatic resuming and restarting of workflows.
+        """
+        return Queue(INTERNAL_QUEUE_NAME)
+
 
 class DBOS:
     """
@@ -489,6 +496,9 @@ class DBOS:
             notification_listener_thread.start()
             self._background_threads.append(notification_listener_thread)
 
+            # Create the internal queue if it has not yet been created
+            self._registry.get_internal_queue()
+
             # Start the queue thread
             evt = threading.Event()
             self.stop_events.append(evt)
@@ -929,11 +939,6 @@ class DBOS:
         """Execute a workflow by ID (for recovery)."""
         return execute_workflow_by_id(_get_dbos_instance(), workflow_id)
 
-    @classmethod
-    def restart_workflow(cls, workflow_id: str) -> None:
-        """Execute a workflow by ID (for recovery)."""
-        execute_workflow_by_id(_get_dbos_instance(), workflow_id, True)
-
     @classmethod
     def recover_pending_workflows(
         cls, executor_ids: List[str] = ["local"]
@@ -954,7 +959,13 @@ class DBOS:
         dbos_logger.info(f"Resuming workflow: {workflow_id}")
         _get_dbos_instance()._sys_db.resume_workflow(workflow_id)
         _get_or_create_dbos_registry().clear_workflow_cancelled(workflow_id)
-        return execute_workflow_by_id(_get_dbos_instance(), workflow_id, False)
+        return cls.retrieve_workflow(workflow_id)
+
+    @classmethod
+    def restart_workflow(cls, workflow_id: str) -> WorkflowHandle[Any]:
+        """Restart a workflow with a new workflow ID"""
+        forked_workflow_id = _get_dbos_instance()._sys_db.fork_workflow(workflow_id)
+        return cls.retrieve_workflow(forked_workflow_id)
 
     @classmethod
     def list_workflows(
@@ -1083,7 +1094,9 @@ class DBOS:
     def span(cls) -> Span:
         """Return the tracing `Span` associated with the current context."""
         ctx = assert_current_dbos_context()
-        return ctx.get_current_span()
+        span = ctx.get_current_span()
+        assert span
+        return span
 
     @classproperty
     def request(cls) -> Optional["Request"]:

{dbos-0.26.0a1 → dbos-0.26.0a5}/dbos/_kafka.py

@@ -115,7 +115,7 @@ def kafka_consumer(
                 _in_order_kafka_queues[topic] = queue
         else:
             global _kafka_queue
-            _kafka_queue = Queue("_dbos_internal_queue")
+            _kafka_queue = dbosreg.get_internal_queue()
         stop_event = threading.Event()
         dbosreg.register_poller(
             stop_event, _kafka_consumer_loop, func, config, topics, stop_event, in_order

{dbos-0.26.0a1 → dbos-0.26.0a5}/dbos/_queue.py

@@ -82,7 +82,8 @@ def queue_thread(stop_event: threading.Event, dbos: "DBOS") -> None:
     while not stop_event.is_set():
         if stop_event.wait(timeout=1):
             return
-        for _, queue in dbos._registry.queue_info_map.items():
+        queues = dict(dbos._registry.queue_info_map)
+        for _, queue in queues.items():
             try:
                 wf_ids = dbos._sys_db.start_queued_workflows(
                     queue, GlobalParams.executor_id, GlobalParams.app_version

{dbos-0.26.0a1 → dbos-0.26.0a5}/dbos/_recovery.py

@@ -17,7 +17,7 @@ if TYPE_CHECKING:
 def _recover_workflow(
     dbos: "DBOS", workflow: GetPendingWorkflowsOutput
 ) -> "WorkflowHandle[Any]":
-    if workflow.queue_name and workflow.queue_name != "_dbos_internal_queue":
+    if workflow.queue_name:
         cleared = dbos._sys_db.clear_queue_assignment(workflow.workflow_uuid)
         if cleared:
             return dbos.retrieve_workflow(workflow.workflow_uuid)

{dbos-0.26.0a1 → dbos-0.26.0a5}/dbos/_scheduler.py

@@ -52,7 +52,7 @@ def scheduled(
             )
 
         global scheduler_queue
-        scheduler_queue = Queue("_dbos_internal_queue")
+        scheduler_queue = dbosreg.get_internal_queue()
         stop_event = threading.Event()
         dbosreg.register_poller(stop_event, scheduler_loop, func, cron, stop_event)
         return func

{dbos-0.26.0a1 → dbos-0.26.0a5}/dbos/_sys_db.py

@@ -4,6 +4,7 @@ import os
 import re
 import threading
 import time
+import uuid
 from enum import Enum
 from typing import (
     TYPE_CHECKING,
@@ -25,7 +26,7 @@ from alembic.config import Config
 from sqlalchemy.exc import DBAPIError
 from sqlalchemy.sql import func
 
-from dbos._utils import GlobalParams
+from dbos._utils import INTERNAL_QUEUE_NAME, GlobalParams
 
 from . import _serialization
 from ._context import get_local_dbos_context
@@ -447,13 +448,12 @@ class SystemDatabase:
                 )
             )
 
-    def resume_workflow(
-        self,
-        workflow_id: str,
-    ) -> None:
+    def resume_workflow(self, workflow_id: str) -> None:
         if self._debug_mode:
             raise Exception("called resume_workflow in debug mode")
         with self.engine.begin() as c:
+            # Execute with snapshot isolation in case of concurrent calls on the same workflow
+            c.execute(sa.text("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ"))
             # Check the status of the workflow. If it is complete, do nothing.
             row = c.execute(
                 sa.select(
@@ -472,12 +472,63 @@ class SystemDatabase:
                     SystemSchema.workflow_queue.c.workflow_uuid == workflow_id
                 )
             )
-            # Set the workflow's status to PENDING and clear its recovery attempts.
+            # Enqueue the workflow on the internal queue
+            c.execute(
+                pg.insert(SystemSchema.workflow_queue).values(
+                    workflow_uuid=workflow_id,
+                    queue_name=INTERNAL_QUEUE_NAME,
+                )
+            )
+            # Set the workflow's status to ENQUEUED 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)
+                .values(status=WorkflowStatusString.ENQUEUED.value, recovery_attempts=0)
+            )
+
+    def fork_workflow(self, original_workflow_id: str) -> str:
+        status = self.get_workflow_status(original_workflow_id)
+        if status is None:
+            raise Exception(f"Workflow {original_workflow_id} not found")
+        inputs = self.get_workflow_inputs(original_workflow_id)
+        if inputs is None:
+            raise Exception(f"Workflow {original_workflow_id} not found")
+        # Generate a random ID for the forked workflow
+        forked_workflow_id = str(uuid.uuid4())
+        with self.engine.begin() as c:
+            # Create an entry for the forked workflow with the same
+            # initial values as the original.
+            c.execute(
+                pg.insert(SystemSchema.workflow_status).values(
+                    workflow_uuid=forked_workflow_id,
+                    status=WorkflowStatusString.ENQUEUED.value,
+                    name=status["name"],
+                    class_name=status["class_name"],
+                    config_name=status["config_name"],
+                    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=INTERNAL_QUEUE_NAME,
+                )
+            )
+            # Copy the original workflow's inputs into the forked workflow
+            c.execute(
+                pg.insert(SystemSchema.workflow_inputs).values(
+                    workflow_uuid=forked_workflow_id,
+                    inputs=_serialization.serialize_args(inputs),
+                )
+            )
+            # Enqueue the forked workflow on the internal queue
+            c.execute(
+                pg.insert(SystemSchema.workflow_queue).values(
+                    workflow_uuid=forked_workflow_id,
+                    queue_name=INTERNAL_QUEUE_NAME,
+                )
             )
+        return forked_workflow_id
 
     def get_workflow_status(
         self, workflow_uuid: str

{dbos-0.26.0a1 → dbos-0.26.0a5}/dbos/_utils.py

@@ -1,5 +1,7 @@
 import os
 
+INTERNAL_QUEUE_NAME = "_dbos_internal_queue"
+
 
 class GlobalParams:
     app_version: str = os.environ.get("DBOS__APPVERSION", "")

{dbos-0.26.0a1 → dbos-0.26.0a5}/pyproject.toml

@@ -28,7 +28,7 @@ dependencies = [
 ]
 requires-python = ">=3.9"
 readme = "README.md"
-version = "0.26.0a1"
+version = "0.26.0a5"
 
 [project.license]
 text = "MIT"

{dbos-0.26.0a1 → dbos-0.26.0a5}/tests/test_admin_server.py

@@ -13,7 +13,7 @@ from requests.exceptions import ConnectionError
 from dbos import DBOS, ConfigFile, DBOSConfig, Queue, SetWorkflowID, _workflow_commands
 from dbos._schemas.system_database import SystemSchema
 from dbos._sys_db import SystemDatabase, WorkflowStatusString
-from dbos._utils import GlobalParams
+from dbos._utils import INTERNAL_QUEUE_NAME, GlobalParams
 
 
 def test_admin_endpoints(dbos: DBOS) -> None:
@@ -41,6 +41,7 @@ def test_admin_endpoints(dbos: DBOS) -> None:
     )
     assert response.status_code == 200
     assert response.json() == [
+        {"name": INTERNAL_QUEUE_NAME},
         {"name": "q1"},
         {"name": "q2", "concurrency": 1},
         {"name": "q3", "concurrency": 1, "workerConcurrency": 1},

{dbos-0.26.0a1 → dbos-0.26.0a5}/tests/test_failures.py

@@ -178,15 +178,8 @@ def test_dead_letter_queue(dbos: DBOS) -> None:
             dead_letter_workflow()
     assert exc_info.errisinstance(DBOSDeadLetterQueueError)
 
-    # Resume the workflow. Verify it returns to PENDING status without error.
+    # Resume the workflow. Verify it can recover again without error.
     resumed_handle = dbos.resume_workflow(wfid)
-    assert (
-        handle.get_status().status
-        == resumed_handle.get_status().status
-        == WorkflowStatusString.PENDING.value
-    )
-
-    # Verify the workflow can recover again without error.
     DBOS.recover_pending_workflows()
 
     # Complete the blocked workflow

{dbos-0.26.0a1 → dbos-0.26.0a5}/tests/test_scheduler.py

@@ -201,9 +201,8 @@ def test_scheduler_oaoo(dbos: DBOS) -> None:
             else:
                 time.sleep(1)
 
-    # Stop all scheduled workflows
-    for event in dbos.stop_events:
-        event.set()
+    # Stop the scheduled workflow
+    dbos.stop_events[0].set()
 
     dbos._sys_db.update_workflow_status(
         {

dbos-0.26.0a5/tests/test_workflow_management.py

@@ -0,0 +1,212 @@
+import threading
+import uuid
+from typing import Callable
+
+import pytest
+
+# Public API
+from dbos import DBOS, Queue, SetWorkflowID
+from dbos._dbos import DBOSConfiguredInstance
+from dbos._error import DBOSWorkflowCancelledError
+from dbos._utils import INTERNAL_QUEUE_NAME
+from tests.conftest import queue_entries_are_cleaned_up
+
+
+def test_cancel_resume(dbos: DBOS) -> None:
+    steps_completed = 0
+    workflow_event = threading.Event()
+    main_thread_event = threading.Event()
+    input = 5
+
+    @DBOS.step()
+    def step_one() -> None:
+        nonlocal steps_completed
+        steps_completed += 1
+
+    @DBOS.step()
+    def step_two() -> None:
+        nonlocal steps_completed
+        steps_completed += 1
+
+    @DBOS.workflow()
+    def simple_workflow(x: int) -> int:
+        step_one()
+        main_thread_event.set()
+        workflow_event.wait()
+        step_two()
+        return x
+
+    # Start the workflow and cancel it.
+    # Verify it stops after step one but before step two
+    wfid = str(uuid.uuid4())
+    with SetWorkflowID(wfid):
+        handle = DBOS.start_workflow(simple_workflow, input)
+    main_thread_event.wait()
+    DBOS.cancel_workflow(wfid)
+    workflow_event.set()
+    with pytest.raises(DBOSWorkflowCancelledError):
+        handle.get_result()
+    assert steps_completed == 1
+
+    # Resume the workflow. Verify it completes successfully.
+    handle = DBOS.resume_workflow(wfid)
+    assert handle.get_result() == input
+    assert steps_completed == 2
+
+    # Resume the workflow again. Verify it does not run again.
+    handle = DBOS.resume_workflow(wfid)
+    assert handle.get_result() == input
+    assert steps_completed == 2
+
+
+def test_cancel_resume_txn(dbos: DBOS) -> None:
+    txn_completed = 0
+    workflow_event = threading.Event()
+    main_thread_event = threading.Event()
+    input = 5
+
+    @DBOS.transaction()
+    def txn_one() -> None:
+        nonlocal txn_completed
+        txn_completed += 1
+
+    @DBOS.transaction()
+    def txn_two() -> None:
+        nonlocal txn_completed
+        txn_completed += 1
+
+    @DBOS.workflow()
+    def simple_workflow(x: int) -> int:
+        txn_one()
+        main_thread_event.set()
+        workflow_event.wait()
+        txn_two()
+        return x
+
+    # Start the workflow and cancel it.
+    # Verify it stops after step one but before step two
+    wfid = str(uuid.uuid4())
+    with SetWorkflowID(wfid):
+        handle = DBOS.start_workflow(simple_workflow, input)
+    main_thread_event.wait()
+    DBOS.cancel_workflow(wfid)
+    workflow_event.set()
+    with pytest.raises(DBOSWorkflowCancelledError):
+        handle.get_result()
+    assert txn_completed == 1
+
+    # Resume the workflow. Verify it completes successfully.
+    handle = DBOS.resume_workflow(wfid)
+    assert handle.get_result() == input
+    assert txn_completed == 2
+
+    # Resume the workflow again. Verify it does not run again.
+    handle = DBOS.resume_workflow(wfid)
+    assert handle.get_result() == input
+    assert txn_completed == 2
+
+
+def test_cancel_resume_queue(dbos: DBOS) -> None:
+    steps_completed = 0
+    workflow_event = threading.Event()
+    main_thread_event = threading.Event()
+    input = 5
+
+    queue = Queue("test_queue")
+
+    @DBOS.step()
+    def step_one() -> None:
+        nonlocal steps_completed
+        steps_completed += 1
+
+    @DBOS.step()
+    def step_two() -> None:
+        nonlocal steps_completed
+        steps_completed += 1
+
+    @DBOS.workflow()
+    def simple_workflow(x: int) -> int:
+        step_one()
+        main_thread_event.set()
+        workflow_event.wait()
+        step_two()
+        return x
+
+    # Start the workflow and cancel it.
+    # Verify it stops after step one but before step two
+    wfid = str(uuid.uuid4())
+    with SetWorkflowID(wfid):
+        handle = queue.enqueue(simple_workflow, input)
+    main_thread_event.wait()
+    DBOS.cancel_workflow(wfid)
+    workflow_event.set()
+    assert steps_completed == 1
+
+    # Resume the workflow. Verify it completes successfully.
+    handle = DBOS.resume_workflow(wfid)
+    assert handle.get_result() == input
+    assert steps_completed == 2
+
+    # Resume the workflow again. Verify it does not run again.
+    handle = DBOS.resume_workflow(wfid)
+    assert handle.get_result() == input
+    assert steps_completed == 2
+
+    # Verify nothing is left on any queue
+    assert queue_entries_are_cleaned_up(dbos)
+
+
+def test_restart(dbos: DBOS) -> None:
+    input = 2
+    multiplier = 5
+
+    @DBOS.dbos_class()
+    class TestClass(DBOSConfiguredInstance):
+
+        def __init__(self, multiplier: int) -> None:
+            self.multiply: Callable[[int], int] = lambda x: x * multiplier
+            super().__init__("test_class")
+
+        @DBOS.workflow()
+        def workflow(self, x: int) -> int:
+            return self.multiply(x)
+
+        @DBOS.step()
+        def step(self, x: int) -> int:
+            return self.multiply(x)
+
+    inst = TestClass(multiplier)
+
+    # Start the workflow, let it finish, restart it.
+    # Verify it returns the same result with a different workflow ID.
+    handle = DBOS.start_workflow(inst.workflow, input)
+    assert handle.get_result() == input * multiplier
+    forked_handle = DBOS.restart_workflow(handle.workflow_id)
+    assert forked_handle.workflow_id != handle.workflow_id
+    assert forked_handle.get_result() == input * multiplier
+
+    # Enqueue the workflow, let it finish, restart it.
+    # Verify it returns the same result with a different workflow ID and queue.
+    queue = Queue("test_queue")
+    handle = queue.enqueue(inst.workflow, input)
+    assert handle.get_result() == input * multiplier
+    forked_handle = DBOS.restart_workflow(handle.workflow_id)
+    assert forked_handle.workflow_id != handle.workflow_id
+    assert forked_handle.get_status().queue_name == INTERNAL_QUEUE_NAME
+    assert forked_handle.get_result() == input * multiplier
+
+    # Enqueue the step, let it finish, restart it.
+    # Verify it returns the same result with a different workflow ID and queue.
+    handle = queue.enqueue(inst.step, input)
+    assert handle.get_result() == input * multiplier
+    forked_handle = DBOS.restart_workflow(handle.workflow_id)
+    assert forked_handle.workflow_id != handle.workflow_id
+    assert forked_handle.get_status().queue_name != handle.get_status().queue_name
+    assert forked_handle.get_result() == input * multiplier
+
+    # Verify restarting a nonexistent workflow throws an exception
+    with pytest.raises(Exception):
+        DBOS.restart_workflow("fake_id")
+
+    # Verify nothing is left on any queue
+    assert queue_entries_are_cleaned_up(dbos)

dbos-0.26.0a1/tests/test_workflow_cancel.py

@@ -1,145 +0,0 @@
-import threading
-import time
-import uuid
-from datetime import datetime, timedelta, timezone
-
-# Public API
-from dbos import (
-    DBOS,
-    ConfigFile,
-    Queue,
-    SetWorkflowID,
-    WorkflowStatusString,
-    _workflow_commands,
-)
-
-
-def test_basic(dbos: DBOS, config: ConfigFile) -> None:
-
-    steps_completed = 0
-
-    @DBOS.step()
-    def step_one() -> None:
-        nonlocal steps_completed
-        steps_completed += 1
-        print("Step one completed!")
-
-    @DBOS.step()
-    def step_two() -> None:
-        nonlocal steps_completed
-        steps_completed += 1
-        print("Step two completed!")
-
-    @DBOS.workflow()
-    def simple_workflow() -> None:
-        step_one()
-        dbos.sleep(1)
-        step_two()
-        print("Executed Simple workflow")
-        return
-
-    # run the workflow
-    simple_workflow()
-    time.sleep(1)  # wait for the workflow to complete
-    assert (
-        steps_completed == 2
-    ), f"Expected steps_completed to be 2, but got {steps_completed}"
-
-
-def test_two_steps_cancel(dbos: DBOS, config: ConfigFile) -> None:
-
-    steps_completed = 0
-
-    @DBOS.step()
-    def step_one() -> None:
-        nonlocal steps_completed
-        steps_completed += 1
-        print("Step one completed!")
-
-    @DBOS.step()
-    def step_two() -> None:
-        nonlocal steps_completed
-        steps_completed += 1
-        print("Step two completed!")
-
-    @DBOS.workflow()
-    def simple_workflow() -> None:
-        step_one()
-        dbos.sleep(2)
-        step_two()
-        print("Executed Simple workflow")
-        return
-
-    # run the workflow
-    wfuuid = str(uuid.uuid4())
-    try:
-        with SetWorkflowID(wfuuid):
-            simple_workflow()
-
-        dbos.cancel_workflow(wfuuid)
-    except Exception as e:
-        # time.sleep(1)  # wait for the workflow to complete
-        assert (
-            steps_completed == 1
-        ), f"Expected steps_completed to be 1, but got {steps_completed}"
-
-    dbos.resume_workflow(wfuuid)
-    time.sleep(1)
-
-    assert (
-        steps_completed == 2
-    ), f"Expected steps_completed to be 2, but got {steps_completed}"
-
-
-def test_two_transactions_cancel(dbos: DBOS, config: ConfigFile) -> None:
-
-    tr_completed = 0
-
-    @DBOS.transaction()
-    def transaction_one() -> None:
-        nonlocal tr_completed
-        tr_completed += 1
-        print("Transaction one completed!")
-
-    @DBOS.transaction()
-    def transaction_two() -> None:
-        nonlocal tr_completed
-        tr_completed += 1
-        print("Step two completed!")
-
-    @DBOS.workflow()
-    def simple_workflow() -> None:
-        transaction_one()
-        dbos.sleep(2)
-        transaction_two()
-        print("Executed Simple workflow")
-        return
-
-    # run the workflow
-    wfuuid = str(uuid.uuid4())
-    try:
-        with SetWorkflowID(wfuuid):
-            simple_workflow()
-
-        dbos.cancel_workflow(wfuuid)
-    except Exception as e:
-        # time.sleep(1)  # wait for the workflow to complete
-        assert (
-            tr_completed == 1
-        ), f"Expected tr_completed to be 1, but got {tr_completed}"
-
-    dbos.resume_workflow(wfuuid)
-    time.sleep(1)
-
-    assert (
-        tr_completed == 2
-    ), f"Expected steps_completed to be 2, but got {tr_completed}"
-
-    # resume it a 2nd time
-
-    dbos.resume_workflow(wfuuid)
-    time.sleep(1)
-
-    assert (
-        tr_completed == 2
-    ), f"Expected steps_completed to be 2, but got {tr_completed}"