#!/usr/bin/python -u
import os, sys, re, tempfile
from optparse import OptionParser
import common
from autotest_lib.client.common_lib import utils
from autotest_lib.database import database_connection
MIGRATE_TABLE = 'migrate_info'
_AUTODIR = os.path.join(os.path.dirname(__file__), '..')
_MIGRATIONS_DIRS = {
'AUTOTEST_WEB': os.path.join(_AUTODIR, 'frontend', 'migrations'),
'TKO': os.path.join(_AUTODIR, 'tko', 'migrations'),
'AUTOTEST_SERVER_DB': os.path.join(_AUTODIR, 'database',
'server_db_migrations'),
}
_DEFAULT_MIGRATIONS_DIR = 'migrations' # use CWD
class Migration(object):
"""Represents a database migration."""
_UP_ATTRIBUTES = ('migrate_up', 'UP_SQL')
_DOWN_ATTRIBUTES = ('migrate_down', 'DOWN_SQL')
def __init__(self, name, version, module):
self.name = name
self.version = version
self.module = module
self._check_attributes(self._UP_ATTRIBUTES)
self._check_attributes(self._DOWN_ATTRIBUTES)
@classmethod
def from_file(cls, filename):
"""Instantiates a Migration from a file.
@param filename: Name of a migration file.
@return An instantiated Migration object.
"""
version = int(filename[:3])
name = filename[:-3]
module = __import__(name, globals(), locals(), [])
return cls(name, version, module)
def _check_attributes(self, attributes):
method_name, sql_name = attributes
assert (hasattr(self.module, method_name) or
hasattr(self.module, sql_name))
def _execute_migration(self, attributes, manager):
method_name, sql_name = attributes
method = getattr(self.module, method_name, None)
if method:
assert callable(method)
method(manager)
else:
sql = getattr(self.module, sql_name)
assert isinstance(sql, basestring)
manager.execute_script(sql)
def migrate_up(self, manager):
"""Performs an up migration (to a newer version).
@param manager: A MigrationManager object.
"""
self._execute_migration(self._UP_ATTRIBUTES, manager)
def migrate_down(self, manager):
"""Performs a down migration (to an older version).
@param manager: A MigrationManager object.
"""
self._execute_migration(self._DOWN_ATTRIBUTES, manager)
class MigrationManager(object):
"""Managest database migrations."""
connection = None
cursor = None
migrations_dir = None
def __init__(self, database_connection, migrations_dir=None, force=False):
self._database = database_connection
self.force = force
# A boolean, this will only be set to True if this migration should be
# simulated rather than actually taken. For use with migrations that
# may make destructive queries
self.simulate = False
self._set_migrations_dir(migrations_dir)
def _set_migrations_dir(self, migrations_dir=None):
config_section = self._config_section()
if migrations_dir is None:
migrations_dir = os.path.abspath(
_MIGRATIONS_DIRS.get(config_section, _DEFAULT_MIGRATIONS_DIR))
self.migrations_dir = migrations_dir
sys.path.append(migrations_dir)
assert os.path.exists(migrations_dir), migrations_dir + " doesn't exist"
def _config_section(self):
return self._database.global_config_section
def get_db_name(self):
"""Gets the database name."""
return self._database.get_database_info()['db_name']
def execute(self, query, *parameters):
"""Executes a database query.
@param query: The query to execute.
@param parameters: Associated parameters for the query.
@return The result of the query.
"""
return self._database.execute(query, parameters)
def execute_script(self, script):
"""Executes a set of database queries.
@param script: A string of semicolon-separated queries.
"""
sql_statements = [statement.strip()
for statement in script.split(';')
if statement.strip()]
for statement in sql_statements:
self.execute(statement)
def check_migrate_table_exists(self):
"""Checks whether the migration table exists."""
try:
self.execute("SELECT * FROM %s" % MIGRATE_TABLE)
return True
except self._database.DatabaseError, exc:
# we can't check for more specifics due to differences between DB
# backends (we can't even check for a subclass of DatabaseError)
return False
def create_migrate_table(self):
"""Creates the migration table."""
if not self.check_migrate_table_exists():
self.execute("CREATE TABLE %s (`version` integer)" %
MIGRATE_TABLE)
else:
self.execute("DELETE FROM %s" % MIGRATE_TABLE)
self.execute("INSERT INTO %s VALUES (0)" % MIGRATE_TABLE)
assert self._database.rowcount == 1
def set_db_version(self, version):
"""Sets the database version.
@param version: The version to which to set the database.
"""
assert isinstance(version, int)
self.execute("UPDATE %s SET version=%%s" % MIGRATE_TABLE,
version)
assert self._database.rowcount == 1
def get_db_version(self):
"""Gets the database version.
@return The database version.
"""
if not self.check_migrate_table_exists():
return 0
rows = self.execute("SELECT * FROM %s" % MIGRATE_TABLE)
if len(rows) == 0:
return 0
assert len(rows) == 1 and len(rows[0]) == 1
return rows[0][0]
def get_migrations(self, minimum_version=None, maximum_version=None):
"""Gets the list of migrations to perform.
@param minimum_version: The minimum database version.
@param maximum_version: The maximum database version.
@return A list of Migration objects.
"""
migrate_files = [filename for filename
in os.listdir(self.migrations_dir)
if re.match(r'^\d\d\d_.*\.py$', filename)]
migrate_files.sort()
migrations = [Migration.from_file(filename)
for filename in migrate_files]
if minimum_version is not None:
migrations = [migration for migration in migrations
if migration.version >= minimum_version]
if maximum_version is not None:
migrations = [migration for migration in migrations
if migration.version <= maximum_version]
return migrations
def do_migration(self, migration, migrate_up=True):
"""Performs a migration.
@param migration: The Migration to perform.
@param migrate_up: Whether to migrate up (if not, then migrates down).
"""
print 'Applying migration %s' % migration.name, # no newline
if migrate_up:
print 'up'
assert self.get_db_version() == migration.version - 1
migration.migrate_up(self)
new_version = migration.version
else:
print 'down'
assert self.get_db_version() == migration.version
migration.migrate_down(self)
new_version = migration.version - 1
self.set_db_version(new_version)
def migrate_to_version(self, version):
"""Performs a migration to a specified version.
@param version: The version to which to migrate the database.
"""
current_version = self.get_db_version()
if current_version == 0 and self._config_section() == 'AUTOTEST_WEB':
self._migrate_from_base()
current_version = self.get_db_version()
if current_version < version:
lower, upper = current_version, version
migrate_up = True
else:
lower, upper = version, current_version
migrate_up = False
migrations = self.get_migrations(lower + 1, upper)
if not migrate_up:
migrations.reverse()
for migration in migrations:
self.do_migration(migration, migrate_up)
assert self.get_db_version() == version
print 'At version', version
def _migrate_from_base(self):
"""Initialize the AFE database.
"""
self.confirm_initialization()
migration_script = utils.read_file(
os.path.join(os.path.dirname(__file__), 'schema_051.sql'))
migration_script = migration_script % (
dict(username=self._database.get_database_info()['username']))
self.execute_script(migration_script)
self.create_migrate_table()
self.set_db_version(51)
def confirm_initialization(self):
"""Confirms with the user that we should initialize the database.
@raises Exception, if the user chooses to abort the migration.
"""
if not self.force:
response = raw_input(
'Your %s database does not appear to be initialized. Do you '
'want to recreate it (this will result in loss of any existing '
'data) (yes/No)? ' % self.get_db_name())
if response != 'yes':
raise Exception('User has chosen to abort migration')
def get_latest_version(self):
"""Gets the latest database version."""
migrations = self.get_migrations()
return migrations[-1].version
def migrate_to_latest(self):
"""Migrates the database to the latest version."""
latest_version = self.get_latest_version()
self.migrate_to_version(latest_version)
def initialize_test_db(self):
"""Initializes a test database."""
db_name = self.get_db_name()
test_db_name = 'test_' + db_name
# first, connect to no DB so we can create a test DB
self._database.connect(db_name='')
print 'Creating test DB', test_db_name
self.execute('CREATE DATABASE ' + test_db_name)
self._database.disconnect()
# now connect to the test DB
self._database.connect(db_name=test_db_name)
def remove_test_db(self):
"""Removes a test database."""
print 'Removing test DB'
self.execute('DROP DATABASE ' + self.get_db_name())
# reset connection back to real DB
self._database.disconnect()
self._database.connect()
def get_mysql_args(self):
"""Returns the mysql arguments as a string."""
return ('-u %(username)s -p%(password)s -h %(host)s %(db_name)s' %
self._database.get_database_info())
def migrate_to_version_or_latest(self, version):
"""Migrates to either a specified version, or the latest version.
@param version: The version to which to migrate the database,
or None in order to migrate to the latest version.
"""
if version is None:
self.migrate_to_latest()
else:
self.migrate_to_version(version)
def do_sync_db(self, version=None):
"""Migrates the database.
@param version: The version to which to migrate the database.
"""
print 'Migration starting for database', self.get_db_name()
self.migrate_to_version_or_latest(version)
print 'Migration complete'
def test_sync_db(self, version=None):
"""Create a fresh database and run all migrations on it.
@param version: The version to which to migrate the database.
"""
self.initialize_test_db()
try:
print 'Starting migration test on DB', self.get_db_name()
self.migrate_to_version_or_latest(version)
# show schema to the user
os.system('mysqldump %s --no-data=true '
'--add-drop-table=false' %
self.get_mysql_args())
finally:
self.remove_test_db()
print 'Test finished successfully'
def simulate_sync_db(self, version=None):
"""Creates a fresh DB, copies existing DB to it, then synchronizes it.
@param version: The version to which to migrate the database.
"""
db_version = self.get_db_version()
# don't do anything if we're already at the latest version
if db_version == self.get_latest_version():
print 'Skipping simulation, already at latest version'
return
# get existing data
self.initialize_and_fill_test_db()
try:
print 'Starting migration test on DB', self.get_db_name()
self.migrate_to_version_or_latest(version)
finally:
self.remove_test_db()
print 'Test finished successfully'
def initialize_and_fill_test_db(self):
"""Initializes and fills up a test database."""
print 'Dumping existing data'
dump_fd, dump_file = tempfile.mkstemp('.migrate_dump')
os.system('mysqldump %s >%s' %
(self.get_mysql_args(), dump_file))
# fill in test DB
self.initialize_test_db()
print 'Filling in test DB'
os.system('mysql %s <%s' % (self.get_mysql_args(), dump_file))
os.close(dump_fd)
os.remove(dump_file)
USAGE = """\
%s [options] sync|test|simulate|safesync [version]
Options:
-d --database Which database to act on
-f --force Don't ask for confirmation
--debug Print all DB queries"""\
% sys.argv[0]
def main():
"""Main function for the migration script."""
parser = OptionParser()
parser.add_option("-d", "--database",
help="which database to act on",
dest="database",
default="AUTOTEST_WEB")
parser.add_option("-f", "--force", help="don't ask for confirmation",
action="store_true")
parser.add_option('--debug', help='print all DB queries',
action='store_true')
(options, args) = parser.parse_args()
manager = get_migration_manager(db_name=options.database,
debug=options.debug, force=options.force)
if len(args) > 0:
if len(args) > 1:
version = int(args[1])
else:
version = None
if args[0] == 'sync':
manager.do_sync_db(version)
elif args[0] == 'test':
manager.simulate=True
manager.test_sync_db(version)
elif args[0] == 'simulate':
manager.simulate=True
manager.simulate_sync_db(version)
elif args[0] == 'safesync':
print 'Simluating migration'
manager.simulate=True
manager.simulate_sync_db(version)
print 'Performing real migration'
manager.simulate=False
manager.do_sync_db(version)
else:
print USAGE
return
print USAGE
def get_migration_manager(db_name, debug, force):
"""Creates a MigrationManager object.
@param db_name: The database name.
@param debug: Whether to print debug messages.
@param force: Whether to force migration without asking for confirmation.
@return A created MigrationManager object.
"""
database = database_connection.DatabaseConnection(db_name)
database.debug = debug
database.reconnect_enabled = False
database.connect()
return MigrationManager(database, force=force)
if __name__ == '__main__':
main()