- add support for autogenerate to include "batch"

2 files changed, 95 insertions, 31 deletions

diff --git a/alembic/autogenerate/api.py b/alembic/autogenerate/api.py
index 3519a6d..3f2bc8b 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 0344b77..5fdfe36 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):