infrahub-server 1.4.0b1__py3-none-any.whl → 1.4.1__py3-none-any.whl

infrahub/core/migrations/graph/{m035_drop_attr_value_index.py → m036_drop_attr_value_index.py}

@@ -18,10 +18,10 @@ if TYPE_CHECKING:
 INDEX_TO_DELETE = IndexItem(name="attr_value", label="AttributeValue", properties=["value"], type=IndexType.RANGE)
 
 
-class Migration035(GraphMigration):
-    name: str = "035_drop_attr_value_index"
+class Migration036(GraphMigration):
+    name: str = "036_drop_attr_value_index"
     queries: Sequence[type[Query]] = []
-    minimum_version: int = 34
+    minimum_version: int = 35
 
     async def execute(self, db: InfrahubDatabase) -> MigrationResult:
         result = MigrationResult()

infrahub/core/migrations/graph/{m036_index_attr_vals.py → m037_index_attr_vals.py}

@@ -441,7 +441,7 @@ REMOVE av_no_index:AttributeValueNonIndexed
         self.add_to_query(query)
 
 
-class Migration036(ArbitraryMigration):
+class Migration037(ArbitraryMigration):
     """
     Update AttributeValue vertices to be AttributeValueIndexed, unless they include values for LARGE_ATTRIBUTE_TYPES
 
@@ -458,8 +458,8 @@ class Migration036(ArbitraryMigration):
     7. Add the index on AttributeValueIndexed again
     """
 
-    name: str = "036_index_attr_vals"
-    minimum_version: int = 35
+    name: str = "037_index_attr_vals"
+    minimum_version: int = 36
 
     async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
         result = MigrationResult()

infrahub/core/migrations/query/node_duplicate.py

@@ -21,6 +21,13 @@ class SchemaNodeInfo(BaseModel):
 
 
 class NodeDuplicateQuery(Query):
+    """
+    Duplicates a Node to use a new kind or inheritance.
+    Creates a copy of each affected Node and sets the new kind/inheritance.
+    Adds duplicate edges to the new Node that match all the active edges on the old Node.
+    Sets all the edges on the old Node to deleted.
+    """
+
     name = "node_duplicate"
     type = QueryType.WRITE
     insert_return: bool = False
@@ -38,11 +45,26 @@ class NodeDuplicateQuery(Query):
 
     def render_match(self) -> str:
         labels_str = ":".join(self.previous_node.labels)
-        query = f"""
+        query = """
         // Find all the active nodes
-        MATCH (node:{labels_str})
+        MATCH (node:%(labels_str)s)
         WITH DISTINCT node
-        """
+        // ----------------
+        // Filter out nodes that have already been migrated
+        // ----------------
+        CALL (node) {
+            WITH labels(node) AS node_labels
+            UNWIND node_labels AS n_label
+            ORDER BY n_label ASC
+            WITH collect(n_label) AS sorted_labels
+
+            RETURN (
+                node.kind = $new_node.kind AND
+                sorted_labels = $new_sorted_labels
+            ) AS already_migrated
+        }
+        WITH node WHERE already_migrated = FALSE
+        """ % {"labels_str": labels_str}
 
         return query
 
@@ -111,6 +133,7 @@ class NodeDuplicateQuery(Query):
 
         self.params["new_node"] = self.new_node.model_dump()
         self.params["previous_node"] = self.previous_node.model_dump()
+        self.params["new_sorted_labels"] = sorted(self.new_node.labels + ["Node"])
 
         self.params["current_time"] = self.at.to_string()
         self.params["branch"] = self.branch.name

infrahub/core/migrations/schema/attribute_kind_update.py

