fail gracefully for batch_alter_table() called in --sql mode

Added an error raise for the condition where
:meth:`.Operations.batch_alter_table` is used in ``--sql`` mode, where the
operation requires table reflection, as is the case when running against
SQLite without giving it a fixed ``Table`` object. Previously the operation
would fail with an internal error.   To get a "move and copy" batch
operation as a SQL script without connecting to a database,
a ``Table`` object should be passed to the
:paramref:`.Operations.batch_alter_table.copy_from` parameter so that
reflection may be skipped.

Change-Id: I2d040e7e5771eefabba1649d71ed451567c25283
Fixes: #1021

diff --git a/alembic/operations/batch.py b/alembic/operations/batch.py
index 7c7de9fbf1cdf24dac091d67960726e43fac454e..71d26816f4dab70d556d602fda7591b0716c83e9 100644 (file)
--- a/alembic/operations/batch.py
+++ b/alembic/operations/batch.py
@@ -119,6 +119,21 @@ class BatchOperationsImpl:
                     existing_table = self.copy_from
                     reflected = False
                 else:
+                    if self.operations.migration_context.as_sql:
+                        raise exc.CommandError(
+                            f"This operation cannot proceed in --sql mode; "
+                            f"batch mode with dialect "
+                            f"{self.operations.migration_context.dialect.name} "  # noqa: E501
+                            f"requires a live database connection with which "
+                            f'to reflect the table "{self.table_name}". '
+                            f"To generate a batch SQL migration script using "
+                            "table "
+                            '"move and copy", a complete Table object '
+                            f'should be passed to the "copy_from" argument '
+                            "of the batch_alter_table() method so that table "
+                            "reflection can be skipped."
+                        )
+
                     existing_table = Table(
                         self.table_name,
                         m1,

diff --git a/alembic/util/sqla_compat.py b/alembic/util/sqla_compat.py
index 9c06e857e495fe022a63dda351bd57504483d63a..179e5482394859a355cdd7413385a4d5a908d7a2 100644 (file)
--- a/alembic/util/sqla_compat.py
+++ b/alembic/util/sqla_compat.py
@@ -114,6 +114,11 @@ def _ensure_scope_for_ddl(
         in_transaction = connection.in_transaction  # type: ignore[union-attr]
     except AttributeError:
         # catch for MockConnection, None
+        in_transaction = None
+        pass
+
+    # yield outside the catch
+    if in_transaction is None:
         yield
     else:
         if not in_transaction():

diff --git a/docs/build/unreleased/1021.rst b/docs/build/unreleased/1021.rst
new file mode 100644 (file)
index 0000000..3d9752c
--- /dev/null
+++ b/docs/build/unreleased/1021.rst
@@ -0,0 +1,13 @@
+.. change::
+    :tags: bug, batch
+    :tickets: 1021
+
+    Added an error raise for the condition where
+    :meth:`.Operations.batch_alter_table` is used in ``--sql`` mode, where the
+    operation requires table reflection, as is the case when running against
+    SQLite without giving it a fixed ``Table`` object. Previously the operation
+    would fail with an internal error.   To get a "move and copy" batch
+    operation as a SQL script without connecting to a database,
+    a ``Table`` object should be passed to the
+    :paramref:`.Operations.batch_alter_table.copy_from` parameter so that
+    reflection may be skipped.
\ No newline at end of file

diff --git a/tests/test_batch.py b/tests/test_batch.py
index adf76cd5952db81479496a27b944a12a64b994ba..2d29f6c6c87d88883ba0bcb45bd354aef21bccd6 100644 (file)
--- a/tests/test_batch.py
+++ b/tests/test_batch.py
@@ -25,19 +25,29 @@ from sqlalchemy.schema import CreateTable
 from sqlalchemy.sql import column
 from sqlalchemy.sql import text
 
+from alembic import command
 from alembic import testing
+from alembic import util
 from alembic.ddl import sqlite
 from alembic.operations import Operations
 from alembic.operations.batch import ApplyBatchImpl
 from alembic.runtime.migration import MigrationContext
+from alembic.script import ScriptDirectory
 from alembic.testing import assert_raises_message
 from alembic.testing import config
 from alembic.testing import eq_
 from alembic.testing import exclusions
+from alembic.testing import expect_raises_message
 from alembic.testing import is_
 from alembic.testing import mock
 from alembic.testing import TestBase
+from alembic.testing.env import _no_sql_testing_config
+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 op_fixture
+from alembic.util import CommandError
 from alembic.util import exc as alembic_exc
 from alembic.util.sqla_compat import _safe_commit_connection_transaction
 from alembic.util.sqla_compat import _select
@@ -2384,3 +2394,164 @@ class BatchRoundTripPostgresqlTest(BatchRoundTripTest):
             ],
             [Boolean],
         )
