infrahub-server 1.2.0rc0__py3-none-any.whl → 1.2.1__py3-none-any.whl

infrahub/graphql/context.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from infrahub.core.constants import GlobalPermissions, InfrahubKind
+from infrahub.core.manager import NodeManager
+from infrahub.exceptions import NodeNotFoundError, ValidationError
+from infrahub.permissions.globals import define_global_permission_from_branch
+
+if TYPE_CHECKING:
+    from .initialization import GraphqlContext
+    from .types.context import ContextInput
+
+
+async def apply_external_context(graphql_context: GraphqlContext, context_input: ContextInput | None) -> None:
+    """Applies context provided by an external mutation to the GraphQL context"""
+    if not context_input or not context_input.account:
+        return
+
+    if graphql_context.active_account_session.account_id == context_input.account.id:
+        # If the account_id from the request context is the same as the current account
+        # there's no point moving forward with other checks to override the current
+        # context we can just continue with what is already there.
+        return
+
+    permission = define_global_permission_from_branch(
+        permission=GlobalPermissions.OVERRIDE_CONTEXT, branch_name=graphql_context.branch.name
+    )
+
+    graphql_context.active_permissions.raise_for_permission(permission=permission)
+
+    try:
+        account = await NodeManager.get_one_by_id_or_default_filter(
+            db=graphql_context.db, id=str(context_input.account.id), kind=InfrahubKind.GENERICACCOUNT
+        )
+    except NodeNotFoundError as exc:
+        raise ValidationError(input_value="Unable to set context for account that doesn't exist") from exc
+
+    graphql_context.active_account_session.account_id = account.id

infrahub/graphql/initialization.py

@@ -9,6 +9,7 @@ from infrahub.context import InfrahubContext
 from infrahub.core import registry
 from infrahub.core.timestamp import Timestamp
 from infrahub.exceptions import InitializationError
+from infrahub.graphql.resolvers.many_relationship import ManyRelationshipResolver
 from infrahub.graphql.resolvers.single_relationship import SingleRelationshipResolver
 from infrahub.permissions import PermissionManager
 
@@ -36,6 +37,7 @@ class GraphqlContext:
     branch: Branch
     types: dict
     single_relationship_resolver: SingleRelationshipResolver
+    many_relationship_resolver: ManyRelationshipResolver
     service: InfrahubServices | None = None
     at: Timestamp | None = None
     related_node_ids: set | None = None