@@ -0,0 +1,156 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Sequence
+
+from infrahub.types import is_large_attribute_type
+
+from ..query import AttributeMigrationQuery
+from ..shared import AttributeSchemaMigration, MigrationResult
+
+if TYPE_CHECKING:
+    from infrahub.core.branch.models import Branch
+    from infrahub.core.timestamp import Timestamp
+    from infrahub.database import InfrahubDatabase
+
+
+class AttributeKindUpdateMigrationQuery(AttributeMigrationQuery):
+    name = "migration_attribute_kind"
+    insert_return = False
+
+    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)
+        needs_index = not is_large_attribute_type(self.migration.new_attribute_schema.kind)
+        self.params["needs_index"] = needs_index
+        self.params["branch"] = self.branch.name
+        self.params["branch_level"] = self.branch.hierarchy_level
+        self.params["at"] = self.at.to_string()
+        self.params["attr_name"] = self.migration.previous_attribute_schema.name
+        new_attr_value_labels = "AttributeValue"
+        if needs_index:
+            new_attr_value_labels += ":AttributeValueIndexed"
+        # ruff: noqa: S608
+        query = """
+// ------------
+// start with all the Attribute vertices we might care about
+// ------------
+MATCH (n:%(schema_kind)s)-[:HAS_ATTRIBUTE]->(attr:Attribute)
+WHERE attr.name = $attr_name
+WITH DISTINCT n, attr
+
+// ------------
+// for each Attribute, find the most recent active edge and AttributeValue vertex that needs to be [un]indexed
+// ------------
+CALL (n, attr) {
+    MATCH (n)-[r1:HAS_ATTRIBUTE]->(attr:Attribute)-[r2:HAS_VALUE]->(av)
+    WHERE all(r IN [r1, r2] WHERE %(branch_filter)s)
+    WITH r2, av, r1.status = "active" AND r2.status = "active" AS is_active
+    ORDER BY r2.branch_level DESC, r2.from DESC, r2.status = "active" DESC, r1.branch_level DESC, r1.from DESC, r1.status = "active" DESC
+    LIMIT 1
+    WITH r2 AS has_value_e, av, "AttributeValueIndexed" IN labels(av) AS is_indexed
+    WHERE is_active AND is_indexed <> $needs_index
+    RETURN has_value_e, av
+}
+
+// ------------
+// check if the correct AttributeValue vertex to use exists
+// create it if not
+// ------------
+WITH DISTINCT av.is_default AS av_is_default, av.value AS av_value
+CALL (av_is_default, av_value) {
+    OPTIONAL MATCH (existing_av:AttributeValue {is_default: av_is_default, value: av_value})
+    WHERE "AttributeValueIndexed" IN labels(existing_av) = $needs_index
+    WITH existing_av WHERE existing_av IS NULL
+    LIMIT 1
+    CREATE (:%(new_attr_value_labels)s {is_default: av_is_default, value: av_value})
+}
+
+// ------------
+// get all the AttributeValue vertices that need to be updated again and run the updates
+// ------------
+WITH 1 AS one
+LIMIT 1
+MATCH (n:%(schema_kind)s)-[:HAS_ATTRIBUTE]->(attr:Attribute)
+WHERE attr.name = $attr_name
+WITH DISTINCT n, attr
+
+// ------------
+// for each Attribute, find the most recent active edge and AttributeValue vertex that needs to be [un]indexed
+// ------------
+CALL (n, attr) {
+    MATCH (n)-[r1:HAS_ATTRIBUTE]->(attr:Attribute)-[r2:HAS_VALUE]->(av)
+    WHERE all(r IN [r1, r2] WHERE %(branch_filter)s)
+    WITH r2, av, r1.status = "active" AND r2.status = "active" AS is_active
+    ORDER BY r2.branch_level DESC, r2.from DESC, r2.status = "active" DESC, r1.branch_level DESC, r1.from DESC, r1.status = "active" DESC
+    LIMIT 1
+    WITH r2 AS has_value_e, av, "AttributeValueIndexed" IN labels(av) AS is_indexed
+    WHERE is_active AND is_indexed <> $needs_index
+    RETURN has_value_e, av
+}
+
+
+// ------------
+// create and update the HAS_VALUE edges
+// ------------
+CALL (attr, has_value_e, av) {
+    // ------------
+    // get the correct AttributeValue vertex b/c it definitely exists now
+    // ------------
+    MATCH (new_av:%(new_attr_value_labels)s {is_default: av.is_default, value: av.value})
+    WHERE "AttributeValueIndexed" IN labels(new_av) = $needs_index
+    LIMIT 1
+
+    // ------------
+    // create the new HAS_VALUE edge
+    // ------------
+    CREATE (attr)-[new_has_value_e:HAS_VALUE]->(new_av)
+    SET new_has_value_e = properties(has_value_e)
+    SET new_has_value_e.status = "active"
+    SET new_has_value_e.branch = $branch
+    SET new_has_value_e.branch_level = $branch_level
+    SET new_has_value_e.from = $at
+    SET new_has_value_e.to = NULL
+
+    // ------------
+    // if we are updating on a branch and the existing edge is on the default branch,
+    // then create a new deleted edge on this branch
+    // ------------
+    WITH attr, has_value_e, av
+    WHERE has_value_e.branch <> $branch
+    CREATE (attr)-[deleted_has_value_e:HAS_VALUE]->(av)
+    SET deleted_has_value_e = properties(has_value_e)
+    SET deleted_has_value_e.status = "deleted"
+    SET deleted_has_value_e.branch = $branch
+    SET deleted_has_value_e.branch_level = $branch_level
+    SET deleted_has_value_e.from = $at
+    SET deleted_has_value_e.to = NULL
+}
+
+// ------------
+// if the existing edge is on the same branch as the update,
+// then set its "to" time
+// ------------
+CALL (has_value_e) {
+    WITH has_value_e
+    WHERE has_value_e.branch = $branch
+    SET has_value_e.to = $at
+}
+        """ % {
+            "schema_kind": self.migration.previous_schema.kind,
+            "branch_filter": branch_filter,
+            "new_attr_value_labels": new_attr_value_labels,
+        }
+        self.add_to_query(query)
+
+
+class AttributeKindUpdateMigration(AttributeSchemaMigration):
+    name: str = "attribute.kind.update"
+    queries: Sequence[type[AttributeMigrationQuery]] = [AttributeKindUpdateMigrationQuery]  # type: ignore[assignment]
+
+    async def execute(self, db: InfrahubDatabase, branch: Branch, at: Timestamp | str | None = None) -> MigrationResult:
+        is_indexed_previous = is_large_attribute_type(self.previous_attribute_schema.kind)
+        is_indexed_new = is_large_attribute_type(self.new_attribute_schema.kind)
+        if is_indexed_previous is is_indexed_new:
+            return MigrationResult()
+
+        return await super().execute(db=db, branch=branch, at=at)

