-import sys
-from mediagoblin.db.sql.base import Session
-from mediagoblin.db.models import MediaEntry, Tag, MediaTag, Collection
-
-from mediagoblin.tools.common import simple_printer
-
-
-class MigrationManager(object):
- """
- Migration handling tool.
-
- Takes information about a database, lets you update the database
- to the latest migrations, etc.
- """
-
- def __init__(self, name, models, migration_registry, session,
- printer=simple_printer):
- """
- Args:
- - name: identifier of this section of the database
- - session: session we're going to migrate
- - migration_registry: where we should find all migrations to
- run
- """
- self.name = unicode(name)
- self.models = models
- self.session = session
- self.migration_registry = migration_registry
- self._sorted_migrations = None
- self.printer = printer
-
- # For convenience
- from mediagoblin.db.models import MigrationData
-
- self.migration_model = MigrationData
- self.migration_table = MigrationData.__table__
-
- @property
- def sorted_migrations(self):
- """
- Sort migrations if necessary and store in self._sorted_migrations
- """
- if not self._sorted_migrations:
- self._sorted_migrations = sorted(
- self.migration_registry.items(),
- # sort on the key... the migration number
- key=lambda migration_tuple: migration_tuple[0])
-
- return self._sorted_migrations
-
- @property
- def migration_data(self):
- """
- Get the migration row associated with this object, if any.
- """
- return self.session.query(
- self.migration_model).filter_by(name=self.name).first()
-
- @property
- def latest_migration(self):
- """
- Return a migration number for the latest migration, or 0 if
- there are no migrations.
- """
- if self.sorted_migrations:
- return self.sorted_migrations[-1][0]
- else:
- # If no migrations have been set, we start at 0.
- return 0
-
- @property
- def database_current_migration(self):
- """
- Return the current migration in the database.
- """
- # If the table doesn't even exist, return None.
- if not self.migration_table.exists(self.session.bind):
- return None
-
- # Also return None if self.migration_data is None.
- if self.migration_data is None:
- return None
-
- return self.migration_data.version
-
- def set_current_migration(self, migration_number=None):
- """
- Set the migration in the database to migration_number
- (or, the latest available)
- """
- self.migration_data.version = migration_number or self.latest_migration
- self.session.commit()
-
- def migrations_to_run(self):
- """
- Get a list of migrations to run still, if any.
-
- Note that this will fail if there's no migration record for
- this class!
- """
- assert self.database_current_migration is not None
-
- db_current_migration = self.database_current_migration
-
- return [
- (migration_number, migration_func)
- for migration_number, migration_func in self.sorted_migrations
- if migration_number > db_current_migration]
-
-
- def init_tables(self):
- """
- Create all tables relative to this package
- """
- # sanity check before we proceed, none of these should be created
- for model in self.models:
- # Maybe in the future just print out a "Yikes!" or something?
- assert not model.__table__.exists(self.session.bind)
-
- self.migration_model.metadata.create_all(
- self.session.bind,
- tables=[model.__table__ for model in self.models])
-
- def create_new_migration_record(self):
- """
- Create a new migration record for this migration set
- """
- migration_record = self.migration_model(
- name=self.name,
- version=self.latest_migration)
- self.session.add(migration_record)
- self.session.commit()
-
- def dry_run(self):
- """
- Print out a dry run of what we would have upgraded.
- """
- if self.database_current_migration is None:
- self.printer(
- u'~> Woulda initialized: %s\n' % self.name_for_printing())
- return u'inited'