From c4abf5a44249fa42ae9c5d5c3035b8258c6d92b6 Mon Sep 17 00:00:00 2001 From: Mike Bayer Date: Mon, 1 Nov 2021 16:36:51 -0400 Subject: [PATCH] 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 --- doc/build/changelog/unreleased_14/7272.rst | 14 ++++++++ lib/sqlalchemy/future/engine.py | 21 +++-------- test/engine/test_execute.py | 41 +++++++++++++++++++--- 3 files changed, 55 insertions(+), 21 deletions(-) create mode 100644 doc/build/changelog/unreleased_14/7272.rst diff --git a/doc/build/changelog/unreleased_14/7272.rst b/doc/build/changelog/unreleased_14/7272.rst new file mode 100644 index 0000000000..a38aacdaa8 --- /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 ab890ca4f4..3235529f73 100644 --- 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 91cc60fc7b..9296164fb3 100644 --- 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 -- 2.47.2