infrahub/core/models.py

@@ -569,7 +569,11 @@ class HashableModel(BaseModel):
 
         for field_name in other.model_fields.keys():
             if not hasattr(self, field_name):
-                setattr(self, field_name, getattr(other, field_name))
+                try:
+                    setattr(self, field_name, getattr(other, field_name))
+                except ValueError:
+                    # handles the case where self and other are different types and other has fields that self does not
+                    pass
                 continue
 
             attr_other = getattr(other, field_name)

infrahub/core/node/resource_manager/ip_address_pool.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import ipaddress
 from typing import TYPE_CHECKING, Any
 
+from infrahub import lock
 from infrahub.core import registry
 from infrahub.core.ipam.reconciler import IpamReconciler
 from infrahub.core.query.ipam import get_ip_addresses
@@ -33,54 +34,55 @@ class CoreIPAddressPool(Node):
         prefixlen: int | None = None,
         at: Timestamp | None = None,
     ) -> Node:
-        # Check if there is already a resource allocated with this identifier
-        # if not, pull all existing prefixes and allocated the next available
-
-        if identifier:
-            query_get = await IPAddressPoolGetReserved.init(db=db, pool_id=self.id, identifier=identifier)
-            await query_get.execute(db=db)
-            result = query_get.get_result()
-
-            if result:
-                address = result.get_node("address")
-                # TODO add support for branch, if the node is reserved with this id in another branch we should return an error
-                node = await registry.manager.get_one(db=db, id=address.get("uuid"), branch=branch)
-
-                if node:
-                    return node
-
-        data = data or {}
-
-        address_type = address_type or data.get("address_type") or self.default_address_type.value  # type: ignore[attr-defined]
-        if not address_type:
-            raise ValueError(
-                f"IPAddressPool: {self.name.value} | "  # type: ignore[attr-defined]
-                "An address_type or a default_value type must be provided to allocate a new IP address"
-            )
-
-        ip_namespace = await self.ip_namespace.get_peer(db=db)  # type: ignore[attr-defined]
-
-        prefixlen = prefixlen or data.get("prefixlen") or self.default_prefix_length.value  # type: ignore[attr-defined]
-
-        next_address = await self.get_next(db=db, prefixlen=prefixlen)
-
-        target_schema = registry.get_node_schema(name=address_type, branch=branch)
-        node = await Node.init(db=db, schema=target_schema, branch=branch, at=at)
-        try:
-            await node.new(db=db, address=str(next_address), ip_namespace=ip_namespace, **data)
-        except ValidationError as exc:
-            raise ValueError(f"IPAddressPool: {self.name.value} | {exc!s}") from exc  # type: ignore[attr-defined]
-        await node.save(db=db, at=at)
-        reconciler = IpamReconciler(db=db, branch=branch)
-        await reconciler.reconcile(ip_value=next_address, namespace=ip_namespace.id, node_uuid=node.get_id())
-
-        if identifier:
-            query_set = await IPAddressPoolSetReserved.init(
-                db=db, pool_id=self.id, identifier=identifier, address_id=node.id, at=at
-            )
-            await query_set.execute(db=db)
-
-        return node
+        async with lock.registry.get(name=self.get_id(), namespace="resource_pool"):
+            # Check if there is already a resource allocated with this identifier
+            # if not, pull all existing prefixes and allocated the next available
+
+            if identifier:
+                query_get = await IPAddressPoolGetReserved.init(db=db, pool_id=self.id, identifier=identifier)
+                await query_get.execute(db=db)
+                result = query_get.get_result()
+
+                if result:
+                    address = result.get_node("address")
+                    # TODO add support for branch, if the node is reserved with this id in another branch we should return an error
+                    node = await registry.manager.get_one(db=db, id=address.get("uuid"), branch=branch)
+
+                    if node:
+                        return node
+
+            data = data or {}
+
+            address_type = address_type or data.get("address_type") or self.default_address_type.value  # type: ignore[attr-defined]
+            if not address_type:
+                raise ValueError(
+                    f"IPAddressPool: {self.name.value} | "  # type: ignore[attr-defined]
+                    "An address_type or a default_value type must be provided to allocate a new IP address"
+                )
+
+            ip_namespace = await self.ip_namespace.get_peer(db=db)  # type: ignore[attr-defined]
+
+            prefixlen = prefixlen or data.get("prefixlen") or self.default_prefix_length.value  # type: ignore[attr-defined]
+
+            next_address = await self.get_next(db=db, prefixlen=prefixlen)
+
+            target_schema = registry.get_node_schema(name=address_type, branch=branch)
+            node = await Node.init(db=db, schema=target_schema, branch=branch, at=at)
+            try:
+                await node.new(db=db, address=str(next_address), ip_namespace=ip_namespace, **data)
+            except ValidationError as exc:
+                raise ValueError(f"IPAddressPool: {self.name.value} | {exc!s}") from exc  # type: ignore[attr-defined]
+            await node.save(db=db, at=at)
+            reconciler = IpamReconciler(db=db, branch=branch)
+            await reconciler.reconcile(ip_value=next_address, namespace=ip_namespace.id, node_uuid=node.get_id())
+
+            if identifier:
+                query_set = await IPAddressPoolSetReserved.init(
+                    db=db, pool_id=self.id, identifier=identifier, address_id=node.id, at=at
+                )
+                await query_set.execute(db=db)
+
+            return node
 
     async def get_next(self, db: InfrahubDatabase, prefixlen: int | None = None) -> IPAddressType:
         resources = await self.resources.get_peers(db=db)  # type: ignore[attr-defined]

