Merge pull request #4874 from opsmill/dga-20241107-pc-merge

Convert proposed change merge to a prefect task

Author: dgarros

backend/infrahub/core/branch/tasks.py

@@ -148,68 +148,71 @@ async def merge_branch(branch: str) -> None:
     await add_branch_tag(branch_name=branch)
     await add_branch_tag(branch_name=registry.default_branch)
 
-    obj = await Branch.get_by_name(db=service.database, name=branch)
-    component_registry = get_component_registry()
-
-    merger: BranchMerger | None = None
-    async with lock.registry.global_graph_lock():
-        diff_coordinator = await component_registry.get_component(DiffCoordinator, db=service.database, branch=obj)
-        diff_merger = await component_registry.get_component(DiffMerger, db=service.database, branch=obj)
-        merger = BranchMerger(
-            db=service.database,
-            diff_coordinator=diff_coordinator,
-            diff_merger=diff_merger,
-            source_branch=obj,
-            service=service,
-        )
-        try:
-            await merger.merge()
-        except Exception as exc:
-            await merger.rollback()
-            raise MergeFailedError(branch_name=branch) from exc
-        await merger.update_schema()
-
-    if merger and merger.migrations:
-        errors = await schema_apply_migrations(
-            message=SchemaApplyMigrationData(
-                branch=merger.destination_branch,
-                new_schema=merger.destination_schema,
-                previous_schema=merger.initial_source_schema,
-                migrations=merger.migrations,
+    async with service.database.start_session() as db:
+        obj = await Branch.get_by_name(db=db, name=branch)
+        component_registry = get_component_registry()
+
+        merger: BranchMerger | None = None
+        async with lock.registry.global_graph_lock():
+            # await update_diff(model=RequestDiffUpdate(branch_name=obj.name))
+
+            diff_coordinator = await component_registry.get_component(DiffCoordinator, db=db, branch=obj)
+            diff_merger = await component_registry.get_component(DiffMerger, db=db, branch=obj)
+            merger = BranchMerger(
+                db=db,
+                diff_coordinator=diff_coordinator,
+                diff_merger=diff_merger,
+                source_branch=obj,
+                service=service,
             )
+            try:
+                await merger.merge()
+            except Exception as exc:
+                await merger.rollback()
+                raise MergeFailedError(branch_name=branch) from exc
+            await merger.update_schema()
+
+        if merger and merger.migrations:
+            errors = await schema_apply_migrations(
+                message=SchemaApplyMigrationData(
+                    branch=merger.destination_branch,
+                    new_schema=merger.destination_schema,
+                    previous_schema=merger.initial_source_schema,
+                    migrations=merger.migrations,
+                )
+            )
+            for error in errors:
+                log.error(error)
+
+        # -------------------------------------------------------------
+        # Trigger the reconciliation of IPAM data after the merge
+        # -------------------------------------------------------------
+        differ = await merger.get_graph_diff()
+        diff_parser = IpamDiffParser(
+            db=db,
+            differ=differ,
+            source_branch_name=obj.name,
+            target_branch_name=registry.default_branch,
+        )
+        ipam_node_details = await diff_parser.get_changed_ipam_node_details()
+        await service.workflow.submit_workflow(
+            workflow=IPAM_RECONCILIATION,
+            parameters={"branch": registry.default_branch, "ipam_node_details": ipam_node_details},
         )
-        for error in errors:
-            log.error(error)
-
-    # -------------------------------------------------------------
-    # Trigger the reconciliation of IPAM data after the merge
-    # -------------------------------------------------------------
-    differ = await merger.get_graph_diff()
-    diff_parser = IpamDiffParser(
-        db=service.database,
-        differ=differ,
-        source_branch_name=obj.name,
-        target_branch_name=registry.default_branch,
-    )
-    ipam_node_details = await diff_parser.get_changed_ipam_node_details()
-    await service.workflow.submit_workflow(
-        workflow=IPAM_RECONCILIATION,
-        parameters={"branch": registry.default_branch, "ipam_node_details": ipam_node_details},
-    )
 
-    # -------------------------------------------------------------
-    # Generate an event to indicate that a branch has been merged
-    # NOTE: we still need to convert this event and potentially pull
-    #   some tasks currently executed based on the event into this workflow
-    # -------------------------------------------------------------
-    log_data = get_log_data()
-    request_id = log_data.get("request_id", "")
-    message = messages.EventBranchMerge(
-        source_branch=obj.name,
-        target_branch=registry.default_branch,
-        meta=Meta(initiator_id=WORKER_IDENTITY, request_id=request_id),
-    )
-    await service.send(message=message)
+        # -------------------------------------------------------------
+        # Generate an event to indicate that a branch has been merged
+        # NOTE: we still need to convert this event and potentially pull
+        #   some tasks currently executed based on the event into this workflow
+        # -------------------------------------------------------------
+        log_data = get_log_data()
+        request_id = log_data.get("request_id", "")
+        message = messages.EventBranchMerge(
+            source_branch=obj.name,
+            target_branch=registry.default_branch,
+            meta=Meta(initiator_id=WORKER_IDENTITY, request_id=request_id),
+        )
+        await service.send(message=message)
 
 
 @flow(name="branch-delete", flow_run_name="Delete branch {branch}")

backend/infrahub/core/constants/__init__.py

@@ -15,6 +15,7 @@
     "SchemaElementPathType",
     "UpdateSupport",
     "UpdateValidationErrorType",
+    "ValidationError",
 ]
 
 
@@ -207,37 +208,6 @@ def from_relationship(cls, relationship: RelationshipCardinality) -> PathType:
         return cls("relationship_many")
 
 
-class ProposedChangeState(InfrahubStringEnum):
-    OPEN = "open"
-    MERGED = "merged"
-    CLOSED = "closed"
-    CANCELED = "canceled"
-
-    def validate_state_check_run(self) -> None:
-        if self == ProposedChangeState.OPEN:
-            return
-
-        raise ValidationError(input_value="Unable to trigger check on proposed changes that aren't in the open state")
-
-    def validate_editability(self) -> None:
-        if self in [ProposedChangeState.CANCELED, ProposedChangeState.MERGED]:
-            raise ValidationError(
-                input_value=f"A proposed change in the {self.value} state is not allowed to be updated"
-            )
-
-    def validate_state_transition(self, updated_state: ProposedChangeState) -> None:
-        if self == ProposedChangeState.OPEN:
-            return
-
-        if self == ProposedChangeState.CLOSED and updated_state not in [
-            ProposedChangeState.CANCELED,
-            ProposedChangeState.OPEN,
-        ]:
-            raise ValidationError(
-                input_value="A closed proposed change is only allowed to transition to the open state"
-            )
-
-
 class HashableModelState(InfrahubStringEnum):
     PRESENT = "present"
     ABSENT = "absent"

backend/infrahub/core/diff/data_check_synchronizer.py

@@ -1,11 +1,12 @@
 from enum import Enum
 
-from infrahub.core.constants import BranchConflictKeep, InfrahubKind, ProposedChangeState
+from infrahub.core.constants import BranchConflictKeep, InfrahubKind
 from infrahub.core.integrity.object_conflict.conflict_recorder import ObjectConflictValidatorRecorder
 from infrahub.core.manager import NodeManager
 from infrahub.core.node import Node
 from infrahub.database import InfrahubDatabase
 from infrahub.exceptions import SchemaNotFoundError
+from infrahub.proposed_change.constants import ProposedChangeState
 
 from .conflicts_extractor import DiffConflictsExtractor
 from .model.path import ConflictSelection, EnrichedDiffConflict, EnrichedDiffRoot

backend/infrahub/core/diff/tasks.py

@@ -6,13 +6,16 @@
 from infrahub.dependencies.registry import get_component_registry
 from infrahub.log import get_logger
 from infrahub.services import services
+from infrahub.workflows.utils import add_branch_tag
 
 log = get_logger()
 
 
-@flow(name="diff-update")
+@flow(name="diff-update", flow_run_name="Update diff for branch {model.branch_name}")
 async def update_diff(model: RequestDiffUpdate) -> None:
     service = services.service
+    await add_branch_tag(branch_name=model.branch_name)
+
     component_registry = get_component_registry()
     base_branch = await registry.get_branch(db=service.database, branch=registry.default_branch)
     diff_branch = await registry.get_branch(db=service.database, branch=model.branch_name)
@@ -28,9 +31,10 @@ async def update_diff(model: RequestDiffUpdate) -> None:
     )
 
 
