Commit | Line | Data |
---|---|---|
59bd06aa | 1 | # GNU MediaGoblin -- federated, autonomous media hosting |
cf29e8a8 | 2 | # Copyright (C) 2011, 2012 MediaGoblin contributors. See AUTHORS. |
59bd06aa E |
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) |