From: Christopher Allan Webber <cwebber@dustycloud.org>
Date: Sat, 9 Jul 2011 19:50:41 +0000 (-0500)
Subject: New migration utility code.... I haven't tested this! ;)
X-Git-Url: https://vcs.fsf.org/?a=commitdiff_plain;h=51dcfb56827fc00eb4761ce5421a23cae3177304;p=mediagoblin.git

New migration utility code.... I haven't tested this! ;)

I think it's looking right though.
 - Provides MigrationManager which should have plenty of utilities for
   doing migrations hopefully correctly :)
 - Provides RegisterMigration which should be able to decorate
   migrations and register them in doing so
---

diff --git a/mediagoblin/db/util.py b/mediagoblin/db/util.py
index 70c37945..ebf8409d 100644
--- a/mediagoblin/db/util.py
+++ b/mediagoblin/db/util.py
@@ -37,6 +37,11 @@ from mongokit import ObjectId
 from mediagoblin.db.indexes import ACTIVE_INDEXES, DEPRECATED_INDEXES
 
 
+################
+# Indexing tools
+################
+
+
 def add_new_indexes(database, active_indexes=ACTIVE_INDEXES):
     """
     Add any new indexes to the database.
@@ -99,3 +104,134 @@ def remove_deprecated_indexes(database, deprecated_indexes=DEPRECATED_INDEXES):
                 indexes_removed.append((collection_name, index_name))
 
     return indexes_removed
+
+
+#################
+# Migration tools
+#################
+
+# The default migration registry...
+# 
+# Don't set this yourself!  RegisterMigration will automatically fill
+# this with stuff via decorating methods in migrations.py
+
+MIGRATIONS = {}
+
+
+class RegisterMigration(object):
+    def __init__(self, migration_number, migration_registry=MIGRATIONS):
+        self.migration_number = migration_number
+        self.migration_registry = migration_registry
+
+    def __call__(self, migration):
+        self.migration_registry[self.migration_number] = migration
+        return migration
+
+
+class MigrationManager(object):
+    """
+    Migration handling tool.
+
+    Takes information about a database, lets you update the database
+    to the latest migrations, etc.
+    """
+    def __init__(self, database, migration_registry=MIGRATIONS):
+        """
+        Args:
+         - database: database we're going to migrate
+         - migration_registry: where we should find all migrations to
+           run
+        """
+        self.database = database
+        self.migration_registry = migration_registry
+        self._sorted_migrations = None
+
+    @property
+    def sorted_migrations(self):
+        """
+        Sort migrations if necessary and store in self._sorted_migrations
+        """
+        if not self._sorted_migrations:
+            self._sorted_migrations = sorted(
+                self.migration_registry.items(),
+                # sort on the key... the migration number
+                key=lambda migration_tuple: migration_tuple[0])
+
+        return self._sorted_migrations
+
+    def latest_migration(self):
+        """
+        Return a tuple like:
+          (migration_number, migration_func)
+
+        Where migration_number is the number of the latest migration
+        and migration func is the actual function that would be run.
+        """
+        return self.sorted_migrations[-1]
+
+    def set_current_migration(self, migration_number=None):
+        """
+        Set the migration in the database to migration_number
+        """
+        # Add the mediagoblin migration if necessary
+        self.database['app_metadata'].update(
+            {'_id': 'mediagoblin'},
+            {'$set': {'current_migration': migration_number}},
+            upsert=True)
+
+    def database_current_migration(self, install_if_missing=True):
+        """
+        Return the current migration in the database.
+        """
+        mgoblin_metadata = self.database['app_metadata'].find_one(
+            {'_id': 'mediagoblin'})
+        if not mgoblin_metadata:
+            if install_if_missing:
+                latest_migration = self.latest_migration()
+                self.set_current_migration(latest_migration)
+                return latest_migration
+            else:
+                return None
+        else:
+            return mgoblin_metadata['current_migration']
+
+    def database_at_latest_migration(self):
+        """
+        See if the database is at the latest migration.
+        Returns a boolean.
+        """
+        current_migration = self.database_current_migration()
+        return current_migration == self.latest_migration()
+
+    def migrations_to_run(self):
+        """
+        Get a list of migrations to run still, if any.
+        """
+        db_current_migration = self.database_current_migration()
+        return [
+            (migration_number, migration_func)
+            for migration_number, migration_func in self.sorted_migrations
+            if migration_number > db_current_migration]
+
+    def iteratively_migrate(self):
+        """
+        Iteratively run all migrations.
+
+        Useful if you need to print some message about each migration
+        after you run it.
+
+        Each time you loop over this, it'll return the migration
+        number and migration function.
+        """
+        for migration_number, migration_func in self.migrations_to_run():
+            migration_func(self.database)
+            self.set_current_migration(migration_number)
+            yield migration_number, migration_func
+        
+    def run_outdated_migrations(self):
+        """
+        Install all migrations that need to be installed, quietly.
+        """
+        for migration_number, migration_func in self.iteratively_migrate():
+            # No need to say anything... we're just migrating iteratively.
+            pass