Integrate Prefect flow run into existing InfrahubTask query

Author: dgarros

backend/infrahub/core/task/task.py

@@ -1,5 +1,16 @@
-from typing import Any, Optional
+from collections import defaultdict
+from typing import TYPE_CHECKING, Any, Optional
 
+from prefect.client.orchestration import get_client
+from prefect.client.schemas.filters import (
+    FlowRunFilter,
+    FlowRunFilterTags,
+    LogFilter,
+    LogFilterFlowRunId,
+)
+from prefect.client.schemas.sorting import (
+    FlowRunSort,
+)
 from pydantic import ConfigDict, Field
 
 from infrahub.core.constants import TaskConclusion
@@ -10,9 +21,15 @@
 from infrahub.core.timestamp import current_timestamp
 from infrahub.database import InfrahubDatabase
 from infrahub.utils import get_nested_dict
+from infrahub.workflows.constants import TAG_NAMESPACE, WorkflowTag
 
 from .task_log import TaskLog
 
+if TYPE_CHECKING:
+    from prefect.client.schemas.objects import Log as PrefectLog
+
+LOG_LEVEL_MAPPING = {10: "debug", 20: "info", 30: "warning", 40: "error", 50: "critical"}
+
 
 class Task(StandardNode):
     model_config = ConfigDict(arbitrary_types_allowed=True)
@@ -91,3 +108,83 @@ async def query(
             )
 
         return {"count": count, "edges": nodes}
+
+
+class NewTask:
+    @classmethod
+    async def query(
+        cls,
+        fields: dict[str, Any],
+        related_nodes: list[str],
+        branch: str | None = None,
+        limit: int | None = None,
+        offset: int = 0,
+    ) -> dict[str, Any]:
+        nodes: list[dict] = []
+        count = None
+
+        log_fields = get_nested_dict(nested_dict=fields, keys=["edges", "node", "logs", "edges", "node"])
+        logs_flow: dict[str, list[PrefectLog]] = defaultdict(list)
+
+        async with get_client(sync_client=False) as client:
+            tags = [TAG_NAMESPACE]
+
+            if branch:
+                tags.append(WorkflowTag.BRANCH.render(identifier=branch))
+
+            # We only support one related node for now, need to investigate HOW (and IF) we can support more
+            if related_nodes:
+                tags.append(WorkflowTag.RELATED_NODE.render(identifier=related_nodes[0]))
+
+            flow_run_filters = FlowRunFilter(
+                tags=FlowRunFilterTags(all_=tags),
+            )
+
+            flows = await client.read_flow_runs(
+                flow_run_filter=flow_run_filters,
+                limit=limit,
+                offset=offset,
+                sort=FlowRunSort.START_TIME_DESC,
+            )
+
+            # For now count will just return the number of objects in the response
+            # it won't work well with pagination but it doesn't look like Prefect provide a good option to count all flows
+            if "count" in fields:
+                count = len(flows)
+
+            if log_fields:
+                flow_ids = [flow.id for flow in flows]
+                all_logs = await client.read_logs(log_filter=LogFilter(flow_run_id=LogFilterFlowRunId(any_=flow_ids)))
+                for log in all_logs:
+                    logs_flow[log.flow_run_id].append(log)
+
+            for flow in flows:
+                logs = []
+                if log_fields:
+                    logs = [
+                        {
+                            "node": {
+                                "message": log.message,
+                                "severity": LOG_LEVEL_MAPPING.get(log.level, "error"),
+                                "timestamp": log.timestamp.to_iso8601_string(),
+                            }
+                        }
+                        for log in logs_flow[flow.id]
+                    ]
+
+                nodes.append(
+                    {
+                        "node": {
+                            "title": flow.name,
+                            "conclusion": flow.state_name,
+                            "related_node": "",
+                            "related_node_kind": "",
+                            "created_at": flow.created.to_iso8601_string(),
+                            "updated_at": flow.updated.to_iso8601_string(),
+                            "id": flow.id,
+                            "logs": {"edges": logs},
+                        }
+                    }
+                )
+
+        return {"count": count or 0, "edges": nodes}

backend/infrahub/graphql/queries/task.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
 
 from graphene import Field, Int, List, ObjectType, String
 from infrahub_sdk.utils import extract_fields_first_node
 
