wipticket_318

4 files changed, 116 insertions, 36 deletions

diff --git a/alembic/autogenerate/api.py b/alembic/autogenerate/api.py
index 811ebf4..b5a6316 100644
--- a/alembic/autogenerate/api.py
+++ b/alembic/autogenerate/api.py
@@ -7,6 +7,7 @@ from . import compare
 from .. import util
 from sqlalchemy.engine.reflection import Inspector
 import contextlib
+import inspect
 
 
 def compare_metadata(context, metadata):
@@ -231,7 +232,11 @@ class AutogenContext(object):
                 "the database for schema information")
 
         if opts is None:
-            opts = migration_context.opts
+            if migration_context is None:
+                opts = {}
+            else:
+                opts = migration_context.opts
+
         self.metadata = metadata = opts.get('target_metadata', None) \
             if metadata is None else metadata
 
@@ -313,19 +318,20 @@ class RevisionContext(object):
             'config': config  # Let templates use config for
                               # e.g. multiple databases
         }
-        self.generated_revisions = [
-            self._default_revision()
-        ]
+        self.generated_revisions = []
+        self.autogenerated = False
+        self.opts = {}
 
     def _to_script(self, migration_script):
         template_args = {}
         for k, v in self.template_args.items():
             template_args.setdefault(k, v)
 
