mediagoblin/db/migration_tools.py

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