From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Sat, 18 Nov 2023 14:58:13 +0000 (-0500)
Subject: add batch template for render create_table_comment, drop_table_comment
X-Git-Tag: rel_1_13_0~7^2
X-Git-Url: http://git.ipfire.org/gitweb/gitweb.cgi?a=commitdiff_plain;h=665f50f742517fd5bae9805a357bf5c44ea2e171;p=thirdparty%2Fsqlalchemy%2Falembic.git

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 9c84cd6c..67cc8c33 100644
--- 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
index 00000000..e744dad5
--- /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 88aa978c..eeeb92ed 100644
--- 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"