From c972e816befb38466196d89eead2e2bc55f1afae Mon Sep 17 00:00:00 2001 From: Mike Bayer Date: Wed, 15 Aug 2012 15:33:52 -0400 Subject: [PATCH] - [feature] Added include_symbol option to EnvironmentContext.configure(), specifies a callable which will include/exclude tables in their entirety from the autogeneration process based on name. #27 --- CHANGES | 6 ++++++ alembic/autogenerate.py | 17 ++++++++++++----- alembic/environment.py | 17 +++++++++++++++++ tests/test_autogenerate.py | 38 +++++++++++++++++++++++++++++++------- 4 files changed, 66 insertions(+), 12 deletions(-) diff --git a/CHANGES b/CHANGES index 6cc63927..1de83319 100644 --- a/CHANGES +++ b/CHANGES @@ -1,5 +1,11 @@ 0.3.6 ===== +- [feature] Added include_symbol option to + EnvironmentContext.configure(), + specifies a callable which will include/exclude tables + in their entirety from the autogeneration process + based on name. #27 + - [feature] Added year, month, day, hour, minute, second variables to file_template. #59 diff --git a/alembic/autogenerate.py b/alembic/autogenerate.py index 2e365cd1..0eb85b06 100644 --- a/alembic/autogenerate.py +++ b/alembic/autogenerate.py @@ -102,9 +102,11 @@ def compare_metadata(context, metadata): # top level def _produce_migration_diffs(context, template_args, - imports, _include_only=()): + imports, include_symbol=None): opts = context.opts metadata = opts['target_metadata'] + include_symbol = opts.get('include_symbol', include_symbol) + if metadata is None: raise util.CommandError( "Can't proceed with --autogenerate option; environment " @@ -116,7 +118,7 @@ def _produce_migration_diffs(context, template_args, diffs = [] _produce_net_changes(connection, metadata, diffs, - autogen_context, _include_only) + autogen_context, include_symbol) template_args[opts['upgrade_token']] = \ _indent(_produce_upgrade_commands(diffs, autogen_context)) template_args[opts['downgrade_token']] = \ @@ -145,17 +147,22 @@ def _indent(text): # walk structures def _produce_net_changes(connection, metadata, diffs, autogen_context, - include_only=None): + include_symbol=None): inspector = Inspector.from_engine(connection) # TODO: not hardcode alembic_version here ? conn_table_names = set(inspector.get_table_names()).\ difference(['alembic_version']) - if include_only: - conn_table_names = conn_table_names.intersection(include_only) + metadata_table_names = OrderedSet([table.name for table in metadata.sorted_tables]) + if include_symbol: + conn_table_names = set(name for name in conn_table_names + if include_symbol(name)) + metadata_table_names = OrderedSet(name for name in metadata_table_names + if include_symbol(name)) + _compare_tables(conn_table_names, metadata_table_names, inspector, metadata, diffs, autogen_context) diff --git a/alembic/environment.py b/alembic/environment.py index a402bde5..82bd7610 100644 --- a/alembic/environment.py +++ b/alembic/environment.py @@ -208,6 +208,7 @@ class EnvironmentContext(object): tag=None, template_args=None, target_metadata=None, + include_symbol=None, compare_type=False, compare_server_default=False, upgrade_token="upgrades", @@ -355,6 +356,21 @@ class EnvironmentContext(object): execute the two defaults on the database side to compare for equivalence. + :param include_symbol: A callable function which, given a table name + and optional schema name, returns ``True`` or ``False``, indicating + if the given table should be considered in the autogenerate sweep. + E.g.:: + + def include_symbol(tablename, schema=None): + return tablename not in ("skip_table_one", "skip_table_two") + + context.configure( + # ... + include_symbol = include_symbol + ) + + .. versionadded:: 0.3.6 + :param upgrade_token: When autogenerate completes, the text of the candidate upgrade operations will be present in this template variable when ``script.py.mako`` is rendered. Defaults to @@ -408,6 +424,7 @@ class EnvironmentContext(object): if template_args and 'template_args' in opts: opts['template_args'].update(template_args) opts['target_metadata'] = target_metadata + opts['include_symbol'] = include_symbol opts['upgrade_token'] = upgrade_token opts['downgrade_token'] = downgrade_token opts['sqlalchemy_module_prefix'] = sqlalchemy_module_prefix diff --git a/tests/test_autogenerate.py b/tests/test_autogenerate.py index ecf2c3ed..78f72a63 100644 --- a/tests/test_autogenerate.py +++ b/tests/test_autogenerate.py @@ -164,9 +164,10 @@ class ImplicitConstraintNoGenTest(AutogenTest, TestCase): def test_boolean_gen_upgrade(self): template_args = {} autogenerate._produce_migration_diffs(self.context, - template_args, set(), _include_only=['sometable']) + template_args, set(), + include_symbol=lambda name: name == 'sometable') eq_( - template_args['upgrades'], + re.sub(r"u'", "'", template_args['upgrades']), "### commands auto generated by Alembic - please adjust! ###\n" " op.create_table('sometable',\n" " sa.Column('id', sa.Integer(), nullable=False),\n" @@ -181,16 +182,17 @@ class ImplicitConstraintNoGenTest(AutogenTest, TestCase): template_args = {} autogenerate._produce_migration_diffs(self.context, - template_args, set(), _include_only=['someothertable']) + template_args, set(), + ) eq_( - template_args['downgrades'], + re.sub(r"u'", "'", template_args['downgrades']), "### commands auto generated by Alembic - please adjust! ###\n" " op.create_table('someothertable',\n" - " sa.Column(u'id', mysql.INTEGER(display_width=11), " + " sa.Column('id', mysql.INTEGER(display_width=11), " "nullable=False),\n" - " sa.Column(u'value', mysql.TINYINT(display_width=1), " + " sa.Column('value', mysql.TINYINT(display_width=1), " "nullable=True),\n" - " sa.PrimaryKeyConstraint(u'id')\n )\n" + " sa.PrimaryKeyConstraint('id')\n )\n" " op.drop_table('sometable')\n" " ### end Alembic commands ###" ) @@ -331,6 +333,28 @@ class AutogenerateDiffTest(AutogenTest, TestCase): op.drop_table('item') ### end Alembic commands ###""") + def test_include_symbol(self): + context = MigrationContext.configure( + connection=self.bind.connect(), + opts={ + 'compare_type': True, + 'compare_server_default': True, + 'target_metadata': self.m2, + 'include_symbol': lambda name, schema=None: + name in ('address', 'order'), + 'upgrade_token': "upgrades", + 'downgrade_token': "downgrades", + 'alembic_module_prefix': 'op.', + 'sqlalchemy_module_prefix': 'sa.', + } + ) + template_args = {} + autogenerate._produce_migration_diffs(context, template_args, set()) + assert "alter_column('user'" not in template_args['upgrades'] + assert "alter_column('user'" not in template_args['downgrades'] + assert "alter_column('order'" in template_args['upgrades'] + assert "alter_column('order'" in template_args['downgrades'] + def test_skip_null_type_comparison_reflected(self): diff = [] autogenerate._compare_type("sometable", "somecol", -- 2.47.2