Ensure proxy transaction still present on exit before closing

Fixed regression caused by the SQLAlchemy 1.4/2.0 compatibility switch
where calling ``.rollback()`` or ``.commit()`` explicitly within the
``context.begin_transaction()`` context manager would cause it to fail when
the block ended, as it did not expect that the transaction was manually
closed.

Change-Id: I4c5ba368bfd480d186276cd75edaac5019130bc6
Fixes: #829

diff --git a/alembic/runtime/migration.py b/alembic/runtime/migration.py
index e8a36e9de2623924b68898e836e9a9f41d25eabd..bf48e2bd9816feccb152ab3366ed186972e710d8 100644 (file)
--- a/alembic/runtime/migration.py
+++ b/alembic/runtime/migration.py
@@ -41,8 +41,9 @@ class _ProxyTransaction(object):
         return self
 
     def __exit__(self, type_, value, traceback):
-        self._proxied_transaction.__exit__(type_, value, traceback)
-        self.migration_context._transaction = None
+        if self._proxied_transaction is not None:
+            self._proxied_transaction.__exit__(type_, value, traceback)
+            self.migration_context._transaction = None
 
 
 class MigrationContext(object):

diff --git a/docs/build/unreleased/829.rst b/docs/build/unreleased/829.rst
new file mode 100644 (file)
index 0000000..530499c
--- /dev/null
+++ b/docs/build/unreleased/829.rst
@@ -0,0 +1,9 @@
+.. change::
+    :tags: bug, environment
+    :tickets: 829
+
+    Fixed regression caused by the SQLAlchemy 1.4/2.0 compatibility switch
+    where calling ``.rollback()`` or ``.commit()`` explicitly within the
+    ``context.begin_transaction()`` context manager would cause it to fail when
+    the block ended, as it did not expect that the transaction was manually
+    closed.

diff --git a/tests/test_environment.py b/tests/test_environment.py
index e90d6fb0df6490df693b1bc904177cc3d0be484e..23e23b361ba806aca28432dd04515d0dac4c8706 100644 (file)
--- a/tests/test_environment.py
+++ b/tests/test_environment.py
@@ -257,6 +257,34 @@ class MigrationTransactionTest(TestBase):
         assert_raises(Exception, go)
         is_false(self.conn.in_transaction())
 
+    def test_proxy_transaction_contextmanager_explicit_rollback(self):
+        context = self._fixture(
+            {"transaction_per_migration": True, "transactional_ddl": True}
+        )
+        proxy = context.begin_transaction(_per_migration=True)
+        is_true(self.conn.in_transaction())
+
+        with proxy:
+            is_true(self.conn.in_transaction())
+            proxy.rollback()
+            is_false(self.conn.in_transaction())
+
+        is_false(self.conn.in_transaction())
+
+    def test_proxy_transaction_contextmanager_explicit_commit(self):
+        context = self._fixture(
+            {"transaction_per_migration": True, "transactional_ddl": True}
+        )
+        proxy = context.begin_transaction(_per_migration=True)
+        is_true(self.conn.in_transaction())
+
+        with proxy:
+            is_true(self.conn.in_transaction())
+            proxy.commit()
+            is_false(self.conn.in_transaction())
+
+        is_false(self.conn.in_transaction())
+
     def test_transaction_per_migration_transactional_ddl(self):
         context = self._fixture(
             {"transaction_per_migration": True, "transactional_ddl": True}