mediagoblin/db/util.py

   1 # GNU MediaGoblin -- federated, autonomous media hosting
   2 # Copyright (C) 2011 Free Software Foundation, Inc
   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.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): pass
 126
 127
 128 MIGRATIONS = {}
 129
 130
 131 class RegisterMigration(object):
 132     """
 133     Tool for registering migrations
 134
 135     Call like:
 136
 137     @RegisterMigration(33)
 138     def update_dwarves(database):
 139         [...]
 140
 141     This will register your migration with the default migration
 142     registry.  Alternately, to specify a very specific
 143     migration_registry, you can pass in that as the second argument.
 144
 145     Note, the number of your migration should NEVER be 0 or less than
 146     0.  0 is the default "no migrations" state!
 147     """
 148     def __init__(self, migration_number, migration_registry=MIGRATIONS):
 149         assert migration_number > 0, "Migration number must be > 0!"
 150         assert not migration_registry.has_key(migration_number), \
 151             "Duplicate migration numbers detected!  That's not allowed!"
 152
 153         self.migration_number = migration_number
 154         self.migration_registry = migration_registry
 155
 156     def __call__(self, migration):
 157         self.migration_registry[self.migration_number] = migration
 158         return migration
 159
 160
 161 class MigrationManager(object):
 162     """
 163     Migration handling tool.
 164
 165     Takes information about a database, lets you update the database
 166     to the latest migrations, etc.
 167     """
 168     def __init__(self, database, migration_registry=MIGRATIONS):
 169         """
 170         Args:
 171          - database: database we're going to migrate
 172          - migration_registry: where we should find all migrations to
 173            run
 174         """
 175         self.database = database
 176         self.migration_registry = migration_registry
 177         self._sorted_migrations = None
 178
 179     def _ensure_current_migration_record(self):
 180         """
 181         If there isn't a database[u'app_metadata'] mediagoblin entry
 182         with the 'current_migration', throw an error.
 183         """
 184         if self.database_current_migration() is None:
 185             raise MissingCurrentMigration(
 186                 "Tried to call function which requires "
 187                 "'current_migration' set in database")
 188
 189     @property
 190     def sorted_migrations(self):
 191         """
 192         Sort migrations if necessary and store in self._sorted_migrations
 193         """
 194         if not self._sorted_migrations:
 195             self._sorted_migrations = sorted(
 196                 self.migration_registry.items(),
 197                 # sort on the key... the migration number
 198                 key=lambda migration_tuple: migration_tuple[0])
 199
 200         return self._sorted_migrations
 201
 202     def latest_migration(self):
 203         """
 204         Return a migration number for the latest migration, or 0 if
 205         there are no migrations.
 206         """
 207         if self.sorted_migrations:
 208             return self.sorted_migrations[-1][0]
 209         else:
 210             # If no migrations have been set, we start at 0.
 211             return 0
 212
 213     def set_current_migration(self, migration_number):
 214         """
 215         Set the migration in the database to migration_number
 216         """
 217         # Add the mediagoblin migration if necessary
 218         self.database[u'app_metadata'].update(
 219             {u'_id': u'mediagoblin'},
 220             {u'$set': {u'current_migration': migration_number}},
 221             upsert=True)
 222
 223     def install_migration_version_if_missing(self):
 224         """
 225         Sets the migration to the latest version if no migration
 226         version at all is set.
 227         """
 228         mgoblin_metadata = self.database[u'app_metadata'].find_one(
 229             {u'_id': u'mediagoblin'})
 230         if not mgoblin_metadata:
 231             latest_migration = self.latest_migration()
 232             self.set_current_migration(latest_migration)
 233
 234     def database_current_migration(self):
 235         """
 236         Return the current migration in the database.
 237         """
 238         mgoblin_metadata = self.database[u'app_metadata'].find_one(
 239             {u'_id': u'mediagoblin'})
 240         if not mgoblin_metadata:
 241             return None
 242         else:
 243             return mgoblin_metadata[u'current_migration']
 244
 245     def database_at_latest_migration(self):
 246         """
 247         See if the database is at the latest migration.
 248         Returns a boolean.
 249         """
 250         current_migration = self.database_current_migration()
 251         return current_migration == self.latest_migration()
 252
 253     def migrations_to_run(self):
 254         """
 255         Get a list of migrations to run still, if any.
 256
 257         Note that calling this will set your migration version to the
 258         latest version if it isn't installed to anything yet!
 259         """
 260         self._ensure_current_migration_record()
 261
 262         db_current_migration = self.database_current_migration()
 263
 264         return [
 265             (migration_number, migration_func)
 266             for migration_number, migration_func in self.sorted_migrations
 267             if migration_number > db_current_migration]
 268
 269     def migrate_new(self, pre_callback=None, post_callback=None):
 270         """
 271         Run all migrations.
 272
 273         Includes two optional args:
 274          - pre_callback: if called, this is a callback on something to
 275            run pre-migration.  Takes (migration_number, migration_func)
 276            as arguments
 277          - pre_callback: if called, this is a callback on something to
 278            run post-migration.  Takes (migration_number, migration_func)
 279            as arguments
 280         """
 281         # If we aren't set to any version number, presume we're at the
 282         # latest (which means we'll do nothing here...)
 283         self.install_migration_version_if_missing()
 284
 285         for migration_number, migration_func in self.migrations_to_run():
 286             if pre_callback:
 287                 pre_callback(migration_number, migration_func)
 288             migration_func(self.database)
 289             self.set_current_migration(migration_number)
 290             if post_callback:
 291                 post_callback(migration_number, migration_func)