infrahub-server 1.1.6__py3-none-any.whl → 1.2.0b1__py3-none-any.whl

infrahub/core/validators/checks_runner.py

@@ -0,0 +1,37 @@
+import asyncio
+from typing import Any, Coroutine
+
+from infrahub_sdk.node import InfrahubNode
+
+from infrahub.core.constants import ValidatorConclusion, ValidatorState
+from infrahub.core.timestamp import Timestamp
+
+
+async def run_checks_and_update_validator(
+    checks: list[Coroutine[Any, None, ValidatorConclusion]], validator: InfrahubNode
+) -> None:
+    """
+    Execute a list of checks coroutines, and set validator fields accordingly.
+    Tasks are retrieved by completion order so as soon as we detect a failing check,
+    we set validator conclusion to failure.
+    """
+
+    # First set validator to in progress, then wait for results
+    validator.state.value = ValidatorState.IN_PROGRESS.value
+    validator.started_at.value = Timestamp().to_string()
+    validator.completed_at.value = ""
+    await validator.save()
+
+    for earliest_task in asyncio.as_completed(checks):
+        result = await earliest_task
+        if validator.conclusion.value != ValidatorConclusion.FAILURE.value and result == ValidatorConclusion.FAILURE:
+            validator.conclusion.value = ValidatorConclusion.FAILURE.value
+            await validator.save()
+            # Continue to iterate to wait for the end of all checks
+
+    validator.state.value = ValidatorState.COMPLETED.value
+    validator.completed_at.value = Timestamp().to_string()
+    if validator.conclusion.value != ValidatorConclusion.FAILURE.value:
+        validator.conclusion.value = ValidatorConclusion.SUCCESS.value
+
+    await validator.save()

infrahub/core/validators/node/generate_profile.py

@@ -28,7 +28,7 @@ class NodeGenerateProfileValidatorQuery(SchemaValidatorQuery):
         super().__init__(*args, **kwargs)
         self.profile_kind = f"Profile{self.node_schema.kind}"
 
-    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:
+    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:  # noqa: ARG002
         branch_filter, branch_params = self.branch.get_query_filter_path(at=self.at)
         self.params.update(branch_params)
 

infrahub/core/validators/node/hierarchy.py

@@ -31,7 +31,7 @@ class NodeHierarchyUpdateValidatorQuery(SchemaValidatorQuery):
         self.check_parent = check_parent
         super().__init__(**kwargs)
 
-    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:
+    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:  # noqa: ARG002
         if self.check_children and self.check_parent:
             raise RuntimeError("Cannot check children and parent at same time")
         if self.check_children:

infrahub/core/validators/query.py

@@ -14,7 +14,7 @@ if TYPE_CHECKING:
 class NodeNotPresentValidatorQuery(SchemaValidatorQuery):
     name: str = "node_not_present_validator"
 
-    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:
+    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:  # noqa: ARG002
         branch_filter, branch_params = self.branch.get_query_filter_path(at=self.at.to_string())
         self.params.update(branch_params)
 

infrahub/core/validators/relationship/count.py

@@ -30,7 +30,7 @@ class RelationshipCountUpdateValidatorQuery(RelationshipSchemaValidatorQuery):
         self.max_count_override = max_count_override
         super().__init__(**kwargs)
 
-    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:
+    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:  # noqa: ARG002
         branch_filter, branch_params = self.branch.get_query_filter_path(at=self.at.to_string(), is_isolated=False)
         self.params.update(branch_params)
 

infrahub/core/validators/relationship/optional.py

@@ -20,7 +20,7 @@ if TYPE_CHECKING:
 class RelationshipOptionalUpdateValidatorQuery(RelationshipSchemaValidatorQuery):
     name = "relationship_constraints_optional_validator"
 
-    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:
+    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:  # noqa: ARG002
         branch_filter, branch_params = self.branch.get_query_filter_path(at=self.at.to_string(), is_isolated=False)
         self.params.update(branch_params)
 

infrahub/core/validators/relationship/peer.py

@@ -21,7 +21,7 @@ if TYPE_CHECKING:
 class RelationshipPeerUpdateValidatorQuery(RelationshipSchemaValidatorQuery):
     name = "relationship_constraints_peer_validator"
 
