Skip to content

Commit 27070dd

Browse files
ogenstadogenstad
committed
Add support for NumberPool attribute kind
1 parent f47464c commit 27070dd

File tree

13 files changed

+375
-13
lines changed

13 files changed

+375
-13
lines changed

‎backend/infrahub/core/node/__init__.py‎

Lines changed: 43 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
RelationshipSchema,
3131
TemplateSchema,
3232
)
33+
from infrahub.core.schema.attribute_parameters import NumberPoolParameters
3334
from infrahub.core.timestamp import Timestamp
3435
from infrahub.exceptions import InitializationError, NodeNotFoundError, PoolExhaustedError, ValidationError
3536
from infrahub.types import ATTRIBUTE_TYPES
@@ -254,6 +255,12 @@ async def handle_pool(self, db: InfrahubDatabase, attribute: BaseAttribute, erro
254255
within the create code.
255256
"""
256257

258+
number_pool_parameters: NumberPoolParameters | None = None
259+
if attribute.schema.kind == "NumberPool" and isinstance(attribute.schema.parameters, NumberPoolParameters):
260+
attribute.from_pool = {"id": attribute.schema.parameters.number_pool_id}
261+
attribute.is_default = False
262+
number_pool_parameters = attribute.schema.parameters
263+
257264
if not attribute.from_pool:
258265
return
259266

@@ -262,6 +269,11 @@ async def handle_pool(self, db: InfrahubDatabase, attribute: BaseAttribute, erro
262269
db=db, id=attribute.from_pool["id"], kind=CoreNumberPool
263270
)
264271
except NodeNotFoundError:
272+
if number_pool_parameters:
273+
await self._create_number_pool(
274+
db=db, attribute=attribute, errors=errors, number_pool_parameters=number_pool_parameters
275+
)
276+
return
265277
errors.append(
266278
ValidationError(
267279
{f"{attribute.name}.from_pool": f"The pool requested {attribute.from_pool} was not found."}
@@ -292,6 +304,33 @@ async def handle_pool(self, db: InfrahubDatabase, attribute: BaseAttribute, erro
292304
)
293305
)
294306

307+
async def _create_number_pool(
308+
self, db: InfrahubDatabase, attribute: BaseAttribute, errors: list, number_pool_parameters: NumberPoolParameters
309+
) -> None:
310+
schema = db.schema.get_node_schema(name="CoreNumberPool", duplicate=False)
311+
312+
pool_node = self._schema.kind
313+
schema_attribute = self._schema.get_attribute(attribute.schema.name)
314+
if schema_attribute.inherited:
315+
for generic_name in self._schema.inherit_from:
316+
generic_node = db.schema.get_generic_schema(name=generic_name, duplicate=False)
317+
if attribute.schema.name in generic_node.attribute_names:
318+
pool_node = generic_node.kind
319+
break
320+
321+
number_pool = await Node.init(db=db, schema=schema, branch=self._branch)
322+
await number_pool.new(
323+
db=db,
324+
id=number_pool_parameters.number_pool_id,
325+
name=f"{pool_node}.{attribute.schema.name} [{number_pool_parameters.number_pool_id}]",
326+
node=pool_node,
327+
node_attribute=attribute.schema.name,
328+
start_range=number_pool_parameters.start_range,
329+
end_range=number_pool_parameters.end_range,
330+
)
331+
await number_pool.save(db=db)
332+
await self.handle_pool(db=db, attribute=attribute, errors=errors)
333+
295334
async def handle_object_template(self, fields: dict, db: InfrahubDatabase, errors: list) -> None:
296335
"""Fill the `fields` parameters with values from an object template if one is in use."""
297336
object_template_field = fields.get(OBJECT_TEMPLATE_RELATIONSHIP_NAME)
@@ -342,7 +381,7 @@ async def handle_object_template(self, fields: dict, db: InfrahubDatabase, error
342381
elif relationship_peers := await relationship.get_peers(db=db):
343382
fields[relationship_name] = [{"id": peer_id} for peer_id in relationship_peers]
344383

345-
async def _process_fields(self, fields: dict, db: InfrahubDatabase) -> None:
384+
async def _process_fields(self, fields: dict, db: InfrahubDatabase) -> None: # noqa: PLR0915
346385
errors = []
347386

348387
if "_source" in fields.keys():
@@ -376,6 +415,9 @@ async def _process_fields(self, fields: dict, db: InfrahubDatabase) -> None:
376415
self._computed_jinja2_attributes.append(mandatory_attr)
377416
continue
378417

418+
if mandatory_attribute.kind == "NumberPool":
419+
continue
420+
379421
errors.append(
380422
ValidationError({mandatory_attr: f"{mandatory_attr} is mandatory for {self.get_kind()}"})
381423
)

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,6 @@ async def get_resource(
3737
db=db, pool_id=self.get_id(), identifier=identifier, reserved=number
3838
)
3939
await query_set.execute(db=db)
40-
4140
return number
4241

4342
async def get_next(self, db: InfrahubDatabase, branch: Branch) -> int:

‎backend/infrahub/core/schema/attribute_parameters.py‎

Lines changed: 32 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,52 @@
11
from __future__ import annotations
22

3-
from pydantic import Field
3+
import sys
4+
from typing import Self
5+
6+
from pydantic import Field, model_validator
47

58
from infrahub.core.constants.schema import UpdateSupport
69
from infrahub.core.models import HashableModel
710

811

912
def get_attribute_parameters_class_for_kind(kind: str) -> type[AttributeParameters]:
10-
return {
13+
param_classes: dict[str, type[AttributeParameters]] = {
14+
"NumberPool": NumberPoolParameters,
1115
"Text": TextAttributeParameters,
1216
"TextArea": TextAttributeParameters,
13-
}.get(kind, AttributeParameters)
17+
}
18+
return param_classes.get(kind, AttributeParameters)
1419

1520

1621
class AttributeParameters(HashableModel):
1722
class Config:
1823
extra = "forbid"
1924

2025

26+
class NumberPoolParameters(AttributeParameters):
27+
end_range: int = Field(
28+
default=sys.maxsize,
29+
description="End range for numbers for the associated NumberPool",
30+
json_schema_extra={"update": UpdateSupport.VALIDATE_CONSTRAINT.value},
31+
)
32+
start_range: int = Field(
33+
default=1,
34+
description="Start range for numbers for the associated NumberPool",
35+
json_schema_extra={"update": UpdateSupport.VALIDATE_CONSTRAINT.value},
36+
)
37+
number_pool_id: str | None = Field(
38+
default=None,
39+
description="The ID of the numberpool associated with this attribute",
40+
json_schema_extra={"update": UpdateSupport.NOT_SUPPORTED.value},
41+
)
42+
43+
@model_validator(mode="after")
44+
def validate_ranges(self) -> Self:
45+
if self.start_range > self.end_range:
46+
raise ValueError("start_range can't be less than end_range")
47+
return self
48+
49+
2150
class TextAttributeParameters(AttributeParameters):
2251
regex: str | None = Field(
2352
default=None,

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

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,12 @@
1212
from infrahub.core.query.attribute import default_attribute_query_filter
1313
from infrahub.types import ATTRIBUTE_KIND_LABELS, ATTRIBUTE_TYPES
1414

15-
from .attribute_parameters import AttributeParameters, TextAttributeParameters, get_attribute_parameters_class_for_kind
15+
from .attribute_parameters import (
16+
AttributeParameters,
17+
NumberPoolParameters,
18+
TextAttributeParameters,
19+
get_attribute_parameters_class_for_kind,
20+
)
1621
from .generated.attribute_schema import GeneratedAttributeSchema
1722

1823
if TYPE_CHECKING:
@@ -25,6 +30,7 @@
2530

2631
def get_attribute_schema_class_for_kind(kind: str) -> type[AttributeSchema]:
2732
attribute_schema_class_by_kind: dict[str, type[AttributeSchema]] = {
33+
"NumberPool": NumberPoolSchema,
2834
"Text": TextAttributeSchema,
2935
"TextArea": TextAttributeSchema,
3036
}
@@ -93,6 +99,16 @@ def set_parameters_type(cls, value: Any, info: ValidationInfo) -> Any:
9399
return expected_parameters_class(**value.model_dump())
94100
return value
95101

102+
@model_validator(mode="after")
103+
def validate_parameters(self) -> Self:
104+
if isinstance(self.parameters, NumberPoolParameters) and not self.kind == "NumberPool":
105+
raise ValueError(f"NumberPoolParameters can't be used as parameters for {self.kind}")
106+
107+
if isinstance(self.parameters, TextAttributeParameters) and self.kind not in ["Text", "TextArea"]:
108+
raise ValueError(f"TextAttributeParameters can't be used as parameters for {self.kind}")
109+
110+
return self
111+
96112
def get_class(self) -> type[BaseAttribute]:
97113
return ATTRIBUTE_TYPES[self.kind].get_infrahub_class()
98114

@@ -185,6 +201,14 @@ async def get_query_filter(
185201
)
186202

187203

204+
class NumberPoolSchema(AttributeSchema):
205+
parameters: NumberPoolParameters = Field(
206+
default_factory=NumberPoolParameters,
207+
description="Extra parameters specific to text attributes",
208+
json_schema_extra={"update": UpdateSupport.VALIDATE_CONSTRAINT.value},
209+
)
210+
211+
188212
class TextAttributeSchema(AttributeSchema):
189213
parameters: TextAttributeParameters = Field(
190214
default_factory=TextAttributeParameters,

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

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@
3131
RelationshipKind,
3232
UpdateSupport,
3333
)
34-
from infrahub.core.schema.attribute_parameters import AttributeParameters
34+
from infrahub.core.schema.attribute_parameters import AttributeParameters, NumberPoolParameters, TextAttributeParameters
3535
from infrahub.core.schema.attribute_schema import AttributeSchema
3636
from infrahub.core.schema.computed_attribute import ComputedAttribute
3737
from infrahub.core.schema.dropdown import DropdownChoice
@@ -48,7 +48,7 @@ class SchemaAttribute(BaseModel):
4848
kind: str
4949
description: str
5050
extra: ExtraField
51-
internal_kind: type[Any] | GenericAlias | None = None
51+
internal_kind: type[Any] | GenericAlias | list[type[Any]] | None = None
5252
regex: str | None = None
5353
unique: bool | None = None
5454
optional: bool | None = None
@@ -94,6 +94,9 @@ def object_kind(self) -> str:
9494
if isinstance(self.internal_kind, GenericAlias):
9595
return str(self.internal_kind)
9696

97+
if isinstance(self.internal_kind, list):
98+
return " | ".join([internal_kind.__name__ for internal_kind in self.internal_kind])
99+
97100
if self.internal_kind and self.kind == "List":
98101
return f"list[{self.internal_kind.__name__}]"
99102

@@ -621,7 +624,7 @@ def to_dict(self) -> dict[str, Any]:
621624
SchemaAttribute(
622625
name="parameters",
623626
kind="JSON",
624-
internal_kind=AttributeParameters,
627+
internal_kind=[AttributeParameters, TextAttributeParameters, NumberPoolParameters],
625628
optional=True,
626629
description="Extra parameters specific to this kind of attribute",
627630
extra={"update": UpdateSupport.VALIDATE_CONSTRAINT},

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,11 @@
88

99
from infrahub.core.constants import AllowOverrideType, BranchSupportType, HashableModelState
1010
from infrahub.core.models import HashableModel
11-
from infrahub.core.schema.attribute_parameters import AttributeParameters # noqa: TC001
11+
from infrahub.core.schema.attribute_parameters import (
12+
AttributeParameters, # noqa: TC001
13+
NumberPoolParameters, # noqa: TC001
14+
TextAttributeParameters, # noqa: TC001
15+
)
1216
from infrahub.core.schema.computed_attribute import ComputedAttribute # noqa: TC001
1317
from infrahub.core.schema.dropdown import DropdownChoice # noqa: TC001
1418

@@ -113,7 +117,7 @@ class GeneratedAttributeSchema(HashableModel):
113117
description="Type of allowed override for the attribute.",
114118
json_schema_extra={"update": "allowed"},
115119
)
116-
parameters: AttributeParameters = Field(
120+
parameters: AttributeParameters | TextAttributeParameters | NumberPoolParameters = Field(
117121
default_factory=AttributeParameters,
118122
description="Extra parameters specific to this kind of attribute",
119123
json_schema_extra={"update": "validate_constraint"},

‎backend/infrahub/core/schema/schema_branch.py‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from collections import defaultdict
66
from itertools import chain, combinations
77
from typing import Any
8+
from uuid import uuid4
89

910
from infrahub_sdk.template import Jinja2Template
1011
from infrahub_sdk.template.exceptions import JinjaTemplateError, JinjaTemplateOperationViolationError
@@ -49,6 +50,7 @@
4950
SchemaRoot,
5051
TemplateSchema,
5152
)
53+
from infrahub.core.schema.attribute_parameters import NumberPoolParameters
5254
from infrahub.core.schema.attribute_schema import get_attribute_schema_class_for_kind
5355
from infrahub.core.schema.definitions.core import core_profile_schema_definition
5456
from infrahub.core.validators import CONSTRAINT_VALIDATOR_MAP
@@ -518,6 +520,7 @@ def process_validate(self) -> None:
518520
self.validate_names()
519521
self.validate_kinds()
520522
self.validate_computed_attributes()
523+
self.validate_attribute_parameters()
521524
self.validate_default_values()
522525
self.validate_count_against_cardinality()
523526
self.validate_identifiers()
@@ -995,6 +998,50 @@ def validate_kinds(self) -> None:
995998
f"{node.kind}: Relationship {rel.name!r} is referring an invalid peer {rel.peer!r}"
996999
) from None
9971000

1001+
def validate_attribute_parameters(self) -> None:
1002+
for name in self.generics.keys():
1003+
generic_schema = self.get_generic(name=name, duplicate=False)
1004+
for attribute in generic_schema.attributes:
1005+
if (
1006+
attribute.kind == "NumberPool"
1007+
and isinstance(attribute.parameters, NumberPoolParameters)
1008+
and not attribute.parameters.number_pool_id
1009+
):
1010+
attribute.parameters.number_pool_id = str(uuid4())
1011+
1012+
for name in self.nodes.keys():
1013+
node_schema = self.get_node(name=name, duplicate=False)
1014+
for attribute in node_schema.attributes:
1015+
if (
1016+
attribute.kind == "NumberPool"
1017+
and isinstance(attribute.parameters, NumberPoolParameters)
1018+
and not attribute.parameters.number_pool_id
1019+
):
1020+
self._validate_number_pool_parameters(
1021+
node_schema=node_schema, attribute=attribute, number_pool_parameters=attribute.parameters
1022+
)
1023+
1024+
def _validate_number_pool_parameters(
1025+
self, node_schema: NodeSchema, attribute: AttributeSchema, number_pool_parameters: NumberPoolParameters
1026+
) -> None:
1027+
if attribute.inherited:
1028+
generics_with_attribute = []
1029+
for generic_name in node_schema.inherit_from:
1030+
generic_schema = self.get_generic(name=generic_name, duplicate=False)
1031+
if attribute.name in generic_schema.attribute_names:
1032+
generic_attribute = generic_schema.get_attribute(name=attribute.name)
1033+
generics_with_attribute.append(generic_schema)
1034+
if isinstance(generic_attribute.parameters, NumberPoolParameters):
1035+
number_pool_parameters.number_pool_id = generic_attribute.parameters.number_pool_id
1036+
1037+
if len(generics_with_attribute) > 1:
1038+
raise ValidationError(
1039+
f"{node_schema.kind}.{attribute.name} is a NumberPool inherited from more than one generic"
1040+
)
1041+
1042+
else:
1043+
number_pool_parameters.number_pool_id = str(uuid4())
1044+
9981045
def validate_computed_attributes(self) -> None:
9991046
self.computed_attributes = ComputedAttributes()
10001047
for name in self.nodes.keys():

‎backend/infrahub/database/__init__.py‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@
3939
from types import TracebackType
4040

4141
from infrahub.core.branch import Branch
42-
from infrahub.core.schema import MainSchemaTypes, NodeSchema
42+
from infrahub.core.schema import GenericSchema, MainSchemaTypes, NodeSchema
4343
from infrahub.core.schema.schema_branch import SchemaBranch
4444

4545
validated_database = {}
@@ -91,6 +91,15 @@ def get_node_schema(self, name: str, branch: Branch | str | None = None, duplica
9191

9292
raise ValueError("The selected node is not of type NodeSchema")
9393

94+
def get_generic_schema(
95+
self, name: str, branch: Branch | str | None = None, duplicate: bool = True
96+
) -> GenericSchema:
97+
schema = self.get(name=name, branch=branch, duplicate=duplicate)
98+
if schema.is_generic_schema:
99+
return schema
100+
101+
raise ValueError("The selected node is not of type GenericSchema")
102+
94103
def set(self, name: str, schema: MainSchemaTypes, branch: str | None = None) -> int:
95104
branch_name = get_branch_name(branch=branch)
96105
if branch_name not in self._db._schemas:

‎backend/infrahub/types.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -235,6 +235,10 @@ class Number(InfrahubDataType):
235235
infrahub = "Integer"
236236

237237

238+
class NumberPool(Number):
239+
label: str = "Number Pool"
240+
241+
238242
class Bandwidth(InfrahubDataType):
239243
label: str = "Bandwidth"
240244
graphql = graphene.Int
@@ -329,6 +333,7 @@ class Any(InfrahubDataType):
329333
"MacAddress": MacAddress,
330334
"Color": Color,
331335
"Number": Number,
336+
"NumberPool": NumberPool,
332337
"Bandwidth": Bandwidth,
333338
"IPHost": IPHost,
334339
"IPNetwork": IPNetwork,
@@ -353,6 +358,7 @@ class Any(InfrahubDataType):
353358
"MacAddress": str, # MAC addresses can be straightforward strings
354359
"Color": str, # Colors often represented as hex strings
355360
"Number": float, # Numbers can be floats for general use
361+
"NumberPool": float, # Numbers can be floats for general use
356362
"Bandwidth": float, # Bandwidth in some units, represented as a float
357363
"IPHost": IPvAnyAddress, # type: ignore[dict-item]
358364
"IPNetwork": str,

‎backend/templates/attributeschema_imports.j2‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@ from typing import Any
55
from infrahub.core.models import HashableModel
66
from infrahub.core.constants import AllowOverrideType, BranchSupportType, ComputedAttributeKind, HashableModelState
77
from infrahub.core.schema.attribute_parameters import AttributeParameters # noqa: TC001
8+
from infrahub.core.schema.attribute_parameters import NumberPoolParameters # noqa: TC001
9+
from infrahub.core.schema.attribute_parameters import TextAttributeParameters # noqa: TC001
810
from infrahub.core.schema.computed_attribute import ComputedAttribute # noqa: TC001
911
from infrahub.core.schema.dropdown import DropdownChoice # noqa: TC001
1012

0 commit comments

Comments
 (0)