@@ -111,6 +113,7 @@ async def prepare_graphql_params(
             db=db,
             branch=branch,
             single_relationship_resolver=SingleRelationshipResolver(),
+            many_relationship_resolver=ManyRelationshipResolver(),
             at=Timestamp(at),
             types=gqlm.get_graphql_types(),
             related_node_ids=set(),

infrahub/graphql/loaders/node.py

@@ -10,6 +10,8 @@ from infrahub.core.node import Node
 from infrahub.core.timestamp import Timestamp
 from infrahub.database import InfrahubDatabase
 
+from .shared import to_frozen_set
+
 
 @dataclass
 class GetManyParams:
@@ -68,15 +70,3 @@ class NodeDataLoader(DataLoader[str, Node | None]):
         for node_id in keys:
             results.append(nodes_by_id.get(node_id, None))
         return results
-
-
-def to_frozen_set(to_freeze: dict[str, Any]) -> frozenset:
-    freezing_dict = {}
-    for k, v in to_freeze.items():
-        if isinstance(v, dict):
-            freezing_dict[k] = to_frozen_set(v)
-        elif isinstance(v, list | set):
-            freezing_dict[k] = frozenset(v)
-        else:
-            freezing_dict[k] = v
-    return frozenset(freezing_dict)

infrahub/graphql/loaders/peers.py

@@ -0,0 +1,77 @@
+from dataclasses import dataclass
+from typing import Any
+
+from aiodataloader import DataLoader
+
+from infrahub.core.branch.models import Branch
+from infrahub.core.manager import NodeManager
+from infrahub.core.relationship.model import Relationship
+from infrahub.core.schema.relationship_schema import RelationshipSchema
+from infrahub.core.timestamp import Timestamp
+from infrahub.database import InfrahubDatabase
+
+from .shared import to_frozen_set
+
+
+@dataclass
+class QueryPeerParams:
+    branch: Branch | str
+    source_kind: str
+    schema: RelationshipSchema
+    filters: dict[str, Any]
+    fields: dict | None = None
+    at: Timestamp | str | None = None
+    branch_agnostic: bool = False
+
+    def __hash__(self) -> int:
+        frozen_fields: frozenset | None = None
+        if self.fields:
+            frozen_fields = to_frozen_set(self.fields)
+        frozen_filters = to_frozen_set(self.filters)
+        timestamp = Timestamp(self.at)
+        branch = self.branch.name if isinstance(self.branch, Branch) else self.branch
+        hash_str = "|".join(
+            [
+                str(hash(frozen_fields)),
+                str(hash(frozen_filters)),
+                timestamp.to_string(),
+                branch,
+                self.schema.name,
+                str(self.source_kind),
+                str(self.branch_agnostic),
+            ]
+        )
+        return hash(hash_str)
+
+
+class PeerRelationshipsDataLoader(DataLoader[str, list[Relationship]]):
+    def __init__(self, db: InfrahubDatabase, query_params: QueryPeerParams, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self.query_params = query_params
+        self.db = db
+
+    async def batch_load_fn(self, keys: list[Any]) -> list[list[Relationship]]:  # pylint: disable=method-hidden
+        async with self.db.start_session() as db:
+            peer_rels = await NodeManager.query_peers(
+                db=db,
+                ids=keys,
+                source_kind=self.query_params.source_kind,
+                schema=self.query_params.schema,
+                filters=self.query_params.filters,
+                fields=self.query_params.fields,
+                at=self.query_params.at,
+                branch=self.query_params.branch,
+                branch_agnostic=self.query_params.branch_agnostic,
+                fetch_peers=True,
+            )
+        peer_rels_by_node_id: dict[str, list[Relationship]] = {}
+        for rel in peer_rels:
+            node_id = rel.node_id
+            if node_id not in peer_rels_by_node_id:
+                peer_rels_by_node_id[node_id] = []
+            peer_rels_by_node_id[node_id].append(rel)
+
+        results = []
+        for node_id in keys:
+            results.append(peer_rels_by_node_id.get(node_id, []))
+        return results

infrahub/graphql/loaders/shared.py

@@ -0,0 +1,13 @@
+from typing import Any
+
+
+def to_frozen_set(to_freeze: dict[str, Any]) -> frozenset:
+    freezing_dict = {}
+    for k, v in to_freeze.items():
+        if isinstance(v, dict):
+            freezing_dict[k] = to_frozen_set(v)
+        elif isinstance(v, list | set):
+            freezing_dict[k] = frozenset(v)
+        else:
+            freezing_dict[k] = v
+    return frozenset(freezing_dict)

infrahub/graphql/manager.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Iterable, Optional, Union
+from typing import TYPE_CHECKING, Any, Iterable
 
 import graphene
 
@@ -38,6 +38,7 @@ from .mutations.repository import InfrahubRepositoryMutation
 from .mutations.resource_manager import (
     InfrahubNumberPoolMutation,
 )
+from .mutations.webhook import InfrahubWebhookMutation
 from .resolvers.resolver import (
     account_resolver,
     ancestors_resolver,
@@ -60,6 +61,7 @@ from .types import (
 )
 from .types.attribute import BaseAttribute as BaseAttributeType
 from .types.attribute import TextAttributeType
+from .types.context import ContextInput
 from .types.event import EVENT_TYPES
 
 if TYPE_CHECKING:
@@ -74,9 +76,7 @@ class DeleteInput(graphene.InputObjectType):
     hfid = graphene.List(of_type=graphene.String, required=False)
 
 
-GraphQLTypes = Union[
-    type[InfrahubMutation], type[BaseAttributeType], type[graphene.Interface], type[graphene.ObjectType]
-]
+GraphQLTypes = type[InfrahubMutation] | type[BaseAttributeType] | type[graphene.Interface] | type[graphene.ObjectType]
 
 
 class OrderInput(graphene.InputObjectType):
@@ -512,6 +512,8 @@ class GraphQLSchemaManager:
                 InfrahubKind.NAMESPACE: InfrahubIPNamespaceMutation,
                 InfrahubKind.NUMBERPOOL: InfrahubNumberPoolMutation,
                 InfrahubKind.MENUITEM: InfrahubCoreMenuMutation,
+                InfrahubKind.STANDARDWEBHOOK: InfrahubWebhookMutation,
+                InfrahubKind.CUSTOMWEBHOOK: InfrahubWebhookMutation,
             }
 
             if isinstance(node_schema, NodeSchema) and node_schema.is_ip_prefix():
@@ -642,7 +644,9 @@ class GraphQLSchemaManager:
         self.set_type(name=type_name, graphql_type=relationship_property)
 
     def generate_graphql_mutations(
-        self, schema: NodeSchema | ProfileSchema | TemplateSchema, base_class: type[InfrahubMutation]
+        self,
+        schema: NodeSchema | ProfileSchema | TemplateSchema,
+        base_class: type[InfrahubMutation],
     ) -> GraphqlMutations:
         graphql_mutation_create_input = self.generate_graphql_mutation_create_input(schema)
         graphql_mutation_update_input = self.generate_graphql_mutation_update_input(schema)
@@ -681,7 +685,7 @@ class GraphQLSchemaManager:
                 slug = InputField(StringAttributeCreate, required=True)
                 description = InputField(StringAttributeCreate, required=False)
         """
-        attrs: dict[str, Union[graphene.String, graphene.InputField]] = {"id": graphene.String(required=False)}
+        attrs: dict[str, graphene.String | graphene.InputField] = {"id": graphene.String(required=False)}
 
         for attr in schema.attributes:
             if attr.read_only:
@@ -717,7 +721,7 @@ class GraphQLSchemaManager:
                 slug = InputField(StringAttributeUpdate, required=False)
                 description = InputField(StringAttributeUpdate, required=False)
         """
-        attrs: dict[str, Union[graphene.String, graphene.InputField]] = {
+        attrs: dict[str, graphene.String | graphene.InputField] = {
             "id": graphene.String(required=False),
             "hfid": graphene.List(of_type=graphene.String, required=False),
         }
@@ -758,7 +762,7 @@ class GraphQLSchemaManager:
                 slug = InputField(StringAttributeUpdate, required=True)
                 description = InputField(StringAttributeUpdate, required=False)
         """
-        attrs: dict[str, Union[graphene.String, graphene.InputField]] = {
+        attrs: dict[str, graphene.String | graphene.InputField] = {
             "id": graphene.String(required=False),
             "hfid": graphene.List(of_type=graphene.String, required=False),
         }
@@ -809,9 +813,7 @@ class GraphQLSchemaManager:
         meta_attrs: dict[str, Any] = {"schema": schema, "name": name, "description": schema.description}
         main_attrs["Meta"] = type("Meta", (object,), meta_attrs)
 
-        args_attrs = {
-            "data": input_type(required=True),
-        }
+        args_attrs = {"data": input_type(required=True), "context": ContextInput(required=False)}
         main_attrs["Arguments"] = type("Arguments", (object,), args_attrs)
 
         return type(name, (base_class,), main_attrs)
@@ -832,9 +834,7 @@ class GraphQLSchemaManager:
         meta_attrs: dict[str, Any] = {"schema": schema, "name": name, "description": schema.description}
         main_attrs["Meta"] = type("Meta", (object,), meta_attrs)
 
-        args_attrs = {
-            "data": input_type(required=True),
-        }
+        args_attrs = {"data": input_type(required=True), "context": ContextInput(required=False)}
         main_attrs["Arguments"] = type("Arguments", (object,), args_attrs)
 
         return type(name, (base_class,), main_attrs)
@@ -851,16 +851,14 @@ class GraphQLSchemaManager:
         meta_attrs = {"schema": schema, "name": name, "description": schema.description}
         main_attrs["Meta"] = type("Meta", (object,), meta_attrs)
 
-        args_attrs: dict[str, Any] = {
-            "data": DeleteInput(required=True),
-        }
+        args_attrs: dict[str, Any] = {"data": DeleteInput(required=True), "context": ContextInput(required=False)}
         main_attrs["Arguments"] = type("Arguments", (object,), args_attrs)
 
         return type(name, (base_class,), main_attrs)
 
     def generate_filters(
         self, schema: MainSchemaTypes, top_level: bool = False, include_properties: bool = True
-    ) -> dict[str, Union[graphene.Scalar, graphene.List]]:
+    ) -> dict[str, graphene.Scalar | graphene.List]:
         """Generate the GraphQL filters for a given Schema object.
 
         The generated filter will be different if we are at the top_level (query)
@@ -925,7 +923,7 @@ class GraphQLSchemaManager:
         self,
         schema: MainSchemaTypes,
         node: type[InfrahubObject],
-        relation_property: Optional[type[InfrahubObject]] = None,
+        relation_property: type[InfrahubObject] | None = None,
         populate_cache: bool = False,
     ) -> type[InfrahubObject]:
         """Generate a edged GraphQL object Type from a Infrahub NodeSchema for pagination."""

infrahub/graphql/mutations/artifact_definition.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
 
 from graphene import InputObjectType, Mutation
 from typing_extensions import Self
@@ -28,7 +28,7 @@ class InfrahubArtifactDefinitionMutation(InfrahubMutationMixin, Mutation):
     def __init_subclass_with_meta__(
         cls,
         schema: NodeSchema,
-        _meta: Optional[Any] = None,
+        _meta: Any | None = None,
         **options: dict[str, Any],
     ) -> None:
         # Make sure schema is a valid NodeSchema Node Class
@@ -48,7 +48,7 @@ class InfrahubArtifactDefinitionMutation(InfrahubMutationMixin, Mutation):
         info: GraphQLResolveInfo,
         data: InputObjectType,
         branch: Branch,
-        database: Optional[InfrahubDatabase] = None,  # noqa: ARG003
+        database: InfrahubDatabase | None = None,  # noqa: ARG003
     ) -> tuple[Node, Self]:
         graphql_context: GraphqlContext = info.context
 
@@ -74,8 +74,8 @@ class InfrahubArtifactDefinitionMutation(InfrahubMutationMixin, Mutation):
         info: GraphQLResolveInfo,
         data: InputObjectType,
         branch: Branch,
-        database: Optional[InfrahubDatabase] = None,  # noqa: ARG003
-        node: Optional[Node] = None,  # noqa: ARG003
+        database: InfrahubDatabase | None = None,  # noqa: ARG003
+        node: Node | None = None,  # noqa: ARG003
     ) -> tuple[Node, Self]:
         graphql_context: GraphqlContext = info.context
 