-    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:
+    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:  # noqa: ARG002
         peer_schema = db.schema.get(name=self.relationship_schema.peer, branch=self.branch)
         allowed_peer_kinds = [peer_schema.kind]
         if isinstance(peer_schema, GenericSchema):

infrahub/core/validators/tasks.py

@@ -13,14 +13,16 @@ from infrahub.core.validators.model import (
     SchemaConstraintValidatorRequest,
 )
 from infrahub.dependencies.registry import get_component_registry
-from infrahub.services import services
+from infrahub.services import InfrahubServices  # noqa: TC001  needed for prefect flow
 from infrahub.workflows.utils import add_tags
 
 from .models.validate_migration import SchemaValidateMigrationData, SchemaValidatorPathResponseData
 
 
 @flow(name="schema_validate_migrations", flow_run_name="Validate schema migrations", persist_result=True)
-async def schema_validate_migrations(message: SchemaValidateMigrationData) -> list[SchemaValidatorPathResponseData]:
+async def schema_validate_migrations(
+    message: SchemaValidateMigrationData, service: InfrahubServices
+) -> list[SchemaValidatorPathResponseData]:
     batch = InfrahubBatch(return_exceptions=True)
     log = get_run_logger()
     await add_tags(branches=[message.branch.name])
@@ -33,7 +35,7 @@ async def schema_validate_migrations(message: SchemaValidateMigrationData) -> li
     # NOTE this task is a good candidate to add a progress bar
     for constraint in message.constraints:
         schema = message.schema_branch.get(name=constraint.path.schema_kind)
-        if not isinstance(schema, (GenericSchema, NodeSchema)):
+        if not isinstance(schema, GenericSchema | NodeSchema):
             continue
         batch.add(
             task=schema_path_validate,
@@ -41,13 +43,14 @@ async def schema_validate_migrations(message: SchemaValidateMigrationData) -> li
             constraint_name=constraint.constraint_name,
             node_schema=schema,
             schema_path=constraint.path,
+            service=service,
         )
 
     results = [result async for _, result in batch.execute()]
     return results
 
 
