infrahub-server 1.3.2__py3-none-any.whl → 1.3.4__py3-none-any.whl

infrahub/core/query/node.py

@@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Any, AsyncIterator, Generator
 from infrahub import config
 from infrahub.core import registry
 from infrahub.core.constants import (
+    GLOBAL_BRANCH_NAME,
     AttributeDBNodeType,
     RelationshipDirection,
     RelationshipHierarchyDirection,
@@ -128,7 +129,7 @@ class NodeCreateAllQuery(NodeQuery):
 
     raise_error_if_empty: bool = True
 
-    async def query_init(self, db: InfrahubDatabase, **kwargs) -> None:  # noqa: ARG002
+    async def query_init(self, db: InfrahubDatabase, **kwargs) -> None:  # noqa: ARG002, PLR0915
         at = self.at or self.node._at
         self.params["uuid"] = self.node.id
         self.params["branch"] = self.branch.name
@@ -151,11 +152,19 @@ class NodeCreateAllQuery(NodeQuery):
             else:
                 attributes.append(attr_data)
 
+        deepest_branch_name = self.branch.name
+        deepest_branch_level = self.branch.hierarchy_level
         relationships: list[RelationshipCreateData] = []
         for rel_name in self.node._relationships:
             rel_manager: RelationshipManager = getattr(self.node, rel_name)
             for rel in rel_manager._relationships:
-                relationships.append(await rel.get_create_data(db=db))
+                rel_create_data = await rel.get_create_data(db=db, at=at)
+                if rel_create_data.peer_branch_level > deepest_branch_level or (
+                    deepest_branch_name == GLOBAL_BRANCH_NAME and rel_create_data.peer_branch == registry.default_branch
+                ):
+                    deepest_branch_name = rel_create_data.peer_branch
+                    deepest_branch_level = rel_create_data.peer_branch_level
+                relationships.append(rel_create_data)
 
         self.params["attrs"] = [attr.model_dump() for attr in attributes]
         self.params["attrs_iphost"] = [attr.model_dump() for attr in attributes_iphost]
@@ -216,6 +225,8 @@ class NodeCreateAllQuery(NodeQuery):
             CREATE (a)-[:HAS_VALUE { branch: attr.branch, branch_level: attr.branch_level, status: attr.status, from: $at }]->(av)
             MERGE (ip:Boolean { value: attr.is_protected })
             MERGE (iv:Boolean { value: attr.is_visible })
+            WITH a, ip, iv
+            LIMIT 1
             CREATE (a)-[:IS_PROTECTED { branch: attr.branch, branch_level: attr.branch_level, status: attr.status, from: $at }]->(ip)
             CREATE (a)-[:IS_VISIBLE { branch: attr.branch, branch_level: attr.branch_level, status: attr.status, from: $at }]->(iv)
             FOREACH ( prop IN attr.source_prop |
@@ -230,9 +241,8 @@ class NodeCreateAllQuery(NodeQuery):
 
         attrs_iphost_query = """
         WITH distinct n
-        UNWIND $attrs_iphost AS attr_iphost
-        CALL (n, attr_iphost) {
-            WITH n, attr_iphost AS attr
+        UNWIND $attrs_iphost AS attr
+        CALL (n, attr) {
             CREATE (a:Attribute { uuid: attr.uuid, name: attr.name, branch_support: attr.branch_support })
             CREATE (n)-[:HAS_ATTRIBUTE { branch: attr.branch, branch_level: attr.branch_level, status: attr.status, from: $at }]->(a)
             MERGE (av:AttributeValue:AttributeIPHost { %(iphost_prop)s })
@@ -241,6 +251,8 @@ class NodeCreateAllQuery(NodeQuery):
             CREATE (a)-[:HAS_VALUE { branch: attr.branch, branch_level: attr.branch_level, status: attr.status, from: $at }]->(av)
             MERGE (ip:Boolean { value: attr.is_protected })
             MERGE (iv:Boolean { value: attr.is_visible })
+            WITH a, ip, iv
+            LIMIT 1
             CREATE (a)-[:IS_PROTECTED { branch: attr.branch, branch_level: attr.branch_level, status: attr.status, from: $at }]->(ip)
             CREATE (a)-[:IS_VISIBLE { branch: attr.branch, branch_level: attr.branch_level, status: attr.status, from: $at }]->(iv)
             FOREACH ( prop IN attr.source_prop |
@@ -256,9 +268,8 @@ class NodeCreateAllQuery(NodeQuery):
 
         attrs_ipnetwork_query = """
         WITH distinct n
-        UNWIND $attrs_ipnetwork AS attr_ipnetwork
-        CALL (n, attr_ipnetwork) {
-            WITH n, attr_ipnetwork AS attr
+        UNWIND $attrs_ipnetwork AS attr
+        CALL (n, attr) {
             CREATE (a:Attribute { uuid: attr.uuid, name: attr.name, branch_support: attr.branch_support })
             CREATE (n)-[:HAS_ATTRIBUTE { branch: attr.branch, branch_level: attr.branch_level, status: attr.status, from: $at }]->(a)
             MERGE (av:AttributeValue:AttributeIPNetwork { %(ipnetwork_prop)s })
@@ -267,6 +278,8 @@ class NodeCreateAllQuery(NodeQuery):
             CREATE (a)-[:HAS_VALUE { branch: attr.branch, branch_level: attr.branch_level, status: attr.status, from: $at }]->(av)
             MERGE (ip:Boolean { value: attr.is_protected })
             MERGE (iv:Boolean { value: attr.is_visible })
+            WITH a, ip, iv
+            LIMIT 1
             CREATE (a)-[:IS_PROTECTED { branch: attr.branch, branch_level: attr.branch_level, status: attr.status, from: $at }]->(ip)
             CREATE (a)-[:IS_VISIBLE { branch: attr.branch, branch_level: attr.branch_level, status: attr.status, from: $at }]->(iv)
             FOREACH ( prop IN attr.source_prop |
@@ -280,16 +293,55 @@ class NodeCreateAllQuery(NodeQuery):
         }
         """ % {"ipnetwork_prop": ", ".join(ipnetwork_prop_list)}
 
+        deepest_branch = await registry.get_branch(db=db, branch=deepest_branch_name)
+        branch_filter, branch_params = deepest_branch.get_query_filter_path(at=self.at)
+        self.params.update(branch_params)
+        self.params["global_branch_name"] = GLOBAL_BRANCH_NAME
+        self.params["default_branch_name"] = registry.default_branch
+
+        dest_node_subquery = """
+        CALL (rel) {
+            MATCH (dest_node:Node { uuid: rel.destination_id })-[r:IS_PART_OF]->(root:Root)
+            WHERE (
+                // if the relationship is on a branch, use the regular filter
+                (rel.peer_branch_level = 2 AND %(branch_filter)s)
+                // simplified filter for the global branch
+                OR (
+                    rel.peer_branch_level = 1
+                    AND rel.peer_branch = $global_branch_name
+                    AND r.branch = $global_branch_name
+                    AND r.from <= $at AND (r.to IS NULL or r.to > $at)
+                )
+                // simplified filter for the default branch
+                OR (
+                    rel.peer_branch_level = 1 AND
+                    rel.peer_branch = $default_branch_name AND
+                    r.branch IN [$default_branch_name, $global_branch_name]
+                    AND r.from <= $at AND (r.to IS NULL or r.to > $at)
+                )
+            )
+            // r.status is a tie-breaker when there are nodes with the same UUID added/deleted at the same time
+            ORDER BY r.branch_level DESC, r.from DESC, r.status ASC
+            WITH dest_node, r
+            LIMIT 1
+            WITH dest_node, r
+            WHERE r.status = "active"
+            RETURN dest_node
+        }
+        """ % {"branch_filter": branch_filter}
+
         rels_bidir_query = """
         WITH distinct n
         UNWIND $rels_bidir AS rel
-        CALL (n, rel) {
-            MERGE (d:Node { uuid: rel.destination_id })
+        %(dest_node_subquery)s
+        CALL (n, rel, dest_node) {
             CREATE (rl:Relationship { uuid: rel.uuid, name: rel.name, branch_support: rel.branch_support })
             CREATE (n)-[:IS_RELATED %(rel_prop)s ]->(rl)
-            CREATE (d)-[:IS_RELATED %(rel_prop)s ]->(rl)
+            CREATE (dest_node)-[:IS_RELATED %(rel_prop)s ]->(rl)
             MERGE (ip:Boolean { value: rel.is_protected })
             MERGE (iv:Boolean { value: rel.is_visible })
+            WITH rl, ip, iv
+            LIMIT 1
             CREATE (rl)-[:IS_PROTECTED { branch: rel.branch, branch_level: rel.branch_level, status: rel.status, from: $at }]->(ip)
             CREATE (rl)-[:IS_VISIBLE { branch: rel.branch, branch_level: rel.branch_level, status: rel.status, from: $at }]->(iv)
             FOREACH ( prop IN rel.source_prop |
@@ -301,19 +353,20 @@ class NodeCreateAllQuery(NodeQuery):
                 CREATE (rl)-[:HAS_OWNER { branch: rel.branch, branch_level: rel.branch_level, status: rel.status, from: $at }]->(peer)
             )
         }
-        """ % {"rel_prop": rel_prop_str}
+        """ % {"rel_prop": rel_prop_str, "dest_node_subquery": dest_node_subquery}
 
         rels_out_query = """
         WITH distinct n
-        UNWIND $rels_out AS rel_out
-        CALL (n, rel_out) {
-            WITH n, rel_out as rel
-            MERGE (d:Node { uuid: rel.destination_id })
+        UNWIND $rels_out AS rel
+        %(dest_node_subquery)s
+        CALL (n, rel, dest_node) {
             CREATE (rl:Relationship { uuid: rel.uuid, name: rel.name, branch_support: rel.branch_support })
             CREATE (n)-[:IS_RELATED %(rel_prop)s ]->(rl)
-            CREATE (d)<-[:IS_RELATED %(rel_prop)s ]-(rl)
+            CREATE (dest_node)<-[:IS_RELATED %(rel_prop)s ]-(rl)
             MERGE (ip:Boolean { value: rel.is_protected })
             MERGE (iv:Boolean { value: rel.is_visible })
+            WITH rl, ip, iv
+            LIMIT 1
             CREATE (rl)-[:IS_PROTECTED { branch: rel.branch, branch_level: rel.branch_level, status: rel.status, from: $at }]->(ip)
             CREATE (rl)-[:IS_VISIBLE { branch: rel.branch, branch_level: rel.branch_level, status: rel.status, from: $at }]->(iv)
             FOREACH ( prop IN rel.source_prop |
@@ -325,19 +378,20 @@ class NodeCreateAllQuery(NodeQuery):
                 CREATE (rl)-[:HAS_OWNER { branch: rel.branch, branch_level: rel.branch_level, status: rel.status, from: $at }]->(peer)
             )
         }
-        """ % {"rel_prop": rel_prop_str}
+        """ % {"rel_prop": rel_prop_str, "dest_node_subquery": dest_node_subquery}
 
         rels_in_query = """
         WITH distinct n
-        UNWIND $rels_in AS rel_in
-        CALL (n, rel_in) {
-            WITH n, rel_in AS rel
-            MERGE (d:Node { uuid: rel.destination_id })
+        UNWIND $rels_in AS rel
+        %(dest_node_subquery)s
+        CALL (n, rel, dest_node) {
             CREATE (rl:Relationship { uuid: rel.uuid, name: rel.name, branch_support: rel.branch_support })
             CREATE (n)<-[:IS_RELATED %(rel_prop)s ]-(rl)
-            CREATE (d)-[:IS_RELATED %(rel_prop)s ]->(rl)
+            CREATE (dest_node)-[:IS_RELATED %(rel_prop)s ]->(rl)
             MERGE (ip:Boolean { value: rel.is_protected })
             MERGE (iv:Boolean { value: rel.is_visible })
+            WITH rl, ip, iv
+            LIMIT 1
             CREATE (rl)-[:IS_PROTECTED { branch: rel.branch, branch_level: rel.branch_level, status: rel.status, from: $at }]->(ip)
             CREATE (rl)-[:IS_VISIBLE { branch: rel.branch, branch_level: rel.branch_level, status: rel.status, from: $at }]->(iv)
             FOREACH ( prop IN rel.source_prop |
@@ -349,7 +403,7 @@ class NodeCreateAllQuery(NodeQuery):
                 CREATE (rl)-[:HAS_OWNER { branch: rel.branch, branch_level: rel.branch_level, status: rel.status, from: $at }]->(peer)
             )
         }
-        """ % {"rel_prop": rel_prop_str}
+        """ % {"rel_prop": rel_prop_str, "dest_node_subquery": dest_node_subquery}
 
         query = f"""
         MATCH (root:Root)

infrahub/core/query/relationship.py

@@ -206,16 +206,18 @@ class RelationshipQuery(Query):
         self.params["source_id"] = self.source_id or self.source.get_id()
         if source_branch.is_global or source_branch.is_default:
             source_query_match = """
-            MATCH (s:Node { uuid: $source_id })
+            MATCH (s:Node { uuid: $source_id })-[source_e:IS_PART_OF {branch: $source_branch, status: "active"}]->(:Root)
+            WHERE source_e.from <= $at AND (source_e.to IS NULL OR source_e.to > $at)
             OPTIONAL MATCH (s)-[delete_edge:IS_PART_OF {status: "deleted", branch: $source_branch}]->(:Root)
             WHERE delete_edge.from <= $at
             WITH *, s WHERE delete_edge IS NULL
             """
             self.params["source_branch"] = source_branch.name
-        source_filter, source_filter_params = source_branch.get_query_filter_path(
-            at=self.at, variable_name="r", params_prefix="src_"
-        )
-        source_query_match = """
+        else:
+            source_filter, source_filter_params = source_branch.get_query_filter_path(
+                at=self.at, variable_name="r", params_prefix="src_"
+            )
+            source_query_match = """
             MATCH (s:Node { uuid: $source_id })
             CALL (s) {
                 MATCH (s)-[r:IS_PART_OF]->(:Root)
@@ -225,15 +227,16 @@ class RelationshipQuery(Query):
                 LIMIT 1
             }
             WITH *, s WHERE s_is_active = TRUE
-            """ % {"source_filter": source_filter}
-        self.params.update(source_filter_params)
+                """ % {"source_filter": source_filter}
+            self.params.update(source_filter_params)
         self.add_to_query(source_query_match)
 
     def add_dest_match_to_query(self, destination_branch: Branch, destination_id: str) -> None:
         self.params["destination_id"] = destination_id
         if destination_branch.is_global or destination_branch.is_default:
             destination_query_match = """
-            MATCH (d:Node { uuid: $destination_id })
+            MATCH (d:Node { uuid: $destination_id })-[dest_e:IS_PART_OF {branch: $destination_branch, status: "active"}]->(:Root)
+            WHERE dest_e.from <= $at AND (dest_e.to IS NULL OR dest_e.to > $at)
             OPTIONAL MATCH (d)-[delete_edge:IS_PART_OF {status: "deleted", branch: $destination_branch}]->(:Root)
             WHERE delete_edge.from <= $at
             WITH *, d WHERE delete_edge IS NULL

infrahub/core/query/resource_manager.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Generator
+
+from pydantic import BaseModel, ConfigDict
 
 from infrahub.core import registry
 from infrahub.core.constants import InfrahubKind, RelationshipStatus
@@ -11,6 +13,13 @@ if TYPE_CHECKING:
     from infrahub.database import InfrahubDatabase
 
 
+class NumberPoolIdentifierData(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    value: int
+    identifier: str
+
+
 class IPAddressPoolGetIdentifiers(Query):
     name = "ipaddresspool_get_identifiers"
     type = QueryType.READ
@@ -158,7 +167,7 @@ class NumberPoolGetReserved(Query):
     def __init__(
         self,
         pool_id: str,
-        identifier: str,
+        identifier: str | None = None,
         **kwargs: dict[str, Any],
     ) -> None:
         self.pool_id = pool_id
@@ -176,26 +185,104 @@ class NumberPoolGetReserved(Query):
 
         self.params.update(branch_params)
 
+        # If identifier is not provided, we return all reservations for the pool
+        identifier_filter = ""
+        if self.identifier:
+            identifier_filter = "r.identifier = $identifier AND "
+            self.params["identifier"] = self.identifier
+
         query = """
         MATCH (pool:%(number_pool)s { uuid: $pool_id })-[r:IS_RESERVED]->(reservation:AttributeValue)
         WHERE
-            r.identifier = $identifier
-            AND
+            %(identifier_filter)s
             %(branch_filter)s
-        """ % {"branch_filter": branch_filter, "number_pool": InfrahubKind.NUMBERPOOL}
+        """ % {
+            "branch_filter": branch_filter,
+            "number_pool": InfrahubKind.NUMBERPOOL,
+            "identifier_filter": identifier_filter,
+        }
         self.add_to_query(query)
-        self.return_labels = ["reservation.value"]
+        self.return_labels = ["reservation.value AS value", "r.identifier AS identifier"]
 
     def get_reservation(self) -> int | None:
         result = self.get_result()
         if result:
-            return result.get_as_optional_type("reservation.value", return_type=int)
+            return result.get_as_optional_type("value", return_type=int)
         return None
 
+    def get_reservations(self) -> Generator[NumberPoolIdentifierData]:
+        for result in self.results:
+            yield NumberPoolIdentifierData.model_construct(
+                value=result.get_as_type("value", return_type=int),
+                identifier=result.get_as_type("identifier", return_type=str),
+            )
+
+
+class PoolChangeReserved(Query):
+    """Change the identifier on all pools.
+    This is useful when a node is being converted to a different type and its ID has changed
+    """
+
+    name = "pool_change_reserved"
+    type = QueryType.WRITE
+
+    def __init__(
+        self,
+        existing_identifier: str,
+        new_identifier: str,
+        **kwargs: dict[str, Any],
+    ) -> None:
+        self.existing_identifier = existing_identifier
+        self.new_identifier = new_identifier
+
+        super().__init__(**kwargs)  # type: ignore[arg-type]
+
+    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:  # noqa: ARG002
+        self.params["new_identifier"] = self.new_identifier
+        self.params["existing_identifier"] = self.existing_identifier
+        self.params["at"] = self.at.to_string()
+
+        branch_filter, branch_params = self.branch.get_query_filter_path(
+            at=self.at.to_string(), branch_agnostic=self.branch_agnostic
+        )
+
+        self.params.update(branch_params)
+
+        global_branch = registry.get_global_branch()
+        self.params["rel_prop"] = {
+            "branch": global_branch.name,
+            "branch_level": global_branch.hierarchy_level,
+            "status": RelationshipStatus.ACTIVE.value,
+            "from": self.at.to_string(),
+            "identifier": self.new_identifier,
+        }
+
+        query = """
+        MATCH (pool:Node)-[r:IS_RESERVED]->(resource)
+        WHERE
+            r.identifier = $existing_identifier
+            AND
+            %(branch_filter)s
+        SET r.to = $at
+        CREATE (pool)-[new_rel:IS_RESERVED $rel_prop]->(resource)
+        """ % {"branch_filter": branch_filter}
+        self.add_to_query(query)
+        self.return_labels = ["pool.uuid AS pool_id", "r", "new_rel"]
+
+
+"""
+Important!: The relationship IS_RESERVED for Number is not being cleaned up when the node or the branch is deleted
+I think this is something we should address in the future.
+It works for now because the query has been updated to match the identifier in IS_RESERVED with the UUID of the related node
+But in the future, if we need to use an identifier that is not the UUID, we will need to clean up the relationships
+This will be especially important as we want to support upsert with NumberPool
+"""
+
 
 class NumberPoolGetUsed(Query):
     name = "number_pool_get_used"
     type = QueryType.READ
+    return_model = NumberPoolIdentifierData
 
     def __init__(
         self,
@@ -219,26 +306,36 @@ class NumberPoolGetUsed(Query):
         self.params["attribute_name"] = self.pool.node_attribute.value
 
         query = """
-        MATCH (pool:%(number_pool)s { uuid: $pool_id })
-        CALL (pool) {
-            MATCH (pool)-[res:IS_RESERVED]->(av:AttributeValue)<-[hv:HAS_VALUE]-(attr:Attribute)
+        MATCH (pool:%(number_pool)s { uuid: $pool_id })-[res:IS_RESERVED]->(av:AttributeValue)
+        WHERE toInteger(av.value) >= $start_range and toInteger(av.value) <= $end_range
+        CALL (pool, res, av) {
+            MATCH (pool)-[res]->(av)<-[hv:HAS_VALUE]-(attr:Attribute)<-[ha:HAS_ATTRIBUTE]-(n:%(node)s)
             WHERE
-                attr.name = $attribute_name
-                AND
-                toInteger(av.value) >= $start_range and toInteger(av.value) <= $end_range
-                AND
-                all(r in [res, hv] WHERE (%(branch_filter)s))
-            RETURN av, (res.status = "active" AND hv.status = "active") AS is_active
+                n.uuid = res.identifier AND
+                attr.name = $attribute_name AND
+                all(r in [res, hv, ha] WHERE (%(branch_filter)s))
+            ORDER BY res.branch_level DESC, hv.branch_level DESC, ha.branch_level DESC, res.from DESC, hv.from DESC, ha.from DESC
+            RETURN (res.status = "active" AND hv.status = "active" AND ha.status = "active") AS is_active
+            LIMIT 1
         }
-        WITH av, is_active
-        WHERE is_active = TRUE
+        WITH av, res, is_active
+        WHERE is_active = True
         """ % {
             "branch_filter": branch_filter,
             "number_pool": InfrahubKind.NUMBERPOOL,
+            "node": self.pool.node.value,
         }
+
         self.add_to_query(query)
-        self.return_labels = ["av.value"]
-        self.order_by = ["av.value"]
+        self.return_labels = ["DISTINCT(av.value) as value", "res.identifier as identifier"]
+        self.order_by = ["value"]
+
+    def iter_results(self) -> Generator[NumberPoolIdentifierData]:
+        for result in self.results:
+            yield self.return_model.model_construct(
+                value=result.get_as_type("value", return_type=int),
+                identifier=result.get_as_type("identifier", return_type=str),
+            )
 
 
 class NumberPoolSetReserved(Query):

infrahub/core/relationship/model.py

@@ -63,9 +63,11 @@ class RelationshipCreateData(BaseModel):
     uuid: str
     name: str
     destination_id: str
-    branch: str | None = None
+    branch: str
     branch_level: int
     branch_support: str | None = None
+    peer_branch: str
+    peer_branch_level: int
     direction: str
     status: str
     is_protected: bool
@@ -420,7 +422,7 @@ class Relationship(FlagPropertyMixin, NodePropertyMixin):
         )
         await delete_query.execute(db=db)
 
-    async def resolve(self, db: InfrahubDatabase) -> None:
+    async def resolve(self, db: InfrahubDatabase, at: Timestamp | None = None) -> None:
         """Resolve the peer of the relationship."""
 
         if self._peer is not None:
@@ -470,7 +472,7 @@ class Relationship(FlagPropertyMixin, NodePropertyMixin):
                 if hfid_str:
                     data_from_pool["identifier"] = f"hfid={hfid_str} rel={self.name}"
 
-            assigned_peer: Node = await pool.get_resource(db=db, branch=self.branch, **data_from_pool)  # type: ignore[attr-defined]
+            assigned_peer: Node = await pool.get_resource(db=db, branch=self.branch, at=at, **data_from_pool)  # type: ignore[attr-defined]
             await self.set_peer(value=assigned_peer)
             self.set_source(value=pool.id)
 
@@ -526,16 +528,19 @@ class Relationship(FlagPropertyMixin, NodePropertyMixin):
 
         return response
 
-    async def get_create_data(self, db: InfrahubDatabase) -> RelationshipCreateData:
+    async def get_create_data(self, db: InfrahubDatabase, at: Timestamp | None = None) -> RelationshipCreateData:
         branch = self.get_branch_based_on_support_type()
 
-        await self.resolve(db=db)
+        await self.resolve(db=db, at=at)
 
         peer = await self.get_peer(db=db)
+        peer_branch = peer.get_branch()
         data = RelationshipCreateData(
             uuid=str(UUIDT()),
             name=self.schema.get_identifier(),
             branch=branch.name,
+            peer_branch=peer_branch.name,
+            peer_branch_level=peer_branch.hierarchy_level,
             destination_id=peer.id,
             status="active",
             direction=self.schema.direction.value,

infrahub/core/schema/__init__.py

@@ -46,6 +46,7 @@ class SchemaExtension(HashableModel):
 
 class SchemaRoot(BaseModel):
     model_config = ConfigDict(extra="forbid")
+
     version: str | None = Field(default=None)
     generics: list[GenericSchema] = Field(default_factory=list)
     nodes: list[NodeSchema] = Field(default_factory=list)
@@ -93,6 +94,10 @@ class SchemaRoot(BaseModel):
         """Return a new `SchemaRoot` after merging `self` with `schema`."""
         return SchemaRoot.model_validate(deep_merge_dict(dicta=self.model_dump(), dictb=schema.model_dump()))
 
+    def duplicate(self) -> SchemaRoot:
+        """Return a duplicate of the current schema."""
+        return SchemaRoot.model_validate(self.model_dump())
+
 
 internal_schema = internal.to_dict()
 

infrahub/core/schema/attribute_parameters.py

@@ -165,3 +165,9 @@ class NumberPoolParameters(AttributeParameters):
         if self.start_range > self.end_range:
             raise ValueError("`start_range` can't be less than `end_range`")
         return self
+
+    def get_pool_size(self) -> int:
+        """
+        Returns the size of the pool based on the defined ranges.
+        """
+        return self.end_range - self.start_range + 1

infrahub/core/schema/attribute_schema.py

@@ -10,6 +10,7 @@ from infrahub import config
 from infrahub.core.constants.schema import UpdateSupport
 from infrahub.core.enums import generate_python_enum
 from infrahub.core.query.attribute import default_attribute_query_filter
+from infrahub.exceptions import InitializationError
 from infrahub.types import ATTRIBUTE_KIND_LABELS, ATTRIBUTE_TYPES
 
 from .attribute_parameters import (
@@ -67,6 +68,11 @@ class AttributeSchema(GeneratedAttributeSchema):
     def is_deprecated(self) -> bool:
         return bool(self.deprecation)
 
+    def get_id(self) -> str:
+        if self.id is None:
+            raise InitializationError("The attribute schema has not been saved yet and doesn't have an id")
+        return self.id
+
     def to_dict(self) -> dict:
         data = self.model_dump(exclude_unset=True, exclude_none=True)
         for field_name, value in data.items():

infrahub/core/schema/manager.py

@@ -535,7 +535,7 @@ class SchemaManager(NodeManager):
         """Delete the node with its attributes and relationships."""
         branch = await registry.get_branch(branch=branch, db=db)
 
-        obj = await self.get_one(id=node.get_id(), branch=branch, db=db)
+        obj = await self.get_one(id=node.get_id(), branch=branch, db=db, prefetch_relationships=True)
         if not obj:
             raise SchemaNotFoundError(
                 branch_name=branch.name,
@@ -544,16 +544,10 @@ class SchemaManager(NodeManager):
             )
 
         # First delete the attributes and the relationships
-        items = await self.get_many(
-            ids=[item.id for item in node.local_attributes + node.local_relationships if item.id],
-            db=db,
-            branch=branch,
-            include_owner=True,
-            include_source=True,
-        )
-
-        for item in items.values():
-            await item.delete(db=db)
+        for attr_schema_node in (await obj.attributes.get_peers(db=db)).values():
+            await attr_schema_node.delete(db=db)
+        for rel_schema_node in (await obj.relationships.get_peers(db=db)).values():
+            await rel_schema_node.delete(db=db)
 
         await obj.delete(db=db)
 

infrahub/core/schema/relationship_schema.py

@@ -9,6 +9,7 @@ from infrahub import config
 from infrahub.core.constants import RelationshipDirection
 from infrahub.core.query import QueryNode, QueryRel, QueryRelDirection
 from infrahub.core.relationship import Relationship
+from infrahub.exceptions import InitializationError
 
 from .generated.relationship_schema import GeneratedRelationshipSchema
 
@@ -57,6 +58,11 @@ class RelationshipSchema(GeneratedRelationshipSchema):
             raise ValueError("RelationshipSchema is not initialized")
         return self.identifier
 
+    def get_id(self) -> str:
+        if not self.id:
+            raise InitializationError("The relationship schema has not been saved yet and doesn't have an id")
+        return self.id
+
     def get_query_arrows(self) -> QueryArrows:
         """Return (in 4 parts) the 2 arrows for the relationship R1 and R2 based on the direction of the relationship."""