infrahub-server 1.1.6__py3-none-any.whl → 1.2.0b1__py3-none-any.whl

infrahub/graphql/mutations/repository.py

@@ -45,7 +45,7 @@ log = get_logger()
 
 class InfrahubRepositoryMutation(InfrahubMutationMixin, Mutation):
     @classmethod
-    def __init_subclass_with_meta__(cls, schema: Optional[NodeSchema] = None, _meta=None, **options):  # pylint: disable=arguments-differ
+    def __init_subclass_with_meta__(cls, schema: Optional[NodeSchema] = None, _meta=None, **options):
         # Make sure schema is a valid NodeSchema Node Class
         if not isinstance(schema, NodeSchema):
             raise ValueError(f"You need to pass a valid NodeSchema in '{cls.__name__}.Meta', received '{schema}'")
@@ -63,9 +63,9 @@ class InfrahubRepositoryMutation(InfrahubMutationMixin, Mutation):
         info: GraphQLResolveInfo,
         data: InputObjectType,
         branch: Branch,
-        database: Optional[InfrahubDatabase] = None,
+        database: Optional[InfrahubDatabase] = None,  # noqa: ARG003
     ):
-        context: GraphqlContext = info.context
+        graphql_context: GraphqlContext = info.context
 
         cleanup_payload(data)
         # Create the object in the database
@@ -74,17 +74,17 @@ class InfrahubRepositoryMutation(InfrahubMutationMixin, Mutation):
 
         # First check the connectivity to the remote repository
         # If the connectivity is not good, we remove the repository to allow the user to add a new one
-        if context.service:
+        if graphql_context.service:
             message = messages.GitRepositoryConnectivity(
                 repository_name=obj.name.value,
                 repository_location=obj.location.value,
             )
-            response = await context.service.message_bus.rpc(
+            response = await graphql_context.service.message_bus.rpc(
                 message=message, response_class=GitRepositoryConnectivityResponse
             )
 
             if response.data.success is False:
-                await obj.delete(db=context.db)
+                await obj.delete(db=graphql_context.db)
                 raise ValidationError(response.data.message)
 
         # If we are in the default branch, we set the sync status to Active
@@ -93,13 +93,13 @@ class InfrahubRepositoryMutation(InfrahubMutationMixin, Mutation):
             obj.internal_status.value = RepositoryInternalStatus.ACTIVE.value
         else:
             obj.internal_status.value = RepositoryInternalStatus.STAGING.value
-        await obj.save(db=context.db)
+        await obj.save(db=graphql_context.db)
 
         # Create the new repository in the filesystem.
         log.info("create_repository", name=obj.name.value)
         authenticated_user = None
-        if context.account_session and context.account_session.authenticated:
-            authenticated_user = context.account_session.account_id
+        if graphql_context.account_session and graphql_context.account_session.authenticated:
+            authenticated_user = graphql_context.account_session.account_id
         if obj.get_kind() == InfrahubKind.READONLYREPOSITORY:
             obj = cast(CoreReadOnlyRepository, obj)
             model = GitRepositoryAddReadOnly(
@@ -112,9 +112,11 @@ class InfrahubRepositoryMutation(InfrahubMutationMixin, Mutation):
                 internal_status=obj.internal_status.value,
                 created_by=authenticated_user,
             )
-            if context.service:
-                await context.service.workflow.submit_workflow(
-                    workflow=GIT_REPOSITORY_ADD_READ_ONLY, parameters={"model": model}
+            if graphql_context.service:
+                await graphql_context.service.workflow.submit_workflow(
+                    workflow=GIT_REPOSITORY_ADD_READ_ONLY,
+                    context=graphql_context.get_context(),
+                    parameters={"model": model},
                 )
 
         else:
@@ -130,9 +132,11 @@ class InfrahubRepositoryMutation(InfrahubMutationMixin, Mutation):
                 created_by=authenticated_user,
             )
 
-            if context.service:
-                await context.service.workflow.submit_workflow(
-                    workflow=GIT_REPOSITORY_ADD, parameters={"model": git_repo_add_model}
+            if graphql_context.service:
+                await graphql_context.service.workflow.submit_workflow(
+                    workflow=GIT_REPOSITORY_ADD,
+                    context=graphql_context.get_context(),
+                    parameters={"model": git_repo_add_model},
                 )
 
         # TODO Validate that the creation of the repository went as expected
@@ -145,15 +149,15 @@ class InfrahubRepositoryMutation(InfrahubMutationMixin, Mutation):
         info: GraphQLResolveInfo,
         data: InputObjectType,
         branch: Branch,
-        database: Optional[InfrahubDatabase] = None,
+        database: Optional[InfrahubDatabase] = None,  # noqa: ARG003
         node: Optional[Node] = None,
     ):
-        context: GraphqlContext = info.context
+        graphql_context: GraphqlContext = info.context
 
         cleanup_payload(data)
         if not node:
             node: CoreReadOnlyRepository | CoreRepository = await NodeManager.get_one_by_id_or_default_filter(
-                db=context.db,
+                db=graphql_context.db,
                 kind=cls._meta.schema.kind,
                 id=data.get("id"),
                 branch=branch,
@@ -161,7 +165,7 @@ class InfrahubRepositoryMutation(InfrahubMutationMixin, Mutation):
                 include_source=True,
             )
         if node.get_kind() != InfrahubKind.READONLYREPOSITORY:
-            return await super().mutate_update(info, data, branch, database=context.db, node=node)
+            return await super().mutate_update(info, data, branch, database=graphql_context.db, node=node)
 
         node = cast(CoreReadOnlyRepository, node)
         current_commit = node.commit.value
@@ -173,7 +177,7 @@ class InfrahubRepositoryMutation(InfrahubMutationMixin, Mutation):
         if data.ref and data.ref.value:
             new_ref = data.ref.value
 
-        obj, result = await super().mutate_update(info, data, branch, database=context.db, node=node)
+        obj, result = await super().mutate_update(info, data, branch, database=graphql_context.db, node=node)
         obj = cast(CoreReadOnlyRepository, obj)
 
         send_update_message = (new_commit and new_commit != current_commit) or (new_ref and new_ref != current_ref)
@@ -196,9 +200,11 @@ class InfrahubRepositoryMutation(InfrahubMutationMixin, Mutation):
             infrahub_branch_name=branch.name,
             infrahub_branch_id=str(branch.get_uuid()),
         )
