Migrate RequestProposedChangeDataIntegrity to prefect

Author: LucasG0

backend/infrahub/message_bus/messages/__init__.py

@@ -17,7 +17,6 @@
 from .git_file_get import GitFileGet, GitFileGetResponse
 from .git_repository_connectivity import GitRepositoryConnectivity
 from .git_repository_importobjects import GitRepositoryImportObjects
-from .proposed_change.request_proposedchange_dataintegrity import RequestProposedChangeDataIntegrity
 from .proposed_change.request_proposedchange_refreshartifacts import RequestProposedChangeRefreshArtifacts
 from .proposed_change.request_proposedchange_repositorychecks import RequestProposedChangeRepositoryChecks
 from .proposed_change.request_proposedchange_rungenerators import RequestProposedChangeRunGenerators
@@ -53,6 +52,7 @@
     "git.diff.names_only": GitDiffNamesOnly,
     "git.file.get": GitFileGet,
     "git.repository.connectivity": GitRepositoryConnectivity,
+    "git.repository.add_read_only": GitRepositoryAddReadOnly,
     "git.repository.import_objects": GitRepositoryImportObjects,
     "schema.migration.path": SchemaMigrationPath,
     "schema.validator.path": SchemaValidatorPath,

backend/infrahub/message_bus/operations/__init__.py

@@ -41,7 +41,6 @@
     "refresh.webhook.configuration": refresh.webhook.configuration,
     "request.generator_definition.check": requests.generator_definition.check,
     "request.artifact_definition.check": requests.artifact_definition.check,
-    "request.proposed_change.data_integrity": requests.proposed_change.data_integrity,
     "request.proposed_change.pipeline": requests.proposed_change.pipeline,
     "request.proposed_change.refresh_artifacts": requests.proposed_change.refresh_artifacts,
     "request.proposed_change.repository_checks": requests.proposed_change.repository_checks,

backend/infrahub/message_bus/operations/requests/proposed_change.py

@@ -32,8 +32,10 @@
     ProposedChangeRepository,
     ProposedChangeSubscriber,
 )
+from infrahub.proposed_change.models import RequestProposedChangeDataIntegrity
 from infrahub.pytest_plugin import InfrahubBackendPlugin
 from infrahub.services import InfrahubServices  # noqa: TCH001
+from infrahub.workflows.catalogue import REQUEST_PROPOSED_CHANGE_DATA_INTEGRITY
 
 if TYPE_CHECKING:
     from infrahub_sdk.node import InfrahubNode
@@ -71,23 +73,6 @@ def log_line(self) -> str:
         return "Doesn't require changes due to no relevant modified kinds or file changes in Git"
 
 
