Send node mutation events after branch merge

Author: ogenstad

backend/infrahub/context.py

@@ -5,10 +5,18 @@
 from infrahub.core.branch import Branch
 
 
-class EventContext(BaseModel):
+class ParentEvent(BaseModel):
+    id: str
     name: str
 
 
+class EventContext(BaseModel):
+    name: str = Field(..., description="The name of the event")
+    id: str = Field(..., description="The ID of the event")
+    parent_id: str | None = Field(default=None)
+    ancestors: list[ParentEvent] = Field(default_factory=list)
+
+
 class BranchContext(BaseModel):
     name: str
     id: str | None = None
@@ -17,8 +25,15 @@ class BranchContext(BaseModel):
 class InfrahubContext(BaseModel):
     branch: BranchContext
     account: AccountSession
-    events: list[EventContext] = Field(default_factory=list)
+    event: EventContext | None = Field(default=None)
 
     @classmethod
     def init(cls, branch: Branch, account: AccountSession) -> Self:
         return cls(branch=BranchContext(name=branch.name, id=str(branch.uuid)), account=account)
+
+    def set_event(self, name: str, id: str) -> None:
+        if self.event:
+            self.event.name = name
+            self.event.id = id
+        else:
+            self.event = EventContext(name=name, id=id)

backend/infrahub/core/branch/tasks.py

@@ -11,6 +11,8 @@
 from infrahub.context import InfrahubContext  # noqa: TC001  needed for prefect flow
 from infrahub.core import registry
 from infrahub.core.branch import Branch
+from infrahub.core.changelog.diff import DiffChangelogCollector
+from infrahub.core.constants import MutationAction
 from infrahub.core.diff.coordinator import DiffCoordinator
 from infrahub.core.diff.ipam_diff_parser import IpamDiffParser
 from infrahub.core.diff.merger.merger import DiffMerger
@@ -23,8 +25,9 @@
 from infrahub.core.validators.models.validate_migration import SchemaValidateMigrationData
 from infrahub.core.validators.tasks import schema_validate_migrations
 from infrahub.dependencies.registry import get_component_registry
-from infrahub.events.branch_action import BranchCreatedEvent, BranchDeletedEvent, BranchRebasedEvent
-from infrahub.events.models import EventMeta
+from infrahub.events.branch_action import BranchCreatedEvent, BranchDeletedEvent, BranchMergedEvent, BranchRebasedEvent
+from infrahub.events.models import EventMeta, InfrahubEvent
+from infrahub.events.node_action import NodeMutatedEvent
 from infrahub.exceptions import BranchNotFoundError, MergeFailedError, ValidationError
 from infrahub.graphql.mutations.models import BranchCreateModel  # noqa: TC001
 from infrahub.log import get_log_data
@@ -169,7 +172,9 @@ async def merge_branch(branch: str, context: InfrahubContext, service: InfrahubS
         await add_tags(branches=[branch, registry.default_branch])
 
         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))
 
         merger: BranchMerger | None = None
         async with lock.registry.global_graph_lock():
@@ -187,13 +192,15 @@ async def merge_branch(branch: str, context: InfrahubContext, service: InfrahubS
                 service=service,
             )
             try:
-                await merger.merge()
+                branch_diff = await merger.merge()
             except Exception as exc:
                 log.exception("Merge failed, beginning rollback")
                 await merger.rollback()
                 raise MergeFailedError(branch_name=branch) from exc
             await merger.update_schema()
 
