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

infrahub/core/node/__init__.py

@@ -26,6 +26,7 @@ from infrahub.core.protocols import CoreNumberPool, CoreObjectTemplate
 from infrahub.core.query.node import NodeCheckIDQuery, NodeCreateAllQuery, NodeDeleteQuery, NodeGetListQuery
 from infrahub.core.schema import (
     AttributeSchema,
+    GenericSchema,
     NodeSchema,
     NonGenericSchemaTypes,
     ProfileSchema,
@@ -82,6 +83,9 @@ class Node(BaseNode, metaclass=BaseNodeMeta):
     def get_schema(self) -> NonGenericSchemaTypes:
         return self._schema
 
+    def get_branch(self) -> Branch:
+        return self._branch
+
     def get_kind(self) -> str:
         """Return the main Kind of the Object."""
         return self._schema.kind
@@ -260,11 +264,9 @@ class Node(BaseNode, metaclass=BaseNodeMeta):
         within the create code.
         """
 
-        number_pool_parameters: NumberPoolParameters | None = None
         if attribute.schema.kind == "NumberPool" and isinstance(attribute.schema.parameters, NumberPoolParameters):
             attribute.from_pool = {"id": attribute.schema.parameters.number_pool_id}
             attribute.is_default = False
-            number_pool_parameters = attribute.schema.parameters
 
         if not attribute.from_pool:
             return
@@ -274,9 +276,9 @@ class Node(BaseNode, metaclass=BaseNodeMeta):
                 db=db, id=attribute.from_pool["id"], kind=CoreNumberPool
             )
         except NodeNotFoundError:
-            if number_pool_parameters:
-                number_pool = await self._fetch_or_create_number_pool(
-                    db=db, attribute=attribute, number_pool_parameters=number_pool_parameters
+            if attribute.schema.kind == "NumberPool" and isinstance(attribute.schema.parameters, NumberPoolParameters):
+                number_pool = await self.fetch_or_create_number_pool(
+                    db=db, schema_node=self._schema, schema_attribute=attribute.schema, branch=self._branch
                 )
 
             else:
@@ -292,7 +294,9 @@ class Node(BaseNode, metaclass=BaseNodeMeta):
             and number_pool.node_attribute.value == attribute.name
         ):
             try:
-                next_free = await number_pool.get_resource(db=db, branch=self._branch, node=self, attribute=attribute)
+                next_free = await number_pool.get_resource(
+                    db=db, branch=self._branch, node=self, attribute=attribute.schema
+                )
             except PoolExhaustedError:
                 errors.append(
                     ValidationError({f"{attribute.name}.from_pool": f"The pool {number_pool.node.value} is exhausted."})
@@ -310,10 +314,28 @@ class Node(BaseNode, metaclass=BaseNodeMeta):
                 )
             )
 
-    async def _fetch_or_create_number_pool(
-        self, db: InfrahubDatabase, attribute: BaseAttribute, number_pool_parameters: NumberPoolParameters
+    @staticmethod
+    async def fetch_or_create_number_pool(
+        db: InfrahubDatabase,
+        schema_node: NodeSchema | GenericSchema,
+        schema_attribute: AttributeSchema,
+        branch: Branch | None = None,
     ) -> CoreNumberPool:
+        """Fetch or create a number pool based on the schema attribute parameters.
+
+        Warning, ideally this method should be outside of the Node class, but it is itself using the Node class to create the pool node.
+        """
+
+        if (
+            schema_attribute.kind != "NumberPool"
+            or not schema_attribute.parameters
+            or not isinstance(schema_attribute.parameters, NumberPoolParameters)
+        ):
+            raise ValueError("Attribute is not of type NumberPool")
+
         number_pool_from_db: CoreNumberPool | None = None
+        number_pool_parameters: NumberPoolParameters = schema_attribute.parameters
+
         lock_definition = NumberPoolLockDefinition(pool_id=str(number_pool_parameters.number_pool_id))
         async with lock.registry.get(
             name=lock_definition.lock_name, namespace=lock_definition.namespace_name, local=False
@@ -322,37 +344,37 @@ class Node(BaseNode, metaclass=BaseNodeMeta):
                 number_pool_from_db = await registry.manager.get_one_by_id_or_default_filter(
                     db=db, id=str(number_pool_parameters.number_pool_id), kind=CoreNumberPool
                 )
+                return number_pool_from_db  # type: ignore[return-value]
+
             except NodeNotFoundError:
                 schema = db.schema.get_node_schema(name="CoreNumberPool", duplicate=False)
 
-                pool_node = self._schema.kind
-                schema_attribute = self._schema.get_attribute(attribute.schema.name)
+                pool_node = schema_node.kind
                 if schema_attribute.inherited:
-                    for generic_name in self._schema.inherit_from:
+                    for generic_name in schema_node.inherit_from:
                         generic_node = db.schema.get_generic_schema(name=generic_name, duplicate=False)
-                        if attribute.schema.name in generic_node.attribute_names:
+                        if schema_attribute.name in generic_node.attribute_names:
                             pool_node = generic_node.kind
                             break
 
-                number_pool = await Node.init(db=db, schema=schema, branch=self._branch)
+                number_pool = await Node.init(db=db, schema=schema, branch=branch)
                 await number_pool.new(
                     db=db,
                     id=number_pool_parameters.number_pool_id,
-                    name=f"{pool_node}.{attribute.schema.name} [{number_pool_parameters.number_pool_id}]",
+                    name=f"{pool_node}.{schema_attribute.name} [{number_pool_parameters.number_pool_id}]",
                     node=pool_node,
-                    node_attribute=attribute.schema.name,
+                    node_attribute=schema_attribute.name,
                     start_range=number_pool_parameters.start_range,
                     end_range=number_pool_parameters.end_range,
                     pool_type=NumberPoolType.SCHEMA.value,
                 )
                 await number_pool.save(db=db)
 
-        # Do a lookup of the number pool to get the correct mapped type from the registry
-        # without this we don't get access to the .get_resource() method.
-        created_pool: CoreNumberPool = number_pool_from_db or await registry.manager.get_one_by_id_or_default_filter(
-            db=db, id=number_pool.id, kind=CoreNumberPool
-        )
-        return created_pool
+                # Do a lookup of the number pool to get the correct mapped type from the registry
+                # without this we don't get access to the .get_resource() method.
+                return await registry.manager.get_one_by_id_or_default_filter(
+                    db=db, id=number_pool.id, kind=CoreNumberPool
+                )
 
     async def handle_object_template(self, fields: dict, db: InfrahubDatabase, errors: list) -> None:
         """Fill the `fields` parameters with values from an object template if one is in use."""
@@ -816,6 +838,7 @@ class Node(BaseNode, metaclass=BaseNodeMeta):
 
         query = await NodeDeleteQuery.init(db=db, node=self, at=delete_at)
         await query.execute(db=db)
+
         self._node_changelog = node_changelog
 
     async def to_graphql(

infrahub/core/node/resource_manager/ip_address_pool.py

@@ -18,6 +18,7 @@ from .. import Node
 if TYPE_CHECKING:
     from infrahub.core.branch import Branch
     from infrahub.core.ipam.constants import IPAddressType
+    from infrahub.core.timestamp import Timestamp
     from infrahub.database import InfrahubDatabase
 
 
@@ -30,6 +31,7 @@ class CoreIPAddressPool(Node):
         data: dict[str, Any] | None = None,
         address_type: str | None = None,
         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
@@ -63,18 +65,18 @@ class CoreIPAddressPool(Node):
         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)
+        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)
+        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
+                db=db, pool_id=self.id, identifier=identifier, address_id=node.id, at=at
             )
             await query_set.execute(db=db)
 

infrahub/core/node/resource_manager/ip_prefix_pool.py

@@ -20,6 +20,7 @@ from .. import Node
 if TYPE_CHECKING:
     from infrahub.core.branch import Branch
     from infrahub.core.ipam.constants import IPNetworkType
+    from infrahub.core.timestamp import Timestamp
     from infrahub.database import InfrahubDatabase
 
 
@@ -33,6 +34,7 @@ class CoreIPPrefixPool(Node):
         prefixlen: int | None = None,
         member_type: str | None = None,
         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
@@ -68,20 +70,21 @@ class CoreIPPrefixPool(Node):
             )
 
         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)
+        node = await Node.init(db=db, schema=target_schema, branch=branch, at=at)
         try:
-            await node.new(db=db, prefix=str(next_prefix), member_type=member_type, ip_namespace=ip_namespace, **data)
+            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)
+        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
+                db=db, pool_id=self.id, identifier=identifier, prefix_id=node.id, at=at
             )
             await query_set.execute(db=db)
 

infrahub/core/node/resource_manager/number_pool.py

@@ -10,16 +10,15 @@ from infrahub.exceptions import PoolExhaustedError
 from .. import Node
 
 if TYPE_CHECKING:
-    from infrahub.core.attribute import BaseAttribute
     from infrahub.core.branch import Branch
+    from infrahub.core.schema import AttributeSchema
+    from infrahub.core.timestamp import Timestamp
     from infrahub.database import InfrahubDatabase
 
 
 class CoreNumberPool(Node):
     def get_attribute_nb_excluded_values(self) -> int:
-        """
-        Returns the number of excluded values for the attribute of the number pool.
-        """
+        """Returns the number of excluded values for the attribute of the number pool."""
 
         pool_node = registry.schema.get(name=self.node.value)  # type: ignore [attr-defined]
         attribute = [attribute for attribute in pool_node.attributes if attribute.name == self.node_attribute.value][0]  # type: ignore [attr-defined]
@@ -34,17 +33,42 @@ class CoreNumberPool(Node):
         res = len(attribute.parameters.get_excluded_single_values()) + sum_excluded_values
         return res
 
+    async def get_used(
+        self,
+        db: InfrahubDatabase,
+        branch: Branch,
+    ) -> list[int]:
+        """Returns a list of used numbers in the pool."""
+
+        query = await NumberPoolGetUsed.init(db=db, branch=branch, pool=self, branch_agnostic=True)
+        await query.execute(db=db)
+        used = [result.value for result in query.iter_results()]
+        return [item for item in used if item is not None]
+
+    async def reserve(self, db: InfrahubDatabase, number: int, identifier: str, at: Timestamp | None = None) -> None:
+        """Reserve a number in the pool for a specific identifier."""
+
+        query = await NumberPoolSetReserved.init(
+            db=db, pool_id=self.get_id(), identifier=identifier, reserved=number, at=at
+        )
+        await query.execute(db=db)
+
     async def get_resource(
         self,
         db: InfrahubDatabase,
         branch: Branch,
         node: Node,
-        attribute: BaseAttribute,
+        attribute: AttributeSchema,
         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 prefixes and allocated the next available
+        # 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)
@@ -54,35 +78,51 @@ class CoreNumberPool(Node):
 
         # 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)
-
-        query_set = await NumberPoolSetReserved.init(
-            db=db, pool_id=self.get_id(), identifier=identifier, reserved=number
-        )
-        await query_set.execute(db=db)
+        await self.reserve(db=db, number=number, identifier=identifier, at=at)
         return number
 
-    async def get_next(self, db: InfrahubDatabase, branch: Branch, attribute: BaseAttribute) -> int:
-        query = await NumberPoolGetUsed.init(db=db, branch=branch, pool=self, branch_agnostic=True)
-        await query.execute(db=db)
-        taken = [result.get_as_optional_type("av.value", return_type=int) for result in query.results]
-        parameters = attribute.schema.parameters
+    async def get_next(self, db: InfrahubDatabase, branch: Branch, attribute: AttributeSchema) -> int:
+        taken = await self.get_used(db=db, branch=branch)
+
         next_number = find_next_free(
             start=self.start_range.value,  # type: ignore[attr-defined]
             end=self.end_range.value,  # type: ignore[attr-defined]
             taken=taken,
-            parameters=parameters if isinstance(parameters, NumberAttributeParameters) else None,
+            parameters=attribute.parameters if isinstance(attribute.parameters, NumberAttributeParameters) else None,
         )
         if next_number is None:
             raise PoolExhaustedError("There are no more values available in this pool.")
 
         return next_number
 
+    async def get_next_many(
+        self, db: InfrahubDatabase, quantity: int, branch: Branch, attribute: AttributeSchema
+    ) -> list[int]:
+        taken = await self.get_used(db=db, branch=branch)
+
+        allocated: list[int] = []
+
+        for _ in range(quantity):
+            next_number = find_next_free(
+                start=self.start_range.value,  # type: ignore[attr-defined]
+                end=self.end_range.value,  # type: ignore[attr-defined]
+                taken=list(set(taken) | set(allocated)),
+                parameters=attribute.parameters
+                if isinstance(attribute.parameters, NumberAttributeParameters)
+                else None,
+            )
+            if next_number is None:
+                raise PoolExhaustedError(
+                    f"There are no more values available in this pool, couldn't allocate {quantity} values, only {len(allocated)} available."
+                )
+
+            allocated.append(next_number)
+
+        return allocated
+
 
-def find_next_free(
-    start: int, end: int, taken: list[int | None], parameters: NumberAttributeParameters | None
-) -> int | None:
-    used_numbers = [number for number in taken if number is not None]
-    used_set = set(used_numbers)
+def find_next_free(start: int, end: int, taken: list[int], parameters: NumberAttributeParameters | None) -> int | None:
+    used_set = set(taken)
 
     for num in range(start, end + 1):
         if num not in used_set:

infrahub/core/node/standard.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import inspect
+from enum import Enum
 from typing import TYPE_CHECKING, Any, Optional, Union, get_args, get_origin
 from uuid import UUID
 
@@ -191,6 +192,9 @@ class StandardNode(BaseModel):
                 continue
 
             attr_value = getattr(self, attr_name)
+            if isinstance(attr_value, Enum):
+                attr_value = attr_value.value
+
             field_type = self.guess_field_type(field)
 
             if attr_value is None:

infrahub/core/query/branch.py

@@ -3,43 +3,12 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Any
 
 from infrahub import config
-from infrahub.core.constants import RelationshipStatus
 from infrahub.core.query import Query, QueryType
 
 if TYPE_CHECKING:
     from infrahub.database import InfrahubDatabase
 
 
-class AddNodeToBranch(Query):
-    name: str = "node_add_to_branch"
-    insert_return: bool = False
-
-    type: QueryType = QueryType.WRITE
-
-    def __init__(self, node_id: int, **kwargs: Any):
-        self.node_id = node_id
-        super().__init__(**kwargs)
-
-    async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:  # noqa: ARG002
-        query = """
-        MATCH (root:Root)
-        MATCH (d) WHERE %(id_func)s(d) = $node_id
-        WITH root,d
-        CREATE (d)-[r:IS_PART_OF { branch: $branch, branch_level: $branch_level, from: $now, status: $status }]->(root)
-        RETURN %(id_func)s(r)
-        """ % {
-            "id_func": db.get_id_function_name(),
-        }
-
-        self.params["node_id"] = db.to_database_id(self.node_id)
-        self.params["now"] = self.at.to_string()
-        self.params["branch"] = self.branch.name
-        self.params["branch_level"] = self.branch.hierarchy_level
-        self.params["status"] = RelationshipStatus.ACTIVE.value
-
-        self.add_to_query(query)
-
-
 class DeleteBranchRelationshipsQuery(Query):
     name: str = "delete_branch_relationships"
     insert_return: bool = False
@@ -52,31 +21,31 @@ class DeleteBranchRelationshipsQuery(Query):
 
     async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:  # noqa: ARG002
         query = """
-        MATCH (s)-[r1]-(d)
-        WHERE r1.branch = $branch_name
-        DELETE r1
-
-        WITH collect(DISTINCT s) + collect(DISTINCT d) AS nodes
-
-        // Collect node IDs for filtering
-        WITH nodes, [n in nodes | n.uuid] as nodes_uuids
-
-        // Also delete agnostic relationships that would not have been deleted above
-        MATCH (s2: Node)-[r2]-(d2)
-        WHERE NOT exists((s2)-[:IS_PART_OF]-(:Root))
-        AND s2.uuid IN nodes_uuids
-        DELETE r2
-
-        WITH nodes, collect(DISTINCT s2) + collect(DISTINCT d2) as additional_nodes
-
-        WITH nodes + additional_nodes as nodes
-
-        // Delete nodes that are no longer connected to any other nodes
-        UNWIND nodes AS n
-        WITH DISTINCT n
-        MATCH (n)
-        WHERE NOT exists((n)--())
-        DELETE n
+// delete all relationships on this branch
+MATCH (s)-[r1]-(d)
+WHERE r1.branch = $branch_name
+CALL (r1) {
+  DELETE r1
+} 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
+} 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
+} IN TRANSACTIONS
         """
         self.params["branch_name"] = self.branch_name
         self.add_to_query(query)