dbos 0.27.0a8__tar.gz → 0.27.0a9__tar.gz

{dbos-0.27.0a8 → dbos-0.27.0a9}/PKG-INFO

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

{dbos-0.27.0a8 → dbos-0.27.0a9}/dbos/_client.py

@@ -243,11 +243,13 @@ class DBOSClient:
     async def cancel_workflow_async(self, workflow_id: str) -> None:
         await asyncio.to_thread(self.cancel_workflow, workflow_id)
 
-    def resume_workflow(self, workflow_id: str) -> None:
+    def resume_workflow(self, workflow_id: str) -> WorkflowHandle[Any]:
         self._sys_db.resume_workflow(workflow_id)
+        return WorkflowHandleClientPolling[Any](workflow_id, self._sys_db)
 
-    async def resume_workflow_async(self, workflow_id: str) -> None:
+    async def resume_workflow_async(self, workflow_id: str) -> WorkflowHandleAsync[Any]:
         await asyncio.to_thread(self.resume_workflow, workflow_id)
+        return WorkflowHandleClientAsyncPolling[Any](workflow_id, self._sys_db)
 
     def list_workflows(
         self,
@@ -361,16 +363,35 @@ class DBOSClient:
     async def list_workflow_steps_async(self, workflow_id: str) -> List[StepInfo]:
         return await asyncio.to_thread(self.list_workflow_steps, workflow_id)
 
-    def fork_workflow(self, workflow_id: str, start_step: int) -> WorkflowHandle[R]:
+    def fork_workflow(
+        self,
+        workflow_id: str,
+        start_step: int,
+        *,
+        application_version: Optional[str] = None,
+    ) -> WorkflowHandle[Any]:
         forked_workflow_id = fork_workflow(
-            self._sys_db, self._app_db, workflow_id, start_step
+            self._sys_db,
+            self._app_db,
+            workflow_id,
+            start_step,
+            application_version=application_version,
         )
-        return WorkflowHandleClientPolling[R](forked_workflow_id, self._sys_db)
+        return WorkflowHandleClientPolling[Any](forked_workflow_id, self._sys_db)
 
     async def fork_workflow_async(
-        self, workflow_id: str, start_step: int
-    ) -> WorkflowHandleAsync[R]:
+        self,
+        workflow_id: str,
+        start_step: int,
+        *,
+        application_version: Optional[str] = None,
+    ) -> WorkflowHandleAsync[Any]:
         forked_workflow_id = await asyncio.to_thread(
-            fork_workflow, self._sys_db, self._app_db, workflow_id, start_step
+            fork_workflow,
+            self._sys_db,
+            self._app_db,
+            workflow_id,
+            start_step,
+            application_version=application_version,
         )
-        return WorkflowHandleClientAsyncPolling[R](forked_workflow_id, self._sys_db)
+        return WorkflowHandleClientAsyncPolling[Any](forked_workflow_id, self._sys_db)

{dbos-0.27.0a8 → dbos-0.27.0a9}/dbos/_dbos.py

@@ -978,7 +978,13 @@ class DBOS:
         return cls.fork_workflow(workflow_id, 1)
 
     @classmethod
-    def fork_workflow(cls, workflow_id: str, start_step: int) -> WorkflowHandle[Any]:
+    def fork_workflow(
+        cls,
+        workflow_id: str,
+        start_step: int,
+        *,
+        application_version: Optional[str] = None,
+    ) -> WorkflowHandle[Any]:
         """Restart a workflow with a new workflow ID from a specific step"""
 
         def fn() -> str:
@@ -988,6 +994,7 @@ class DBOS:
                 _get_dbos_instance()._app_db,
                 workflow_id,
                 start_step,