+        changelog_collector = DiffChangelogCollector(diff=branch_diff, branch=obj, db=db)
+        node_events = changelog_collector.collect_changelogs()
         if merger and merger.migrations:
             errors = await schema_apply_migrations(
                 message=SchemaApplyMigrationData(
@@ -242,6 +249,24 @@ async def merge_branch(branch: str, context: InfrahubContext, service: InfrahubS
         )
         await service.message_bus.send(message=message)
 
+        events: list[InfrahubEvent] = [merge_event]
+
+        for action, node_changelog in node_events:
+            meta = EventMeta.from_parent(parent=merge_event)
+            mutate_event = NodeMutatedEvent(
+                kind=node_changelog.node_kind,
+                node_id=node_changelog.node_id,
+                data=node_changelog,
+                action=MutationAction.from_diff_action(diff_action=action),
+                fields=node_changelog.updated_fields,
+                meta=meta,
+            )
+            mutate_event.set_context_branch(branch=default_branch)
+            events.append(mutate_event)
+
+        for event in events:
+            await service.event.send(event=event)
+
 
 @flow(name="branch-delete", flow_run_name="Delete branch {branch}")
 async def delete_branch(branch: str, context: InfrahubContext, service: InfrahubServices) -> None:

backend/infrahub/core/changelog/diff.py

@@ -50,6 +50,15 @@ def _populate_diff_nodes(self) -> None:
     def get_node(self, node_id: str) -> NodeInDiff:
         return self._diff_nodes[node_id]
 
+    def get_peer_kind(self, peer_id: str, node_kind: str, relationship_name: str) -> str:
+        """If the peer kind doesn't exist in the diff use the peer kind from the schema"""
+        try:
+            return self.get_node(node_id=peer_id).kind
+        except KeyError:
+            schema = self._db.schema.get(node_kind, branch=self._branch, duplicate=False)
+            rel_schema = schema.get_relationship(name=relationship_name)
+            return rel_schema.peer
+
     def _process_node(self, node: EnrichedDiffNode) -> NodeChangelog:
         node_changelog = NodeChangelog(node_id=node.uuid, node_kind=node.kind, display_label=node.label)
         schema = self._db.schema.get(node_changelog.node_kind, branch=self._branch, duplicate=False)
@@ -64,8 +73,15 @@ def _process_node(self, node: EnrichedDiffNode) -> NodeChangelog:
     def _process_node_attribute(
         self, node: NodeChangelog, attribute: EnrichedDiffAttribute, schema: MainSchemaTypes
     ) -> None:
-        schema_attribute = schema.get_attribute(name=attribute.name)
-        changelog_attribute = AttributeChangelog(name=attribute.name, kind=schema_attribute.kind)
+        try:
+            schema_attribute = schema.get_attribute(name=attribute.name)
+            attribute_kind = schema_attribute.kind
+        except ValueError:
+            # This would currently happen if there has been a schema migration as part of the merge
+            # then we don't have access to the attribute kind
+            attribute_kind = "n/a"
+
+        changelog_attribute = AttributeChangelog(name=attribute.name, kind=attribute_kind)
         for attr_property in attribute.properties:
             match attr_property.property_type:
                 case DatabaseEdgeType.HAS_VALUE:
@@ -117,10 +133,18 @@ def _process_node_cardinality_one_relationship(
                     case DatabaseEdgeType.IS_RELATED:
                         if rel_prop.new_value:
                             changelog_rel.peer_id = rel_prop.new_value
-                            changelog_rel.peer_kind = self.get_node(node_id=rel_prop.new_value).kind
+                            changelog_rel.peer_kind = self.get_peer_kind(
+                                peer_id=rel_prop.new_value,
+                                node_kind=node.node_kind,
+                                relationship_name=relationship.name,
+                            )
                         if rel_prop.previous_value:
                             changelog_rel.peer_id_previous = rel_prop.previous_value
-                            changelog_rel.peer_kind_previous = self.get_node(node_id=rel_prop.previous_value).kind
+                            changelog_rel.peer_kind_previous = self.get_peer_kind(
+                                peer_id=rel_prop.previous_value,
+                                node_kind=node.node_kind,
+                                relationship_name=relationship.name,
+                            )
                     case DatabaseEdgeType.IS_PROTECTED:
                         changelog_rel.add_property(
                             name="is_protected",
@@ -161,7 +185,11 @@ def _process_node_cardinality_many_relationship(
         changelog_rel = RelationshipCardinalityManyChangelog(name=relationship.name)
         for peer in relationship.relationships:
             peer_log = RelationshipPeerChangelog(
-                peer_id=peer.peer_id, peer_kind=self.get_node(node_id=peer.peer_id).kind, peer_status=peer.action
+                peer_id=peer.peer_id,
+                peer_kind=self.get_peer_kind(
+                    peer_id=peer.peer_id, node_kind=node.node_kind, relationship_name=relationship.name
+                ),
+                peer_status=peer.action,
             )
             for peer_prop in peer.properties:
                 match peer_prop.property_type:

backend/infrahub/core/constants/__init__.py

@@ -180,6 +180,18 @@ class MutationAction(InfrahubStringEnum):
     UPDATED = "updated"
     UNDEFINED = "undefined"
 
+    @classmethod
+    def from_diff_action(cls, diff_action: DiffAction) -> MutationAction:
+        match diff_action:
+            case DiffAction.ADDED:
+                return MutationAction.CREATED
+            case DiffAction.REMOVED:
+                return MutationAction.DELETED
+            case DiffAction.UPDATED:
+                return MutationAction.UPDATED
+            case DiffAction.UNCHANGED:
+                return MutationAction.UNDEFINED
+
 
 class PathResourceType(InfrahubStringEnum):
     SCHEMA = "schema"

backend/infrahub/events/branch_action.py

@@ -70,6 +70,25 @@ def event_name(self) -> str:
         return f"{EVENT_NAMESPACE}.branch.created"
 
 
+class BranchMergedEvent(InfrahubEvent):
+    """Event generated when a branch has been merged"""
+
+    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,
+        }
+
+    def get_messages(self) -> list[InfrahubMessage]:
+        return []
+
+    @computed_field
+    def event_name(self) -> str:
+        return f"{EVENT_NAMESPACE}.branch.merged"
+
+
 class BranchRebasedEvent(InfrahubEvent):
     """Event generated when a branch has been rebased"""
 

backend/infrahub/events/models.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
-from typing import Any, cast, final
+from copy import deepcopy
+from typing import Any, Self, cast, final
 from uuid import UUID, uuid4
 
-from pydantic import BaseModel, Field, computed_field
+from pydantic import BaseModel, Field, PrivateAttr, computed_field, model_validator
 
 from infrahub import __version__
 from infrahub.auth import AccountSession, AuthType
@@ -45,6 +46,7 @@ class EventMeta(BaseModel):
 
     parent: UUID | None = Field(default=None, description="The UUID of the parent event if applicable")
     ancestors: list[ParentEvent] = Field(default_factory=list, description="Any event used to trigger this event")
+    _created_with_context: bool = PrivateAttr(default=False)
 
     def get_branch_id(self) -> str:
         if self.context.branch.id:
@@ -123,10 +125,20 @@ def from_parent(cls, parent: InfrahubEvent) -> EventMeta:
             initiator_id=parent.meta.initiator_id,
             account_id=parent.meta.account_id,
             level=parent.meta.level + 1,
-            context=parent.meta.context,
+            context=deepcopy(parent.meta.context),
             ancestors=[ParentEvent(id=parent.get_id(), name=parent.get_name())] + parent.meta.ancestors,
         )
 
+    @classmethod
+    def from_context(cls, context: InfrahubContext, branch: Branch | None = None) -> EventMeta:
+        # Create a copy of the context so local changes aren't brought back to a parent object
+        meta = cls(context=deepcopy(context))
+        meta._created_with_context = True
+        if branch:
+            meta.context.branch.name = branch.name
+            meta.context.branch.id = str(branch.get_uuid())
+        return meta
+
 
 class InfrahubEvent(BaseModel):
     meta: EventMeta = Field(..., description="Metadata for the event")
@@ -180,3 +192,10 @@ def get_message_meta(self) -> Meta:
     @computed_field
     def event_name(self) -> str:
         raise NotImplementedError("The event name has not been defined")
+
+    @model_validator(mode="after")
+    def update_context(self) -> Self:
+        """Update the context object using this event provided that the meta data was created with a context."""
+        if self.meta._created_with_context:
+            self.meta.context.set_event(self.get_name(), id=self.get_id())
+        return self

backend/infrahub/events/node_action.py

@@ -2,6 +2,7 @@
 
 from pydantic import Field, computed_field
 
+from infrahub.core.branch import Branch
 from infrahub.core.changelog.models import NodeChangelog
 from infrahub.core.constants import MutationAction
 from infrahub.message_bus import InfrahubMessage
@@ -93,6 +94,10 @@ def get_messages(self) -> list[InfrahubMessage]:
             # )
         ]
 
+    def set_context_branch(self, branch: Branch) -> None:
+        self.meta.context.branch.id = str(branch.get_uuid())
+        self.meta.context.branch.name = branch.name
+
 
 class NodeCreatedEvent(NodeMutatedEvent):
     action: MutationAction = MutationAction.CREATED

backend/infrahub/services/adapters/event/__init__.py

@@ -30,6 +30,7 @@ async def _send_bus(self, event: InfrahubEvent) -> None:
 
     async def _send_prefect(self, event: InfrahubEvent) -> None:
         emit_event(
+            id=event.meta.id,
             event=event.get_name(),
             resource=event.get_resource(),
             related=event.get_related(),

backend/tests/helpers/utils.py

@@ -1,3 +1,7 @@
+import time
+from pathlib import Path
+
+import httpx
 import pytest
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.waiting_utils import wait_for_logs
@@ -38,17 +42,38 @@ def start_prefect_server_container(
     if not INFRAHUB_USE_TEST_CONTAINERS:
         return None
 
+    prefect_base = Path(Path(__file__).parent.resolve() / "./../../infrahub/prefect_server")
     container = (
         DockerContainer(image="prefecthq/prefect:3.1.14-python3.12")
-        .with_command("prefect server start --host 0.0.0.0 --ui")
+        .with_command("uvicorn --host 0.0.0.0 --port 4200 --factory prefect_server.app:create_infrahub_prefect")
         .with_exposed_ports(PORT_PREFECT)
+        .with_volume_mapping(host=str(prefect_base), container="/opt/prefect/prefect_server", mode="ro")
+        .with_env(key="PREFECT_SERVER_SERVICES_EVENT_PERSISTER_FLUSH_INTERVAL", value="1")
     )
 
     def cleanup() -> None:
         container.stop()
 
     container.start()
-    wait_for_logs(container, "Configure Prefect to communicate with the server")
+
+    mapped_port = get_exposed_port(container, PORT_PREFECT)
+    # As our entrypoint doesn't print out any logs on startup we can't "wait for logs"
+    wait_for_prefect(port=mapped_port)
     request.addfinalizer(cleanup)
 
-    return {PORT_PREFECT: get_exposed_port(container, PORT_PREFECT)}
+    return {PORT_PREFECT: mapped_port}
+
+
+url = "http://localhost:52879/api/admin/version"  # Replace with your target URL
+
+
+def wait_for_prefect(port: int) -> None:
+    for _ in range(120):
+        try:
+            response = httpx.get(f"http://localhost:{port}/api/admin/version", timeout=1)
+            if response.status_code == 200:
+                return
+        except httpx.HTTPError:
+            time.sleep(1)
+
+    pytest.fail(reason="Prefect didn't start in an orderly fashion")