-from infrahub.core.task import Task as TaskNode
-from infrahub.graphql.types import TaskNodes
+from infrahub.core.task.task import NewTask as TaskNewNode
+from infrahub.core.task.task import Task as TaskNode
+from infrahub.graphql.types.task import TaskNodes
 
 if TYPE_CHECKING:
     from graphql import GraphQLResolveInfo
@@ -24,30 +25,53 @@ async def resolve(
         info: GraphQLResolveInfo,
         limit: int = 10,
         offset: int = 0,
-        ids: Optional[list] = None,
-        related_node__ids: Optional[list] = None,
+        ids: list | None = None,
+        branch: str | None = None,
+        related_node__ids: list | None = None,
     ) -> dict[str, Any]:
         related_nodes = related_node__ids or []
         ids = ids or []
-        return await Tasks.query(info=info, limit=limit, offset=offset, ids=ids, related_nodes=related_nodes)
+        return await Tasks.query(
+            info=info, branch=branch, limit=limit, offset=offset, ids=ids, related_nodes=related_nodes
+        )
 
     @classmethod
     async def query(
-        cls, info: GraphQLResolveInfo, limit: int, offset: int, related_nodes: list[str], ids: list[str]
+        cls,
+        info: GraphQLResolveInfo,
+        limit: int,
+        offset: int,
+        related_nodes: list[str],
+        ids: list[str],
+        branch: str | None,
     ) -> dict[str, Any]:
         context: GraphqlContext = info.context
         fields = await extract_fields_first_node(info)
 
-        return await TaskNode.query(
-            db=context.db, fields=fields, limit=limit, offset=offset, ids=ids, related_nodes=related_nodes
+        # During the migration, query both Prefect and Infrahub to get the list of tasks
+        if not branch:
+            infrahub_tasks = await TaskNode.query(
+                db=context.db, fields=fields, limit=limit, offset=offset, ids=ids, related_nodes=related_nodes
+            )
+        else:
+            infrahub_tasks = {}
+
+        prefect_tasks = await TaskNewNode.query(
+            fields=fields, branch=branch, related_nodes=related_nodes, limit=limit, offset=offset
         )
 
+        return {
+            "count": infrahub_tasks.get("count", 0) + prefect_tasks.get("count", 0),
+            "edges": infrahub_tasks.get("edges", []) + prefect_tasks.get("edges", []),
+        }
+
 
 Task = Field(
     Tasks,
     resolver=Tasks.resolve,
     limit=Int(required=False),
     offset=Int(required=False),
     related_node__ids=List(String),
+    branch=String(required=False),
     ids=List(String),
 )

backend/infrahub/graphql/schema.py

@@ -50,9 +50,9 @@
     InfrahubSearchAnywhere,
     InfrahubStatus,
     Relationship,
-    Task,
 )
 from .queries.diff.tree import DiffTreeQuery, DiffTreeSummaryQuery
+from .queries.task import Task
 
 
 class InfrahubBaseQuery(ObjectType):

backend/infrahub/workflows/constants.py

@@ -13,6 +13,7 @@ class WorkflowTag(InfrahubStringEnum):
     BRANCH = "branch/{identifier}"
     WORKFLOWTYPE = "workflow-type/{identifier}"
     DATABASE_CHANGE = "database-change"
+    RELATED_NODE = "node/{identifier}"
 
     def render(self, identifier: str | None = None) -> str:
         if identifier is None:

backend/infrahub/workflows/models.py

@@ -12,7 +12,7 @@
 from infrahub import __version__
 from infrahub.core.constants import BranchSupportType
 
-from .constants import WorkflowTag, WorkflowType
+from .constants import TAG_NAMESPACE, WorkflowTag, WorkflowType
 
 TASK_RESULT_STORAGE_NAME = "infrahub-storage"
 
@@ -61,7 +61,7 @@ def to_deployment(self) -> dict[str, Any]:
         return payload
 
     def get_tags(self) -> list[str]:
-        tags: list[str] = [WorkflowTag.WORKFLOWTYPE.render(identifier=self.type.value)]
+        tags: list[str] = [TAG_NAMESPACE, WorkflowTag.WORKFLOWTYPE.render(identifier=self.type.value)]
         tags += [tag.render() for tag in self.tags]
         return tags