Add process_revision_directives param to command.revision()

This allows programmatic use of command.revision() to
specify an ad-hoc process_revision_directives callable,
rather than having it placed via the env.py script.

Co-authored-by: Tyson Holub
Change-Id: Ief1c11fd2a6f10e712851145d6d190a3b167817c
Pull-request: https://github.com/zzzeek/alembic/pull/35

6 files changed, 107 insertions, 9 deletions

diff --git a/alembic/autogenerate/api.py b/alembic/autogenerate/api.py
index 61f4036..42c12c0 100644
--- a/alembic/autogenerate/api.py
+++ b/alembic/autogenerate/api.py
@@ -255,7 +255,7 @@ class AutogenContext(object):
         self.metadata = metadata = opts.get('target_metadata', None) \
             if metadata is None else metadata
 
-        if metadata is None and \
+        if autogenerate and metadata is None and \
                 migration_context is not None and \
                 migration_context.script is not None:
             raise util.CommandError(
@@ -325,10 +325,12 @@ class RevisionContext(object):
     """Maintains configuration and state that's specific to a revision
     file generation operation."""
 
-    def __init__(self, config, script_directory, command_args):
+    def __init__(self, config, script_directory, command_args,
+                 process_revision_directives=None):
         self.config = config
         self.script_directory = script_directory
         self.command_args = command_args
+        self.process_revision_directives = process_revision_directives
         self.template_args = {
             'config': config  # Let templates use config for
                               # e.g. multiple databases
@@ -404,6 +406,10 @@ class RevisionContext(object):
             compare._populate_migration_script(
                 autogen_context, migration_script)
 
+        if self.process_revision_directives:
+            self.process_revision_directives(
+                migration_context, rev, self.generated_revisions)
+
         hook = migration_context.opts['process_revision_directives']
         if hook:
             hook(migration_context, rev, self.generated_revisions)
diff --git a/alembic/command.py b/alembic/command.py
index 0034c18..9a333a0 100644
--- a/alembic/command.py
+++ b/alembic/command.py
@@ -68,7 +68,8 @@ def init(config, directory, template='generic'):
 def revision(
         config, message=None, autogenerate=False, sql=False,
         head="head", splice=False, branch_label=None,
-        version_path=None, rev_id=None, depends_on=None):
+        version_path=None, rev_id=None, depends_on=None,
+        process_revision_directives=None):
     """Create a new revision file."""
 
     script_directory = ScriptDirectory.from_config(config)
@@ -80,7 +81,8 @@ def revision(
         version_path=version_path, rev_id=rev_id, depends_on=depends_on
     )
     revision_context = autogen.RevisionContext(
-        config, script_directory, command_args)
+        config, script_directory, command_args,
+        process_revision_directives=process_revision_directives)
 
     environment = util.asbool(
         config.get_main_option("revision_environment")
diff --git a/alembic/config.py b/alembic/config.py
index ada59dd..3774cd3 100644
--- a/alembic/config.py
+++ b/alembic/config.py
@@ -452,8 +452,8 @@ class CommandLine(object):
 
         try:
             fn(config,
-               *[getattr(options, k) for k in positional],
-               **dict((k, getattr(options, k)) for k in kwarg)
+               *[getattr(options, k, None) for k in positional],
+               **dict((k, getattr(options, k, None)) for k in kwarg)
                )
         except util.CommandError as e:
             if options.raiseerr:
diff --git a/docs/build/changelog.rst b/docs/build/changelog.rst
index ca6834d..caf3f90 100644
--- a/docs/build/changelog.rst
+++ b/docs/build/changelog.rst
@@ -16,6 +16,16 @@ Changelog
       for the inner type as well as the outer type.  Pull request courtesy
       Paul Brackin.
 
+    .. change:: process_revision_directives_command
+      :tags: feature, autogenerate
+
+      Added a keyword argument ``process_revision_directives`` to the
+      :func:`.command.revision` API call.  This function acts in the
+      same role as the environment-level call, and allows API use of the
+      command to drop in an ad-hoc directive process function.  This
+      function can be used among other things to place a complete
+      :class:`.MigrationScript` structure in place.
+
     .. change:: fk_schema_compare
       :tags: bug, operations
 
diff --git a/tests/test_command.py b/tests/test_command.py
index 55152c9..58a827f 100644
--- a/tests/test_command.py
+++ b/tests/test_command.py
@@ -1,6 +1,7 @@
 from alembic import command
 from io import TextIOWrapper, BytesIO
 from alembic.script import ScriptDirectory
+from alembic import config
 from alembic.testing.fixtures import TestBase, capture_context_buffer
 from alembic.testing.env import staging_env, _sqlite_testing_config, \
     three_rev_fixture, clear_staging_env, _no_sql_testing_config, \
@@ -603,3 +604,45 @@ class EditTest(TestBase):
         with mock.patch('alembic.util.edit') as edit:
             command.edit(self.cfg, "current")
             edit.assert_called_with(expected_call_arg)
+
+
+class CommandLineTest(TestBase):
+    @classmethod
+    def setup_class(cls):
+        cls.env = staging_env()
+        cls.cfg = _sqlite_testing_config()
+        cls.a, cls.b, cls.c = three_rev_fixture(cls.cfg)
+
+    def test_run_cmd_args_missing(self):
+        canary = mock.Mock()
+
+        orig_revision = command.revision
+
+        # the command function has "process_revision_directives"
+        # however the ArgumentParser does not.  ensure things work
+        def revision(
+            config, message=None, autogenerate=False, sql=False,
+            head="head", splice=False, branch_label=None,
+            version_path=None, rev_id=None, depends_on=None,
+            process_revision_directives=None
+        ):
+            canary(
+                config, message=message
+            )
+
+        revision.__module__ = 'alembic.command'
+
+        # CommandLine() pulls the function into the ArgumentParser
+        # and needs the full signature, so we can't patch the "revision"
+        # command normally as ArgumentParser gives us no way to get to it.
+        config.command.revision = revision
+        try:
+            commandline = config.CommandLine()
+            options = commandline.parser.parse_args(["revision", "-m", "foo"])
+            commandline.run_cmd(self.cfg, options)
+        finally:
+            config.command.revision = orig_revision
+        eq_(
+            canary.mock_calls,
+            [mock.call(self.cfg, message="foo")]
+        )
diff --git a/tests/test_script_production.py b/tests/test_script_production.py
index 1f1b3ea..3703014 100644
--- a/tests/test_script_production.py
+++ b/tests/test_script_production.py
@@ -222,8 +222,6 @@ class RevisionCommandTest(TestBase):
         )
 
 
