infrahub-server 1.7.0rc0__py3-none-any.whl → 1.7.2__py3-none-any.whl

infrahub/core/node/lock_utils.py

@@ -1,6 +1,7 @@
 import hashlib
 from typing import TYPE_CHECKING
 
+from infrahub.core.constants import RelationshipCardinality
 from infrahub.core.node import Node
 from infrahub.core.schema import GenericSchema
 from infrahub.core.schema.schema_branch import SchemaBranch
@@ -10,6 +11,7 @@ if TYPE_CHECKING:
 
 
 RESOURCE_POOL_LOCK_NAMESPACE = "resource_pool"
+RELATIONSHIP_COUNT_LOCK_NAMESPACE = "relationship_count"
 
 
 def _get_kinds_to_lock_on_object_mutation(kind: str, schema_branch: SchemaBranch) -> list[str]:
@@ -55,7 +57,8 @@ def _hash(value: str) -> str:
 def get_lock_names_on_object_mutation(node: Node, schema_branch: SchemaBranch) -> list[str]:
     """
     Return lock names for object on which we want to avoid concurrent mutation (create/update).
-    Lock names include kind, some generic kinds, resource pool ids, and values of attributes of corresponding uniqueness constraints.
+    Lock names include kind, some generic kinds, resource pool ids, peer ids for cardinality one relationships,
+    and values of attributes of corresponding uniqueness constraints.
     """
 
     lock_names: set[str] = set()
@@ -66,13 +69,31 @@ def get_lock_names_on_object_mutation(node: Node, schema_branch: SchemaBranch) -
         if attribute is not None and getattr(attribute, "from_pool", None) and "id" in attribute.from_pool:
             lock_names.add(f"{RESOURCE_POOL_LOCK_NAMESPACE}.{attribute.from_pool['id']}")
 
-    # Check if relationships allocate resources
+    # Check if relationships allocate resources or have cardinality one constraint
     for rel_name in node._relationships:
         rel_manager: RelationshipManager = getattr(node, rel_name)
         for rel in rel_manager._relationships:
             if rel.from_pool and "id" in rel.from_pool:
                 lock_names.add(f"{RESOURCE_POOL_LOCK_NAMESPACE}.{rel.from_pool['id']}")
 
+            peer_id = rel.peer_id
+            if not peer_id or not rel.schema.identifier:
+                continue
+
+            # Check if this node's relationship has cardinality one or max/min_count constraint
+            # This prevents concurrent updates to the same node's constrained relationship
+            if rel.schema.cardinality == RelationshipCardinality.ONE or rel.schema.max_count or rel.schema.min_count:
+                lock_names.add(f"{RELATIONSHIP_COUNT_LOCK_NAMESPACE}.{rel.schema.identifier}.{node.id}")
+
+            # Check if the peer has count constraints on the reverse relationship
+            # This includes cardinality one, max_count, and min_count constraints
+            peer_schema = schema_branch.get(name=rel.schema.peer, duplicate=False)
+            peer_rel = peer_schema.get_relationship_by_identifier(id=rel.schema.identifier, raise_on_error=False)
+            if peer_rel and (
+                peer_rel.cardinality == RelationshipCardinality.ONE or peer_rel.max_count or peer_rel.min_count
+            ):
+                lock_names.add(f"{RELATIONSHIP_COUNT_LOCK_NAMESPACE}.{rel.schema.identifier}.{peer_id}")
+
     lock_kinds = _get_kinds_to_lock_on_object_mutation(node.get_kind(), schema_branch)
     for kind in lock_kinds:
         schema = schema_branch.get(name=kind, duplicate=False)

infrahub/core/node/resource_manager/ip_address_pool.py

@@ -6,13 +6,12 @@ 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
+from infrahub.core.ipam.resource_allocator import IPAMResourceAllocator
 from infrahub.core.query.resource_manager import (
     IPAddressPoolGetReserved,
     IPAddressPoolSetReserved,
 )
 from infrahub.exceptions import PoolExhaustedError, ValidationError
