Memoize elements in rewriter; use correct iteration

Modified the logic of the :class:`.Rewriter` object such that it keeps a
memoization of which directives it has processed, so that it can ensure it
processes a particular directive only once, and additionally fixed
:class:`.Rewriter` so that it functions correctly for multiple-pass
autogenerate schemes, such as the one illustrated in the "multidb"
template.  By tracking which directives have been processed, a
multiple-pass scheme which calls upon the :class:`.Rewriter` multiple times
for the same structure as elements are added can work without running
duplicate operations on the same elements more than once.

Change-Id: I2cef13d51912f9d86ddd99b60e4d5b96dbf680ff
Fixes: #505

diff --git a/alembic/autogenerate/rewriter.py b/alembic/autogenerate/rewriter.py
index 1e9522bd401975abe8927072fec7f5e24540ad34..90a931f22ea96acc91c44e28dbcc54630d802f21 100644 (file)
--- a/alembic/autogenerate/rewriter.py
+++ b/alembic/autogenerate/rewriter.py
@@ -94,10 +94,16 @@ class Rewriter(object):
             _rewriter = None
             yield directive
         else:
-            for r_directive in util.to_list(
-                _rewriter(context, revision, directive)
-            ):
-                yield r_directive
+            if self in directive._mutations:
+                yield directive
+            else:
+                for r_directive in util.to_list(
+                    _rewriter(context, revision, directive)
+                ):
+                    r_directive._mutations = r_directive._mutations.union(
+                        [self]
+                    )
+                    yield r_directive
 
     def __call__(self, context, revision, directives):
         self.process_revision_directives(context, revision, directives)
@@ -108,7 +114,7 @@ class Rewriter(object):
     def _traverse_script(self, context, revision, directive):
         upgrade_ops_list = []
         for upgrade_ops in directive.upgrade_ops_list:
-            ret = self._traverse_for(context, revision, directive.upgrade_ops)
+            ret = self._traverse_for(context, revision, upgrade_ops)
             if len(ret) != 1:
                 raise ValueError(
                     "Can only return single object for UpgradeOps traverse"
@@ -118,9 +124,7 @@ class Rewriter(object):
 
         downgrade_ops_list = []
         for downgrade_ops in directive.downgrade_ops_list:
-            ret = self._traverse_for(
-                context, revision, directive.downgrade_ops
-            )
+            ret = self._traverse_for(context, revision, downgrade_ops)
             if len(ret) != 1:
                 raise ValueError(
                     "Can only return single object for DowngradeOps traverse"

diff --git a/alembic/command.py b/alembic/command.py
index c02db62a8d7d0111040ad6b09f9bc918238f0fa1..f9eb4fd2a6336f85e65d308b778c746b08739b65 100644 (file)
--- a/alembic/command.py
+++ b/alembic/command.py
@@ -199,6 +199,11 @@ def revision(
         ):
             script_directory.run_env()
 
+        # the revision_context now has MigrationScript structure(s) present.
+        # these could theoretically be further processed / rewritten *here*,
+        # in addition to the hooks present within each run_migrations() call,
+        # or at the end of env.py run_migrations_online().
+
     scripts = [script for script in revision_context.generate_scripts()]
     if len(scripts) == 1:
         return scripts[0]

diff --git a/alembic/operations/ops.py b/alembic/operations/ops.py
index c3eb469a466ff08ea528a513ab43b6c724bdebec..8f00b0c3a9e523f0f9fea82b35a9ba4578adbd3d 100644 (file)
--- a/alembic/operations/ops.py
+++ b/alembic/operations/ops.py
@@ -34,6 +34,8 @@ class MigrateOperation(object):
         """
         return {}
 
+    _mutations = frozenset()
+
 
 class AddConstraintOp(MigrateOperation):
     """Represent an add constraint operation."""

diff --git a/docs/build/unreleased/505.rst b/docs/build/unreleased/505.rst
new file mode 100644 (file)
index 0000000..44816a2
--- /dev/null
+++ b/docs/build/unreleased/505.rst
@@ -0,0 +1,13 @@
+.. change::
+    :tags: bug, autogenerate
+    :tickets: 505
+
+    Modified the logic of the :class:`.Rewriter` object such that it keeps a
+    memoization of which directives it has processed, so that it can ensure it
+    processes a particular directive only once, and additionally fixed
+    :class:`.Rewriter` so that it functions correctly for multiple-pass
+    autogenerate schemes, such as the one illustrated in the "multidb"
+    template.  By tracking which directives have been processed, a
+    multiple-pass scheme which calls upon the :class:`.Rewriter` multiple times
+    for the same structure as elements are added can work without running
+    duplicate operations on the same elements more than once.

diff --git a/tests/test_script_production.py b/tests/test_script_production.py
index 05d01335432e63578b6092ef0cb392e0215379b2..220793c8a3b667892600be9340ac6eccad5f29e7 100644 (file)
--- a/tests/test_script_production.py
+++ b/tests/test_script_production.py
@@ -976,6 +976,82 @@ class RewriterTest(TestBase):
             "    # ### end Alembic commands ###",
         )
 
+    def test_multiple_passes_with_mutations(self):
+        writer1 = autogenerate.Rewriter()
+
+        @writer1.rewrites(ops.CreateTableOp)
+        def rewrite_alter_column(context, revision, op):
+            op.table_name += "_pass"
+            return 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()
+        writer1(ctx, rev, directives)
+
+        directives[0].upgrade_ops_list.extend(
+            [
+                ops.UpgradeOps(
+                    ops=[
+                        ops.CreateTableOp(
+                            "another_test_table",
+                            [sa.Column("id", sa.Integer(), primary_key=True)],
+                        )
+                    ]
+                ),
+                ops.UpgradeOps(
+                    ops=[
+                        ops.CreateTableOp(
+                            "third_test_table",
+                            [sa.Column("id", sa.Integer(), primary_key=True)],
+                        )
+                    ]
+                ),
+            ]
+        )
+
+        writer1(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_pass',\n"
+            "    sa.Column('id', sa.Integer(), nullable=False),\n"
+            "    sa.PrimaryKeyConstraint('id')\n"
+            "    )\n"
+            "    # ### end Alembic commands ###",
+        )
+        eq_(
+            autogenerate.render_python_code(directives[0].upgrade_ops_list[1]),
+            "# ### commands auto generated by Alembic - please adjust! ###\n"
+            "    op.create_table('another_test_table_pass',\n"
+            "    sa.Column('id', sa.Integer(), nullable=False),\n"
+            "    sa.PrimaryKeyConstraint('id')\n"
+            "    )\n"
+            "    # ### end Alembic commands ###",
+        )
+        eq_(
+            autogenerate.render_python_code(directives[0].upgrade_ops_list[2]),
+            "# ### commands auto generated by Alembic - please adjust! ###\n"
+            "    op.create_table('third_test_table_pass',\n"
+            "    sa.Column('id', sa.Integer(), nullable=False),\n"
+            "    sa.PrimaryKeyConstraint('id')\n"
+            "    )\n"
+            "    # ### end Alembic commands ###",
+        )
+
 
 class MultiDirRevisionCommandTest(TestBase):
     def setUp(self):