infrahub/graphql/mutations/branch.py

@@ -9,6 +9,8 @@ from typing_extensions import Self
 
 from infrahub.core.branch import Branch
 from infrahub.database import retry_db_transaction
+from infrahub.graphql.context import apply_external_context
+from infrahub.graphql.types.context import ContextInput
 from infrahub.log import get_logger
 from infrahub.workflows.catalogue import (
     BRANCH_CREATE,
@@ -44,6 +46,7 @@ class BranchCreateInput(InputObjectType):
 class BranchCreate(Mutation):
     class Arguments:
         data = BranchCreateInput(required=True)
+        context = ContextInput(required=False)
         background_execution = Boolean(required=False, deprecation_reason="Please use `wait_until_completion` instead")
         wait_until_completion = Boolean(required=False)
 
@@ -58,6 +61,7 @@ class BranchCreate(Mutation):
         root: dict,  # noqa: ARG003
         info: GraphQLResolveInfo,
         data: BranchCreateInput,
+        context: ContextInput | None = None,
         background_execution: bool = False,
         wait_until_completion: bool = True,
     ) -> Self:
@@ -65,6 +69,7 @@ class BranchCreate(Mutation):
         task: dict | None = None
 
         model = BranchCreateModel(**data)
+        await apply_external_context(graphql_context=graphql_context, context_input=context)
 
         if background_execution or not wait_until_completion:
             workflow = await graphql_context.active_service.workflow.submit_workflow(
@@ -96,6 +101,7 @@ class BranchUpdateInput(InputObjectType):
 class BranchDelete(Mutation):
     class Arguments:
         data = BranchNameInput(required=True)
+        context = ContextInput(required=False)
         wait_until_completion = Boolean(required=False)
 
     ok = Boolean()
@@ -107,10 +113,12 @@ class BranchDelete(Mutation):
         root: dict,  # noqa: ARG003
         info: GraphQLResolveInfo,
         data: BranchNameInput,
+        context: ContextInput | None = None,
         wait_until_completion: bool = True,
     ) -> Self:
         graphql_context: GraphqlContext = info.context
         obj = await Branch.get_by_name(db=graphql_context.db, name=str(data.name))
