diff options
Diffstat (limited to 'alembic/runtime')
-rw-r--r-- | alembic/runtime/environment.py | 21 | ||||
-rw-r--r-- | alembic/runtime/migration.py | 112 |
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) |