From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Sat, 17 Apr 2021 16:53:22 +0000 (-0400)
Subject: Ensure proxy transaction still present on exit before closing
X-Git-Tag: rel_1_6_0~6^2
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=7ed092df5c4a16990942333a8a43743d02ce18ff;p=thirdparty%2Fsqlalchemy%2Falembic.git

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 e8a36e9d..bf48e2bd 100644
--- 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
index 00000000..530499c1
--- /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 e90d6fb0..23e23b36 100644
--- 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}