- add the autogenerate customize API

- add tests for extension APIs
- add a render helper function
- lots of docs

11 files changed, 572 insertions, 29 deletions

diff --git a/alembic/autogenerate/__init__.py b/alembic/autogenerate/__init__.py
index 81e1ddc..4272a7e 100644
--- a/alembic/autogenerate/__init__.py
+++ b/alembic/autogenerate/__init__.py
@@ -1,6 +1,7 @@
 from .api import ( # noqa
-    compare_metadata, _render_migration_diffs
+    compare_metadata, _render_migration_diffs,
+    produce_migrations, render_python_code
     )
 from .compare import _produce_net_changes  # noqa
 from .generate import RevisionContext  # noqa
-from .render import render_op_text  # noqa
\ No newline at end of file
+from .render import render_op_text, renderers  # noqa
\ No newline at end of file
diff --git a/alembic/autogenerate/api.py b/alembic/autogenerate/api.py
index 5fdc5b1..cff977b 100644
--- a/alembic/autogenerate/api.py
+++ b/alembic/autogenerate/api.py
@@ -94,7 +94,7 @@ def compare_metadata(context, metadata):
     .. seealso::
 
         :func:`.produce_migrations` - produces a :class:`.MigrationScript`
-         structure based on metadata comparison.
+        structure based on metadata comparison.
 
     """
 
@@ -115,14 +115,21 @@ def compare_metadata(context, metadata):
 
 
 def produce_migrations(context, metadata):
-    """Produce a :class:`.MigrationScript` structure based on schema comparison.
+    """Produce a :class:`.MigrationScript` structure based on schema
+    comparison.
+
+    This function does essentially what :func:`.compare_metadata` does,
+    but then runs the resulting list of diffs to produce the full
+    :class:`.MigrationScript` object.   For an example of what this looks like,
+    see the example in :ref:`customizing_revision`.
 
     .. versionadded:: 0.8.0
 
     .. seealso::
 
         :func:`.compare_metadata` - returns more fundamental "diff"
-         data from comparing a schema.
+        data from comparing a schema.
+
     """
 
     autogen_context = _autogen_context(context, metadata=metadata)
@@ -141,6 +148,34 @@ def produce_migrations(context, metadata):
     return migration_script
 
 
+def render_python_code(
+    up_or_down_op,
+    sqlalchemy_module_prefix='sa.',
+    alembic_module_prefix='op.',
+    imports=(),
+    render_item=None,
+):
+    """Render Python code given an :class:`.UpgradeOps` or
+    :class:`.DowngradeOps` object.
+
+    This is a convenience function that can be used to test the
+    autogenerate output of a user-defined :class:`.MigrationScript` structure.
+
+    """
+    autogen_context = {
+        'opts': {
+            'sqlalchemy_module_prefix': sqlalchemy_module_prefix,
+            'alembic_module_prefix': alembic_module_prefix,
+            'render_item': render_item,
+        },
+        'imports': set(imports)
+    }
+    return render._indent(render._render_cmd_body(
+        up_or_down_op, autogen_context))
+
+
+
+
 def _render_migration_diffs(context, template_args, imports):
     """legacy, used by test_autogen_composition at the moment"""
 
diff --git a/alembic/autogenerate/generate.py b/alembic/autogenerate/generate.py
index f6094c3..c686156 100644
--- a/alembic/autogenerate/generate.py
+++ b/alembic/autogenerate/generate.py
@@ -54,14 +54,23 @@ class RevisionContext(object):
         compare._produce_net_changes(autogen_context, diffs)
 
         migration_script = self.generated_revisions[0]
-        migration_script._autogen_context = autogen_context
 
         compose._to_migration_script(autogen_context, migration_script, diffs)
 
-        # DO THE HOOK HERE!!
+        hook = context.opts.get('process_revision_directives', None)
+        if hook:
+            hook(context, rev, self.generated_revisions)
+
+        for migration_script in self.generated_revisions:
+            migration_script._autogen_context = autogen_context
 
     def run_no_autogenerate(self, rev, context):
