Fix for Pydantic 2.x (#468)

edgedb · Jan 9, 2024 · c6581cb · c6581cb
1 parent b4d91b9
commit c6581cb
Show file tree

Hide file tree

Showing 9 changed files with 51 additions and 4 deletions.
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
@@ -88,9 +88,6 @@ jobs:
         LOOP_IMPL: ${{ matrix.loop }}
         SERVER_VERSION: ${{ matrix.edgedb-version }}
       run: |
-        if [ "${SERVER_VERSION}" = "nightly" ]; then
-          export EDGEDB_TEST_CODEGEN_ASSERT_SUFFIX=.assert4
-        fi
         if [ "${LOOP_IMPL}" = "uvloop" ]; then
             env USE_UVLOOP=1 python -m unittest -v tests.suite
         else

diff --git a/edgedb/codegen/generator.py b/edgedb/codegen/generator.py
@@ -97,8 +97,15 @@
 
 PYDANTIC_MIXIN = """\
 class NoPydanticValidation:
+    @classmethod
+    def __get_pydantic_core_schema__(cls, _source_type, _handler):
+        # Pydantic 2.x
+        from pydantic_core.core_schema import any_schema
+        return any_schema()
+
     @classmethod
     def __get_validators__(cls):
+        # Pydantic 1.x
         from pydantic.dataclasses import dataclass as pydantic_dataclass
         pydantic_dataclass(cls)
         cls.__pydantic_model__.__get_validators__ = lambda: []

diff --git a/tests/codegen/test-project1/select_optional_json_async_edgeql.py.assert b/tests/codegen/test-project1/select_optional_json_async_edgeql.py.assert
@@ -10,8 +10,15 @@ import uuid
 
 
 class NoPydanticValidation:
+    @classmethod
+    def __get_pydantic_core_schema__(cls, _source_type, _handler):
+        # Pydantic 2.x
+        from pydantic_core.core_schema import any_schema
+        return any_schema()
+
     @classmethod
     def __get_validators__(cls):
+        # Pydantic 1.x
         from pydantic.dataclasses import dataclass as pydantic_dataclass
         pydantic_dataclass(cls)
         cls.__pydantic_model__.__get_validators__ = lambda: []

diff --git a/tests/codegen/test-project2/object/link_prop_async_edgeql.py.assert b/tests/codegen/test-project2/object/link_prop_async_edgeql.py.assert
@@ -11,8 +11,15 @@ import uuid
 
 
 class NoPydanticValidation:
+    @classmethod
+    def __get_pydantic_core_schema__(cls, _source_type, _handler):
+        # Pydantic 2.x
+        from pydantic_core.core_schema import any_schema
+        return any_schema()
+
     @classmethod
     def __get_validators__(cls):
+        # Pydantic 1.x
         from pydantic.dataclasses import dataclass as pydantic_dataclass
         pydantic_dataclass(cls)
         cls.__pydantic_model__.__get_validators__ = lambda: []

diff --git a/tests/codegen/test-project2/object/select_object_async_edgeql.py.assert b/tests/codegen/test-project2/object/select_object_async_edgeql.py.assert
@@ -10,8 +10,15 @@ import uuid
 
 
 class NoPydanticValidation:
+    @classmethod
+    def __get_pydantic_core_schema__(cls, _source_type, _handler):
+        # Pydantic 2.x
+        from pydantic_core.core_schema import any_schema
+        return any_schema()
+
     @classmethod
     def __get_validators__(cls):
+        # Pydantic 1.x
         from pydantic.dataclasses import dataclass as pydantic_dataclass
         pydantic_dataclass(cls)
         cls.__pydantic_model__.__get_validators__ = lambda: []

diff --git a/tests/codegen/test-project2/object/select_objects_async_edgeql.py.assert b/tests/codegen/test-project2/object/select_objects_async_edgeql.py.assert
@@ -10,8 +10,15 @@ import uuid
 
 
 class NoPydanticValidation:
+    @classmethod
+    def __get_pydantic_core_schema__(cls, _source_type, _handler):
+        # Pydantic 2.x
+        from pydantic_core.core_schema import any_schema
+        return any_schema()
+
     @classmethod
     def __get_validators__(cls):
+        # Pydantic 1.x
         from pydantic.dataclasses import dataclass as pydantic_dataclass
         pydantic_dataclass(cls)
         cls.__pydantic_model__.__get_validators__ = lambda: []

diff --git a/tests/codegen/test-project2/parpkg/select_args_async_edgeql.py.assert b/tests/codegen/test-project2/parpkg/select_args_async_edgeql.py.assert
@@ -10,8 +10,15 @@ import uuid
 
 
 class NoPydanticValidation:
+    @classmethod
+    def __get_pydantic_core_schema__(cls, _source_type, _handler):
+        # Pydantic 2.x
+        from pydantic_core.core_schema import any_schema
+        return any_schema()
+
     @classmethod
     def __get_validators__(cls):
+        # Pydantic 1.x
         from pydantic.dataclasses import dataclass as pydantic_dataclass
         pydantic_dataclass(cls)
         cls.__pydantic_model__.__get_validators__ = lambda: []

diff --git a/tests/codegen/test-project2/parpkg/subpkg/my_query_async_edgeql.py.assert b/tests/codegen/test-project2/parpkg/subpkg/my_query_async_edgeql.py.assert
@@ -16,8 +16,15 @@ MyScalar = int
 
 
 class NoPydanticValidation:
+    @classmethod
+    def __get_pydantic_core_schema__(cls, _source_type, _handler):
+        # Pydantic 2.x
+        from pydantic_core.core_schema import any_schema
+        return any_schema()
+
     @classmethod
     def __get_validators__(cls):
+        # Pydantic 1.x
         from pydantic.dataclasses import dataclass as pydantic_dataclass
         pydantic_dataclass(cls)
         cls.__pydantic_model__.__get_validators__ = lambda: []

diff --git a/tests/test_codegen.py b/tests/test_codegen.py
@@ -27,7 +27,8 @@
 from edgedb import _testbase as tb
 
 
-ASSERT_SUFFIX = os.environ.get("EDGEDB_TEST_CODEGEN_ASSERT_SUFFIX", ".assert")
+# Use ".assert" for EdgeDB 3.x and lower
+ASSERT_SUFFIX = os.environ.get("EDGEDB_TEST_CODEGEN_ASSERT_SUFFIX", ".assert4")
 
 
 class TestCodegen(tb.AsyncQueryTestCase):