+        await apply_external_context(graphql_context=graphql_context, context_input=context)
 
         if wait_until_completion:
             await graphql_context.active_service.workflow.execute_workflow(
@@ -127,15 +135,23 @@ class BranchDelete(Mutation):
 class BranchUpdate(Mutation):
     class Arguments:
         data = BranchUpdateInput(required=True)
+        context = ContextInput(required=False)
 
     ok = Boolean()
 
     @classmethod
     @retry_db_transaction(name="branch_update")
-    async def mutate(cls, root: dict, info: GraphQLResolveInfo, data: BranchNameInput) -> Self:  # noqa: ARG003
+    async def mutate(
+        cls,
+        root: dict,  # noqa: ARG003
+        info: GraphQLResolveInfo,
+        data: BranchNameInput,
+        context: ContextInput | None = None,
+    ) -> Self:
         graphql_context: GraphqlContext = info.context
 
         obj = await Branch.get_by_name(db=graphql_context.db, name=data["name"])
+        await apply_external_context(graphql_context=graphql_context, context_input=context)
 
         to_extract = ["description"]
         for field_name in to_extract:
@@ -151,6 +167,7 @@ class BranchUpdate(Mutation):
 class BranchRebase(Mutation):
     class Arguments:
         data = BranchNameInput(required=True)
+        context = ContextInput(required=False)
         wait_until_completion = Boolean(required=False)
 
     ok = Boolean()
@@ -163,11 +180,13 @@ class BranchRebase(Mutation):
         root: dict,  # noqa: ARG003
         info: GraphQLResolveInfo,
         data: BranchNameInput,
+        context: ContextInput | None = None,
         wait_until_completion: bool = True,
     ) -> Self:
         graphql_context: GraphqlContext = info.context
 
         obj = await Branch.get_by_name(db=graphql_context.db, name=str(data.name))
