Merge pull request #5831 from opsmill/pog-branch-events

Refactor branch events with regards to impacted branch

Author: ogenstad

backend/infrahub/core/branch/tasks.py

@@ -159,7 +159,9 @@ async def rebase_branch(branch: str, context: InfrahubContext, service: Infrahub
     # TODO Add account information
     await service.event.send(
         event=BranchRebasedEvent(
-            branch_name=obj.name, branch_id=str(obj.uuid), meta=EventMeta(branch=obj, context=context)
+            branch_name=obj.name,
+            branch_id=str(obj.uuid),
+            meta=EventMeta.from_context(context=context, branch=registry.get_global_branch()),
         )
     )
 
@@ -174,7 +176,11 @@ async def merge_branch(branch: str, context: InfrahubContext, service: InfrahubS
         obj = await Branch.get_by_name(db=db, name=branch)
         default_branch = await registry.get_branch(db=db, branch=registry.default_branch)
         component_registry = get_component_registry()
-        merge_event = BranchMergedEvent(meta=EventMeta.from_context(context=context, branch=obj))
+        merge_event = BranchMergedEvent(
+            branch_name=obj.name,
+            branch_id=str(obj.get_uuid()),
+            meta=EventMeta.from_context(context=context, branch=registry.get_global_branch()),
+        )
 
         merger: BranchMerger | None = None
         async with lock.registry.global_graph_lock():
@@ -280,7 +286,7 @@ async def delete_branch(branch: str, context: InfrahubContext, service: Infrahub
             branch_name=branch,
             branch_id=str(obj.uuid),
             sync_with_git=obj.sync_with_git,
-            meta=EventMeta(branch=obj, context=context),
+            meta=EventMeta.from_context(context=context, branch=registry.get_global_branch()),
         )
 
         await service.workflow.submit_workflow(
@@ -348,9 +354,7 @@ async def create_branch(model: BranchCreateModel, context: InfrahubContext, serv
             branch_name=obj.name,
             branch_id=str(obj.uuid),
             sync_with_git=obj.sync_with_git,
-            meta=EventMeta(
-                branch=obj, account_id=context.account.account_id, initiator_id=WORKER_IDENTITY, context=context
-            ),
+            meta=EventMeta.from_context(context=context, branch=registry.get_global_branch()),
         )
         await service.event.send(event=event)
 

backend/infrahub/events/branch_action.py

@@ -73,12 +73,14 @@ def event_name(self) -> str:
 class BranchMergedEvent(InfrahubEvent):
     """Event generated when a branch has been merged"""
 
+    branch_name: str = Field(..., description="The name of the branch")
+    branch_id: str = Field(..., description="The ID of the branch")
+
     def get_resource(self) -> dict[str, str]:
         return {
-            "prefect.resource.id": f"infrahub.branch.{self.meta.get_branch_id()}",
-            "infrahub.node.kind": "Branch",
-            "infrahub.node.id": self.meta.get_branch_id(),
-            "infrahub.node.label": self.meta.context.branch.name,
+            "prefect.resource.id": f"infrahub.branch.{self.branch_name}",
+            "infrahub.branch.id": self.branch_id,
+            "infrahub.branch.name": self.branch_name,
         }
 
     def get_messages(self) -> list[InfrahubMessage]:
@@ -99,6 +101,7 @@ def get_resource(self) -> dict[str, str]:
         return {
             "prefect.resource.id": f"infrahub.branch.{self.branch_name}",
             "infrahub.branch.id": self.branch_id,
+            "infrahub.branch.name": self.branch_name,
         }
 
     def get_messages(self) -> list[InfrahubMessage]:

backend/infrahub/graphql/queries/event.py

@@ -2,11 +2,11 @@
 
 from typing import TYPE_CHECKING, Any
 
-from graphene import Boolean, DateTime, Field, Int, List, NonNull, ObjectType, String
+from graphene import Argument, Boolean, DateTime, Field, Int, List, NonNull, ObjectType, String
 from infrahub_sdk.utils import extract_fields_first_node
 
 from infrahub.exceptions import ValidationError
-from infrahub.graphql.types.event import EventNodes
+from infrahub.graphql.types.event import EventNodes, EventTypeFilter
 from infrahub.task_manager.event import PrefectEvent
 from infrahub.task_manager.models import InfrahubEventFilter
 
@@ -32,6 +32,7 @@ async def resolve(
         ids: list[str] | None = None,
         branches: list[str] | None = None,
         event_type: list[str] | None = None,
+        event_type_filter: dict[str, Any] | None = None,
         related_node__ids: list[str] | None = None,
         primary_node__ids: list[str] | None = None,
         parent__ids: list[str] | None = None,
@@ -49,6 +50,7 @@ async def resolve(
             account__ids=account__ids,
             has_children=has_children,
             event_type=event_type,
+            event_type_filter=event_type_filter,
             related_node__ids=related_node__ids,
             primary_node__ids=primary_node__ids,
             parent__ids=parent__ids,
@@ -93,6 +95,7 @@ async def query(
     level=Int(required=False),
     has_children=Boolean(required=False, description="Filter events based on if they can have children or not"),
     event_type=List(NonNull(String), description="Filter events that match a specific type"),
+    event_type_filter=Argument(EventTypeFilter, required=False, description="Filters specific to a given event_type"),
     primary_node__ids=List(
         NonNull(String), description="Filter events where the primary node id is within indicated node ids"
     ),

backend/infrahub/graphql/types/event.py

@@ -2,7 +2,7 @@
 
 from typing import TYPE_CHECKING, Any
 
-from graphene import Boolean, DateTime, Field, Int, Interface, List, NonNull, ObjectType, String
+from graphene import Boolean, DateTime, Field, InputObjectType, Int, Interface, List, NonNull, ObjectType, String
 from graphene.types.generic import GenericScalar
 
 from .common import RelatedNode
@@ -21,16 +21,25 @@ class InfrahubMutatedAttribute(ObjectType):
 
 
 class EventNodeInterface(Interface):
-    id = String(required=True)
-    event = String(required=True)
-    branch = String(required=False)
-    account_id = String(required=False)
-    occurred_at = DateTime(required=True)
-    level = Int(required=True)
-    primary_node = Field(RelatedNode, required=False)
-    related_nodes = List(NonNull(RelatedNode), required=True)
-    has_children = Boolean(required=True)
-    parent_id = String(required=False)
+    id = String(required=True, description="The ID of the event.")
+    event = String(required=True, description="The name of the event.")
+    branch = String(required=False, description="The branch where the event occurred.")
+    account_id = String(required=False, description="The account ID that triggered the event.")
+    occurred_at = DateTime(required=True, description="The timestamp when the event occurred.")
+    level = Int(
+        required=True,
+        description="The level of the event 0 is a root level event, the child events will have 1 and grand children 2.",
+    )
+    primary_node = Field(
+        RelatedNode, required=False, description="The primary Infrahub node this event is associated with."
+    )
+    related_nodes = List(
+        NonNull(RelatedNode), required=True, description="Related Infrahub nodes this event is associated with."
+    )
+    has_children = Boolean(
+        required=True, description="Indicates if the event is expected to have child events under it"
+    )
+    parent_id = String(required=False, description="The event ID of the direct parent to this event.")
 
     @classmethod
     def resolve_type(
@@ -47,27 +56,49 @@ class EventNodes(ObjectType):
     node = Field(EventNodeInterface)
 
 
+class BranchEventTypeFilter(InputObjectType):
+    branches = List(NonNull(String), required=True, description="Name of impacted branches")
+
+
+class EventTypeFilter(InputObjectType):
+    branch_merged = Field(
+        BranchEventTypeFilter, required=False, description="Filters specific to infrahub.branch.merged events"
+    )
+
+
 # ---------------------------------------
 # Branch events
 # ---------------------------------------
 class BranchCreatedEvent(ObjectType):
     class Meta:
         interfaces = (EventNodeInterface,)
 
+    created_branch = String(required=True, description="The name of the branch that was created")
     payload = Field(GenericScalar, required=True)
 
 
+class BranchMergedEvent(ObjectType):
+    class Meta:
+        interfaces = (EventNodeInterface,)
+
+    source_branch = String(required=True, description="The name of the branch that was merged into the default branch")
+
+
 class BranchRebasedEvent(ObjectType):
     class Meta:
         interfaces = (EventNodeInterface,)
 
+    rebased_branch = String(
+        required=True, description="The name of the branch that was rebased and aligned with the default branch"
+    )
     payload = Field(GenericScalar, required=True)
 
 
 class BranchDeletedEvent(ObjectType):
     class Meta:
         interfaces = (EventNodeInterface,)
 
+    deleted_branch = String(required=True, description="The name of the branch that was deleted")
     payload = Field(GenericScalar, required=True)
 
 
@@ -94,6 +125,7 @@ class Meta:
     "infrahub.node.updated": NodeMutatedEvent,
     "infrahub.node.deleted": NodeMutatedEvent,
     "infrahub.branch.created": BranchCreatedEvent,
+    "infrahub.branch.merged": BranchMergedEvent,
     "infrahub.branch.rebased": BranchRebasedEvent,
     "infrahub.branch.deleted": BranchDeletedEvent,
     "undefined": StandardEvent,

backend/infrahub/task_manager/event.py

@@ -6,6 +6,7 @@
 from prefect.events.schemas.events import Event as PrefectEventModel
 from pydantic import BaseModel, Field, TypeAdapter
 
+from infrahub.core.constants import GLOBAL_BRANCH_NAME
 from infrahub.log import get_logger
 from infrahub.utils import get_nested_dict
 
@@ -22,6 +23,8 @@ def get_branch(self) -> str | None:
                 continue
             if "infrahub.resource.label" not in resource:
                 continue
+            if resource.get("infrahub.resource.label") == GLOBAL_BRANCH_NAME:
+                return None
             return resource.get("infrahub.resource.label")
         return None
 
@@ -107,10 +110,34 @@ def _return_node_mutation(self) -> dict[str, Any]:
 
         return {"attributes": attributes}
 
+    def _get_branch_name_from_resource(self) -> str:
+        return self.resource.get("infrahub.branch.name") or ""
+
+    def _return_branch_created(self) -> dict[str, Any]:
+        return {"created_branch": self._get_branch_name_from_resource()}
+
+    def _return_branch_deleted(self) -> dict[str, Any]:
+        return {"deleted_branch": self._get_branch_name_from_resource()}
+
+    def _return_branch_merged(self) -> dict[str, Any]:
+        return {"source_branch": self._get_branch_name_from_resource()}
+
+    def _return_branch_rebased(self) -> dict[str, Any]:
+        return {"rebased_branch": self._get_branch_name_from_resource()}
+
     def _return_event_specifics(self) -> dict[str, Any]:
+        """Return event specific data based on the type of event being processed"""
         match self.event:
             case "infrahub.node.created" | "infrahub.node.updated" | "infrahub.node.deleted":
                 return self._return_node_mutation()
+            case "infrahub.branch.created":
+                return self._return_branch_created()
+            case "infrahub.branch.deleted":
+                return self._return_branch_deleted()
+            case "infrahub.branch.merged":
+                return self._return_branch_merged()
+            case "infrahub.branch.rebased":
+                return self._return_branch_rebased()
 
         return {}
 

backend/infrahub/task_manager/models.py

@@ -2,7 +2,7 @@
 
 import uuid
 from collections import defaultdict
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 from uuid import UUID
 
 from prefect.client.schemas.objects import Log as PrefectLog  # noqa: TC002
@@ -122,7 +122,19 @@ def add_event_id_filter(self, ids: list[str] | None = None) -> None:
         if ids:
             self.id = EventIDFilter(id=[uuid.UUID(id) for id in ids])
 
-    def add_event_type_filter(self, event_type: list[str] | None = None) -> None:
+    def add_event_type_filter(
+        self, event_type: list[str] | None = None, event_type_filter: dict[str, Any] | None = None
+    ) -> None:
+        event_type = event_type or []
+        event_type_filter = event_type_filter or {}
+
+        if branch_merged := event_type_filter.get("branch_merged"):
+            branches: list[str] = branch_merged.get("branches") or []
+            if "infrahub.branch.created" not in event_type:
+                event_type.append("infrahub.branch.merged")
+            if branches:
+                self.resource = EventResourceFilter(labels=ResourceSpecification({"infrahub.branch.name": branches}))
+
         if event_type:
             self.event = EventNameFilter(name=event_type)
 
@@ -159,6 +171,7 @@ def from_filters(
         parent__ids: list[str] | None = None,
         primary_node__ids: list[str] | None = None,
         event_type: list[str] | None = None,
+        event_type_filter: dict[str, Any] | None = None,
         branches: list[str] | None = None,
         level: int | None = None,
         has_children: bool | None = None,
@@ -179,7 +192,7 @@ def from_filters(
 
         filters.add_event_filter(level=level, has_children=has_children)
         filters.add_event_id_filter(ids=ids)
-        filters.add_event_type_filter(event_type=event_type)
+        filters.add_event_type_filter(event_type=event_type, event_type_filter=event_type_filter)
         filters.add_branch_filter(branches=branches)
         filters.add_account_filter(account__ids=account__ids)
         filters.add_parent_filter(parent__ids=parent__ids)

backend/tests/integration/proposed_change/test_proposed_change_conflict.py

@@ -251,7 +251,10 @@ async def test_happy_pipeline(self, db: InfrahubDatabase, happy_data_branch: str
 
         for _ in range(10):
             merge_event = await client.execute_graphql(
-                query=QUERY_EVENT, variables={"branch": happy_data_branch, "event_type": "infrahub.branch.merged"}
+                query=QUERY_EVENT,
+                variables={
+                    "event_type_filter": {"branch_merged": {"branches": happy_data_branch}},
+                },
             )
             if merge_event["InfrahubEvent"]["count"] == 1:
                 break
@@ -322,11 +325,13 @@ async def test_connectivity(self, db: InfrahubDatabase, initial_dataset: str, cl
     $branch: [String!],
     $parent__ids: [String!],
     $event_type: [String!]
+    $event_type_filter: EventTypeFilter
 ) {
   InfrahubEvent(
     branches: $branch,
     parent__ids: $parent__ids
     event_type: $event_type
+    event_type_filter: $event_type_filter
   ) {
     count
     edges {