Commit | Line | Data |
---|---|---|
70b44584 CAW |
1 | # GNU MediaGoblin -- federated, autonomous media hosting |
2 | # Copyright (C) 2011 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 | ||
def13c54 | 17 | |
705689b9 CAW |
18 | import sys |
19 | ||
20 | def _simple_printer(string): | |
21 | """ | |
22 | Prints a string, but without an auto \n at the end. | |
23 | """ | |
24 | sys.stdout.write(string) | |
25 | sys.stdout.flush() | |
26 | ||
27 | ||
70b44584 | 28 | class MigrationManager(object): |
def13c54 CAW |
29 | """ |
30 | Migration handling tool. | |
31 | ||
32 | Takes information about a database, lets you update the database | |
33 | to the latest migrations, etc. | |
34 | """ | |
35 | ||
705689b9 CAW |
36 | def __init__(self, name, models, migration_registry, database, |
37 | printer=_simple_printer): | |
def13c54 CAW |
38 | """ |
39 | Args: | |
40 | - name: identifier of this section of the database | |
41 | - database: database we're going to migrate | |
42 | - migration_registry: where we should find all migrations to | |
43 | run | |
44 | """ | |
45 | self.name = name | |
46 | self.models = models | |
47 | self.database = database | |
48 | self.migration_registry = migration_registry | |
49 | self._sorted_migrations = None | |
705689b9 | 50 | self.printer = printer |
def13c54 CAW |
51 | |
52 | # For convenience | |
53 | from mediagoblin.db.sql.models import MigrationData | |
54 | ||
55 | self.migration_model = MigrationData | |
56 | self.migration_table = MigrationData.__table__ | |
57 | ||
58 | @property | |
59 | def sorted_migrations(self): | |
60 | """ | |
61 | Sort migrations if necessary and store in self._sorted_migrations | |
62 | """ | |
63 | if not self._sorted_migrations: | |
64 | self._sorted_migrations = sorted( | |
65 | self.migration_registry.items(), | |
66 | # sort on the key... the migration number | |
67 | key=lambda migration_tuple: migration_tuple[0]) | |
68 | ||
69 | return self._sorted_migrations | |
70 | ||
3635ccdf CAW |
71 | @property |
72 | def migration_data(self): | |
73 | """ | |
74 | Get the migration row associated with this object, if any. | |
75 | """ | |
851df621 CAW |
76 | return self.database.query( |
77 | self.migration_model).filter_by(name=self.name).first() | |
3635ccdf | 78 | |
def13c54 CAW |
79 | def latest_migration(self): |
80 | """ | |
81 | Return a migration number for the latest migration, or 0 if | |
82 | there are no migrations. | |
83 | """ | |
84 | if self.sorted_migrations: | |
85 | return self.sorted_migrations[-1][0] | |
86 | else: | |
87 | # If no migrations have been set, we start at 0. | |
88 | return 0 | |
89 | ||
90 | def database_current_migration(self): | |
91 | """ | |
92 | Return the current migration in the database. | |
93 | """ | |
3635ccdf | 94 | return self.migration_data.version |
def13c54 CAW |
95 | |
96 | def set_current_migration(self, migration_number): | |
97 | """ | |
98 | Set the migration in the database to migration_number | |
99 | """ | |
3635ccdf CAW |
100 | self.migration_data = migration_number |
101 | self.database.commit() | |
def13c54 CAW |
102 | |
103 | def migrations_to_run(self): | |
104 | """ | |
105 | Get a list of migrations to run still, if any. | |
106 | ||
3635ccdf CAW |
107 | Note that this will fail if there's no migration record for |
108 | this class! | |
def13c54 | 109 | """ |
cbf29f2d | 110 | assert self.database_current_migration is not None |
3635ccdf CAW |
111 | |
112 | db_current_migration = self.database_current_migration() | |
113 | ||
114 | return [ | |
115 | (migration_number, migration_func) | |
116 | for migration_number, migration_func in self.sorted_migrations | |
117 | if migration_number > db_current_migration] | |
118 | ||
def13c54 | 119 | |
705689b9 | 120 | def init_tables(self): |
8bf3f63a CAW |
121 | """ |
122 | Create all tables relative to this package | |
123 | """ | |
124 | # sanity check before we proceed, none of these should be created | |
125 | for model in self.models: | |
b0ec21bf | 126 | # Maybe in the future just print out a "Yikes!" or something? |
8bf3f63a CAW |
127 | assert not model.__table__.exists(self.database) |
128 | ||
129 | self.migration_model.metadata.create_all( | |
130 | self.database, | |
131 | tables=[model.__table__ for model in self.models]) | |
705689b9 CAW |
132 | |
133 | def create_new_migration_record(self): | |
b0ec21bf CAW |
134 | """ |
135 | Create a new migration record for this migration set | |
136 | """ | |
23f4c6b2 | 137 | migration_record = self.migration_model( |
b0ec21bf CAW |
138 | name=self.name, |
139 | version=self.latest_migration()) | |
23f4c6b2 | 140 | self.database.add(migration_record) |
09dcc34c | 141 | self.database.commit() |
705689b9 CAW |
142 | |
143 | def dry_run(self): | |
144 | """ | |
145 | Print out a dry run of what we would have upgraded. | |
146 | """ | |
147 | if self.database_current_migration() is None: | |
148 | self.printer( | |
149 | u'~> Woulda initialized: %s\n' % self.name_for_printing()) | |
150 | return u'inited' | |
151 | ||
152 | migrations_to_run = self.migrations_to_run() | |
153 | if migrations_to_run: | |
154 | self.printer( | |
155 | u'~> Woulda updated %s:\n' % self.name_for_printing()) | |
156 | ||
157 | for migration_number, migration_func in migrations_to_run(): | |
158 | self.printer( | |
159 | u' + Would update %s, "%s"\n' % ( | |
160 | migration_number, migration_func.func_name)) | |
161 | ||
162 | return u'migrated' | |
163 | ||
164 | def name_for_printing(self): | |
165 | if self.name == u'__main__': | |
166 | return u"main mediagoblin tables" | |
167 | else: | |
168 | # TODO: Use the friendlier media manager "human readable" name | |
169 | return u'media type "%s"' % self.name | |
170 | ||
4c869057 | 171 | def init_or_migrate(self): |
705689b9 CAW |
172 | """ |
173 | Initialize the database or migrate if appropriate. | |
174 | ||
175 | Returns information about whether or not we initialized | |
176 | ('inited'), migrated ('migrated'), or did nothing (None) | |
177 | """ | |
8bf3f63a CAW |
178 | assure_migrations_table_setup(self.database) |
179 | ||
def13c54 CAW |
180 | # Find out what migration number, if any, this database data is at, |
181 | # and what the latest is. | |
705689b9 | 182 | migration_number = self.database_current_migration() |
def13c54 CAW |
183 | |
184 | # Is this our first time? Is there even a table entry for | |
185 | # this identifier? | |
def13c54 CAW |
186 | # If so: |
187 | # - create all tables | |
188 | # - create record in migrations registry | |
189 | # - print / inform the user | |
190 | # - return 'inited' | |
705689b9 CAW |
191 | if migration_number is None: |
192 | self.printer(u"-> Initializing %s... " % self.name_for_printing()) | |
def13c54 | 193 | |
705689b9 CAW |
194 | self.init_tables() |
195 | # auto-set at latest migration number | |
196 | self.create_new_migration_record() | |
197 | ||
198 | self.printer(u"done.\n") | |
199 | return u'inited' | |
def13c54 | 200 | |
705689b9 CAW |
201 | # Run migrations, if appropriate. |
202 | migrations_to_run = self.migrations_to_run() | |
203 | if migrations_to_run: | |
204 | self.printer( | |
205 | u'~> Updating %s:\n' % self.name_for_printing()) | |
206 | for migration_number, migration_func in migrations_to_run(): | |
207 | self.printer( | |
a315962f | 208 | u' + Running migration %s, "%s"... ' % ( |
705689b9 | 209 | migration_number, migration_func.func_name)) |
a315962f CAW |
210 | migration_func(self.database) |
211 | self.printer('done.') | |
705689b9 CAW |
212 | |
213 | return u'migrated' | |
70b44584 | 214 | |
a315962f CAW |
215 | # Otherwise return None. Well it would do this anyway, but |
216 | # for clarity... ;) | |
217 | return None | |
218 | ||
70b44584 CAW |
219 | |
220 | class RegisterMigration(object): | |
221 | """ | |
222 | Tool for registering migrations | |
223 | ||
224 | Call like: | |
225 | ||
226 | @RegisterMigration(33) | |
227 | def update_dwarves(database): | |
228 | [...] | |
229 | ||
230 | This will register your migration with the default migration | |
231 | registry. Alternately, to specify a very specific | |
232 | migration_registry, you can pass in that as the second argument. | |
233 | ||
234 | Note, the number of your migration should NEVER be 0 or less than | |
235 | 0. 0 is the default "no migrations" state! | |
236 | """ | |
237 | def __init__(self, migration_number, migration_registry): | |
238 | assert migration_number > 0, "Migration number must be > 0!" | |
239 | assert migration_number not in migration_registry, \ | |
240 | "Duplicate migration numbers detected! That's not allowed!" | |
241 | ||
242 | self.migration_number = migration_number | |
243 | self.migration_registry = migration_registry | |
244 | ||
245 | def __call__(self, migration): | |
246 | self.migration_registry[self.migration_number] = migration | |
247 | return migration | |
248 | ||
249 | ||
250 | def assure_migrations_table_setup(db): | |
251 | """ | |
252 | Make sure the migrations table is set up in the database. | |
253 | """ | |
254 | from mediagoblin.db.sql.models import MigrationData | |
255 | ||
256 | if not MigrationData.__table__.exists(db): | |
257 | MigrationData.metadata.create_all( | |
258 | db, tables=[MigrationData.__table__]) |