-        if context.service:
-            await context.service.workflow.submit_workflow(
-                workflow=GIT_REPOSITORIES_PULL_READ_ONLY, parameters={"model": model}
+        if graphql_context.service:
+            await graphql_context.service.workflow.submit_workflow(
+                workflow=GIT_REPOSITORIES_PULL_READ_ONLY,
+                context=graphql_context.get_context(),
+                parameters={"model": model},
             )
         return obj, result
 
@@ -226,15 +232,15 @@ class ProcessRepository(Mutation):
     @classmethod
     async def mutate(
         cls,
-        root: dict,  # pylint: disable=unused-argument
+        root: dict,  # noqa: ARG003
         info: GraphQLResolveInfo,
         data: IdentifierInput,
     ) -> dict[str, bool]:
-        context: GraphqlContext = info.context
-        branch = context.branch
+        graphql_context: GraphqlContext = info.context
+        branch = graphql_context.branch
         repository_id = str(data.id)
         repo: CoreReadOnlyRepository | CoreRepository = await NodeManager.get_one_by_id_or_default_filter(
-            db=context.db,
+            db=graphql_context.db,
             kind=InfrahubKind.GENERICREPOSITORY,
             id=str(data.id),
             branch=branch,
@@ -247,8 +253,8 @@ class ProcessRepository(Mutation):
             commit=str(repo.commit.value),
             infrahub_branch_name=branch.name,
         )
-        workflow = await context.active_service.workflow.submit_workflow(
-            workflow=GIT_REPOSITORIES_IMPORT_OBJECTS, parameters={"model": model}
+        workflow = await graphql_context.active_service.workflow.submit_workflow(
+            workflow=GIT_REPOSITORIES_IMPORT_OBJECTS, context=graphql_context.get_context(), parameters={"model": model}
         )
         task = {"id": workflow.id}
         return cls(ok=True, task=task)
@@ -264,15 +270,15 @@ class ValidateRepositoryConnectivity(Mutation):
     @classmethod
     async def mutate(
         cls,
-        root: dict,  # pylint: disable=unused-argument
+        root: dict,  # noqa: ARG003
         info: GraphQLResolveInfo,
         data: IdentifierInput,
     ) -> dict[str, Any]:
-        context: GraphqlContext = info.context
-        branch = context.branch
+        graphql_context: GraphqlContext = info.context
+        branch = graphql_context.branch
         repository_id = str(data.id)
         repo: CoreReadOnlyRepository | CoreRepository = await NodeManager.get_one_by_id_or_default_filter(
-            db=context.db,
+            db=graphql_context.db,
             kind=InfrahubKind.GENERICREPOSITORY,
             id=repository_id,
             branch=branch,
@@ -282,8 +288,8 @@ class ValidateRepositoryConnectivity(Mutation):
             repository_name=str(repo.name.value),
             repository_location=str(repo.location.value),
         )
-        if context.service:
-            response = await context.service.message_bus.rpc(
+        if graphql_context.service:
+            response = await graphql_context.service.message_bus.rpc(
                 message=message, response_class=GitRepositoryConnectivityResponse
             )
 

infrahub/graphql/mutations/resource_manager.py

@@ -61,11 +61,11 @@ class IPPrefixPoolGetResource(Mutation):
     @classmethod
     async def mutate(
         cls,
-        root: dict,  # pylint: disable=unused-argument
+        root: dict,  # noqa: ARG003
         info: GraphQLResolveInfo,
         data: InputObjectType,
     ) -> Self:
-        context: GraphqlContext = info.context
+        graphql_context: GraphqlContext = info.context
 
         member_type = data.get("member_type", None)
         allowed_member_types = [t.value for t in PrefixMemberType]
@@ -73,15 +73,15 @@ class IPPrefixPoolGetResource(Mutation):
             raise QueryValidationError(f"Invalid member_type value, allowed values are {allowed_member_types}")
 
         obj: CoreIPPrefixPool = await registry.manager.find_object(  # type: ignore[assignment]
-            db=context.db,
+            db=graphql_context.db,
             kind=InfrahubKind.IPPREFIXPOOL,
             id=data.get("id"),
             hfid=data.get("hfid"),
-            branch=context.branch,
+            branch=graphql_context.branch,
         )
         resource = await obj.get_resource(
-            db=context.db,
-            branch=context.branch,
+            db=graphql_context.db,
+            branch=graphql_context.branch,
             identifier=data.get("identifier", None),
             prefixlen=data.get("prefix_length", None),
             member_type=member_type,
@@ -95,8 +95,8 @@ class IPPrefixPoolGetResource(Mutation):
                 "id": resource.id,
                 "kind": resource.get_kind(),
                 "identifier": data.get("identifier", None),
-                "display_label": await resource.render_display_label(db=context.db),
-                "branch": context.branch.name,
+                "display_label": await resource.render_display_label(db=graphql_context.db),
+                "branch": graphql_context.branch.name,
             },
         }
 
@@ -113,26 +113,26 @@ class IPAddressPoolGetResource(Mutation):
     @classmethod
     async def mutate(
         cls,
-        root: dict,  # pylint: disable=unused-argument
+        root: dict,  # noqa: ARG003
         info: GraphQLResolveInfo,
         data: dict[str, Any],
     ) -> Self:
-        context: GraphqlContext = info.context
+        graphql_context: GraphqlContext = info.context
 
         obj: CoreIPAddressPool = await registry.manager.find_object(
-            db=context.db,
+            db=graphql_context.db,
             kind=InfrahubKind.IPADDRESSPOOL,
             id=data.get("id"),
             hfid=data.get("hfid"),
-            branch=context.branch,
+            branch=graphql_context.branch,
         )
         resource = await obj.get_resource(
-            db=context.db,
-            branch=context.branch,
-            identifier=data.get("identifier", None),
-            prefixlen=data.get("prefix_length", None),
-            address_type=data.get("address_type", None),
-            data=data.get("data", None),
+            db=graphql_context.db,
+            branch=graphql_context.branch,
+            identifier=data.get("identifier"),
+            prefixlen=data.get("prefix_length"),
+            address_type=data.get("address_type"),
+            data=data.get("data"),
         )
 
         result = {
@@ -140,9 +140,9 @@ class IPAddressPoolGetResource(Mutation):
             "node": {
                 "id": resource.id,
                 "kind": resource.get_kind(),
-                "identifier": data.get("identifier", None),
-                "display_label": await resource.render_display_label(db=context.db),
-                "branch": context.branch.name,
+                "identifier": data.get("identifier"),
+                "display_label": await resource.render_display_label(db=graphql_context.db),
+                "branch": graphql_context.branch.name,
             },
         }
 
@@ -151,7 +151,7 @@ class IPAddressPoolGetResource(Mutation):
 
 class InfrahubNumberPoolMutation(InfrahubMutationMixin, Mutation):
     @classmethod
-    def __init_subclass_with_meta__(  # pylint: disable=arguments-differ
+    def __init_subclass_with_meta__(
         cls,
         schema: NodeSchema | None = None,
         _meta: InfrahubMutationOptions | None = None,
@@ -174,7 +174,7 @@ class InfrahubNumberPoolMutation(InfrahubMutationMixin, Mutation):
         info: GraphQLResolveInfo,
         data: InputObjectType,
         branch: Branch,
-        database: InfrahubDatabase | None = None,
+        database: InfrahubDatabase | None = None,  # noqa: ARG003
     ) -> Any:
         try:
             pool_node = registry.get_node_schema(name=data["node"].value)
@@ -204,16 +204,16 @@ class InfrahubNumberPoolMutation(InfrahubMutationMixin, Mutation):
         info: GraphQLResolveInfo,
         data: InputObjectType,
         branch: Branch,
-        database: InfrahubDatabase | None = None,
+        database: InfrahubDatabase | None = None,  # noqa: ARG003
         node: Node | None = None,
     ) -> tuple[Node, Self]:
         if (data.get("node") and data.get("node").value) or (
             data.get("node_attribute") and data.get("node_attribute").value
         ):
             raise ValidationError(input_value="The fields 'node' or 'node_attribute' can't be changed.")
-        context: GraphqlContext = info.context
+        graphql_context: GraphqlContext = info.context
 
-        async with context.db.start_transaction() as dbt:
+        async with graphql_context.db.start_transaction() as dbt:
             number_pool, result = await super().mutate_update(
                 info=info, data=data, branch=branch, database=dbt, node=node
             )

infrahub/graphql/mutations/schema.py

@@ -13,6 +13,8 @@ from infrahub.database import InfrahubDatabase, retry_db_transaction
 from infrahub.events import EventMeta
 from infrahub.events.schema_action import SchemaUpdatedEvent
 from infrahub.exceptions import ValidationError
+from infrahub.graphql.context import apply_external_context
+from infrahub.graphql.types.context import ContextInput
 from infrahub.log import get_log_data, get_logger
 from infrahub.worker import WORKER_IDENTITY
 
@@ -21,6 +23,7 @@ from ..types import DropdownFields
 if TYPE_CHECKING:
     from graphql import GraphQLResolveInfo
 
+    from infrahub.context import InfrahubContext
     from infrahub.core.branch import Branch
     from infrahub.services import InfrahubServices
 
@@ -50,6 +53,7 @@ class SchemaDropdownAddInput(SchemaDropdownRemoveInput):
 class SchemaDropdownAdd(Mutation):
     class Arguments:
         data = SchemaDropdownAddInput(required=True)
+        context = ContextInput(required=False)
 
     ok = Boolean()
     object = Field(DropdownFields)
@@ -58,13 +62,16 @@ class SchemaDropdownAdd(Mutation):
     @retry_db_transaction(name="schema_dropdown_add")
     async def mutate(
         cls,
-        root: dict,  # pylint: disable=unused-argument
+        root: dict,  # noqa: ARG003
         info: GraphQLResolveInfo,
         data: SchemaDropdownAddInput,
+        context: ContextInput | None = None,
     ) -> Self:
-        context: GraphqlContext = info.context
+        graphql_context: GraphqlContext = info.context
 
-        kind = context.db.schema.get(name=str(data.kind), branch=context.branch.name)
+        await apply_external_context(graphql_context=graphql_context, context_input=context)
+
+        kind = graphql_context.db.schema.get(name=str(data.kind), branch=graphql_context.branch.name)
         attribute = str(data.attribute)
         validate_kind_dropdown(kind=kind, attribute=attribute)
         dropdown = str(data.dropdown)
@@ -80,13 +87,14 @@ class SchemaDropdownAdd(Mutation):
 
         await update_registry(
             kind=kind,
-            branch=context.branch,
-            db=context.db,
-            account_id=context.active_account_session.account_id,
-            service=context.active_service,
+            branch=graphql_context.branch,
+            db=graphql_context.db,
+            account_id=graphql_context.active_account_session.account_id,
+            service=graphql_context.active_service,
+            context=graphql_context.get_context(),
         )
 
-        kind = context.db.schema.get(name=str(data.kind), branch=context.branch.name)
+        kind = graphql_context.db.schema.get(name=str(data.kind), branch=graphql_context.branch.name)
         attrib = kind.get_attribute(attribute)
         dropdown_entry = {}
         success = False
@@ -107,6 +115,7 @@ class SchemaDropdownAdd(Mutation):
 class SchemaDropdownRemove(Mutation):
     class Arguments:
         data = SchemaDropdownRemoveInput(required=True)
+        context = ContextInput(required=False)
 
     ok = Boolean()
 
@@ -114,19 +123,24 @@ class SchemaDropdownRemove(Mutation):
     @retry_db_transaction(name="schema_dropdown_remove")
     async def mutate(
         cls,
-        root: dict,  # pylint: disable=unused-argument
+        root: dict,  # noqa: ARG003
         info: GraphQLResolveInfo,
         data: SchemaDropdownRemoveInput,
+        context: ContextInput | None = None,
     ) -> dict[str, bool]:
-        context: GraphqlContext = info.context
+        graphql_context: GraphqlContext = info.context
 
-        kind = context.db.schema.get(name=str(data.kind), branch=context.branch.name)
+        kind = graphql_context.db.schema.get(name=str(data.kind), branch=graphql_context.branch.name)
+        await apply_external_context(graphql_context=graphql_context, context_input=context)
 
         attribute = str(data.attribute)
         validate_kind_dropdown(kind=kind, attribute=attribute)
         dropdown = str(data.dropdown)
         nodes_with_dropdown = await NodeManager.query(
-            db=context.db, schema=kind.kind, filters={f"{attribute}__value": dropdown}, branch=context.branch
+            db=graphql_context.db,
+            schema=kind.kind,
+            filters={f"{attribute}__value": dropdown},
+            branch=graphql_context.branch,
         )
         if nodes_with_dropdown:
             raise ValidationError(f"There are still {kind.kind} objects using this dropdown")
@@ -143,10 +157,11 @@ class SchemaDropdownRemove(Mutation):
 
         await update_registry(
             kind=kind,
-            branch=context.branch,
-            db=context.db,
-            account_id=context.active_account_session.account_id,
-            service=context.active_service,
+            branch=graphql_context.branch,
+            db=graphql_context.db,
+            account_id=graphql_context.active_account_session.account_id,
+            service=graphql_context.active_service,
+            context=graphql_context.get_context(),
         )
 
         return {"ok": True}
@@ -155,6 +170,7 @@ class SchemaDropdownRemove(Mutation):
 class SchemaEnumAdd(Mutation):
     class Arguments:
         data = SchemaEnumInput(required=True)
+        context = ContextInput(required=False)
 
     ok = Boolean()
 
@@ -162,13 +178,15 @@ class SchemaEnumAdd(Mutation):
     @retry_db_transaction(name="schema_dropdown_add")
     async def mutate(
         cls,
-        root: dict,  # pylint: disable=unused-argument
+        root: dict,  # noqa: ARG003
         info: GraphQLResolveInfo,
         data: SchemaEnumInput,
+        context: ContextInput | None = None,
     ) -> dict[str, bool]:
-        context: GraphqlContext = info.context
+        graphql_context: GraphqlContext = info.context
 
-        kind = context.db.schema.get(name=str(data.kind), branch=context.branch.name)
+        kind = graphql_context.db.schema.get(name=str(data.kind), branch=graphql_context.branch.name)
+        await apply_external_context(graphql_context=graphql_context, context_input=context)
 
         attribute = str(data.attribute)
         enum = str(data.enum)
@@ -184,10 +202,11 @@ class SchemaEnumAdd(Mutation):
 
         await update_registry(
             kind=kind,
-            branch=context.branch,
-            db=context.db,
-            account_id=context.active_account_session.account_id,
-            service=context.active_service,
+            branch=graphql_context.branch,
+            db=graphql_context.db,
+            account_id=graphql_context.active_account_session.account_id,
+            service=graphql_context.active_service,
+            context=graphql_context.get_context(),
         )
 
         return {"ok": True}
@@ -196,6 +215,7 @@ class SchemaEnumAdd(Mutation):
 class SchemaEnumRemove(Mutation):
     class Arguments:
         data = SchemaEnumInput(required=True)
+        context = ContextInput(required=False)
 
     ok = Boolean()
 
@@ -203,19 +223,24 @@ class SchemaEnumRemove(Mutation):
     @retry_db_transaction(name="schema_enum_remove")
     async def mutate(
         cls,
-        root: dict,  # pylint: disable=unused-argument
+        root: dict,  # noqa: ARG003
         info: GraphQLResolveInfo,
         data: SchemaEnumInput,
+        context: ContextInput | None = None,
     ) -> dict[str, bool]:
-        context: GraphqlContext = info.context
+        graphql_context: GraphqlContext = info.context
 
-        kind = context.db.schema.get(name=str(data.kind), branch=context.branch.name)
+        kind = graphql_context.db.schema.get(name=str(data.kind), branch=graphql_context.branch.name)
+        await apply_external_context(graphql_context=graphql_context, context_input=context)
 
         attribute = str(data.attribute)
         enum = str(data.enum)
         validate_kind_enum(kind=kind, attribute=attribute)
         nodes_with_enum = await NodeManager.query(
-            db=context.db, schema=kind.kind, filters={f"{attribute}__value": enum}, branch=context.branch
+            db=graphql_context.db,
+            schema=kind.kind,
+            filters={f"{attribute}__value": enum},
+            branch=graphql_context.branch,
         )
         if nodes_with_enum:
             raise ValidationError(f"There are still {kind.kind} objects using this enum")
@@ -232,10 +257,11 @@ class SchemaEnumRemove(Mutation):
 
         await update_registry(
             kind=kind,
-            branch=context.branch,
-            db=context.db,
-            account_id=context.active_account_session.account_id,
-            service=context.active_service,
+            branch=graphql_context.branch,
+            db=graphql_context.db,
+            account_id=graphql_context.active_account_session.account_id,
+            service=graphql_context.active_service,
+            context=graphql_context.get_context(),
         )
 
         return {"ok": True}
@@ -268,7 +294,12 @@ def validate_kind(kind: Union[GenericSchema, NodeSchema], attribute: str) -> Non
 
 
 async def update_registry(
-    kind: NodeSchema, db: InfrahubDatabase, branch: Branch, account_id: str, service: InfrahubServices
+    kind: NodeSchema,
+    db: InfrahubDatabase,
+    branch: Branch,
+    account_id: str,
+    service: InfrahubServices,
+    context: InfrahubContext,
 ) -> None:
     async with lock.registry.global_schema_lock():
         branch_schema = registry.schema.get_schema_branch(name=branch.name)
@@ -296,12 +327,14 @@ async def update_registry(
             log_data = get_log_data()
             request_id = log_data.get("request_id", "")
             event = SchemaUpdatedEvent(
-                branch=branch.name,
+                branch_name=branch.name,
                 schema_hash=branch.active_schema_hash.main,
                 meta=EventMeta(
                     initiator_id=WORKER_IDENTITY,
                     request_id=request_id,
                     account_id=account_id,
+                    branch=branch,
+                    context=context,
                 ),
             )
             await service.event.send(event=event)

infrahub/graphql/mutations/tasks.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from prefect import flow
 
+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.diff.coordinator import DiffCoordinator
@@ -13,15 +14,14 @@ from infrahub.core.validators.models.validate_migration import SchemaValidateMig
 from infrahub.core.validators.tasks import schema_validate_migrations
 from infrahub.dependencies.registry import get_component_registry
 from infrahub.exceptions import ValidationError
-from infrahub.services import services
+from infrahub.services import InfrahubServices  # noqa: TC001  needed for prefect flow
 from infrahub.workflows.catalogue import BRANCH_MERGE
-from infrahub.workflows.utils import add_branch_tag
+from infrahub.workflows.utils import add_tags
 
 
 @flow(name="merge-branch-mutation", flow_run_name="Merge branch graphQL mutation")
-async def merge_branch_mutation(branch: str) -> None:
-    service = services.service
-    await add_branch_tag(branch_name=branch)
+async def merge_branch_mutation(branch: str, context: InfrahubContext, service: InfrahubServices) -> None:
+    await add_tags(branches=[branch])
 
     async with service.database.start_session() as db:
         obj = await Branch.get_by_name(db=db, name=branch)
@@ -31,14 +31,17 @@ async def merge_branch_mutation(branch: str) -> None:
         diff_coordinator = await component_registry.get_component(DiffCoordinator, db=db, branch=obj)
         diff_repository = await component_registry.get_component(DiffRepository, db=db, branch=obj)
         diff_merger = await component_registry.get_component(DiffMerger, db=db, branch=obj)
-        enriched_diff = await diff_coordinator.update_branch_diff_and_return(base_branch=base_branch, diff_branch=obj)
-        if enriched_diff.get_all_conflicts():
+        enriched_diff_metadata = await diff_coordinator.update_branch_diff(base_branch=base_branch, diff_branch=obj)
+        async for _ in diff_repository.get_all_conflicts_for_diff(
+            diff_branch_name=enriched_diff_metadata.diff_branch_name, diff_id=enriched_diff_metadata.uuid
+        ):
+            # if there are any conflicts, raise the error
             raise ValidationError(
                 f"Branch {obj.name} contains conflicts with the default branch."
                 " Please create a Proposed Change to resolve the conflicts or manually update them before merging."
             )
         node_diff_field_summaries = await diff_repository.get_node_field_summaries(
-            diff_branch_name=enriched_diff.diff_branch_name, diff_id=enriched_diff.uuid
+            diff_branch_name=enriched_diff_metadata.diff_branch_name, diff_id=enriched_diff_metadata.uuid
         )
 
         merger = BranchMerger(
@@ -57,10 +60,13 @@ async def merge_branch_mutation(branch: str) -> None:
 
         if constraints:
             responses = await schema_validate_migrations(
-                message=SchemaValidateMigrationData(branch=obj, schema_branch=candidate_schema, constraints=constraints)
+                message=SchemaValidateMigrationData(
+                    branch=obj, schema_branch=candidate_schema, constraints=constraints
+                ),
+                service=service,
             )
             error_messages = [violation.message for response in responses for violation in response.violations]
             if error_messages:
                 raise ValidationError(",\n".join(error_messages))
 
-        await service.workflow.execute_workflow(workflow=BRANCH_MERGE, parameters={"branch": obj.name})
+        await service.workflow.execute_workflow(workflow=BRANCH_MERGE, context=context, parameters={"branch": obj.name})

infrahub/graphql/parser.py

@@ -162,7 +162,7 @@ class GraphQLExtractor:
                         FieldNode(
                             kind="field",
                             name=NameNode(kind="name", value=sub_node.key),
-                            selection_set=SelectionSetNode(selections=tuple([sub_node.node])),
+                            selection_set=SelectionSetNode(selections=(sub_node.node,)),
                         )
                     )
                     selection_set.selections = tuple(selections)