-from infrahub.pools.address import get_available
 
 from .. import Node
 from ..lock_utils import RESOURCE_POOL_LOCK_NAMESPACE
@@ -88,6 +87,7 @@ class CoreIPAddressPool(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]
         ip_namespace = await self.ip_namespace.get_peer(db=db)  # type: ignore[attr-defined]
+        allocator = IPAMResourceAllocator(db=db, namespace=ip_namespace, branch=self._branch, branch_agnostic=True)
 
         try:
             weighted_resources = sorted(resources.values(), key=lambda r: r.allocation_weight.value or 0, reverse=True)
@@ -101,18 +101,12 @@ class CoreIPAddressPool(Node):
             if not ip_prefix.prefixlen <= prefix_length <= ip_prefix.max_prefixlen:
                 raise ValidationError(input_value="Invalid prefix length for current selected prefix")
 
-            addresses = await get_ip_addresses(
-                db=db, ip_prefix=ip_prefix, namespace=ip_namespace, branch=self._branch, branch_agnostic=True
-            )
-
-            available = get_available(
-                network=ip_prefix,
-                addresses=[ip.address for ip in addresses],
+            next_address = await allocator.get_next_address(
+                ip_prefix=ip_prefix,
                 is_pool=resource.is_pool.value,  # type: ignore[attr-defined]
             )
 
-            if available:
-                next_address = available.iter_cidrs()[0]
+            if next_address:
                 return ipaddress.ip_interface(f"{next_address.ip}/{prefix_length}")
 
         raise PoolExhaustedError("There are no more addresses available in this pool.")

infrahub/core/node/resource_manager/ip_prefix_pool.py

@@ -3,18 +3,15 @@ from __future__ import annotations
 import ipaddress
 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
+from infrahub.core.ipam.resource_allocator import IPAMResourceAllocator
 from infrahub.core.query.resource_manager import (
     PrefixPoolGetReserved,
     PrefixPoolSetReserved,
 )
 from infrahub.exceptions import ValidationError
-from infrahub.pools.prefix import get_next_available_prefix
 
 from .. import Node
 from ..lock_utils import RESOURCE_POOL_LOCK_NAMESPACE
@@ -96,6 +93,7 @@ class CoreIPPrefixPool(Node):
     async def get_next(self, db: InfrahubDatabase, prefixlen: int) -> IPNetworkType:
         resources = await self.resources.get_peers(db=db)  # type: ignore[attr-defined]
         ip_namespace = await self.ip_namespace.get_peer(db=db)  # type: ignore[attr-defined]
+        allocator = IPAMResourceAllocator(db=db, namespace=ip_namespace, branch=self._branch, branch_agnostic=True)
 
         try:
             weighted_resources = sorted(resources.values(), key=lambda r: r.allocation_weight.value or 0, reverse=True)
@@ -103,23 +101,9 @@ class CoreIPPrefixPool(Node):
             weighted_resources = list(resources.values())
 
         for resource in weighted_resources:
-            subnets = await get_subnets(
-                db=db,
-                ip_prefix=ipaddress.ip_network(resource.prefix.value),  # type: ignore[attr-defined]
-                namespace=ip_namespace,
-                branch=self._branch,
-                branch_agnostic=True,
-            )
-
-            pool = IPSet([resource.prefix.value])
-            for subnet in subnets:
-                pool.remove(addr=str(subnet.prefix))
-
-            try:
-                prefix_ver = ipaddress.ip_network(resource.prefix.value).version
-                next_available = get_next_available_prefix(pool=pool, prefix_length=prefixlen, prefix_ver=prefix_ver)
+            resource_prefix = ipaddress.ip_network(resource.prefix.value)  # type: ignore[attr-defined]
+            next_available = await allocator.get_next_prefix(ip_prefix=resource_prefix, target_prefix_length=prefixlen)
+            if next_available:
                 return next_available
-            except ValueError:
-                continue
 
         raise IndexError("No more resources available")

infrahub/core/node/resource_manager/number_pool.py

@@ -4,7 +4,12 @@ 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.query.resource_manager import (
+    NumberPoolGetFree,
+    NumberPoolGetReserved,
+    NumberPoolGetUsed,
+    NumberPoolSetReserved,
+)
 from infrahub.core.schema.attribute_parameters import NumberAttributeParameters
 from infrahub.exceptions import PoolExhaustedError
 
@@ -47,6 +52,28 @@ class CoreNumberPool(Node):
         used = [result.value for result in query.iter_results()]
         return [item for item in used if item is not None]
 
+    async def get_free(
+        self, db: InfrahubDatabase, branch: Branch, min_value: int | None = None, max_value: int | None = None
+    ) -> int | None:
+        """Returns the next free number in the pool.
+
+        Args:
+            db: Database connection.
+            branch: Branch to query.
+            min_value: Minimum value to start searching from.
+            max_value: Maximum value to search up to.
+
+        Returns:
+            The next free number, or None if no free numbers are available.
+        """
+
+        query = await NumberPoolGetFree.init(
+            db=db, branch=branch, pool=self, branch_agnostic=True, min_value=min_value, max_value=max_value
+        )
+        await query.execute(db=db)
+
+        return query.get_result_value()
+
     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."""
 
@@ -85,51 +112,94 @@ class CoreNumberPool(Node):
             return number
 
     async def get_next(self, db: InfrahubDatabase, branch: Branch, attribute: AttributeSchema) -> int:
-        taken = await self.get_used(db=db, branch=branch)
+        """Get the next available number from the pool.
 
-        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=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
+        Args:
+            db: Database connection.
+            branch: Branch to query.
+            attribute: Attribute schema that may contain NumberAttributeParameters constraints.
 
-    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)
+        Returns:
+            The next available number that satisfies all constraints.
 
