Merge "add batch template for render create_table_comment, drop_table…

…_comment" into main
sqlalchemy · Nov 20, 2023 · 7007547 · 7007547
2 parents 70e1fba + 665f50f
commit 7007547
Show file tree

Hide file tree

Showing 3 changed files with 100 additions and 15 deletions.
diff --git 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
@@ -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
@@ -2070,6 +2070,38 @@ def test_render_create_table_comment_op_with_schema(self):
             ")",
         )
 
+    @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 @@ def test_render_drop_table_comment_op(self):
             ")",
         )
 
+    @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"