Extend InfrahubTask and add InfrahubTaskBranchStatus

dgarros · dgarros · commit 7b5be6dd6793 · 2024-11-04T19:40:44.000+01:00
diff --git a/backend/infrahub/core/task/task.py b/backend/infrahub/core/task/task.py
@@ -1,16 +1,5 @@
-from collections import defaultdict
-from typing import TYPE_CHECKING, Any, Optional
+from typing import 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
@@ -21,15 +10,9 @@
 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)
@@ -55,11 +38,13 @@ async def query(
         cls,
         db: InfrahubDatabase,
         fields: dict[str, Any],
-        limit: int,
-        offset: int,
-        ids: list[str],
-        related_nodes: list[str],
+        limit: int | None = None,
+        offset: int | None = None,
+        ids: list[str] | None = None,
+        related_nodes: list[str] | None = None,
     ) -> dict[str, Any]:
+        ids = ids or []
+        related_nodes = related_nodes or []
         log_fields = get_nested_dict(nested_dict=fields, keys=["edges", "node", "logs", "edges", "node"])
         count = None
         if "count" in fields:
@@ -108,83 +93,3 @@ 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}
diff --git a/backend/infrahub/graphql/queries/task.py b/backend/infrahub/graphql/queries/task.py
@@ -4,10 +4,12 @@
 
 from graphene import Field, Int, List, ObjectType, String
 from infrahub_sdk.utils import extract_fields_first_node
+from prefect.client.schemas.objects import StateType
 
-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
+from infrahub.task_manager.task import PrefectTask
+from infrahub.workflows.constants import WorkflowTag
 
 if TYPE_CHECKING:
     from graphql import GraphQLResolveInfo
@@ -35,15 +37,28 @@ async def resolve(
             info=info, branch=branch, limit=limit, offset=offset, ids=ids, related_nodes=related_nodes
         )
 
+    @staticmethod
+    async def resolve_branch_status(
+        root: dict,  # pylint: disable=unused-argument
+        info: GraphQLResolveInfo,
+        branch: str,
+    ) -> dict[str, Any]:
+        statuses: list[StateType] = [StateType.PENDING, StateType.RUNNING, StateType.CANCELLING, StateType.SCHEDULED]
+        tags: list[str] = [WorkflowTag.DATABASE_CHANGE.render()]
+
+        return await Tasks.query(info=info, branch=branch, statuses=statuses, tags=tags)
+
     @classmethod
     async def query(
         cls,
         info: GraphQLResolveInfo,
-        limit: int,
-        offset: int,
-        related_nodes: list[str],
-        ids: list[str],
-        branch: str | None,
+        related_nodes: list[str] | None = None,
+        ids: list[str] | None = None,
+        statuses: list[StateType] | None = None,
+        tags: list[str] | None = None,
+        branch: str | None = None,
+        limit: int | None = None,
+        offset: int | None = None,
     ) -> dict[str, Any]:
         context: GraphqlContext = info.context
         fields = await extract_fields_first_node(info)
