infrahub-server 1.2.11__py3-none-any.whl → 1.3.0b1__py3-none-any.whl

infrahub/core/query/relationship.py

@@ -217,8 +217,7 @@ class RelationshipQuery(Query):
         )
         source_query_match = """
             MATCH (s:Node { uuid: $source_id })
-            CALL {
-                WITH s
+            CALL (s) {
                 MATCH (s)-[r:IS_PART_OF]->(:Root)
                 WHERE %(source_filter)s
                 RETURN r.status = "active" AS s_is_active
@@ -246,8 +245,7 @@ class RelationshipQuery(Query):
             )
             destination_query_match = """
             MATCH (d:Node { uuid: $destination_id })
-            CALL {
-                WITH d
+            CALL (d) {
                 MATCH (d)-[r:IS_PART_OF]->(:Root)
                 WHERE %(destination_filter)s
                 RETURN r.status = "active" AS d_is_active
@@ -524,8 +522,7 @@ class RelationshipDeleteQuery(RelationshipQuery):
         CREATE (s)%(r1)s(rl)
         CREATE (rl)%(r2)s(d)
         WITH rl
-        CALL {
-            WITH rl
+        CALL (rl) {
             MATCH (rl)-[edge:IS_VISIBLE]->(visible)
             WHERE %(rel_filter)s AND edge.status = "active"
             WITH rl, edge, visible
@@ -536,8 +533,7 @@ class RelationshipDeleteQuery(RelationshipQuery):
             WHERE edge.branch = $branch
             SET edge.to = $at
         }
-        CALL {
-            WITH rl
+        CALL (rl) {
             MATCH (rl)-[edge:IS_PROTECTED]->(protected)
             WHERE %(rel_filter)s AND edge.status = "active"
             WITH rl, edge, protected
@@ -548,8 +544,7 @@ class RelationshipDeleteQuery(RelationshipQuery):
             WHERE edge.branch = $branch
             SET edge.to = $at
         }
-        CALL {
-            WITH rl
+        CALL (rl) {
             MATCH (rl)-[edge:HAS_OWNER]->(owner_node)
             WHERE %(rel_filter)s AND edge.status = "active"
             WITH rl, edge, owner_node
@@ -560,8 +555,7 @@ class RelationshipDeleteQuery(RelationshipQuery):
             WHERE edge.branch = $branch
             SET edge.to = $at
         }
-        CALL {
-            WITH rl
+        CALL (rl) {
             MATCH (rl)-[edge:HAS_SOURCE]->(source_node)
             WHERE %(rel_filter)s AND edge.status = "active"
             WITH rl, edge, source_node
@@ -656,8 +650,7 @@ class RelationshipGetPeerQuery(Query):
         MATCH (source_node:Node)%(arrow_left_start)s[:IS_RELATED]%(arrow_left_end)s(rl:Relationship { name: $rel_identifier })
         WHERE source_node.uuid IN $source_ids
         WITH DISTINCT source_node, rl
-        CALL {
-            WITH rl, source_node
+        CALL (rl, source_node) {
             MATCH path = (source_node)%(path)s(peer:Node)
             WHERE
                 $source_kind IN LABELS(source_node) AND
@@ -726,22 +719,19 @@ class RelationshipGetPeerQuery(Query):
             with_str = ", ".join(
                 [f"{subquery_result_name} as {label}" if label == "peer" else label for label in self.return_labels]
             )
-            self.add_subquery(subquery=subquery, with_clause=with_str)
-
+            self.add_subquery(subquery=subquery, node_alias="peer", with_clause=with_str)
         # ----------------------------------------------------------------------------
         # QUERY Properties
         # ----------------------------------------------------------------------------
         query = """
-        CALL {
-            WITH rl
+        CALL (rl) {
             MATCH (rl)-[r:IS_VISIBLE]-(is_visible)
             WHERE %(branch_filter)s
             RETURN r AS rel_is_visible, is_visible
             ORDER BY r.branch_level DESC, r.from DESC, r.status ASC
             LIMIT 1
         }
-        CALL {
-            WITH rl
+        CALL (rl) {
             MATCH (rl)-[r:IS_PROTECTED]-(is_protected)
             WHERE %(branch_filter)s
             RETURN r AS rel_is_protected, is_protected
@@ -758,8 +748,7 @@ class RelationshipGetPeerQuery(Query):
         # We must query them one by one otherwise the second one won't return
         for node_prop in ["source", "owner"]:
             query = """
-        CALL {
-            WITH rl
+        CALL (rl) {
             OPTIONAL MATCH (rl)-[r:HAS_%(node_prop_type)s]-(%(node_prop)s)
             WHERE %(branch_filter)s
             RETURN r AS rel_%(node_prop)s, %(node_prop)s
@@ -800,7 +789,7 @@ class RelationshipGetPeerQuery(Query):
                 self.order_by.append(subquery_result_name)
                 self.params.update(subquery_params)
 
-                self.add_subquery(subquery=subquery)
+                self.add_subquery(subquery=subquery, node_alias="peer")
 
                 order_cnt += 1
 
@@ -939,8 +928,7 @@ class RelationshipGetByIdentifierQuery(Query):
         query = """
         MATCH (rl:Relationship)
         WHERE rl.name IN $identifiers
-        CALL {
-            WITH rl
+        CALL (rl) {
             MATCH (src:Node)-[r1:IS_RELATED]-(rl:Relationship)-[r2:IS_RELATED]-(dst:Node)
             WHERE (size($full_identifiers) = 0 OR [src.kind, rl.name, dst.kind] in $full_identifiers)
             AND NOT src.namespace IN $excluded_namespaces
@@ -1003,8 +991,7 @@ class RelationshipCountPerNodeQuery(Query):
         query = """
         MATCH (peer_node:Node)%(path)s(rl:Relationship { name: $rel_identifier })
         WHERE peer_node.uuid IN $peer_ids AND %(branch_filter)s
-        CALL {
-            WITH rl
+        CALL (rl) {
             MATCH path = (peer_node:Node)%(path)s(rl)
             WHERE peer_node.uuid IN $peer_ids AND %(branch_filter)s
             RETURN peer_node as peer, r as r1
@@ -1082,8 +1069,7 @@ class RelationshipDeleteAllQuery(Query):
         for arrow_left, arrow_right in (("<-", "-"), ("-", "->")):
             for edge_type in edge_types:
                 sub_query = """
-                    CALL {
-                        WITH rl
+                    CALL (rl) {
                         MATCH (rl)%(arrow_left)s[active_edge:%(edge_type)s]%(arrow_right)s(n)
                         WHERE %(active_rel_filter)s AND active_edge.status ="active"
                         CREATE (rl)%(arrow_left)s[deleted_edge:%(edge_type)s $rel_prop]%(arrow_right)s(n)
@@ -1103,8 +1089,7 @@ class RelationshipDeleteAllQuery(Query):
 
         # We only want to return uuid/kind of `Node` connected through `IS_RELATED` edges.
         query += """
-        CALL {
-            WITH rl
+        CALL (rl) {
             MATCH (rl)-[active_edge:IS_RELATED]->(n)
             WHERE %(active_rel_filter)s
             WITH rl, active_edge, n
@@ -1123,7 +1108,6 @@ class RelationshipDeleteAllQuery(Query):
                 "outbound" as rel_direction
 
             UNION
-
             WITH rl
             MATCH (rl)<-[active_edge:IS_RELATED]-(n)
             WHERE %(active_rel_filter)s
@@ -1193,3 +1177,40 @@ class RelationshipDeleteAllQuery(Query):
             changelog_mapper.delete_relationship(peer_id=peer_uuid, peer_kind=kind, rel_schema=deleted_rel_schema)
 
         return [changelog_mapper.changelog for changelog_mapper in rel_identifier_to_changelog_mapper.values()]
+
+
+class GetAllPeersIds(Query):
+    """
+    Return all peers ids connected to input node. Some peers can be excluded using `exclude_identifiers`.
+    """
+
+    name = "get_peers_ids"
+    type: QueryType = QueryType.READ
+    insert_return = False
+
+    def __init__(self, node_id: str, exclude_identifiers: list[str], **kwargs):
+        self.node_id = node_id
+        self.exclude_identifiers = exclude_identifiers
+        super().__init__(**kwargs)
+
+    async def query_init(self, db: InfrahubDatabase, **kwargs) -> None:  # noqa: ARG002
+        self.params["source_id"] = kwargs["node_id"]
+        self.params["branch"] = self.branch.name
+        self.params["exclude_identifiers"] = self.exclude_identifiers
+
+        active_rel_filter, rel_params = self.branch.get_query_filter_path(
+            at=self.at, variable_name="e1", branch_agnostic=self.branch_agnostic
+        )
+        self.params.update(rel_params)
+
+        query = """
+            MATCH (node:Node { uuid: $source_id })-[e1:IS_RELATED]-(rl:Relationship)-[e2:IS_RELATED]-(peer:Node)
+            WHERE %(active_rel_filter)s AND peer.uuid <> node.uuid AND NOT (rl.name IN $exclude_identifiers)
+            WITH DISTINCT(peer.uuid) as uuid
+            RETURN uuid
+        """ % {"active_rel_filter": active_rel_filter}
+
+        self.add_to_query(query)
+
+    def get_peers_uuids(self) -> list[str]:
+        return [row.data["uuid"] for row in self.results]  # type: ignore

infrahub/core/query/resource_manager.py

@@ -220,8 +220,7 @@ class NumberPoolGetUsed(Query):
 
         query = """
         MATCH (pool:%(number_pool)s { uuid: $pool_id })
-        CALL {
-            WITH pool
+        CALL (pool) {
             MATCH (pool)-[res:IS_RESERVED]->(av:AttributeValue)<-[hv:HAS_VALUE]-(attr:Attribute)
             WHERE
                 attr.name = $attribute_name

infrahub/core/query/subquery.py

@@ -61,7 +61,7 @@ async def build_subquery_filter(
     where_str = " AND ".join(field_where)
     branch_level_str = "reduce(br_lvl = 0, r in relationships(path) | br_lvl + r.branch_level)"
     froms_str = db.render_list_comprehension(items="relationships(path)", item_name="from")
-    to_return = f"{node_alias} AS {prefix}"
+    to_return = f"{prefix}"
     with_extra = ""
     final_with_extra = ""
     is_isnull = filter_name == "isnull"
@@ -82,7 +82,6 @@ async def build_subquery_filter(
         elif field is not None and field.is_attribute:
             is_active_filter = "(latest_node_details[2]).value = 'NULL'"
     query = f"""
-    WITH {node_alias}
     {match} path = {filter_str}
     WHERE {where_str}
     WITH
@@ -94,7 +93,7 @@ async def build_subquery_filter(
     ORDER BY branch_level DESC, froms[-1] DESC, froms[-2] DESC
     WITH head(collect([is_active, {node_alias}{with_extra}])) AS latest_node_details
     WHERE {is_active_filter}
-    WITH latest_node_details[1] AS {node_alias}{final_with_extra}
+    WITH latest_node_details[1] AS {prefix}{final_with_extra}
     RETURN {to_return}
     """
     return query, params, prefix
@@ -174,7 +173,6 @@ async def build_subquery_order(
         to_return_str_parts.append(f"CASE WHEN is_active = TRUE THEN {expression} ELSE NULL END AS {alias}")
     to_return_str = ", ".join(to_return_str_parts)
     query = f"""
-    WITH {node_alias}
     OPTIONAL MATCH path = {filter_str}
     WHERE {where_str}
     WITH {with_str_to_alias}

infrahub/core/relationship/model.py

@@ -789,6 +789,9 @@ class RelationshipManager:
 
         return len(self._relationships)
 
+    def validate(self) -> None:
+        self._relationships.validate()
+
     @overload
     async def get_peer(
         self,

infrahub/core/schema/__init__.py

@@ -21,7 +21,8 @@ from .profile_schema import ProfileSchema
 from .relationship_schema import RelationshipSchema
 from .template_schema import TemplateSchema
 
-MainSchemaTypes: TypeAlias = NodeSchema | GenericSchema | ProfileSchema | TemplateSchema
+NonGenericSchemaTypes: TypeAlias = NodeSchema | ProfileSchema | TemplateSchema
+MainSchemaTypes: TypeAlias = NonGenericSchemaTypes | GenericSchema
 
 
 # -----------------------------------------------------

infrahub/core/schema/attribute_parameters.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+import sys
+from typing import Self
+
+from pydantic import ConfigDict, Field, model_validator
+
+from infrahub import config
+from infrahub.core.constants.schema import UpdateSupport
+from infrahub.core.models import HashableModel
+
+
+def get_attribute_parameters_class_for_kind(kind: str) -> type[AttributeParameters]:
+    param_classes: dict[str, type[AttributeParameters]] = {
+        "NumberPool": NumberPoolParameters,
+        "Text": TextAttributeParameters,
+        "TextArea": TextAttributeParameters,
+        "Number": NumberAttributeParameters,
+    }
+    return param_classes.get(kind, AttributeParameters)
+
+
+class AttributeParameters(HashableModel):
+    model_config = ConfigDict(extra="forbid")
+
+
+class TextAttributeParameters(AttributeParameters):
+    regex: str | None = Field(
+        default=None,
+        description="Regular expression that attribute value must match if defined",
+        json_schema_extra={"update": UpdateSupport.VALIDATE_CONSTRAINT.value},
+    )
+    min_length: int | None = Field(
+        default=None,
+        description="Set a minimum number of characters allowed.",
+        json_schema_extra={"update": UpdateSupport.VALIDATE_CONSTRAINT.value},
+    )
+    max_length: int | None = Field(
+        default=None,
+        description="Set a maximum number of characters allowed.",
+        json_schema_extra={"update": UpdateSupport.VALIDATE_CONSTRAINT.value},
+    )
+
+    @model_validator(mode="after")
+    def validate_min_max(self) -> Self:
+        if (
+            config.SETTINGS.initialized
+            and config.SETTINGS.main.schema_strict_mode
+            and self.min_length is not None
+            and self.max_length is not None
+        ):
+            if self.min_length > self.max_length:
+                raise ValueError(
+                    "`max_length` can't be less than `min_length` when the schema is configured with strict mode"
+                )
+
+        return self
+
+
+class NumberAttributeParameters(AttributeParameters):
+    min_value: int | None = Field(
+        default=None,
+        description="Set a minimum value allowed.",
+        json_schema_extra={"update": UpdateSupport.VALIDATE_CONSTRAINT.value},
+    )
+    max_value: int | None = Field(
+        default=None,
+        description="Set a maximum value allowed.",
+        json_schema_extra={"update": UpdateSupport.VALIDATE_CONSTRAINT.value},
+    )
+    excluded_values: str | None = Field(
+        default=None,
+        description="List of values or range of values not allowed for the attribute, format is: '100,150-200,280,300-400'",
+        pattern=r"^(\d+(?:-\d+)?)(?:,\d+(?:-\d+)?)*$",
+        json_schema_extra={"update": UpdateSupport.VALIDATE_CONSTRAINT.value},
+    )
+
+    @model_validator(mode="after")
+    def validate_ranges(self) -> Self:
+        ranges = self.get_excluded_ranges()
+        for i, (start_range_1, end_range_1) in enumerate(ranges):
+            if start_range_1 > end_range_1:
+                raise ValueError("`start_range` can't be less than `end_range`")
+
+            # Check for overlapping ranges
+            for start_range_2, end_range_2 in ranges[i + 1 :]:
+                if not (end_range_1 < start_range_2 or start_range_1 > end_range_2):
+                    raise ValueError("Excluded ranges cannot overlap")
+
+        return self
+
+    @model_validator(mode="after")
+    def validate_min_max(self) -> Self:
+        if (
+            config.SETTINGS.initialized
+            and config.SETTINGS.main.schema_strict_mode
+            and self.min_value is not None
+            and self.max_value is not None
+        ):
+            if self.min_value > self.max_value:
+                raise ValueError(
+                    "`max_value` can't be less than `min_value` when the schema is configured with strict mode"
+                )
+
+        return self
+
+    def get_excluded_single_values(self) -> list[int]:
+        if not self.excluded_values:
+            return []
+
+        results = [int(value) for value in self.excluded_values.split(",") if "-" not in value]
+        return results
+
+    def get_excluded_ranges(self) -> list[tuple[int, int]]:
+        if not self.excluded_values:
+            return []
+
+        ranges = []
+        for value in self.excluded_values.split(","):
+            if "-" in value:
+                start, end = map(int, value.split("-"))
+                ranges.append((start, end))
+
+        return ranges
+
+    def is_valid_value(self, value: int) -> bool:
+        if self.min_value is not None and value < self.min_value:
+            return False
+        if self.max_value is not None and value > self.max_value:
+            return False
+        if value in self.get_excluded_single_values():
+            return False
+        for start, end in self.get_excluded_ranges():
+            if start <= value <= end:
+                return False
+        return True
+
+
+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

infrahub/core/schema/attribute_schema.py

@@ -2,15 +2,23 @@ from __future__ import annotations
 
 import enum
 from enum import Enum
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Self
 
-from pydantic import field_validator, model_validator
+from pydantic import Field, ValidationInfo, field_validator, model_validator
 
 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.types import ATTRIBUTE_KIND_LABELS, ATTRIBUTE_TYPES
 
+from .attribute_parameters import (
+    AttributeParameters,
+    NumberAttributeParameters,
+    NumberPoolParameters,
+    TextAttributeParameters,
+    get_attribute_parameters_class_for_kind,
+)
 from .generated.attribute_schema import GeneratedAttributeSchema
 
 if TYPE_CHECKING:
@@ -21,6 +29,16 @@ if TYPE_CHECKING:
     from infrahub.database import InfrahubDatabase
 
 
+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,
+        "Number": NumberAttributeSchema,
+    }
+    return attribute_schema_class_by_kind.get(kind, AttributeSchema)
+
+
 class AttributeSchema(GeneratedAttributeSchema):
     _sort_by: list[str] = ["name"]
     _enum_class: type[enum.Enum] | None = None
@@ -53,16 +71,46 @@ class AttributeSchema(GeneratedAttributeSchema):
 
     @model_validator(mode="before")
     @classmethod
-    def validate_dropdown_choices(cls, values: dict[str, Any]) -> dict[str, Any]:
+    def validate_dropdown_choices(cls, values: Any) -> Any:
         """Validate that choices are defined for a dropdown but not for other kinds."""
-        if values.get("kind") != "Dropdown" and values.get("choices"):
-            raise ValueError(f"Can only specify 'choices' for kind=Dropdown: {values['kind']}")
-
-        if values.get("kind") == "Dropdown" and not values.get("choices"):
+        if isinstance(values, dict):
+            kind = values.get("kind")
+            choices = values.get("choices")
+        elif isinstance(values, AttributeSchema):
+            kind = values.kind
+            choices = values.choices
+        else:
+            return values
+        if kind != "Dropdown" and choices:
+            raise ValueError(f"Can only specify 'choices' for kind=Dropdown: {kind}")
+
+        if kind == "Dropdown" and not choices:
             raise ValueError("The property 'choices' is required for kind=Dropdown")
 
         return values
 
+    @field_validator("parameters", mode="before")
+    @classmethod
+    def set_parameters_type(cls, value: Any, info: ValidationInfo) -> Any:
+        """Override parameters class if using base AttributeParameters class and should be using a subclass"""
+        kind = info.data["kind"]
+        expected_parameters_class = get_attribute_parameters_class_for_kind(kind=kind)
+        if value is None:
+            return expected_parameters_class()
+        if not isinstance(value, expected_parameters_class) and isinstance(value, AttributeParameters):
+            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()
 
@@ -106,7 +154,7 @@ class AttributeSchema(GeneratedAttributeSchema):
 
     def to_node(self) -> dict[str, Any]:
         fields_to_exclude = {"id", "state", "filters"}
-        fields_to_json = {"computed_attribute"}
+        fields_to_json = {"computed_attribute", "parameters"}
         data = self.model_dump(exclude=fields_to_exclude | fields_to_json)
 
         for field_name in fields_to_json:
@@ -117,6 +165,15 @@ class AttributeSchema(GeneratedAttributeSchema):
 
         return data
 
+    def get_regex(self) -> str | None:
+        return self.regex
+
+    def get_min_length(self) -> int | None:
+        return self.min_length
+
+    def get_max_length(self) -> int | None:
+        return self.max_length
+
     async def get_query_filter(
         self,
         name: str,
@@ -144,3 +201,49 @@ class AttributeSchema(GeneratedAttributeSchema):
             partial_match=partial_match,
             support_profiles=support_profiles,
         )
+
+
+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,
+        description="Extra parameters specific to text attributes",
+        json_schema_extra={"update": UpdateSupport.VALIDATE_CONSTRAINT.value},
+    )
+
+    @model_validator(mode="after")
+    def reconcile_parameters(self) -> Self:
+        if self.regex != self.parameters.regex:
+            final_regex = self.parameters.regex if self.parameters.regex is not None else self.regex
+            self.regex = self.parameters.regex = final_regex
+        if self.min_length != self.parameters.min_length:
+            final_min_length = self.parameters.min_length if self.parameters.min_length is not None else self.min_length
+            self.min_length = self.parameters.min_length = final_min_length
+        if self.max_length != self.parameters.max_length:
+            final_max_length = self.parameters.max_length if self.parameters.max_length is not None else self.max_length
+            self.max_length = self.parameters.max_length = final_max_length
+        return self
+
+    def get_regex(self) -> str | None:
+        return self.parameters.regex
+
+    def get_min_length(self) -> int | None:
+        return self.parameters.min_length
+
+    def get_max_length(self) -> int | None:
+        return self.parameters.max_length
+
+
+class NumberAttributeSchema(AttributeSchema):
+    parameters: NumberAttributeParameters = Field(
+        default_factory=NumberAttributeParameters,
+        description="Extra parameters specific to number attributes",
+        json_schema_extra={"update": UpdateSupport.VALIDATE_CONSTRAINT.value},
+    )

infrahub/core/schema/basenode_schema.py

@@ -13,7 +13,7 @@ from pydantic import field_validator
 from infrahub.core.constants import RelationshipCardinality, RelationshipKind
 from infrahub.core.models import HashableModel, HashableModelDiff
 
-from .attribute_schema import AttributeSchema
+from .attribute_schema import AttributeSchema, get_attribute_schema_class_for_kind
 from .generated.base_node_schema import GeneratedBaseNodeSchema
 from .relationship_schema import RelationshipSchema
 
@@ -74,6 +74,30 @@ class BaseNodeSchema(GeneratedBaseNodeSchema):
         Be careful hash generated from hash() have a salt by default and they will not be the same across run"""
         return hash(self.get_hash())
 
+    @field_validator("attributes", mode="before")
+    @classmethod
+    def set_attribute_type(cls, raw_attributes: Any) -> Any:
+        if not isinstance(raw_attributes, list):
+            return raw_attributes
+        attribute_schemas_with_types: list[Any] = []
+        for raw_attr in raw_attributes:
+            if not isinstance(raw_attr, (dict, AttributeSchema)):
+                attribute_schemas_with_types.append(raw_attr)
+                continue
+            if isinstance(raw_attr, dict):
+                kind = raw_attr.get("kind")
+                attribute_type_class = get_attribute_schema_class_for_kind(kind=kind)
+                attribute_schemas_with_types.append(attribute_type_class(**raw_attr))
+                continue
+
+            expected_attr_schema_class = get_attribute_schema_class_for_kind(kind=raw_attr.kind)
+            if not isinstance(raw_attr, expected_attr_schema_class):
+                final_attr = expected_attr_schema_class(**raw_attr.model_dump())
+            else:
+                final_attr = raw_attr
+            attribute_schemas_with_types.append(final_attr)
+        return attribute_schemas_with_types
+
     def to_dict(self) -> dict:
         data = self.model_dump(
             exclude_unset=True, exclude_none=True, exclude_defaults=True, exclude={"attributes", "relationships"}