diff options
-rw-r--r-- | alembic/command.py | 7 | ||||
-rw-r--r-- | alembic/script.py | 77 | ||||
-rw-r--r-- | alembic/util.py | 23 | ||||
-rw-r--r-- | templates/generic/script.py.mako | 4 | ||||
-rw-r--r-- | templates/multidb/script.py.mako | 21 | ||||
-rw-r--r-- | templates/pylons/script.py.mako | 4 |
6 files changed, 128 insertions, 8 deletions
diff --git a/alembic/command.py b/alembic/command.py index c9a80ab..8278d6e 100644 --- a/alembic/command.py +++ b/alembic/command.py @@ -2,6 +2,7 @@ from alembic.script import ScriptDirectory from alembic import options, util import os import sys +import uuid def list_templates(opts): """List available templates""" @@ -28,12 +29,13 @@ def init(opts): util.status("Creating directory %s" % os.path.abspath(dir_), os.makedirs, dir_) - script = ScriptDirectory(dir_) + script = ScriptDirectory(dir_, opts) template_dir = os.path.join(opts.get_template_directory(), opts.cmd_line_options.template) if not os.access(template_dir, os.F_OK): opts.err("No such template %r" % opts.cmd_line_options.template) + for file_ in os.listdir(template_dir): if file_ == 'alembic.ini.mako': config_file = os.path.abspath(opts.cmd_line_options.config) @@ -57,6 +59,9 @@ def init(opts): def revision(opts): """Create a new revision file.""" + + script = ScriptDirectory.from_options(opts) + script.generate_rev(uuid.uuid4()) def upgrade(opts): """Upgrade to the latest version.""" diff --git a/alembic/script.py b/alembic/script.py index 385041e..8cac691 100644 --- a/alembic/script.py +++ b/alembic/script.py @@ -1,15 +1,52 @@ import os from alembic import util import shutil +import re +import inspect + +_uuid_re = re.compile(r'[a-z0-9]{16}') +_mod_def_re = re.compile(r'(upgrade|downgrade)_([a-z0-9]{16})') class ScriptDirectory(object): - def __init__(self, dir): + def __init__(self, dir, options): self.dir = dir + self.options = otions @classmethod def from_options(cls, options): - return ScriptDirectory(options.get_main_option('script_location')) + return ScriptDirectory( + options.get_main_option('script_location'), + options) + @util.memoized_property + def _revision_map(self): + for file_ in os.listdir(self.dir): + script = Script.from_file(self.dir, file_) + if script is None: + continue + map_[script.upgrade] = script + return map_ + + def _get_head(self): + # TODO: keep map sorted chronologically + + for script in self._revision_map.values(): + if script.upgrade is None \ + and script.downgrade in self._revision_map: + return script + else: + return None + + def _get_origin(self): + # TODO: keep map sorted chronologically + + for script in self._revision_map.values(): + if script.downgrade is None \ + and script.upgrade in self._revision_map: + return script + else: + return None + def generate_template(self, src, dest, **kw): util.status("Generating %s" % os.path.abspath(src), util.template_to_file, @@ -22,4 +59,40 @@ class ScriptDirectory(object): util.status("Generating %s" % os.path.abspath(dest), shutil.copy, src, dest) + + + def generate_rev(self, revid): + current_head = self._get_head() + self.generate_template( + os.path.join(self.dir, "script.py.mako", + up_revision=revid, + down_revision=current_head.upgrade if current_head else None + ) + ) + +class Script(object): + def __init__(self, module): + self.module = module + self.upgrade = self.downgrade = None + for name in dir(module): + m = _mod_def_re.match(name) + if not m: + continue + fn = getattr(module, name) + if not inspect.isfunction(fn): + continue + if m.group(1) == 'upgrade': + self.upgrade = m.group(2) + elif m.group(1) == 'downgrade': + self.downgrade = m.group(2) + if not self.downgrade and not self.upgrade: + raise Exception("Script %s has no upgrade or downgrade path" % module) + + @classmethod + def from_path(cls, dir_, filename): + if not _uuid_re.match(filename): + return None + + module = util.load_python_file(dir_, filename) + return Script(module)
\ No newline at end of file diff --git a/alembic/util.py b/alembic/util.py index 88085ee..df5bfd8 100644 --- a/alembic/util.py +++ b/alembic/util.py @@ -3,6 +3,7 @@ import sys import os import textwrap from sqlalchemy import util +import imp NO_VALUE = util.symbol("NO_VALUE") @@ -40,3 +41,25 @@ def msg(msg, newline=True): for line in lines[0:-1]: sys.stdout.write(" " +line + "\n") sys.stdout.write(" " + lines[-1] + ("\n" if newline else "")) + +def load_python_file(dir_, filename): + """Load a file from the given path as a Python module.""" + + module_id = re.sub(r'\W', "_", filename) + path = os.path.join(dir_, filename) + module = imp.load_source(module_id, path, open(path, 'rb')) + del sys.modules[module_id] + return module + +class memoized_property(object): + """A read-only @property that is only evaluated once.""" + def __init__(self, fget, doc=None): + self.fget = fget + self.__doc__ = doc or fget.__doc__ + self.__name__ = fget.__name__ + + def __get__(self, obj, cls): + if obj is None: + return None + obj.__dict__[self.__name__] = result = self.fget(obj) + return result diff --git a/templates/generic/script.py.mako b/templates/generic/script.py.mako index 4181f3a..596027f 100644 --- a/templates/generic/script.py.mako +++ b/templates/generic/script.py.mako @@ -3,5 +3,9 @@ from alembic.op import * def upgrade_${up_revision}(): pass +% if down_revision: def downgrade_${down_revision}(): pass +% else: +# this is the origin node, no downgrade ! +% endif diff --git a/templates/multidb/script.py.mako b/templates/multidb/script.py.mako index ee22014..112b026 100644 --- a/templates/multidb/script.py.mako +++ b/templates/multidb/script.py.mako @@ -1,9 +1,20 @@ from alembic.op import * -% for engine in engines: -def upgrade_${engine}_${up_revision}(): - pass +def upgrade_${up_revision}(engine): + eval("upgrade_%s_${up_revision}" % engine.name)() -def downgrade_${engine}_${down_revision}(): - pass +% if down_revision: +def downgrade_${down_revision}(engine): + eval("upgrade_%s_${down_revision}" % engine.name)() +% else: +# this is the origin node, no downgrade ! +% endif + + +% for engine in ["engine1", "engine2"]: + def upgrade_${engine}_${up_revision}(): + pass + + def downgrade_${engine}_${down_revision}(): + pass % endfor
\ No newline at end of file diff --git a/templates/pylons/script.py.mako b/templates/pylons/script.py.mako index 4181f3a..596027f 100644 --- a/templates/pylons/script.py.mako +++ b/templates/pylons/script.py.mako @@ -3,5 +3,9 @@ from alembic.op import * def upgrade_${up_revision}(): pass +% if down_revision: def downgrade_${down_revision}(): pass +% else: +# this is the origin node, no downgrade ! +% endif |