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': ['index_identifier1', 'index_identifier2']}
  90
  91     Returns:
  92       A list of indexes removed in form ('collection', 'index_name')
  93     """
  94     indexes_removed = []
  95
  96     for collection_name, index_names in deprecated_indexes.iteritems():
  97         collection = database[collection_name]
  98         collection_indexes = collection.index_information().keys()
  99
 100         for index_name in index_names:
 101             if index_name in collection_indexes:
 102                 collection.drop_index(index_name)
 103
 104                 indexes_removed.append((collection_name, index_name))
 105
 106     return indexes_removed
 107
 108
 109 #################
 110 # Migration tools
 111 #################
 112
 113 # The default migration registry...
 114 #
 115 # Don't set this yourself!  RegisterMigration will automatically fill
 116 # this with stuff via decorating methods in migrations.py
 117
 118 MIGRATIONS = {}
 119
 120
 121 class RegisterMigration(object):
 122     """
 123     Tool for registering migrations
 124
 125     Call like:
 126
 127     @RegisterMigration(33)
 128     def update_dwarves(database):
 129         [...]
 130
 131     This will register your migration with the default migration
 132     registry.  Alternately, to specify a very specific
 133     migration_registry, you can pass in that as the second argument.
 134
 135     Note, the number of your migration should NEVER be 0 or less than
 136     0.  0 is the default "no migrations" state!
 137     """
 138     def __init__(self, migration_number, migration_registry=MIGRATIONS):
 139         assert migration_number > 0, "Migration number must be > 0!"
 140
 141         self.migration_number = migration_number
 142         self.migration_registry = migration_registry
 143
 144     def __call__(self, migration):
 145         self.migration_registry[self.migration_number] = migration
 146         return migration
 147
 148
 149 class MigrationManager(object):
 150     """
 151     Migration handling tool.
 152
 153     Takes information about a database, lets you update the database
 154     to the latest migrations, etc.
 155     """
 156     def __init__(self, database, migration_registry=MIGRATIONS):
 157         """
 158         Args:
 159          - database: database we're going to migrate
 160          - migration_registry: where we should find all migrations to
 161            run
 162         """
 163         self.database = database
 164         self.migration_registry = migration_registry
 165         self._sorted_migrations = None
 166
 167     @property
 168     def sorted_migrations(self):
 169         """
 170         Sort migrations if necessary and store in self._sorted_migrations
 171         """
 172         if not self._sorted_migrations:
 173             self._sorted_migrations = sorted(
 174                 self.migration_registry.items(),
 175                 # sort on the key... the migration number
 176                 key=lambda migration_tuple: migration_tuple[0])
 177
 178         return self._sorted_migrations
 179
 180     def latest_migration(self):
 181         """
 182         Return a migration number for the latest migration, or 0 if
 183         there are no migrations.
 184         """
 185         if self.sorted_migrations:
 186             return self.sorted_migrations[-1][0]
 187         else:
 188             # If no migrations have been set, we start at 0.
 189             return 0
 190
 191     def set_current_migration(self, migration_number):
 192         """
 193         Set the migration in the database to migration_number
 194         """
 195         # Add the mediagoblin migration if necessary
 196         self.database[u'app_metadata'].update(
 197             {u'_id': u'mediagoblin'},
 198             {u'$set': {u'current_migration': migration_number}},
 199             upsert=True)
 200
 201     def install_migration_version_if_missing(self):
 202         """
 203         Sets the migration to the latest version if no migration
 204         version at all is set.
 205         """
 206         mgoblin_metadata = self.database[u'app_metadata'].find_one(
 207             {u'_id': u'mediagoblin'})
 208         if not mgoblin_metadata:
 209             latest_migration = self.latest_migration()
 210             self.set_current_migration(latest_migration)
 211
 212     def database_current_migration(self):
 213         """
 214         Return the current migration in the database.
 215         """
 216         mgoblin_metadata = self.database[u'app_metadata'].find_one(
 217             {u'_id': u'mediagoblin'})
 218         if not mgoblin_metadata:
 219             return None
 220         else:
 221             return mgoblin_metadata[u'current_migration']
 222
 223     def database_at_latest_migration(self):
 224         """
 225         See if the database is at the latest migration.
 226         Returns a boolean.
 227         """
 228         current_migration = self.database_current_migration()
 229         return current_migration == self.latest_migration()
 230
 231     def migrations_to_run(self):
 232         """
 233         Get a list of migrations to run still, if any.
 234
 235         Note that calling this will set your migration version to the
 236         latest version if it isn't installed to anything yet!
 237         """
 238         # If we aren't set to any version number, presume we're at the
 239         # latest (which means we'll do nothing here...)
 240         self.install_migration_version_if_missing()
 241
 242         db_current_migration = self.database_current_migration()
 243
 244         return [
 245             (migration_number, migration_func)
 246             for migration_number, migration_func in self.sorted_migrations
 247             if migration_number > db_current_migration]
 248
 249     def migrate_new(self, pre_callback=None, post_callback=None):
 250         """
 251         Run all migrations.
 252
 253         Includes two optional args:
 254          - pre_callback: if called, this is a callback on something to
 255            run pre-migration.  Takes (migration_number, migration_func)
 256            as arguments
 257          - pre_callback: if called, this is a callback on something to
 258            run post-migration.  Takes (migration_number, migration_func)
 259            as arguments
 260         """
 261         for migration_number, migration_func in self.migrations_to_run():
 262             if pre_callback:
 263                 pre_callback(migration_number, migration_func)
 264             migration_func(self.database)
 265             self.set_current_migration(migration_number)
 266             if post_callback:
 267                 post_callback(migration_number, migration_func)