move sqlmodel to pydantic v2 and remove compat code

Author: masenf

pyi_hashes.json

@@ -1,5 +1,5 @@
 {
-  "reflex/__init__.pyi": "cc4f461d8244f0f372b7607eb1edd146",
+  "reflex/__init__.pyi": "2fa0051c43f2d3d10114283480c666fb",
   "reflex/components/__init__.pyi": "ac05995852baa81062ba3d18fbc489fb",
   "reflex/components/base/__init__.pyi": "16e47bf19e0d62835a605baa3d039c5a",
   "reflex/components/base/app_wrap.pyi": "22e94feaa9fe675bcae51c412f5b67f1",

reflex/__init__.py

@@ -323,7 +323,7 @@
         "SessionStorage",
     ],
     "middleware": ["middleware", "Middleware"],
-    "model": ["asession", "session", "Model"],
+    "model": ["asession", "session", "Model", "ModelRegistry"],
     "page": ["page"],
     "state": [
         "var",

reflex/istate/proxy.py

@@ -18,7 +18,7 @@
 from reflex.base import Base
 from reflex.utils import prerequisites
 from reflex.utils.exceptions import ImmutableStateError
-from reflex.utils.serializers import serializer
+from reflex.utils.serializers import can_serialize, serialize, serializer
 from reflex.vars.base import Var
 
 if TYPE_CHECKING:
@@ -689,7 +689,10 @@ def serialize_mutable_proxy(mp: MutableProxy):
     Returns:
         The wrapped object.
     """
-    return mp.__wrapped__
+    obj = mp.__wrapped__
+    if can_serialize(type(obj)):
+        return serialize(obj)
+    return obj
 
 
 _orig_json_encoder_default = json.JSONEncoder.default

reflex/model.py

@@ -10,6 +10,7 @@
 from reflex.environment import environment
 from reflex.utils import console
 from reflex.utils.compat import sqlmodel_field_has_primary_key
+from reflex.utils.serializers import serializer
 
 if TYPE_CHECKING:
     import sqlalchemy
@@ -254,8 +255,6 @@ def get_metadata(cls) -> sqlalchemy.MetaData:
     import sqlmodel
     from alembic.runtime.migration import MigrationContext
     from alembic.script.base import Script
-    from pydantic import ConfigDict
-    from sqlmodel._compat import IS_PYDANTIC_V2
     from sqlmodel.ext.asyncio.session import AsyncSession
 
     _AsyncSessionLocal: dict[str | None, sqlalchemy.ext.asyncio.async_sessionmaker] = {}
@@ -312,91 +311,60 @@ async def get_db_status() -> dict[str, bool]:
 
         return {"db": status}
 
+    @serializer
+    def serialize_sqlmodel(m: sqlmodel.SQLModel) -> dict[str, Any]:
+        """Serialize a SQLModel object to a dictionary.
+
+        Args:
+            m: The SQLModel object to serialize.
+
+        Returns:
+            The serialized object as a dictionary.
+        """
+        base_fields = m.model_dump()
+        relationships = {}
+        # SQLModel relationships do not appear in __fields__, but should be included if present.
+        for name in m.__sqlmodel_relationships__:
+            with suppress(
+                sqlalchemy.orm.exc.DetachedInstanceError  # This happens when the relationship was never loaded and the session is closed.
+            ):
+                relationships[name] = getattr(m, name)
+        return {
+            **base_fields,
+            **relationships,
+        }
+
     class Model(sqlmodel.SQLModel):
         """Base class to define a table in the database."""
 
         # The primary key for the table.
         id: int | None = sqlmodel.Field(default=None, primary_key=True)
 
-        if IS_PYDANTIC_V2:
-            model_config = ConfigDict(  # pyright: ignore [reportAssignmentType]
-                arbitrary_types_allowed=True,
-                extra="allow",
-                use_enum_values=True,
-                from_attributes=True,
-            )
-        else:
-
-            class Config:  # pyright: ignore [reportIncompatibleVariableOverride]
-                """Pydantic V1 config."""
-
-                arbitrary_types_allowed = True
-                use_enum_values = True
-                extra = "allow"
-                orm_mode = True
+        model_config = {  # pyright: ignore [reportAssignmentType]
+            "arbitrary_types_allowed": True,
+            "use_enum_values": True,
+            "extra": "allow",
+        }
 
-        def __init_subclass__(cls):
+        @classmethod
+        def __pydantic_init_subclass__(cls):
             """Drop the default primary key field if any primary key field is defined."""
             non_default_primary_key_fields = [
                 field_name
-                for field_name, field in cls.__fields__.items()
-                if field_name != "id" and sqlmodel_field_has_primary_key(field)
+                for field_name, field_info in cls.model_fields.items()
+                if field_name != "id" and sqlmodel_field_has_primary_key(field_info)
             ]
             if non_default_primary_key_fields:
-                cls.__fields__.pop("id", None)
-
-            super().__init_subclass__()
-
-        @classmethod
-        def _dict_recursive(cls, value: Any):
-            """Recursively serialize the relationship object(s).
-
-            Args:
-                value: The value to serialize.
-
-            Returns:
-                The serialized value.
-            """
-            if hasattr(value, "dict"):
-                return value.dict()
-            if isinstance(value, list):
-                return [cls._dict_recursive(item) for item in value]
-            return value
-
-        def dict(self, **kwargs):
-            """Convert the object to a dictionary.
-
-            Args:
-                kwargs: Ignored but needed for compatibility.
-
-            Returns:
-                The object as a dictionary.
-            """
-            base_fields = {name: getattr(self, name) for name in self.__fields__}
-            relationships = {}
-            # SQLModel relationships do not appear in __fields__, but should be included if present.
-            for name in self.__sqlmodel_relationships__:
-                with suppress(
-                    sqlalchemy.orm.exc.DetachedInstanceError  # This happens when the relationship was never loaded and the session is closed.
-                ):
-                    relationships[name] = self._dict_recursive(getattr(self, name))
-            return {
-                **base_fields,
-                **relationships,
-            }
-
-        def json(self) -> str:
-            """Convert the object to a json string.
-
-            Returns:
-                The object as a json string.
-            """
-            from reflex.utils.serializers import serialize
-
-            return self.__config__.json_dumps(
-                self.dict(),
-                default=serialize,
-            )
+                cls.model_fields.pop("id", None)
+                console.deprecate(
+                    feature_name="Overriding default primary key",
+                    reason=(
+                        "Register sqlmodel.SQLModel classes with `@rx.ModelRegistry.register`"
+                    ),
+                    deprecation_version="0.8.0",
+                    removal_version="0.9.0",
+                )
+            super().__pydantic_init_subclass__()
 
         @staticmethod
         def create_all():

reflex/utils/compat.py

@@ -1,6 +1,9 @@
 """Compatibility hacks and helpers."""
 
-from typing import Any
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pydantic.fields import FieldInfo
 
 
 async def windows_hot_reload_lifespan_hack():
@@ -27,17 +30,17 @@ async def windows_hot_reload_lifespan_hack():
         pass
 
 
-def sqlmodel_field_has_primary_key(field: Any) -> bool:
+def sqlmodel_field_has_primary_key(field_info: "FieldInfo") -> bool:
     """Determines if a field is a primary.
 
     Args:
-        field: a rx.model field
+        field_info: a rx.model field
 
     Returns:
-        If field is a primary key (Bool)
+        If field_info is a primary key (Bool)
     """
-    if getattr(field.field_info, "primary_key", None) is True:
+    if getattr(field_info, "primary_key", None) is True:
         return True
-    if getattr(field.field_info, "sa_column", None) is None:
+    if getattr(field_info, "sa_column", None) is None:
         return False
-    return bool(getattr(field.field_info.sa_column, "primary_key", None))
+    return bool(getattr(field_info.sa_column, "primary_key", None))  # pyright: ignore[reportAttributeAccessIssue]

tests/units/test_model.py

@@ -53,7 +53,7 @@ def test_default_primary_key(model_default_primary: Model):
     Args:
         model_default_primary: Fixture.
     """
-    assert "id" in type(model_default_primary).__fields__
+    assert "id" in type(model_default_primary).model_fields
 
 
 def test_custom_primary_key(model_custom_primary: Model):
@@ -62,7 +62,7 @@ def test_custom_primary_key(model_custom_primary: Model):
     Args:
         model_custom_primary: Fixture.
     """
-    assert "id" not in type(model_custom_primary).__fields__
+    assert "id" not in type(model_custom_primary).model_fields
 
 
 @pytest.mark.filterwarnings(