+                application_version=application_version,
             )
 
         new_id = _get_dbos_instance()._sys_db.call_function_as_step(

{dbos-0.27.0a8 → dbos-0.27.0a9}/dbos/_sys_db.py

@@ -544,15 +544,17 @@ class SystemDatabase:
             # 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(
+            status_row = c.execute(
                 sa.select(
                     SystemSchema.workflow_status.c.status,
                 ).where(SystemSchema.workflow_status.c.workflow_uuid == workflow_id)
             ).fetchone()
+            if status_row is None:
+                return
+            status = status_row[0]
             if (
-                row is None
-                or row[0] == WorkflowStatusString.SUCCESS.value
-                or row[0] == WorkflowStatusString.ERROR.value
+                status == WorkflowStatusString.SUCCESS.value
+                or status == WorkflowStatusString.ERROR.value
             ):
                 return
             # Remove the workflow from the queues table so resume can safely be called on an ENQUEUED workflow
@@ -592,7 +594,12 @@ class SystemDatabase:
             return max_function_id
 
     def fork_workflow(
-        self, original_workflow_id: str, forked_workflow_id: str, start_step: int = 1
+        self,
+        original_workflow_id: str,
+        forked_workflow_id: str,
+        start_step: int,
+        *,
+        application_version: Optional[str],
     ) -> str:
 
         status = self.get_workflow_status(original_workflow_id)
@@ -612,7 +619,11 @@ class SystemDatabase:
                     name=status["name"],
                     class_name=status["class_name"],
                     config_name=status["config_name"],
-                    application_version=status["app_version"],
+                    application_version=(
+                        application_version
+                        if application_version is not None
+                        else status["app_version"]
+                    ),
                     application_id=status["app_id"],
                     request=status["request"],
                     authenticated_user=status["authenticated_user"],

{dbos-0.27.0a8 → dbos-0.27.0a9}/dbos/_workflow_commands.py

@@ -1,6 +1,7 @@
 import uuid
 from typing import List, Optional
 
+from dbos._context import get_local_dbos_context
 from dbos._error import DBOSException
 
 from ._app_db import ApplicationDatabase
@@ -103,6 +104,8 @@ def fork_workflow(
     app_db: ApplicationDatabase,
     workflow_id: str,
     start_step: int,
+    *,
+    application_version: Optional[str],
 ) -> str:
     def get_max_function_id(workflow_uuid: str) -> int:
         max_transactions = app_db.get_max_function_id(workflow_uuid) or 0
@@ -114,7 +117,17 @@ def fork_workflow(
         raise DBOSException(
             f"Cannot fork workflow {workflow_id} from step {start_step}. The workflow has {max_function_id} steps."
         )
-    forked_workflow_id = str(uuid.uuid4())
+    ctx = get_local_dbos_context()
+    if ctx is not None and len(ctx.id_assigned_for_next_workflow) > 0:
+        forked_workflow_id = ctx.id_assigned_for_next_workflow
+        ctx.id_assigned_for_next_workflow = ""
+    else:
+        forked_workflow_id = str(uuid.uuid4())
     app_db.clone_workflow_transactions(workflow_id, forked_workflow_id, start_step)
-    sys_db.fork_workflow(workflow_id, forked_workflow_id, start_step)
+    sys_db.fork_workflow(
+        workflow_id,
+        forked_workflow_id,
+        start_step,
+        application_version=application_version,
+    )
     return forked_workflow_id

{dbos-0.27.0a8 → dbos-0.27.0a9}/pyproject.toml

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

{dbos-0.27.0a8 → dbos-0.27.0a9}/tests/test_client.py

@@ -439,12 +439,17 @@ def test_client_fork(dbos: DBOS, client: DBOSClient) -> None:
     assert handle.get_result() == input * 2
     assert len(client.list_workflow_steps(handle.workflow_id)) == 2
 
-    forked_handle: WorkflowHandle[int] = client.fork_workflow(handle.workflow_id, 1)
-    assert forked_handle.workflow_id != handle.workflow_id
+    fork_id = str(uuid.uuid4())
+    with SetWorkflowID(fork_id):
+        forked_handle: WorkflowHandle[int] = client.fork_workflow(handle.workflow_id, 1)
+    assert forked_handle.workflow_id == fork_id
     assert forked_handle.get_result() == input * 2
 
     forked_handle = client.fork_workflow(handle.workflow_id, 2)
-    assert forked_handle.workflow_id != handle.workflow_id
+    assert (
+        forked_handle.workflow_id != handle.workflow_id
+        and forked_handle.workflow_id != fork_id
+    )
     assert forked_handle.get_result() == input * 2
 
     assert len(client.list_workflows()) == 3

{dbos-0.27.0a8 → dbos-0.27.0a9}/tests/test_workflow_management.py

@@ -8,7 +8,7 @@ import pytest
 from dbos import DBOS, Queue, SetWorkflowID
 from dbos._dbos import DBOSConfiguredInstance
 from dbos._error import DBOSException, DBOSWorkflowCancelledError
-from dbos._utils import INTERNAL_QUEUE_NAME
+from dbos._utils import INTERNAL_QUEUE_NAME, GlobalParams
 from tests.conftest import queue_entries_are_cleaned_up
 
 
@@ -54,6 +54,7 @@ def test_cancel_resume(dbos: DBOS) -> None:
 
     # Resume the workflow. Verify it completes successfully.
     handle = DBOS.resume_workflow(wfid)
+    assert handle.get_status().app_version == GlobalParams.app_version
     assert handle.get_result() == input
     assert steps_completed == 2
 
@@ -218,7 +219,7 @@ def test_restart(dbos: DBOS) -> None:
     assert queue_entries_are_cleaned_up(dbos)
 
 
-def test_restart_fromsteps_stepsonly(
+def test_fork_steps(
     dbos: DBOS,
 ) -> None:
 
@@ -229,47 +230,45 @@ def test_restart_fromsteps_stepsonly(
     stepFiveCount = 0
 
     @DBOS.workflow()
-    def simple_workflow() -> None:
-        stepOne()
-        stepTwo()
-        stepThree()
-        stepFour()
-        stepFive()
-        return
+    def simple_workflow(x: int) -> int:
+        return stepOne(x) + stepTwo(x) + stepThree(x) + stepFour(x) + stepFive(x)
 
     @DBOS.step()
-    def stepOne() -> None:
+    def stepOne(x: int) -> int:
         nonlocal stepOneCount
         stepOneCount += 1
-        return
+        return x + 1
 
     @DBOS.step()
-    def stepTwo() -> None:
+    def stepTwo(x: int) -> int:
         nonlocal stepTwoCount
         stepTwoCount += 1
-        return
+        return x + 2
 
     @DBOS.step()
-    def stepThree() -> None:
+    def stepThree(x: int) -> int:
         nonlocal stepThreeCount
         stepThreeCount += 1
-        return
+        return x + 3
 
     @DBOS.step()
-    def stepFour() -> None:
+    def stepFour(x: int) -> int:
         nonlocal stepFourCount
         stepFourCount += 1
-        return
+        return x + 4
 
     @DBOS.step()
-    def stepFive() -> None:
+    def stepFive(x: int) -> int:
         nonlocal stepFiveCount
         stepFiveCount += 1
-        return
+        return x + 5
+
+    input = 1
+    output = 5 * input + 15
 
     wfid = str(uuid.uuid4())
     with SetWorkflowID(wfid):
-        simple_workflow()
+        assert simple_workflow(input) == output
 
     assert stepOneCount == 1
     assert stepTwoCount == 1
@@ -277,9 +276,12 @@ def test_restart_fromsteps_stepsonly(
     assert stepFourCount == 1
     assert stepFiveCount == 1
 
-    forked_handle = DBOS.fork_workflow(wfid, 3)
-    assert forked_handle.workflow_id != wfid
-    forked_handle.get_result()
+    fork_id = str(uuid.uuid4())
+    with SetWorkflowID(fork_id):
+        forked_handle = DBOS.fork_workflow(wfid, 3)
+    assert forked_handle.workflow_id == fork_id
+    assert forked_handle.get_status().app_version == GlobalParams.app_version
+    assert forked_handle.get_result() == output
 
     assert stepOneCount == 1
     assert stepTwoCount == 1
@@ -289,7 +291,7 @@ def test_restart_fromsteps_stepsonly(
 
     forked_handle = DBOS.fork_workflow(wfid, 5)
     assert forked_handle.workflow_id != wfid
-    forked_handle.get_result()
+    assert forked_handle.get_result() == output
 
     assert stepOneCount == 1
     assert stepTwoCount == 1
@@ -299,7 +301,7 @@ def test_restart_fromsteps_stepsonly(
 
     forked_handle = DBOS.fork_workflow(wfid, 1)
     assert forked_handle.workflow_id != wfid
-    forked_handle.get_result()
+    assert forked_handle.get_result() == output
 
     assert stepOneCount == 2
     assert stepTwoCount == 2
@@ -589,3 +591,42 @@ def test_restart_fromsteps_childwf(
     assert stepOneCount == 2
     assert childwfCount == 3
     assert stepThreeCount == 4
+
+
+def test_fork_version(
+    dbos: DBOS,
+) -> None:
+
+    stepOneCount = 0
+    stepTwoCount = 0
+
+    @DBOS.workflow()
+    def simple_workflow(x: int) -> int:
+        return stepOne(x) + stepTwo(x)
+
+    @DBOS.step()
+    def stepOne(x: int) -> int:
+        nonlocal stepOneCount
+        stepOneCount += 1
+        return x + 1
+
+    @DBOS.step()
+    def stepTwo(x: int) -> int:
+        nonlocal stepTwoCount
+        stepTwoCount += 1
+        return x + 2
+
+    input = 1
+    output = 2 * input + 3
+
+    workflow_id = str(uuid.uuid4())
+    with SetWorkflowID(workflow_id):
+        assert simple_workflow(input) == output
+
+    # Fork the workflow with a different version. Verify it is set to that version.
+    new_version = "my_new_version"
+    handle = DBOS.fork_workflow(workflow_id, 2, application_version=new_version)
+    assert handle.get_status().app_version == new_version
+    # Set the global version to this new version, verify the workflow completes
+    GlobalParams.app_version = new_version
+    assert handle.get_result() == output