Merge pull request #5729 from opsmill/pog-member-of-group-event-improvements

Improve relationship group events for member of groups

Author: ogenstad

backend/infrahub/graphql/mutations/relationship.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from enum import Enum
 from typing import TYPE_CHECKING, Self
 
 from graphene import Boolean, InputField, InputObjectType, List, Mutation, String
@@ -9,6 +10,7 @@
 from infrahub.core.changelog.models import NodeChangelog
 from infrahub.core.constants import InfrahubKind, MutationAction, RelationshipCardinality
 from infrahub.core.manager import NodeManager
+from infrahub.core.query.node import NodeGetKindQuery
 from infrahub.core.query.relationship import (
     RelationshipGetPeerQuery,
     RelationshipPeerData,
@@ -33,6 +35,12 @@
 RELATIONSHIP_PEERS_TO_IGNORE = [InfrahubKind.NODE]
 
 
+class GroupUpdateType(str, Enum):
+    NONE = "none"
+    MEMBERS = "members"
+    MEMBER_OF_GROUPS = "member_of_groups"
+
+
 class RelationshipNodesInput(InputObjectType):
     id = InputField(String(required=True), description="ID of the node at the source of the relationship")
     name = InputField(String(required=True), description="Name of the relationship to add or remove nodes")
@@ -65,16 +73,24 @@ async def mutate(  # noqa: PLR0915
             raise NodeNotFoundError(node_type="node", identifier=input_id, branch_name=graphql_context.branch.name)
 
         # Check if the name of the relationship provided exist for this node and is of cardinality Many
-        if relationship_name not in source._schema.relationship_names:
+        if relationship_name not in source.get_schema().relationship_names:
             raise ValidationError(
                 {"name": f"'{relationship_name}' is not a valid relationship for '{source.get_kind()}'"}
             )
 
-        rel_schema = source._schema.get_relationship(name=relationship_name)
+        rel_schema = source.get_schema().get_relationship(name=relationship_name)
         if rel_schema.cardinality != RelationshipCardinality.MANY:
             raise ValidationError({"name": f"'{relationship_name}' must be a relationship of cardinality Many"})
 
-        is_group_event = rel_schema.identifier == "group_member" and "CoreGroup" in source._schema.inherit_from
+        group_event_type = GroupUpdateType.NONE
+        if rel_schema.identifier == "group_member":
+            if "CoreGroup" in source.get_schema().inherit_from and relationship_name == "members":
+                # Updating members of a group
+                group_event_type = GroupUpdateType.MEMBERS
+
+            elif relationship_name == "member_of_groups":
+                # Modifying the membership of the current node
+                group_event_type = GroupUpdateType.MEMBER_OF_GROUPS
 
         # Query the node in the database and validate that all of them exist and are if the correct kind
         node_ids: list[str] = [node_data["id"] for node_data in data.get("nodes") if "id" in node_data]
@@ -93,7 +109,9 @@ async def mutate(  # noqa: PLR0915
             if rel_schema.peer not in node.get_labels():
                 raise ValidationError(f"{node_id!r} {node.get_kind()!r} is not a valid peer for '{rel_schema.peer}'")
 
-            peer_relationships = [rel for rel in node._schema.relationships if rel.identifier == rel_schema.identifier]
+            peer_relationships = [
+                rel for rel in node.get_schema().relationships if rel.identifier == rel_schema.identifier
+            ]
             if (
                 rel_schema.identifier
                 and len(peer_relationships) == 1
@@ -120,7 +138,7 @@ async def mutate(  # noqa: PLR0915
         await query.execute(db=graphql_context.db)
         existing_peers: dict[str, RelationshipPeerData] = {str(peer.peer_id): peer for peer in query.get_peers()}
         async with graphql_context.db.start_transaction() as db:
-            members = []
+            peers: list[EventNode] = []
             if cls.__name__ == "RelationshipAdd":
                 for node_data in data.get("nodes"):
                     # Instantiate and resolve a relationship
@@ -130,7 +148,7 @@ async def mutate(  # noqa: PLR0915
                     await rel.resolve(db=db)
                     # Save it only if it does not exist
                     if rel.get_peer_id() not in existing_peers.keys():
-                        members.append(EventNode(id=rel.get_peer_id(), kind=rel.get_peer_kind()))
+                        peers.append(EventNode(id=rel.get_peer_id(), kind=rel.get_peer_kind()))
                         node_changelog.create_relationship(relationship=rel)
                         await rel.save(db=db)
 
@@ -142,13 +160,13 @@ async def mutate(  # noqa: PLR0915
                         # it would be more query efficient
                         rel = Relationship(schema=rel_schema, branch=graphql_context.branch, node=source)
                         await rel.load(db=db, data=existing_peers[node_data.get("id")])
-                        members.append(EventNode(id=rel.get_peer_id(), kind=rel.get_peer_kind()))
+                        peers.append(EventNode(id=rel.get_peer_id(), kind=rel.get_peer_kind()))
                         node_changelog.delete_relationship(relationship=rel)
                         await rel.delete(db=db)
 
         if config.SETTINGS.broker.enable and graphql_context.background:
             event = NodeMutatedEvent(
-                kind=source._schema.kind,
+                kind=source.get_schema().kind,
                 node_id=source.id,
                 data=node_changelog,
                 action=MutationAction.UPDATED,
@@ -158,18 +176,43 @@ async def mutate(  # noqa: PLR0915
                 ),
             )
             graphql_context.background.add_task(graphql_context.active_service.event.send, event)
-            if is_group_event:
+            if group_event_type == GroupUpdateType.MEMBERS:
                 if cls.__name__ == "RelationshipAdd":
                     group_add_event = GroupMemberAddedEvent(
-                        node_id=source.id, kind=source._schema.kind, members=members
+                        node_id=source.id, kind=source.get_schema().kind, members=peers
                     )
                     graphql_context.background.add_task(graphql_context.active_service.event.send, group_add_event)
                 elif cls.__name__ == "RelationshipRemove":
                     group_remove_event = GroupMemberRemovedEvent(
-                        node_id=source.id, kind=source._schema.kind, members=members
+                        node_id=source.id, kind=source.get_schema().kind, members=peers
                     )
                     graphql_context.background.add_task(graphql_context.active_service.event.send, group_remove_event)
-
+            elif group_event_type == GroupUpdateType.MEMBER_OF_GROUPS:
+                group_ids = [node.id for node in peers]
+                async with graphql_context.db.start_session() as db:
+                    node_kind_query = await NodeGetKindQuery.init(db=db, branch=graphql_context.branch, ids=group_ids)
+                    await node_kind_query.execute(db=db)
+                    node_kind_map = await node_kind_query.get_node_kind_map()
+
+                    for node_id, node_kind in node_kind_map.items():
+                        if cls.__name__ == "RelationshipAdd":
+                            group_add_event = GroupMemberAddedEvent(
+                                node_id=node_id,
+                                kind=node_kind,
+                                members=[EventNode(id=source.get_id(), kind=source.get_kind())],
+                            )
+                            graphql_context.background.add_task(
+                                graphql_context.active_service.event.send, group_add_event
+                            )
+                        elif cls.__name__ == "RelationshipRemove":
+                            group_remove_event = GroupMemberRemovedEvent(
+                                node_id=node_id,
+                                kind=node_kind,
+                                members=[EventNode(id=source.get_id(), kind=source.get_kind())],
+                            )
+                            graphql_context.background.add_task(
+                                graphql_context.active_service.event.send, group_remove_event
+                            )
         return cls(ok=True)  # type: ignore[call-arg]
 
 

backend/tests/unit/graphql/test_mutation_relationship.py

@@ -385,7 +385,9 @@ async def test_relationship_groups_add(
         g1.id,
         g2.id,
     )
-
+    memory_event = MemoryInfrahubEvent()
+    service = await InfrahubServices.new(event=memory_event)
+    gql_params = await prepare_graphql_params(db=db, include_subscription=False, branch=default_branch, service=service)
     result = await graphql(
         schema=gql_params.schema,
         source=query,
@@ -404,6 +406,18 @@ async def test_relationship_groups_add(
     members = await group2.members.get(db=db)
     assert len(members) == 2
 
+    assert gql_params.context.background
+    await gql_params.context.background()
+
+    node_event = memory_event.events[0]
+    group_event = memory_event.events[1]
+    assert isinstance(node_event, NodeMutatedEvent)
+    assert isinstance(group_event, GroupMemberAddedEvent)
+    # While we mutated the relationship for c3 we expect the group member event to reflect that of the g1
+    # group as c3 was already a member of g2 we don't see an event for that entry
+    assert group_event.node_id == g1.id
+    assert [member.id for member in group_event.members] == [c3.id]
+
 
 async def test_relationship_groups_remove(
     db: InfrahubDatabase,
@@ -479,6 +493,10 @@ async def test_relationship_groups_remove(
         g1.id,
         g2.id,
     )
+    memory_event = MemoryInfrahubEvent()
+    service = await InfrahubServices.new(event=memory_event)
+
+    gql_params = await prepare_graphql_params(db=db, include_subscription=False, branch=default_branch, service=service)
 
     result = await graphql(
         schema=gql_params.schema,
@@ -498,6 +516,17 @@ async def test_relationship_groups_remove(
     members = await group2.members.get(db=db)
     assert len(members) == 1
 
+    assert gql_params.context.background
+    await gql_params.context.background()
+
+    node_event = memory_event.events[0]
+    group_event = memory_event.events[1]
+    assert isinstance(node_event, NodeMutatedEvent)
+    assert isinstance(group_event, GroupMemberRemovedEvent)
+    # The c3 node is not member of g1 so we only expect to see a group event for the g2 group
+    assert group_event.node_id == g2.id
+    assert [member.id for member in group_event.members] == [c3.id]
+
 
 async def test_relationship_groups_add_remove(db: InfrahubDatabase, default_branch: Branch, car_person_generics_data):
     c1 = car_person_generics_data["c1"]