Merge pull request #6507 from opsmill/pog-attributekind-numberpool-IFC-1419

Add support for NumberPool attribute kind

Author: ogenstad

backend/infrahub/core/node/__init__.py

@@ -30,6 +30,7 @@
     RelationshipSchema,
     TemplateSchema,
 )
+from infrahub.core.schema.attribute_parameters import NumberPoolParameters
 from infrahub.core.timestamp import Timestamp
 from infrahub.exceptions import InitializationError, NodeNotFoundError, PoolExhaustedError, ValidationError
 from infrahub.types import ATTRIBUTE_TYPES
@@ -254,6 +255,12 @@ async def handle_pool(self, db: InfrahubDatabase, attribute: BaseAttribute, erro
         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
 
@@ -262,12 +269,18 @@ async def handle_pool(self, db: InfrahubDatabase, attribute: BaseAttribute, erro
                 db=db, id=attribute.from_pool["id"], kind=CoreNumberPool
             )
         except NodeNotFoundError:
-            errors.append(
-                ValidationError(
-                    {f"{attribute.name}.from_pool": f"The pool requested {attribute.from_pool} was not found."}
+            if number_pool_parameters:
+                number_pool = await self._create_number_pool(
+                    db=db, attribute=attribute, number_pool_parameters=number_pool_parameters
                 )
-            )
-            return
+
+            else:
+                errors.append(
+                    ValidationError(
+                        {f"{attribute.name}.from_pool": f"The pool requested {attribute.from_pool} was not found."}
+                    )
+                )
+                return
 
         if (
             number_pool.node.value in [self._schema.kind] + self._schema.inherit_from
@@ -292,6 +305,35 @@ async def handle_pool(self, db: InfrahubDatabase, attribute: BaseAttribute, erro
                 )
             )
 
+    async def _create_number_pool(
+        self, db: InfrahubDatabase, attribute: BaseAttribute, number_pool_parameters: NumberPoolParameters
+    ) -> CoreNumberPool:
+        schema = db.schema.get_node_schema(name="CoreNumberPool", duplicate=False)
+
+        pool_node = self._schema.kind
+        schema_attribute = self._schema.get_attribute(attribute.schema.name)
+        if schema_attribute.inherited:
+            for generic_name in self._schema.inherit_from:
+                generic_node = db.schema.get_generic_schema(name=generic_name, duplicate=False)
+                if attribute.schema.name in generic_node.attribute_names:
+                    pool_node = generic_node.kind
+                    break
+
+        number_pool = await Node.init(db=db, schema=schema, branch=self._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}]",
+            node=pool_node,
+            node_attribute=attribute.schema.name,
+            start_range=number_pool_parameters.start_range,
+            end_range=number_pool_parameters.end_range,
+        )
+        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.
+        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."""
         object_template_field = fields.get(OBJECT_TEMPLATE_RELATIONSHIP_NAME)
@@ -342,7 +384,7 @@ async def handle_object_template(self, fields: dict, db: InfrahubDatabase, error
             elif relationship_peers := await relationship.get_peers(db=db):
                 fields[relationship_name] = [{"id": peer_id} for peer_id in relationship_peers]
 
-    async def _process_fields(self, fields: dict, db: InfrahubDatabase) -> None:
+    async def _process_fields(self, fields: dict, db: InfrahubDatabase) -> None:  # noqa: PLR0915
         errors = []
 
         if "_source" in fields.keys():
@@ -376,6 +418,9 @@ async def _process_fields(self, fields: dict, db: InfrahubDatabase) -> None:
                             self._computed_jinja2_attributes.append(mandatory_attr)
                             continue
 
+                        if mandatory_attribute.kind == "NumberPool":
+                            continue
+
                     errors.append(
                         ValidationError({mandatory_attr: f"{mandatory_attr} is mandatory for {self.get_kind()}"})
                     )

backend/infrahub/core/node/resource_manager/number_pool.py

@@ -37,7 +37,6 @@ async def get_resource(
             db=db, pool_id=self.get_id(), identifier=identifier, reserved=number
         )
         await query_set.execute(db=db)
-
         return number
 
     async def get_next(self, db: InfrahubDatabase, branch: Branch) -> int:

backend/infrahub/core/schema/attribute_parameters.py

@@ -1,23 +1,52 @@
 from __future__ import annotations
 
-from pydantic import Field
+import sys
+from typing import Self
+
+from pydantic import Field, model_validator
 
 from infrahub.core.constants.schema import UpdateSupport
 from infrahub.core.models import HashableModel
 
 
 def get_attribute_parameters_class_for_kind(kind: str) -> type[AttributeParameters]:
-    return {
+    param_classes: dict[str, type[AttributeParameters]] = {
+        "NumberPool": NumberPoolParameters,
         "Text": TextAttributeParameters,
         "TextArea": TextAttributeParameters,
-    }.get(kind, AttributeParameters)
+    }
+    return param_classes.get(kind, AttributeParameters)
 
 
 class AttributeParameters(HashableModel):
     class Config:
         extra = "forbid"
 
 
+class NumberPoolParameters(AttributeParameters):
+    end_range: int = Field(
+        default=sys.maxsize,
+        description="End range for numbers for the associated NumberPool",
+        json_schema_extra={"update": UpdateSupport.VALIDATE_CONSTRAINT.value},
+    )
+    start_range: int = Field(
+        default=1,
+        description="Start range for numbers for the associated NumberPool",
+        json_schema_extra={"update": UpdateSupport.VALIDATE_CONSTRAINT.value},
+    )
+    number_pool_id: str | None = Field(
+        default=None,
+        description="The ID of the numberpool associated with this attribute",
+        json_schema_extra={"update": UpdateSupport.NOT_SUPPORTED.value},
+    )
+
+    @model_validator(mode="after")
+    def validate_ranges(self) -> Self:
+        if self.start_range > self.end_range:
+            raise ValueError("start_range can't be less than end_range")
+        return self
+
+
 class TextAttributeParameters(AttributeParameters):
     regex: str | None = Field(
         default=None,

backend/infrahub/core/schema/attribute_schema.py

@@ -12,7 +12,12 @@
 from infrahub.core.query.attribute import default_attribute_query_filter
 from infrahub.types import ATTRIBUTE_KIND_LABELS, ATTRIBUTE_TYPES
 
-from .attribute_parameters import AttributeParameters, TextAttributeParameters, get_attribute_parameters_class_for_kind
+from .attribute_parameters import (
+    AttributeParameters,
+    NumberPoolParameters,
+    TextAttributeParameters,
+    get_attribute_parameters_class_for_kind,
+)
 from .generated.attribute_schema import GeneratedAttributeSchema
 
 if TYPE_CHECKING:
@@ -25,6 +30,7 @@
 
 def get_attribute_schema_class_for_kind(kind: str) -> type[AttributeSchema]:
     attribute_schema_class_by_kind: dict[str, type[AttributeSchema]] = {
+        "NumberPool": NumberPoolSchema,
         "Text": TextAttributeSchema,
         "TextArea": TextAttributeSchema,
     }
@@ -93,6 +99,16 @@ def set_parameters_type(cls, value: Any, info: ValidationInfo) -> Any:
             return expected_parameters_class(**value.model_dump())
         return value
 
+    @model_validator(mode="after")
+    def validate_parameters(self) -> Self:
+        if isinstance(self.parameters, NumberPoolParameters) and not self.kind == "NumberPool":
+            raise ValueError(f"NumberPoolParameters can't be used as parameters for {self.kind}")
+
+        if isinstance(self.parameters, TextAttributeParameters) and self.kind not in ["Text", "TextArea"]:
+            raise ValueError(f"TextAttributeParameters can't be used as parameters for {self.kind}")
+
+        return self
+
     def get_class(self) -> type[BaseAttribute]:
         return ATTRIBUTE_TYPES[self.kind].get_infrahub_class()
 
@@ -185,6 +201,14 @@ async def get_query_filter(
         )
 
 
+class NumberPoolSchema(AttributeSchema):
+    parameters: NumberPoolParameters = Field(
+        default_factory=NumberPoolParameters,
+        description="Extra parameters specific to NumberPool attributes",
+        json_schema_extra={"update": UpdateSupport.VALIDATE_CONSTRAINT.value},
+    )
+
+
 class TextAttributeSchema(AttributeSchema):
     parameters: TextAttributeParameters = Field(
         default_factory=TextAttributeParameters,

backend/infrahub/core/schema/definitions/internal.py

@@ -31,7 +31,7 @@
     RelationshipKind,
     UpdateSupport,
 )
-from infrahub.core.schema.attribute_parameters import AttributeParameters
+from infrahub.core.schema.attribute_parameters import AttributeParameters, NumberPoolParameters, TextAttributeParameters
 from infrahub.core.schema.attribute_schema import AttributeSchema
 from infrahub.core.schema.computed_attribute import ComputedAttribute
 from infrahub.core.schema.dropdown import DropdownChoice
@@ -48,7 +48,7 @@ class SchemaAttribute(BaseModel):
     kind: str
     description: str
     extra: ExtraField
-    internal_kind: type[Any] | GenericAlias | None = None
+    internal_kind: type[Any] | GenericAlias | list[type[Any]] | None = None
     regex: str | None = None
     unique: bool | None = None
     optional: bool | None = None
@@ -94,6 +94,9 @@ def object_kind(self) -> str:
         if isinstance(self.internal_kind, GenericAlias):
             return str(self.internal_kind)
 
+        if isinstance(self.internal_kind, list):
+            return " | ".join([internal_kind.__name__ for internal_kind in self.internal_kind])
+
         if self.internal_kind and self.kind == "List":
             return f"list[{self.internal_kind.__name__}]"
 
@@ -621,7 +624,7 @@ def to_dict(self) -> dict[str, Any]:
         SchemaAttribute(
             name="parameters",
             kind="JSON",
-            internal_kind=AttributeParameters,
+            internal_kind=[AttributeParameters, TextAttributeParameters, NumberPoolParameters],
             optional=True,
             description="Extra parameters specific to this kind of attribute",
             extra={"update": UpdateSupport.VALIDATE_CONSTRAINT},

backend/infrahub/core/schema/generated/attribute_schema.py

@@ -8,7 +8,11 @@
 
 from infrahub.core.constants import AllowOverrideType, BranchSupportType, HashableModelState
 from infrahub.core.models import HashableModel
-from infrahub.core.schema.attribute_parameters import AttributeParameters  # noqa: TC001
+from infrahub.core.schema.attribute_parameters import (
+    AttributeParameters,  # noqa: TC001
+    NumberPoolParameters,  # noqa: TC001
+    TextAttributeParameters,  # noqa: TC001
+)
 from infrahub.core.schema.computed_attribute import ComputedAttribute  # noqa: TC001
 from infrahub.core.schema.dropdown import DropdownChoice  # noqa: TC001
 
@@ -113,7 +117,7 @@ class GeneratedAttributeSchema(HashableModel):
         description="Type of allowed override for the attribute.",
         json_schema_extra={"update": "allowed"},
     )
-    parameters: AttributeParameters = Field(
+    parameters: AttributeParameters | TextAttributeParameters | NumberPoolParameters = Field(
         default_factory=AttributeParameters,
         description="Extra parameters specific to this kind of attribute",
         json_schema_extra={"update": "validate_constraint"},

backend/infrahub/core/schema/schema_branch.py

@@ -5,6 +5,7 @@
 from collections import defaultdict
 from itertools import chain, combinations
 from typing import Any
+from uuid import uuid4
 
 from infrahub_sdk.template import Jinja2Template
 from infrahub_sdk.template.exceptions import JinjaTemplateError, JinjaTemplateOperationViolationError
@@ -49,6 +50,7 @@
     SchemaRoot,
     TemplateSchema,
 )
+from infrahub.core.schema.attribute_parameters import NumberPoolParameters
 from infrahub.core.schema.attribute_schema import get_attribute_schema_class_for_kind
 from infrahub.core.schema.definitions.core import core_profile_schema_definition
 from infrahub.core.validators import CONSTRAINT_VALIDATOR_MAP
@@ -518,6 +520,7 @@ def process_validate(self) -> None:
         self.validate_names()
         self.validate_kinds()
         self.validate_computed_attributes()
+        self.validate_attribute_parameters()
         self.validate_default_values()
         self.validate_count_against_cardinality()
         self.validate_identifiers()
@@ -995,6 +998,50 @@ def validate_kinds(self) -> None:
                         f"{node.kind}: Relationship {rel.name!r} is referring an invalid peer {rel.peer!r}"
                     ) from None
 
+    def validate_attribute_parameters(self) -> None:
+        for name in self.generics.keys():
+            generic_schema = self.get_generic(name=name, duplicate=False)
+            for attribute in generic_schema.attributes:
+                if (
+                    attribute.kind == "NumberPool"
+                    and isinstance(attribute.parameters, NumberPoolParameters)
+                    and not attribute.parameters.number_pool_id
+                ):
+                    attribute.parameters.number_pool_id = str(uuid4())
+
+        for name in self.nodes.keys():
+            node_schema = self.get_node(name=name, duplicate=False)
+            for attribute in node_schema.attributes:
+                if (
+                    attribute.kind == "NumberPool"
+                    and isinstance(attribute.parameters, NumberPoolParameters)
+                    and not attribute.parameters.number_pool_id
+                ):
+                    self._validate_number_pool_parameters(
+                        node_schema=node_schema, attribute=attribute, number_pool_parameters=attribute.parameters
+                    )
+
+    def _validate_number_pool_parameters(
+        self, node_schema: NodeSchema, attribute: AttributeSchema, number_pool_parameters: NumberPoolParameters
+    ) -> None:
+        if attribute.inherited:
+            generics_with_attribute = []
+            for generic_name in node_schema.inherit_from:
+                generic_schema = self.get_generic(name=generic_name, duplicate=False)
+                if attribute.name in generic_schema.attribute_names:
+                    generic_attribute = generic_schema.get_attribute(name=attribute.name)
+                    generics_with_attribute.append(generic_schema)
+                    if isinstance(generic_attribute.parameters, NumberPoolParameters):
+                        number_pool_parameters.number_pool_id = generic_attribute.parameters.number_pool_id
+
+            if len(generics_with_attribute) > 1:
+                raise ValidationError(
+                    f"{node_schema.kind}.{attribute.name} is a NumberPool inherited from more than one generic"
+                )
+
+        else:
+            number_pool_parameters.number_pool_id = str(uuid4())
+
     def validate_computed_attributes(self) -> None:
         self.computed_attributes = ComputedAttributes()
         for name in self.nodes.keys():

backend/infrahub/database/__init__.py

@@ -39,7 +39,7 @@
     from types import TracebackType
 
     from infrahub.core.branch import Branch
-    from infrahub.core.schema import MainSchemaTypes, NodeSchema
+    from infrahub.core.schema import GenericSchema, MainSchemaTypes, NodeSchema
     from infrahub.core.schema.schema_branch import SchemaBranch
 
 validated_database = {}
@@ -91,6 +91,15 @@ def get_node_schema(self, name: str, branch: Branch | str | None = None, duplica
 
         raise ValueError("The selected node is not of type NodeSchema")
 
+    def get_generic_schema(
+        self, name: str, branch: Branch | str | None = None, duplicate: bool = True
+    ) -> GenericSchema:
+        schema = self.get(name=name, branch=branch, duplicate=duplicate)
+        if schema.is_generic_schema:
+            return schema
+
+        raise ValueError("The selected node is not of type GenericSchema")
+
     def set(self, name: str, schema: MainSchemaTypes, branch: str | None = None) -> int:
         branch_name = get_branch_name(branch=branch)
         if branch_name not in self._db._schemas: