- autogen composition back up

7 files changed, 95 insertions, 36 deletions

diff --git a/alembic/autogenerate/__init__.py b/alembic/autogenerate/__init__.py
index c33d2d8..3866337 100644
--- a/alembic/autogenerate/__init__.py
+++ b/alembic/autogenerate/__init__.py
@@ -1,5 +1,5 @@
 from .api import ( # noqa
-    compare_metadata, _produce_net_changes
+    compare_metadata, _produce_net_changes, _render_migration_diffs
     )
 from .generate import RevisionContext  # noqa
 from .render import render_op_text  # noqa
\ No newline at end of file
diff --git a/alembic/autogenerate/api.py b/alembic/autogenerate/api.py
index 6c30b13..a1602ec 100644
--- a/alembic/autogenerate/api.py
+++ b/alembic/autogenerate/api.py
@@ -116,6 +116,31 @@ def compare_metadata(context, metadata):
     return diffs
 
 
+def _render_migration_diffs(context, template_args, imports):
+    """legacy, used by test_autogen_composition at the moment"""
+
+    from ..operations import ops
+    from . import render
+
+    autogen_context = _autogen_context(context, imports=imports)
+    diffs = []
+
+    _produce_net_changes(autogen_context, diffs)
+
+    migration_script = ops.MigrationScript(
+        rev_id=None,
+        imports=imports,
+        upgrade_ops=ops.UpgradeOps([]),
+        downgrade_ops=ops.DowngradeOps([]),
+    )
+
+    compose._to_migration_script(autogen_context, migration_script, diffs)
+
+    render._render_migration_script(
+        autogen_context, migration_script, template_args
+    )
+
+
 def _autogen_context(
     context, imports=None, metadata=None, include_symbol=None,
         include_object=None, include_schemas=False):
diff --git a/alembic/autogenerate/compose.py b/alembic/autogenerate/compose.py
index f0ec4e5..b42b505 100644
--- a/alembic/autogenerate/compose.py
+++ b/alembic/autogenerate/compose.py
@@ -33,29 +33,31 @@ def _to_updown_op(autogen_context, diffs, op_container, type_):
 
     dest = [op_container.ops]
 
-    for (schema, table), subdiffs in _group_diffs_by_table(diffs):
-        if table is not None:
+    for (schema, tablename), subdiffs in _group_diffs_by_table(diffs):
+        subdiffs = list(subdiffs)
+        if tablename is not None:
             table_ops = []
-            op = ops.ModifyTableOps(table.name, table_ops, schema=table.schema)
+            op = ops.ModifyTableOps(tablename, table_ops, schema=schema)
             dest[-1].append(op)
-            dest.append(ops)
+            dest.append(table_ops)
         for diff in subdiffs:
             _produce_command(autogen_context, diff, dest[-1], type_)
-        dest.pop(-1)
+        if tablename is not None:
+            dest.pop(-1)
 
 
-def _produce_command(autogen_context, diff, ops, updown):
+def _produce_command(autogen_context, diff, op_list, updown):
     if isinstance(diff, tuple):
-        _produce_adddrop_command(updown, diff, autogen_context)
+        _produce_adddrop_command(updown, diff, op_list, autogen_context)
     else:
-        _produce_modify_command(updown, diff, autogen_context)
+        _produce_modify_command(updown, diff, op_list, autogen_context)
 
 
-def _produce_adddrop_command(updown, diff, autogen_context):
+def _produce_adddrop_command(updown, diff, op_list, autogen_context):
     cmd_type = diff[0]
     adddrop, cmd_type = cmd_type.split("_")
 
-    cmd_args = diff[1:] + (autogen_context,)
+    cmd_args = diff[1:]
 
     _commands = {
         "table": (ops.DropTableOp.from_table, ops.CreateTableOp.from_table),
@@ -78,12 +80,12 @@ def _produce_adddrop_command(updown, diff, autogen_context):
     ) or (
         updown == "downgrade" and adddrop == "remove"
     ):
-        return cmd_callables[1](*cmd_args)
+        op_list.append(cmd_callables[1](*cmd_args))
     else:
-        return cmd_callables[0](*cmd_args)
+        op_list.append(cmd_callables[0](*cmd_args))
 
 
-def _produce_modify_command(updown, diffs, autogen_context):
+def _produce_modify_command(updown, diffs, op_list, autogen_context):
     sname, tname, cname = diffs[0][1:4]
     kw = {}
 
@@ -107,14 +109,16 @@ def _produce_modify_command(updown, diffs, autogen_context):
             kw[new_kw] = diff[-2]
             kw[old_kw] = diff[-1]
 
-    if "nullable" in kw:
+    if "modify_nullable" in kw:
         kw.pop("existing_nullable", None)
-    if "server_default" in kw:
+    if "modify_server_default" in kw:
         kw.pop("existing_server_default", None)
 
-    return ops.AlterColumnOp(
-        tname, cname, schema=sname,
-        **kw
+    op_list.append(
+        ops.AlterColumnOp(
+            tname, cname, schema=sname,
+            **kw
+        )
     )
 
 
diff --git a/alembic/autogenerate/render.py b/alembic/autogenerate/render.py
index b540551..e4fb8ac 100644
--- a/alembic/autogenerate/render.py
+++ b/alembic/autogenerate/render.py
@@ -54,11 +54,14 @@ def _render_cmd_body(op_container, autogen_context):
         "please adjust! ###"
     )
 
-    for op in op_container:
-        lines = render_op(autogen_context, op)
+    if not op_container.ops:
+        printer.writeline("pass")
+    else:
+        for op in op_container.ops:
+            lines = render_op(autogen_context, op)
 
-        for line in lines:
-            printer.writeline(line)
+            for line in lines:
+                printer.writeline(line)
 
     printer.writeline("### end Alembic commands ###")
 
@@ -94,6 +97,7 @@ def _render_modify_table(autogen_context, op):
         if render_as_batch:
             del autogen_context['batch_prefix']
             lines.append("")
+        return lines
     else:
         return [
             "pass"
@@ -102,12 +106,14 @@ def _render_modify_table(autogen_context, op):
 
 @renderers.dispatch_for(ops.CreateTableOp)
 def _add_table(autogen_context, op):
+    table = op.to_table()
+
     args = [col for col in
-            [_render_column(col, autogen_context) for col in op.columns]
+            [_render_column(col, autogen_context) for col in table.columns]
             if col] + \
         sorted([rcons for rcons in
                 [_render_constraint(cons, autogen_context) for cons in
-                 op.constraints]
+                 table.constraints]
                 if rcons is not None
                 ])
 
@@ -295,7 +301,7 @@ def _drop_column(autogen_context, op):
         "cname": _ident(column_name),
         "schema": _ident(schema)
     }
-    return text
+    return [text]
 
 
 @renderers.dispatch_for(ops.AlterColumnOp)
@@ -349,7 +355,7 @@ def _alter_column(autogen_context, op):
     if schema and "batch_prefix" not in autogen_context:
         text += ",\n%sschema=%r" % (indent, schema)
     text += ")"
-    return text
+    return [text]
 
 
 ################################################################
diff --git a/alembic/operations/base.py b/alembic/operations/base.py
index b6fa7d4..2d69708 100644
--- a/alembic/operations/base.py
+++ b/alembic/operations/base.py
@@ -1,6 +1,7 @@
 from contextlib import contextmanager
 
 from sqlalchemy import schema as sa_schema
+from sqlalchemy import sql
 
 from .. import util
 from ..util import sqla_compat
@@ -804,9 +805,12 @@ class Operations(object):
             object is returned.
 
         """
-        table = self.schema_obj.table(table_name, *columns, **kw)
-        self.impl.create_table(table)
-        return table
+        op = ops.CreateTableOp(
+            table_name,
+            columns,
+            **kw
+        )
+        return self.invoke(op)
 
     @util._with_legacy_names([('name', 'table_name')])
     def drop_table(self, table_name, schema=None, **kw):
diff --git a/alembic/operations/ops.py b/alembic/operations/ops.py
index 1e27654..eae9e34 100644
--- a/alembic/operations/ops.py
+++ b/alembic/operations/ops.py
@@ -214,23 +214,33 @@ class DropIndexOp(MigrateOperation):
 
 
 class CreateTableOp(MigrateOperation):
-    def __init__(self, table_name, columns, constraints, schema=None, **kw):
+    def __init__(
+            self, table_name, columns, schema=None, _orig_table=None, **kw):
         self.table_name = table_name
         self.columns = columns
-        self.constraints = constraints
         self.schema = schema
         self.kw = kw
+        self._orig_table = _orig_table
 
     @classmethod
     def from_table(cls, table):
         return CreateTableOp(
             table.name,
-            list(table.c),
-            list(table.constraints),
+            list(table.c) + list(table.constraints),
             schema=table.schema,
+            _orig_table=table,
             **table.kwargs
         )
 
+    def to_table(self, migration_context=None):
+        if self._orig_table is not None:
+            return self._orig_table
+        schema_obj = schemaobj.SchemaObjects(migration_context)
+
+        return schema_obj.table(
+            self.table_name, *self.columns, schema=self.schema, **self.kw
+        )
+
 
 class DropTableOp(MigrateOperation):
     def __init__(self, table_name, schema=None, table_kw=None):
@@ -363,8 +373,10 @@ class MigrationScript(MigrateOperation):
     """
 
     def __init__(
-            self, rev_id, message, imports, upgrade_ops, downgrade_ops,
-            head, splice, branch_label, version_path):
+            self, rev_id, upgrade_ops, downgrade_ops,
+            message=None,
+            imports=None, head=None, splice=None,
+            branch_label=None, version_path=None):
         self.rev_id = rev_id
         self.message = message
         self.imports = imports
diff --git a/alembic/operations/toimpl.py b/alembic/operations/toimpl.py
index 28d4feb..1c9d808 100644
--- a/alembic/operations/toimpl.py
+++ b/alembic/operations/toimpl.py
@@ -82,3 +82,11 @@ def drop_column(operations, operation):
 def create_index(operations, operation):
     idx = operation.to_index(operations.migration_context)
     operations.impl.create_index(idx)
+
+
+@ops.to_impl.dispatch_for(ops.CreateTableOp)
+def create_table(operations, operation):
+    table = operation.to_table(operations.migration_context)
+    operations.impl.create_table(table)
+
+