From d7abcd45de6f44254959dc6fba085a11d965bf62 Mon Sep 17 00:00:00 2001
From: Masen Furer <m_github@0x26.net>
Date: Thu, 11 Apr 2024 13:42:30 -0700
Subject: [PATCH] Force pydantic v1 for sqlmodel compatibility (#3026)

---
 .github/workflows/unit_tests.yml      |  6 ++++
 reflex/compiler/utils.py              |  2 +-
 reflex/model.py                       |  2 +-
 reflex/utils/compat.py                | 43 +++++++++++++++++++++++++++
 reflex/utils/types.py                 |  2 +-
 tests/components/core/test_foreach.py |  7 ++++-
 6 files changed, 58 insertions(+), 4 deletions(-)
 create mode 100644 reflex/utils/compat.py

diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml
index ba0f47caa..528dd3dff 100644
--- a/.github/workflows/unit_tests.yml
+++ b/.github/workflows/unit_tests.yml
@@ -76,4 +76,10 @@ jobs:
           export PYTHONUNBUFFERED=1
           export REDIS_URL=redis://localhost:6379
           poetry run pytest tests --cov --no-cov-on-fail --cov-report=
+      # Change to explicitly install v1 when reflex-hosting-cli is compatible with v2
+      - name: Run unit tests w/ pydantic v2
+        run: |
+          export PYTHONUNBUFFERED=1
+          poetry run pip install "pydantic>2"
+          poetry run pytest tests --cov --no-cov-on-fail --cov-report=
       - run: poetry run coverage html
diff --git a/reflex/compiler/utils.py b/reflex/compiler/utils.py
index ab5bf9650..14d7d4d36 100644
--- a/reflex/compiler/utils.py
+++ b/reflex/compiler/utils.py
@@ -11,7 +11,7 @@ try:
     # reflex-hosting-cli tools are compatible with pydantic v2
 
     if not TYPE_CHECKING:
-        import pydantic.v1.fields as ModelField
+        from pydantic.v1.fields import ModelField
     else:
         raise ModuleNotFoundError
 except ModuleNotFoundError:
diff --git a/reflex/model.py b/reflex/model.py
index 40dbc212d..dcc971707 100644
--- a/reflex/model.py
+++ b/reflex/model.py
@@ -16,12 +16,12 @@ import alembic.script
 import alembic.util
 import sqlalchemy
 import sqlalchemy.orm
-import sqlmodel
 
 from reflex import constants
 from reflex.base import Base
 from reflex.config import get_config
 from reflex.utils import console
+from reflex.utils.compat import sqlmodel
 
 
 def get_engine(url: str | None = None):
diff --git a/reflex/utils/compat.py b/reflex/utils/compat.py
new file mode 100644
index 000000000..0b5ad3ad8
--- /dev/null
+++ b/reflex/utils/compat.py
@@ -0,0 +1,43 @@
+"""Compatibility hacks and helpers."""
+
+import contextlib
+import sys
+
+
+@contextlib.contextmanager
+def pydantic_v1_patch():
+    """A context manager that patches the Pydantic module to mimic v1 behaviour.
+
+    Yields:
+        None when the Pydantic module is patched.
+    """
+    patched_modules = [
+        "pydantic",
+        "pydantic.fields",
+        "pydantic.errors",
+        "pydantic.main",
+    ]
+    originals = {module: sys.modules.get(module) for module in patched_modules}
+    try:
+        import pydantic.v1  # type: ignore
+
+        sys.modules["pydantic.fields"] = pydantic.v1.fields  # type: ignore
+        sys.modules["pydantic.main"] = pydantic.v1.main  # type: ignore
+        sys.modules["pydantic.errors"] = pydantic.v1.errors  # type: ignore
+        sys.modules["pydantic"] = pydantic.v1
+        yield
+    except (ImportError, AttributeError):
+        # pydantic v1 is already installed
+        yield
+    finally:
+        # Restore the original Pydantic module
+        for k, original in originals.items():
+            if k in sys.modules:
+                if original:
+                    sys.modules[k] = original
+                else:
+                    del sys.modules[k]
+
+
+with pydantic_v1_patch():
+    import sqlmodel as sqlmodel
diff --git a/reflex/utils/types.py b/reflex/utils/types.py
index 8f7bf059a..bf1755246 100644
--- a/reflex/utils/types.py
+++ b/reflex/utils/types.py
@@ -29,7 +29,7 @@ try:
     # reflex-hosting-cli tools are compatible with pydantic v2
 
     if not TYPE_CHECKING:
-        import pydantic.v1.fields as ModelField
+        from pydantic.v1.fields import ModelField
     else:
         raise ModuleNotFoundError
 except ModuleNotFoundError:
diff --git a/tests/components/core/test_foreach.py b/tests/components/core/test_foreach.py
index c87d08b25..bfd97265e 100644
--- a/tests/components/core/test_foreach.py
+++ b/tests/components/core/test_foreach.py
@@ -1,12 +1,17 @@
 from typing import Dict, List, Set, Tuple
 
 import pytest
-from pydantic import ValidationError
 
 from reflex.components import box, foreach, text, theme
 from reflex.components.core import Foreach
 from reflex.state import BaseState
 
+try:
+    # When pydantic v2 is installed
+    from pydantic.v1 import ValidationError  # type: ignore
+except ImportError:
+    from pydantic import ValidationError
+
 
 class ForEachState(BaseState):
     """A state for testing the ForEach component."""