Commit | Line | Data |
---|---|---|
a050e776 E |
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 | from mediagoblin.tools.common import simple_printer | |
c4466cb4 | 18 | from sqlalchemy import Table |
a050e776 | 19 | |
7e4a87dc CAW |
20 | class TableAlreadyExists(Exception): |
21 | pass | |
22 | ||
a050e776 E |
23 | |
24 | class MigrationManager(object): | |
25 | """ | |
26 | Migration handling tool. | |
27 | ||
28 | Takes information about a database, lets you update the database | |
29 | to the latest migrations, etc. | |
30 | """ | |
31 | ||
08cd10d8 | 32 | def __init__(self, name, models, foundations, migration_registry, session, |
a050e776 E |
33 | printer=simple_printer): |
34 | """ | |
35 | Args: | |
36 | - name: identifier of this section of the database | |
37 | - session: session we're going to migrate | |
38 | - migration_registry: where we should find all migrations to | |
39 | run | |
40 | """ | |
41 | self.name = unicode(name) | |
42 | self.models = models | |
08cd10d8 | 43 | self.foundations = foundations |
a050e776 E |
44 | self.session = session |
45 | self.migration_registry = migration_registry | |
46 | self._sorted_migrations = None | |
47 | self.printer = printer | |
48 | ||
49 | # For convenience | |
50 | from mediagoblin.db.models import MigrationData | |
51 | ||
52 | self.migration_model = MigrationData | |
53 | self.migration_table = MigrationData.__table__ | |
54 | ||
55 | @property | |
56 | def sorted_migrations(self): | |
57 | """ | |
58 | Sort migrations if necessary and store in self._sorted_migrations | |
59 | """ | |
60 | if not self._sorted_migrations: | |
61 | self._sorted_migrations = sorted( | |
62 | self.migration_registry.items(), | |
63 | # sort on the key... the migration number | |
64 | key=lambda migration_tuple: migration_tuple[0]) | |
65 | ||
66 | return self._sorted_migrations | |
67 | ||
68 | @property | |
69 | def migration_data(self): | |
70 | """ | |
71 | Get the migration row associated with this object, if any. | |
72 | """ | |
73 | return self.session.query( | |
74 | self.migration_model).filter_by(name=self.name).first() | |
75 | ||
76 | @property | |
77 | def latest_migration(self): | |
78 | """ | |
79 | Return a migration number for the latest migration, or 0 if | |
80 | there are no migrations. | |
81 | """ | |
82 | if self.sorted_migrations: | |
83 | return self.sorted_migrations[-1][0] | |
84 | else: | |
85 | # If no migrations have been set, we start at 0. | |
86 | return 0 | |
87 | ||
88 | @property | |
89 | def database_current_migration(self): | |
90 | """ | |
91 | Return the current migration in the database. | |
92 | """ | |
93 | # If the table doesn't even exist, return None. | |
94 | if not self.migration_table.exists(self.session.bind): | |
95 | return None | |
96 | ||
97 | # Also return None if self.migration_data is None. | |
98 | if self.migration_data is None: | |
99 | return None | |
100 | ||
101 | return self.migration_data.version | |
102 | ||
103 | def set_current_migration(self, migration_number=None): | |
104 | """ | |
105 | Set the migration in the database to migration_number | |
106 | (or, the latest available) | |
107 | """ | |
108 | self.migration_data.version = migration_number or self.latest_migration | |
109 | self.session.commit() | |
110 | ||
111 | def migrations_to_run(self): | |
112 | """ | |
113 | Get a list of migrations to run still, if any. | |
114 | ||
115 | Note that this will fail if there's no migration record for | |
116 | this class! | |
117 | """ | |
118 | assert self.database_current_migration is not None | |
119 | ||
120 | db_current_migration = self.database_current_migration | |
121 | ||
122 | return [ | |
123 | (migration_number, migration_func) | |
124 | for migration_number, migration_func in self.sorted_migrations | |
125 | if migration_number > db_current_migration] | |
126 | ||
127 | ||
128 | def init_tables(self): | |
129 | """ | |
130 | Create all tables relative to this package | |
131 | """ | |
132 | # sanity check before we proceed, none of these should be created | |
133 | for model in self.models: | |
134 | # Maybe in the future just print out a "Yikes!" or something? | |
7e4a87dc CAW |
135 | if model.__table__.exists(self.session.bind): |
136 | raise TableAlreadyExists( | |
137 | u"Intended to create table '%s' but it already exists" % | |
138 | model.__table__.name) | |
a050e776 E |
139 | |
140 | self.migration_model.metadata.create_all( | |
141 | self.session.bind, | |
142 | tables=[model.__table__ for model in self.models]) | |
143 | ||
f2b2008d | 144 | def populate_table_foundations(self): |
145 | """ | |
146 | Create the table foundations (default rows) as layed out in FOUNDATIONS | |
147 | in mediagoblin.db.models | |
148 | """ | |
08cd10d8 | 149 | for Model, rows in self.foundations.items(): |
e1561d04 | 150 | self.printer(u' + Laying foundations for %s table\n' % |
151 | (Model.__name__)) | |
f2b2008d | 152 | for parameters in rows: |
08cd10d8 | 153 | new_row = Model(**parameters) |
e1561d04 | 154 | self.session.add(new_row) |
f2b2008d | 155 | |
a050e776 E |
156 | def create_new_migration_record(self): |
157 | """ | |
158 | Create a new migration record for this migration set | |
159 | """ | |
160 | migration_record = self.migration_model( | |
161 | name=self.name, | |
162 | version=self.latest_migration) | |
163 | self.session.add(migration_record) | |
164 | self.session.commit() | |
165 | ||
166 | def dry_run(self): | |
167 | """ | |
168 | Print out a dry run of what we would have upgraded. | |
169 | """ | |
170 | if self.database_current_migration is None: | |
171 | self.printer( | |
172 | u'~> Woulda initialized: %s\n' % self.name_for_printing()) | |
173 | return u'inited' | |
174 | ||
175 | migrations_to_run = self.migrations_to_run() | |
176 | if migrations_to_run: | |
177 | self.printer( | |
178 | u'~> Woulda updated %s:\n' % self.name_for_printing()) | |
179 | ||
180 | for migration_number, migration_func in migrations_to_run(): | |
181 | self.printer( | |
182 | u' + Would update %s, "%s"\n' % ( | |
183 | migration_number, migration_func.func_name)) | |
184 | ||
185 | return u'migrated' | |
186 | ||
187 | def name_for_printing(self): | |
188 | if self.name == u'__main__': | |
189 | return u"main mediagoblin tables" | |
190 | else: | |
003ea474 | 191 | return u'plugin "%s"' % self.name |
a050e776 E |
192 | |
193 | def init_or_migrate(self): | |
194 | """ | |
195 | Initialize the database or migrate if appropriate. | |
196 | ||
197 | Returns information about whether or not we initialized | |
198 | ('inited'), migrated ('migrated'), or did nothing (None) | |
199 | """ | |
200 | assure_migrations_table_setup(self.session) | |
201 | ||
202 | # Find out what migration number, if any, this database data is at, | |
203 | # and what the latest is. | |
204 | migration_number = self.database_current_migration | |
205 | ||
206 | # Is this our first time? Is there even a table entry for | |
207 | # this identifier? | |
208 | # If so: | |
209 | # - create all tables | |
210 | # - create record in migrations registry | |
211 | # - print / inform the user | |
212 | # - return 'inited' | |
213 | if migration_number is None: | |
214 | self.printer(u"-> Initializing %s... " % self.name_for_printing()) | |
215 | ||
216 | self.init_tables() | |
217 | # auto-set at latest migration number | |
f2b2008d | 218 | self.create_new_migration_record() |
a050e776 | 219 | self.printer(u"done.\n") |
e1561d04 | 220 | self.populate_table_foundations() |
a050e776 E |
221 | self.set_current_migration() |
222 | return u'inited' | |
223 | ||
224 | # Run migrations, if appropriate. | |
225 | migrations_to_run = self.migrations_to_run() | |
226 | if migrations_to_run: | |
227 | self.printer( | |
228 | u'-> Updating %s:\n' % self.name_for_printing()) | |
229 | for migration_number, migration_func in migrations_to_run: | |
230 | self.printer( | |
231 | u' + Running migration %s, "%s"... ' % ( | |
232 | migration_number, migration_func.func_name)) | |
233 | migration_func(self.session) | |
234 | self.set_current_migration(migration_number) | |
235 | self.printer('done.\n') | |
236 | ||
237 | return u'migrated' | |
238 | ||
239 | # Otherwise return None. Well it would do this anyway, but | |
240 | # for clarity... ;) | |
241 | return None | |
242 | ||
243 | ||
244 | class RegisterMigration(object): | |
245 | """ | |
246 | Tool for registering migrations | |
247 | ||
248 | Call like: | |
249 | ||
250 | @RegisterMigration(33) | |
251 | def update_dwarves(database): | |
252 | [...] | |
253 | ||
254 | This will register your migration with the default migration | |
255 | registry. Alternately, to specify a very specific | |
256 | migration_registry, you can pass in that as the second argument. | |
257 | ||
258 | Note, the number of your migration should NEVER be 0 or less than | |
259 | 0. 0 is the default "no migrations" state! | |
260 | """ | |
261 | def __init__(self, migration_number, migration_registry): | |
262 | assert migration_number > 0, "Migration number must be > 0!" | |
263 | assert migration_number not in migration_registry, \ | |
264 | "Duplicate migration numbers detected! That's not allowed!" | |
265 | ||
266 | self.migration_number = migration_number | |
267 | self.migration_registry = migration_registry | |
268 | ||
269 | def __call__(self, migration): | |
270 | self.migration_registry[self.migration_number] = migration | |
271 | return migration | |
272 | ||
273 | ||
274 | def assure_migrations_table_setup(db): | |
275 | """ | |
276 | Make sure the migrations table is set up in the database. | |
277 | """ | |
278 | from mediagoblin.db.models import MigrationData | |
279 | ||
280 | if not MigrationData.__table__.exists(db.bind): | |
281 | MigrationData.metadata.create_all( | |
282 | db.bind, tables=[MigrationData.__table__]) | |
c4466cb4 E |
283 | |
284 | ||
285 | def inspect_table(metadata, table_name): | |
286 | """Simple helper to get a ref to an already existing table""" | |
287 | return Table(table_name, metadata, autoload=True, | |
288 | autoload_with=metadata.bind) |