- add support for autogenerate to include "batch"

diff --git a/alembic/autogenerate/api.py b/alembic/autogenerate/api.py
index 3519a6d1f25c41ce88be2e048d8cfa4e9d5cf938..3f2bc8b441a94547f8e5453dc6a66979f630b8f1 100644 (file)
--- 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 0344b7741763623b62e71b0be10f49ff10100f2e..5fdfe3662c04e8d35bf1cca51bf1d68874d55ea1 100644 (file)
--- 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 96e12f43c55492b62a1d7ba691e0da380772a779..26a148d34865a5c45a2ec73ab7feb314caea1bda 100644 (file)
--- 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 59718e891ebda27b9b1f6362b7fddfeda77fe55b..5dd9c6fb9dd3802e08f7eac8708531a186d8ea8a 100644 (file)
--- 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 10032a0f23825dbff6cfc0496661643dce73a425..0cf8bb6238846351176a1c41078bd95f420ccb5e 100644 (file)
--- 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()