-        pass
+        hook = context.opts.get('process_revision_directives', None)
+        if hook:
+            hook(context, rev, self.generated_revisions)
+
+        for migration_script in self.generated_revisions:
+            migration_script._autogen_context = None
 
     def _default_revision(self):
         op = ops.MigrationScript(
diff --git a/alembic/autogenerate/render.py b/alembic/autogenerate/render.py
index 873db47..c3f3df1 100644
--- a/alembic/autogenerate/render.py
+++ b/alembic/autogenerate/render.py
@@ -68,7 +68,7 @@ def _render_cmd_body(op_container, autogen_context):
 
 def render_op(autogen_context, op):
     renderer = renderers.dispatch(op)
-    lines = renderer(autogen_context, op)
+    lines = util.to_list(renderer(autogen_context, op))
     return lines
 
 
@@ -130,7 +130,7 @@ def _add_table(autogen_context, op):
     for k in sorted(op.kw):
         text += ",\n%s=%r" % (k.replace(" ", "_"), op.kw[k])
     text += "\n)"
-    return [text]
+    return text
 
 
 @renderers.dispatch_for(ops.DropTableOp)
@@ -142,7 +142,7 @@ def _drop_table(autogen_context, op):
     if op.schema:
         text += ", schema=%r" % _ident(op.schema)
     text += ")"
-    return [text]
+    return text
 
 
 @renderers.dispatch_for(ops.CreateIndexOp)
@@ -175,7 +175,7 @@ def _add_index(autogen_context, op):
                  for key, val in index.kwargs.items()]))
         if len(index.kwargs) else ''
     }
-    return [text]
+    return text
 
 
 @renderers.dispatch_for(ops.DropIndexOp)
@@ -195,7 +195,7 @@ def _drop_index(autogen_context, op):
         'schema': ((", schema=%r" % _ident(op.schema))
                    if op.schema else '')
     }
-    return [text]
+    return text
 
 
 @renderers.dispatch_for(ops.CreateUniqueConstraintOp)
@@ -224,10 +224,10 @@ def _add_fk_constraint(autogen_context, op):
             if value is not None:
                 args.append("%s=%r" % (k, value))
 
-    return ["%(prefix)screate_foreign_key(%(args)s)" % {
+    return "%(prefix)screate_foreign_key(%(args)s)" % {
         'prefix': _alembic_autogenerate_prefix(autogen_context),
         'args': ", ".join(args)
-    }]
+    }
 
 
 @renderers.dispatch_for(ops.CreatePrimaryKeyOp)
@@ -259,7 +259,7 @@ def _drop_constraint(autogen_context, op):
         'schema': (", schema='%s'" % _ident(op.schema))
         if op.schema else '',
     }
-    return [text]
+    return text
 
 
 @renderers.dispatch_for(ops.AddColumnOp)
@@ -279,7 +279,7 @@ def _add_column(autogen_context, op):
         "column": _render_column(column, autogen_context),
         "schema": schema
     }
-    return [text]
+    return text
 
 
 @renderers.dispatch_for(ops.DropColumnOp)
@@ -301,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)
@@ -355,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
 
 
 class _f_name(object):
diff --git a/alembic/operations/ops.py b/alembic/operations/ops.py
index 9f8176f..1a38d07 100644
--- a/alembic/operations/ops.py
+++ b/alembic/operations/ops.py
@@ -18,6 +18,8 @@ class MigrateOperation(object):
 
         :ref:`operation_plugins`
 
+        :ref:`customizing_revision`
+
     """
 
 
@@ -1496,7 +1498,7 @@ class ExecuteSQLOp(MigrateOperation):
 
 class OpContainer(MigrateOperation):
     """Represent a sequence of operations operation."""
-    def __init__(self, ops):
+    def __init__(self, ops=()):
         self.ops = ops
 
 
@@ -1511,12 +1513,24 @@ class ModifyTableOps(OpContainer):
 
 class UpgradeOps(OpContainer):
     """contains a sequence of operations that would apply to the
