Query forked_from (#502)

Instead of including `forked_to` in workflow status (which is awkward as
it's not a database column), allow querying by `forked_from` to trace
fork lineage.

Author: kraftp

dbos/_admin_server.py

@@ -338,6 +338,7 @@ def _handle_workflows(self, filters: Dict[str, Any]) -> None:
             end_time=filters.get("end_time"),
             status=filters.get("status"),
             app_version=filters.get("application_version"),
+            forked_from=filters.get("forked_from"),
             name=filters.get("workflow_name"),
             limit=filters.get("limit"),
             offset=filters.get("offset"),
@@ -364,6 +365,7 @@ def _handle_queued_workflows(self, filters: Dict[str, Any]) -> None:
             start_time=filters.get("start_time"),
             end_time=filters.get("end_time"),
             status=filters.get("status"),
+            forked_from=filters.get("forked_from"),
             name=filters.get("workflow_name"),
             limit=filters.get("limit"),
             offset=filters.get("offset"),

dbos/_conductor/conductor.py

@@ -233,6 +233,7 @@ def run(self) -> None:
                                     end_time=body["end_time"],
                                     status=body["status"],
                                     app_version=body["application_version"],
+                                    forked_from=body["forked_from"],
                                     name=body["workflow_name"],
                                     limit=body["limit"],
                                     offset=body["offset"],
@@ -267,6 +268,7 @@ def run(self) -> None:
                                     start_time=q_body["start_time"],
                                     end_time=q_body["end_time"],
                                     status=q_body["status"],
+                                    forked_from=q_body["forked_from"],
                                     name=q_body["workflow_name"],
                                     limit=q_body["limit"],
                                     offset=q_body["offset"],

dbos/_conductor/protocol.py

@@ -118,6 +118,7 @@ class ListWorkflowsBody(TypedDict, total=False):
     end_time: Optional[str]
     status: Optional[str]
     application_version: Optional[str]
+    forked_from: Optional[str]
     limit: Optional[int]
     offset: Optional[int]
     sort_desc: bool
@@ -149,7 +150,6 @@ class WorkflowsOutput:
     Priority: Optional[str]
     QueuePartitionKey: Optional[str]
     ForkedFrom: Optional[str]
-    ForkedTo: Optional[list[str]]
 
     @classmethod
     def from_workflow_information(cls, info: WorkflowStatus) -> "WorkflowsOutput":
@@ -199,7 +199,6 @@ def from_workflow_information(cls, info: WorkflowStatus) -> "WorkflowsOutput":
             Priority=priority_str,
             QueuePartitionKey=info.queue_partition_key,
             ForkedFrom=info.forked_from,
-            ForkedTo=info.forked_to,
         )
 
 
@@ -254,6 +253,7 @@ class ListQueuedWorkflowsBody(TypedDict, total=False):
     start_time: Optional[str]
     end_time: Optional[str]
     status: Optional[str]
+    forked_from: Optional[str]
     queue_name: Optional[str]
     limit: Optional[int]
     offset: Optional[int]

dbos/_dbos.py

@@ -1127,6 +1127,7 @@ def list_workflows(
         end_time: Optional[str] = None,
         name: Optional[str] = None,
         app_version: Optional[str] = None,
+        forked_from: Optional[str] = None,
         user: Optional[str] = None,
         queue_name: Optional[str] = None,
         limit: Optional[int] = None,
@@ -1145,6 +1146,7 @@ def fn() -> List[WorkflowStatus]:
                 end_time=end_time,
                 name=name,
                 app_version=app_version,
+                forked_from=forked_from,
                 user=user,
                 limit=limit,
                 offset=offset,
@@ -1169,6 +1171,7 @@ async def list_workflows_async(
         end_time: Optional[str] = None,
         name: Optional[str] = None,
         app_version: Optional[str] = None,
+        forked_from: Optional[str] = None,
         user: Optional[str] = None,
         limit: Optional[int] = None,
         offset: Optional[int] = None,
@@ -1186,6 +1189,7 @@ async def list_workflows_async(
             end_time=end_time,
             name=name,
             app_version=app_version,
+            forked_from=forked_from,
             user=user,
             limit=limit,
             offset=offset,
@@ -1201,6 +1205,7 @@ def list_queued_workflows(
         *,
         queue_name: Optional[str] = None,
         status: Optional[Union[str, List[str]]] = None,
+        forked_from: Optional[str] = None,
         start_time: Optional[str] = None,
         end_time: Optional[str] = None,
         name: Optional[str] = None,
@@ -1214,6 +1219,7 @@ def fn() -> List[WorkflowStatus]:
                 _get_dbos_instance()._sys_db,
                 queue_name=queue_name,
                 status=status,
+                forked_from=forked_from,
                 start_time=start_time,
                 end_time=end_time,
                 name=name,
@@ -1233,6 +1239,7 @@ async def list_queued_workflows_async(
         *,
         queue_name: Optional[str] = None,
         status: Optional[Union[str, List[str]]] = None,
+        forked_from: Optional[str] = None,
         start_time: Optional[str] = None,
         end_time: Optional[str] = None,
         name: Optional[str] = None,
@@ -1246,6 +1253,7 @@ async def list_queued_workflows_async(
             cls.list_queued_workflows,
             queue_name=queue_name,
             status=status,
+            forked_from=forked_from,
             start_time=start_time,
             end_time=end_time,
             name=name,

dbos/_migration.py

@@ -218,6 +218,7 @@ def get_dbos_migration_three(schema: str) -> str:
 def get_dbos_migration_four(schema: str) -> str:
     return f"""
 ALTER TABLE \"{schema}\".workflow_status ADD COLUMN forked_from TEXT;
+CREATE INDEX "idx_workflow_status_forked_from" ON \"{schema}\"."workflow_status" ("forked_from")
 """
 
 
@@ -334,6 +335,7 @@ def get_sqlite_timestamp_expr() -> str:
 
 sqlite_migration_four = """
 ALTER TABLE workflow_status ADD COLUMN forked_from TEXT;
+CREATE INDEX "idx_workflow_status_forked_from" ON "workflow_status" ("forked_from")
 """
 
 sqlite_migration_five = """

dbos/_sys_db.py

@@ -106,7 +106,7 @@ class WorkflowStatus:
     updated_at: Optional[int]
     # If this workflow was enqueued, on which queue
     queue_name: Optional[str]
-    # The executor to most recently executed this workflow
+    # The executor to most recently execute this workflow
     executor_id: Optional[str]
     # The application version on which this workflow was started
     app_version: Optional[str]
@@ -122,8 +122,6 @@ class WorkflowStatus:
     queue_partition_key: Optional[str]
     # If this workflow was forked from another, that workflow's ID.
     forked_from: Optional[str]
-    # If this workflow was forked to others, those workflows' IDs
-    forked_to: Optional[list[str]]
 
     # INTERNAL FIELDS
 
@@ -213,6 +211,8 @@ def __init__(self) -> None:
         self.status: Optional[List[str]] = None
         # The application version that ran this workflow.
         self.application_version: Optional[str] = None
+        # Get workflows forked from this workflow ID.
+        self.forked_from: Optional[str] = None
         # Return up to this many workflows IDs. IDs are ordered by workflow creation time.
         self.limit: Optional[int] = None
         # Offset into the matching records for pagination
@@ -244,9 +244,9 @@ class StepInfo(TypedDict):
     error: Optional[Exception]
     # If the step starts or retrieves the result of a workflow, its ID
     child_workflow_id: Optional[str]
-    # The UNIX epoch timestamp at which this step started
+    # The Unix epoch timestamp at which this step started
     started_at_epoch_ms: Optional[int]
-    # The UNIX epoch timestamp at which this step completed
+    # The Unix epoch timestamp at which this step completed
     completed_at_epoch_ms: Optional[int]
 
 
@@ -931,6 +931,10 @@ def get_workflows(
                 SystemSchema.workflow_status.c.application_version
                 == input.application_version
             )
+        if input.forked_from:
+            query = query.where(
+                SystemSchema.workflow_status.c.forked_from == input.forked_from
+            )
         if input.workflow_ids:
             query = query.where(
                 SystemSchema.workflow_status.c.workflow_uuid.in_(input.workflow_ids)
@@ -997,29 +1001,6 @@ def get_workflows(
 
             workflow_ids.append(info.workflow_id)
             infos.append(info)
-
-        # Calculate forked_to relationships
-        if workflow_ids:
-            with self.engine.begin() as c:
-                forked_to_query = sa.select(
-                    SystemSchema.workflow_status.c.forked_from,
-                    SystemSchema.workflow_status.c.workflow_uuid,
-                ).where(SystemSchema.workflow_status.c.forked_from.in_(workflow_ids))
-                forked_to_rows = c.execute(forked_to_query).fetchall()
-
-            # Build a mapping of fork-parent workflow ID to list of fork-child workflow IDs
-            forked_to_map: Dict[str, List[str]] = {}
-            for row in forked_to_rows:
-                parent_id = row[0]
-                child_id = row[1]
-                if parent_id not in forked_to_map:
-                    forked_to_map[parent_id] = []
-                forked_to_map[parent_id].append(child_id)
-
-            # Populate the forked_to field for each workflow
-            for info in infos:
-                info.forked_to = forked_to_map.get(info.workflow_id, None)
-
         return infos
 
     def get_pending_workflows(

dbos/_workflow_commands.py

@@ -26,6 +26,7 @@ def list_workflows(
     end_time: Optional[str] = None,
     name: Optional[str] = None,
     app_version: Optional[str] = None,
+    forked_from: Optional[str] = None,
     user: Optional[str] = None,
     queue_name: Optional[str] = None,
     limit: Optional[int] = None,
@@ -42,6 +43,7 @@ def list_workflows(
     input.end_time = end_time
     input.status = status if status is None or isinstance(status, list) else [status]
     input.application_version = app_version
+    input.forked_from = forked_from
     input.queue_name = queue_name
     input.limit = limit
     input.name = name
@@ -61,6 +63,7 @@ def list_queued_workflows(
     *,
     queue_name: Optional[str] = None,
     status: Optional[Union[str, List[str]]] = None,
+    forked_from: Optional[str] = None,
     start_time: Optional[str] = None,
     end_time: Optional[str] = None,
     name: Optional[str] = None,
@@ -73,6 +76,7 @@ def list_queued_workflows(
     input.start_time = start_time
     input.end_time = end_time
     input.status = status if status is None or isinstance(status, list) else [status]
+    input.forked_from = forked_from
     input.limit = limit
     input.name = name
     input.offset = offset

tests/test_admin_server.py

@@ -527,7 +527,6 @@ def test_workflow_2(my_time: datetime) -> str:
     assert workflows[0]["Priority"] == "0"
     assert workflows[0]["QueuePartitionKey"] is None
     assert workflows[0]["ForkedFrom"] is None
-    assert workflows[0]["ForkedTo"] is None
 
     # Only load input and output as requested
     filters = {
@@ -600,10 +599,8 @@ def test_workflow_2(my_time: datetime) -> str:
     assert len(workflows) == 2
     assert workflows[0]["WorkflowUUID"] == handle_1.workflow_id
     assert workflows[0]["ForkedFrom"] == None
-    assert workflows[0]["ForkedTo"] == [handle_3.workflow_id]
     assert workflows[1]["WorkflowUUID"] == handle_3.workflow_id
     assert workflows[1]["ForkedFrom"] == handle_1.workflow_id
-    assert workflows[1]["ForkedTo"] == None
 
     filters = {
         "end_time": (datetime.now(timezone.utc) - timedelta(minutes=10)).isoformat()
@@ -621,6 +618,14 @@ def test_workflow_2(my_time: datetime) -> str:
     workflows = response.json()
     assert len(workflows) == 0
 
+    filters = {
+        "forked_from": "not-an-id",
+    }
+    response = requests.post("http://localhost:3001/workflows", json=filters, timeout=5)
+    assert response.status_code == 200
+    workflows = response.json()
+    assert len(workflows) == 0
+
     filters = {
         "status": ["SUCCESS", "CANCELLED"],
     }

tests/test_workflow_management.py

@@ -262,8 +262,9 @@ def stepFive(x: int) -> int:
     assert stepFourCount == 3
     assert stepFiveCount == 4
 
-    handle: WorkflowHandle[int] = DBOS.retrieve_workflow(wfid)
-    assert handle.get_status().forked_to == [fork_id, fork_id_2, fork_id_3]
+    forks = DBOS.list_workflows(forked_from=wfid)
+    assert len(forks) == 3
+    assert [f.workflow_id for f in forks] == [fork_id, fork_id_2, fork_id_3]
 
 
 def test_restart_fromsteps_transactionsonly(
@@ -357,11 +358,6 @@ def trFive() -> None:
     assert trThreeCount == 3
     assert trFourCount == 4
     assert trFiveCount == 4
-    assert DBOS.retrieve_workflow(wfid).get_status().forked_to == [
-        fork_id_one,
-        fork_id_two,
-        fork_id_three,
-    ]
 
 
 def test_restart_fromsteps_steps_tr(