IFC-1053 Add mutation to trigger run of a generator (#5791)

gmazoyer · web-flow · commit 504b8f13fee6 · 2025-02-24T15:53:10.000+01:00
diff --git a/backend/infrahub/generators/models.py b/backend/infrahub/generators/models.py
@@ -30,6 +30,7 @@ class RequestGeneratorDefinitionRun(BaseModel):
 
     generator_definition: ProposedChangeGeneratorDefinition = Field(..., description="The Generator Definition")
     branch: str = Field(..., description="The branch to target")
+    target_members: list[str] = Field(default_factory=list, description="List of targets to run the generator for")
 
 
 class ProposedChangeGeneratorDefinition(BaseModel):
diff --git a/backend/infrahub/generators/tasks.py b/backend/infrahub/generators/tasks.py
@@ -1,11 +1,15 @@
 from __future__ import annotations
 
+import asyncio
+from typing import TYPE_CHECKING, Any
+
 from infrahub_sdk.exceptions import ModuleImportError
 from infrahub_sdk.node import InfrahubNode
 from infrahub_sdk.protocols import CoreGeneratorInstance
 from infrahub_sdk.schema.repository import InfrahubGeneratorDefinitionConfig
-from prefect import flow, task
+from prefect import State, flow, task
 from prefect.cache_policies import NONE
+from prefect.states import Completed, Failed
 
 from infrahub import lock
 from infrahub.context import InfrahubContext  # noqa: TC001 needed for prefect flow
@@ -21,6 +25,9 @@
 from infrahub.workflows.catalogue import REQUEST_GENERATOR_DEFINITION_RUN, REQUEST_GENERATOR_RUN
 from infrahub.workflows.utils import add_tags
 
+if TYPE_CHECKING:
+    from collections.abc import Coroutine
+
 
 @flow(
     name="generator-run",
@@ -154,7 +161,7 @@ async def run_generator_definition(branch: str, context: InfrahubContext, servic
 )
 async def request_generator_definition_run(
     model: RequestGeneratorDefinitionRun, context: InfrahubContext, service: InfrahubServices
-) -> None:
+) -> State[Any]:
     await add_tags(branches=[model.branch], nodes=[model.generator_definition.definition_id])
 
     group = await service.client.get(
@@ -190,8 +197,13 @@ async def request_generator_definition_run(
             raise_when_missing=True,
         )
 
+    tasks: list[Coroutine[Any, Any, Any]] = []
     for relationship in group.members.peers:
         member = relationship.peer
+
+        if model.target_members and member.id not in model.target_members:
+            continue
+
         generator_instance = instance_by_member.get(member.id)
         request_generator_run_model = RequestGeneratorRun(
             generator_definition=model.generator_definition,
@@ -206,6 +218,14 @@ async def request_generator_definition_run(
             target_id=member.id,
             target_name=member.display_label,
         )
-        await service.workflow.submit_workflow(
-            workflow=REQUEST_GENERATOR_RUN, context=context, parameters={"model": request_generator_run_model}
+        tasks.append(
+            service.workflow.execute_workflow(
+                workflow=REQUEST_GENERATOR_RUN, context=context, parameters={"model": request_generator_run_model}
+            )
         )
+
+    try:
+        await asyncio.gather(*tasks)
+        return Completed(message=f"Successfully run {len(tasks)} generators")
+    except Exception as exc:
+        return Failed(message="One or more generators failed", error=exc)
diff --git a/backend/infrahub/graphql/mutations/generator.py b/backend/infrahub/graphql/mutations/generator.py
@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from graphene import Boolean, Field, InputField, InputObjectType, List, Mutation, NonNull, String
+
+from infrahub.core.manager import NodeManager
+from infrahub.generators.models import ProposedChangeGeneratorDefinition, RequestGeneratorDefinitionRun
+from infrahub.graphql.types.task import TaskInfo
+from infrahub.workflows.catalogue import REQUEST_GENERATOR_DEFINITION_RUN
+
+if TYPE_CHECKING:
+    from graphql import GraphQLResolveInfo
+
+    from ..initialization import GraphqlContext
+
+
+class GeneratorDefinitionRequestRunInput(InputObjectType):
+    id = InputField(String(required=True), description="ID of the generator definition to run")
+    nodes = InputField(List(of_type=NonNull(String)), description="ID list of targets to run the generator for")
+
+
+class GeneratorDefinitionRequestRun(Mutation):
+    class Arguments:
+        data = GeneratorDefinitionRequestRunInput(required=True)
+        wait_until_completion = Boolean(required=False)
+
+    ok = Boolean()
+    task = Field(TaskInfo, required=False)
+
+    @classmethod
+    async def mutate(
+        cls,
+        root: dict,  # noqa: ARG003
+        info: GraphQLResolveInfo,
+        data: GeneratorDefinitionRequestRunInput,
+        wait_until_completion: bool = True,
+    ) -> GeneratorDefinitionRequestRun:
+        graphql_context: GraphqlContext = info.context
+        db = graphql_context.db
+
+        generator_definition = await NodeManager.get_one(
+            id=str(data.id), db=db, branch=graphql_context.branch, prefetch_relationships=True, raise_on_error=True
+        )
+        query = await generator_definition.query.get_peer(db=db)
+        repository = await generator_definition.repository.get_peer(db=db)
+        group = await generator_definition.targets.get_peer(db=db)
+
+        request_model = RequestGeneratorDefinitionRun(
+            generator_definition=ProposedChangeGeneratorDefinition(
+                definition_id=generator_definition.id,
+                definition_name=generator_definition.name.value,
+                class_name=generator_definition.class_name.value,
+                file_path=generator_definition.file_path.value,
+                query_name=query.name.value,
+                query_models=query.models.value,
+                repository_id=repository.id,
+                parameters=generator_definition.parameters.value,
+                group_id=group.id,
+                convert_query_response=generator_definition.convert_query_response.value or False,
+            ),
+            branch=graphql_context.branch.name,
+            target_members=data.get("nodes", []),
+        )
+
+        if not wait_until_completion:
+            workflow = await graphql_context.active_service.workflow.submit_workflow(
+                workflow=REQUEST_GENERATOR_DEFINITION_RUN,
+                context=graphql_context.get_context(),
+                parameters={"model": request_model},
+            )
+            return cls(ok=True, task={"id": workflow.id})
+
+        await graphql_context.active_service.workflow.execute_workflow(
+            workflow=REQUEST_GENERATOR_DEFINITION_RUN,
+            context=graphql_context.get_context(),
+            parameters={"model": request_model},
+        )
+        return cls(ok=True)
diff --git a/backend/infrahub/graphql/schema.py b/backend/infrahub/graphql/schema.py
@@ -18,6 +18,7 @@
 from .mutations.computed_attribute import UpdateComputedAttribute
 from .mutations.diff import DiffUpdateMutation
 from .mutations.diff_conflict import ResolveDiffConflict
+from .mutations.generator import GeneratorDefinitionRequestRun
 from .mutations.proposed_change import ProposedChangeMerge, ProposedChangeRequestRunCheck
 from .mutations.relationship import (
     RelationshipAdd,
@@ -83,6 +84,7 @@ class InfrahubBaseMutation(ObjectType):
     InfrahubAccountTokenDelete = InfrahubAccountTokenDelete.Field()
     CoreProposedChangeRunCheck = ProposedChangeRequestRunCheck.Field()
     CoreProposedChangeMerge = ProposedChangeMerge.Field()
+    CoreGeneratorDefinitionRun = GeneratorDefinitionRequestRun.Field()
 
     IPPrefixPoolGetResource = IPPrefixPoolGetResource.Field()
     IPAddressPoolGetResource = IPAddressPoolGetResource.Field()
diff --git a/backend/tests/functional/generator/test_mutation_generator.py b/backend/tests/functional/generator/test_mutation_generator.py
@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+from infrahub_sdk.graphql import Mutation
+from infrahub_sdk.protocols import CoreGeneratorDefinition
+from tests.constants import TestKind
+from tests.helpers.file_repo import FileRepo
+from tests.helpers.schema import CAR_SCHEMA, load_schema
+from tests.helpers.test_app import TestInfrahubApp
+
+from infrahub.core.constants import InfrahubKind
+from infrahub.core.node import Node
+from infrahub.services.adapters.cache.redis import RedisCache
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from infrahub_sdk import InfrahubClient
+    from tests.adapters.message_bus import BusSimulator
+
+    from infrahub.database import InfrahubDatabase
+
+
+class TestMutationGenerator(TestInfrahubApp):
+    @pytest.fixture(scope="class")
+    async def initial_dataset(
+        self,
+        db: InfrahubDatabase,
+        initialize_registry: None,
+        git_repos_source_dir_module_scope: Path,
+        client: InfrahubClient,
+        bus_simulator: BusSimulator,
+        prefect_test_fixture,
+    ) -> None:
+        await load_schema(db, schema=CAR_SCHEMA)
+
+        bus_simulator.service._cache = RedisCache()
+
+        john = await Node.init(schema=TestKind.PERSON, db=db)
+        await john.new(db=db, name="John", height=175, age=25, description="The famous Joe Doe")
+        await john.save(db=db)
+        koenigsegg = await Node.init(schema=TestKind.MANUFACTURER, db=db)
+        await koenigsegg.new(db=db, name="Koenigsegg")
+        await koenigsegg.save(db=db)
+        people = await Node.init(schema=InfrahubKind.STANDARDGROUP, db=db)
+        await people.new(db=db, name="people", members=[john])
+        await people.save(db=db)
+
+        jesko = await Node.init(schema=TestKind.CAR, db=db)
+        await jesko.new(
+            db=db,
+            name="Jesko",
+            color="Red",
+            description="A limited production mid-engine sports car",
+            owner=john,
+            manufacturer=koenigsegg,
+        )
+        await jesko.save(db=db)
+
+        branch1 = await client.branch.create(branch_name="branch1")
+
+        FileRepo(name="car-dealership", sources_directory=git_repos_source_dir_module_scope)
+        client_repository = await client.create(
+            kind=InfrahubKind.REPOSITORY,
+            data={"name": "car-dealership", "location": f"{git_repos_source_dir_module_scope}/car-dealership"},
+            branch=branch1.name,
+        )
+        await client_repository.save()
+
+        richard = await Node.init(schema=TestKind.PERSON, db=db, branch=branch1.name)
+        await richard.new(db=db, name="Richard", height=180, description="The less famous Richard Doe")
+        await richard.save(db=db)
+
+    async def test_execute_generator(self, db: InfrahubDatabase, initial_dataset: None, client: InfrahubClient) -> None:
+        generator = await client.get(kind=CoreGeneratorDefinition, branch="branch1", name__value="cartags")
+        mutation = Mutation(
+            mutation="CoreGeneratorDefinitionRun", input_data={"data": {"id": generator.id}}, query={"ok": None}
+        )
+        response = await client.execute_graphql(query=mutation.render(), branch_name="branch1")
+        assert response["CoreGeneratorDefinitionRun"]["ok"]
+
+    async def test_execute_generator_background(
+        self, db: InfrahubDatabase, initial_dataset: None, client: InfrahubClient
+    ) -> None:
+        generator = await client.get(kind=CoreGeneratorDefinition, branch="branch1", name__value="cartags")
+        mutation = Mutation(
+            mutation="CoreGeneratorDefinitionRun",
+            input_data={"data": {"id": generator.id}, "wait_until_completion": False},
+            query={"ok": None, "task": {"id": None}},
+        )
+        response = await client.execute_graphql(query=mutation.render(), branch_name="branch1")
+        assert response["CoreGeneratorDefinitionRun"]["ok"]
+        assert response["CoreGeneratorDefinitionRun"]["task"]["id"]
diff --git a/backend/tests/unit/conftest.py b/backend/tests/unit/conftest.py
@@ -1205,7 +1205,7 @@ async def car_person_data_generic(db: InfrahubDatabase, register_core_models_sch
     """
 
     q1 = await Node.init(db=db, schema=InfrahubKind.GRAPHQLQUERY)
-    await q1.new(db=db, name="query01", query=query)
+    await q1.new(db=db, name="query01", query=query, models=["TestPerson"])
     await q1.save(db=db)
 
     r1 = await Node.init(db=db, schema=InfrahubKind.REPOSITORY)
diff --git a/backend/tests/unit/graphql/test_mutation_generator.py b/backend/tests/unit/graphql/test_mutation_generator.py
@@ -0,0 +1,112 @@
+from unittest.mock import call, patch
+
+import pytest
+
+from infrahub.auth import AccountSession, AuthType
+from infrahub.context import InfrahubContext
+from infrahub.core.branch import Branch
+from infrahub.core.constants import InfrahubKind
+from infrahub.core.node import Node
+from infrahub.database import InfrahubDatabase
+from infrahub.generators.models import ProposedChangeGeneratorDefinition, RequestGeneratorDefinitionRun
+from infrahub.graphql.initialization import prepare_graphql_params
+from infrahub.services import InfrahubServices
+from infrahub.services.adapters.workflow.local import WorkflowLocalExecution
+from infrahub.workflows.catalogue import REQUEST_GENERATOR_DEFINITION_RUN
+from tests.adapters.message_bus import BusRecorder
+from tests.helpers.graphql import graphql
+
+
+@pytest.fixture
+async def group1(db: InfrahubDatabase, car_person_data_generic: dict[str, Node]) -> Node:
+    g1 = await Node.init(db=db, schema=InfrahubKind.STANDARDGROUP)
+    await g1.new(db=db, name="group1", members=[car_person_data_generic["c1"], car_person_data_generic["c2"]])
+    await g1.save(db=db)
+    return g1
+
+
+@pytest.fixture
+async def definition1(db: InfrahubDatabase, car_person_data_generic: dict[str, Node], group1: Node) -> Node:
+    gd1 = await Node.init(db=db, schema=InfrahubKind.GENERATORDEFINITION)
+    await gd1.new(
+        db=db,
+        name="generatordef01",
+        query=str(car_person_data_generic["q1"].id),
+        repository=str(car_person_data_generic["r1"].id),
+        file_path="generator01.py",
+        class_name="Generator01",
+        targets=str(group1.id),
+        parameters={"value": {"name": "name__value"}},
+    )
+    await gd1.save(db=db)
+    return gd1
+
+
+async def test_run_generator_definition(
+    db: InfrahubDatabase,
+    default_branch: Branch,
+    register_core_models_schema,
+    car_person_data_generic,
+    create_test_admin: Node,
+    definition1: Node,
+):
+    query = """
+    mutation {
+        CoreGeneratorDefinitionRun(data: { id: "%s" }, wait_until_completion: false) {
+            ok
+        }
+    }
+    """ % (definition1.id)
+    recorder = BusRecorder()
+    service = await InfrahubServices.new(message_bus=recorder, workflow=WorkflowLocalExecution())
+
+    account_session = AccountSession(
+        authenticated=True, account_id=create_test_admin.id, session_id=None, auth_type=AuthType.API
+    )
+    gql_params = await prepare_graphql_params(
+        db=db, include_subscription=False, branch=default_branch, service=service, account_session=account_session
+    )
+
+    with patch(
+        "infrahub.services.adapters.workflow.local.WorkflowLocalExecution.submit_workflow"
+    ) as mock_submit_workflow:
+        result = await graphql(
+            schema=gql_params.schema,
+            source=query,
+            context_value=gql_params.context,
+            root_value=None,
+            variable_values={},
+        )
+
+        assert not result.errors
+        assert result.data
+        assert result.data["CoreGeneratorDefinitionRun"]["ok"]
+
+        context = InfrahubContext.init(branch=default_branch, account=account_session)
+        query = await definition1.query.get_peer(db=db)
+        repository = await definition1.repository.get_peer(db=db)
+        group = await definition1.targets.get_peer(db=db)
+        expected_calls = [
+            call(
+                workflow=REQUEST_GENERATOR_DEFINITION_RUN,
+                parameters={
+                    "model": RequestGeneratorDefinitionRun(
+                        generator_definition=ProposedChangeGeneratorDefinition(
+                            definition_id=definition1.id,
+                            definition_name=definition1.name.value,
+                            class_name=definition1.class_name.value,
+                            file_path=definition1.file_path.value,
+                            query_name=query.name.value,
+                            query_models=query.models.value,
+                            repository_id=repository.id,
+                            parameters=definition1.parameters.value,
+                            group_id=group.id,
+                            convert_query_response=definition1.convert_query_response.value,
+                        ),
+                        branch=context.branch.name,
+                    )
+                },
+                context=context,
+            ),
+        ]
+        mock_submit_workflow.assert_has_calls(expected_calls)