-    'upgrade' stream of a script."""
+    'upgrade' stream of a script.
+
+    .. seealso::
+
+        :ref:`customizing_revision`
+
+    """
 
 
 class DowngradeOps(OpContainer):
     """contains a sequence of operations that would apply to the
-    'downgrade' stream of a script."""
+    'downgrade' stream of a script.
+
+    .. seealso::
+
+        :ref:`customizing_revision`
+
+    """
 
 
 class MigrationScript(MigrateOperation):
@@ -1528,6 +1542,10 @@ class MigrationScript(MigrateOperation):
     A normal :class:`.MigrationScript` object would contain a single
     :class:`.UpgradeOps` and a single :class:`.DowngradeOps` directive.
 
+    .. seealso::
+
+        :ref:`customizing_revision`
+
     """
 
     def __init__(
diff --git a/alembic/runtime/environment.py b/alembic/runtime/environment.py
index ffaa205..3b04fea 100644
--- a/alembic/runtime/environment.py
+++ b/alembic/runtime/environment.py
@@ -290,6 +290,7 @@ class EnvironmentContext(util.ModuleClsProxy):
                   include_symbol=None,
                   include_object=None,
                   include_schemas=False,
+                  process_revision_directives=None,
                   compare_type=False,
                   compare_server_default=False,
                   render_item=None,
@@ -653,6 +654,43 @@ class EnvironmentContext(util.ModuleClsProxy):
 
             :ref:`autogen_module_prefix`
 
+        :param process_revision_directives: a callable function that will
+         be passed a structure representing the end result of an autogenerate
+         or plain "revision" operation, which can be manipulated to affect
+         how the ``alembic revision`` command ultimately outputs new
+         revision scripts.   The structure of the callable is::
+
+            def process_revision_directives(context, revision, directives):
+                pass
+
+         The ``directives`` parameter is a Python list containing
+         a single :class:`.MigrationScript` directive, which represents
+         the revision file to be generated.    This list as well as its
+         contents may be freely modified to produce any set of commands.
+         The section :ref:`customizing_revision` shows an example of
+         doing this.  The ``context`` parameter is the
+         :class:`.MigrationContext` in use,
+         and ``revision`` is a tuple of revision identifiers representing the
+         current revision of the database.
+
+         The callable is invoked at all times when the ``--autogenerate``
+         option is passed to ``alembic revision``.  If ``--autogenerate``
+         is not passed, the callable is invoked only if the
+         ``revision_environment`` variable is set to True in the Alembic
+         configuration, in which case the given ``directives`` collection
+         will contain empty :class:`.UpgradeOps` and :class:`.DowngradeOps`
+         collections for ``.upgrade_ops`` and ``.downgrade_ops``.  The
+         ``--autogenerate`` option itself can be inferred by inspecting
+         ``context.config.cmd_opts.autogenerate``.
+
+
+         .. versionadded:: 0.8.0
+
+         .. seealso::
+
+             :ref:`customizing_revision`
+
+
         Parameters specific to individual backends:
 
         :param mssql_batch_separator: The "batch separator" which will
@@ -693,6 +731,8 @@ class EnvironmentContext(util.ModuleClsProxy):
         opts['alembic_module_prefix'] = alembic_module_prefix
         opts['user_module_prefix'] = user_module_prefix
         opts['literal_binds'] = literal_binds
+        opts['process_revision_directives'] = process_revision_directives
+
         if render_item is not None:
             opts['render_item'] = render_item
         if compare_type is not None:
diff --git a/alembic/util/__init__.py b/alembic/util/__init__.py
index 7dfbfae..bd7196c 100644
--- a/alembic/util/__init__.py
+++ b/alembic/util/__init__.py
@@ -1,5 +1,5 @@
 from .langhelpers import (  # noqa
-    asbool, rev_id, to_tuple, memoized_property,
+    asbool, rev_id, to_tuple, to_list, memoized_property,
     immutabledict, _with_legacy_names, Dispatcher, ModuleClsProxy)
 from .messaging import (  # noqa
     write_outstream, status, err, obfuscate_url_pw, warn, msg, format_as_comma)
diff --git a/alembic/util/langhelpers.py b/alembic/util/langhelpers.py
index cfcca85..904848c 100644
--- a/alembic/util/langhelpers.py
+++ b/alembic/util/langhelpers.py
@@ -131,6 +131,17 @@ def rev_id():
     return hex(val)[2:-1]
 
 
+def to_list(x, default=None):
+    if x is None:
+        return default
+    elif isinstance(x, string_types):
+        return [x]
+    elif isinstance(x, collections.Iterable):
+        return list(x)
+    else:
+        raise ValueError("Don't know how to turn %r into a list" % x)
+
+
 def to_tuple(x, default=None):
     if x is None:
         return default
diff --git a/docs/build/api.rst b/docs/build/api.rst
index 8dbb88c..3fc8ca0 100644
--- a/docs/build/api.rst
+++ b/docs/build/api.rst
@@ -169,12 +169,16 @@ The migration operations present on :class:`.Operations` are themselves
 delivered via operation objects that represent an operation and its
 arguments.   All operations descend from the :class:`.MigrateOperation`
 class, and are registered with the :class:`.Operations` class using
-the :meth:`.Operations.register_operation` class decorator.
+the :meth:`.Operations.register_operation` class decorator.  The
+:class:`.MigrateOperation` objects also serve as the basis for how the
+autogenerate system renders new migration scripts.
 
 .. seealso::
 
     :ref:`operation_plugins`
 
+    :ref:`customizing_revision`
+
 The built-in operation objects are listed below.
 
 .. automodule:: alembic.operations.ops
@@ -266,11 +270,235 @@ management, used exclusively by :class:`.ScriptDirectory`.
 Autogeneration
 ==============
 
-Alembic 0.3 introduces a small portion of the autogeneration system
-as a public API.
+The autogenerate system has two areas of API that are public:
+
+1. The ability to do a "diff" of a :class:`.MetaData` object against
+   a database, and receive a data structure back.  This structure
+   is available either as a rudimentary list of changes, or as
+   a :class:`.MigrateOperation` structure.
+
+2. The ability to alter how the ``alembic revision`` command generates
+   revision scripts, including support for multiple revision scripts
+   generated in one pass.
+
+Getting Diffs
+-------------
 
 .. autofunction:: alembic.autogenerate.compare_metadata
 
+.. autofunction:: alembic.autogenerate.produce_migrations
+
+.. _customizing_revision:
+
+Customizing Revision Generation
+-------------------------------
+
+.. versionadded:: 0.8.0 - the ``alembic revision`` system is now customizable.
+
+The ``alembic revision`` command, also available programmatically
+via :func:`.command.revision`, essentially produces a single migration
+script after being run.  Whether or not the ``--autogenerate`` option
+was specified basically determines if this script is a blank revision
+script with empty ``upgrade()`` and ``downgrade()`` functions, or was
+produced with alembic operation directives as the result of autogenerate.
+
+In either case, the system creates a full plan of what is to be done
+in the form of a :class:`.MigrateOperation` structure, which is then
+used to produce the script.
+
+For example, suppose we ran ``alembic revision --autogenerate``, and the
+end result was that it produced a new revision ``'eced083f5df'``
+with the following contents::
+
+    """create the organization table."""
+
+    # revision identifiers, used by Alembic.
+    revision = 'eced083f5df'
+    down_revision = 'beafc7d709f'
+
+    from alembic import op
+    import sqlalchemy as sa
+
+
+    def upgrade():
+        op.create_table(
+            'organization',
+            sa.Column('id', sa.Integer(), primary_key=True),
+            sa.Column('name', sa.String(50), nullable=False)
+        )
+        op.add_column(
+            'user',
+            sa.Column('organization_id', sa.Integer())
+        )
+        op.create_foreign_key(
+            'org_fk', 'user', 'organization', ['organization_id'], ['id']
+        )
+
+    def downgrade():
+        op.drop_constraint('org_fk', 'user')
+        op.drop_column('user', 'organization_id')
+        op.drop_table('organization')
+
+The above script is generated by a :class:`.MigrateOperation` structure
+that looks like this::
+
+    from alembic.operations import ops
+    import sqlalchemy as sa
+
+    migration_script = ops.MigrationScript(
+        'eced083f5df',
+        ops.UpgradeOps(
+            ops=[
+                ops.CreateTableOp(
+                    'organization',
+                    [
+                        sa.Column('id', sa.Integer(), primary_key=True),
+                        sa.Column('name', sa.String(50), nullable=False)
+                    ]
+                ),
+                ops.ModifyTableOps(
+                    'user',
+                    ops=[
+                        ops.AddColumnOp(
+                            'user',
+                            sa.Column('organization_id', sa.Integer())
+                        ),
+                        ops.CreateForeignKeyOp(
+                            'org_fk', 'user', 'organization',
+                            ['organization_id'], ['id']
+                        )
+                    ]
+                )
+            ]
+        ),
+        ops.DowngradeOps(
+            ops=[
+                ops.ModifyTableOps(
+                    'user',
+                    ops=[
+                        ops.DropConstraintOp('org_fk', 'user'),
+                        ops.DropColumnOp('user', 'organization_id')
+                    ]
+                ),
+                ops.DropTableOp('organization')
+            ]
+        ),
+        message='create the organization table.'
+    )
+
+When we deal with a :class:`.MigrationScript` structure, we can render
+the upgrade/downgrade sections into strings for debugging purposes
+using the :func:`.render_python_code` helper function::
+
+    from alembic.autogenerate import render_python_code
+    print(render_python_code(migration_script.upgrade_ops))
+
+Renders::
+
+    ### commands auto generated by Alembic - please adjust! ###
+        op.create_table('organization',
+        sa.Column('id', sa.Integer(), nullable=False),
+        sa.Column('name', sa.String(length=50), nullable=False),
+        sa.PrimaryKeyConstraint('id')
+        )
+        op.add_column('user', sa.Column('organization_id', sa.Integer(), nullable=True))
+        op.create_foreign_key('org_fk', 'user', 'organization', ['organization_id'], ['id'])
+        ### end Alembic commands ###
+
+Given that structures like the above are used to generate new revision
+files, and that we'd like to be able to alter these as they are created,
+we then need a system to access this structure when the
+:func:`.command.revision` command is used.  The
+:paramref:`.EnvironmentContext.configure.process_revision_directives`
+parameter gives us a way to alter this.   This is a function that
+is passed the above structure as generated by Alembic, giving us a chance
+to alter it.
+For example, if we wanted to put all the "upgrade" operations into
+a certain branch, and we wanted our script to not have any "downgrade"
+operations at all, we could build an extension as follows, illustrated
+within an ``env.py`` script::
+
+    def process_revision_directives(context, revision, directives):
+        script = directives[0]
+
+        # set specific branch
+        script.head = "mybranch@head"
+
+        # erase downgrade operations
+        script.downgrade_ops.ops[:] = []
+
+    # ...
+
+    def run_migrations_online():
+
+        # ...
+        with engine.connect() as connection:
+
+            context.configure(
+                connection=connection,
+                target_metadata=target_metadata,
+                process_revision_directives=process_revision_directives)
+
+            with context.begin_transaction():
+                context.run_migrations()
+
+Above, the ``directives`` argument is a Python list.  We may alter the
+given structure within this list in-place, or replace it with a new
+structure consisting of zero or more :class:`.MigrationScript` directives.
+The :func:`.command.revision` command will then produce scripts corresponding
+to whatever is in this list.
+
+.. autofunction:: alembic.autogenerate.render_python_code
+
+Autogenerating Custom Operation Directives
+------------------------------------------
+
+In the section :ref:`operation_plugins`, we talked about adding new
+subclasses of :class:`.MigrateOperation` in order to add new ``op.``
+directives.  In the preceding section :ref:`customizing_revision`, we
+also learned that these same :class:`.MigrateOperation` structures are at
+the base of how the autogenerate system knows what Python code to render.
+How to connect these two systems, so that our own custom operation
+directives can be used?  First off, we'd probably be implementing
+a :paramref:`.EnvironmentContext.configure.process_revision_directives`
+plugin as described previously, so that we can add our own directives
+to the autogenerate stream.  What if we wanted to add our ``CreateSequenceOp``
+to the autogenerate structure?  We basically need to define an autogenerate
+renderer for it, as follows::
+
+    # note: this is a continuation of the example from the
+    # "Operation Plugins" section
+
+    from alembic.autogenerate import renderers
+
+    @renderers.dispatch_for(CreateSequenceOp)
+    def render_create_sequence(autogen_context, op):
+        return "op.create_sequence(%r, **%r)" % (
+            op.sequence_name,
+            op.kw
+        )
+
+With our render function established, we can our ``CreateSequenceOp``
+generated in an autogenerate context using the :func:`.render_python_code`
+debugging function in conjunction with an :class:`.UpgradeOps` structure::
+
+    from alembic.operations import ops
+    from alembic.autogenerate import render_python_code
+
+    upgrade_ops = ops.UpgradeOps(
+        ops=[
+            CreateSequenceOp("my_seq")
+        ]
+    )
+
+    print(render_python_code(upgrade_ops))
+
+Which produces::
+
+    ### commands auto generated by Alembic - please adjust! ###
+        op.create_sequence('my_seq', **{})
+        ### end Alembic commands ###
+
 DDL Internals
 =============
 
diff --git a/tests/test_op.py b/tests/test_op.py
index e8caf99..9c14e49 100644
--- a/tests/test_op.py
+++ b/tests/test_op.py
@@ -851,4 +851,32 @@ class SQLModeOpTest(TestBase):
         context.assert_(
             "CREATE TABLE some_table (id INTEGER NOT NULL, st_id INTEGER, "
             "PRIMARY KEY (id), FOREIGN KEY(st_id) REFERENCES some_table (id))"
-        )
\ No newline at end of file
+        )
+
+
+class CustomOpTest(TestBase):
+    def test_custom_op(self):
+        from alembic.operations import Operations, MigrateOperation
+
+        @Operations.register_operation("create_sequence")
+        class CreateSequenceOp(MigrateOperation):
+            """Create a SEQUENCE."""
+
+            def __init__(self, sequence_name, **kw):
+                self.sequence_name = sequence_name
+                self.kw = kw
+
+            @classmethod
+            def create_sequence(cls, operations, sequence_name, **kw):
+                """Issue a "CREATE SEQUENCE" instruction."""
+
+                op = CreateSequenceOp(sequence_name, **kw)
+                return operations.invoke(op)
+
+        @Operations.implementation_for(CreateSequenceOp)
+        def create_sequence(operations, operation):
+            operations.execute("CREATE SEQUENCE %s" % operation.sequence_name)
+
+        context = op_fixture()
+        op.create_sequence('foob')
+        context.assert_("CREATE SEQUENCE foob")
diff --git a/tests/test_script_production.py b/tests/test_script_production.py
index 1f380ab..3ce6200 100644
--- a/tests/test_script_production.py
+++ b/tests/test_script_production.py
@@ -1,15 +1,20 @@
 from alembic.testing.fixtures import TestBase
