implement custom setstate to work around implicit type/comparator

Fixed issue where unpickling of a :class:`_schema.Column` or other
:class:`_sql.ColumnElement` would fail to restore the correct "comparator"
object, which is used to generate SQL expressions specific to the type
object.

Fixes: #10213
Change-Id: I74e805024bcc0d93d549bd94757c2865b3117d72

diff --git a/doc/build/changelog/unreleased_14/10213.rst b/doc/build/changelog/unreleased_14/10213.rst
new file mode 100644 (file)
index 0000000..96c17b1
--- /dev/null
+++ b/doc/build/changelog/unreleased_14/10213.rst
@@ -0,0 +1,9 @@
+.. change::
+    :tags: bug, sql
+    :tickets: 10213
+    :versions: 2.0.20
+
+    Fixed issue where unpickling of a :class:`_schema.Column` or other
+    :class:`_sql.ColumnElement` would fail to restore the correct "comparator"
+    object, which is used to generate SQL expressions specific to the type
+    object.

diff --git a/lib/sqlalchemy/sql/elements.py b/lib/sqlalchemy/sql/elements.py
index 75857fcfc445d4fb529eb5b1ad06d815dd524d81..a6b4c8ab16f1e9afd020fb26df3efe8a60fddc24 100644 (file)
--- a/lib/sqlalchemy/sql/elements.py
+++ b/lib/sqlalchemy/sql/elements.py
@@ -1591,6 +1591,9 @@ class ColumnElement(
         else:
             return comparator_factory(self)
 
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+
     def __getattr__(self, key: str) -> Any:
         try:
             return getattr(self.comparator, key)

diff --git a/lib/sqlalchemy/sql/type_api.py b/lib/sqlalchemy/sql/type_api.py
index a1588dc174f448aff96dcec42961f76f6d851bf9..2be397288ec437886a8c7c54deacda3e379f8cc4 100644 (file)
--- a/lib/sqlalchemy/sql/type_api.py
+++ b/lib/sqlalchemy/sql/type_api.py
@@ -236,9 +236,6 @@ class TypeEngine(Visitable, Generic[_T]):
 
             return op, self.type
 
-        def __reduce__(self) -> Any:
-            return _reconstitute_comparator, (self.expr,)
-
     hashable = True
     """Flag, if False, means values from this type aren't hashable.
 
@@ -2309,10 +2306,6 @@ class Variant(TypeDecorator[_T]):
         )
 
 
-def _reconstitute_comparator(expression: Any) -> Any:
-    return expression.comparator
-
-
 @overload
 def to_instance(typeobj: Union[Type[_TE], _TE], *arg: Any, **kw: Any) -> _TE:
     ...

diff --git a/test/sql/test_operators.py b/test/sql/test_operators.py
index eafc2c3a76f40065688daa1ddf4b8ffb16358e14..9ae827e3776e92507dc4d8d4b5ca426240569eea 100644 (file)
--- a/test/sql/test_operators.py
+++ b/test/sql/test_operators.py
@@ -72,6 +72,7 @@ from sqlalchemy.types import DateTime
 from sqlalchemy.types import Indexable
 from sqlalchemy.types import JSON
 from sqlalchemy.types import MatchType
+from sqlalchemy.types import NullType
 from sqlalchemy.types import TypeDecorator
 from sqlalchemy.types import TypeEngine
 from sqlalchemy.types import UserDefinedType
@@ -2614,6 +2615,22 @@ class ComparisonOperatorTest(fixtures.TestBase, testing.AssertsCompiledSQL):
         clause = tuple_(1, 2, 3)
         eq_(str(clause), str(pickle.loads(pickle.dumps(clause))))
 
+    @testing.combinations(Integer(), String(), JSON(), argnames="typ")
+    @testing.variation("eval_first", [True, False])
+    def test_pickle_comparator(self, typ, eval_first):
+        """test #10213"""
+
+        table1 = Table("t", MetaData(), Column("x", typ))
+        t1 = table1.c.x
+
+        if eval_first:
+            t1.comparator
+
+        t1p = pickle.loads(pickle.dumps(table1.c.x))
+
+        is_not(t1p.comparator.__class__, NullType.Comparator)
+        is_(t1.comparator.__class__, t1p.comparator.__class__)
+
     @testing.combinations(
         (operator.lt, "<", ">"),
         (operator.gt, ">", "<"),