Merge pull request #4549 from opsmill/dga-20241006-convert-flow

Convert all operations to prefect flows

Author: dgarros

backend/infrahub/message_bus/operations/__init__.py

@@ -1,6 +1,7 @@
 from typing import Optional
 
 import ujson
+from prefect import Flow
 
 from infrahub.message_bus import RPCErrorResponse, messages
 from infrahub.message_bus.operations import (
@@ -81,12 +82,17 @@
 }
 
 
-async def execute_message(routing_key: str, message_body: bytes, service: InfrahubServices) -> Optional[MessageTTL]:
+async def execute_message(
+    routing_key: str, message_body: bytes, service: InfrahubServices, skip_flow: bool = False
+) -> Optional[MessageTTL]:
     message_data = ujson.loads(message_body)
     message = messages.MESSAGE_MAP[routing_key](**message_data)
     message.set_log_data(routing_key=routing_key)
     try:
-        await COMMAND_MAP[routing_key](message=message, service=service)
+        func = COMMAND_MAP[routing_key]
+        if skip_flow and isinstance(func, Flow):
+            func = func.fn
+        await func(message=message, service=service)
     except Exception as exc:  # pylint: disable=broad-except
         if message.reply_requested:
             response = RPCErrorResponse(errors=[str(exc)], initial_message=message.model_dump())

backend/infrahub/message_bus/operations/check/artifact.py

@@ -1,5 +1,7 @@
 from typing import Union
 
+from prefect import flow
+
 from infrahub.core.constants import InfrahubKind, ValidatorConclusion
 from infrahub.core.timestamp import Timestamp
 from infrahub.git.repository import InfrahubReadOnlyRepository, InfrahubRepository
@@ -12,6 +14,7 @@
 log = get_logger()
 
 
+@flow(name="git-repository-check-artifact-create")
 async def create(message: messages.CheckArtifactCreate, service: InfrahubServices) -> None:
     log.debug("Creating artifact", message=message)
     validator = await service.client.get(

backend/infrahub/message_bus/operations/check/generator.py

@@ -3,6 +3,7 @@
 from infrahub_sdk import InfrahubNode
 from infrahub_sdk.exceptions import ModuleImportError
 from infrahub_sdk.schema import InfrahubGeneratorDefinitionConfig
+from prefect import flow
 
 from infrahub import lock
 from infrahub.core.constants import GeneratorInstanceStatus, InfrahubKind, ValidatorConclusion
@@ -16,6 +17,7 @@
 # pylint: disable=duplicate-code
 
 
+@flow(name="git-repository-check-generator-run")
 async def run(message: messages.CheckGeneratorRun, service: InfrahubServices) -> None:
     repository = await get_initialized_repo(
         repository_id=message.repository_id,

backend/infrahub/message_bus/operations/check/repository.py

@@ -1,6 +1,7 @@
 from typing import List
 
 from infrahub_sdk import UUIDT
+from prefect import flow
 
 from infrahub import lock
 from infrahub.core.constants import InfrahubKind
@@ -16,6 +17,7 @@
 log = get_logger()
 
 
+@flow(name="git-repository-check-definition")
 async def check_definition(message: messages.CheckRepositoryCheckDefinition, service: InfrahubServices) -> None:
     definition = await service.client.get(
         kind=InfrahubKind.CHECKDEFINITION, id=message.check_definition_id, branch=message.branch_name
@@ -141,6 +143,7 @@ async def check_definition(message: messages.CheckRepositoryCheckDefinition, ser
             await service.send(message=event)
 
 
+@flow(name="git-repository-check-merge-conflict")
 async def merge_conflicts(message: messages.CheckRepositoryMergeConflicts, service: InfrahubServices) -> None:
     """Runs a check to see if there are merge conflicts between two branches."""
     log.info(
@@ -222,6 +225,7 @@ async def merge_conflicts(message: messages.CheckRepositoryMergeConflicts, servi
     )
 
 
+@flow(name="git-repository-user-check")
 async def user_check(message: messages.CheckRepositoryUserCheck, service: InfrahubServices) -> None:
     validator = await service.client.get(kind=InfrahubKind.USERVALIDATOR, id=message.validator_id)
     await validator.checks.fetch()

backend/infrahub/message_bus/operations/event/branch.py

@@ -1,5 +1,7 @@
 from typing import List
 
+from prefect import flow
+
 from infrahub.core import registry
 from infrahub.core.diff.model.path import BranchTrackingId
 from infrahub.core.diff.repository.repository import DiffRepository
@@ -11,6 +13,7 @@
 log = get_logger()
 
 
+@flow(name="event-branch-create")
 async def create(message: messages.EventBranchCreate, service: InfrahubServices) -> None:
     log.info("run_message", branch=message.branch)
 
@@ -23,6 +26,7 @@ async def create(message: messages.EventBranchCreate, service: InfrahubServices)
         await service.send(message=event)
 
 
+@flow(name="event-branch-delete")
 async def delete(message: messages.EventBranchDelete, service: InfrahubServices) -> None:
     log.info("Branch was deleted", branch=message.branch)
 
@@ -36,6 +40,7 @@ async def delete(message: messages.EventBranchDelete, service: InfrahubServices)
         await service.send(message=event)
 
 
+@flow(name="branch-event-merge")
 async def merge(message: messages.EventBranchMerge, service: InfrahubServices) -> None:
     log.info("Branch merged", source_branch=message.source_branch, target_branch=message.target_branch)
 
@@ -64,6 +69,7 @@ async def merge(message: messages.EventBranchMerge, service: InfrahubServices) -
         await service.send(message=event)
 
 
+@flow(name="event-branch-rebased")
 async def rebased(message: messages.EventBranchRebased, service: InfrahubServices) -> None:
     log.info("Branch rebased", branch=message.branch)
 

backend/infrahub/message_bus/operations/event/node.py

@@ -1,5 +1,7 @@
 from typing import List
 
+from prefect import flow
+
 from infrahub.core.constants import InfrahubKind
 from infrahub.log import get_logger
 from infrahub.message_bus import InfrahubMessage, messages
@@ -8,6 +10,7 @@
 log = get_logger()
 
 
+@flow(name="event-node-mutated")
 async def mutated(
     message: messages.EventNodeMutated,
     service: InfrahubServices,

backend/infrahub/message_bus/operations/event/schema.py

@@ -1,10 +1,13 @@
+from prefect import flow
+
 from infrahub.log import get_logger
 from infrahub.message_bus import messages
 from infrahub.services import InfrahubServices
 
 log = get_logger()
 
 
+@flow(name="event-schema-update")
 async def update(message: messages.EventSchemaUpdate, service: InfrahubServices) -> None:
     log.info("run_message", branch=message.branch)
 

backend/infrahub/message_bus/operations/event/worker.py

@@ -1,7 +1,10 @@
+from prefect import flow
+
 from infrahub.message_bus import messages
 from infrahub.services import InfrahubServices
 
 
+@flow(name="event-worker-newprimary-api")
 async def new_primary_api(message: messages.EventWorkerNewPrimaryAPI, service: InfrahubServices) -> None:
     service.log.info("api_worker promoted to primary", worker_id=message.worker_id)
 

backend/infrahub/message_bus/operations/finalize/validator.py

@@ -1,3 +1,5 @@
+from prefect import flow
+
 from infrahub import config
 from infrahub.core.timestamp import Timestamp
 from infrahub.log import get_logger
@@ -8,6 +10,7 @@
 log = get_logger()
 
 
+@flow(name="validator-finalize-execution")
 async def execution(message: messages.FinalizeValidatorExecution, service: InfrahubServices) -> None:
     """Monitors the status of checks associated with a validator and finalizes the conclusion of the validator
 

backend/infrahub/message_bus/operations/git/branch.py

@@ -1,3 +1,5 @@
+from prefect import flow
+
 from infrahub import lock
 from infrahub.git.repository import InfrahubRepository
 from infrahub.log import get_logger
@@ -7,6 +9,7 @@
 log = get_logger()
 
 
+@flow(name="git-repository-branch-create")
 async def create(message: messages.GitBranchCreate, service: InfrahubServices) -> None:
     log.info("creating branch in repository", branch=message.branch, repository=message.repository_name)
     repo = await InfrahubRepository.init(id=message.repository_id, name=message.repository_name, client=service.client)