-from alembic.testing import eq_, ne_, is_, assert_raises_message
+from alembic.testing import eq_, ne_, assert_raises_message
 from alembic.testing.env import clear_staging_env, staging_env, \
     _get_staging_directory, _no_sql_testing_config, env_file_fixture, \
     script_file_fixture, _testing_config, _sqlite_testing_config, \
-    three_rev_fixture, _multi_dir_testing_config
+    three_rev_fixture, _multi_dir_testing_config, write_script,\
+    _sqlite_file_db
 from alembic import command
 from alembic.script import ScriptDirectory
 from alembic.environment import EnvironmentContext
+from alembic.testing import mock
 from alembic import util
+from alembic.operations import ops
 import os
 import datetime
+import sqlalchemy as sa
+from sqlalchemy.engine.reflection import Inspector
 
 env, abc, def_ = None, None, None
 
@@ -214,6 +219,174 @@ class RevisionCommandTest(TestBase):
         )
 
 
+class CustomizeRevisionTest(TestBase):
+    def setUp(self):
+        self.env = staging_env()
+        self.cfg = _multi_dir_testing_config()
+        self.cfg.set_main_option("revision_environment", "true")
+
+        script = ScriptDirectory.from_config(self.cfg)
+        # MARKMARK
+        self.model1 = util.rev_id()
+        self.model2 = util.rev_id()
+        self.model3 = util.rev_id()
+        for model, name in [
+            (self.model1, "model1"),
+            (self.model2, "model2"),
+            (self.model3, "model3"),
+        ]:
+            script.generate_revision(
+                model, name, refresh=True,
+                version_path=os.path.join(_get_staging_directory(), name),
+                head="base")
+
+            write_script(script, model, """\
+"%s"
+revision = '%s'
+down_revision = None
+branch_labels = ['%s']
+
+from alembic import op
+
+def upgrade():
+    pass
+
+def downgrade():
+    pass
+
+""" % (name, model, name))
+
+    def tearDown(self):
+        clear_staging_env()
+
+    def _env_fixture(self, fn, target_metadata):
+        self.engine = engine = _sqlite_file_db()
+
+        def run_env(self):
+            from alembic import context
+
+            with engine.connect() as connection:
+                context.configure(
+                    connection=connection,
+                    target_metadata=target_metadata,
+                    process_revision_directives=fn)
+                with context.begin_transaction():
+                    context.run_migrations()
+
+        return mock.patch(
+            "alembic.script.base.ScriptDirectory.run_env",
+            run_env
+        )
+
+    def test_new_locations_no_autogen(self):
+        m = sa.MetaData()
+
+        def process_revision_directives(context, rev, generate_revisions):
+            generate_revisions[:] = [
+                ops.MigrationScript(
+                    util.rev_id(),
+                    ops.UpgradeOps(),
+                    ops.DowngradeOps(),
+                    version_path=os.path.join(
+                        _get_staging_directory(), "model1"),
+                    head="model1@head"
+                ),
+                ops.MigrationScript(
+                    util.rev_id(),
+                    ops.UpgradeOps(),
+                    ops.DowngradeOps(),
+                    version_path=os.path.join(
+                        _get_staging_directory(), "model2"),
+                    head="model2@head"
+                ),
+                ops.MigrationScript(
+                    util.rev_id(),
+                    ops.UpgradeOps(),
+                    ops.DowngradeOps(),
+                    version_path=os.path.join(
+                        _get_staging_directory(), "model3"),
+                    head="model3@head"
+                ),
+            ]
+
+        with self._env_fixture(process_revision_directives, m):
+            revs = command.revision(self.cfg, message="some message")
+
+        script = ScriptDirectory.from_config(self.cfg)
+
+        for rev, model in [
+            (revs[0], "model1"),
+            (revs[1], "model2"),
+            (revs[2], "model3"),
+        ]:
+            rev_script = script.get_revision(rev.revision)
+            eq_(
+                rev_script.path,
+                os.path.abspath(os.path.join(
+                    _get_staging_directory(), model,
+                    "%s_.py" % (rev_script.revision, )
+                ))
+            )
+            assert os.path.exists(rev_script.path)
+
+    def test_autogen(self):
+        m = sa.MetaData()
+        sa.Table('t', m, sa.Column('x', sa.Integer))
+
+        def process_revision_directives(context, rev, generate_revisions):
+            existing_upgrades = generate_revisions[0].upgrade_ops
+            existing_downgrades = generate_revisions[0].downgrade_ops
+
+            # model1 will run the upgrades, e.g. create the table,
+            # model2 will run the downgrades as upgrades, e.g. drop
+            # the table again
+
+            generate_revisions[:] = [
+                ops.MigrationScript(
+                    util.rev_id(),
+                    existing_upgrades,
+                    ops.DowngradeOps(),
+                    version_path=os.path.join(
+                        _get_staging_directory(), "model1"),
+                    head="model1@head"
+                ),
+                ops.MigrationScript(
+                    util.rev_id(),
+                    existing_downgrades,
+                    ops.DowngradeOps(),
+                    version_path=os.path.join(
+                        _get_staging_directory(), "model2"),
+                    head="model2@head"
+                )
+            ]
+
+        with self._env_fixture(process_revision_directives, m):
+            command.upgrade(self.cfg, "heads")
+
+            eq_(
+                Inspector.from_engine(self.engine).get_table_names(),
+                ["alembic_version"]
+            )
+
+            command.revision(
+                self.cfg, message="some message",
+                autogenerate=True)
+
+            command.upgrade(self.cfg, "model1@head")
+
+            eq_(
+                Inspector.from_engine(self.engine).get_table_names(),
+                ["alembic_version", "t"]
+            )
+
+            command.upgrade(self.cfg, "model2@head")
+
+            eq_(
+                Inspector.from_engine(self.engine).get_table_names(),
+                ["alembic_version"]
+            )
+
+
 class MultiDirRevisionCommandTest(TestBase):
     def setUp(self):
         self.env = staging_env()