Use v2 of pydantic API

Author: mattkjames7

gqlalchemy/__init__.py

@@ -14,7 +14,7 @@
 
 import warnings
 
-from pydantic.v1 import validator  # noqa F401
+from pydantic import validator  # noqa F401
 
 from gqlalchemy.models import (  # noqa F401
     MemgraphConstraintExists,

gqlalchemy/connection.py

@@ -119,7 +119,7 @@ def _create_connection(self) -> Connection:
 def _convert_memgraph_value(value: Any) -> Any:
     """Converts Memgraph objects to custom Node/Relationship objects."""
     if isinstance(value, mgclient.Relationship):
-        return Relationship.parse_obj(
+        return Relationship.model_validate(
             {
                 "_type": value.type,
                 "_id": value.id,
@@ -130,7 +130,7 @@ def _convert_memgraph_value(value: Any) -> Any:
         )
 
     if isinstance(value, mgclient.Node):
-        return Node.parse_obj(
+        return Node.model_validate(
             {
                 "_id": value.id,
                 "_labels": set(value.labels),
@@ -139,7 +139,7 @@ def _convert_memgraph_value(value: Any) -> Any:
         )
 
     if isinstance(value, mgclient.Path):
-        return Path.parse_obj(
+        return Path.model_validate(
             {
                 "_nodes": list([_convert_memgraph_value(node) for node in value.nodes]),
                 "_relationships": list([_convert_memgraph_value(rel) for rel in value.relationships]),
@@ -192,7 +192,7 @@ def _create_connection(self):
 def _convert_neo4j_value(value: Any) -> Any:
     """Converts Neo4j objects to custom Node/Relationship objects."""
     if isinstance(value, Neo4jRelationship):
-        return Relationship.parse_obj(
+        return Relationship.model_validate(
             {
                 "_type": value.type,
                 "_id": value.id,
@@ -203,7 +203,7 @@ def _convert_neo4j_value(value: Any) -> Any:
         )
 
     if isinstance(value, Neo4jNode):
-        return Node.parse_obj(
+        return Node.model_validate(
             {
                 "_id": value.id,
                 "_labels": set(value.labels),
@@ -212,7 +212,7 @@ def _convert_neo4j_value(value: Any) -> Any:
         )
 
     if isinstance(value, Neo4jPath):
-        return Path.parse_obj(
+        return Path.model_validate(
             {
                 "_nodes": list([_convert_neo4j_value(node) for node in value.nodes]),
                 "_relationships": list([_convert_neo4j_value(rel) for rel in value.relationships]),

gqlalchemy/models.py

@@ -18,9 +18,9 @@
 from datetime import datetime, date, time, timedelta
 from enum import Enum
 import json
-from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union
+from typing import Any, ClassVar, Dict, Iterable, List, Optional, Set, Tuple, Union
 
-from pydantic.v1 import BaseModel, Extra, Field, PrivateAttr  # noqa F401
+from pydantic import BaseModel, ConfigDict, Field as PydanticField, PrivateAttr  # noqa F401
 
 from gqlalchemy.exceptions import (
     GQLAlchemyError,
@@ -48,6 +48,64 @@ class DatetimeKeywords(Enum):
 }
 
 
+def _get_model_fields(model: Any) -> Dict[str, Any]:
+    model_fields = getattr(model, "model_fields", None)
+    if model_fields is not None:
+        return model_fields
+    return getattr(model, "__fields__", {})
+
+
+def _get_field_attrs(field: Any) -> Dict[str, Any]:
+    attrs = getattr(field, "json_schema_extra", None)
+    if isinstance(attrs, dict):
+        return attrs
+
+    field_info = getattr(field, "field_info", None)
+    if field_info is not None:
+        extra = getattr(field_info, "extra", None)
+        if isinstance(extra, dict):
+            return extra
+        json_schema_extra = getattr(field_info, "json_schema_extra", None)
+        if isinstance(json_schema_extra, dict):
+            return json_schema_extra
+
+    return {}
+
+
+def _set_field_attrs(field: Any, attrs: Dict[str, Any]) -> None:
+    copied_attrs = dict(attrs)
+    if hasattr(field, "json_schema_extra"):
+        field.json_schema_extra = copied_attrs
+        return
+
+    field_info = getattr(field, "field_info", None)
+    if field_info is not None and hasattr(field_info, "extra"):
+        field_info.extra = copied_attrs
+
+
+def _get_field_type_name(field: Any) -> str:
+    if hasattr(field, "type_"):
+        return field.type_.__name__
+
+    annotation = getattr(field, "annotation", None)
+    return getattr(annotation, "__name__", str(annotation))
+
+
+def Field(default=..., **kwargs):  # noqa N802
+    """Pydantic Field wrapper that stores custom OGM metadata in json_schema_extra."""
+    json_schema_extra = kwargs.pop("json_schema_extra", None)
+    extras = dict(json_schema_extra or {})
+
+    for attr in ("index", "exists", "unique", "db", "on_disk", "label"):
+        if attr in kwargs:
+            extras[attr] = kwargs.pop(attr)
+
+    if extras:
+        kwargs["json_schema_extra"] = extras
+
+    return PydanticField(default, **kwargs)
+
+
 def _format_timedelta(duration: timedelta) -> str:
     days = int(duration.total_seconds() // 86400)
     remainder_sec = duration.total_seconds() - days * 86400
@@ -352,10 +410,16 @@ def to_cypher(self) -> str:
 
 
 class GraphObject(BaseModel):
-    _subtypes_: Dict = dict()
+    _subtypes_: ClassVar[Dict[str, Any]] = {}
+    model_config = ConfigDict(extra="allow", coerce_numbers_to_str=True)
 
-    class Config:
-        extra = Extra.allow
+    def __init__(self, **data):
+        # Preserve explicitly provided underscore-prefixed attributes.
+        private_data = {key: value for key, value in data.items() if key.startswith("_")}
+        super().__init__(**data)
+        for key, value in private_data.items():
+            if key in self.__private_attributes__:
+                setattr(self, key, value)
 
     def __init_subclass__(cls, type=None, label=None, labels=None, index=None, db=None):
         """Stores the subclass by type if type is specified, or by class name
@@ -401,12 +465,21 @@ def _convert_to_real_type_(cls, data):
 
         return sub(**data)
 
+    @classmethod
+    def model_validate(cls, obj, *args, **kwargs):
+        """Used to convert a dictionary object into the appropriate GraphObject."""
+        if isinstance(obj, cls):
+            return obj
+        if isinstance(obj, dict):
+            return cls._convert_to_real_type_(obj)
+        return super().model_validate(obj, *args, **kwargs)
+
     @classmethod
     def parse_obj(cls, obj):
         """Used to convert a dictionary object into the appropriate
         GraphObject.
         """
-        return cls._convert_to_real_type_(obj)
+        return cls.model_validate(obj)
 
     def escape_value(
         self, value: Union[None, bool, int, float, str, list, dict, datetime, timedelta, date, time]
@@ -463,7 +536,7 @@ def _get_cypher_field_assignment_block(self, variable_name: str, operator: str)
                 "user.name = 'John' AND user.age = 34"
         """
         cypher_fields = []
-        for field in self.__fields__:
+        for field in _get_model_fields(type(self)):
             value = getattr(self, field)
             if value is not None:
                 cypher_fields.append(f"{variable_name}.{field} = {self.escape_value(value)}")
@@ -493,8 +566,8 @@ def _get_cypher_fields_xor_block(self, variable_name: str) -> str:
     def _get_cypher_set_properties(self, variable_name: str) -> str:
         """Returns a cypher set properties block."""
         cypher_set_properties = []
-        for field in self.__fields__:
-            attributes = self.__fields__[field].field_info.extra
+        for field, field_definition in _get_model_fields(type(self)).items():
+            attributes = _get_field_attrs(field_definition)
             value = getattr(self, field)
             if value is not None and not attributes.get("on_disk", False):
                 cypher_set_properties.append(f" SET {variable_name}.{field} = {self.escape_value(value)}")
@@ -520,7 +593,7 @@ def __init__(self, **data):
 
     @property
     def _properties(self) -> Dict[str, Any]:  # noqa: F811
-        return {k: v for k, v in dict(self).items() if not k.startswith("_") and k != "labels"}
+        return {k: v for k, v in self.model_dump().items() if not k.startswith("_") and k != "labels"}
 
     def __str__(self) -> str:
         return f"<GraphObject id={self._id} properties={self._properties}>"
@@ -539,8 +612,9 @@ def __new__(mcs, name, bases, namespace, **kwargs):  # noqa C901
         def field_in_superclass(field, constraint):
             nonlocal bases
             for base in bases:
-                if field in base.__fields__:
-                    attrs = base.__fields__[field].field_info.extra
+                base_fields = _get_model_fields(base)
+                if field in base_fields:
+                    attrs = _get_field_attrs(base_fields[field])
                     if constraint in attrs:
                         return base
 
@@ -569,9 +643,10 @@ def get_base_labels() -> Set[str]:
             index = MemgraphIndex(cls.label)
             db.create_index(index)
 
-        for field in cls.__fields__:
-            attrs = cls.__fields__[field].field_info.extra
-            field_type = cls.__fields__[field].type_.__name__
+        cls_fields = _get_model_fields(cls)
+        for field, field_definition in cls_fields.items():
+            attrs = _get_field_attrs(field_definition)
+            field_type = _get_field_type_name(field_definition)
             label = attrs.get("label", cls.label)
             skip_constraints = False
 
@@ -582,7 +657,8 @@ def get_base_labels() -> Set[str]:
                 if constraint in attrs and db is None:
                     base = field_in_superclass(field, constraint)
                     if base is not None:
-                        cls.__fields__[field].field_info.extra = base.__fields__[field].field_info.extra
+                        base_fields = _get_model_fields(base)
+                        _set_field_attrs(field_definition, _get_field_attrs(base_fields[field]))
                         skip_constraints = True
                         break
 
@@ -634,8 +710,8 @@ def __str__(self) -> str:
     def _get_cypher_unique_fields_or_block(self, variable_name: str) -> str:
         """Get's a cypher assignment block using the unique fields."""
         cypher_unique_fields = []
-        for field in self.__fields__:
-            attrs = self.__fields__[field].field_info.extra
+        for field, field_definition in _get_model_fields(type(self)).items():
+            attrs = _get_field_attrs(field_definition)
             if "unique" in attrs:
                 value = getattr(self, field)
                 if value is not None:
@@ -645,8 +721,8 @@ def _get_cypher_unique_fields_or_block(self, variable_name: str) -> str:
 
     def has_unique_fields(self) -> bool:
         """Returns True if the Node has any unique fields."""
-        for field in self.__fields__:
-            if "unique" in self.__fields__[field].field_info.extra:
+        for field, field_definition in _get_model_fields(type(self)).items():
+            if "unique" in _get_field_attrs(field_definition):
                 if getattr(self, field) is not None:
                     return True
         return False
@@ -665,7 +741,7 @@ def save(self, db: "Database") -> "Node":  # noqa F821
         Null properties are ignored.
         """
         node = db.save_node(self)
-        for field in self.__fields__:
+        for field in _get_model_fields(type(self)):
             setattr(self, field, getattr(node, field))
         self._id = node._id
         return self
@@ -681,7 +757,7 @@ def load(self, db: "Database") -> "Node":  # noqa F821
         If no node is found or no properties are set it raises a GQLAlchemyError.
         """
         node = db.load_node(self)
-        for field in self.__fields__:
+        for field in _get_model_fields(type(self)):
             setattr(self, field, getattr(node, field))
         self._id = node._id
         return self
@@ -754,7 +830,7 @@ def save(self, db: "Database") -> "Relationship":  # noqa F821
         relationship, use `load_relationship` first.
         """
         relationship = db.save_relationship(self)
-        for field in self.__fields__:
+        for field in _get_model_fields(type(self)):
             setattr(self, field, getattr(relationship, field))
         self._id = relationship._id
         return self
@@ -770,7 +846,7 @@ def load(self, db: "Database") -> "Relationship":  # noqa F821
         multiple relationships like that in Memgraph, throws GQLAlchemyError.
         """
         relationship = db.load_relationship(self)
-        for field in self.__fields__:
+        for field in _get_model_fields(type(self)):
             setattr(self, field, getattr(relationship, field))
         self._id = relationship._id
         return self

gqlalchemy/vendors/memgraph.py

@@ -35,6 +35,8 @@
     MemgraphTrigger,
     Node,
     Relationship,
+    _get_field_attrs,
+    _get_model_fields,
 )
 from gqlalchemy.vendors.database_client import DatabaseClient
 from gqlalchemy.graph_algorithms.query_modules import QueryModule
@@ -346,9 +348,9 @@ def _save_node_properties_on_disk(self, node: Node, result: Node) -> Node:
         """Saves all on_disk properties to the on disk database attached to
         the database.
         """
-        for field in node.__fields__:
+        for field, field_definition in _get_model_fields(type(node)).items():
             value = getattr(node, field, None)
-            if value is not None and "on_disk" in node.__fields__[field].field_info.extra:
+            if value is not None and "on_disk" in _get_field_attrs(field_definition):
                 if self.on_disk_db is None:
                     raise GQLAlchemyOnDiskPropertyDatabaseNotDefinedError()
                 self.on_disk_db.save_node_property(result._id, field, value)
@@ -381,9 +383,9 @@ def load_node(self, node: Node) -> Optional[Node]:
 
     def _load_node_properties_on_disk(self, result: Node) -> Node:
         """Loads all on_disk properties from the on disk database."""
-        for field in result.__fields__:
+        for field, field_definition in _get_model_fields(type(result)).items():
             value = getattr(result, field, None)
-            if "on_disk" in result.__fields__[field].field_info.extra:
+            if "on_disk" in _get_field_attrs(field_definition):
                 if self.on_disk_db is None:
                     raise GQLAlchemyOnDiskPropertyDatabaseNotDefinedError()
                 try:
@@ -420,9 +422,9 @@ def _load_relationship_properties_on_disk(self, result: Relationship) -> Relatio
         Memgraph().init_disk_storage() throws a
         GQLAlchemyOnDiskPropertyDatabaseNotDefinedError.
         """
-        for field in result.__fields__:
+        for field, field_definition in _get_model_fields(type(result)).items():
             value = getattr(result, field, None)
-            if "on_disk" in result.__fields__[field].field_info.extra:
+            if "on_disk" in _get_field_attrs(field_definition):
                 if self.on_disk_db is None:
                     raise GQLAlchemyOnDiskPropertyDatabaseNotDefinedError()
                 try:
@@ -456,9 +458,9 @@ def _save_relationship_properties_on_disk(self, relationship: Relationship, resu
         added with Memgraph().init_disk_storage(db). If OnDiskPropertyDatabase
         is not defined raises GQLAlchemyOnDiskPropertyDatabaseNotDefinedError.
         """
-        for field in relationship.__fields__:
+        for field, field_definition in _get_model_fields(type(relationship)).items():
             value = getattr(relationship, field, None)
-            if value is not None and "on_disk" in relationship.__fields__[field].field_info.extra:
+            if value is not None and "on_disk" in _get_field_attrs(field_definition):
                 if self.on_disk_db is None:
                     raise GQLAlchemyOnDiskPropertyDatabaseNotDefinedError()
                 self.on_disk_db.save_relationship_property(result._id, field, value)