+
+
+class OfflineTest(TestBase):
+    @testing.fixture
+    def no_reflect_batch_fixture(self):
+        staging_env()
+
+        def go():
+            self.cfg = cfg = _no_sql_testing_config(dialect="sqlite")
+
+            self.a = a = util.rev_id()
+
+            script = ScriptDirectory.from_config(cfg)
+            script.generate_revision(
+                a, "revision a", refresh=True, head="base"
+            )
+            write_script(
+                script,
+                a,
+                """\
+    "Rev A"
+    revision = '%s'
+    down_revision = None
+
+    from alembic import op
+    from sqlalchemy import Column
+    from sqlalchemy import Integer
+    from sqlalchemy import String, Table, MetaData
+
+    some_table_up = Table(
+        "some_table", MetaData(),
+        Column('id', Integer),
+        Column('bar', String)
+    )
+
+    some_table_down = Table(
+        "some_table", MetaData(),
+        Column('id', Integer),
+        Column('foo', Integer)
+    )
+
+    def upgrade():
+        with op.batch_alter_table("some_table", copy_from=some_table_up) as batch_op:
+            batch_op.add_column(Column('foo', Integer))
+            batch_op.drop_column('bar')
+
+    def downgrade():
+        with op.batch_alter_table("some_table", copy_from=some_table_down) as batch_op:
+            batch_op.drop_column('foo')
+            batch_op.add_column(Column('bar', String))
+
+    """  # noqa: E501
+                % a,
+            )
+
+        yield go
+        clear_staging_env()
+
+    @testing.fixture
+    def batch_fixture(self):
+        staging_env()
+
+        def go(dialect):
+            self.cfg = cfg = _no_sql_testing_config(dialect=dialect)
+
+            self.a = a = util.rev_id()
+
+            script = ScriptDirectory.from_config(cfg)
+            script.generate_revision(
+                a, "revision a", refresh=True, head="base"
+            )
+            write_script(
+                script,
+                a,
+                """\
+    "Rev A"
+    revision = '%s'
+    down_revision = None
+
+    from alembic import op
+    from sqlalchemy import Column
+    from sqlalchemy import Integer
+    from sqlalchemy import String
+
+    def upgrade():
+        with op.batch_alter_table("some_table") as batch_op:
+            batch_op.add_column(Column('foo', Integer))
+            batch_op.drop_column('bar')
+
+    def downgrade():
+        with op.batch_alter_table("some_table") as batch_op:
+            batch_op.drop_column('foo')
+            batch_op.add_column(Column('bar', String))
+
+    """
+                % a,
+            )
+
+        yield go
+        clear_staging_env()
+
+    def test_upgrade_non_batch(self, batch_fixture):
+        batch_fixture("postgresql")
+
+        with capture_context_buffer() as buf:
+            command.upgrade(self.cfg, self.a, sql=True)
+
+        assert re.search(
+            r"ALTER TABLE some_table ADD COLUMN foo INTEGER", buf.getvalue()
+        )
+
+    def test_downgrade_non_batch(self, batch_fixture):
+        batch_fixture("postgresql")
+
+        with capture_context_buffer() as buf:
+            command.downgrade(self.cfg, f"{self.a}:base", sql=True)
+        assert re.search(
+            r"ALTER TABLE some_table DROP COLUMN foo", buf.getvalue()
+        )
+
+    def test_upgrade_batch_fails_gracefully(self, batch_fixture):
+        batch_fixture("sqlite")
+
+        with expect_raises_message(
+            CommandError,
+            "This operation cannot proceed in --sql mode; batch mode with "
+            "dialect sqlite requires a live database connection with which "
+            'to reflect the table "some_table"',
+        ):
+            command.upgrade(self.cfg, self.a, sql=True)
+
+    def test_downgrade_batch_fails_gracefully(self, batch_fixture):
+        batch_fixture("sqlite")
+
+        with expect_raises_message(
+            CommandError,
+            "This operation cannot proceed in --sql mode; batch mode with "
+            "dialect sqlite requires a live database connection with which "
+            'to reflect the table "some_table"',
+        ):
+            command.downgrade(self.cfg, f"{self.a}:base", sql=True)
+
+    def test_upgrade_batch_no_reflection(self, no_reflect_batch_fixture):
+        no_reflect_batch_fixture()
+
+        with capture_context_buffer() as buf:
+            command.upgrade(self.cfg, self.a, sql=True)
+
+        assert re.search(
+            r"CREATE TABLE _alembic_tmp_some_table", buf.getvalue()
+        )
+
+    def test_downgrade_batch_no_reflection(self, no_reflect_batch_fixture):
+        no_reflect_batch_fixture()
+
+        with capture_context_buffer() as buf:
+            command.downgrade(self.cfg, f"{self.a}:base", sql=True)
+
+        assert re.search(
+            r"CREATE TABLE _alembic_tmp_some_table", buf.getvalue()
+        )