From 0f5a1df5179ebd79464429b7eff0303b3c43c884 Mon Sep 17 00:00:00 2001
From: Damien Garros <dgarros@gmail.com>
Date: Tue, 14 Oct 2025 05:56:44 +0200
Subject: [PATCH] Extend GeneratorDefinition with flags to control execution

---
 infrahub_sdk/ctl/generator.py     |  4 ++++
 infrahub_sdk/generator.py         |  8 +++++++-
 infrahub_sdk/protocols.py         | 14 ++++++++++++++
 infrahub_sdk/schema/repository.py |  8 ++++++++
 4 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/infrahub_sdk/ctl/generator.py b/infrahub_sdk/ctl/generator.py
index c0f60c52..75354ee1 100644
--- a/infrahub_sdk/ctl/generator.py
+++ b/infrahub_sdk/ctl/generator.py
@@ -64,6 +64,8 @@ async def run(
             branch=branch or "",
             params=variables_dict,
             convert_query_response=generator_config.convert_query_response,
+            execute_in_proposed_change=generator_config.execute_in_proposed_change,
+            execute_after_merge=generator_config.execute_after_merge,
             infrahub_node=InfrahubNode,
         )
         await generator._init_client.schema.all(branch=generator.branch_name)
@@ -93,6 +95,8 @@ async def run(
                 branch=branch or "",
                 params=params,
                 convert_query_response=generator_config.convert_query_response,
+                execute_in_proposed_change=generator_config.execute_in_proposed_change,
+                execute_after_merge=generator_config.execute_after_merge,
                 infrahub_node=InfrahubNode,
             )
             data = execute_graphql_query(
diff --git a/infrahub_sdk/generator.py b/infrahub_sdk/generator.py
index 3c9d26d7..20b7dc88 100644
--- a/infrahub_sdk/generator.py
+++ b/infrahub_sdk/generator.py
@@ -26,6 +26,8 @@ def __init__(
         generator_instance: str = "",
         params: dict | None = None,
         convert_query_response: bool = False,
+        execute_in_proposed_change: bool = True,
+        execute_after_merge: bool = True,
         logger: logging.Logger | None = None,
         request_context: RequestContext | None = None,
     ) -> None:
@@ -44,6 +46,8 @@ def __init__(
         self._client: InfrahubClient | None = None
         self.logger = logger if logger else logging.getLogger("infrahub.tasks")
         self.request_context = request_context
+        self.execute_in_proposed_change = execute_in_proposed_change
+        self.execute_after_merge = execute_after_merge
 
     @property
     def subscribers(self) -> list[str] | None:
@@ -81,8 +85,10 @@ async def run(self, identifier: str, data: dict | None = None) -> None:
         unpacked = data.get("data") or data
         await self.process_nodes(data=unpacked)
 
+        group_type = "CoreGeneratorGroup" if self.execute_after_merge else "CoreGeneratorAwareGroup"
+
         async with self._init_client.start_tracking(
-            identifier=identifier, params=self.params, delete_unused_nodes=True, group_type="CoreGeneratorGroup"
+            identifier=identifier, params=self.params, delete_unused_nodes=True, group_type=group_type
         ) as self.client:
             await self.generate(data=unpacked)
 
diff --git a/infrahub_sdk/protocols.py b/infrahub_sdk/protocols.py
index 6abf6a2b..b3752bed 100644
--- a/infrahub_sdk/protocols.py
+++ b/infrahub_sdk/protocols.py
@@ -131,6 +131,7 @@ class CoreGenericRepository(CoreNode):
     queries: RelationshipManager
     checks: RelationshipManager
     generators: RelationshipManager
+    groups_objects: RelationshipManager
 
 
 class CoreGroup(CoreNode):
@@ -355,6 +356,10 @@ class CoreGeneratorAction(CoreAction):
     generator: RelatedNode
 
 
+class CoreGeneratorAwareGroup(CoreGroup):
+    pass
+
+
 class CoreGeneratorCheck(CoreCheck):
     instance: String
 
@@ -366,6 +371,8 @@ class CoreGeneratorDefinition(CoreTaskTarget):
     file_path: String
     class_name: String
     convert_query_response: BooleanOptional
+    execute_in_proposed_change: BooleanOptional
+    execute_after_merge: BooleanOptional
     query: RelatedNode
     repository: RelatedNode
     targets: RelatedNode
@@ -681,6 +688,7 @@ class CoreGenericRepositorySync(CoreNodeSync):
     queries: RelationshipManagerSync
     checks: RelationshipManagerSync
     generators: RelationshipManagerSync
+    groups_objects: RelationshipManagerSync
 
 
 class CoreGroupSync(CoreNodeSync):
@@ -905,6 +913,10 @@ class CoreGeneratorActionSync(CoreActionSync):
     generator: RelatedNodeSync
 
 
+class CoreGeneratorAwareGroupSync(CoreGroupSync):
+    pass
+
+
 class CoreGeneratorCheckSync(CoreCheckSync):
     instance: String
 
@@ -916,6 +928,8 @@ class CoreGeneratorDefinitionSync(CoreTaskTargetSync):
     file_path: String
     class_name: String
     convert_query_response: BooleanOptional
+    execute_in_proposed_change: BooleanOptional
+    execute_after_merge: BooleanOptional
     query: RelatedNodeSync
     repository: RelatedNodeSync
     targets: RelatedNodeSync
diff --git a/infrahub_sdk/schema/repository.py b/infrahub_sdk/schema/repository.py
index 832651c9..7a793d01 100644
--- a/infrahub_sdk/schema/repository.py
+++ b/infrahub_sdk/schema/repository.py
@@ -96,6 +96,14 @@ class InfrahubGeneratorDefinitionConfig(InfrahubRepositoryConfigElement):
         default=False,
         description="Decide if the generator should convert the result of the GraphQL query to SDK InfrahubNode objects.",
     )
+    execute_in_proposed_change: bool = Field(
+        default=True,
+        description="Decide if the generator should execute in a proposed change.",
+    )
+    execute_after_merge: bool = Field(
+        default=True,
+        description="Decide if the generator should execute after a merge.",
+    )
 
     def load_class(self, import_root: str | None = None, relative_path: str | None = None) -> type[InfrahubGenerator]:
         module = import_module(module_path=self.file_path, import_root=import_root, relative_path=relative_path)