infrahub/core/node/resource_manager/ip_prefix_pool.py

@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any
 
 from netaddr import IPSet
 
+from infrahub import lock
 from infrahub.core import registry
 from infrahub.core.ipam.reconciler import IpamReconciler
 from infrahub.core.query.ipam import get_subnets
@@ -36,59 +37,60 @@ class CoreIPPrefixPool(Node):
         prefix_type: str | None = None,
         at: Timestamp | None = None,
     ) -> Node:
-        # Check if there is already a resource allocated with this identifier
-        # if not, pull all existing prefixes and allocated the next available
-        if identifier:
-            query_get = await PrefixPoolGetReserved.init(db=db, pool_id=self.id, identifier=identifier)
-            await query_get.execute(db=db)
-            result = query_get.get_result()
-            if result:
-                prefix = result.get_node("prefix")
-                # TODO add support for branch, if the node is reserved with this id in another branch we should return an error
-                node = await registry.manager.get_one(db=db, id=prefix.get("uuid"), branch=branch)
-                if node:
-                    return node
-
-        ip_namespace = await self.ip_namespace.get_peer(db=db)  # type: ignore[attr-defined]
-
-        data = data or {}
-
-        prefixlen = prefixlen or data.get("prefixlen", None) or self.default_prefix_length.value  # type: ignore[attr-defined]
-        if not prefixlen:
-            raise ValueError(
-                f"IPPrefixPool: {self.name.value} | "  # type: ignore[attr-defined]
-                "A prefixlen or a default_value must be provided to allocate a new prefix"
-            )
-
-        next_prefix = await self.get_next(db=db, prefixlen=prefixlen)
-
-        prefix_type = prefix_type or data.get("prefix_type", None) or self.default_prefix_type.value  # type: ignore[attr-defined]
-        if not prefix_type:
-            raise ValueError(
-                f"IPPrefixPool: {self.name.value} | "  # type: ignore[attr-defined]
-                "A prefix_type or a default_value type must be provided to allocate a new prefix"
-            )
-
-        member_type = member_type or data.get("member_type", None) or self.default_member_type.value.value  # type: ignore[attr-defined]
-        data["member_type"] = member_type
-
-        target_schema = registry.get_node_schema(name=prefix_type, branch=branch)
-        node = await Node.init(db=db, schema=target_schema, branch=branch, at=at)
-        try:
-            await node.new(db=db, prefix=str(next_prefix), ip_namespace=ip_namespace, **data)
-        except ValidationError as exc:
-            raise ValueError(f"IPPrefixPool: {self.name.value} | {exc!s}") from exc  # type: ignore[attr-defined]
-        await node.save(db=db, at=at)
-        reconciler = IpamReconciler(db=db, branch=branch)
-        await reconciler.reconcile(ip_value=next_prefix, namespace=ip_namespace.id, node_uuid=node.get_id())
-
-        if identifier:
-            query_set = await PrefixPoolSetReserved.init(
-                db=db, pool_id=self.id, identifier=identifier, prefix_id=node.id, at=at
-            )
-            await query_set.execute(db=db)
-
-        return node
+        async with lock.registry.get(name=self.get_id(), namespace="resource_pool"):
+            # Check if there is already a resource allocated with this identifier
+            # if not, pull all existing prefixes and allocated the next available
+            if identifier:
+                query_get = await PrefixPoolGetReserved.init(db=db, pool_id=self.id, identifier=identifier)
+                await query_get.execute(db=db)
+                result = query_get.get_result()
+                if result:
+                    prefix = result.get_node("prefix")
+                    # TODO add support for branch, if the node is reserved with this id in another branch we should return an error
+                    node = await registry.manager.get_one(db=db, id=prefix.get("uuid"), branch=branch)
+                    if node:
+                        return node
+
+            ip_namespace = await self.ip_namespace.get_peer(db=db)  # type: ignore[attr-defined]
+
+            data = data or {}
+
+            prefixlen = prefixlen or data.get("prefixlen", None) or self.default_prefix_length.value  # type: ignore[attr-defined]
+            if not prefixlen:
+                raise ValueError(
+                    f"IPPrefixPool: {self.name.value} | "  # type: ignore[attr-defined]
+                    "A prefixlen or a default_value must be provided to allocate a new prefix"
+                )
+
+            next_prefix = await self.get_next(db=db, prefixlen=prefixlen)
+
+            prefix_type = prefix_type or data.get("prefix_type", None) or self.default_prefix_type.value  # type: ignore[attr-defined]
+            if not prefix_type:
+                raise ValueError(
+                    f"IPPrefixPool: {self.name.value} | "  # type: ignore[attr-defined]
+                    "A prefix_type or a default_value type must be provided to allocate a new prefix"
+                )
+
+            member_type = member_type or data.get("member_type", None) or self.default_member_type.value.value  # type: ignore[attr-defined]
+            data["member_type"] = member_type
+
+            target_schema = registry.get_node_schema(name=prefix_type, branch=branch)
+            node = await Node.init(db=db, schema=target_schema, branch=branch, at=at)
+            try:
+                await node.new(db=db, prefix=str(next_prefix), ip_namespace=ip_namespace, **data)
+            except ValidationError as exc:
+                raise ValueError(f"IPPrefixPool: {self.name.value} | {exc!s}") from exc  # type: ignore[attr-defined]
+            await node.save(db=db, at=at)
+            reconciler = IpamReconciler(db=db, branch=branch)
+            await reconciler.reconcile(ip_value=next_prefix, namespace=ip_namespace.id, node_uuid=node.get_id())
+
+            if identifier:
+                query_set = await PrefixPoolSetReserved.init(
+                    db=db, pool_id=self.id, identifier=identifier, prefix_id=node.id, at=at
+                )
+                await query_set.execute(db=db)
+
+            return node
 
     async def get_next(self, db: InfrahubDatabase, prefixlen: int) -> IPNetworkType:
         resources = await self.resources.get_peers(db=db)  # type: ignore[attr-defined]

