1 files changed, 54 insertions, 2 deletions
diff --git a/mediagoblin/db/migration_tools.py b/mediagoblin/db/migration_tools.py
index e39070c3..ab4487d2 100644
--- a/mediagoblin/db/migration_tools.py
+++ b/mediagoblin/db/migration_tools.py
@@ -14,14 +14,66 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+from __future__ import unicode_literals
+
+import logging
+import os
+
+from alembic import command
+from alembic.config import Config
+from alembic.migration import MigrationContext
+
+from mediagoblin.db.base import Base
 from mediagoblin.tools.common import simple_printer
 from sqlalchemy import Table
 from sqlalchemy.sql import select
 
+log = logging.getLogger(__name__)
+
+
 class TableAlreadyExists(Exception):
     pass
 
 
+class AlembicMigrationManager(object):
+
+    def __init__(self, session):
+        root_dir = os.path.abspath(os.path.dirname(os.path.dirname(
+            os.path.dirname(__file__))))
+        alembic_cfg_path = os.path.join(root_dir, 'alembic.ini')
+        self.alembic_cfg = Config(alembic_cfg_path)
+        self.session = session
+
+    def get_current_revision(self):
+        context = MigrationContext.configure(self.session.bind)
+        return context.get_current_revision()
+
+    def upgrade(self, version):
+        return command.upgrade(self.alembic_cfg, version or 'head')
+
+    def downgrade(self, version):
+        if isinstance(version, int) or version is None or version.isdigit():
+            version = 'base'
+        return command.downgrade(self.alembic_cfg, version)
+
+    def stamp(self, revision):
+        return command.stamp(self.alembic_cfg, revision=revision)
+
+    def init_tables(self):
+        Base.metadata.create_all(self.session.bind)
+        # load the Alembic configuration and generate the
+        # version table, "stamping" it with the most recent rev:
+        command.stamp(self.alembic_cfg, 'head')
+
+    def init_or_migrate(self, version=None):
+        if self.get_current_revision() is None:
+            log.info('Initializing tables and stamping it with '
+                     'the most recent migration...')
+            self.init_tables()
+        else:
+            self.upgrade(version)
+
+
 class MigrationManager(object):
     """
     Migration handling tool.
@@ -39,7 +91,7 @@ class MigrationManager(object):
          - migration_registry: where we should find all migrations to
            run
         """
-        self.name = unicode(name)
+        self.name = name
         self.models = models
         self.foundations = foundations
         self.session = session
@@ -230,7 +282,7 @@ class MigrationManager(object):
             for migration_number, migration_func in migrations_to_run:
                 self.printer(
                     u'   + Running migration %s, "%s"... ' % (
-                        migration_number, migration_func.func_name))
+                        migration_number, migration_func.__name__))
                 migration_func(self.session)
                 self.set_current_migration(migration_number)
                 self.printer('done.\n')