From: Sebastian Bayer (AE/MFD2-SO) <sebastian.bayer@de.bosch.com>
Date: Thu, 26 Oct 2023 13:18:48 +0000 (-0400)
Subject: Implement ExecuteSQLOp.to_diff_tuple
X-Git-Tag: rel_1_12_1~1^2
X-Git-Url: http://git.ipfire.org/?a=commitdiff_plain;h=1033b9998f3fe7177d320e4d105c5b43ed4ca50b;p=thirdparty%2Fsqlalchemy%2Falembic.git

Implement ExecuteSQLOp.to_diff_tuple

Repaired :class:`.ExecuteSQLOp` so that it can participate in "diff"
operations; while this object is typically not present in a reflected
operation stream, custom hooks may be adding this construct where it needs
to have the correct ``to_diff_tuple()`` method.  Pull request courtesy
Sebastian Bayer.

Fixes: #1335
Closes: #1336
Pull-request: https://github.com/sqlalchemy/alembic/pull/1336
Pull-request-sha: c5e217e99d0fae87ccc89fc57c71e1c07c57ea97

Change-Id: I303f1a6b1b7bad52f00c6183b77f5e7fe6d0ef37
---

diff --git a/alembic/operations/ops.py b/alembic/operations/ops.py
index fe681217..711d7aba 100644
--- a/alembic/operations/ops.py
+++ b/alembic/operations/ops.py
@@ -2539,6 +2539,9 @@ class ExecuteSQLOp(MigrateOperation):
             operations, sqltext, execution_options=execution_options
         )
 
+    def to_diff_tuple(self) -> Tuple[str, Union[Executable, str]]:
+        return ("execute", self.sqltext)
+
 
 class OpContainer(MigrateOperation):
     """Represent a sequence of operations operation."""
diff --git a/docs/build/unreleased/1335.rst b/docs/build/unreleased/1335.rst
new file mode 100644
index 00000000..ce4ab573
--- /dev/null
+++ b/docs/build/unreleased/1335.rst
@@ -0,0 +1,9 @@
+.. change::
+    :tags: bug, operations
+    :tickets: 1335
+
+    Repaired :class:`.ExecuteSQLOp` so that it can participate in "diff"
+    operations; while this object is typically not present in a reflected
+    operation stream, custom hooks may be adding this construct where it needs
+    to have the correct ``to_diff_tuple()`` method.  Pull request courtesy
+    Sebastian Bayer.
diff --git a/tests/test_autogen_diffs.py b/tests/test_autogen_diffs.py
index 2a622aee..197bf39a 100644
--- a/tests/test_autogen_diffs.py
+++ b/tests/test_autogen_diffs.py
@@ -696,6 +696,17 @@ class AutogenerateDiffTest(ModelOne, AutogenTest, TestBase):
             ],
         )
 
+    def test_add_execute_sql_op(self):
+        uo = ops.UpgradeOps(ops=[])
+        autogenerate._produce_net_changes(self.autogen_context, uo)
+
+        uo.ops.append(ops.ExecuteSQLOp("STATEMENT"))
+
+        diffs = uo.as_diffs()
+
+        eq_(diffs[-1][0], "execute")
+        eq_(diffs[-1][1], "STATEMENT")
+
 
 class AutogenerateDiffTestWSchema(ModelOne, AutogenTest, TestBase):
     __only_on__ = "postgresql"
diff --git a/tests/test_script_production.py b/tests/test_script_production.py
index d50f7f50..ccc1a107 100644
--- a/tests/test_script_production.py
+++ b/tests/test_script_production.py
@@ -1090,6 +1090,50 @@ class RewriterTest(TestBase):
             "    # ### end Alembic commands ###",
         )
 
+    def test_add_execute_sql(self):
+        writer = autogenerate.Rewriter()
+
+        @writer.rewrites(ops.CreateTableOp)
+        def rewriter_execute_sql(context, revision, op):
+            execute_op = ops.ExecuteSQLOp(sqltext="STATEMENT")
+            return [op, execute_op]
+
+        directives = [
+            ops.MigrationScript(
+                util.rev_id(),
+                ops.UpgradeOps(
+                    ops=[
+                        ops.CreateTableOp(
+                            "test_table",
+                            [sa.Column("id", sa.Integer(), primary_key=True)],
+                        )
+                    ]
+                ),
+                ops.DowngradeOps(ops=[]),
+            )
+        ]
+
+        ctx, rev = mock.Mock(), mock.Mock()
+        writer(ctx, rev, directives)
+
+        eq_(
+            autogenerate.render_python_code(directives[0].upgrade_ops_list[0]),
+            "# ### commands auto generated by Alembic - please adjust! ###\n"
+            "    op.create_table('test_table',\n"
+            "    sa.Column('id', sa.Integer(), nullable=False),\n"
+            "    sa.PrimaryKeyConstraint('id')\n"
+            "    )\n"
+            "    op.execute('STATEMENT')\n"
+            "    # ### end Alembic commands ###",
+        )
+
+        diffs = directives[0].upgrade_ops_list[0].as_diffs()
+
+        eq_(diffs[0][0], "add_table")
+
+        eq_(diffs[1][0], "execute")
+        eq_(diffs[1][1], "STATEMENT")
+
 
 class MultiDirRevisionCommandTest(TestBase):
     def setUp(self):