+        await apply_external_context(graphql_context=graphql_context, context_input=context)
         task: dict | None = None
 
         if wait_until_completion:
@@ -192,6 +211,7 @@ class BranchRebase(Mutation):
 class BranchValidate(Mutation):
     class Arguments:
         data = BranchNameInput(required=True)
+        context = ContextInput(required=False)
         wait_until_completion = Boolean(required=False)
 
     ok = Boolean()
@@ -205,11 +225,13 @@ class BranchValidate(Mutation):
         root: dict,  # noqa: ARG003
         info: GraphQLResolveInfo,
         data: BranchNameInput,
+        context: ContextInput | None = None,
         wait_until_completion: bool = True,
     ) -> Self:
         graphql_context: GraphqlContext = info.context
 
         obj = await Branch.get_by_name(db=graphql_context.db, name=str(data.name))
+        await apply_external_context(graphql_context=graphql_context, context_input=context)
         task: dict | None = None
         ok = True
 
@@ -231,6 +253,7 @@ class BranchValidate(Mutation):
 class BranchMerge(Mutation):
     class Arguments:
         data = BranchNameInput(required=True)
+        context = ContextInput(required=False)
         wait_until_completion = Boolean(required=False)
 
     ok = Boolean()
@@ -243,11 +266,13 @@ class BranchMerge(Mutation):
         root: dict,  # noqa: ARG003
         info: GraphQLResolveInfo,
         data: BranchNameInput,
+        context: ContextInput | None = None,
         wait_until_completion: bool = True,
     ) -> Self:
         branch_name = data["name"]
         task: dict | None = None
         graphql_context: GraphqlContext = info.context
