Merge pull request #6475 from opsmill/pog-global-trigger-locks

Add central locking options for updating triggers

Author: ogenstad

backend/infrahub/actions/tasks.py

@@ -2,13 +2,11 @@
 
 from infrahub_sdk.graphql import Mutation
 from prefect import flow
-from prefect.client.orchestration import get_client
 
-from infrahub import lock
 from infrahub.context import InfrahubContext  # noqa: TC001  needed for prefect flow
 from infrahub.services import InfrahubServices  # noqa: TC001  needed for prefect flow
 from infrahub.trigger.models import TriggerType
-from infrahub.trigger.setup import setup_triggers
+from infrahub.trigger.setup import setup_triggers_specific
 from infrahub.workflows.utils import add_tags
 
 from .gather import gather_trigger_action_rules
@@ -101,14 +99,9 @@ async def run_generator_group_event(
 async def configure_action_rules(
     service: InfrahubServices,
 ) -> None:
-    async with lock.registry.get(name="configure-action-rules", namespace="trigger-rules", local=False):
-        triggers = await gather_trigger_action_rules(db=service.database)
-        async with get_client(sync_client=False) as prefect_client:
-            await setup_triggers(
-                client=prefect_client,
-                triggers=triggers,
-                trigger_type=TriggerType.ACTION_TRIGGER_RULE,
-            )  # type: ignore[misc]
+    await setup_triggers_specific(
+        gatherer=gather_trigger_action_rules, trigger_type=TriggerType.ACTION_TRIGGER_RULE, db=service.database
+    )  # type: ignore[misc]
 
 
 async def _run_generator(

backend/infrahub/computed_attribute/gather.py

@@ -100,7 +100,9 @@ async def gather_python_transform_attributes(
     name="gather-trigger-computed-attribute-jinja2",
     cache_policy=NONE,
 )
-async def gather_trigger_computed_attribute_jinja2() -> list[ComputedAttrJinja2TriggerDefinition]:
+async def gather_trigger_computed_attribute_jinja2(
+    db: InfrahubDatabase | None = None,  # noqa: ARG001 Needed to have a common function signature for gathering functions
+) -> list[ComputedAttrJinja2TriggerDefinition]:
     log = get_run_logger()
 
     # Build a list of all branches to process based on which branch is different from main

backend/infrahub/computed_attribute/tasks.py

@@ -15,7 +15,7 @@
 from infrahub.git.repository import get_initialized_repo
 from infrahub.services import InfrahubServices  # noqa: TC001  needed for prefect flow
 from infrahub.trigger.models import TriggerSetupReport, TriggerType
-from infrahub.trigger.setup import setup_triggers
+from infrahub.trigger.setup import setup_triggers, setup_triggers_specific
 from infrahub.workflows.catalogue import (
     COMPUTED_ATTRIBUTE_PROCESS_JINJA2,
     COMPUTED_ATTRIBUTE_PROCESS_TRANSFORM,
@@ -304,36 +304,30 @@ async def computed_attribute_setup_jinja2(
             await add_tags(branches=[branch_name])
             await wait_for_schema_to_converge(branch_name=branch_name, component=service.component, db=db, log=log)
 
-        triggers = await gather_trigger_computed_attribute_jinja2()
-
+        report: TriggerSetupReport = await setup_triggers_specific(
+            gatherer=gather_trigger_computed_attribute_jinja2, trigger_type=TriggerType.COMPUTED_ATTR_JINJA2
+        )  # type: ignore[misc]
         # Configure all ComputedAttrJinja2Trigger in Prefect
-        async with get_client(sync_client=False) as prefect_client:
-            report: TriggerSetupReport = await setup_triggers(
-                client=prefect_client,
-                triggers=triggers,
-                trigger_type=TriggerType.COMPUTED_ATTR_JINJA2,
-                force_update=False,
-            )  # type: ignore[misc]
 
-            # Since we can have multiple trigger per NodeKind
-            # we need to extract the list of unique node that should be processed
-            unique_nodes: set[tuple[str, str, str]] = {
-                (trigger.branch, trigger.computed_attribute.kind, trigger.computed_attribute.attribute.name)  # type: ignore[attr-defined]
-                for trigger in report.updated + report.created
-            }
-            for branch, kind, attribute_name in unique_nodes:
-                if event_name != BranchDeletedEvent.event_name and branch == branch_name:
-                    await service.workflow.submit_workflow(
-                        workflow=TRIGGER_UPDATE_JINJA_COMPUTED_ATTRIBUTES,
-                        context=context,
-                        parameters={
-                            "branch_name": branch,
-                            "computed_attribute_name": attribute_name,
-                            "computed_attribute_kind": kind,
-                        },
-                    )
-
-        log.info(f"{len(triggers)} Computed Attribute for Jinja2 automation configuration completed")
+        # Since we can have multiple trigger per NodeKind
+        # we need to extract the list of unique node that should be processed
+        unique_nodes: set[tuple[str, str, str]] = {
+            (trigger.branch, trigger.computed_attribute.kind, trigger.computed_attribute.attribute.name)  # type: ignore[attr-defined]
+            for trigger in report.updated + report.created
+        }
+        for branch, kind, attribute_name in unique_nodes:
+            if event_name != BranchDeletedEvent.event_name and branch == branch_name:
+                await service.workflow.submit_workflow(
+                    workflow=TRIGGER_UPDATE_JINJA_COMPUTED_ATTRIBUTES,
+                    context=context,
+                    parameters={
+                        "branch_name": branch,
+                        "computed_attribute_name": attribute_name,
+                        "computed_attribute_kind": kind,
+                    },
+                )
+
+        log.info(f"{report.in_use_count} Computed Attribute for Jinja2 automation configuration completed")
 
 
 @flow(

backend/infrahub/trigger/models.py

@@ -28,6 +28,10 @@ class TriggerSetupReport(BaseModel):
     deleted: list[Automation] = Field(default_factory=list)
     unchanged: list[TriggerDefinition] = Field(default_factory=list)
 
+    @property
+    def in_use_count(self) -> int:
+        return len(self.created + self.updated + self.unchanged)
+
 
 class TriggerType(str, Enum):
     ACTION_TRIGGER_RULE = "action_trigger_rule"

backend/infrahub/trigger/setup.py

@@ -1,12 +1,14 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Awaitable, Callable
 
 from prefect import get_run_logger, task
 from prefect.automations import AutomationCore
 from prefect.cache_policies import NONE
-from prefect.client.orchestration import PrefectClient
+from prefect.client.orchestration import PrefectClient, get_client
 from prefect.client.schemas.filters import DeploymentFilter, DeploymentFilterName
 from prefect.events.schemas.automations import Automation
 
+from infrahub import lock
+from infrahub.database import InfrahubDatabase
 from infrahub.trigger.models import TriggerDefinition
 
 from .models import TriggerSetupReport, TriggerType
@@ -27,6 +29,28 @@ def compare_automations(target: AutomationCore, existing: Automation) -> bool:
     return target_dump == existing_dump
 
 
+@task(name="trigger-setup-specific", task_run_name="Setup triggers of a specific kind", cache_policy=NONE)  # type: ignore[arg-type]
+async def setup_triggers_specific(
+    gatherer: Callable[[InfrahubDatabase | None], Awaitable[list[TriggerDefinition]]],
+    trigger_type: TriggerType,
+    db: InfrahubDatabase | None = None,
+) -> TriggerSetupReport:
+    async with lock.registry.get(
+        name=f"configure-action-rules-{trigger_type.value}", namespace="trigger-rules", local=False
+    ):
+        if db:
+            async with db.start_session(read_only=True) as dbs:
+                triggers = await gatherer(dbs)
+        else:
+            triggers = await gatherer(db)
+        async with get_client(sync_client=False) as prefect_client:
+            return await setup_triggers(
+                client=prefect_client,
+                triggers=triggers,
+                trigger_type=trigger_type,
+            )  # type: ignore[misc]
+
+
 @task(name="trigger-setup", task_run_name="Setup triggers", cache_policy=NONE)  # type: ignore[arg-type]
 async def setup_triggers(
     client: PrefectClient,

backend/infrahub/webhook/tasks.py

@@ -14,7 +14,7 @@
 from infrahub.message_bus.types import KVTTL
 from infrahub.services import InfrahubServices  # noqa: TC001  needed for prefect flow
 from infrahub.trigger.models import TriggerType
-from infrahub.trigger.setup import setup_triggers
+from infrahub.trigger.setup import setup_triggers_specific
 from infrahub.workflows.utils import add_tags
 
 from .gather import gather_trigger_webhook
@@ -114,14 +114,10 @@ async def configure_webhook_all(service: InfrahubServices) -> None:
     async with service.database.start_session(read_only=True) as db:
         triggers = await gather_trigger_webhook(db=db)
 
-    async with get_client(sync_client=False) as prefect_client:
-        await setup_triggers(
-            client=prefect_client,
-            triggers=triggers,
-            trigger_type=TriggerType.WEBHOOK,
-        )  # type: ignore[misc]
-
     log.info(f"{len(triggers)} Webhooks automation configuration completed")
+    await setup_triggers_specific(
+        gatherer=gather_trigger_webhook, db=service.database, trigger_type=TriggerType.WEBHOOK
+    )  # type: ignore[misc]
 
 
 @flow(name="webhook-setup-automation-one", flow_run_name="Configurate webhook for {webhook_name}")