mediagoblin/db/mongo/util.py

   1 # GNU MediaGoblin -- federated, autonomous media hosting
   2 # Copyright (C) 2011, 2012 MediaGoblin contributors.  See AUTHORS.
   3 #
   4 # This program is free software: you can redistribute it and/or modify
   5 # it under the terms of the GNU Affero General Public License as published by
   6 # the Free Software Foundation, either version 3 of the License, or
   7 # (at your option) any later version.
   8 #
   9 # This program is distributed in the hope that it will be useful,
  10 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  11 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  12 # GNU Affero General Public License for more details.
  13 #
  14 # You should have received a copy of the GNU Affero General Public License
  15 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
  16
  17 """
  18 Utilities for database operations.
  19
  20 Some note on migration and indexing tools:
  21
  22 We store information about what the state of the database is in the
  23 'mediagoblin' document of the 'app_metadata' collection.  Keys in that
  24 document relevant to here:
  25
  26  - 'migration_number': The integer representing the current state of
  27    the migrations
  28 """
  29
  30 import copy
  31
  32 # Imports that other modules might use
  33 from pymongo import ASCENDING, DESCENDING
  34 from pymongo.errors import InvalidId
  35 from mongokit import ObjectId
  36
  37 from mediagoblin.db.mongo.indexes import ACTIVE_INDEXES, DEPRECATED_INDEXES
  38
  39
  40 ################
  41 # Indexing tools
  42 ################
  43
  44
  45 def add_new_indexes(database, active_indexes=ACTIVE_INDEXES):
  46     """
  47     Add any new indexes to the database.
  48
  49     Args:
  50      - database: pymongo or mongokit database instance.
  51      - active_indexes: indexes to possibly add in the pattern of:
  52        {'collection_name': {
  53             'identifier': {
  54                 'index': [index_foo_goes_here],
  55                 'unique': True}}
  56        where 'index' is the index to add and all other options are
  57        arguments for collection.create_index.
  58
  59     Returns:
  60       A list of indexes added in form ('collection', 'index_name')
  61     """
  62     indexes_added = []
  63
  64     for collection_name, indexes in active_indexes.iteritems():
  65         collection = database[collection_name]
  66         collection_indexes = collection.index_information().keys()
  67
  68         for index_name, index_data in indexes.iteritems():
  69             if not index_name in collection_indexes:
  70                 # Get a copy actually so we don't modify the actual
  71                 # structure
  72                 index_data = copy.copy(index_data)
  73                 index = index_data.pop('index')
  74                 collection.create_index(
  75                     index, name=index_name, **index_data)
  76
  77                 indexes_added.append((collection_name, index_name))
  78
  79     return indexes_added
  80
  81
  82 def remove_deprecated_indexes(database, deprecated_indexes=DEPRECATED_INDEXES):
  83     """
  84     Remove any deprecated indexes from the database.
  85
  86     Args:
  87      - database: pymongo or mongokit database instance.
  88      - deprecated_indexes: the indexes to deprecate in the pattern of:
  89        {'collection_name': {
  90             'identifier': {
  91                 'index': [index_foo_goes_here],
  92                 'unique': True}}
  93
  94        (... although we really only need the 'identifier' here, as the
  95        rest of the information isn't used in this case.  But it's kept
  96        around so we can remember what it was)
  97
  98     Returns:
  99       A list of indexes removed in form ('collection', 'index_name')
 100     """
 101     indexes_removed = []
 102
 103     for collection_name, indexes in deprecated_indexes.iteritems():
 104         collection = database[collection_name]
 105         collection_indexes = collection.index_information().keys()
 106
 107         for index_name, index_data in indexes.iteritems():
 108             if index_name in collection_indexes:
 109                 collection.drop_index(index_name)
 110
 111                 indexes_removed.append((collection_name, index_name))
 112
 113     return indexes_removed
 114
 115
 116 #################
 117 # Migration tools
 118 #################
 119
 120 # The default migration registry...
 121 #
 122 # Don't set this yourself!  RegisterMigration will automatically fill
 123 # this with stuff via decorating methods in migrations.py
 124
 125 class MissingCurrentMigration(Exception):
 126     pass
 127
 128
 129 MIGRATIONS = {}
 130
 131
 132 class RegisterMigration(object):
 133     """
 134     Tool for registering migrations
 135
 136     Call like:
 137
 138     @RegisterMigration(33)
 139     def update_dwarves(database):
 140         [...]
 141
 142     This will register your migration with the default migration
 143     registry.  Alternately, to specify a very specific
 144     migration_registry, you can pass in that as the second argument.
 145
 146     Note, the number of your migration should NEVER be 0 or less than
 147     0.  0 is the default "no migrations" state!
 148     """
 149     def __init__(self, migration_number, migration_registry=MIGRATIONS):
 150         assert migration_number > 0, "Migration number must be > 0!"
 151         assert migration_number not in migration_registry, \
 152             "Duplicate migration numbers detected!  That's not allowed!"
 153
 154         self.migration_number = migration_number
 155         self.migration_registry = migration_registry
 156
 157     def __call__(self, migration):
 158         self.migration_registry[self.migration_number] = migration
 159         return migration
 160
 161
 162 class MigrationManager(object):
 163     """
 164     Migration handling tool.
 165
 166     Takes information about a database, lets you update the database
 167     to the latest migrations, etc.
 168     """
 169     def __init__(self, database, migration_registry=MIGRATIONS):
 170         """
 171         Args:
 172          - database: database we're going to migrate
 173          - migration_registry: where we should find all migrations to
 174            run
 175         """
 176         self.database = database
 177         self.migration_registry = migration_registry
 178         self._sorted_migrations = None
 179
 180     def _ensure_current_migration_record(self):
 181         """
 182         If there isn't a database[u'app_metadata'] mediagoblin entry
 183         with the 'current_migration', throw an error.
 184         """
 185         if self.database_current_migration() is None:
 186             raise MissingCurrentMigration(
 187                 "Tried to call function which requires "
 188                 "'current_migration' set in database")
 189
 190     @property
 191     def sorted_migrations(self):
 192         """
 193         Sort migrations if necessary and store in self._sorted_migrations
 194         """
 195         if not self._sorted_migrations:
 196             self._sorted_migrations = sorted(
 197                 self.migration_registry.items(),
 198                 # sort on the key... the migration number
 199                 key=lambda migration_tuple: migration_tuple[0])
 200
 201         return self._sorted_migrations
 202
 203     def latest_migration(self):
 204         """
 205         Return a migration number for the latest migration, or 0 if
 206         there are no migrations.
 207         """
 208         if self.sorted_migrations:
 209             return self.sorted_migrations[-1][0]
 210         else:
 211             # If no migrations have been set, we start at 0.
 212             return 0
 213
 214     def set_current_migration(self, migration_number):
 215         """
 216         Set the migration in the database to migration_number
 217         """
 218         # Add the mediagoblin migration if necessary
 219         self.database[u'app_metadata'].update(
 220             {u'_id': u'mediagoblin'},
 221             {u'$set': {u'current_migration': migration_number}},
 222             upsert=True)
 223
 224     def install_migration_version_if_missing(self):
 225         """
 226         Sets the migration to the latest version if no migration
 227         version at all is set.
 228         """
 229         mgoblin_metadata = self.database[u'app_metadata'].find_one(
 230             {u'_id': u'mediagoblin'})
 231         if not mgoblin_metadata:
 232             latest_migration = self.latest_migration()
 233             self.set_current_migration(latest_migration)
 234
 235     def database_current_migration(self):
 236         """
 237         Return the current migration in the database.
 238         """
 239         mgoblin_metadata = self.database[u'app_metadata'].find_one(
 240             {u'_id': u'mediagoblin'})
 241         if not mgoblin_metadata:
 242             return None
 243         else:
 244             return mgoblin_metadata[u'current_migration']
 245
 246     def database_at_latest_migration(self):
 247         """
 248         See if the database is at the latest migration.
 249         Returns a boolean.
 250         """
 251         current_migration = self.database_current_migration()
 252         return current_migration == self.latest_migration()
 253
 254     def migrations_to_run(self):
 255         """
 256         Get a list of migrations to run still, if any.
 257
 258         Note that calling this will set your migration version to the
 259         latest version if it isn't installed to anything yet!
 260         """
 261         self._ensure_current_migration_record()
 262
 263         db_current_migration = self.database_current_migration()
 264
 265         return [
 266             (migration_number, migration_func)
 267             for migration_number, migration_func in self.sorted_migrations
 268             if migration_number > db_current_migration]
 269
 270     def migrate_new(self, pre_callback=None, post_callback=None):
 271         """
 272         Run all migrations.
 273
 274         Includes two optional args:
 275          - pre_callback: if called, this is a callback on something to
 276            run pre-migration.  Takes (migration_number, migration_func)
 277            as arguments
 278          - pre_callback: if called, this is a callback on something to
 279            run post-migration.  Takes (migration_number, migration_func)
 280            as arguments
 281         """
 282         # If we aren't set to any version number, presume we're at the
 283         # latest (which means we'll do nothing here...)
 284         self.install_migration_version_if_missing()
 285
 286         for migration_number, migration_func in self.migrations_to_run():
 287             if pre_callback:
 288                 pre_callback(migration_number, migration_func)
 289             migration_func(self.database)
 290             self.set_current_migration(migration_number)
 291             if post_callback:
 292                 post_callback(migration_number, migration_func)