Use request context from SDK within pipeline

Author: ogenstad

backend/infrahub/context.py

@@ -1,3 +1,4 @@
+from infrahub_sdk.context import ContextAccount, RequestContext
 from pydantic import BaseModel, Field
 from typing_extensions import Self
 
@@ -37,3 +38,6 @@ def set_event(self, name: str, id: str) -> None:
             self.event.id = id
         else:
             self.event = EventContext(name=name, id=id)
+
+    def to_request_context(self) -> RequestContext:
+        return RequestContext(account=ContextAccount(id=self.account.account_id))

backend/infrahub/graphql/context.py

@@ -17,6 +17,12 @@ async def apply_external_context(graphql_context: GraphqlContext, context_input:
     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
     )

backend/infrahub/message_bus/messages/check_generator_run.py

@@ -1,5 +1,6 @@
 from pydantic import Field
 
+from infrahub.context import InfrahubContext
 from infrahub.generators.models import ProposedChangeGeneratorDefinition
 from infrahub.message_bus import InfrahubMessage
 
@@ -22,3 +23,4 @@ class CheckGeneratorRun(InfrahubMessage):
     variables: dict = Field(..., description="Input variables when running the generator")
     validator_id: str = Field(..., description="The ID of the validator")
     proposed_change: str | None = Field(None, description="The unique ID of the Proposed Change")
+    context: InfrahubContext = Field(..., description="The Infrahub context")

backend/infrahub/message_bus/operations/check/generator.py

@@ -72,6 +72,7 @@ async def run(message: messages.CheckGeneratorRun, service: InfrahubServices) ->
             convert_query_response=generator_definition.convert_query_response,
             infrahub_node=InfrahubNode,
         )
+        generator._init_client.request_context = message.context.to_request_context()
         await generator.run(identifier=generator_definition.name)
         generator_instance.status.value = GeneratorInstanceStatus.READY.value
     except ModuleImportError as exc:

backend/infrahub/message_bus/operations/requests/generator_definition.py

@@ -87,6 +87,7 @@ async def check(message: messages.RequestGeneratorDefinitionCheck, service: Infr
             log.info(f"Trigger execution of {message.generator_definition.definition_name} for {member.display_label}")
             events.append(
                 messages.CheckGeneratorRun(
+                    context=message.context,
                     generator_definition=message.generator_definition,
                     generator_instance=generator_instance,
                     commit=repository.source_commit,

backend/infrahub/tasks/artifact.py

@@ -39,6 +39,6 @@ async def define_artifact(
                         "content_type": model.content_type,
                     },
                 )
-                await artifact.save()
+                await artifact.save(request_context=model.context.to_request_context())
                 created = True
     return artifact, created

backend/tests/integration/proposed_change/test_proposed_change_conflict.py

@@ -9,11 +9,12 @@
 from uuid import uuid4
 
 import pytest
+from infrahub_sdk.context import ContextAccount, RequestContext
 from infrahub_sdk.exceptions import GraphQLError
 from infrahub_sdk.protocols import CoreDataCheck, CoreProposedChange
 
 from infrahub.core.constants import InfrahubKind, ValidatorConclusion
-from infrahub.core.initialization import create_branch
+from infrahub.core.initialization import create_account, create_branch
 from infrahub.core.manager import NodeManager
 from infrahub.core.merge import BranchMerger
 from infrahub.core.node import Node
@@ -203,11 +204,14 @@ async def test_conflict_pipeline(
         assert john.description.value == "Oh boy"  # type: ignore[attr-defined]
 
     async def test_happy_pipeline(self, db: InfrahubDatabase, happy_data_branch: str, client: InfrahubClient) -> None:
+        proposed_change_user = await create_account(db=db, name="jimmy-change-user", password="Password123")
         proposed_change_create = await client.create(
             kind=CoreProposedChange,
             data={"source_branch": happy_data_branch, "destination_branch": "main", "name": "happy-test"},
         )
-        await proposed_change_create.save()
+        await proposed_change_create.save(
+            request_context=RequestContext(account=ContextAccount(id=proposed_change_user.id))
+        )
 
         # -------------------------------------------------
         # Ensure that all validators have been executed and aren't reporting errors
@@ -292,6 +296,7 @@ async def test_happy_pipeline(self, db: InfrahubDatabase, happy_data_branch: str
         assert artifact_events["InfrahubEvent"]["count"] > 0
         latest_artifact_event = artifact_events["InfrahubEvent"]["edges"][0]["node"]
         assert sorted(latest_artifact_event.keys()) == [
+            "account_id",
             "artifact_definition_id",
             "branch",
             "checksum",
@@ -348,6 +353,14 @@ async def test_happy_pipeline(self, db: InfrahubDatabase, happy_data_branch: str
             "CoreUserValidator",
         ]
 
+        pr_account_events = await client.execute_graphql(
+            query=QUERY_EVENT,
+            variables={"account__ids": [proposed_change_user.id]},
+        )
+        pr_account_events_types = {event["node"]["event"] for event in pr_account_events["InfrahubEvent"]["edges"]}
+        assert "infrahub.validator.passed" in pr_account_events_types
+        assert "infrahub.artifact.updated" in pr_account_events_types
+
     async def test_merge_failure(
         self,
         db: InfrahubDatabase,
@@ -389,6 +402,7 @@ async def test_connectivity(self, db: InfrahubDatabase, initial_dataset: str, cl
     $branch: [String!],
     $parent__ids: [String!],
     $event_type: [String!]
+    $account__ids: [String!],
     $related_node__ids: [String!],
     $event_type_filter: EventTypeFilter
 ) {
@@ -397,6 +411,7 @@ async def test_connectivity(self, db: InfrahubDatabase, initial_dataset: str, cl
     parent__ids: $parent__ids
     event_type: $event_type
     event_type_filter: $event_type_filter
+    account__ids: $account__ids
     related_node__ids: $related_node__ids
   ) {
     count
@@ -409,6 +424,7 @@ async def test_connectivity(self, db: InfrahubDatabase, initial_dataset: str, cl
         parent_id
         level
         occurred_at
+        account_id
         primary_node {
           id
           kind
@@ -425,6 +441,14 @@ async def test_connectivity(self, db: InfrahubDatabase, initial_dataset: str, cl
           checksum_previous
           checksum
         }
+       ... on NodeMutatedEvent {
+          id
+          attributes {
+            name
+            value
+            value_previous
+          }
+        }
       }
     }
   }