+        await apply_external_context(graphql_context=graphql_context, context_input=context)
 
         if wait_until_completion:
             await graphql_context.active_service.workflow.execute_workflow(

infrahub/graphql/mutations/computed_attribute.py

@@ -5,13 +5,15 @@ from typing import TYPE_CHECKING, Any
 from graphene import Boolean, InputObjectType, Mutation, String
 
 from infrahub.core.account import ObjectPermission
-from infrahub.core.constants import ComputedAttributeKind, MutationAction, PermissionAction, PermissionDecision
+from infrahub.core.constants import ComputedAttributeKind, PermissionAction, PermissionDecision
 from infrahub.core.manager import NodeManager
 from infrahub.core.registry import registry
 from infrahub.database import retry_db_transaction
 from infrahub.events import EventMeta
-from infrahub.events.node_action import NodeMutatedEvent
+from infrahub.events.node_action import NodeUpdatedEvent
 from infrahub.exceptions import NodeNotFoundError, ValidationError
+from infrahub.graphql.context import apply_external_context
+from infrahub.graphql.types.context import ContextInput
 from infrahub.log import get_log_data
 from infrahub.worker import WORKER_IDENTITY
 
@@ -31,6 +33,7 @@ class InfrahubComputedAttributeUpdateInput(InputObjectType):
 class UpdateComputedAttribute(Mutation):
     class Arguments:
         data = InfrahubComputedAttributeUpdateInput(required=True)
+        context = ContextInput(required=False)
 
     ok = Boolean()
 
@@ -41,6 +44,7 @@ class UpdateComputedAttribute(Mutation):
         _: dict,
         info: GraphQLResolveInfo,
         data: InfrahubComputedAttributeUpdateInput,
+        context: ContextInput | None = None,
     ) -> UpdateComputedAttribute:
         graphql_context: GraphqlContext = info.context
         node_schema = registry.schema.get_node_schema(
@@ -63,6 +67,7 @@ class UpdateComputedAttribute(Mutation):
                 else PermissionDecision.ALLOW_OTHER.value,
             )
         )
