expose on_version_apply callback to context users

Change-Id: I694e26f7d161dcaf4f035277c8317ff6ffe41680
Pull-request: https://bitbucket.org/zzzeek/alembic/pull-requests/67

2 files changed, 131 insertions, 2 deletions

diff --git a/alembic/runtime/environment.py b/alembic/runtime/environment.py
index edabc06..613b745 100644
--- a/alembic/runtime/environment.py
+++ b/alembic/runtime/environment.py
@@ -307,6 +307,7 @@ class EnvironmentContext(util.ModuleClsProxy):
                   alembic_module_prefix="op.",
                   sqlalchemy_module_prefix="sa.",
                   user_module_prefix=None,
+                  on_version_apply=None,
                   **kw
                   ):
         """Configure a :class:`.MigrationContext` within this
@@ -416,6 +417,24 @@ class EnvironmentContext(util.ModuleClsProxy):
             flag and additionally established that the Alembic version table
             has a primary key constraint by default.
 
+        :param on_version_apply: a callable or collection of callables to be
+            run for each migration step.
+            The callables will be run in the order they are given, once for
+            each migration step, after the respective operation has been
+            applied but before its transaction is finalized.
+            Each callable accepts no positional arguments and the following
+            keyword arguments:
+
+            * ``ctx``: the :class:`.MigrationContext` running the migration,
+            * ``step``: a :class:`.MigrationInfo` representing the
+              step currently being applied,
+            * ``heads``: a collection of version strings representing the
+              current heads,
+            * ``run_args``: the ``**kwargs`` passed to :meth:`.run_migrations`.
+
+            .. versionadded:: 0.9.3
+
+
         Parameters specific to the autogenerate feature, when
         ``alembic revision`` is run with the ``--autogenerate`` feature:
 
@@ -732,7 +751,6 @@ class EnvironmentContext(util.ModuleClsProxy):
 
              :paramref:`.command.revision.process_revision_directives`
 
-
         Parameters specific to individual backends:
 
         :param mssql_batch_separator: The "batch separator" which will
@@ -774,6 +792,7 @@ class EnvironmentContext(util.ModuleClsProxy):
         opts['user_module_prefix'] = user_module_prefix
         opts['literal_binds'] = literal_binds
         opts['process_revision_directives'] = process_revision_directives
+        opts['on_version_apply'] = util.to_tuple(on_version_apply, default=())
 
         if render_item is not None:
             opts['render_item'] = render_item
diff --git a/alembic/runtime/migration.py b/alembic/runtime/migration.py
index 13e6ebe..5b95208 100644
--- a/alembic/runtime/migration.py
+++ b/alembic/runtime/migration.py
@@ -70,6 +70,7 @@ class MigrationContext(object):
         transactional_ddl = opts.get("transactional_ddl")
         self._transaction_per_migration = opts.get(
             "transaction_per_migration", False)
+        self.on_version_apply_callbacks = opts.get('on_version_apply', ())
 
         if as_sql:
             self.connection = self._stdout_connection(connection)
@@ -334,6 +335,11 @@ class MigrationContext(object):
                 # and row-targeted updates and deletes, it's simpler for now
                 # just to run the operations on every version
                 head_maintainer.update_to_step(step)
+                for callback in self.on_version_apply_callbacks:
+                    callback(ctx=self,
+                             step=step.info,
+                             heads=set(head_maintainer.heads),
+                             run_args=kw)
 
             if not starting_in_transaction and not self.as_sql and \
                 not self.impl.transactional_ddl and \
@@ -534,6 +540,95 @@ class HeadMaintainer(object):
             self._update_version(from_, to_)
 
 
+class MigrationInfo(object):
+    """Exposes information about a migration step to a callback listener.
+
+    The :class:`.MigrationInfo` object is available exclusively for the
+    benefit of the :paramref:`.EnvironmentContext.on_version_apply`
+    callback hook.
+
+    .. versionadded:: 0.9.3
+
+    """
+
+    is_upgrade = None
+    """True/False: indicates whether this operation ascends or descends the
+    version tree."""
+
+    is_stamp = None
+    """True/False: indicates whether this operation is a stamp (i.e. whether
+    it results in any actual database operations)."""
+
+    up_revision_id = None
+    """Version string corresponding to :attr:`.Revision.revision`."""
+
+    down_revision_ids = None
+    """Tuple of strings representing the base revisions of this migration step.
+
+    If empty, this represents a root revision; otherwise, the first item
+    corresponds to :attr:`.Revision.down_revision`, and the rest are inferred
+    from dependencies.
+    """
+
+    revision_map = None
+    """The revision map inside of which this operation occurs."""
+
+    def __init__(self, revision_map, is_upgrade, is_stamp, up_revision,
+                 down_revisions):
+        self.revision_map = revision_map
+        self.is_upgrade = is_upgrade
+        self.is_stamp = is_stamp
+        self.up_revision_id = up_revision
+        self.down_revision_ids = util.to_tuple(down_revisions)
+
+    @property
+    def is_migration(self):
+        """True/False: indicates whether this operation is a migration.
+
+        At present this is true if and only the migration is not a stamp.
+        If other operation types are added in the future, both this attribute
+        and :attr:`~.MigrationInfo.is_stamp` will be false.
+        """
+        return not self.is_stamp
+
+    @property
+    def source_revision_ids(self):
+        """Active revisions before this migration step is applied."""
+        revs = self.down_revision_ids if self.is_upgrade \
+            else self.up_revision_id
+        return util.to_tuple(revs, default=())
+
+    @property
+    def destination_revision_ids(self):
+        """Active revisions after this migration step is applied."""
+        revs = self.up_revision_id if self.is_upgrade \
+            else self.down_revision_ids
+        return util.to_tuple(revs, default=())
+
+    @property
+    def up_revision(self):
+        """Get :attr:`~MigrationInfo.up_revision_id` as a :class:`.Revision`."""
+        return self.revision_map.get_revision(self.up_revision_id)
+
+    @property
+    def down_revisions(self):
+        """Get :attr:`~MigrationInfo.down_revision_ids` as a tuple of
+        :class:`Revisions <.Revision>`."""
+        return self.revision_map.get_revisions(self.down_revision_ids)
+
+    @property
+    def source_revisions(self):
+        """Get :attr:`~MigrationInfo.source_revision_ids` as a tuple of
+        :class:`Revisions <.Revision>`."""
+        return self.revision_map.get_revisions(self.source_revision_ids)
+
+    @property
+    def destination_revisions(self):
+        """Get :attr:`~MigrationInfo.destination_revision_ids` as a tuple of
+        :class:`Revisions <.Revision>`."""
+        return self.revision_map.get_revisions(self.destination_revision_ids)
+
+
 class MigrationStep(object):
     @property
     def name(self):
@@ -759,14 +854,22 @@ class RevisionStep(MigrationStep):
     def insert_version_num(self):
         return self.revision.revision
 
+    @property
+    def info(self):
+        return MigrationInfo(revision_map=self.revision_map,
+                             up_revision=self.revision.revision,
+                             down_revisions=self.revision._all_down_revisions,
+                             is_upgrade=self.is_upgrade, is_stamp=False)
+
 
 class StampStep(MigrationStep):
-    def __init__(self, from_, to_, is_upgrade, branch_move):
+    def __init__(self, from_, to_, is_upgrade, branch_move, revision_map=None):
         self.from_ = util.to_tuple(from_, default=())
         self.to_ = util.to_tuple(to_, default=())
         self.is_upgrade = is_upgrade
         self.branch_move = branch_move
         self.migration_fn = self.stamp_revision
+        self.revision_map = revision_map
 
     doc = None
 
@@ -836,3 +939,10 @@ class StampStep(MigrationStep):
 
     def should_unmerge_branches(self, heads):
         return len(self.to_) > 1
+
+    @property
+    def info(self):
+        up, down = (self.to_, self.from_) if self.is_upgrade \
+            else (self.from_, self.to_)
+        return MigrationInfo(self.revision_map, up, down, self.is_upgrade,
+                             True)