Update prefect query functions

utils/prefect.py

@@ -5,6 +5,7 @@
 from prefect.client.schemas.filters import (
     FlowRunFilter,
     FlowRunFilterName,
+    FlowRunFilterParentFlowRunId,
     FlowRunFilterTags,
 )
 
@@ -58,26 +59,45 @@ def get_flow_run_name(flow_run_id):
     return asyncio.run(_get_name(flow_run_id))
 
 
-async def _flow_run_query(tags, flow_run_name=None):
-    flow_runs_by_name = []
+async def _flow_run_query(
+    tags=None, flow_run_name=None, parent_flow_run_id=None, sort="START_TIME_DESC"
+):
+    flow_run_filter_parent_flow_run_id = (
+        FlowRunFilterParentFlowRunId(any_=[parent_flow_run_id])
+        if parent_flow_run_id
+        else None
+    )
     async with get_client() as client:
         flow_runs = await client.read_flow_runs(
             flow_run_filter=FlowRunFilter(
                 name=FlowRunFilterName(like_=flow_run_name),
+                parent_flow_run_id=flow_run_filter_parent_flow_run_id,
                 tags=FlowRunFilterTags(all_=tags),
             ),
-            sort="START_TIME_DESC",
+            sort=sort,
         )
-        for flow_run in flow_runs:
-            if flow_run.state_name == "Failed":
-                flow_name = f"❌ {flow_run.name}"
-            elif flow_run.state_name == "Completed":
-                flow_name = f"✅ {flow_run.name}"
-            else:
-                flow_name = f"🕑 {flow_run.name}"
-            flow_runs_by_name.append({"label": flow_name, "value": str(flow_run.id)})
-        return flow_runs_by_name
+        return flow_runs
 
 
-def query_flow_run(tags, flow_run_name=None):
-    return asyncio.run(_flow_run_query(tags, flow_run_name))
+def get_flow_runs_by_name(flow_run_name=None, tags=None):
+    flow_runs_by_name = []
+    flow_runs = asyncio.run(_flow_run_query(tags, flow_run_name=flow_run_name))
+    for flow_run in flow_runs:
+        if flow_run.state_name == "Failed":
+            flow_name = f"❌ {flow_run.name}"
+        elif flow_run.state_name == "Completed":
+            flow_name = f"✅ {flow_run.name}"
+        else:
+            flow_name = f"🕑 {flow_run.name}"
+        flow_runs_by_name.append({"label": flow_name, "value": str(flow_run.id)})
+    return flow_runs_by_name
+
+
+def get_children_flow_run_ids(parent_flow_run_id, sort="START_TIME_ASC"):
+    children_flow_runs = asyncio.run(
+        _flow_run_query(parent_flow_run_id=parent_flow_run_id, sort=sort)
+    )
+    children_flow_run_ids = [
+        str(children_flow_run.id) for children_flow_run in children_flow_runs
+    ]
+    return children_flow_run_ids