use full context manager flow for future.Engine.begin()

Fixed issue in future :class:`_future.Engine` where calling upon
:meth:`_future.Engine.begin` and entering the context manager would not
close the connection if the actual BEGIN operation failed for some reason,
such as an event handler raising an exception; this use case failed to be
tested for the future version of the engine. Note that the "future" context
managers which handle ``begin()`` blocks in Core and ORM don't actually run
the "BEGIN" operation until the context managers are actually entered. This
is different from the legacy version which runs the "BEGIN" operation up
front.

Fixes: #7272
Change-Id: I9667ac0861a9e007c4b3dfcf0fcf0829038a8711

diff --git a/doc/build/changelog/unreleased_14/7272.rst b/doc/build/changelog/unreleased_14/7272.rst
new file mode 100644 (file)
index 0000000..a38aacd
--- /dev/null
+++ b/doc/build/changelog/unreleased_14/7272.rst
@@ -0,0 +1,14 @@
+.. change::
+    :tags: bug, engine
+    :tickets: 7272
+    :versions: 2.0.0b1
+
+    Fixed issue in future :class:`_future.Engine` where calling upon
+    :meth:`_future.Engine.begin` and entering the context manager would not
+    close the connection if the actual BEGIN operation failed for some reason,
+    such as an event handler raising an exception; this use case failed to be
+    tested for the future version of the engine. Note that the "future" context
+    managers which handle ``begin()`` blocks in Core and ORM don't actually run
+    the "BEGIN" operation until the context managers are actually entered. This
+    is different from the legacy version which runs the "BEGIN" operation up
+    front.

diff --git a/lib/sqlalchemy/future/engine.py b/lib/sqlalchemy/future/engine.py
index ab890ca4f4c283ded2d5e582ccb33df09b1364bf..3235529f7362f430e7a77b6aeeda5168cf474f05 100644 (file)
--- a/lib/sqlalchemy/future/engine.py
+++ b/lib/sqlalchemy/future/engine.py
@@ -353,21 +353,7 @@ class Engine(_LegacyEngine):
             execution_options=legacy_engine._execution_options,
         )
 
-    class _trans_ctx(object):
-        def __init__(self, conn):
-            self.conn = conn
-
-        def __enter__(self):
-            self.transaction = self.conn.begin()
-            self.transaction.__enter__()
-            return self.conn
-
-        def __exit__(self, type_, value, traceback):
-            try:
-                self.transaction.__exit__(type_, value, traceback)
-            finally:
-                self.conn.close()
-
+    @util.contextmanager
     def begin(self):
         """Return a :class:`_future.Connection` object with a transaction
         begun.
@@ -390,8 +376,9 @@ class Engine(_LegacyEngine):
             :meth:`_future.Connection.begin`
 
         """
-        conn = self.connect()
-        return self._trans_ctx(conn)
+        with self.connect() as conn:
+            with conn.begin():
+                yield conn
 
     def connect(self):
         """Return a new :class:`_future.Connection` object.

diff --git a/test/engine/test_execute.py b/test/engine/test_execute.py
index 91cc60fc7bcd5897ac14f5b6f24f75e08e4647cf..9296164fb33f043f7cc389a0296d959a06673147 100644 (file)
--- a/test/engine/test_execute.py
+++ b/test/engine/test_execute.py
@@ -809,16 +809,42 @@ class ConvenienceExecuteTest(fixtures.TablesTest):
         testing.run_as_contextmanager(ctx, fn, 5, value=8)
         self._assert_fn(5, value=8)
 
-    def test_transaction_engine_ctx_begin_fails(self):
+    def test_transaction_engine_ctx_begin_fails_dont_enter_enter(self):
+        """test #7272"""
         engine = engines.testing_engine()
 
         mock_connection = Mock(
             return_value=Mock(begin=Mock(side_effect=Exception("boom")))
         )
-        engine._connection_cls = mock_connection
-        assert_raises(Exception, engine.begin)
+        with mock.patch.object(engine, "_connection_cls", mock_connection):
+            if testing.requires.legacy_engine.enabled:
+                with expect_raises_message(Exception, "boom"):
+                    engine.begin()
+            else:
+                # context manager isn't entered, doesn't actually call
+                # connect() or connection.begin()
+                engine.begin()
 
-        eq_(mock_connection.return_value.close.mock_calls, [call()])
+        if testing.requires.legacy_engine.enabled:
+            eq_(mock_connection.return_value.close.mock_calls, [call()])
+        else:
+            eq_(mock_connection.return_value.close.mock_calls, [])
+
+    def test_transaction_engine_ctx_begin_fails_include_enter(self):
+        """test #7272"""
+        engine = engines.testing_engine()
+
+        close_mock = Mock()
+        with mock.patch.object(
+            engine._connection_cls,
+            "begin",
+            Mock(side_effect=Exception("boom")),
+        ), mock.patch.object(engine._connection_cls, "close", close_mock):
+            with expect_raises_message(Exception, "boom"):
+                with engine.begin():
+                    pass
+
+        eq_(close_mock.mock_calls, [call()])
 
     def test_transaction_engine_ctx_rollback(self):
         fn = self._trans_rollback_fn()
@@ -862,6 +888,7 @@ class ConvenienceExecuteTest(fixtures.TablesTest):
             fn(conn, 5, value=8)
         self._assert_fn(5, value=8)
 
+    @testing.requires.legacy_engine
     def test_connect_as_ctx_noautocommit(self):
         fn = self._trans_fn()
         self._assert_no_data()
@@ -873,6 +900,12 @@ class ConvenienceExecuteTest(fixtures.TablesTest):
             self._assert_no_data()
 
 
+class FutureConvenienceExecuteTest(
+    fixtures.FutureEngineMixin, ConvenienceExecuteTest
+):
+    __backend__ = True
+
+
 class CompiledCacheTest(fixtures.TestBase):
     __backend__ = True