From b6a4e7dd8f9eb274e6420f8385ab91dc09cc27e7 Mon Sep 17 00:00:00 2001 From: Mike Bayer Date: Wed, 18 Sep 2019 21:47:27 -0400 Subject: [PATCH] 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 --- alembic/autogenerate/rewriter.py | 20 +++++---- alembic/command.py | 5 +++ alembic/operations/ops.py | 2 + docs/build/unreleased/505.rst | 13 ++++++ tests/test_script_production.py | 76 ++++++++++++++++++++++++++++++++ 5 files changed, 108 insertions(+), 8 deletions(-) create mode 100644 docs/build/unreleased/505.rst diff --git a/alembic/autogenerate/rewriter.py b/alembic/autogenerate/rewriter.py index 1e9522bd..90a931f2 100644 --- 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 c02db62a..f9eb4fd2 100644 --- 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 c3eb469a..8f00b0c3 100644 --- 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 index 00000000..44816a28 --- /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 05d01335..220793c8 100644 --- 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): -- 2.47.2