From ed3428ec969d12befe7bbbe5c8b6b34ffea29e2e Mon Sep 17 00:00:00 2001
From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Mon, 4 Oct 2021 15:15:16 -0400
Subject: [PATCH] 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
---
 alembic/runtime/migration.py  | 8 ++++++++
 docs/build/unreleased/944.rst | 8 ++++++++
 tests/test_postgresql.py      | 5 +++++
 3 files changed, 21 insertions(+)
 create mode 100644 docs/build/unreleased/944.rst

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 @@ def upgrade():
             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 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 @@ def test_alter_enum(self, migration_context):
                 )
 
 
+class PGAutocommitBlockTestFuture(FutureEngineMixin, PGAutocommitBlockTest):
+    pass
+
+
 class PGOfflineEnumTest(TestBase):
     def setUp(self):
         staging_env()