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 c64e91f85eb1a38b2fb2ac0acde44032bbde7eda..466264efcfc706e98f06bf28865a51b173dc0856 100644 (file)
--- 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 (file)
index 0000000..79b846e
--- /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 b41c383be6046f0354ca0fe508b28bb8516333cf..dc0f5441a4b7d94d8f5bc07e541a68a5d26a560f 100644 (file)
--- 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()