-
-
 class CustomizeRevisionTest(TestBase):
     def setUp(self):
         self.env = staging_env()
@@ -412,6 +410,45 @@ def downgrade():
                 ["alembic_version"]
             )
 
+    def test_programmatic_command_option(self):
+
+        def process_revision_directives(context, rev, generate_revisions):
+            generate_revisions[0].message = "test programatic"
+            generate_revisions[0].upgrade_ops = ops.UpgradeOps(
+                ops=[
+                    ops.CreateTableOp(
+                        'test_table',
+                        [
+                            sa.Column('id', sa.Integer(), primary_key=True),
+                            sa.Column('name', sa.String(50), nullable=False)
+                        ]
+                    ),
+                ]
+            )
+            generate_revisions[0].downgrade_ops = ops.DowngradeOps(
+                ops=[
+                    ops.DropTableOp('test_table')
+                ]
+            )
+
+        with self._env_fixture(None, None):
+            rev = command.revision(
+                self.cfg,
+                head="model1@head",
+                process_revision_directives=process_revision_directives)
+
+        result = open(rev.path).read()
+        assert ("""
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('test_table',
+    sa.Column('id', sa.Integer(), nullable=False),
+    sa.Column('name', sa.String(length=50), nullable=False),
+    sa.PrimaryKeyConstraint('id')
+    )
+    # ### end Alembic commands ###
+""") in result
+
 
 class ScriptAccessorTest(TestBase):
     def test_upgrade_downgrade_ops_list_accessors(self):
@@ -900,4 +937,4 @@ def downgrade():
             eq_(
                 [rev.revision for rev in script.walk_revisions()],
                 [self.model1, self.model2, self.model3]
-            )
\ No newline at end of file
+            )