Merge pull request #4842 from opsmill/lgu-migrate-req-gen-def-run

Migrate RequestGeneratorDefinitionRun to prefect

Author: LucasG0

backend/infrahub/generators/models.py

@@ -1,8 +1,8 @@
-from typing import Optional
+from __future__ import annotations
 
-from pydantic import BaseModel, Field
+from typing import Optional
 
-from infrahub.message_bus.types import ProposedChangeGeneratorDefinition
+from pydantic import BaseModel, ConfigDict, Field
 
 
 class RequestGeneratorRun(BaseModel):
@@ -21,3 +21,25 @@ class RequestGeneratorRun(BaseModel):
     target_name: str = Field(..., description="Name of the generator target")
     query: str = Field(..., description="The name of the query to use when collecting data")
     variables: dict = Field(..., description="Input variables when running the generator")
+
+
+class RequestGeneratorDefinitionRun(BaseModel):
+    """Sent to trigger a Generator to run on a specific branch."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    generator_definition: ProposedChangeGeneratorDefinition = Field(..., description="The Generator Definition")
+    branch: str = Field(..., description="The branch to target")
+
+
+class ProposedChangeGeneratorDefinition(BaseModel):
+    definition_id: str
+    definition_name: str
+    query_name: str
+    convert_query_response: bool
+    query_models: list[str]
+    repository_id: str
+    class_name: str
+    file_path: str
+    parameters: dict
+    group_id: str

backend/infrahub/generators/tasks.py

@@ -8,12 +8,15 @@
 
 from infrahub import lock
 from infrahub.core.constants import GeneratorInstanceStatus, InfrahubKind
-from infrahub.generators.models import RequestGeneratorRun
+from infrahub.generators.models import (
+    ProposedChangeGeneratorDefinition,
+    RequestGeneratorDefinitionRun,
+    RequestGeneratorRun,
+)
 from infrahub.git.base import extract_repo_file_information
 from infrahub.git.repository import get_initialized_repo
-from infrahub.message_bus import messages
-from infrahub.message_bus.types import ProposedChangeGeneratorDefinition
 from infrahub.services import InfrahubServices, services
+from infrahub.workflows.catalogue import REQUEST_GENERATOR_DEFINITION_RUN, REQUEST_GENERATOR_RUN
 
 
 @flow(name="generator-run")
@@ -130,9 +133,61 @@ async def run_generator_definition(branch: str) -> None:
         for generator in generators
     ]
 
-    events = [
-        messages.RequestGeneratorDefinitionRun(branch=branch, generator_definition=generator_definition)
-        for generator_definition in generator_definitions
-    ]
-    for event in events:
-        await service.send(message=event)
+    for generator_definition in generator_definitions:
+        model = RequestGeneratorDefinitionRun(branch=branch, generator_definition=generator_definition)
+        await service.workflow.submit_workflow(workflow=REQUEST_GENERATOR_DEFINITION_RUN, parameters={"model": model})
+
+
+@flow(name="request_generator_definition_run")
+async def request_generator_definition_run(model: RequestGeneratorDefinitionRun) -> None:
+    service = services.service
+    async with service.task_report(
+        title="Executing Generator",
+        related_node=model.generator_definition.definition_id,
+    ) as task_report:
+        group = await service.client.get(
+            kind=InfrahubKind.GENERICGROUP,
+            prefetch_relationships=True,
+            populate_store=True,
+            id=model.generator_definition.group_id,
+            branch=model.branch,
+        )
+        await group.members.fetch()
+
+        existing_instances = await service.client.filters(
+            kind=InfrahubKind.GENERATORINSTANCE,
+            definition__ids=[model.generator_definition.definition_id],
+            include=["object"],
+            branch=model.branch,
+        )
+        instance_by_member = {}
+        for instance in existing_instances:
+            instance_by_member[instance.object.peer.id] = instance.id
+
+        repository = await service.client.get(
+            kind=InfrahubKind.REPOSITORY, branch=model.branch, id=model.generator_definition.repository_id
+        )
+
+        for relationship in group.members.peers:
+            member = relationship.peer
+            generator_instance = instance_by_member.get(member.id)
+            request_generator_run_model = RequestGeneratorRun(
+                generator_definition=model.generator_definition,
+                commit=repository.commit.value,
+                generator_instance=generator_instance,
+                repository_id=repository.id,
+                repository_name=repository.name.value,
+                repository_kind=repository.typename,
+                branch_name=model.branch,
+                query=model.generator_definition.query_name,
+                variables=member.extract(params=model.generator_definition.parameters),
+                target_id=member.id,
+                target_name=member.display_label,
+            )
+            await service.workflow.submit_workflow(
+                workflow=REQUEST_GENERATOR_RUN, parameters={"model": request_generator_run_model}
+            )
+
+        await task_report.info(
+            event=f"Generator triggered for {len(group.members.peers)} members in {group.name.value}."
+        )

backend/infrahub/message_bus/messages/__init__.py

@@ -30,7 +30,6 @@
 from .refresh_webhook_configuration import RefreshWebhookConfiguration
 from .request_artifactdefinition_check import RequestArtifactDefinitionCheck
 from .request_generatordefinition_check import RequestGeneratorDefinitionCheck
-from .request_generatordefinition_run import RequestGeneratorDefinitionRun
 from .request_graphqlquerygroup_update import RequestGraphQLQueryGroupUpdate
 from .request_proposedchange_pipeline import RequestProposedChangePipeline
 from .request_repository_checks import RequestRepositoryChecks
@@ -67,7 +66,6 @@
     "refresh.webhook.configuration": RefreshWebhookConfiguration,
     "request.artifact_definition.check": RequestArtifactDefinitionCheck,
     "request.generator_definition.check": RequestGeneratorDefinitionCheck,
-    "request.generator_definition.run": RequestGeneratorDefinitionRun,
     "request.graphql_query_group.update": RequestGraphQLQueryGroupUpdate,
     "request.proposed_change.data_integrity": RequestProposedChangeDataIntegrity,
     "request.proposed_change.pipeline": RequestProposedChangePipeline,

backend/infrahub/message_bus/messages/check_generator_run.py

@@ -2,8 +2,8 @@
 
 from pydantic import Field
 
+from infrahub.generators.models import ProposedChangeGeneratorDefinition
 from infrahub.message_bus import InfrahubMessage
-from infrahub.message_bus.types import ProposedChangeGeneratorDefinition
 
 
 class CheckGeneratorRun(InfrahubMessage):

backend/infrahub/message_bus/messages/request_generatordefinition_check.py

@@ -1,7 +1,8 @@
 from pydantic import ConfigDict, Field
 
+from infrahub.generators.models import ProposedChangeGeneratorDefinition
 from infrahub.message_bus import InfrahubMessage
-from infrahub.message_bus.types import ProposedChangeBranchDiff, ProposedChangeGeneratorDefinition
+from infrahub.message_bus.types import ProposedChangeBranchDiff
 
 
 class RequestGeneratorDefinitionCheck(InfrahubMessage):

backend/infrahub/message_bus/messages/request_generatordefinition_run.py

@@ -42,7 +42,6 @@
     "refresh.registry.rebased_branch": refresh.registry.rebased_branch,
     "refresh.webhook.configuration": refresh.webhook.configuration,
     "request.generator_definition.check": requests.generator_definition.check,
-    "request.generator_definition.run": requests.generator_definition.run,
     "request.graphql_query_group.update": requests.graphql_query_group.update,
     "request.artifact_definition.check": requests.artifact_definition.check,
     "request.proposed_change.data_integrity": requests.proposed_change.data_integrity,

backend/infrahub/message_bus/operations/__init__.py

@@ -5,11 +5,9 @@
 
 from infrahub.core.constants import InfrahubKind, ValidatorConclusion, ValidatorState
 from infrahub.core.timestamp import Timestamp
-from infrahub.generators.models import RequestGeneratorRun
 from infrahub.message_bus import InfrahubMessage, Meta, messages
 from infrahub.message_bus.types import KVTTL
 from infrahub.services import InfrahubServices
-from infrahub.workflows.catalogue import REQUEST_GENERATOR_RUN
 
 
 @flow(name="generator-definition-check")
@@ -131,64 +129,6 @@ async def check(message: messages.RequestGeneratorDefinitionCheck, service: Infr
             await service.send(message=event)
 
 
-@flow(name="generator-definition-run")
-async def run(message: messages.RequestGeneratorDefinitionRun, service: InfrahubServices) -> None:
-    async with service.task_report(
-        title="Executing Generator",
-        related_node=message.generator_definition.definition_id,
-    ) as task_report:
-        service.log.info(
-            "Received request to run generator",
-            branch=message.branch,
-            generator_definition=message.generator_definition.definition_id,
-        )
-
-        group = await service.client.get(
-            kind=InfrahubKind.GENERICGROUP,
-            prefetch_relationships=True,
-            populate_store=True,
-            id=message.generator_definition.group_id,
-            branch=message.branch,
-        )
-        await group.members.fetch()
-
-        existing_instances = await service.client.filters(
-            kind=InfrahubKind.GENERATORINSTANCE,
-            definition__ids=[message.generator_definition.definition_id],
-            include=["object"],
-            branch=message.branch,
-        )
-        instance_by_member = {}
-        for instance in existing_instances:
-            instance_by_member[instance.object.peer.id] = instance.id
-
-        repository = await service.client.get(
-            kind=InfrahubKind.REPOSITORY, branch=message.branch, id=message.generator_definition.repository_id
-        )
-
-        for relationship in group.members.peers:
-            member = relationship.peer
-            generator_instance = instance_by_member.get(member.id)
-            model = RequestGeneratorRun(
-                generator_definition=message.generator_definition,
-                commit=repository.commit.value,
-                generator_instance=generator_instance,
-                repository_id=repository.id,
-                repository_name=repository.name.value,
-                repository_kind=repository.typename,
-                branch_name=message.branch,
-                query=message.generator_definition.query_name,
-                variables=member.extract(params=message.generator_definition.parameters),
-                target_id=member.id,
-                target_name=member.display_label,
-            )
-            await service.workflow.submit_workflow(workflow=REQUEST_GENERATOR_RUN, parameters={"model": model})
-
-        await task_report.info(
-            event=f"Generator triggered for {len(group.members.peers)} members in {group.name.value}."
-        )
-
-
 def _run_generator(instance_id: Optional[str], managed_branch: bool, impacted_instances: list[str]) -> bool:
     """Returns a boolean to indicate if a generator instance needs to be executed
     Will return true if:

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

