From: Mike Bayer Date: Fri, 7 Nov 2014 00:37:49 +0000 (-0500) Subject: - add support for autogenerate to include "batch" X-Git-Tag: rel_0_7_0~54 X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=0933f08dc5877e7e1bb2597d8f9326a29d3e822a;p=thirdparty%2Fsqlalchemy%2Falembic.git - add support for autogenerate to include "batch" --- diff --git a/alembic/autogenerate/api.py b/alembic/autogenerate/api.py index 3519a6d1..3f2bc8b4 100644 --- a/alembic/autogenerate/api.py +++ b/alembic/autogenerate/api.py @@ -2,6 +2,7 @@ automatically.""" import logging +import itertools import re from ..compat import StringIO @@ -245,17 +246,30 @@ def _produce_net_changes(connection, metadata, diffs, autogen_context, def _produce_upgrade_commands(diffs, autogen_context): - if diffs: - for diff in diffs: - yield _invoke_command("upgrade", diff, autogen_context) - else: - yield "pass" + return _produce_commands("upgrade", diffs, autogen_context) def _produce_downgrade_commands(diffs, autogen_context): + return _produce_commands("downgrade", diffs, autogen_context) + + +def _produce_commands(type_, diffs, autogen_context): + opts = autogen_context['opts'] + render_as_batch = opts.get('render_as_batch', False) + if diffs: - for diff in reversed(diffs): - yield _invoke_command("downgrade", diff, autogen_context) + if type_ == 'downgrade': + diffs = reversed(diffs) + for (schema, table), subdiffs in _group_diffs_by_table(diffs): + if table is not None and render_as_batch: + yield "with op.batch_alter_table"\ + "(%r, schema=%r) as batch_op:" % (table, schema) + autogen_context['batch_prefix'] = 'batch_op.' + for diff in subdiffs: + yield _invoke_command(type_, diff, autogen_context) + if table is not None and render_as_batch: + del autogen_context['batch_prefix'] + yield "" else: yield "pass" @@ -321,3 +335,23 @@ def _invoke_modify_command(updown, args, autogen_context): if "server_default" in kw: kw.pop("existing_server_default", None) return _modify_col(tname, cname, autogen_context, schema=sname, **kw) + + +def _group_diffs_by_table(diffs): + _adddrop = { + "table": lambda diff: (None, None), + "column": lambda diff: (diff[0], diff[1]), + "index": lambda diff: (diff[0].table.schema, diff[0].table.name), + "constraint": lambda diff: (diff[0].table.schema, diff[0].table.name) + } + + def _derive_table(diff): + if isinstance(diff, tuple): + cmd_type = diff[0] + adddrop, cmd_type = cmd_type.split("_") + return _adddrop[cmd_type](diff[1:]) + else: + sname, tname = diff[0][1:3] + return sname, tname + + return itertools.groupby(diffs, _derive_table) diff --git a/alembic/autogenerate/render.py b/alembic/autogenerate/render.py index 0344b774..5fdfe366 100644 --- a/alembic/autogenerate/render.py +++ b/alembic/autogenerate/render.py @@ -165,19 +165,23 @@ def _add_unique_constraint(constraint, autogen_context): def _uq_constraint(constraint, autogen_context, alter): opts = [] + + has_batch = 'batch_prefix' in autogen_context + if constraint.deferrable: opts.append(("deferrable", str(constraint.deferrable))) if constraint.initially: opts.append(("initially", str(constraint.initially))) - if alter and constraint.table.schema: + if not has_batch and alter and constraint.table.schema: opts.append(("schema", str(constraint.table.schema))) if not alter and constraint.name: opts.append( ("name", _render_gen_name(autogen_context, constraint.name))) if alter: - args = [repr(_render_gen_name(autogen_context, constraint.name)), - repr(constraint.table.name)] + args = [repr(_render_gen_name(autogen_context, constraint.name))] + if not has_batch: + args += [repr(constraint.table.name)] args.append(repr([col.name for col in constraint.columns])) args.extend(["%s=%r" % (k, v) for k, v in opts]) return "%(prefix)screate_unique_constraint(%(args)s)" % { @@ -224,38 +228,55 @@ def _drop_constraint(constraint, autogen_context): Generate Alembic operations for the ALTER TABLE ... DROP CONSTRAINT of a :class:`~sqlalchemy.schema.UniqueConstraint` instance. """ - text = "%(prefix)sdrop_constraint"\ - "(%(name)r, '%(table_name)s'%(schema)s)" % { - 'prefix': _alembic_autogenerate_prefix(autogen_context), - 'name': _render_gen_name(autogen_context, constraint.name), - 'table_name': constraint.table.name, - 'schema': (", schema='%s'" % constraint.table.schema) - if constraint.table.schema else '', - } + if 'batch_prefix' in autogen_context: + template = "%(prefix)sdrop_constraint"\ + "(%(name)r)" + else: + template = "%(prefix)sdrop_constraint"\ + "(%(name)r, '%(table_name)s'%(schema)s)" + + text = template % { + 'prefix': _alembic_autogenerate_prefix(autogen_context), + 'name': _render_gen_name(autogen_context, constraint.name), + 'table_name': constraint.table.name, + 'schema': (", schema='%s'" % constraint.table.schema) + if constraint.table.schema else '', + } return text def _add_column(schema, tname, column, autogen_context): - text = "%(prefix)sadd_column(%(tname)r, %(column)s" % { + if 'batch_prefix' in autogen_context: + template = "%(prefix)sadd_column(%(column)s)" + else: + template = "%(prefix)sadd_column(%(tname)r, %(column)s" + if schema: + template += ", schema=%(schema)r" + template += ")" + text = template % { "prefix": _alembic_autogenerate_prefix(autogen_context), "tname": tname, - "column": _render_column(column, autogen_context) + "column": _render_column(column, autogen_context), + "schema": schema } - if schema: - text += ", schema=%r" % schema - text += ")" return text def _drop_column(schema, tname, column, autogen_context): - text = "%(prefix)sdrop_column(%(tname)r, %(cname)r" % { + if 'batch_prefix' in autogen_context: + template = "%(prefix)sdrop_column(%(cname)r)" + else: + template = "%(prefix)sdrop_column(%(tname)r, %(cname)r" + if schema: + template += ", schema=%(schema)r" + template += ")" + + text = template % { "prefix": _alembic_autogenerate_prefix(autogen_context), "tname": tname, - "cname": column.name + "cname": column.name, + "schema": schema } - if schema: - text += ", schema=%r" % schema - text += ")" return text @@ -269,7 +290,13 @@ def _modify_col(tname, cname, existing_server_default=False, schema=None): indent = " " * 11 - text = "%(prefix)salter_column(%(tname)r, %(cname)r" % { + + if 'batch_prefix' in autogen_context: + template = "%(prefix)salter_column(%(cname)r" + else: + template = "%(prefix)salter_column(%(tname)r, %(cname)r" + + text = template % { 'prefix': _alembic_autogenerate_prefix( autogen_context), 'tname': tname, @@ -297,7 +324,7 @@ def _modify_col(tname, cname, autogen_context) text += ",\n%sexisting_server_default=%s" % ( indent, rendered) - if schema: + if schema and "batch_prefix" not in autogen_context: text += ",\n%sschema=%r" % (indent, schema) text += ")" return text @@ -316,7 +343,10 @@ def _sqlalchemy_autogenerate_prefix(autogen_context): def _alembic_autogenerate_prefix(autogen_context): - return autogen_context['opts']['alembic_module_prefix'] or '' + if 'batch_prefix' in autogen_context: + return autogen_context['batch_prefix'] + else: + return autogen_context['opts']['alembic_module_prefix'] or '' def _user_defined_render(type_, object_, autogen_context): diff --git a/alembic/batch.py b/alembic/batch.py index 96e12f43..26a148d3 100644 --- a/alembic/batch.py +++ b/alembic/batch.py @@ -1,7 +1,8 @@ class BatchOperationsImpl(object): - def __init__(self, operations, table_name, recreate): + def __init__(self, operations, table_name, schema, recreate): self.operations = operations self.table_name = table_name + self.schema = schema self.recreate = recreate self.batch = [] @@ -33,6 +34,7 @@ class BatchOperationsImpl(object): ) def add_column(self, *arg, **kw): + # TODO: omit table and schema names from all commands self.batch.append( ("add_column", arg, kw) ) diff --git a/alembic/environment.py b/alembic/environment.py index 59718e89..5dd9c6fb 100644 --- a/alembic/environment.py +++ b/alembic/environment.py @@ -264,6 +264,7 @@ class EnvironmentContext(object): starting_rev=None, tag=None, template_args=None, + render_as_batch=False, target_metadata=None, include_symbol=None, include_object=None, @@ -538,6 +539,16 @@ class EnvironmentContext(object): :paramref:`.EnvironmentContext.configure.include_object` + :param render_as_batch: if True, commands which alter elements + within a table will be placed under a ``with batch_alter_table():`` + directive, so that batch migrations will take place. + + .. versionadded:: 0.7.0 + + .. seealso:: + + :ref:`batch_migrations` + :param include_schemas: If True, autogenerate will scan across all schemas located by the SQLAlchemy :meth:`~sqlalchemy.engine.reflection.Inspector.get_schema_names` @@ -664,6 +675,7 @@ class EnvironmentContext(object): opts['include_symbol'] = include_symbol opts['include_object'] = include_object opts['include_schemas'] = include_schemas + opts['render_as_batch'] = True #render_as_batch opts['upgrade_token'] = upgrade_token opts['downgrade_token'] = downgrade_token opts['sqlalchemy_module_prefix'] = sqlalchemy_module_prefix diff --git a/alembic/operations.py b/alembic/operations.py index 10032a0f..0cf8bb62 100644 --- a/alembic/operations.py +++ b/alembic/operations.py @@ -190,8 +190,8 @@ class Operations(object): rel_t.append_column(sa_schema.Column(cname, NULLTYPE)) @contextmanager - def batch_alter_table(self, table_name, recreate=None): - impl = batch.BatchOperationImpl(self, table_name, recreate) + def batch_alter_table(self, table_name, schema=None, recreate=None): + impl = batch.BatchOperationImpl(self, table_name, schema, recreate) batch_op = Operations(self.migration_context, impl=impl) yield batch_op impl.flush()