dbos 0.16.0a2__tar.gz → 0.17.0a2__tar.gz

{dbos-0.16.0a2 → dbos-0.17.0a2}/PKG-INFO

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

{dbos-0.16.0a2 → dbos-0.17.0a2}/dbos/_core.py

@@ -719,7 +719,7 @@ def decorate_step(
                 finally:
                     dbos._sys_db.record_operation_result(step_output)
 
-            def check_existing_result() -> Optional[R]:
+            def check_existing_result() -> Optional[str]:
                 ctx = assert_current_dbos_context()
                 recorded_output = dbos._sys_db.check_operation_execution(
                     ctx.workflow_id, ctx.function_id
@@ -734,9 +734,7 @@ def decorate_step(
                         )
                         raise deserialized_error
                     elif recorded_output["output"] is not None:
-                        return cast(
-                            R, _serialization.deserialize(recorded_output["output"])
-                        )
+                        return recorded_output["output"]
                     else:
                         raise Exception("Output and error are both None")
                 else:

{dbos-0.16.0a2 → dbos-0.17.0a2}/dbos/_dbos.py

@@ -45,7 +45,7 @@ from ._core import (
     start_workflow,
     workflow_wrapper,
 )
-from ._queue import Queue, _queue_thread
+from ._queue import Queue, queue_thread
 from ._recovery import recover_pending_workflows, startup_recovery_thread
 from ._registrations import (
     DEFAULT_MAX_RECOVERY_ATTEMPTS,
@@ -283,6 +283,7 @@ class DBOS:
         self.flask: Optional["Flask"] = flask
         self._executor_field: Optional[ThreadPoolExecutor] = None
         self._background_threads: List[threading.Thread] = []
+        self._executor_id: str = os.environ.get("DBOS__VMID", "local")
 
         # If using FastAPI, set up middleware and lifecycle events
         if self.fastapi is not None:
@@ -383,7 +384,7 @@ class DBOS:
             evt = threading.Event()
             self.stop_events.append(evt)
             bg_queue_thread = threading.Thread(
-                target=_queue_thread, args=(evt, self), daemon=True
+                target=queue_thread, args=(evt, self), daemon=True
             )
             bg_queue_thread.start()
             self._background_threads.append(bg_queue_thread)

{dbos-0.16.0a2 → dbos-0.17.0a2}/dbos/_outcome.py

@@ -4,6 +4,8 @@ import inspect
 import time
 from typing import Any, Callable, Coroutine, Optional, Protocol, TypeVar, Union, cast
 
+from . import _serialization
+
 T = TypeVar("T")
 R = TypeVar("R")
 
@@ -28,7 +30,7 @@ class Outcome(Protocol[T]):
         exceeded_retries: Callable[[int], BaseException],
     ) -> "Outcome[T]": ...
 
-    def intercept(self, interceptor: Callable[[], Optional[T]]) -> "Outcome[T]": ...
+    def intercept(self, interceptor: Callable[[], Optional[str]]) -> "Outcome[T]": ...
 
     def __call__(self) -> Union[T, Coroutine[Any, Any, T]]: ...
 
@@ -58,11 +60,15 @@ class Immediate(Outcome[T]):
         return Immediate(lambda: before()(self._func))
 
     @staticmethod
-    def _intercept(func: Callable[[], T], interceptor: Callable[[], Optional[T]]) -> T:
+    def _intercept(
+        func: Callable[[], T], interceptor: Callable[[], Optional[str]]
+    ) -> T:
         intercepted = interceptor()
-        return intercepted if intercepted else func()
+        return (
+            cast(T, _serialization.deserialize(intercepted)) if intercepted else func()
+        )
 
-    def intercept(self, interceptor: Callable[[], Optional[T]]) -> "Immediate[T]":
+    def intercept(self, interceptor: Callable[[], Optional[str]]) -> "Immediate[T]":
         return Immediate[T](lambda: Immediate._intercept(self._func, interceptor))
 
     @staticmethod
@@ -151,12 +157,16 @@ class Pending(Outcome[T]):
     @staticmethod
     async def _intercept(
         func: Callable[[], Coroutine[Any, Any, T]],
-        interceptor: Callable[[], Optional[T]],
+        interceptor: Callable[[], Optional[str]],
     ) -> T:
         intercepted = await asyncio.to_thread(interceptor)
-        return intercepted if intercepted else await func()
+        return (
+            cast(T, _serialization.deserialize(intercepted))
+            if intercepted
+            else await func()
+        )
 
-    def intercept(self, interceptor: Callable[[], Optional[T]]) -> "Pending[T]":
+    def intercept(self, interceptor: Callable[[], Optional[str]]) -> "Pending[T]":
         return Pending[T](lambda: Pending._intercept(self._func, interceptor))
 
     @staticmethod

{dbos-0.16.0a2 → dbos-0.17.0a2}/dbos/_queue.py

@@ -51,13 +51,13 @@ class Queue:
         return start_workflow(dbos, func, self.name, False, *args, **kwargs)
 
 
-def _queue_thread(stop_event: threading.Event, dbos: "DBOS") -> None:
+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():
             try:
-                wf_ids = dbos._sys_db.start_queued_workflows(queue)
+                wf_ids = dbos._sys_db.start_queued_workflows(queue, dbos._executor_id)
                 for id in wf_ids:
                     execute_workflow_by_id(dbos, id)
             except Exception:

{dbos-0.16.0a2 → dbos-0.17.0a2}/dbos/_sys_db.py

@@ -1104,7 +1104,7 @@ class SystemDatabase:
                 .on_conflict_do_nothing()
             )
 
-    def start_queued_workflows(self, queue: "Queue") -> List[str]:
+    def start_queued_workflows(self, queue: "Queue", executor_id: str) -> List[str]:
         start_time_ms = int(time.time() * 1000)
         if queue.limiter is not None:
             limiter_period_ms = int(queue.limiter["period"] * 1000)
@@ -1159,7 +1159,7 @@ class SystemDatabase:
                     if len(ret_ids) + num_recent_queries >= queue.limiter["limit"]:
                         break
 
-                # To start a function, first set its status to PENDING
+                # To start a function, first set its status to PENDING and update its executor ID
                 c.execute(
                     SystemSchema.workflow_status.update()
                     .where(SystemSchema.workflow_status.c.workflow_uuid == id)
@@ -1167,7 +1167,10 @@ class SystemDatabase:
                         SystemSchema.workflow_status.c.status
                         == WorkflowStatusString.ENQUEUED.value
                     )
-                    .values(status=WorkflowStatusString.PENDING.value)
+                    .values(
+                        status=WorkflowStatusString.PENDING.value,
+                        executor_id=executor_id,
+                    )
                 )
 
                 # Then give it a start time

