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
  20 class TableAlreadyExists(Exception):
  21     pass
  22
  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
  32     def __init__(self, name, models, migration_registry, session,
  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
  43         self.session = session
  44         self.migration_registry = migration_registry
  45         self._sorted_migrations = None
  46         self.printer = printer
  47
  48         # For convenience
  49         from mediagoblin.db.models import MigrationData
  50
  51         self.migration_model = MigrationData
  52         self.migration_table = MigrationData.__table__
  53
  54     @property
  55     def sorted_migrations(self):
  56         """
  57         Sort migrations if necessary and store in self._sorted_migrations
  58         """
  59         if not self._sorted_migrations:
  60             self._sorted_migrations = sorted(
  61                 self.migration_registry.items(),
  62                 # sort on the key... the migration number
  63                 key=lambda migration_tuple: migration_tuple[0])
  64
  65         return self._sorted_migrations
  66
  67     @property
  68     def migration_data(self):
  69         """
  70         Get the migration row associated with this object, if any.
  71         """
  72         return self.session.query(
  73             self.migration_model).filter_by(name=self.name).first()
  74
  75     @property
  76     def latest_migration(self):
  77         """
  78         Return a migration number for the latest migration, or 0 if
  79         there are no migrations.
  80         """
  81         if self.sorted_migrations:
  82             return self.sorted_migrations[-1][0]
  83         else:
  84             # If no migrations have been set, we start at 0.
  85             return 0
  86
  87     @property
  88     def database_current_migration(self):
  89         """
  90         Return the current migration in the database.
  91         """
  92         # If the table doesn't even exist, return None.
  93         if not self.migration_table.exists(self.session.bind):
  94             return None
  95
  96         # Also return None if self.migration_data is None.
  97         if self.migration_data is None:
  98             return None
  99
 100         return self.migration_data.version
 101
 102     def set_current_migration(self, migration_number=None):
 103         """
 104         Set the migration in the database to migration_number
 105         (or, the latest available)
 106         """
 107         self.migration_data.version = migration_number or self.latest_migration
 108         self.session.commit()
 109
 110     def migrations_to_run(self):
 111         """
 112         Get a list of migrations to run still, if any.
 113
 114         Note that this will fail if there's no migration record for
 115         this class!
 116         """
 117         assert self.database_current_migration is not None
 118
 119         db_current_migration = self.database_current_migration
 120
 121         return [
 122             (migration_number, migration_func)
 123             for migration_number, migration_func in self.sorted_migrations
 124             if migration_number > db_current_migration]
 125
 126
 127     def init_tables(self):
 128         """
 129         Create all tables relative to this package
 130         """
 131         # sanity check before we proceed, none of these should be created
 132         for model in self.models:
 133             # Maybe in the future just print out a "Yikes!" or something?
 134             if model.__table__.exists(self.session.bind):
 135                 raise TableAlreadyExists(
 136                     u"Intended to create table '%s' but it already exists" %
 137                     model.__table__.name)
 138
 139         self.migration_model.metadata.create_all(
 140             self.session.bind,
 141             tables=[model.__table__ for model in self.models])
 142
 143     def populate_table_foundations(self):
 144         """
 145         Create the table foundations (default rows) as layed out in FOUNDATIONS
 146             in mediagoblin.db.models
 147         """
 148         from mediagoblin.db.models import FOUNDATIONS as MAIN_FOUNDATIONS
 149         for Model, rows in MAIN_FOUNDATIONS.items():
 150             print u'\n--> Laying foundations for %s table' % Model.__name__
 151             for parameters in rows:
 152                 row = Model(**parameters)
 153                 row.save()
 154
 155     def create_new_migration_record(self):
 156         """
 157         Create a new migration record for this migration set
 158         """
 159         migration_record = self.migration_model(
 160             name=self.name,
 161             version=self.latest_migration)
 162         self.session.add(migration_record)
 163         self.session.commit()
 164
 165     def dry_run(self):
 166         """
 167         Print out a dry run of what we would have upgraded.
 168         """
 169         if self.database_current_migration is None:
 170             self.printer(
 171                     u'~> Woulda initialized: %s\n' % self.name_for_printing())
 172             return u'inited'
 173
 174         migrations_to_run = self.migrations_to_run()
 175         if migrations_to_run:
 176             self.printer(
 177                 u'~> Woulda updated %s:\n' % self.name_for_printing())
 178
 179             for migration_number, migration_func in migrations_to_run():
 180                 self.printer(
 181                     u'   + Would update %s, "%s"\n' % (
 182                         migration_number, migration_func.func_name))
 183
 184             return u'migrated'
 185
 186     def name_for_printing(self):
 187         if self.name == u'__main__':
 188             return u"main mediagoblin tables"
 189         else:
 190             return u'plugin "%s"' % self.name
 191
 192     def init_or_migrate(self):
 193         """
 194         Initialize the database or migrate if appropriate.
 195
 196         Returns information about whether or not we initialized
 197         ('inited'), migrated ('migrated'), or did nothing (None)
 198         """
 199         assure_migrations_table_setup(self.session)
 200
 201         # Find out what migration number, if any, this database data is at,
 202         # and what the latest is.
 203         migration_number = self.database_current_migration
 204
 205         # Is this our first time?  Is there even a table entry for
 206         # this identifier?
 207         # If so:
 208         #  - create all tables
 209         #  - create record in migrations registry
 210         #  - print / inform the user
 211         #  - return 'inited'
 212         if migration_number is None:
 213             self.printer(u"-> Initializing %s... " % self.name_for_printing())
 214
 215             self.init_tables()
 216             # auto-set at latest migration number
 217             self.create_new_migration_record()
 218             if self.name==u'__main__':
 219                 self.populate_table_foundations()
 220
 221             self.printer(u"done.\n")
 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)