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

9 files changed, 288 insertions, 63 deletions

diff --git a/lib/sqlalchemy/engine/base.py b/lib/sqlalchemy/engine/base.py
index 293dc21b4..663482b1f 100644
--- a/lib/sqlalchemy/engine/base.py
+++ b/lib/sqlalchemy/engine/base.py
@@ -13,6 +13,7 @@ from .interfaces import Connectable
 from .interfaces import ExceptionContext
 from .util import _distill_params
 from .util import _distill_params_20
+from .util import TransactionalContext
 from .. import exc
 from .. import inspection
 from .. import log
@@ -60,6 +61,9 @@ class Connection(Connectable):
     _is_future = False
     _sqla_logger_namespace = "sqlalchemy.engine.Connection"
 
+    # used by sqlalchemy.engine.util.TransactionalContext
+    _trans_context_manager = None
+
     def __init__(
         self,
         engine,
@@ -1683,6 +1687,9 @@ class Connection(Connectable):
         ):
             self._invalid_transaction()
 
+        elif self._trans_context_manager:
+            TransactionalContext._trans_ctx_check(self)
+
         if self._is_future and self._transaction is None:
             self._autobegin()
 
@@ -2182,7 +2189,7 @@ class ExceptionContextImpl(ExceptionContext):
         self.invalidate_pool_on_disconnect = invalidate_pool_on_disconnect
 
 
