From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Wed, 9 Aug 2023 14:17:35 +0000 (-0400)
Subject: implement custom setstate to work around implicit type/comparator
X-Git-Tag: rel_2_0_20~11^2
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=ccefbc643ee0e74f69977b34b91a9a93c7f77f09;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git

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
index 0000000000..96c17b1946
--- /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 75857fcfc4..a6b4c8ab16 100644
--- 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 a1588dc174..2be397288e 100644
--- 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 eafc2c3a76..9ae827e377 100644
--- 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, ">", "<"),