From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Mon, 4 Oct 2021 19:15:16 +0000 (-0400)
Subject: frame a transaction around autocommit
X-Git-Tag: rel_1_7_4~1^2
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=ed3428ec969d12befe7bbbe5c8b6b34ffea29e2e;p=thirdparty%2Fsqlalchemy%2Falembic.git

frame a transaction around autocommit

Fixed issue where the :meth:`.MigrationContext.autocommit_block` feature
would fail to function when using a SQLAlchemy engine using 2.0 future
mode.

Change-Id: I851573424c7cde2595ae22c816ec6580d7cab248
Fixes: #944
---

diff --git a/alembic/runtime/migration.py b/alembic/runtime/migration.py
index c64e91f8..466264ef 100644
--- a/alembic/runtime/migration.py
+++ b/alembic/runtime/migration.py
@@ -335,11 +335,19 @@ class MigrationContext:
             self.connection = (
                 self.impl.connection
             ) = base_connection.execution_options(isolation_level="AUTOCOMMIT")
+
+            # sqlalchemy future mode will "autobegin" in any case, so take
+            # control of that "transaction" here
+            fake_trans: Optional[Transaction] = self.connection.begin()
+        else:
+            fake_trans = None
         try:
             yield
         finally:
             if not self.as_sql:
                 assert self.connection is not None
+                if fake_trans is not None:
+                    fake_trans.commit()
                 self.connection.execution_options(
                     isolation_level=current_level
                 )
diff --git a/docs/build/unreleased/944.rst b/docs/build/unreleased/944.rst
new file mode 100644
index 00000000..79b846e9
--- /dev/null
+++ b/docs/build/unreleased/944.rst
@@ -0,0 +1,8 @@
+.. change::
+    :tags: bug, environment
+    :tickets: 944
+
+    Fixed issue where the :meth:`.MigrationContext.autocommit_block` feature
+    would fail to function when using a SQLAlchemy engine using 2.0 future
+    mode.
+
diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py
index b41c383b..dc0f5441 100644
--- a/tests/test_postgresql.py
+++ b/tests/test_postgresql.py
@@ -48,6 +48,7 @@ from alembic.testing.env import clear_staging_env
 from alembic.testing.env import staging_env
 from alembic.testing.env import write_script
 from alembic.testing.fixtures import capture_context_buffer
+from alembic.testing.fixtures import FutureEngineMixin
 from alembic.testing.fixtures import op_fixture
 from alembic.testing.fixtures import TablesTest
 from alembic.testing.fixtures import TestBase
@@ -457,6 +458,10 @@ class PGAutocommitBlockTest(TestBase):
                 )
 
 
+class PGAutocommitBlockTestFuture(FutureEngineMixin, PGAutocommitBlockTest):
+    pass
+
+
 class PGOfflineEnumTest(TestBase):
     def setUp(self):
         staging_env()