@@ -56,12 +71,20 @@ async def query(
         else:
             infrahub_tasks = {}
 
-        prefect_tasks = await TaskNewNode.query(
-            fields=fields, branch=branch, related_nodes=related_nodes, limit=limit, offset=offset
+        prefect_tasks = await PrefectTask.query(
+            db=context.db,
+            fields=fields,
+            branch=branch,
+            statuses=statuses,
+            tags=tags,
+            related_nodes=related_nodes,
+            limit=limit,
+            offset=offset,
         )
-
+        infrahub_count = infrahub_tasks.get("count", None)
+        prefect_count = prefect_tasks.get("count", None)
         return {
-            "count": infrahub_tasks.get("count", 0) + prefect_tasks.get("count", 0),
+            "count": (infrahub_count or 0) + (prefect_count or 0),
             "edges": infrahub_tasks.get("edges", []) + prefect_tasks.get("edges", []),
         }
 
@@ -75,3 +98,10 @@ async def query(
     branch=String(required=False),
     ids=List(String),
 )
+
+TaskBranchStatus = Field(
+    Tasks,
+    resolver=Tasks.resolve_branch_status,
+    branch=String(required=False),
+    description="Return the list of all pending or running tasks that can modify the data, for a given branch",
+)
diff --git a/backend/infrahub/graphql/schema.py b/backend/infrahub/graphql/schema.py
@@ -52,7 +52,7 @@
     Relationship,
 )
 from .queries.diff.tree import DiffTreeQuery, DiffTreeSummaryQuery
-from .queries.task import Task
+from .queries.task import Task, TaskBranchStatus
 
 
 class InfrahubBaseQuery(ObjectType):
@@ -72,6 +72,7 @@ class InfrahubBaseQuery(ObjectType):
     InfrahubSearchAnywhere = InfrahubSearchAnywhere
 
     InfrahubTask = Task
+    InfrahubTaskBranchStatus = TaskBranchStatus
 
     IPAddressGetNextAvailable = InfrahubIPAddressGetNextAvailable
     IPPrefixGetNextAvailable = InfrahubIPPrefixGetNextAvailable
diff --git a/backend/infrahub/graphql/types/task.py b/backend/infrahub/graphql/types/task.py
@@ -1,21 +1,31 @@
 from __future__ import annotations
 
-from graphene import Field, ObjectType, String
+from graphene import Enum, Field, Float, List, ObjectType, String
+from graphene.types.generic import GenericScalar
+from prefect.client.schemas.objects import StateType
 
 from .task_log import TaskLogEdge
 
+TaskState = Enum.from_enum(StateType)
+
 
 class Task(ObjectType):
     id = String(required=True)
     title = String(required=True)
     conclusion = String(required=True)
+    state = TaskState(required=False)
+    progress = Float(required=False)
+    branch = String(required=False)
     created_at = String(required=True)
     updated_at = String(required=True)
+    parameters = GenericScalar(required=False)
+    tags = List(String, required=False)
+    start_time = String(required=False)
 
 
 class TaskNode(Task):
-    related_node = String(required=True)
-    related_node_kind = String(required=True)
+    related_node = String(required=False)
+    related_node_kind = String(required=False)
     logs = Field(TaskLogEdge)
 
 
diff --git a/backend/infrahub/task_manager/__init__.py b/backend/infrahub/task_manager/__init__.py
diff --git a/backend/infrahub/task_manager/constants.py b/backend/infrahub/task_manager/constants.py
@@ -0,0 +1,15 @@
+from infrahub.core.constants import TaskConclusion
+
+LOG_LEVEL_MAPPING = {10: "debug", 20: "info", 30: "warning", 40: "error", 50: "critical"}
+
+CONCLUSION_STATE_MAPPING = {
+    "Scheduled": TaskConclusion.UNKNOWN,
+    "Pending": TaskConclusion.UNKNOWN,
+    "Running": TaskConclusion.UNKNOWN,
+    "Completed": TaskConclusion.SUCCESS,
+    "Failed": TaskConclusion.FAILURE,
+    "Cancelled": TaskConclusion.FAILURE,
+    "Crashed": TaskConclusion.FAILURE,
+    "Paused": TaskConclusion.UNKNOWN,
+    "Cancelling": TaskConclusion.FAILURE,
+}
diff --git a/backend/infrahub/task_manager/models.py b/backend/infrahub/task_manager/models.py
@@ -0,0 +1,36 @@
+from collections import defaultdict
+from typing import DefaultDict
+from uuid import UUID
+
+from prefect.client.schemas.objects import Log as PrefectLog
+from pydantic import BaseModel, Field
+
+from .constants import LOG_LEVEL_MAPPING
+
+
+class RelatedNodesInfo(BaseModel):
+    id: dict[UUID, str] = Field(default_factory=dict)
+    kind: dict[UUID, str | None] = Field(default_factory=dict)
+
+    def get_unique_related_node_ids(self) -> list[str]:
+        return list(set(list(self.id.values())))
+
+
+class FlowLogs(BaseModel):
+    logs: DefaultDict[UUID, list[PrefectLog]] = Field(default_factory=lambda: defaultdict(list))
+
+    def to_graphql(self, flow_id: UUID) -> list[dict]:
+        return [
+            {
+                "node": {
+                    "message": log.message,
+                    "severity": LOG_LEVEL_MAPPING.get(log.level, "error"),
+                    "timestamp": log.timestamp.to_iso8601_string(),
+                }
+            }
+            for log in self.logs[flow_id]
+        ]
+
+
+class FlowProgress(BaseModel):
+    data: dict[UUID, float] = Field(default_factory=dict)
diff --git a/backend/infrahub/task_manager/task.py b/backend/infrahub/task_manager/task.py
diff --git a/backend/tests/unit/graphql/queries/test_task.py b/backend/tests/unit/graphql/queries/test_task.py
