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     def __init__(self, migration_number, migration_registry=MIGRATIONS):
 126         self.migration_number = migration_number
 127         self.migration_registry = migration_registry
 128
 129     def __call__(self, migration):
 130         self.migration_registry[self.migration_number] = migration
 131         return migration
 132
 133
 134 class MigrationManager(object):
 135     """
 136     Migration handling tool.
 137
 138     Takes information about a database, lets you update the database
 139     to the latest migrations, etc.
 140     """
 141     def __init__(self, database, migration_registry=MIGRATIONS):
 142         """
 143         Args:
 144          - database: database we're going to migrate
 145          - migration_registry: where we should find all migrations to
 146            run
 147         """
 148         self.database = database
 149         self.migration_registry = migration_registry
 150         self._sorted_migrations = None
 151
 152     @property
 153     def sorted_migrations(self):
 154         """
 155         Sort migrations if necessary and store in self._sorted_migrations
 156         """
 157         if not self._sorted_migrations:
 158             self._sorted_migrations = sorted(
 159                 self.migration_registry.items(),
 160                 # sort on the key... the migration number
 161                 key=lambda migration_tuple: migration_tuple[0])
 162
 163         return self._sorted_migrations
 164
 165     def latest_migration(self):
 166         """
 167         Return a migration number for the latest migration, or 0 if
 168         there are no migrations.
 169         """
 170         if self.sorted_migrations:
 171             return self.sorted_migrations[-1][0]
 172         else:
 173             # If no migrations have been set, we start at 0.
 174             return 0
 175
 176     def set_current_migration(self, migration_number=None):
 177         """
 178         Set the migration in the database to migration_number
 179         """
 180         # Add the mediagoblin migration if necessary
 181         self.database['app_metadata'].update(
 182             {'_id': 'mediagoblin'},
 183             {'$set': {'current_migration': migration_number}},
 184             upsert=True)
 185
 186     def database_current_migration(self, install_if_missing=True):
 187         """
 188         Return the current migration in the database.
 189         """
 190         mgoblin_metadata = self.database['app_metadata'].find_one(
 191             {'_id': 'mediagoblin'})
 192         if not mgoblin_metadata:
 193             if install_if_missing:
 194                 latest_migration = self.latest_migration()
 195                 self.set_current_migration(latest_migration)
 196                 return latest_migration
 197             else:
 198                 return None
 199         else:
 200             return mgoblin_metadata['current_migration']
 201
 202     def database_at_latest_migration(self):
 203         """
 204         See if the database is at the latest migration.
 205         Returns a boolean.
 206         """
 207         current_migration = self.database_current_migration()
 208         return current_migration == self.latest_migration()
 209
 210     def migrations_to_run(self):
 211         """
 212         Get a list of migrations to run still, if any.
 213         """
 214         db_current_migration = self.database_current_migration()
 215         return [
 216             (migration_number, migration_func)
 217             for migration_number, migration_func in self.sorted_migrations
 218             if migration_number > db_current_migration]
 219
 220     def iteratively_migrate(self):
 221         """
 222         Iteratively run all migrations.
 223
 224         Useful if you need to print some message about each migration
 225         after you run it.
 226
 227         Each time you loop over this, it'll return the migration
 228         number and migration function.
 229         """
 230         for migration_number, migration_func in self.migrations_to_run():
 231             migration_func(self.database)
 232             self.set_current_migration(migration_number)
 233             yield migration_number, migration_func
 234
 235     def run_outdated_migrations(self):
 236         """
 237         Install all migrations that need to be installed, quietly.
 238         """
 239         for migration_number, migration_func in self.iteratively_migrate():
 240             # No need to say anything... we're just migrating quietly.
 241             pass