dbos 0.26.0a14__tar.gz → 0.26.0a15__tar.gz

{dbos-0.26.0a14 → dbos-0.26.0a15}/PKG-INFO

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

{dbos-0.26.0a14 → dbos-0.26.0a15}/dbos/_admin_server.py

@@ -7,6 +7,7 @@ from functools import partial
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from typing import TYPE_CHECKING, Any, List, TypedDict
 
+from ._error import DBOSException
 from ._logger import dbos_logger
 from ._recovery import recover_pending_workflows
 from ._utils import GlobalParams
@@ -22,6 +23,7 @@ _workflow_queues_metadata_path = "/dbos-workflow-queues-metadata"
 # /workflows/:workflow_id/resume
 # /workflows/:workflow_id/restart
 # /workflows/:workflow_id/steps
+# /workflows/:workflow_id/fork
 
 
 class AdminServer:
@@ -123,6 +125,9 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
             restart_match = re.match(
                 r"^/workflows/(?P<workflow_id>[^/]+)/restart$", self.path
             )
+            fork_match = re.match(
+                r"^/workflows/(?P<workflow_id>[^/]+)/fork$", self.path
+            )
             resume_match = re.match(
                 r"^/workflows/(?P<workflow_id>[^/]+)/resume$", self.path
             )
@@ -130,7 +135,23 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
                 r"^/workflows/(?P<workflow_id>[^/]+)/cancel$", self.path
             )
 
-            if restart_match:
+            if fork_match:
+                workflow_id = fork_match.group("workflow_id")
+                try:
+                    data = json.loads(post_data.decode("utf-8"))
+                    start_step: int = data.get("start_step", 1)
+                    self._handle_fork(workflow_id, start_step)
+                except (json.JSONDecodeError, AttributeError) as e:
+                    self.send_response(500)
+                    self.send_header("Content-Type", "application/json")
+                    self.end_headers()
+                    self.wfile.write(
+                        json.dumps({"error": f"Invalid JSON input: {str(e)}"}).encode(
+                            "utf-8"
+                        )
+                    )
+                    return
+            elif restart_match:
                 workflow_id = restart_match.group("workflow_id")
                 self._handle_restart(workflow_id)
             elif resume_match:
@@ -147,10 +168,33 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
         return  # Disable admin server request logging
 
     def _handle_restart(self, workflow_id: str) -> None:
-        self.dbos.restart_workflow(workflow_id)
-        print("Restarting workflow", workflow_id)
-        self.send_response(204)
-        self._end_headers()
+        try:
+            print(f"Restarting workflow {workflow_id}")
+            self.dbos.restart_workflow(workflow_id)
+            self.send_response(204)
+            self._end_headers()
+        except DBOSException as e:
+            print(f"Error restarting workflow: {e}")
+            self.send_response(500)
+            response_body = json.dumps({"error": str(e)}).encode("utf-8")
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(response_body)))
+            self.end_headers()
+            self.wfile.write(response_body)
+
+    def _handle_fork(self, workflow_id: str, start_step: int) -> None:
+        try:
+            self.dbos.fork_workflow(workflow_id, start_step)
+            self.send_response(204)
+            self._end_headers()
+        except DBOSException as e:
+            print(f"Error forking workflow: {e}")
+            self.send_response(500)
+            response_body = json.dumps({"error": str(e)}).encode("utf-8")
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(response_body)))
+            self.end_headers()
+            self.wfile.write(response_body)
 
     def _handle_resume(self, workflow_id: str) -> None:
         print("Resuming workflow", workflow_id)

{dbos-0.26.0a14 → dbos-0.26.0a15}/dbos/_app_db.py

@@ -228,3 +228,59 @@ class ApplicationDatabase:
             )
             for row in rows
         ]