infrahub/core/node/resource_manager/number_pool.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
+from infrahub import lock
 from infrahub.core import registry
 from infrahub.core.query.resource_manager import NumberPoolGetReserved, NumberPoolGetUsed, NumberPoolSetReserved
 from infrahub.core.schema.attribute_parameters import NumberAttributeParameters
@@ -62,24 +63,25 @@ class CoreNumberPool(Node):
         identifier: str | None = None,
         at: Timestamp | None = None,
     ) -> int:
-        # NOTE: ideally we should use the HFID as the identifier (if available)
-        # one of the challenge with using the HFID is that it might change over time
-        # so we need to ensure that the identifier is stable, or we need to handle the case where the identifier changes
-        identifier = identifier or node.get_id()
-
-        # Check if there is already a resource allocated with this identifier
-        # if not, pull all existing number and allocate the next available
-        # TODO add support for branch, if the node is reserved with this id in another branch we should return an error
-        query_get = await NumberPoolGetReserved.init(db=db, branch=branch, pool_id=self.id, identifier=identifier)
-        await query_get.execute(db=db)
-        reservation = query_get.get_reservation()
-        if reservation is not None:
-            return reservation
-
-        # If we have not returned a value we need to find one if avaiable
-        number = await self.get_next(db=db, branch=branch, attribute=attribute)
-        await self.reserve(db=db, number=number, identifier=identifier, at=at)
-        return number
+        async with lock.registry.get(name=self.get_id(), namespace="resource_pool"):
+            # NOTE: ideally we should use the HFID as the identifier (if available)
+            # one of the challenge with using the HFID is that it might change over time
+            # so we need to ensure that the identifier is stable, or we need to handle the case where the identifier changes
+            identifier = identifier or node.get_id()
+
+            # Check if there is already a resource allocated with this identifier
+            # if not, pull all existing number and allocate the next available
+            # TODO add support for branch, if the node is reserved with this id in another branch we should return an error
+            query_get = await NumberPoolGetReserved.init(db=db, branch=branch, pool_id=self.id, identifier=identifier)
+            await query_get.execute(db=db)
+            reservation = query_get.get_reservation()
+            if reservation is not None:
+                return reservation
+
+            # If we have not returned a value we need to find one if avaiable
+            number = await self.get_next(db=db, branch=branch, attribute=attribute)
+            await self.reserve(db=db, number=number, identifier=identifier, at=at)
+            return number
 
     async def get_next(self, db: InfrahubDatabase, branch: Branch, attribute: AttributeSchema) -> int:
         taken = await self.get_used(db=db, branch=branch)

