Merge pull request #7433 from opsmill/pog-refactor-prefect-setup

Refactor recurring setup jobs for Prefect automations

Author: ogenstad

backend/infrahub/core/initialization.py

@@ -137,7 +137,8 @@ async def add_indexes(db: InfrahubDatabase) -> None:
     await index_manager.add()
 
 
-async def initialization(db: InfrahubDatabase, add_database_indexes: bool = False) -> None:
+async def initialization(db: InfrahubDatabase, add_database_indexes: bool = False) -> bool:
+    """Run initialization and setup, returns a boolean to indicate if it's the initial setup."""
     if config.SETTINGS.database.db_type == config.DatabaseType.MEMGRAPH:
         session = await db.session()
         await session.run(query="SET DATABASE SETTING 'log.level' TO 'INFO'")
@@ -148,6 +149,7 @@ async def initialization(db: InfrahubDatabase, add_database_indexes: bool = Fals
     # Initialize the database and Load the Root node
     # ---------------------------------------------------
     async with lock.registry.initialization():
+        first_time_initialization = len(await Root.get_list(db=db)) == 0
         log.debug("Checking Root Node")
         await initialize_registry(db=db, initialize=True)
 
@@ -210,6 +212,7 @@ async def initialization(db: InfrahubDatabase, add_database_indexes: bool = Fals
     ip_namespace = await get_default_ipnamespace(db=db)
     if ip_namespace:
         registry.default_ipnamespace = ip_namespace.id
+    return first_time_initialization
 
 
 async def create_root_node(db: InfrahubDatabase) -> Root:

backend/infrahub/core/schema/schema_branch.py

@@ -540,8 +540,8 @@ def process_validate(self) -> None:
         self.validate_identifiers()
         self.sync_uniqueness_constraints_and_unique_attributes()
         self.validate_uniqueness_constraints()
-        self.validate_display_label()
         self.validate_display_labels()
+        self.validate_display_label()
         self.validate_order_by()
         self.validate_default_filters()
         self.validate_parent_component()
@@ -789,6 +789,7 @@ def validate_display_label(self) -> None:
                     )
                 self.set(name=name, schema=update_candidate)
 
+            node_schema = self.get(name=name, duplicate=False)
             if not node_schema.display_label:
                 continue
 

backend/infrahub/server.py

@@ -84,13 +84,13 @@ async def app_initialization(application: FastAPI, enable_scheduler: bool = True
     initialize_lock(service=service)
     # We must initialize DB after initialize lock and initialize lock depends on cache initialization
     async with application.state.db.start_session() as db:
-        await initialization(db=db, add_database_indexes=True)
+        is_initial_setup = await initialization(db=db, add_database_indexes=True)
 
     async with database.start_session() as dbs:
         await validate_graph_version(db=dbs)
 
     # Initialize the workflow after the registry has been setup
-    await service.initialize_workflow()
+    await service.initialize_workflow(is_initial_setup=is_initial_setup)
 
     application.state.service = service
     application.state.response_delay = config.SETTINGS.miscellaneous.response_delay

backend/infrahub/services/__init__.py

@@ -112,15 +112,15 @@ async def new(
 
         return service
 
-    async def initialize_workflow(self) -> None:
+    async def initialize_workflow(self, is_initial_setup: bool = False) -> None:
         if self.workflow is not None and isinstance(self.workflow, WorkflowWorkerExecution):
             assert self.component is not None
             # Ideally `WorkflowWorkerExecution.initialize` would be directly part of WorkflowWorkerExecution
             # constructor but this requires some redesign as it depends on InfrahubComponent which is instantiated
             # after workflow instantiation.
             await self.component.refresh_heartbeat()
             is_primary = await self.component.is_primary_gunicorn_worker()
-            await self.workflow.initialize(component_is_primary_server=is_primary)
+            await self.workflow.initialize(component_is_primary_server=is_primary, is_initial_setup=is_initial_setup)
 
     @property
     def component(self) -> InfrahubComponent:

backend/infrahub/services/adapters/workflow/worker.py

@@ -6,7 +6,7 @@
 from prefect.deployments import run_deployment
 
 from infrahub.workers.utils import inject_context_parameter
-from infrahub.workflows.initialization import setup_task_manager
+from infrahub.workflows.initialization import setup_task_manager, setup_task_manager_identifiers
 from infrahub.workflows.models import WorkflowInfo
 
 from . import InfrahubWorkflow, Return
@@ -20,10 +20,13 @@
 
 class WorkflowWorkerExecution(InfrahubWorkflow):
     @staticmethod
-    async def initialize(component_is_primary_server: bool) -> None:
+    async def initialize(component_is_primary_server: bool, is_initial_setup: bool = False) -> None:
         if component_is_primary_server:
             await setup_task_manager()
 
+        if is_initial_setup:
+            await setup_task_manager_identifiers()
+
     @overload
     async def execute_workflow(
         self,

backend/infrahub/trigger/tasks.py

@@ -7,6 +7,7 @@
     gather_trigger_computed_attribute_python,
 )
 from infrahub.display_labels.gather import gather_trigger_display_labels_jinja2
+from infrahub.hfid.gather import gather_trigger_hfid
 from infrahub.trigger.catalogue import builtin_triggers
 from infrahub.webhook.gather import gather_trigger_webhook
 from infrahub.workers.dependencies import get_database
@@ -20,6 +21,7 @@ async def trigger_configure_all() -> None:
     async with database.start_session() as db:
         webhook_trigger = await gather_trigger_webhook(db=db)
         display_label_triggers = await gather_trigger_display_labels_jinja2()
+        human_friendly_id_triggers = await gather_trigger_hfid()
         computed_attribute_j2_triggers = await gather_trigger_computed_attribute_jinja2()
         (
             computed_attribute_python_triggers,
@@ -31,6 +33,7 @@ async def trigger_configure_all() -> None:
             + computed_attribute_python_triggers
             + computed_attribute_python_query_triggers
             + display_label_triggers
+            + human_friendly_id_triggers
             + builtin_triggers
             + webhook_trigger
             + action_rules

backend/infrahub/workflows/initialization.py

@@ -76,6 +76,11 @@ async def setup_task_manager() -> None:
         await setup_triggers(
             client=client, triggers=builtin_triggers, trigger_type=TriggerType.BUILTIN, force_update=True
         )
+
+
+@flow(name="task-manager-identifiers", flow_run_name="Setup Task Manager Display Labels and HFID")
+async def setup_task_manager_identifiers() -> None:
+    async with get_client(sync_client=False) as client:
         display_label_triggers = await gather_trigger_display_labels_jinja2()
         await setup_triggers(
             client=client,