add batch template for render create_table_comment, drop_table_comment

Fixed autogenerate issue where ``create_table_comment()`` and
``drop_table_comment()`` rendering in a batch table modify would include
the "table" and "schema" arguments, which are not accepted in batch as
these are already part of the top level block.

Fixes: #1361
Change-Id: I12a01ee2a69ba89138dfda7bfda04bbc83929550

diff --git a/alembic/autogenerate/render.py b/alembic/autogenerate/render.py
index 9c84cd6c51745ed9ab68412e4a2d9f8bb3f08612..67cc8c33dda0fda1df3576e5810c3d767ee1dd2e 100644 (file)
--- a/alembic/autogenerate/render.py
+++ b/alembic/autogenerate/render.py
@@ -164,14 +164,22 @@ def _render_modify_table(
 def _render_create_table_comment(
     autogen_context: AutogenContext, op: ops.CreateTableCommentOp
 ) -> str:
-    templ = (
-        "{prefix}create_table_comment(\n"
-        "{indent}'{tname}',\n"
-        "{indent}{comment},\n"
-        "{indent}existing_comment={existing},\n"
-        "{indent}schema={schema}\n"
-        ")"
-    )
+    if autogen_context._has_batch:
+        templ = (
+            "{prefix}create_table_comment(\n"
+            "{indent}{comment},\n"
+            "{indent}existing_comment={existing}\n"
+            ")"
+        )
+    else:
+        templ = (
+            "{prefix}create_table_comment(\n"
+            "{indent}'{tname}',\n"
+            "{indent}{comment},\n"
+            "{indent}existing_comment={existing},\n"
+            "{indent}schema={schema}\n"
+            ")"
+        )
     return templ.format(
         prefix=_alembic_autogenerate_prefix(autogen_context),
         tname=op.table_name,
@@ -188,13 +196,20 @@ def _render_create_table_comment(
 def _render_drop_table_comment(
     autogen_context: AutogenContext, op: ops.DropTableCommentOp
 ) -> str:
-    templ = (
-        "{prefix}drop_table_comment(\n"
-        "{indent}'{tname}',\n"
-        "{indent}existing_comment={existing},\n"
-        "{indent}schema={schema}\n"
-        ")"
-    )
+    if autogen_context._has_batch:
+        templ = (
+            "{prefix}drop_table_comment(\n"
+            "{indent}existing_comment={existing}\n"
+            ")"
+        )
+    else:
+        templ = (
+            "{prefix}drop_table_comment(\n"
+            "{indent}'{tname}',\n"
+            "{indent}existing_comment={existing},\n"
+            "{indent}schema={schema}\n"
+            ")"
+        )
     return templ.format(
         prefix=_alembic_autogenerate_prefix(autogen_context),
         tname=op.table_name,

diff --git a/docs/build/unreleased/1361.rst b/docs/build/unreleased/1361.rst
new file mode 100644 (file)
index 0000000..e744dad
--- /dev/null
+++ b/docs/build/unreleased/1361.rst
@@ -0,0 +1,8 @@
+.. change::
+    :tags: bug, autogenerate
+    :tickets: 1361
+
+    Fixed autogenerate issue where ``create_table_comment()`` and
+    ``drop_table_comment()`` rendering in a batch table modify would include
+    the "table" and "schema" arguments, which are not accepted in batch as
+    these are already part of the top level block.

diff --git a/tests/test_autogen_render.py b/tests/test_autogen_render.py
index 88aa978cc93654f74437ad33431de54aad4a5f02..eeeb92ed1c3d9ad5d353aff0bc81f4ec882f073c 100644 (file)
--- a/tests/test_autogen_render.py
+++ b/tests/test_autogen_render.py
@@ -2070,6 +2070,38 @@ class AutogenRenderTest(TestBase):
             ")",
         )
 
+    @testing.combinations((True,), (False,), argnames="use_schema")
+    def test_render_create_table_comment_op_batch(self, use_schema):
+        """test #1361"""
+        uo = ops.UpgradeOps(
+            ops=[
+                ops.ModifyTableOps(
+                    "table_name",
+                    schema="SomeSchema" if use_schema else None,
+                    ops=[
+                        ops.CreateTableCommentOp(
+                            "table_name",
+                            "comment",
+                            schema="SomeSchema" if use_schema else None,
+                        )
+                    ],
+                )
+            ]
+        )
+
+        eq_(
+            autogenerate.render_python_code(uo, render_as_batch=True),
+            "# ### commands auto generated by Alembic - please adjust! ###\n"
+            "    with op.batch_alter_table('table_name', "
+            f"schema={repr('SomeSchema' if use_schema else None)}) "
+            "as batch_op:\n"
+            "        batch_op.create_table_comment(\n"
+            "        'comment',\n"
+            "        existing_comment=None\n"
+            "    )\n\n"
+            "    # ### end Alembic commands ###",
+        )
+
     def test_render_drop_table_comment_op(self):
         op_obj = ops.DropTableCommentOp("table_name")
         eq_ignore_whitespace(
@@ -2081,6 +2113,36 @@ class AutogenRenderTest(TestBase):
             ")",
         )
 
+    @testing.combinations((True,), (False,), argnames="use_schema")
+    def test_render_drop_table_comment_op_batch(self, use_schema):
+        """test #1361"""
+        uo = ops.UpgradeOps(
+            ops=[
+                ops.ModifyTableOps(
+                    "table_name",
+                    schema="SomeSchema" if use_schema else None,
+                    ops=[
+                        ops.DropTableCommentOp(
+                            "table_name",
+                            schema="SomeSchema" if use_schema else None,
+                        )
+                    ],
+                )
+            ]
+        )
+
+        eq_(
+            autogenerate.render_python_code(uo, render_as_batch=True),
+            "# ### commands auto generated by Alembic - please adjust! ###\n"
+            "    with op.batch_alter_table('table_name', "
+            f"schema={repr('SomeSchema' if use_schema else None)}) "
+            "as batch_op:\n"
+            "        batch_op.drop_table_comment(\n"
+            "        existing_comment=None\n"
+            "    )\n\n"
+            "    # ### end Alembic commands ###",
+        )
+
     def test_render_drop_table_comment_op_existing_with_quote(self):
         op_obj = ops.DropTableCommentOp(
             "table_name", existing_comment="This was john's comment"