-class Transaction(object):
+class Transaction(TransactionalContext):
     """Represent a database transaction in progress.
 
     The :class:`.Transaction` object is procured by
@@ -2324,21 +2331,14 @@ class Transaction(object):
         finally:
             assert not self.is_active
 
-    def __enter__(self):
-        return self
+    def _get_subject(self):
+        return self.connection
 
-    def __exit__(self, type_, value, traceback):
-        if type_ is None and self.is_active:
-            try:
-                self.commit()
-            except:
-                with util.safe_reraise():
-                    self.rollback()
-        else:
-            if self._deactivated_from_connection:
-                self.close()
-            else:
-                self.rollback()
+    def _transaction_is_active(self):
+        return self.is_active
+
+    def _transaction_is_closed(self):
+        return not self._deactivated_from_connection
 
 
 class MarkerTransaction(Transaction):
@@ -2368,6 +2368,10 @@ class MarkerTransaction(Transaction):
         )
 
         self.connection = connection
+
+        if connection._trans_context_manager:
+            TransactionalContext._trans_ctx_check(connection)
+
         if connection._nested_transaction is not None:
             self._transaction = connection._nested_transaction
         else:
@@ -2429,6 +2433,8 @@ class RootTransaction(Transaction):
 
     def __init__(self, connection):
         assert connection._transaction is None
+        if connection._trans_context_manager:
+            TransactionalContext._trans_ctx_check(connection)
         self.connection = connection
         self._connection_begin_impl()
         connection._transaction = self
@@ -2564,6 +2570,8 @@ class NestedTransaction(Transaction):
 
     def __init__(self, connection):
         assert connection._transaction is not None
+        if connection._trans_context_manager:
+            TransactionalContext._trans_ctx_check(connection)
         self.connection = connection
         self._savepoint = self.connection._savepoint_impl()
         self.is_active = True
@@ -2935,16 +2943,12 @@ class Engine(Connectable, log.Identified):
             self.close_with_result = close_with_result
 
         def __enter__(self):
+            self.transaction.__enter__()
             return self.conn
 
         def __exit__(self, type_, value, traceback):
             try:
-                if type_ is not None:
-                    if self.transaction.is_active:
-                        self.transaction.rollback()
-                else:
-                    if self.transaction.is_active:
-                        self.transaction.commit()
+                self.transaction.__exit__(type_, value, traceback)
             finally:
                 if not self.close_with_result:
                     self.conn.close()
diff --git a/lib/sqlalchemy/engine/util.py b/lib/sqlalchemy/engine/util.py
index ede263198..17e3510aa 100644
--- a/lib/sqlalchemy/engine/util.py
+++ b/lib/sqlalchemy/engine/util.py
@@ -153,3 +153,82 @@ def _distill_params_20(params):
         return (params,), _no_kw
     else:
         raise exc.ArgumentError("mapping or sequence expected for parameters")
+
+
+class TransactionalContext(object):
+    """Apply Python context manager behavior to transaction objects.
+
+    Performs validation to ensure the subject of the transaction is not
+    used if the transaction were ended prematurely.
+
+    """
+
+    _trans_subject = None
+
+    def _transaction_is_active(self):
+        raise NotImplementedError()
+
+    def _transaction_is_closed(self):
+        raise NotImplementedError()
+
+    def _get_subject(self):
+        raise NotImplementedError()
+
+    @classmethod
+    def _trans_ctx_check(cls, subject):
+        trans_context = subject._trans_context_manager
+        if trans_context:
+            if not trans_context._transaction_is_active():
+                raise exc.InvalidRequestError(
+                    "Can't operate on closed transaction inside context "
+                    "manager.  Please complete the context manager "
+                    "before emitting further commands."
+                )
+
+    def __enter__(self):
+        subject = self._get_subject()
+
+        # none for outer transaction, may be non-None for nested
+        # savepoint, legacy nesting cases
+        trans_context = subject._trans_context_manager
+        self._outer_trans_ctx = trans_context
+
+        self._trans_subject = subject
+        subject._trans_context_manager = self
+        return self
+
+    def __exit__(self, type_, value, traceback):
+        subject = self._trans_subject
+
+        # simplistically we could assume that
+        # "subject._trans_context_manager is self".  However, any calling
+        # code that is manipulating __exit__ directly would break this
+        # assumption.  alembic context manager
+        # is an example of partial use that just calls __exit__ and
+        # not __enter__ at the moment.  it's safe to assume this is being done
+        # in the wild also
+        out_of_band_exit = (
+            subject is None or subject._trans_context_manager is not self
+        )
+
+        if type_ is None and self._transaction_is_active():
+            try:
+                self.commit()
+            except:
+                with util.safe_reraise():
+                    self.rollback()
+            finally:
+                if not out_of_band_exit:
+                    subject._trans_context_manager = self._outer_trans_ctx
+                self._trans_subject = self._outer_trans_ctx = None
+        else:
+            try:
+                if not self._transaction_is_active():
+                    if not self._transaction_is_closed():
+                        self.close()
+                else:
+                    self.rollback()
+            finally:
+                if not out_of_band_exit:
+                    subject._trans_context_manager = self._outer_trans_ctx
+                self._trans_subject = self._outer_trans_ctx = None
diff --git a/lib/sqlalchemy/ext/asyncio/base.py b/lib/sqlalchemy/ext/asyncio/base.py
index fa8c5006e..d11b059fd 100644
--- a/lib/sqlalchemy/ext/asyncio/base.py
+++ b/lib/sqlalchemy/ext/asyncio/base.py
@@ -5,14 +5,14 @@ from . import exc as async_exc
 
 class StartableContext(abc.ABC):
     @abc.abstractmethod
-    async def start(self) -> "StartableContext":
+    async def start(self, is_ctxmanager=False) -> "StartableContext":
         pass
 
     def __await__(self):
         return self.start().__await__()
 
     async def __aenter__(self):
-        return await self.start()
+        return await self.start(is_ctxmanager=True)
 
     @abc.abstractmethod
     async def __aexit__(self, type_, value, traceback):
diff --git a/lib/sqlalchemy/ext/asyncio/engine.py b/lib/sqlalchemy/ext/asyncio/engine.py
index c637b3d90..17ddb614a 100644
--- a/lib/sqlalchemy/ext/asyncio/engine.py
+++ b/lib/sqlalchemy/ext/asyncio/engine.py
@@ -101,7 +101,7 @@ class AsyncConnection(ProxyComparable, StartableContext, AsyncConnectable):
         self.sync_engine = async_engine.sync_engine
         self.sync_connection = sync_connection
 
-    async def start(self):
+    async def start(self, is_ctxmanager=False):
         """Start this :class:`_asyncio.AsyncConnection` object's context
         outside of using a Python ``with:`` block.
 
@@ -518,19 +518,15 @@ class AsyncEngine(ProxyComparable, AsyncConnectable):
         def __init__(self, conn):
             self.conn = conn
 
-        async def start(self):
-            await self.conn.start()
+        async def start(self, is_ctxmanager=False):
+            await self.conn.start(is_ctxmanager=is_ctxmanager)
             self.transaction = self.conn.begin()
             await self.transaction.__aenter__()
 
             return self.conn
 
         async def __aexit__(self, type_, value, traceback):
