Implement ExecuteSQLOp.to_diff_tuple, fixes #1335

alembic/operations/ops.py

@@ -2539,6 +2539,9 @@ def batch_execute(
             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."""

docs/build/unreleased/1335.rst

@@ -0,0 +1,6 @@
+.. change::
+    :tags: usecase
+    :tickets: 1335
+
+    The check command now works when there are custom SQL statements added to the
+    diff stream via :class:`.ExecuteSQLOp`.

tests/test_autogen_diffs.py

@@ -696,6 +696,17 @@ def test_dont_barf_on_already_reflected(self):
             ],
         )
 
+    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"

tests/test_script_production.py

@@ -1090,6 +1090,50 @@ def rewrite_alter_column(context, revision, op):
             "    # ### 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):