unify transactional context managers

Applied consistent behavior to the use case of
calling ``.commit()`` or ``.rollback()`` inside of an existing
``.begin()`` context manager, with the addition of potentially
emitting SQL within the block subsequent to the commit or rollback.
This change continues upon the change first added in
:ticket:`6155` where the use case of calling "rollback" inside of
a ``.begin()`` contextmanager block was proposed:

* calling ``.commit()`` or ``.rollback()`` will now be allowed
without error or warning within all scopes, including
that of legacy and future :class:`_engine.Engine`, ORM
:class:`_orm.Session`, asyncio :class:`.AsyncEngine`.  Previously,
the :class:`_orm.Session` disallowed this.

* The remaining scope of the context manager is then closed;
when the block ends, a check is emitted to see if the transaction
was already ended, and if so the block returns without action.

* It will now raise **an error** if subsequent SQL of any kind
is emitted within the block, **after** ``.commit()`` or
``.rollback()`` is called.   The block should be closed as
the state of the executable object would otherwise be undefined
in this state.

Fixes: #6288
Change-Id: I8b21766ae430f0fa1ac5ef689f4c0fb19fc84336

1 files changed, 17 insertions, 15 deletions

diff --git a/lib/sqlalchemy/orm/session.py b/lib/sqlalchemy/orm/session.py
index a3ec360d0..cdf3a1585 100644
--- a/lib/sqlalchemy/orm/session.py
+++ b/lib/sqlalchemy/orm/session.py
@@ -31,6 +31,7 @@ from .. import engine
 from .. import exc as sa_exc
 from .. import sql
 from .. import util
+from ..engine.util import TransactionalContext
 from ..inspection import inspect
 from ..sql import coercions
 from ..sql import dml
@@ -475,7 +476,7 @@ class ORMExecuteState(util.MemoizedSlots):
         ]
 
 
-class SessionTransaction(object):
+class SessionTransaction(TransactionalContext):
     """A :class:`.Session`-level transaction.
 
     :class:`.SessionTransaction` is produced from the
@@ -523,6 +524,8 @@ class SessionTransaction(object):
         nested=False,
         autobegin=False,
     ):
+        TransactionalContext._trans_ctx_check(session)
+
         self.session = session
         self._connections = {}
         self._parent = parent
@@ -927,21 +930,14 @@ class SessionTransaction(object):
         self.session = None
         self._connections = None
 
-    def __enter__(self):
-        return self
+    def _get_subject(self):
+        return self.session
 
-    def __exit__(self, type_, value, traceback):
-        self._assert_active(deactive_ok=True, prepared_ok=True)
-        if self.session._transaction is None:
-            return
-        if type_ is None:
-            try:
-                self.commit()
-            except:
-                with util.safe_reraise():
-                    self.rollback()
-        else:
-            self.rollback()
+    def _transaction_is_active(self):
+        return self._state is ACTIVE
+
+    def _transaction_is_closed(self):
+        return self._state is CLOSED
 
 
 class Session(_SessionClassMethods):
@@ -1154,6 +1150,9 @@ class Session(_SessionClassMethods):
 
         _sessions[self.hash_key] = self
 
+    # used by sqlalchemy.engine.util.TransactionalContext
+    _trans_context_manager = None
+
     connection_callable = None
 
     def __enter__(self):
@@ -1252,6 +1251,7 @@ class Session(_SessionClassMethods):
 
     def _autobegin(self):
         if not self.autocommit and self._transaction is None:
+
             trans = SessionTransaction(self, autobegin=True)
             assert self._transaction is trans
             return True
@@ -1520,6 +1520,8 @@ class Session(_SessionClassMethods):
         )
 
     def _connection_for_bind(self, engine, execution_options=None, **kw):
+        TransactionalContext._trans_ctx_check(self)
+
         if self._transaction is not None or self._autobegin():
             return self._transaction._connection_for_bind(
                 engine, execution_options