-@flow(name="proposed-changed-data-integrity")
-async def data_integrity(message: messages.RequestProposedChangeDataIntegrity, service: InfrahubServices) -> None:
-    """Triggers a data integrity validation check on the provided proposed change to start."""
-    async with service.task_report(
-        related_node=message.proposed_change,
-        title="Data Integrity",
-    ):
-        log.info(f"Got a request to process data integrity defined in proposed_change: {message.proposed_change}")
-
-        destination_branch = await registry.get_branch(db=service.database, branch=message.destination_branch)
-        source_branch = await registry.get_branch(db=service.database, branch=message.source_branch)
-        component_registry = get_component_registry()
-        async with service.database.start_transaction() as dbt:
-            diff_coordinator = await component_registry.get_component(DiffCoordinator, db=dbt, branch=source_branch)
-            await diff_coordinator.update_branch_diff(base_branch=destination_branch, diff_branch=source_branch)
-
-
 @flow(name="proposed-changed-pipeline")
 async def pipeline(message: messages.RequestProposedChangePipeline, service: InfrahubServices) -> None:
     async with service.task_report(
@@ -164,14 +149,15 @@ async def pipeline(message: messages.RequestProposedChangePipeline, service: Inf
             branch=message.source_branch
         ):
             await task_report.info("Adding Data Integrity job", proposed_change=message.proposed_change)
-            events.append(
-                messages.RequestProposedChangeDataIntegrity(
-                    proposed_change=message.proposed_change,
-                    source_branch=message.source_branch,
-                    source_branch_sync_with_git=message.source_branch_sync_with_git,
-                    destination_branch=message.destination_branch,
-                    branch_diff=branch_diff,
-                )
+            model = RequestProposedChangeDataIntegrity(
+                proposed_change=message.proposed_change,
+                source_branch=message.source_branch,
+                source_branch_sync_with_git=message.source_branch_sync_with_git,
+                destination_branch=message.destination_branch,
+                branch_diff=branch_diff,
+            )
+            await service.workflow.submit_workflow(
+                workflow=REQUEST_PROPOSED_CHANGE_DATA_INTEGRITY, parameters={"model": model}
             )
 
         if message.check_type in [CheckType.REPOSITORY, CheckType.USER]:

backend/infrahub/message_bus/messages/proposed_change/request_proposedchange_dataintegrity.py renamed to backend/infrahub/proposed_change/models.py

@@ -1,4 +1,4 @@
-from .base_with_diff import BaseProposedChangeWithDiffMessage
+from infrahub.message_bus.messages.proposed_change.base_with_diff import BaseProposedChangeWithDiffMessage
 
 
 class RequestProposedChangeDataIntegrity(BaseProposedChangeWithDiffMessage):

backend/infrahub/proposed_change/tasks.py

@@ -1,13 +1,19 @@
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 from infrahub_sdk.protocols import CoreProposedChange
 from prefect import flow, task
 from prefect.logging import get_run_logger
 
+from infrahub.core import registry
 from infrahub.core.constants import ProposedChangeState
-from infrahub.services import (
-    services,
-)
+from infrahub.core.diff.coordinator import DiffCoordinator
+from infrahub.dependencies.registry import get_component_registry
+from infrahub.services import services
+
+if TYPE_CHECKING:
+    from infrahub.proposed_change.models import RequestProposedChangeDataIntegrity
 
 
 @flow(name="proposed-changes-cancel-branch", description="Cancel all Proposed change associated with a branch.")
@@ -39,3 +45,26 @@ async def cancel_proposed_change(proposed_change: CoreProposedChange) -> None:
     proposed_change = await service.client.get(kind=CoreProposedChange, id=proposed_change.id)
     proposed_change.state.value = ProposedChangeState.CANCELED.value
     await proposed_change.save()
+
+
+@flow(
+    name="proposed-changed-data-integrity",
+    flow_run_name="Triggers data integrity check on proposed change {model.proposed_change}",
+)
+async def run_proposed_change_data_integrity_check(model: RequestProposedChangeDataIntegrity) -> None:
+    """Triggers a data integrity validation check on the provided proposed change to start."""
+
+    service = services.service
+    async with service.task_report(
+        related_node=model.proposed_change,
+        title="Data Integrity",
+    ):
+        log = get_run_logger()
+        log.info(f"Got a request to process data integrity defined in proposed_change: {model.proposed_change}")
+
+        destination_branch = await registry.get_branch(db=service.database, branch=model.destination_branch)
+        source_branch = await registry.get_branch(db=service.database, branch=model.source_branch)
+        component_registry = get_component_registry()
+        async with service.database.start_transaction() as dbt:
+            diff_coordinator = await component_registry.get_component(DiffCoordinator, db=dbt, branch=source_branch)
+            await diff_coordinator.update_branch_diff(base_branch=destination_branch, diff_branch=source_branch)

backend/infrahub/workflows/catalogue.py

@@ -242,6 +242,13 @@
     function="process_transform",
 )
 
+REQUEST_PROPOSED_CHANGE_DATA_INTEGRITY = WorkflowDefinition(
+    name="proposed-changed-data-integrity",
+    type=WorkflowType.INTERNAL,
+    module="infrahub.proposed_change.tasks",
+    function="run_proposed_change_data_integrity_check",
+)
+
 
 worker_pools = [INFRAHUB_WORKER_POOL]