-        allocated: list[int] = []
+        Raises:
+            PoolExhaustedError: If no valid numbers are available in the pool.
+        """
+        parameters = attribute.parameters if isinstance(attribute.parameters, NumberAttributeParameters) else None
 
-        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."
-                )
+        # Extract exclusion constraints from the attribute parameters
+        excluded_values: set[int] = set()
+        excluded_ranges: list[tuple[int, int]] = []
 
-            allocated.append(next_number)
+        if parameters:
+            excluded_values = set(parameters.get_excluded_single_values())
+            excluded_ranges = parameters.get_excluded_ranges()
 
-        return allocated
+        # Compute effective range by combining pool range with min/max constraints
+        pool_start = self.start_range.value  # type: ignore[attr-defined]
+        pool_end = self.end_range.value  # type: ignore[attr-defined]
 
+        effective_start = pool_start
+        effective_end = pool_end
 
-def find_next_free(start: int, end: int, taken: list[int], parameters: NumberAttributeParameters | None) -> int | None:
-    used_set = set(taken)
+        if parameters:
+            if parameters.min_value is not None:
+                effective_start = max(effective_start, parameters.min_value)
+            if parameters.max_value is not None:
+                effective_end = min(effective_end, parameters.max_value)
 
-    for num in range(start, end + 1):
-        if num not in used_set:
-            if parameters is None or parameters.is_valid_value(num):
-                return num
+        # Check if the effective range is valid
+        if effective_start > effective_end:
+            raise PoolExhaustedError("There are no more values available in this pool.")
 
-    return None
+        def skip_excluded(value: int) -> int | None:
+            """Skip past any excluded values/ranges starting from value.
+
+            Returns the next non-excluded value, or None if we exceed effective_end.
+            """
+            current = value
+            while current <= effective_end:
+                # Check if in an excluded range and skip past it
+                in_range = False
+                for range_start, range_end in excluded_ranges:
+                    if range_start <= current <= range_end:
+                        current = range_end + 1
+                        in_range = True
+                        break
+                if in_range:
+                    continue
+
+                # Check if it's an excluded single value
+                if current in excluded_values:
+                    current += 1
+                    continue
+
+                # Found a non-excluded value
+                return current
+
+            return None
+
+        # Skip any excluded values at the start
+        first_valid = skip_excluded(effective_start)
+        if first_valid is None:
+            raise PoolExhaustedError("There are no more values available in this pool.")
+        min_value = first_valid
+
+        # Re-run the query until we find a non-excluded value or exhaust the pool
+        while True:
+            candidate = await self.get_free(db=db, branch=branch, min_value=min_value, max_value=effective_end)
+            if candidate is None:
+                raise PoolExhaustedError("There are no more values available in this pool.")
+
+            # Check if candidate is excluded (single value or range)
+            next_valid = skip_excluded(candidate)
+            if next_valid is None:
+                raise PoolExhaustedError("There are no more values available in this pool.")
+
+            if next_valid != candidate:
+                # Candidate was excluded, re-query starting from next valid point
+                min_value = next_valid
+                continue
+
+            # Candidate passed all checks
+            return candidate

infrahub/core/query/__init__.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from collections import defaultdict
 from dataclasses import dataclass, field
 from enum import Enum
+from functools import lru_cache
 from typing import TYPE_CHECKING, Any, Callable, Generator, Iterator, TypeVar
 
 import ujson
@@ -417,6 +418,11 @@ class Query:
         Right now it's mainly used to add more labels to the metrics."""
         return {}
 
