From 5811276bb7515af3418a6d20f5213d658e320121 Mon Sep 17 00:00:00 2001
From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Thu, 4 Jan 2018 14:09:32 -0500
Subject: [PATCH] Check for object was expunged before restoring after pk
 switch + rollback

Fixed bug where an object that is expunged during a rollback of
a nested or subtransaction which also had its primary key mutated
would not be correctly removed from the session, causing subsequent
issues in using the session.

Change-Id: I57e2888902015d67ee11857e44382818f1d2f8bc
Fixes: #4151
---
 doc/build/changelog/unreleased_11/4151.rst |  9 +++++++
 lib/sqlalchemy/orm/session.py              | 14 +++++++----
 test/orm/test_transaction.py               | 28 ++++++++++++++++++++++
 3 files changed, 47 insertions(+), 4 deletions(-)
 create mode 100644 doc/build/changelog/unreleased_11/4151.rst

diff --git a/doc/build/changelog/unreleased_11/4151.rst b/doc/build/changelog/unreleased_11/4151.rst
new file mode 100644
index 0000000000..f498959493
--- /dev/null
+++ b/doc/build/changelog/unreleased_11/4151.rst
@@ -0,0 +1,9 @@
+.. change::
+    :tags: bug, orm
+    :tickets: 4151
+    :versions: 1.2.1
+
+    Fixed bug where an object that is expunged during a rollback of
+    a nested or subtransaction which also had its primary key mutated
+    would not be correctly removed from the session, causing subsequent
+    issues in using the session.
diff --git a/lib/sqlalchemy/orm/session.py b/lib/sqlalchemy/orm/session.py
index bd0bf91972..6c91df866a 100644
--- a/lib/sqlalchemy/orm/session.py
+++ b/lib/sqlalchemy/orm/session.py
@@ -339,14 +339,20 @@ class SessionTransaction(object):
         """
         assert self._is_transaction_boundary
 
-        self.session._expunge_states(
-            set(self._new).union(self.session._new),
-            to_transient=True)
+        to_expunge = set(self._new).union(self.session._new)
+        self.session._expunge_states(to_expunge, to_transient=True)
 
         for s, (oldkey, newkey) in self._key_switches.items():
+            # we probably can do this conditionally based on
+            # if we expunged or not, but safe_discard does that anyway
             self.session.identity_map.safe_discard(s)
+
+            # restore the old key
             s.key = oldkey
-            self.session.identity_map.replace(s)
+
+            # now restore the object, but only if we didn't expunge
+            if s not in to_expunge:
+                self.session.identity_map.replace(s)
 
         for s in set(self._deleted).union(self.session._deleted):
             self.session._update_impl(s, revert_deletion=True)
diff --git a/test/orm/test_transaction.py b/test/orm/test_transaction.py
index c639418d9d..c7e81705e7 100644
--- a/test/orm/test_transaction.py
+++ b/test/orm/test_transaction.py
@@ -1720,6 +1720,34 @@ class NaturalPKRollbackTest(fixtures.MappedTest):
         assert s.identity_map[identity_key(User, ('u1',))] is u1
         assert s.identity_map[identity_key(User, ('u2',))] is u2
 
+    @testing.requires.savepoints
+    def test_key_replaced_by_update_nested(self):
+        users, User = self.tables.users, self.classes.User
+
+        mapper(User, users)
+
+        u1 = User(name='u1')
+
+        s = Session()
+        s.add(u1)
+        s.commit()
+
+        with s.begin_nested():
+            u2 = User(name='u2')
+            s.add(u2)
+            s.flush()
+
+            u2.name = 'u3'
+
+        s.rollback()
+
+        assert u1 in s
+        assert u2 not in s
+
+        u1.name = 'u5'
+
+        s.commit()
+
     def test_multiple_key_replaced_by_update(self):
         users, User = self.tables.users, self.classes.User
 
-- 
2.47.3