-@flow(name="diff-refresh")
+@flow(name="diff-refresh", flow_run_name="Recreate diff for branch {model.branch_name}")
 async def refresh_diff(model: RequestDiffRefresh) -> None:
     service = services.service
+    await add_branch_tag(branch_name=model.branch_name)
 
     component_registry = get_component_registry()
     base_branch = await registry.get_branch(db=service.database, branch=registry.default_branch)

backend/infrahub/core/relationship/model.py

@@ -843,23 +843,23 @@ async def get_peer(
     async def get_peers(
         self,
         db: InfrahubDatabase,
-        branch_agnostic: bool,
         peer_type: type[PeerType],
+        branch_agnostic: bool = ...,
     ) -> Mapping[str, PeerType]: ...
 
     @overload
     async def get_peers(
         self,
         db: InfrahubDatabase,
-        branch_agnostic: bool,
         peer_type: Literal[None] = None,
+        branch_agnostic: bool = ...,
     ) -> Mapping[str, Node]: ...
 
     async def get_peers(
         self,
         db: InfrahubDatabase,
-        branch_agnostic: bool = False,
         peer_type: type[PeerType] | None = None,  # pylint: disable=unused-argument
+        branch_agnostic: bool = False,
     ) -> Mapping[str, Node | PeerType]:
         rels = await self.get_relationships(db=db, branch_agnostic=branch_agnostic)
         peer_ids = [rel.peer_id for rel in rels if rel.peer_id]

backend/infrahub/core/schema/definitions/core.py

@@ -17,7 +17,6 @@
     InfrahubKind,
     PermissionAction,
     PermissionDecision,
-    ProposedChangeState,
     RelationshipDeleteBehavior,
     RepositoryInternalStatus,
     RepositoryOperationalStatus,
@@ -26,6 +25,7 @@
     ValidatorConclusion,
     ValidatorState,
 )
