From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Wed, 2 Aug 2023 17:34:03 +0000 (-0400)
Subject: ensure collection adapter is serialized
X-Git-Tag: rel_2_0_20~22^2
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=b62f791bd4306b636c01707b9c715aa14e6e0903;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git

ensure collection adapter is serialized

Fixed issue where dictionary-based collections such as
:func:`_orm.attribute_keyed_dict` did not fully pickle/unpickle correctly,
leading to issues when attempting to mutate such a collection after
unpickling.

Fixes: #10175
Change-Id: I1281f8695a0c95a20cab9449ee9c5d2db0d544fe
---

diff --git a/doc/build/changelog/unreleased_20/10175.rst b/doc/build/changelog/unreleased_20/10175.rst
new file mode 100644
index 0000000000..856b55bf62
--- /dev/null
+++ b/doc/build/changelog/unreleased_20/10175.rst
@@ -0,0 +1,9 @@
+.. change::
+    :tags: bug, orm
+    :tickets: 10175
+
+    Fixed issue where dictionary-based collections such as
+    :func:`_orm.attribute_keyed_dict` did not fully pickle/unpickle correctly,
+    leading to issues when attempting to mutate such a collection after
+    unpickling.
+
diff --git a/lib/sqlalchemy/orm/mapped_collection.py b/lib/sqlalchemy/orm/mapped_collection.py
index fb6b05d78e..9e479d0d30 100644
--- a/lib/sqlalchemy/orm/mapped_collection.py
+++ b/lib/sqlalchemy/orm/mapped_collection.py
@@ -23,6 +23,7 @@ from typing import Union
 
 from . import base
 from .collections import collection
+from .collections import collection_adapter
 from .. import exc as sa_exc
 from .. import util
 from ..sql import coercions
@@ -33,6 +34,7 @@ from ..util.typing import Literal
 if TYPE_CHECKING:
     from . import AttributeEventToken
     from . import Mapper
+    from .collections import CollectionAdapter
     from ..sql.elements import ColumnElement
 
 _KT = TypeVar("_KT", bound=Any)
@@ -376,19 +378,31 @@ class KeyFuncDict(Dict[_KT, _VT]):
 
     @classmethod
     def _unreduce(
-        cls, keyfunc: _F, values: Dict[_KT, _KT]
+        cls,
+        keyfunc: _F,
+        values: Dict[_KT, _KT],
+        adapter: Optional[CollectionAdapter] = None,
     ) -> "KeyFuncDict[_KT, _KT]":
         mp: KeyFuncDict[_KT, _KT] = KeyFuncDict(keyfunc)
         mp.update(values)
+        # note that the adapter sets itself up onto this collection
+        # when its `__setstate__` method is called
         return mp
 
     def __reduce__(
         self,
     ) -> Tuple[
         Callable[[_KT, _KT], KeyFuncDict[_KT, _KT]],
-        Tuple[Any, Union[Dict[_KT, _KT], Dict[_KT, _KT]]],
+        Tuple[Any, Union[Dict[_KT, _KT], Dict[_KT, _KT]], CollectionAdapter],
     ]:
-        return (KeyFuncDict._unreduce, (self.keyfunc, dict(self)))
+        return (
+            KeyFuncDict._unreduce,
+            (
+                self.keyfunc,
+                dict(self),
+                collection_adapter(self),
+            ),
+        )
 
     @util.preload_module("sqlalchemy.orm.attributes")
     def _raise_for_unpopulated(
diff --git a/test/orm/test_pickled.py b/test/orm/test_pickled.py
index 87461bc102..96dec4a60b 100644
--- a/test/orm/test_pickled.py
+++ b/test/orm/test_pickled.py
@@ -10,6 +10,7 @@ from sqlalchemy import testing
 from sqlalchemy.orm import aliased
 from sqlalchemy.orm import attributes
 from sqlalchemy.orm import clear_mappers
+from sqlalchemy.orm import collections
 from sqlalchemy.orm import exc as orm_exc
 from sqlalchemy.orm import lazyload
 from sqlalchemy.orm import relationship
@@ -22,6 +23,7 @@ from sqlalchemy.orm.collections import column_keyed_dict
 from sqlalchemy.testing import assert_raises_message
 from sqlalchemy.testing import eq_
 from sqlalchemy.testing import fixtures
+from sqlalchemy.testing import is_not_none
 from sqlalchemy.testing.fixtures import fixture_session
 from sqlalchemy.testing.pickleable import Address
 from sqlalchemy.testing.pickleable import AddressWMixin
@@ -592,6 +594,8 @@ class PickleTest(fixtures.MappedTest):
             eq_(u1.addresses, repickled.addresses)
             eq_(repickled.addresses["email1"], Address(email_address="email1"))
 
+            is_not_none(collections.collection_adapter(repickled.addresses))
+
     def test_column_mapped_collection(self):
         users, addresses = self.tables.users, self.tables.addresses
 
@@ -618,6 +622,8 @@ class PickleTest(fixtures.MappedTest):
             eq_(u1.addresses, repickled.addresses)
             eq_(repickled.addresses["email1"], Address(email_address="email1"))
 
+            is_not_none(collections.collection_adapter(repickled.addresses))
+
     def test_composite_column_mapped_collection(self):
         users, addresses = self.tables.users, self.tables.addresses
 
@@ -646,6 +652,7 @@ class PickleTest(fixtures.MappedTest):
                 repickled.addresses[(1, "email1")],
                 Address(id=1, email_address="email1"),
             )
+            is_not_none(collections.collection_adapter(repickled.addresses))
 
 
 class OptionsTest(_Polymorphic):