-        if migration_script._autogen_context is not None:
-            render._render_migration_script(
-                migration_script._autogen_context, migration_script,
-                template_args
+        if self.autogenerated:
+            autogen_context = AutogenContext(None, opts=self.opts)
+
+            render._render_python_into_templatevars(
+                autogen_context, migration_script, template_args
             )
 
         return self.script_directory.generate_revision(
@@ -339,7 +345,7 @@ class RevisionContext(object):
             depends_on=migration_script.depends_on,
             **template_args)
 
-    def run_autogenerate(self, rev, context):
+    def run_autogenerate(self, rev, migration_context):
         if self.command_args['sql']:
             raise util.CommandError(
                 "Using --sql with --autogenerate does not make any sense")
@@ -347,41 +353,63 @@ class RevisionContext(object):
                 set(self.script_directory.get_revisions("heads")):
             raise util.CommandError("Target database is not up to date.")
 
-        autogen_context = AutogenContext(context)
+        self.autogenerated = True
+        self.opts = migration_context.opts
 
-        migration_script = self.generated_revisions[0]
+        autogen_context = AutogenContext(migration_context)
 
-        compare._populate_migration_script(autogen_context, migration_script)
+        upgrade_token = migration_context.opts['upgrade_token']
+        downgrade_token = migration_context.opts['downgrade_token']
 
-        hook = context.opts.get('process_revision_directives', None)
-        if hook:
-            hook(context, rev, self.generated_revisions)
+        if not self.generated_revisions:
+            self.generated_revisions.append(
+                self._default_revision(
+                    autogen_context,
+                    upgrade_token, downgrade_token
+                )
+            )
+            migration_script = self.generated_revisions[0]
+        else:
+            migration_script = self.generated_revisions[0]
+            migration_script.append_upgrade_downgrade_ops(
+                ops.UpgradeOps([], upgrade_token=upgrade_token),
+                ops.DowngradeOps([], downgrade_token=downgrade_token)
+            )
 
-        for migration_script in self.generated_revisions:
-            migration_script._autogen_context = autogen_context
+        compare._populate_migration_script(
+            autogen_context, migration_script)
 
-    def run_no_autogenerate(self, rev, context):
-        hook = context.opts.get('process_revision_directives', None)
+        hook = migration_context.opts['process_revision_directives']
         if hook:
-            hook(context, rev, self.generated_revisions)
+            hook(migration_context, rev, self.generated_revisions)
 
-        for migration_script in self.generated_revisions:
-            migration_script._autogen_context = None
+    def run_no_autogenerate(self, rev, migration_context):
+        if not self.generated_revisions:
+            self.generated_revisions.append(
+                self._default_revision(None)
+            )
+        self.opts = migration_context.opts
+        hook = migration_context.opts['process_revision_directives']
+        if hook:
+            hook(migration_context, rev, self.generated_revisions)
 
-    def _default_revision(self):
+    def _default_revision(
+            self, autogen_context, upgrade_token, downgrade_token):
         op = ops.MigrationScript(
             rev_id=self.command_args['rev_id'] or util.rev_id(),
             message=self.command_args['message'],
-            imports=set(),
-            upgrade_ops=ops.UpgradeOps([]),
-            downgrade_ops=ops.DowngradeOps([]),
+            imports=autogen_context._imports if autogen_context else set(),
+            upgrade_ops=ops.UpgradeOps(
+                [], upgrade_token=upgrade_token),
+            downgrade_ops=ops.DowngradeOps(
+                [], downgrade_token=downgrade_token),
             head=self.command_args['head'],
             splice=self.command_args['splice'],
             branch_label=self.command_args['branch_label'],
             version_path=self.command_args['version_path'],
             depends_on=self.command_args['depends_on']
         )
-        op._autogen_context = None
+        op._autogen_contexts = []
         return op
 
     def generate_scripts(self):
diff --git a/alembic/autogenerate/compare.py b/alembic/autogenerate/compare.py
index cc12dae..d97bc62 100644
--- a/alembic/autogenerate/compare.py
+++ b/alembic/autogenerate/compare.py
@@ -16,8 +16,15 @@ log = logging.getLogger(__name__)
 
 
 def _populate_migration_script(autogen_context, migration_script):
-    _produce_net_changes(autogen_context, migration_script.upgrade_ops)
-    migration_script.upgrade_ops.reverse_into(migration_script.downgrade_ops)
+    if isinstance(migration_script.upgrade_ops, list):
+        upgrade_ops = migration_script.upgrade_ops[-1]
+        downgrade_ops = migration_script.downgrade_ops[-1]
+    else:
+        upgrade_ops = migration_script.upgrade_ops
+        downgrade_ops = migration_script.downgrade_ops
+
+    _produce_net_changes(autogen_context, upgrade_ops)
+    upgrade_ops.reverse_into(downgrade_ops)
 
 
 comparators = util.Dispatcher(uselist=True)
diff --git a/alembic/autogenerate/render.py b/alembic/autogenerate/render.py
index 422506d..a90f929 100644
--- a/alembic/autogenerate/render.py
+++ b/alembic/autogenerate/render.py
@@ -29,13 +29,19 @@ def _indent(text):
     return text
 
 
-def _render_migration_script(autogen_context, migration_script, template_args):
-    opts = autogen_context.opts
-    imports = autogen_context._imports
-    template_args[opts['upgrade_token']] = _indent(_render_cmd_body(
-        migration_script.upgrade_ops, autogen_context))
-    template_args[opts['downgrade_token']] = _indent(_render_cmd_body(
-        migration_script.downgrade_ops, autogen_context))
+def _render_python_into_templatevars(
+        autogen_context, migration_script, template_args):
+    imports = migration_script.imports
+
+    upgrade_ops_list = util.to_list(migration_script.upgrade_ops)
+    downgrade_ops_list = util.to_list(migration_script.downgrade_ops)
+
+    for upgrade_ops, downgrade_ops in zip(
+            upgrade_ops_list, downgrade_ops_list):
+        template_args[upgrade_ops.upgrade_token] = _indent(
+            _render_cmd_body(upgrade_ops, autogen_context))
+        template_args[downgrade_ops.downgrade_token] = _indent(
+            _render_cmd_body(downgrade_ops, autogen_context))
     template_args['imports'] = "\n".join(sorted(imports))
 
 
diff --git a/alembic/operations/ops.py b/alembic/operations/ops.py
index 7bdbb1f..7383d38 100644
--- a/alembic/operations/ops.py
+++ b/alembic/operations/ops.py
@@ -1875,6 +1875,10 @@ class UpgradeOps(OpContainer):
 
     """
 
+    def __init__(self, ops=(), upgrade_token="upgrades"):
+        super(UpgradeOps, self).__init__(ops=ops)
+        self.upgrade_token = upgrade_token
+
     def reverse_into(self, downgrade_ops):
         downgrade_ops.ops[:] = list(reversed(
             [op.reverse() for op in self.ops]
@@ -1895,6 +1899,10 @@ class DowngradeOps(OpContainer):
 
     """
 
+    def __init__(self, ops=(), downgrade_token="downgrades"):
+        super(DowngradeOps, self).__init__(ops=ops)
+        self.downgrade_token = downgrade_token
+
     def reverse(self):
         return UpgradeOps(
             ops=list(reversed(
@@ -1911,6 +1919,18 @@ class MigrationScript(MigrateOperation):
 
     A normal :class:`.MigrationScript` object would contain a single
     :class:`.UpgradeOps` and a single :class:`.DowngradeOps` directive.
+    These are accessible via the ``.upgrade_ops`` and ``.downgrade_ops``
+    attributes.
+
+    In the case of an autogenerate option that runs multiple times,
+    such as the multiple database example in the "multidb" template,
+    the ``.upgrade_ops`` and ``.downgrade_ops`` will be present as lists.
+    Each :class:`.UpgradeOps` and :class:`.DowngradeOps` will refer to
+    individual "upgrade_token" and "downgrade_token" values.
+
+    .. versionchanged:: 0.8.1 the ``.upgrade_ops`` and ``.downgrade_ops``
+       attributes are converted into lists if multiple autogenerate
+       passes proceed on the same :class:`.MigrationScript` object.
 
     .. seealso::
 
@@ -1933,3 +1953,22 @@ class MigrationScript(MigrateOperation):
         self.depends_on = depends_on
         self.upgrade_ops = upgrade_ops
         self.downgrade_ops = downgrade_ops
+
+    def append_upgrade_downgrade_ops(self, upgrade_ops, downgrade_ops):
+        """add the given :class:`.UpgradeOps` and :class:`.DowngradeOps`
+        to this :class:`.MigrationScript`.
+
+        This will convert the existing ``.upgrade_ops`` and ``.downgrade_ops``
+        elements into lists if they are not already, and append the new
+        entries.
+
+        """
+        if not isinstance(self.upgrade_ops, list):
+            assert isinstance(self.upgrade_ops, UpgradeOps)
+            self.upgrade_ops = [self.upgrade_ops]
+        if not isinstance(self.downgrade_ops, list):
+            assert isinstance(self.downgrade_ops, DowngradeOps)
+            self.downgrade_ops = [self.downgrade_ops]
+
+        self.upgrade_ops.append(upgrade_ops)
+        self.downgrade_ops.append(downgrade_ops)