+
+    def get_max_function_id(self, workflow_uuid: str) -> Optional[int]:
+        with self.engine.begin() as conn:
+            max_function_id_row = conn.execute(
+                sa.select(
+                    sa.func.max(ApplicationSchema.transaction_outputs.c.function_id)
+                ).where(
+                    ApplicationSchema.transaction_outputs.c.workflow_uuid
+                    == workflow_uuid
+                )
+            ).fetchone()
+
+            max_function_id = max_function_id_row[0] if max_function_id_row else None
+
+            return max_function_id
+
+    def clone_workflow_transactions(
+        self, src_workflow_id: str, forked_workflow_id: str, start_step: int
+    ) -> None:
+        """
+        Copies all steps from dbos.transctions_outputs where function_id < input function_id
+        into a new workflow_uuid. Returns the new workflow_uuid.
+        """
+
+        with self.engine.begin() as conn:
+
+            insert_stmt = sa.insert(ApplicationSchema.transaction_outputs).from_select(
+                [
+                    "workflow_uuid",
+                    "function_id",
+                    "output",
+                    "error",
+                    "txn_id",
+                    "txn_snapshot",
+                    "executor_id",
+                    "function_name",
+                ],
+                sa.select(
+                    sa.literal(forked_workflow_id).label("workflow_uuid"),
+                    ApplicationSchema.transaction_outputs.c.function_id,
+                    ApplicationSchema.transaction_outputs.c.output,
+                    ApplicationSchema.transaction_outputs.c.error,
+                    ApplicationSchema.transaction_outputs.c.txn_id,
+                    ApplicationSchema.transaction_outputs.c.txn_snapshot,
+                    ApplicationSchema.transaction_outputs.c.executor_id,
+                    ApplicationSchema.transaction_outputs.c.function_name,
+                ).where(
+                    (
+                        ApplicationSchema.transaction_outputs.c.workflow_uuid
+                        == src_workflow_id
+                    )
+                    & (ApplicationSchema.transaction_outputs.c.function_id < start_step)
+                ),
+            )
+
+            conn.execute(insert_stmt)

{dbos-0.26.0a14 → dbos-0.26.0a15}/dbos/_dbos.py

@@ -960,14 +960,45 @@ class DBOS:
     def restart_workflow(cls, workflow_id: str) -> WorkflowHandle[Any]:
         """Restart a workflow with a new workflow ID"""
 
+        return cls.fork_workflow(workflow_id, 1)
+
+    @classmethod
+    def fork_workflow(
+        cls, workflow_id: str, start_step: int = 1
+    ) -> WorkflowHandle[Any]:
+        """Restart a workflow with a new workflow ID"""
+
+        def get_max_function_id(workflow_uuid: str) -> int:
+            max_transactions = (
+                _get_dbos_instance()._app_db.get_max_function_id(workflow_uuid) or 0
+            )
+            max_operations = (
+                _get_dbos_instance()._sys_db.get_max_function_id(workflow_uuid) or 0
+            )
+            return max(max_transactions, max_operations)
+
+        max_function_id = get_max_function_id(workflow_id)
+        if max_function_id > 0 and start_step > max_function_id:
+            raise DBOSException(
+                f"Cannot fork workflow {workflow_id} at step {start_step}. The workflow has  {max_function_id} steps."
+            )
+
         def fn() -> str:
-            dbos_logger.info(f"Restarting workflow: {workflow_id}")
-            return _get_dbos_instance()._sys_db.fork_workflow(workflow_id)
+            forked_workflow_id = str(uuid.uuid4())
+            dbos_logger.info(f"Forking workflow: {workflow_id} from step {start_step}")
+
+            _get_dbos_instance()._app_db.clone_workflow_transactions(
+                workflow_id, forked_workflow_id, start_step
+            )
+
+            return _get_dbos_instance()._sys_db.fork_workflow(
+                workflow_id, forked_workflow_id, start_step
+            )
 
-        forked_workflow_id = _get_dbos_instance()._sys_db.call_function_as_step(
-            fn, "DBOS.restartWorkflow"
+        new_id = _get_dbos_instance()._sys_db.call_function_as_step(
+            fn, "DBOS.forkWorkflow"
         )
-        return cls.retrieve_workflow(forked_workflow_id)
+        return cls.retrieve_workflow(new_id)
 
     @classmethod
     def list_workflows(

{dbos-0.26.0a14 → dbos-0.26.0a15}/dbos/_sys_db.py

@@ -489,15 +489,29 @@ class SystemDatabase:
                 .values(status=WorkflowStatusString.ENQUEUED.value, recovery_attempts=0)
             )
 
-    def fork_workflow(self, original_workflow_id: str) -> str:
+    def get_max_function_id(self, workflow_uuid: str) -> Optional[int]:
+        with self.engine.begin() as conn:
+            max_function_id_row = conn.execute(
+                sa.select(
+                    sa.func.max(SystemSchema.operation_outputs.c.function_id)
+                ).where(SystemSchema.operation_outputs.c.workflow_uuid == workflow_uuid)
+            ).fetchone()
+
+            max_function_id = max_function_id_row[0] if max_function_id_row else None
+
+            return max_function_id
+
+    def fork_workflow(
+        self, original_workflow_id: str, forked_workflow_id: str, start_step: int = 1
+    ) -> 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.
@@ -524,6 +538,37 @@ class SystemDatabase:
                     inputs=_serialization.serialize_args(inputs),
                 )
             )