{dbos-0.16.0a2 → dbos-0.17.0a2}/pyproject.toml

@@ -23,7 +23,7 @@ dependencies = [
 ]
 requires-python = ">=3.9"
 readme = "README.md"
-version = "0.16.0a2"
+version = "0.17.0a2"
 
 [project.license]
 text = "MIT"

{dbos-0.16.0a2 → dbos-0.17.0a2}/tests/test_dbos.py

@@ -374,6 +374,64 @@ def test_recovery_workflow(dbos: DBOS) -> None:
     assert stat.recovery_attempts == 1
 
 
+def test_recovery_workflow_step(dbos: DBOS) -> None:
+    step_counter: int = 0
+    wf_counter: int = 0
+
+    @DBOS.workflow()
+    def test_workflow(var: str, var2: str) -> str:
+        nonlocal wf_counter
+        wf_counter += 1
+        test_step(var2)
+        return var
+
+    @DBOS.step()
+    def test_step(var2: str) -> None:
+        nonlocal step_counter
+        step_counter += 1
+        print(f"I'm a test_step {var2}!")
+        return
+
+    wfuuid = str(uuid.uuid4())
+    with SetWorkflowID(wfuuid):
+        assert test_workflow("bob", "bob") == "bob"
+
+    dbos._sys_db.wait_for_buffer_flush()
+    # Change the workflow status to pending
+    dbos._sys_db.update_workflow_status(
+        {
+            "workflow_uuid": wfuuid,
+            "status": "PENDING",
+            "name": test_workflow.__qualname__,
+            "class_name": None,
+            "config_name": None,
+            "output": None,
+            "error": None,
+            "executor_id": None,
+            "app_id": None,
+            "app_version": None,
+            "request": None,
+            "recovery_attempts": None,
+            "authenticated_user": None,
+            "authenticated_roles": None,
+            "assumed_role": None,
+            "queue_name": None,
+        }
+    )
+
+    # Recovery should execute the workflow again but skip the transaction
+    workflow_handles = DBOS.recover_pending_workflows()
+    assert len(workflow_handles) == 1
+    assert workflow_handles[0].get_result() == "bob"
+    assert wf_counter == 2
+    assert step_counter == 1
+
+    # Test that there was a recovery attempt of this
+    stat = workflow_handles[0].get_status()
+    assert stat
+    assert stat.recovery_attempts == 1
+
+
 def test_recovery_temp_workflow(dbos: DBOS) -> None:
     txn_counter: int = 0