Skip to content

Commit

Permalink
- [feature] Added include_symbol option to
Browse files Browse the repository at this point in the history
  EnvironmentContext.configure(),
  specifies a callable which will include/exclude tables
  in their entirety from the autogeneration process
  based on name.  #27
  • Loading branch information
zzzeek committed Aug 15, 2012
1 parent 64a58cc commit c972e81
Show file tree
Hide file tree
Showing 4 changed files with 66 additions and 12 deletions.
6 changes: 6 additions & 0 deletions CHANGES
Original file line number Diff line number Diff line change
@@ -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

Expand Down
17 changes: 12 additions & 5 deletions alembic/autogenerate.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 "
Expand All @@ -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']] = \
Expand Down Expand Up @@ -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)

Expand Down
17 changes: 17 additions & 0 deletions alembic/environment.py
Original file line number Diff line number Diff line change
Expand Up @@ -208,6 +208,7 @@ def configure(self,
tag=None,
template_args=None,
target_metadata=None,
include_symbol=None,
compare_type=False,
compare_server_default=False,
upgrade_token="upgrades",
Expand Down Expand Up @@ -355,6 +356,21 @@ def my_compare_server_default(context, inspected_column,
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
Expand Down Expand Up @@ -408,6 +424,7 @@ def my_compare_server_default(context, inspected_column,
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
Expand Down
38 changes: 31 additions & 7 deletions tests/test_autogenerate.py
Original file line number Diff line number Diff line change
Expand Up @@ -164,9 +164,10 @@ def _get_model_schema(cls):
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"
Expand All @@ -181,16 +182,17 @@ def test_boolean_gen_downgrade(self):

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 ###"
)
Expand Down Expand Up @@ -331,6 +333,28 @@ def test_render_diffs(self):
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",
Expand Down

0 comments on commit c972e81

Please sign in to comment.