+
+            if start_step > 1:
+
+                # Copy the original workflow's outputs into the forked workflow
+                insert_stmt = sa.insert(SystemSchema.operation_outputs).from_select(
+                    [
+                        "workflow_uuid",
+                        "function_id",
+                        "output",
+                        "error",
+                        "function_name",
+                        "child_workflow_id",
+                    ],
+                    sa.select(
+                        sa.literal(forked_workflow_id).label("workflow_uuid"),
+                        SystemSchema.operation_outputs.c.function_id,
+                        SystemSchema.operation_outputs.c.output,
+                        SystemSchema.operation_outputs.c.error,
+                        SystemSchema.operation_outputs.c.function_name,
+                        SystemSchema.operation_outputs.c.child_workflow_id,
+                    ).where(
+                        (
+                            SystemSchema.operation_outputs.c.workflow_uuid
+                            == original_workflow_id
+                        )
+                        & (SystemSchema.operation_outputs.c.function_id < start_step)
+                    ),
+                )
+
+                c.execute(insert_stmt)
+
             # Enqueue the forked workflow on the internal queue
             c.execute(
                 pg.insert(SystemSchema.workflow_queue).values(

{dbos-0.26.0a14 → dbos-0.26.0a15}/dbos/_workflow_commands.py

@@ -1,4 +1,5 @@
 import json
+import uuid
 from typing import Any, List, Optional
 
 from . import _serialization

{dbos-0.26.0a14 → dbos-0.26.0a15}/dbos/cli/cli.py

@@ -433,13 +433,60 @@ def restart(
     ] = 3001,
 ) -> None:
     response = requests.post(
-        f"http://{host}:{port}/workflows/{uuid}/restart", json=[], timeout=5
+        f"http://{host}:{port}/workflows/{uuid}/restart",
+        json=[],
+        timeout=5,
     )
 
     if response.status_code == 204:
         print(f"Workflow {uuid} has been restarted")
     else:
-        print(f"Failed to resume workflow {uuid}. Status code: {response.status_code}")
+        error_message = response.json().get("error", "Unknown error")
+        print(
+            f"Failed to restart workflow {uuid}. "
+            f"Status code: {response.status_code}. "
+            f"Error: {error_message}"
+        )
+
+
+@workflow.command(
+    help="fork a workflow from the beginning with a new id and from a step"
+)
+def fork(
+    uuid: Annotated[str, typer.Argument()],
+    host: Annotated[
+        typing.Optional[str],
+        typer.Option("--host", "-H", help="Specify the admin host"),
+    ] = "localhost",
+    port: Annotated[
+        typing.Optional[int],
+        typer.Option("--port", "-p", help="Specify the admin port"),
+    ] = 3001,
+    step: Annotated[
+        typing.Optional[int],
+        typer.Option(
+            "--step",
+            "-s",
+            help="Restart from this step (default: first step)",
+        ),
+    ] = 1,
+) -> None:
+    print(f"Forking workflow {uuid} from step {step}")
+    response = requests.post(
+        f"http://{host}:{port}/workflows/{uuid}/fork",
+        json={"start_step": step},
+        timeout=5,
+    )
+
+    if response.status_code == 204:
+        print(f"Workflow {uuid} has been forked")
+    else:
+        error_message = response.json().get("error", "Unknown error")
+        print(
+            f"Failed to fork workflow {uuid}. "
+            f"Status code: {response.status_code}. "
+            f"Error: {error_message}"
+        )
 
 
 @queue.command(name="list", help="List enqueued functions for your application")

{dbos-0.26.0a14 → dbos-0.26.0a15}/pyproject.toml

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