@@ -22,13 +22,13 @@
 from infrahub.core.validators.checker import schema_validators_checker
 from infrahub.core.validators.determiner import ConstraintValidatorDeterminer
 from infrahub.dependencies.registry import get_component_registry
+from infrahub.generators.models import ProposedChangeGeneratorDefinition
 from infrahub.git.repository import InfrahubRepository, get_initialized_repo
 from infrahub.log import get_logger
 from infrahub.message_bus import InfrahubMessage, messages
 from infrahub.message_bus.types import (
     ProposedChangeArtifactDefinition,
     ProposedChangeBranchDiff,
-    ProposedChangeGeneratorDefinition,
     ProposedChangeRepository,
     ProposedChangeSubscriber,
 )

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

@@ -107,19 +107,6 @@ def transform_location(self) -> str:
         raise ValueError("Invalid kind for Transform")
 
 
-class ProposedChangeGeneratorDefinition(BaseModel):
-    definition_id: str
-    definition_name: str
-    query_name: str
-    convert_query_response: bool
-    query_models: list[str]
-    repository_id: str
-    class_name: str
-    file_path: str
-    parameters: dict
-    group_id: str
-
-
 class ProposedChangeBranchDiff(BaseModel):
     diff_summary: list[NodeDiff] = Field(default_factory=list, description="The DiffSummary between two branches")
     repositories: list[ProposedChangeRepository] = Field(default_factory=list)