-            if type_ is not None:
-                await self.transaction.rollback()
-            else:
-                if self.transaction.is_active:
-                    await self.transaction.commit()
+            await self.transaction.__aexit__(type_, value, traceback)
             await self.conn.close()
 
     def __init__(self, sync_engine: Engine):
@@ -678,7 +674,7 @@ class AsyncTransaction(ProxyComparable, StartableContext):
 
         await greenlet_spawn(self._sync_transaction().commit)
 
-    async def start(self):
+    async def start(self, is_ctxmanager=False):
         """Start this :class:`_asyncio.AsyncTransaction` object's context
         outside of using a Python ``with:`` block.
 
@@ -689,17 +685,14 @@ class AsyncTransaction(ProxyComparable, StartableContext):
             if self.nested
             else self.connection._sync_connection().begin
         )
+        if is_ctxmanager:
+            self.sync_transaction.__enter__()
         return self
 
     async def __aexit__(self, type_, value, traceback):
-        if type_ is None and self.is_active:
-            try:
-                await self.commit()
-            except:
-                with util.safe_reraise():
-                    await self.rollback()
-        else:
-            await self.rollback()
+        await greenlet_spawn(
+            self._sync_transaction().__exit__, type_, value, traceback
+        )
 
 
 def _get_sync_engine_or_connection(async_engine):
diff --git a/lib/sqlalchemy/ext/asyncio/session.py b/lib/sqlalchemy/ext/asyncio/session.py
index d8a5673eb..1b61d6ee3 100644
--- a/lib/sqlalchemy/ext/asyncio/session.py
+++ b/lib/sqlalchemy/ext/asyncio/session.py
@@ -399,15 +399,17 @@ class AsyncSessionTransaction(StartableContext):
 
         await greenlet_spawn(self._sync_transaction().commit)
 
-    async def start(self):
+    async def start(self, is_ctxmanager=False):
         self.sync_transaction = await greenlet_spawn(
             self.session.sync_session.begin_nested
             if self.nested
             else self.session.sync_session.begin
         )
+        if is_ctxmanager:
+            self.sync_transaction.__enter__()
         return self
 
     async def __aexit__(self, type_, value, traceback):
-        return await greenlet_spawn(
+        await greenlet_spawn(
             self._sync_transaction().__exit__, type_, value, traceback
         )
diff --git a/lib/sqlalchemy/future/engine.py b/lib/sqlalchemy/future/engine.py
index cee17a432..ab890ca4f 100644
--- a/lib/sqlalchemy/future/engine.py
+++ b/lib/sqlalchemy/future/engine.py
@@ -359,16 +359,12 @@ class Engine(_LegacyEngine):
 
         def __enter__(self):
             self.transaction = self.conn.begin()
+            self.transaction.__enter__()
             return self.conn
 
         def __exit__(self, type_, value, traceback):
             try:
-                if type_ is not None:
-                    if self.transaction.is_active:
-                        self.transaction.rollback()
-                else:
-                    if self.transaction.is_active:
-                        self.transaction.commit()
+                self.transaction.__exit__(type_, value, traceback)
             finally:
                 self.conn.close()
 
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
diff --git a/lib/sqlalchemy/testing/fixtures.py b/lib/sqlalchemy/testing/fixtures.py
index c3eb1b363..581bc8bec 100644
--- a/lib/sqlalchemy/testing/fixtures.py
+++ b/lib/sqlalchemy/testing/fixtures.py
@@ -147,6 +147,150 @@ class TestBase(object):
         else:
             drop_all_tables_from_metadata(metadata, config.db)
 
+    @config.fixture(
+        params=[
+            (rollback, second_operation, begin_nested)
+            for rollback in (True, False)
+            for second_operation in ("none", "execute", "begin")
+            for begin_nested in (
+                True,
+                False,
+            )
+        ]
+    )
+    def trans_ctx_manager_fixture(self, request, metadata):
+        rollback, second_operation, begin_nested = request.param
+
+        from sqlalchemy import Table, Column, Integer, func, select
+        from . import eq_
+
+        t = Table("test", metadata, Column("data", Integer))
+        eng = getattr(self, "bind", None) or config.db
+
+        t.create(eng)
+
+        def run_test(subject, trans_on_subject, execute_on_subject):
+            with subject.begin() as trans:
+
+                if begin_nested:
+                    if not config.requirements.savepoints.enabled:
+                        config.skip_test("savepoints not enabled")
+                    if execute_on_subject:
+                        nested_trans = subject.begin_nested()
+                    else:
+                        nested_trans = trans.begin_nested()
+
+                    with nested_trans:
+                        if execute_on_subject:
+                            subject.execute(t.insert(), {"data": 10})
+                        else:
+                            trans.execute(t.insert(), {"data": 10})
+
+                        # for nested trans, we always commit/rollback on the
+                        # "nested trans" object itself.
+                        # only Session(future=False) will affect savepoint
+                        # transaction for session.commit/rollback
+
+                        if rollback:
+                            nested_trans.rollback()
+                        else:
+                            nested_trans.commit()
+
+                        if second_operation != "none":
+                            with assertions.expect_raises_message(
+                                sa.exc.InvalidRequestError,
+                                "Can't operate on closed transaction "
+                                "inside context "
+                                "manager.  Please complete the context "
+                                "manager "
+                                "before emitting further commands.",
+                            ):
+                                if second_operation == "execute":
+                                    if execute_on_subject:
+                                        subject.execute(
+                                            t.insert(), {"data": 12}
+                                        )
+                                    else:
+                                        trans.execute(t.insert(), {"data": 12})
+                                elif second_operation == "begin":
+                                    if execute_on_subject:
+                                        subject.begin_nested()
+                                    else:
+                                        trans.begin_nested()
+
+                    # outside the nested trans block, but still inside the
+                    # transaction block, we can run SQL, and it will be
+                    # committed
+                    if execute_on_subject:
+                        subject.execute(t.insert(), {"data": 14})
+                    else:
+                        trans.execute(t.insert(), {"data": 14})
+
+                else:
+                    if execute_on_subject:
+                        subject.execute(t.insert(), {"data": 10})
+                    else:
+                        trans.execute(t.insert(), {"data": 10})
+
+                    if trans_on_subject:
+                        if rollback:
+                            subject.rollback()
+                        else:
+                            subject.commit()
+                    else:
+                        if rollback:
+                            trans.rollback()
+                        else:
+                            trans.commit()
+
+                    if second_operation != "none":
+                        with assertions.expect_raises_message(
+                            sa.exc.InvalidRequestError,
+                            "Can't operate on closed transaction inside "
+                            "context "
+                            "manager.  Please complete the context manager "
+                            "before emitting further commands.",
+                        ):
+                            if second_operation == "execute":
+                                if execute_on_subject:
+                                    subject.execute(t.insert(), {"data": 12})
+                                else:
+                                    trans.execute(t.insert(), {"data": 12})
+                            elif second_operation == "begin":
+                                if hasattr(trans, "begin"):
+                                    trans.begin()
+                                else:
+                                    subject.begin()
+                            elif second_operation == "begin_nested":
+                                if execute_on_subject:
+                                    subject.begin_nested()
+                                else:
+                                    trans.begin_nested()
+
+            expected_committed = 0
+            if begin_nested:
+                # begin_nested variant, we inserted a row after the nested
+                # block
+                expected_committed += 1
+            if not rollback:
+                # not rollback variant, our row inserted in the target
+                # block itself would be committed
+                expected_committed += 1
+
+            if execute_on_subject:
+                eq_(
+                    subject.scalar(select(func.count()).select_from(t)),
+                    expected_committed,
+                )
+            else:
+                with subject.connect() as conn:
+                    eq_(
+                        conn.scalar(select(func.count()).select_from(t)),
+                        expected_committed,
+                    )
+
+        return run_test
+
 
 _connection_fixture_connection = None
 
diff --git a/lib/sqlalchemy/util/concurrency.py b/lib/sqlalchemy/util/concurrency.py
index e26f305d9..60db9cfff 100644
--- a/lib/sqlalchemy/util/concurrency.py
+++ b/lib/sqlalchemy/util/concurrency.py
@@ -24,6 +24,11 @@ if not have_greenlet:
     asyncio = None  # noqa F811
 
     def _not_implemented():
+        # this conditional is to prevent pylance from considering
+        # greenlet_spawn() etc as "no return" and dimming out code below it
+        if have_greenlet:
+            return None
+
         if not compat.py3k:
             raise ValueError("Cannot use this function in py2.")
         else: