Migrate RequestGeneratorDefinitionRun to prefect

Author: LucasG0

backend/infrahub/generators/models.py

@@ -1,6 +1,6 @@
 from typing import Optional
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
 
 from infrahub.message_bus.types import ProposedChangeGeneratorDefinition
 
@@ -21,3 +21,12 @@ 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")

backend/infrahub/generators/tasks.py

@@ -8,12 +8,12 @@
 
 from infrahub import lock
 from infrahub.core.constants import GeneratorInstanceStatus, InfrahubKind
-from infrahub.generators.models import RequestGeneratorRun
+from infrahub.generators.models import 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 +130,66 @@ 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, service: InfrahubServices) -> None:
+    async with service.task_report(
+        title="Executing Generator",
+        related_node=model.generator_definition.definition_id,
+    ) as task_report:
+        service.log.info(
+            "Received request to run generator",
+            branch=model.branch,
+            generator_definition=model.generator_definition.definition_id,
+        )
+
+        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/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

@@ -85,6 +85,14 @@
     function="run_generator",
 )
 
+REQUEST_GENERATOR_DEFINITION_RUN = WorkflowDefinition(
+    name="request_generator_definition_run",
+    type=WorkflowType.INTERNAL,
+    module="infrahub.generators.tasks",
+    function="request_generator_definition_run",
+    branch_support=BranchSupportType.AWARE,
+)
+
 REQUEST_ARTIFACT_GENERATE = WorkflowDefinition(
     name="artifact-generate",
     type=WorkflowType.INTERNAL,
@@ -204,4 +212,5 @@
     GIT_REPOSITORIES_MERGE,
     TRIGGER_GENERATOR_DEFINITION_RUN,
     BRANCH_CANCEL_PROPOSED_CHANGES,
+    REQUEST_GENERATOR_DEFINITION_RUN,
 ]

backend/infrahub/workflows/catalogue.py

@@ -571,21 +571,6 @@ For more detailed explanations on how to use these events within Infrahub, see t
 | **source_branch_sync_with_git** | Indicates if the source branch should sync with git | boolean | None |
 | **destination_branch** | The target branch | string | None |
 <!-- vale on -->
-<!-- vale off -->
-#### Event request.generator_definition.run
-<!-- vale on -->
-
-**Description**: Sent to trigger a Generator to run on a specific branch.
-
-**Priority**: 3
-
-<!-- vale off -->
-| Key | Description | Type | Default Value |
-|-----|-------------|------|---------------|
-| **meta** | Meta properties for the message | N/A | None |
-| **generator_definition** | The Generator Definition | N/A | None |
-| **branch** | The branch to target | string | None |
-<!-- vale on -->
 
 <!-- vale off -->
 ### Request Graphql Query Group
@@ -1414,22 +1399,6 @@ For more detailed explanations on how to use these events within Infrahub, see t
 | **source_branch_sync_with_git** | Indicates if the source branch should sync with git | boolean | None |
 | **destination_branch** | The target branch | string | None |
 <!-- vale on -->
-<!-- vale off -->
-#### Event request.generator_definition.run
-<!-- vale on -->
-
-**Description**: Sent to trigger a Generator to run on a specific branch.
-
-**Priority**: 3
-
-
-<!-- vale off -->
-| Key | Description | Type | Default Value |
-|-----|-------------|------|---------------|
-| **meta** | Meta properties for the message | N/A | None |
-| **generator_definition** | The Generator Definition | N/A | None |
-| **branch** | The branch to target | string | None |
-<!-- vale on -->
 
 <!-- vale off -->
 ### Request Graphql Query Group