+        await apply_external_context(graphql_context=graphql_context, context_input=context)
 
         if not (
             target_node := await NodeManager.get_one(
@@ -90,12 +95,11 @@ class UpdateComputedAttribute(Mutation):
             log_data = get_log_data()
             request_id = log_data.get("request_id", "")
 
-            event = NodeMutatedEvent(
+            event = NodeUpdatedEvent(
                 kind=node_schema.kind,
                 node_id=target_node.get_id(),
-                data=target_node.node_changelog.model_dump(),
+                changelog=target_node.node_changelog.model_dump(),
                 fields=[str(data.attribute)],
-                action=MutationAction.UPDATED,
                 meta=EventMeta(
                     context=graphql_context.get_context(),
                     initiator_id=WORKER_IDENTITY,

infrahub/graphql/mutations/diff.py

@@ -1,6 +1,6 @@
 from typing import TYPE_CHECKING
 
-from graphene import Boolean, DateTime, InputObjectType, Mutation, String
+from graphene import Boolean, DateTime, Field, InputObjectType, Mutation, String
 from graphql import GraphQLResolveInfo
 
 from infrahub.core import registry
@@ -9,11 +9,14 @@ from infrahub.core.diff.model.path import NameTrackingId
 from infrahub.core.diff.models import RequestDiffUpdate
 from infrahub.core.diff.repository.repository import DiffRepository
 from infrahub.core.timestamp import Timestamp
-from infrahub.database import retry_db_transaction
 from infrahub.dependencies.registry import get_component_registry
 from infrahub.exceptions import ValidationError
+from infrahub.graphql.context import apply_external_context
+from infrahub.graphql.types.context import ContextInput
 from infrahub.workflows.catalogue import DIFF_UPDATE
 
+from ..types.task import TaskInfo
+
 if TYPE_CHECKING:
     from ..initialization import GraphqlContext
 
@@ -23,24 +26,32 @@ class DiffUpdateInput(InputObjectType):
     name = String(required=False)
     from_time = DateTime(required=False)
     to_time = DateTime(required=False)
-    wait_for_completion = Boolean(required=False)
+    wait_for_completion = Boolean(required=False, deprecation_reason="Please use `wait_until_completion` instead")
 
 
 class DiffUpdateMutation(Mutation):
     class Arguments:
         data = DiffUpdateInput(required=True)
+        context = ContextInput(required=False)
+        wait_until_completion = Boolean(required=False)
 
     ok = Boolean()
+    task = Field(TaskInfo, required=False)
 
     @classmethod
-    @retry_db_transaction(name="diff_update")
     async def mutate(
         cls,
         root: dict,  # noqa: ARG003
         info: GraphQLResolveInfo,
         data: DiffUpdateInput,
-    ) -> dict[str, bool]:
+        context: ContextInput | None = None,
+        wait_until_completion: bool = False,
+    ) -> dict[str, bool | dict[str, str]]:
         graphql_context: GraphqlContext = info.context
+        await apply_external_context(graphql_context=graphql_context, context_input=context)
+
+        if data.wait_for_completion is True:
+            wait_until_completion = True
 
         from_timestamp_str = DateTime.serialize(data.from_time) if data.from_time else None
         to_timestamp_str = DateTime.serialize(data.to_time) if data.to_time else None
@@ -55,11 +66,11 @@ class DiffUpdateMutation(Mutation):
         )
 
         tracking_id = NameTrackingId(name=data.name)
-        existing_diffs_metatdatas = await diff_repository.get_roots_metadata(
+        existing_diffs_metadatas = await diff_repository.get_roots_metadata(
             diff_branch_names=[diff_branch.name], base_branch_names=[base_branch.name], tracking_id=tracking_id
         )
-        if existing_diffs_metatdatas:
-            metadata = existing_diffs_metatdatas[0]
+        if existing_diffs_metadatas:
+            metadata = existing_diffs_metadatas[0]
             from_time = Timestamp(from_timestamp_str) if from_timestamp_str else None
             to_time = Timestamp(to_timestamp_str) if to_timestamp_str else None
             branched_from_timestamp = Timestamp(diff_branch.get_branched_from())
@@ -70,7 +81,7 @@ class DiffUpdateMutation(Mutation):
             if to_time and to_time < metadata.to_time:
                 raise ValidationError(f"to_time must be null or greater than or equal to {metadata.to_time}")
 
-        if data.wait_for_completion is True:
+        if wait_until_completion is True:
             diff_coordinator = await component_registry.get_component(
                 DiffCoordinator, db=graphql_context.db, branch=diff_branch
             )
@@ -91,8 +102,9 @@ class DiffUpdateMutation(Mutation):
             to_time=to_timestamp_str,
         )
         if graphql_context.service:
-            await graphql_context.service.workflow.submit_workflow(
-                workflow=DIFF_UPDATE, context=graphql_context.get_context(), parameters={"model": model}
+            workflow = await graphql_context.service.workflow.submit_workflow(
+                workflow=DIFF_UPDATE, parameters={"model": model}
             )
+            return {"ok": True, "task": {"id": str(workflow.id)}}
 
         return {"ok": True}

infrahub/graphql/mutations/diff_conflict.py

@@ -11,7 +11,9 @@ from infrahub.core.manager import NodeManager
 from infrahub.database import retry_db_transaction
 from infrahub.dependencies.registry import get_component_registry
 from infrahub.exceptions import ProcessingError
+from infrahub.graphql.context import apply_external_context
 from infrahub.graphql.enums import ConflictSelection as GraphQlConflictSelection
+from infrahub.graphql.types.context import ContextInput
 
 if TYPE_CHECKING:
     from graphql import GraphQLResolveInfo
@@ -29,6 +31,7 @@ class ResolveDiffConflictInput(InputObjectType):
 class ResolveDiffConflict(Mutation):
     class Arguments:
         data = ResolveDiffConflictInput(required=True)
+        context = ContextInput(required=False)
 
     ok = Boolean()
 
@@ -39,8 +42,10 @@ class ResolveDiffConflict(Mutation):
         root: dict,  # noqa: ARG003
         info: GraphQLResolveInfo,
         data: ResolveDiffConflictInput,
+        context: ContextInput | None = None,
     ) -> ResolveDiffConflict:
         graphql_context: GraphqlContext = info.context
+        await apply_external_context(graphql_context=graphql_context, context_input=context)
 
         component_registry = get_component_registry()
         diff_repo = await component_registry.get_component(