beginning to build the revision system

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