From cf7981f60d485f17465f44c6ff651ae283ade377 Mon Sep 17 00:00:00 2001
From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Fri, 12 Dec 2014 19:59:11 -0500
Subject: - Added new method :meth:`.Session.invalidate`, functions similarly
 to :meth:`.Session.close`, except also calls :meth:`.Connection.invalidate`
 on all connections, guaranteeing that they will not be returned to the
 connection pool.  This is useful in situations e.g. dealing with gevent
 timeouts when it is not safe to use the connection further, even for
 rollbacks. references #3258

---
 doc/build/changelog/changelog_09.rst | 12 +++++++++++
 lib/sqlalchemy/orm/session.py        | 42 ++++++++++++++++++++++++++++++++++--
 test/orm/test_session.py             |  3 +++
 test/orm/test_transaction.py         | 17 +++++++++++++++
 4 files changed, 72 insertions(+), 2 deletions(-)

diff --git a/doc/build/changelog/changelog_09.rst b/doc/build/changelog/changelog_09.rst
index f10d48273..419827959 100644
--- a/doc/build/changelog/changelog_09.rst
+++ b/doc/build/changelog/changelog_09.rst
@@ -13,6 +13,18 @@
 .. changelog::
     :version: 0.9.9
 
+    .. change::
+        :tags: enhancement, orm
+        :versions: 1.0.0
+
+        Added new method :meth:`.Session.invalidate`, functions similarly
+        to :meth:`.Session.close`, except also calls
+        :meth:`.Connection.invalidate`
+        on all connections, guaranteeing that they will not be returned to
+        the connection pool.  This is useful in situations e.g. dealing
+        with gevent timeouts when it is not safe to use the connection further,
+        even for rollbacks.
+
     .. change::
         :tags: bug, examples
         :versions: 1.0.0
diff --git a/lib/sqlalchemy/orm/session.py b/lib/sqlalchemy/orm/session.py
index d40d28154..507e99b2e 100644
--- a/lib/sqlalchemy/orm/session.py
+++ b/lib/sqlalchemy/orm/session.py
@@ -435,11 +435,13 @@ class SessionTransaction(object):
 
         self.session.dispatch.after_rollback(self.session)
 
-    def close(self):
+    def close(self, invalidate=False):
         self.session.transaction = self._parent
         if self._parent is None:
             for connection, transaction, autoclose in \
                     set(self._connections.values()):
+                if invalidate:
+                    connection.invalidate()
                 if autoclose:
                     connection.close()
                 else:
@@ -1000,10 +1002,46 @@ class Session(_SessionClassMethods):
         not use any connection resources until they are first needed.
 
         """
+        self._close_impl(invalidate=False)
+
+    def invalidate(self):
+        """Close this Session, using connection invalidation.
+
+        This is a variant of :meth:`.Session.close` that will additionally
+        ensure that the :meth:`.Connection.invalidate` method will be called
+        on all :class:`.Connection` objects.  This can be called when
+        the database is known to be in a state where the connections are
+        no longer safe to be used.
+
+        E.g.::
+
+            try:
+                sess = Session()
+                sess.add(User())
+                sess.commit()
+            except gevent.Timeout:
+                sess.invalidate()
+                raise
+            except:
+                sess.rollback()
+                raise
+
+        This clears all items and ends any transaction in progress.
+
+        If this session were created with ``autocommit=False``, a new
+        transaction is immediately begun.  Note that this new transaction does
+        not use any connection resources until they are first needed.
+
+        .. versionadded:: 0.9.9
+
+        """
+        self._close_impl(invalidate=True)
+
+    def _close_impl(self, invalidate):
         self.expunge_all()
         if self.transaction is not None:
             for transaction in self.transaction._iterate_parents():
-                transaction.close()
+                transaction.close(invalidate)
 
     def expunge_all(self):
         """Remove all object instances from this ``Session``.
diff --git a/test/orm/test_session.py b/test/orm/test_session.py
index b81c03f88..2aa0cd3eb 100644
--- a/test/orm/test_session.py
+++ b/test/orm/test_session.py
@@ -1364,6 +1364,9 @@ class DisposedStates(fixtures.MappedTest):
     def test_close(self):
         self._test_session().close()
 
+    def test_invalidate(self):
+        self._test_session().invalidate()
+
     def test_expunge_all(self):
         self._test_session().expunge_all()
 
diff --git a/test/orm/test_transaction.py b/test/orm/test_transaction.py
index ba31e4c7d..1d7e8e693 100644
--- a/test/orm/test_transaction.py
+++ b/test/orm/test_transaction.py
@@ -184,6 +184,23 @@ class SessionTransactionTest(FixtureTest):
         assert users.count().scalar() == 1
         assert addresses.count().scalar() == 1
 
+    @testing.requires.independent_connections
+    def test_invalidate(self):
+        User, users = self.classes.User, self.tables.users
+        mapper(User, users)
+        sess = Session()
+        u = User(name='u1')
+        sess.add(u)
+        sess.flush()
+        c1 = sess.connection(User)
+
+        sess.invalidate()
+        assert c1.invalidated
+
+        eq_(sess.query(User).all(), [])
+        c2 = sess.connection(User)
+        assert not c2.invalidated
+
     def test_subtransaction_on_noautocommit(self):
         User, users = self.classes.User, self.tables.users
 
-- 
cgit v1.2.1