-@task(
+@task(  # type: ignore[arg-type]
     name="schema-path-validate",
     task_run_name="Validate schema path {constraint_name} in {branch.name}",
     description="Validate if a given migration is compatible with the existing data",
@@ -59,9 +62,8 @@ async def schema_path_validate(
     constraint_name: str,
     node_schema: NodeSchema | GenericSchema,
     schema_path: SchemaPath,
+    service: InfrahubServices,
 ) -> SchemaValidatorPathResponseData:
-    service = services.service
-
     async with service.database.start_session() as db:
         constraint_request = SchemaConstraintValidatorRequest(
             branch=branch,

infrahub/core/validators/uniqueness/query.py

@@ -30,7 +30,7 @@ class NodeUniqueAttributeConstraintQuery(Query):
     def get_context(self) -> dict[str, str]:
         return {"kind": self.query_request.kind}
 
-    async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:
+    async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:  # noqa: ARG002
         branch_filter, branch_params = self.branch.get_query_filter_path(at=self.at.to_string(), is_isolated=False)
         self.params.update(branch_params)
         from_times = db.render_list_comprehension(items="relationships(potential_path)", item_name="from")
@@ -56,6 +56,7 @@ class NodeUniqueAttributeConstraintQuery(Query):
         relationship_names = set()
         relationship_attr_paths = []
         relationship_only_attr_paths = []
+        relationship_only_attr_values = []
         relationship_attr_paths_with_value = []
         for rel_path in self.query_request.relationship_attribute_paths:
             relationship_names.add(rel_path.identifier)
@@ -67,6 +68,8 @@ class NodeUniqueAttributeConstraintQuery(Query):
                 relationship_attr_paths.append((rel_path.identifier, rel_path.attribute_name))
             else:
                 relationship_only_attr_paths.append(rel_path.identifier)
+                if rel_path.value:
+                    relationship_only_attr_values.append(rel_path.value)
 
         if (
             not attr_paths
@@ -89,34 +92,37 @@ class NodeUniqueAttributeConstraintQuery(Query):
                 "relationship_attr_paths": relationship_attr_paths,
                 "relationship_attr_paths_with_value": relationship_attr_paths_with_value,
                 "relationship_only_attr_paths": relationship_only_attr_paths,
+                "relationship_only_attr_values": relationship_only_attr_values,
                 "min_count_required": self.min_count_required,
             }
         )
 
         attr_paths_subquery = """
-        WITH start_node
-        MATCH attr_path = (start_node)-[:HAS_ATTRIBUTE]->(attr:Attribute)-[r:HAS_VALUE]->(attr_value:AttributeValue)
+        MATCH attr_path = (start_node:%(node_kind)s)-[:HAS_ATTRIBUTE]->(attr:Attribute)-[r:HAS_VALUE]->(attr_value:AttributeValue)
         WHERE attr.name in $attribute_names
             AND ([attr.name, type(r)] in $attr_paths
             OR [attr.name, type(r), attr_value.value] in $attr_paths_with_value)
-        RETURN attr_path as potential_path, NULL as rel_identifier, attr.name as potential_attr, attr_value.value as potential_attr_value
-        """
+        RETURN start_node, attr_path as potential_path, NULL as rel_identifier, attr.name as potential_attr, attr_value.value as potential_attr_value
+        """ % {"node_kind": self.query_request.kind}
 
         relationship_attr_paths_with_value_subquery = """
-        WITH start_node
-        MATCH rel_path = (start_node)-[:IS_RELATED]-(relationship_node:Relationship)-[:IS_RELATED]-(related_n:Node)-[:HAS_ATTRIBUTE]->(rel_attr:Attribute)-[:HAS_VALUE]->(rel_attr_value:AttributeValue)
+        MATCH rel_path = (start_node:%(node_kind)s)-[:IS_RELATED]-(relationship_node:Relationship)-[:IS_RELATED]-(related_n:Node)-[:HAS_ATTRIBUTE]->(rel_attr:Attribute)-[:HAS_VALUE]->(rel_attr_value:AttributeValue)
         WHERE relationship_node.name in $relationship_names
             AND ([relationship_node.name, rel_attr.name] in $relationship_attr_paths
             OR [relationship_node.name, rel_attr.name, rel_attr_value.value] in $relationship_attr_paths_with_value)
-        RETURN rel_path as potential_path, relationship_node.name as rel_identifier, rel_attr.name as potential_attr, rel_attr_value.value as potential_attr_value
-        """
+        RETURN start_node, rel_path as potential_path, relationship_node.name as rel_identifier, rel_attr.name as potential_attr, rel_attr_value.value as potential_attr_value
+        """ % {"node_kind": self.query_request.kind}
 
         relationship_only_attr_paths_subquery = """
-        WITH start_node
-        MATCH rel_path = (start_node)-[:IS_RELATED]-(relationship_node:Relationship)-[:IS_RELATED]-(related_n:Node)
-        WHERE relationship_node.name in $relationship_only_attr_paths
-        RETURN rel_path as potential_path, relationship_node.name as rel_identifier, "id" as potential_attr, related_n.uuid as potential_attr_value
-        """
+        MATCH rel_path = (start_node:%(node_kind)s)-[:IS_RELATED]-(relationship_node:Relationship)-[:IS_RELATED]-(related_n:Node)
+        WHERE %(rel_node_filter)s relationship_node.name in $relationship_only_attr_paths
+        RETURN start_node, rel_path as potential_path, relationship_node.name as rel_identifier, "id" as potential_attr, related_n.uuid as potential_attr_value
+        """ % {
+            "node_kind": self.query_request.kind,
+            "rel_node_filter": "related_n.uuid IN $relationship_only_attr_values AND "
+            if relationship_only_attr_values
+            else "",
+        }
 
         select_subqueries = []
         if attr_paths or attr_paths_with_value:
@@ -130,8 +136,6 @@ class NodeUniqueAttributeConstraintQuery(Query):
 
         # ruff: noqa: E501
         query = """
-        // group by node
-        MATCH (start_node:%(node_kind)s)
         // get attributes for node and its relationships
         CALL {
             %(select_subqueries_str)s
@@ -201,7 +205,6 @@ class NodeUniqueAttributeConstraintQuery(Query):
             attr_value,
             relationship_identifier
         """ % {
-            "node_kind": self.query_request.kind,
             "select_subqueries_str": select_subqueries_str,
             "branch_filter": branch_filter,
             "from_times": from_times,

infrahub/database/__init__.py

@@ -173,6 +173,19 @@ class InfrahubDatabase:
         elif self.db_type == DatabaseType.MEMGRAPH:
             self.manager = DatabaseManagerMemgraph(db=self)
 
+    def __del__(self) -> None:
+        if not self._session or not self._is_session_local or self._session.closed():
+            return
+
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            loop = None
+        if loop and loop.is_running():
+            loop.create_task(self._session.close())
+        else:
+            asyncio.run(self._session.close())
+
     @property
     def is_session(self) -> bool:
         if self._mode == InfrahubDatabaseMode.SESSION:
@@ -307,7 +320,7 @@ class InfrahubDatabase:
         params: dict[str, Any] | None = None,
         name: str = "undefined",
         context: dict[str, str] | None = None,
-        type: QueryType | None = None,  # pylint: disable=redefined-builtin
+        type: QueryType | None = None,
     ) -> list[Record]:
         results, _ = await self.execute_query_with_metadata(
             query=query, params=params, name=name, context=context, type=type
@@ -320,7 +333,7 @@ class InfrahubDatabase:
         params: dict[str, Any] | None = None,
         name: str = "undefined",
         context: dict[str, str] | None = None,
-        type: QueryType | None = None,  # pylint: disable=redefined-builtin
+        type: QueryType | None = None,
     ) -> tuple[list[Record], dict[str, Any]]:
         with trace.get_tracer(__name__).start_as_current_span("execute_db_query_with_metadata") as span:
             span.set_attribute("query", query)
@@ -501,6 +514,7 @@ def retry_db_transaction(
                         if exc.code != "Neo.ClientError.Statement.EntityNotFound":
                             raise exc
                     retry_time: float = random.randrange(100, 500) / 1000
+                    log.exception("Retry handler caught database error")
                     log.info(
                         f"Retrying database transaction, attempt {attempt}/{config.SETTINGS.database.retry_limit}",
                         retry_time=retry_time,

infrahub/database/memgraph.py

@@ -21,7 +21,7 @@ class IndexNodeMemgraph(IndexItem):
 
 
 class IndexManagerMemgraph(IndexManagerBase):
-    def init(self, nodes: list[IndexItem], rels: list[IndexItem]) -> None:
+    def init(self, nodes: list[IndexItem], rels: list[IndexItem]) -> None:  # noqa: ARG002
         self.nodes = [IndexNodeMemgraph(**item.model_dump()) for item in nodes]
         self.initialized = True
 

infrahub/dependencies/builder/constraint/grouped/node_runner.py

@@ -2,7 +2,6 @@ from infrahub.core.constraint.node.runner import NodeConstraintRunner
 from infrahub.dependencies.interface import DependencyBuilder, DependencyBuilderContext
 
 from ..node.grouped_uniqueness import NodeGroupedUniquenessConstraintDependency
-from ..node.uniqueness import NodeAttributeUniquenessConstraintDependency
 from ..relationship_manager.count import RelationshipCountConstraintDependency
 from ..relationship_manager.peer_kind import RelationshipPeerKindConstraintDependency
 from ..relationship_manager.profiles_kind import RelationshipProfilesKindConstraintDependency
@@ -15,7 +14,6 @@ class NodeConstraintRunnerDependency(DependencyBuilder[NodeConstraintRunner]):
             db=context.db,
             branch=context.branch,
             node_constraints=[
-                NodeAttributeUniquenessConstraintDependency.build(context=context),
                 NodeGroupedUniquenessConstraintDependency.build(context=context),
             ],
             relationship_manager_constraints=[

infrahub/dependencies/builder/diff/combiner.py

@@ -4,5 +4,5 @@ from infrahub.dependencies.interface import DependencyBuilder, DependencyBuilder
 
 class DiffCombinerDependency(DependencyBuilder[DiffCombiner]):
     @classmethod
-    def build(cls, context: DependencyBuilderContext) -> DiffCombiner:
+    def build(cls, context: DependencyBuilderContext) -> DiffCombiner:  # noqa: ARG003
         return DiffCombiner()

infrahub/dependencies/builder/diff/conflicts_enricher.py

@@ -4,5 +4,5 @@ from infrahub.dependencies.interface import DependencyBuilder, DependencyBuilder
 
 class DiffConflictsEnricherDependency(DependencyBuilder[ConflictsEnricher]):
     @classmethod
-    def build(cls, context: DependencyBuilderContext) -> ConflictsEnricher:
+    def build(cls, context: DependencyBuilderContext) -> ConflictsEnricher:  # noqa: ARG003
         return ConflictsEnricher()

infrahub/dependencies/builder/diff/coordinator.py

@@ -8,7 +8,6 @@ from .conflicts_enricher import DiffConflictsEnricherDependency
 from .data_check_synchronizer import DiffDataCheckSynchronizerDependency
 from .enricher.aggregated import DiffAggregatedEnricherDependency
 from .enricher.labels import DiffLabelsEnricherDependency
-from .enricher.summary_counts import DiffSummaryCountsEnricherDependency
 from .repository import DiffRepositoryDependency
 
 
@@ -22,7 +21,6 @@ class DiffCoordinatorDependency(DependencyBuilder[DiffCoordinator]):
             diff_enricher=DiffAggregatedEnricherDependency.build(context=context),
             conflicts_enricher=DiffConflictsEnricherDependency.build(context=context),
             labels_enricher=DiffLabelsEnricherDependency.build(context=context),
-            summary_counts_enricher=DiffSummaryCountsEnricherDependency.build(context=context),
             data_check_synchronizer=DiffDataCheckSynchronizerDependency.build(context=context),
             conflict_transferer=DiffConflictTransfererDependency.build(context=context),
         )

infrahub/dependencies/builder/diff/deserializer.py

@@ -1,8 +1,10 @@
 from infrahub.core.diff.repository.deserializer import EnrichedDiffDeserializer
 from infrahub.dependencies.interface import DependencyBuilder, DependencyBuilderContext
 
+from .parent_node_adder import DiffParentNodeAdderDependency
+
 
 class DiffDeserializerDependency(DependencyBuilder[EnrichedDiffDeserializer]):
     @classmethod
-    def build(cls, context: DependencyBuilderContext) -> EnrichedDiffDeserializer:
-        return EnrichedDiffDeserializer()
+    def build(cls, context: DependencyBuilderContext) -> EnrichedDiffDeserializer:  # noqa: ARG003
+        return EnrichedDiffDeserializer(parent_adder=DiffParentNodeAdderDependency.build(context=context))

infrahub/dependencies/builder/diff/enricher/hierarchy.py

@@ -1,8 +1,10 @@
 from infrahub.core.diff.enricher.hierarchy import DiffHierarchyEnricher
 from infrahub.dependencies.interface import DependencyBuilder, DependencyBuilderContext
 
+from ..parent_node_adder import DiffParentNodeAdderDependency
+
 
 class DiffHierarchyEnricherDependency(DependencyBuilder[DiffHierarchyEnricher]):
     @classmethod
     def build(cls, context: DependencyBuilderContext) -> DiffHierarchyEnricher:
-        return DiffHierarchyEnricher(db=context.db)
+        return DiffHierarchyEnricher(db=context.db, parent_adder=DiffParentNodeAdderDependency.build(context=context))

infrahub/dependencies/builder/diff/enricher/summary_counts.py

@@ -4,5 +4,5 @@ from infrahub.dependencies.interface import DependencyBuilder, DependencyBuilder
 
 class DiffSummaryCountsEnricherDependency(DependencyBuilder[DiffSummaryCountsEnricher]):
     @classmethod
-    def build(cls, context: DependencyBuilderContext) -> DiffSummaryCountsEnricher:
+    def build(cls, context: DependencyBuilderContext) -> DiffSummaryCountsEnricher:  # noqa: ARG003
         return DiffSummaryCountsEnricher()

infrahub/dependencies/builder/diff/parent_node_adder.py

@@ -0,0 +1,8 @@
+from infrahub.core.diff.parent_node_adder import DiffParentNodeAdder
+from infrahub.dependencies.interface import DependencyBuilder, DependencyBuilderContext
+
+
+class DiffParentNodeAdderDependency(DependencyBuilder[DiffParentNodeAdder]):
+    @classmethod
+    def build(cls, context: DependencyBuilderContext) -> DiffParentNodeAdder:  # noqa: ARG003
+        return DiffParentNodeAdder()

infrahub/events/artifact_action.py

@@ -0,0 +1,76 @@
+from typing import Any
+
+from pydantic import Field, computed_field
+
+from infrahub.message_bus import InfrahubMessage
+
+from .constants import EVENT_NAMESPACE
+from .models import InfrahubEvent
+
+
+class ArtifactEvent(InfrahubEvent):
+    """Event generated when an artifact has been created or updated."""
+
+    node_id: str = Field(..., description="The ID of the artifact")
+    artifact_definition_id: str = Field(..., description="The ID of the artifact definition")
+    target_id: str = Field(..., description="The ID of the target of the artifact")
+    target_kind: str = Field(..., description="The kind of the target of the artifact")
+    checksum: str = Field(..., description="The current checksum of the artifact")
+    checksum_previous: str | None = Field(default=None, description="The previous checksum of the artifact")
+    storage_id: str = Field(..., description="The current storage id of the artifact")
+    storage_id_previous: str | None = Field(default=None, description="The previous storage id of the artifact")
+    created: bool = Field(..., description="Indicates if the artifact was created with this event or already existed")
+
+    def get_related(self) -> list[dict[str, str]]:
+        related = super().get_related()
+        related.append(
+            {
+                "prefect.resource.id": self.target_id,
+                "prefect.resource.role": "infrahub.related.node",
+                "infrahub.node.kind": self.target_kind,
+            }
+        )
+        related.append(
+            {
+                "prefect.resource.id": self.target_id,
+                "prefect.resource.role": "infrahub.artifact",
+                "infrahub.artifact.checksum": self.checksum,
+                "infrahub.artifact.checksum_previous": self.checksum_previous or "",
+                "infrahub.artifact.storage_id": self.storage_id,
+                "infrahub.artifact.storage_id_previous": self.storage_id_previous or "",
+                "infrahub.artifact.artifact_definition_id": self.artifact_definition_id,
+            }
+        )
+
+        return related
+
+    def get_resource(self) -> dict[str, str]:
+        return {
+            "prefect.resource.id": self.node_id,
+            "infrahub.node.kind": "CoreArtifact",
+            "infrahub.node.id": self.node_id,
+            "infrahub.branch.name": self.meta.context.branch.name,
+        }
+
+    def get_payload(self) -> dict[str, Any]:
+        return {
+            "node_id": self.node_id,
+            "artifact_definition_id": self.artifact_definition_id,
+            "target_id": self.target_id,
+            "target_kind": self.target_kind,
+            "checksum": self.checksum,
+            "checksum_previous": self.checksum_previous,
+            "storage_id": self.storage_id,
+            "storage_id_previous": self.storage_id_previous,
+        }
+
+    def get_messages(self) -> list[InfrahubMessage]:
+        return []
+
+    @computed_field
+    def event_name(self) -> str:
+        match self.created:
+            case True:
+                return f"{EVENT_NAMESPACE}.artifact.created"
+            case False:
+                return f"{EVENT_NAMESPACE}.artifact.updated"

infrahub/events/branch_action.py

@@ -1,26 +1,25 @@
-from pydantic import Field
+from pydantic import Field, computed_field
 
 from infrahub.message_bus import InfrahubMessage
 from infrahub.message_bus.messages.refresh_registry_branches import RefreshRegistryBranches
 from infrahub.message_bus.messages.refresh_registry_rebasedbranch import RefreshRegistryRebasedBranch
 
-from .models import InfrahubBranchEvent
+from .constants import EVENT_NAMESPACE
+from .models import InfrahubEvent
 
 
-class BranchDeleteEvent(InfrahubBranchEvent):
+class BranchDeletedEvent(InfrahubEvent):
     """Event generated when a branch has been deleted"""
 
+    branch_name: str = Field(..., description="The name of the branch")
     branch_id: str = Field(..., description="The ID of the mutated node")
     sync_with_git: bool = Field(..., description="Indicates if the branch was extended to Git")
 
-    def get_name(self) -> str:
-        return f"{self.get_event_namespace()}.branch.deleted"
-
     def get_resource(self) -> dict[str, str]:
         return {
-            "prefect.resource.id": f"infrahub.branch.{self.branch}",
+            "prefect.resource.id": f"infrahub.branch.{self.branch_name}",
             "infrahub.branch.id": self.branch_id,
-            "infrahub.branch.name": self.branch,
+            "infrahub.branch.name": self.branch_name,
         }
 
     def get_messages(self) -> list[InfrahubMessage]:
@@ -35,21 +34,23 @@ class BranchDeleteEvent(InfrahubBranchEvent):
         ]
         return events
 
+    @computed_field
+    def event_name(self) -> str:
+        return f"{EVENT_NAMESPACE}.branch.deleted"
+
 
-class BranchCreateEvent(InfrahubBranchEvent):
+class BranchCreatedEvent(InfrahubEvent):
     """Event generated when a branch has been created"""
 
-    branch_id: str = Field(..., description="The ID of the mutated node")
+    branch_name: str = Field(..., description="The name of the branch")
+    branch_id: str = Field(..., description="The ID of the branch")
     sync_with_git: bool = Field(..., description="Indicates if the branch was extended to Git")
 
-    def get_name(self) -> str:
-        return f"{self.get_event_namespace()}.branch.created"
-
     def get_resource(self) -> dict[str, str]:
         return {
-            "prefect.resource.id": f"infrahub.branch.{self.branch}",
+            "prefect.resource.id": f"infrahub.branch.{self.branch_name}",
             "infrahub.branch.id": self.branch_id,
-            "infrahub.branch.name": self.branch,
+            "infrahub.branch.name": self.branch_name,
         }
 
     def get_messages(self) -> list[InfrahubMessage]:
@@ -64,19 +65,43 @@ class BranchCreateEvent(InfrahubBranchEvent):
         ]
         return events
 
+    @computed_field
+    def event_name(self) -> str:
+        return f"{EVENT_NAMESPACE}.branch.created"
+
+
+class BranchMergedEvent(InfrahubEvent):
+    """Event generated when a branch has been merged"""
+
+    branch_name: str = Field(..., description="The name of the branch")
+    branch_id: str = Field(..., description="The ID of the branch")
+
+    def get_resource(self) -> dict[str, str]:
+        return {
+            "prefect.resource.id": f"infrahub.branch.{self.branch_name}",
+            "infrahub.branch.id": self.branch_id,
+            "infrahub.branch.name": self.branch_name,
+        }
+
+    def get_messages(self) -> list[InfrahubMessage]:
+        return []
 
-class BranchRebaseEvent(InfrahubBranchEvent):
+    @computed_field
+    def event_name(self) -> str:
+        return f"{EVENT_NAMESPACE}.branch.merged"
+
+
+class BranchRebasedEvent(InfrahubEvent):
     """Event generated when a branch has been rebased"""
 
     branch_id: str = Field(..., description="The ID of the mutated node")
-
-    def get_name(self) -> str:
-        return f"{self.get_event_namespace()}.branch.rebased"
+    branch_name: str = Field(..., description="The name of the branch")
 
     def get_resource(self) -> dict[str, str]:
         return {
-            "prefect.resource.id": f"infrahub.branch.{self.branch}",
+            "prefect.resource.id": f"infrahub.branch.{self.branch_name}",
             "infrahub.branch.id": self.branch_id,
+            "infrahub.branch.name": self.branch_name,
         }
 
     def get_messages(self) -> list[InfrahubMessage]:
@@ -85,6 +110,10 @@ class BranchRebaseEvent(InfrahubBranchEvent):
             #     branch=self.branch,
             #     meta=self.get_message_meta(),
             # ),
-            RefreshRegistryRebasedBranch(branch=self.branch),
+            RefreshRegistryRebasedBranch(branch=self.branch_name),
         ]
         return events
+
+    @computed_field
+    def event_name(self) -> str:
+        return f"{EVENT_NAMESPACE}.branch.rebased"