+from infrahub.proposed_change.constants import ProposedChangeState
 
 # pylint: disable=too-many-lines
 

backend/infrahub/database/__init__.py

@@ -349,6 +349,8 @@ async def execute_query_with_metadata(
 
             with QUERY_EXECUTION_METRICS.labels(**labels).time():
                 response = await self.run_query(query=query, params=params, name=name)
+                if response is None:
+                    return [], {}
                 results = [item async for item in response]
                 return results, response._metadata or {}
 

backend/infrahub/graphql/mutations/branch.py

@@ -3,7 +3,7 @@
 from typing import TYPE_CHECKING, Any
 
 import pydantic
-from graphene import Boolean, Field, InputField, InputObjectType, Mutation, ObjectType, String
+from graphene import Boolean, Field, InputField, InputObjectType, Mutation, String
 from infrahub_sdk.utils import extract_fields, extract_fields_first_node
 from opentelemetry import trace
 from typing_extensions import Self
@@ -25,6 +25,7 @@
 )
 
 from ..types import BranchType
+from ..types.task import TaskInfo
 
 if TYPE_CHECKING:
     from graphql import GraphQLResolveInfo
@@ -37,10 +38,6 @@
 log = get_logger()
 
 
-class TaskInfo(ObjectType):
-    id = Field(String)
-
-
 class BranchCreateInput(InputObjectType):
     id = String(required=False)
     name = String(required=True)