Skip to content

Commit c972e81

Browse files
committed
- [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
1 parent 64a58cc commit c972e81

File tree

4 files changed

+66
-12
lines changed

4 files changed

+66
-12
lines changed

CHANGES

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,11 @@
11
0.3.6
22
=====
3+
- [feature] Added include_symbol option to
4+
EnvironmentContext.configure(),
5+
specifies a callable which will include/exclude tables
6+
in their entirety from the autogeneration process
7+
based on name. #27
8+
39
- [feature] Added year, month, day, hour, minute, second
410
variables to file_template. #59
511

alembic/autogenerate.py

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -102,9 +102,11 @@ def compare_metadata(context, metadata):
102102
# top level
103103

104104
def _produce_migration_diffs(context, template_args,
105-
imports, _include_only=()):
105+
imports, include_symbol=None):
106106
opts = context.opts
107107
metadata = opts['target_metadata']
108+
include_symbol = opts.get('include_symbol', include_symbol)
109+
108110
if metadata is None:
109111
raise util.CommandError(
110112
"Can't proceed with --autogenerate option; environment "
@@ -116,7 +118,7 @@ def _produce_migration_diffs(context, template_args,
116118

117119
diffs = []
118120
_produce_net_changes(connection, metadata, diffs,
119-
autogen_context, _include_only)
121+
autogen_context, include_symbol)
120122
template_args[opts['upgrade_token']] = \
121123
_indent(_produce_upgrade_commands(diffs, autogen_context))
122124
template_args[opts['downgrade_token']] = \
@@ -145,17 +147,22 @@ def _indent(text):
145147
# walk structures
146148

147149
def _produce_net_changes(connection, metadata, diffs, autogen_context,
148-
include_only=None):
150+
include_symbol=None):
149151
inspector = Inspector.from_engine(connection)
150152
# TODO: not hardcode alembic_version here ?
151153
conn_table_names = set(inspector.get_table_names()).\
152154
difference(['alembic_version'])
153-
if include_only:
154-
conn_table_names = conn_table_names.intersection(include_only)
155+
155156

156157
metadata_table_names = OrderedSet([table.name
157158
for table in metadata.sorted_tables])
158159

160+
if include_symbol:
161+
conn_table_names = set(name for name in conn_table_names
162+
if include_symbol(name))
163+
metadata_table_names = OrderedSet(name for name in metadata_table_names
164+
if include_symbol(name))
165+
159166
_compare_tables(conn_table_names, metadata_table_names,
160167
inspector, metadata, diffs, autogen_context)
161168

alembic/environment.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -208,6 +208,7 @@ def configure(self,
208208
tag=None,
209209
template_args=None,
210210
target_metadata=None,
211+
include_symbol=None,
211212
compare_type=False,
212213
compare_server_default=False,
213214
upgrade_token="upgrades",
@@ -355,6 +356,21 @@ def my_compare_server_default(context, inspected_column,
355356
execute
356357
the two defaults on the database side to compare for equivalence.
357358
359+
:param include_symbol: A callable function which, given a table name
360+
and optional schema name, returns ``True`` or ``False``, indicating
361+
if the given table should be considered in the autogenerate sweep.
362+
E.g.::
363+
364+
def include_symbol(tablename, schema=None):
365+
return tablename not in ("skip_table_one", "skip_table_two")
366+
367+
context.configure(
368+
# ...
369+
include_symbol = include_symbol
370+
)
371+
372+
.. versionadded:: 0.3.6
373+
358374
:param upgrade_token: When autogenerate completes, the text of the
359375
candidate upgrade operations will be present in this template
360376
variable when ``script.py.mako`` is rendered. Defaults to
@@ -408,6 +424,7 @@ def my_compare_server_default(context, inspected_column,
408424
if template_args and 'template_args' in opts:
409425
opts['template_args'].update(template_args)
410426
opts['target_metadata'] = target_metadata
427+
opts['include_symbol'] = include_symbol
411428
opts['upgrade_token'] = upgrade_token
412429
opts['downgrade_token'] = downgrade_token
413430
opts['sqlalchemy_module_prefix'] = sqlalchemy_module_prefix

tests/test_autogenerate.py

Lines changed: 31 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -164,9 +164,10 @@ def _get_model_schema(cls):
164164
def test_boolean_gen_upgrade(self):
165165
template_args = {}
166166
autogenerate._produce_migration_diffs(self.context,
167-
template_args, set(), _include_only=['sometable'])
167+
template_args, set(),
168+
include_symbol=lambda name: name == 'sometable')
168169
eq_(
169-
template_args['upgrades'],
170+
re.sub(r"u'", "'", template_args['upgrades']),
170171
"### commands auto generated by Alembic - please adjust! ###\n"
171172
" op.create_table('sometable',\n"
172173
" sa.Column('id', sa.Integer(), nullable=False),\n"
@@ -181,16 +182,17 @@ def test_boolean_gen_downgrade(self):
181182

182183
template_args = {}
183184
autogenerate._produce_migration_diffs(self.context,
184-
template_args, set(), _include_only=['someothertable'])
185+
template_args, set(),
186+
)
185187
eq_(
186-
template_args['downgrades'],
188+
re.sub(r"u'", "'", template_args['downgrades']),
187189
"### commands auto generated by Alembic - please adjust! ###\n"
188190
" op.create_table('someothertable',\n"
189-
" sa.Column(u'id', mysql.INTEGER(display_width=11), "
191+
" sa.Column('id', mysql.INTEGER(display_width=11), "
190192
"nullable=False),\n"
191-
" sa.Column(u'value', mysql.TINYINT(display_width=1), "
193+
" sa.Column('value', mysql.TINYINT(display_width=1), "
192194
"nullable=True),\n"
193-
" sa.PrimaryKeyConstraint(u'id')\n )\n"
195+
" sa.PrimaryKeyConstraint('id')\n )\n"
194196
" op.drop_table('sometable')\n"
195197
" ### end Alembic commands ###"
196198
)
@@ -331,6 +333,28 @@ def test_render_diffs(self):
331333
op.drop_table('item')
332334
### end Alembic commands ###""")
333335

336+
def test_include_symbol(self):
337+
context = MigrationContext.configure(
338+
connection=self.bind.connect(),
339+
opts={
340+
'compare_type': True,
341+
'compare_server_default': True,
342+
'target_metadata': self.m2,
343+
'include_symbol': lambda name, schema=None:
344+
name in ('address', 'order'),
345+
'upgrade_token': "upgrades",
346+
'downgrade_token': "downgrades",
347+
'alembic_module_prefix': 'op.',
348+
'sqlalchemy_module_prefix': 'sa.',
349+
}
350+
)
351+
template_args = {}
352+
autogenerate._produce_migration_diffs(context, template_args, set())
353+
assert "alter_column('user'" not in template_args['upgrades']
354+
assert "alter_column('user'" not in template_args['downgrades']
355+
assert "alter_column('order'" in template_args['upgrades']
356+
assert "alter_column('order'" in template_args['downgrades']
357+
334358
def test_skip_null_type_comparison_reflected(self):
335359
diff = []
336360
autogenerate._compare_type("sometable", "somecol",

0 commit comments

Comments
 (0)