Improve Fork and Tracing (#432)

- Addresses #431
- Makes the fork CLI command consistent with the internal API

Author: kraftp

dbos/_fastapi.py

@@ -83,4 +83,6 @@ async def dbos_fastapi_middleware(
                     response = await call_next(request)
             else:
                 response = await call_next(request)
+            if hasattr(response, "status_code"):
+                DBOS.span.set_attribute("responseCode", response.status_code)
         return response

dbos/cli/cli.py

@@ -14,6 +14,7 @@
 from rich.prompt import IntPrompt
 from typing_extensions import Annotated, List
 
+from dbos._context import SetWorkflowID
 from dbos._debug import debug_workflow, parse_start_command
 from dbos.cli.migration import grant_dbos_schema_permissions, migrate_dbos_databases
 
@@ -567,7 +568,9 @@ def resume(
     start_client(db_url=db_url).resume_workflow(workflow_id=workflow_id)
 
 
-@workflow.command(help="Restart a workflow from the beginning with a new id")
+@workflow.command(
+    help="[DEPRECATED - Use fork instead] Restart a workflow from the beginning with a new id"
+)
 def restart(
     workflow_id: Annotated[str, typer.Argument()],
     db_url: Annotated[
@@ -600,6 +603,22 @@ def fork(
             help="Restart from this step",
         ),
     ] = 1,
+    forked_workflow_id: Annotated[
+        typing.Optional[str],
+        typer.Option(
+            "--forked-workflow-id",
+            "-f",
+            help="Custom ID for the forked workflow",
+        ),
+    ] = None,
+    application_version: Annotated[
+        typing.Optional[str],
+        typer.Option(
+            "--application-version",
+            "-v",
+            help="Custom application version for the forked workflow",
+        ),
+    ] = None,
     db_url: Annotated[
         typing.Optional[str],
         typer.Option(
@@ -609,11 +628,21 @@ def fork(
         ),
     ] = None,
 ) -> None:
-    status = (
-        start_client(db_url=db_url)
-        .fork_workflow(workflow_id=workflow_id, start_step=step)
-        .get_status()
-    )
+    client = start_client(db_url=db_url)
+
+    if forked_workflow_id is not None:
+        with SetWorkflowID(forked_workflow_id):
+            status = client.fork_workflow(
+                workflow_id=workflow_id,
+                start_step=step,
+                application_version=application_version,
+            ).get_status()
+    else:
+        status = client.fork_workflow(
+            workflow_id=workflow_id,
+            start_step=step,
+            application_version=application_version,
+        ).get_status()
     print(jsonpickle.encode(status, unpicklable=False))
 
 

tests/test_package.py

@@ -275,3 +275,103 @@ def test_workflow_commands(postgres_db_engine: sa.Engine) -> None:
         assert isinstance(fork_wf_data, dict)
         assert fork_wf_data["workflow_id"] != wf_id
         assert fork_wf_data["status"] == "ENQUEUED"
+
+        # fork the workflow with custom forked workflow ID
+        custom_fork_id = "custom-fork-id-12345"
+        output = subprocess.check_output(
+            [
+                "dbos",
+                "workflow",
+                "fork",
+                wf_id,
+                "--step",
+                "3",
+                "--forked-workflow-id",
+                custom_fork_id,
+            ],
+            cwd=temp_path,
+            env=env,
+        )
+        custom_fork_data = json.loads(output)
+        assert isinstance(custom_fork_data, dict)
+        assert custom_fork_data["workflow_id"] == custom_fork_id
+        assert custom_fork_data["status"] == "ENQUEUED"
+
+        # verify the forked workflow data with get command
+        output = subprocess.check_output(
+            ["dbos", "workflow", "get", custom_fork_id, "--db-url", db_url],
+            cwd=temp_path,
+        )
+        custom_fork_get_data = json.loads(output)
+        assert isinstance(custom_fork_get_data, dict)
+        assert custom_fork_get_data["workflow_id"] == custom_fork_id
+
+        # fork the workflow with custom application version
+        output = subprocess.check_output(
+            [
+                "dbos",
+                "workflow",
+                "fork",
+                wf_id,
+                "--step",
+                "2",
+                "--application-version",
+                "test-version",
+            ],
+            cwd=temp_path,
+            env=env,
+        )
+        version_fork_data = json.loads(output)
+        assert isinstance(version_fork_data, dict)
+        assert version_fork_data["workflow_id"] != wf_id
+        assert version_fork_data["status"] == "ENQUEUED"
+
+        # verify the forked workflow data with get command and check application version
+        output = subprocess.check_output(
+            [
+                "dbos",
+                "workflow",
+                "get",
+                version_fork_data["workflow_id"],
+                "--db-url",
+                db_url,
+            ],
+            cwd=temp_path,
+        )
+        version_fork_get_data = json.loads(output)
+        assert isinstance(version_fork_get_data, dict)
+        assert version_fork_get_data["workflow_id"] == version_fork_data["workflow_id"]
+        assert version_fork_get_data["app_version"] == "test-version"
+
+        # fork the workflow with both custom ID and application version
+        custom_fork_id2 = "custom-fork-with-version-67890"
+        output = subprocess.check_output(
+            [
+                "dbos",
+                "workflow",
+                "fork",
+                wf_id,
+                "--step",
+                "4",
+                "--forked-workflow-id",
+                custom_fork_id2,
+                "--application-version",
+                "v2.0.0",
+            ],
+            cwd=temp_path,
+            env=env,
+        )
+        combined_fork_data = json.loads(output)
+        assert isinstance(combined_fork_data, dict)
+        assert combined_fork_data["workflow_id"] == custom_fork_id2
+        assert combined_fork_data["status"] == "ENQUEUED"
+
+        # verify the forked workflow data with get command and check both ID and application version
+        output = subprocess.check_output(
+            ["dbos", "workflow", "get", custom_fork_id2, "--db-url", db_url],
+            cwd=temp_path,
+        )
+        combined_fork_get_data = json.loads(output)
+        assert isinstance(combined_fork_get_data, dict)
+        assert combined_fork_get_data["workflow_id"] == custom_fork_id2
+        assert combined_fork_get_data["app_version"] == "v2.0.0"

tests/test_spans.py

@@ -244,6 +244,8 @@ def test_workflow_endpoint() -> str:
 
     assert spans[0].name == test_workflow_endpoint.__qualname__
     assert spans[1].name == "/wf"
+    assert spans[1].attributes is not None
+    assert spans[1].attributes["responseCode"] == 200
 
     assert spans[0].parent.span_id == spans[1].context.span_id  # type: ignore
     assert spans[1].parent == None