Query result enum normalization: apply convert_input_enums_for_model to node and relationship properties during rehydration, test coverage for extended enum round-trip persistence

Author: FanaticPythoner

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "kuzualchemy"
-version = "0.3.27"
+version = "0.3.28"
 description = "SQLAlchemy-like ORM for Kuzu graph database"
 readme = "README.md"
 license = { file = "LICENSE" }

src/kuzualchemy/kuzu_query.py

@@ -27,6 +27,7 @@
 )
 from .kuzu_query_builder import QueryState, JoinClause, CypherQueryBuilder
 from .kuzu_query_fields import QueryField, ModelFieldAccessor
+from .enum_normalization import convert_input_enums_for_model
 from .uuid_normalization import _NULL_UUID
 
 logger = logging.getLogger(__name__)
@@ -1102,6 +1103,7 @@ def _assert_uuid_seq(seq_value: Any, field_name: str) -> None:
                         _assert_uuid_seq(v, k)
 
                     clean_props[k] = v
+                clean_props = convert_input_enums_for_model(model_class=node_cls, values=clean_props)
 
                 # 6) Use model_construct for fast instantiation (skip validation - data from DB is valid)
                 inst = node_cls.model_construct(**clean_props)
@@ -1178,6 +1180,7 @@ def _assert_uuid_seq(seq_value: Any, field_name: str) -> None:
                     rel_props.pop(DDLConstants.REL_FROM_NODE_FIELD)
                 if DDLConstants.REL_TO_NODE_FIELD in rel_props:
                     rel_props.pop(DDLConstants.REL_TO_NODE_FIELD)
+                rel_props = convert_input_enums_for_model(model_class=result_model_class, values=rel_props)
 
                 # Construct relationship instance using model_construct for speed (data from DB is valid)
                 instance = result_model_class.model_construct(from_node=from_node, to_node=to_node, **rel_props)

tests/test_basemodel_enum_conversion.py

@@ -9,11 +9,21 @@
 from __future__ import annotations
 
 import pytest
+from pathlib import Path
 from enum import Enum, IntEnum, StrEnum
 from typing import Optional, Union, List, Tuple
-from pydantic import ValidationError
 
-from kuzualchemy import BaseModel, kuzu_node, kuzu_field, KuzuDataType
+from kuzualchemy import (
+    BaseModel,
+    KuzuDataType,
+    KuzuSession,
+    KuzuRelationshipBase,
+    kuzu_field,
+    kuzu_int8enum,
+    kuzu_node,
+    kuzu_relationship,
+)
+from kuzualchemy.test_utilities import initialize_schema
 
 
 class StatusEnum(Enum):
@@ -52,6 +62,15 @@ class StatusStrEnum(StrEnum):
     PENDING = "pending"
 
 
+class CanonicalStageEnum(IntEnum):
+    BASELINE = 1
+
+
+@kuzu_int8enum(base_enum=CanonicalStageEnum)
+class ExtendedStageEnum:
+    EXPERIMENTAL = 9
+
+
 class PrimaryUnionEnum(StrEnum):
     """First enum candidate in a multi-enum Union field."""
     FIRST = "first"
@@ -411,6 +430,56 @@ def test_validator_error_message_formatting(self):
         assert "Valid names:" in error_message
         assert "valid values:" in error_message
 
+    def test_node_and_relationship_query_round_trip_rehydrates_extended_enums(self, tmp_path: Path):
+        @kuzu_node("EnumRoundTripNode")
+        class EnumRoundTripNode(BaseModel):
+            id: int = kuzu_field(kuzu_type=KuzuDataType.INT32, primary_key=True)
+            stage: ExtendedStageEnum = kuzu_field(kuzu_type=KuzuDataType.INT8)
+
+        @kuzu_relationship("ENUM_ROUND_TRIP", pairs=[(EnumRoundTripNode, EnumRoundTripNode)])
+        class EnumRoundTripRelationship(KuzuRelationshipBase):
+            stage: ExtendedStageEnum = kuzu_field(kuzu_type=KuzuDataType.INT8)
+
+        db_path = tmp_path / "enum_round_trip.kuzu"
+        session = KuzuSession(db_path=db_path)
+
+        from kuzualchemy.kuzu_orm import get_ddl_for_node, get_ddl_for_relationship
+
+        ddl = "\n".join(
+            [
+                get_ddl_for_node(EnumRoundTripNode),
+                get_ddl_for_relationship(EnumRoundTripRelationship),
+            ]
+        )
+        initialize_schema(session, ddl=ddl)
+
+        source = EnumRoundTripNode(id=1, stage=CanonicalStageEnum.BASELINE)
+        target = EnumRoundTripNode(id=2, stage=9)
+        session.add(source)
+        session.add(target)
+        session.commit()
+
+        source_db = session.query(EnumRoundTripNode).filter_by(id=1).first()
+        target_db = session.query(EnumRoundTripNode).filter_by(id=2).first()
+
+        assert source_db is not None
+        assert target_db is not None
+        assert source_db.stage is ExtendedStageEnum.BASELINE
+        assert target_db.stage is ExtendedStageEnum.EXPERIMENTAL
+
+        session.create_relationship(
+            EnumRoundTripRelationship,
+            source_db.id,
+            target_db.id,
+            stage=9,
+        )
+        session.commit()
+
+        relationships = session.query(EnumRoundTripRelationship).all()
+
+        assert len(relationships) == 1
+        assert relationships[0].stage is ExtendedStageEnum.EXPERIMENTAL
+
 
 # Additional tests for list/tuple of Enum conversions
 from typing import List, Tuple