diff options
Diffstat (limited to 'alembic')
| -rw-r--r-- | alembic/autogenerate/api.py | 82 | ||||
| -rw-r--r-- | alembic/autogenerate/compare.py | 11 | ||||
| -rw-r--r-- | alembic/autogenerate/render.py | 20 | ||||
| -rw-r--r-- | alembic/operations/ops.py | 39 |
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) |