+    @staticmethod
+    @lru_cache(maxsize=1024)
+    def _split_query_lines(query: str) -> list[str]:
+        return [line.strip() for line in query.split("\n") if line.strip()]
+
     def add_to_query(self, query: str | list[str]) -> None:
         """Add a new section at the end of the query.
 
@@ -427,7 +433,7 @@ class Query:
             for item in query:
                 self.add_to_query(query=item)
         else:
-            self.query_lines.extend([line.strip() for line in query.split("\n") if line.strip()])
+            self.query_lines.extend(self._split_query_lines(query=query))
 
     def add_subquery(self, subquery: str, node_alias: str, with_clause: str | None = None) -> None:
         self.add_to_query(f"CALL ({node_alias}) {{")

infrahub/core/query/branch.py

@@ -24,15 +24,31 @@ class DeleteBranchRelationshipsQuery(Query):
     async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:  # noqa: ARG002
         query = """
 // --------------
-// for every Node created on this branch (it's about to be deleted), find any agnostic relationships
-// connected to the Node and delete them
+// for every Node that only exists on this branch (it's about to be deleted),
+// find any agnostic relationships or attributes connected to the Node and delete them
 // --------------
 OPTIONAL MATCH (:Root)<-[e:IS_PART_OF {status: "active"}]-(n:Node)
 WHERE e.branch = $branch_name
+// does the node only exist on this branch?
 CALL (n) {
+    OPTIONAL MATCH (n)-[ipo:IS_PART_OF {status: "active"}]->(:Root)
+    WHERE ipo.branch <> $branch_name
+    LIMIT 1
+    RETURN ipo IS NOT NULL AS node_exists_on_other_branch
+}
+// if so, delete any linked agnostic relationships or attributes
+CALL (n, node_exists_on_other_branch) {
+    WITH n, node_exists_on_other_branch
+    WHERE node_exists_on_other_branch = FALSE
     OPTIONAL MATCH (n)-[:IS_RELATED {branch: $global_branch_name}]-(rel:Relationship)
     DETACH DELETE rel
 } IN TRANSACTIONS OF 500 ROWS
+CALL (n, node_exists_on_other_branch) {
+    WITH n, node_exists_on_other_branch
+    WHERE node_exists_on_other_branch = FALSE
+    OPTIONAL MATCH (n)-[:HAS_ATTRIBUTE {branch: $global_branch_name}]-(attr:Attribute)
+    DETACH DELETE attr
+} IN TRANSACTIONS OF 500 ROWS
 
 // reduce the results to a single row
 WITH 1 AS one