infrahub/core/query/branch.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Any
 
 from infrahub import config
+from infrahub.core.constants import GLOBAL_BRANCH_NAME
 from infrahub.core.query import Query, QueryType
 
 if TYPE_CHECKING:
@@ -21,33 +22,49 @@ class DeleteBranchRelationshipsQuery(Query):
 
     async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:  # noqa: ARG002
         query = """
-// delete all relationships on this branch
-MATCH (s)-[r1]-(d)
-WHERE r1.branch = $branch_name
-CALL (r1) {
-  DELETE r1
+// --------------
+// for every Node created on this branch (it's about to be deleted), find any agnostic relationships
+// connected to the Node and delete them
+// --------------
+OPTIONAL MATCH (:Root)<-[e:IS_PART_OF {status: "active"}]-(n:Node)
+WHERE e.branch = $branch_name
+CALL (n) {
+    OPTIONAL MATCH (n)-[:IS_RELATED {branch: $global_branch_name}]-(rel:Relationship)
+    DETACH DELETE rel
 } IN TRANSACTIONS
 
-// check for any orphaned Node vertices and delete them
-WITH collect(DISTINCT s.uuid) + collect(DISTINCT d.uuid) AS nodes_uuids
-MATCH (s2:Node)-[r2]-(d2)
-WHERE NOT exists((s2)-[:IS_PART_OF]-(:Root))
-AND s2.uuid IN nodes_uuids
-CALL (r2) {
-  DELETE r2
+// reduce the results to a single row
+WITH 1 AS one
+LIMIT 1
+
+// --------------
+// for every edge on this branch, delete it
+// --------------
+MATCH (s)-[r]->(d)
+WHERE r.branch = $branch_name
+CALL (r) {
+    DELETE r
 } IN TRANSACTIONS
 
-// reduce results to a single row
-WITH 1 AS one LIMIT 1
-
-// find any orphaned vertices and delete them
-MATCH (n)
-WHERE NOT exists((n)--())
-CALL (n) {
-  DELETE n
+// --------------
+// get the database IDs of every vertex linked to a deleted edge
+// --------------
+WITH DISTINCT elementId(s) AS s_id, elementId(d) AS d_id
+WITH collect(s_id) + collect(d_id) AS vertex_ids
+UNWIND vertex_ids AS vertex_id
+
+// --------------
+// delete any vertices that are now orphaned
+// --------------
+CALL (vertex_id) {
+    MATCH (n)
+    WHERE elementId(n) = vertex_id
+    AND NOT exists((n)--())
+    DELETE n
 } IN TRANSACTIONS
         """
         self.params["branch_name"] = self.branch_name
+        self.params["global_branch_name"] = GLOBAL_BRANCH_NAME
         self.add_to_query(query)