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
(cherry picked from commit c4abf5a44249fa42ae9c5d5c3035b8258c6d92b6)

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 4a14cbcca0c9f2e2e1038f9628ad1cbf06fd160a..22731f5d082d83abc93e0ed20a2e2138aae12aff 100644 (file)
--- a/test/engine/test_execute.py
+++ b/test/engine/test_execute.py
@@ -813,16 +813,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()
@@ -867,6 +893,7 @@ class ConvenienceExecuteTest(fixtures.TablesTest):
         self._assert_fn(5, value=8)
 
     @testing.fails_on("mysql+oursql", "oursql bug ?  getting wrong rowcount")
+    @testing.requires.legacy_engine
     def test_connect_as_ctx_noautocommit(self):
         fn = self._trans_fn()
         self._assert_no_data()
@@ -878,6 +905,12 @@ class ConvenienceExecuteTest(fixtures.TablesTest):
             self._assert_no_data()
 
 
+class FutureConvenienceExecuteTest(
+    fixtures.FutureEngineMixin, ConvenienceExecuteTest
+):
+    __backend__ = True
+
+
 class CompiledCacheTest(fixtures.TestBase):
     __backend__ = True