Caution the admins about deleting the users' media though.
[mediagoblin.git] / mediagoblin / db / util.py
index 5f0699c1d0d960807f1e5ff356f801d7d6bcd7b5..0f3220d2d89019846f63720d511f164d347a6bf3 100644 (file)
@@ -86,18 +86,25 @@ def remove_deprecated_indexes(database, deprecated_indexes=DEPRECATED_INDEXES):
     Args:
      - database: pymongo or mongokit database instance.
      - deprecated_indexes: the indexes to deprecate in the pattern of:
-       {'collection': ['index_identifier1', 'index_identifier2']}
+       {'collection_name': {
+            'identifier': {
+                'index': [index_foo_goes_here],
+                'unique': True}}
+
+       (... although we really only need the 'identifier' here, as the
+       rest of the information isn't used in this case.  But it's kept
+       around so we can remember what it was)
 
     Returns:
       A list of indexes removed in form ('collection', 'index_name')
     """
     indexes_removed = []
 
-    for collection_name, index_names in deprecated_indexes.iteritems():
+    for collection_name, indexes in deprecated_indexes.iteritems():
         collection = database[collection_name]
         collection_indexes = collection.index_information().keys()
 
-        for index_name in index_names:
+        for index_name, index_data in indexes.iteritems():
             if index_name in collection_indexes:
                 collection.drop_index(index_name)
 
@@ -115,6 +122,9 @@ def remove_deprecated_indexes(database, deprecated_indexes=DEPRECATED_INDEXES):
 # Don't set this yourself!  RegisterMigration will automatically fill
 # this with stuff via decorating methods in migrations.py
 
+class MissingCurrentMigration(Exception): pass
+
+
 MIGRATIONS = {}
 
 
@@ -136,6 +146,10 @@ class RegisterMigration(object):
     0.  0 is the default "no migrations" state!
     """
     def __init__(self, migration_number, migration_registry=MIGRATIONS):
+        assert migration_number > 0, "Migration number must be > 0!"
+        assert not migration_registry.has_key(migration_number), \
+            "Duplicate migration numbers detected!  That's not allowed!"
+
         self.migration_number = migration_number
         self.migration_registry = migration_registry
 
@@ -162,6 +176,16 @@ class MigrationManager(object):
         self.migration_registry = migration_registry
         self._sorted_migrations = None
 
+    def _ensure_current_migration_record(self):
+        """
+        If there isn't a database[u'app_metadata'] mediagoblin entry
+        with the 'current_migration', throw an error.
+        """
+        if self.database_current_migration() is None:
+            raise MissingCurrentMigration(
+                "Tried to call function which requires "
+                "'current_migration' set in database")
+
     @property
     def sorted_migrations(self):
         """
@@ -186,31 +210,37 @@ class MigrationManager(object):
             # If no migrations have been set, we start at 0.
             return 0
 
-    def set_current_migration(self, migration_number=None):
+    def set_current_migration(self, migration_number):
         """
         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}},
+        self.database[u'app_metadata'].update(
+            {u'_id': u'mediagoblin'},
+            {u'$set': {u'current_migration': migration_number}},
             upsert=True)
 
-    def database_current_migration(self, install_if_missing=True):
+    def install_migration_version_if_missing(self):
+        """
+        Sets the migration to the latest version if no migration
+        version at all is set.
+        """
+        mgoblin_metadata = self.database[u'app_metadata'].find_one(
+            {u'_id': u'mediagoblin'})
+        if not mgoblin_metadata:
+            latest_migration = self.latest_migration()
+            self.set_current_migration(latest_migration)
+
+    def database_current_migration(self):
         """
         Return the current migration in the database.
         """
-        mgoblin_metadata = self.database['app_metadata'].find_one(
-            {'_id': 'mediagoblin'})
+        mgoblin_metadata = self.database[u'app_metadata'].find_one(
+            {u'_id': u'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
+            return None
         else:
-            return mgoblin_metadata['current_migration']
+            return mgoblin_metadata[u'current_migration']
 
     def database_at_latest_migration(self):
         """
@@ -223,32 +253,39 @@ class MigrationManager(object):
     def migrations_to_run(self):
         """
         Get a list of migrations to run still, if any.
+        
+        Note that calling this will set your migration version to the
+        latest version if it isn't installed to anything yet!
         """
+        self._ensure_current_migration_record()
+
         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):
+    def migrate_new(self, pre_callback=None, post_callback=None):
         """
-        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.
+        Run all migrations.
+
+        Includes two optional args:
+         - pre_callback: if called, this is a callback on something to
+           run pre-migration.  Takes (migration_number, migration_func)
+           as arguments
+         - pre_callback: if called, this is a callback on something to
+           run post-migration.  Takes (migration_number, migration_func)
+           as arguments
         """
+        # If we aren't set to any version number, presume we're at the
+        # latest (which means we'll do nothing here...)
+        self.install_migration_version_if_missing()
+
         for migration_number, migration_func in self.migrations_to_run():
+            if pre_callback:
+                pre_callback(migration_number, migration_func)
             migration_func(self.database)
             self.set_current_migration(migration_number)
-            yield migration_number, migration_func
-        
-    def run_new_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 quietly.
-            pass
+            if post_callback:
+                post_callback(migration_number, migration_func)