database/migrate.py

   1 #!/usr/bin/python -u
   2
   3 import os, sys, re, subprocess, tempfile
   4 from optparse import OptionParser
   5 import common
   6 import MySQLdb, MySQLdb.constants.ER
   7 from autotest_lib.client.common_lib import global_config, utils
   8 from autotest_lib.database import database_connection
   9
  10 MIGRATE_TABLE = 'migrate_info'
  11
  12 _AUTODIR = os.path.join(os.path.dirname(__file__), '..')
  13 _MIGRATIONS_DIRS = {
  14     'AUTOTEST_WEB' : os.path.join(_AUTODIR, 'frontend', 'migrations'),
  15     'TKO' : os.path.join(_AUTODIR, 'tko', 'migrations'),
  16 }
  17 _DEFAULT_MIGRATIONS_DIR = 'migrations' # use CWD
  18
  19 class Migration(object):
  20     _UP_ATTRIBUTES = ('migrate_up', 'UP_SQL')
  21     _DOWN_ATTRIBUTES = ('migrate_down', 'DOWN_SQL')
  22
  23     def __init__(self, name, version, module):
  24         self.name = name
  25         self.version = version
  26         self.module = module
  27         self._check_attributes(self._UP_ATTRIBUTES)
  28         self._check_attributes(self._DOWN_ATTRIBUTES)
  29
  30
  31     @classmethod
  32     def from_file(cls, filename):
  33         version = int(filename[:3])
  34         name = filename[:-3]
  35         module = __import__(name, globals(), locals(), [])
  36         return cls(name, version, module)
  37
  38
  39     def _check_attributes(self, attributes):
  40         method_name, sql_name = attributes
  41         assert (hasattr(self.module, method_name) or
  42                 hasattr(self.module, sql_name))
  43
  44
  45     def _execute_migration(self, attributes, manager):
  46         method_name, sql_name = attributes
  47         method = getattr(self.module, method_name, None)
  48         if method:
  49             assert callable(method)
  50             method(manager)
  51         else:
  52             sql = getattr(self.module, sql_name)
  53             assert isinstance(sql, basestring)
  54             manager.execute_script(sql)
  55
  56
  57     def migrate_up(self, manager):
  58         self._execute_migration(self._UP_ATTRIBUTES, manager)
  59
  60
  61     def migrate_down(self, manager):
  62         self._execute_migration(self._DOWN_ATTRIBUTES, manager)
  63
  64
  65 class MigrationManager(object):
  66     connection = None
  67     cursor = None
  68     migrations_dir = None
  69
  70     def __init__(self, database_connection, migrations_dir=None, force=False):
  71         self._database = database_connection
  72         self.force = force
  73         # A boolean, this will only be set to True if this migration should be
  74         # simulated rather than actually taken. For use with migrations that
  75         # may make destructive queries
  76         self.simulate = False
  77         self._set_migrations_dir(migrations_dir)
  78
  79
  80     def _set_migrations_dir(self, migrations_dir=None):
  81         config_section = self._config_section()
  82         if migrations_dir is None:
  83             migrations_dir = os.path.abspath(
  84                 _MIGRATIONS_DIRS.get(config_section, _DEFAULT_MIGRATIONS_DIR))
  85         self.migrations_dir = migrations_dir
  86         sys.path.append(migrations_dir)
  87         assert os.path.exists(migrations_dir), migrations_dir + " doesn't exist"
  88
  89
  90     def _config_section(self):
  91         return self._database.global_config_section
  92
  93
  94     def get_db_name(self):
  95         return self._database.get_database_info()['db_name']
  96
  97
  98     def execute(self, query, *parameters):
  99         return self._database.execute(query, parameters)
 100
 101
 102     def execute_script(self, script):
 103         sql_statements = [statement.strip()
 104                           for statement in script.split(';')
 105                           if statement.strip()]
 106         for statement in sql_statements:
 107             self.execute(statement)
 108
 109
 110     def check_migrate_table_exists(self):
 111         try:
 112             self.execute("SELECT * FROM %s" % MIGRATE_TABLE)
 113             return True
 114         except self._database.DatabaseError, exc:
 115             # we can't check for more specifics due to differences between DB
 116             # backends (we can't even check for a subclass of DatabaseError)
 117             return False
 118
 119
 120     def create_migrate_table(self):
 121         if not self.check_migrate_table_exists():
 122             self.execute("CREATE TABLE %s (`version` integer)" %
 123                          MIGRATE_TABLE)
 124         else:
 125             self.execute("DELETE FROM %s" % MIGRATE_TABLE)
 126         self.execute("INSERT INTO %s VALUES (0)" % MIGRATE_TABLE)
 127         assert self._database.rowcount == 1
 128
 129
 130     def set_db_version(self, version):
 131         assert isinstance(version, int)
 132         self.execute("UPDATE %s SET version=%%s" % MIGRATE_TABLE,
 133                      version)
 134         assert self._database.rowcount == 1
 135
 136
 137     def get_db_version(self):
 138         if not self.check_migrate_table_exists():
 139             return 0
 140         rows = self.execute("SELECT * FROM %s" % MIGRATE_TABLE)
 141         if len(rows) == 0:
 142             return 0
 143         assert len(rows) == 1 and len(rows[0]) == 1
 144         return rows[0][0]
 145
 146
 147     def get_migrations(self, minimum_version=None, maximum_version=None):
 148         migrate_files = [filename for filename
 149                          in os.listdir(self.migrations_dir)
 150                          if re.match(r'^\d\d\d_.*\.py$', filename)]
 151         migrate_files.sort()
 152         migrations = [Migration.from_file(filename)
 153                       for filename in migrate_files]
 154         if minimum_version is not None:
 155             migrations = [migration for migration in migrations
 156                           if migration.version >= minimum_version]
 157         if maximum_version is not None:
 158             migrations = [migration for migration in migrations
 159                           if migration.version <= maximum_version]
 160         return migrations
 161
 162
 163     def do_migration(self, migration, migrate_up=True):
 164         print 'Applying migration %s' % migration.name, # no newline
 165         if migrate_up:
 166             print 'up'
 167             assert self.get_db_version() == migration.version - 1
 168             migration.migrate_up(self)
 169             new_version = migration.version
 170         else:
 171             print 'down'
 172             assert self.get_db_version() == migration.version
 173             migration.migrate_down(self)
 174             new_version = migration.version - 1
 175         self.set_db_version(new_version)
 176
 177
 178     def migrate_to_version(self, version):
 179         current_version = self.get_db_version()
 180         if current_version == 0 and self._config_section() == 'AUTOTEST_WEB':
 181             self._migrate_from_base()
 182             current_version = self.get_db_version()
 183
 184         if current_version < version:
 185             lower, upper = current_version, version
 186             migrate_up = True
 187         else:
 188             lower, upper = version, current_version
 189             migrate_up = False
 190
 191         migrations = self.get_migrations(lower + 1, upper)
 192         if not migrate_up:
 193             migrations.reverse()
 194         for migration in migrations:
 195             self.do_migration(migration, migrate_up)
 196
 197         assert self.get_db_version() == version
 198         print 'At version', version
 199
 200
 201     def _migrate_from_base(self):
 202         self.confirm_initialization()
 203
 204         migration_script = utils.read_file(
 205                 os.path.join(os.path.dirname(__file__), 'schema_051.sql'))
 206         migration_script = migration_script % (
 207                 dict(username=self._database.get_database_info()['username']))
 208         self.execute_script(migration_script)
 209
 210         self.create_migrate_table()
 211         self.set_db_version(51)
 212
 213
 214     def confirm_initialization(self):
 215         if not self.force:
 216             response = raw_input(
 217                 'Your %s database does not appear to be initialized.  Do you '
 218                 'want to recreate it (this will result in loss of any existing '
 219                 'data) (yes/No)? ' % self.get_db_name())
 220             if response != 'yes':
 221                 raise Exception('User has chosen to abort migration')
 222
 223
 224     def get_latest_version(self):
 225         migrations = self.get_migrations()
 226         return migrations[-1].version
 227
 228
 229     def migrate_to_latest(self):
 230         latest_version = self.get_latest_version()
 231         self.migrate_to_version(latest_version)
 232
 233
 234     def initialize_test_db(self):
 235         db_name = self.get_db_name()
 236         test_db_name = 'test_' + db_name
 237         # first, connect to no DB so we can create a test DB
 238         self._database.connect(db_name='')
 239         print 'Creating test DB', test_db_name
 240         self.execute('CREATE DATABASE ' + test_db_name)
 241         self._database.disconnect()
 242         # now connect to the test DB
 243         self._database.connect(db_name=test_db_name)
 244
 245
 246     def remove_test_db(self):
 247         print 'Removing test DB'
 248         self.execute('DROP DATABASE ' + self.get_db_name())
 249         # reset connection back to real DB
 250         self._database.disconnect()
 251         self._database.connect()
 252
 253
 254     def get_mysql_args(self):
 255         return ('-u %(username)s -p%(password)s -h %(host)s %(db_name)s' %
 256                 self._database.get_database_info())
 257
 258
 259     def migrate_to_version_or_latest(self, version):
 260         if version is None:
 261             self.migrate_to_latest()
 262         else:
 263             self.migrate_to_version(version)
 264
 265
 266     def do_sync_db(self, version=None):
 267         print 'Migration starting for database', self.get_db_name()
 268         self.migrate_to_version_or_latest(version)
 269         print 'Migration complete'
 270
 271
 272     def test_sync_db(self, version=None):
 273         """\
 274         Create a fresh DB and run all migrations on it.
 275         """
 276         self.initialize_test_db()
 277         try:
 278             print 'Starting migration test on DB', self.get_db_name()
 279             self.migrate_to_version_or_latest(version)
 280             # show schema to the user
 281             os.system('mysqldump %s --no-data=true '
 282                       '--add-drop-table=false' %
 283                       self.get_mysql_args())
 284         finally:
 285             self.remove_test_db()
 286         print 'Test finished successfully'
 287
 288
 289     def simulate_sync_db(self, version=None):
 290         """\
 291         Create a fresh DB, copy the existing DB to it, and then
 292         try to synchronize it.
 293         """
 294         db_version = self.get_db_version()
 295         # don't do anything if we're already at the latest version
 296         if db_version == self.get_latest_version():
 297             print 'Skipping simulation, already at latest version'
 298             return
 299         # get existing data
 300         self.initialize_and_fill_test_db()
 301         try:
 302             print 'Starting migration test on DB', self.get_db_name()
 303             self.migrate_to_version_or_latest(version)
 304         finally:
 305             self.remove_test_db()
 306         print 'Test finished successfully'
 307
 308
 309     def initialize_and_fill_test_db(self):
 310         print 'Dumping existing data'
 311         dump_fd, dump_file = tempfile.mkstemp('.migrate_dump')
 312         os.system('mysqldump %s >%s' %
 313                   (self.get_mysql_args(), dump_file))
 314         # fill in test DB
 315         self.initialize_test_db()
 316         print 'Filling in test DB'
 317         os.system('mysql %s <%s' % (self.get_mysql_args(), dump_file))
 318         os.close(dump_fd)
 319         os.remove(dump_file)
 320
 321
 322 USAGE = """\
 323 %s [options] sync|test|simulate|safesync [version]
 324 Options:
 325     -d --database   Which database to act on
 326     -a --action     Which action to perform"""\
 327     % sys.argv[0]
 328
 329
 330 def main():
 331     parser = OptionParser()
 332     parser.add_option("-d", "--database",
 333                       help="which database to act on",
 334                       dest="database",
 335                       default="AUTOTEST_WEB")
 336     parser.add_option("-a", "--action", help="what action to perform",
 337                       dest="action")
 338     parser.add_option("-f", "--force", help="don't ask for confirmation",
 339                       action="store_true")
 340     parser.add_option('--debug', help='print all DB queries',
 341                       action='store_true')
 342     (options, args) = parser.parse_args()
 343     manager = get_migration_manager(db_name=options.database,
 344                                     debug=options.debug, force=options.force)
 345
 346     if len(args) > 0:
 347         if len(args) > 1:
 348             version = int(args[1])
 349         else:
 350             version = None
 351         if args[0] == 'sync':
 352             manager.do_sync_db(version)
 353         elif args[0] == 'test':
 354             manager.simulate=True
 355             manager.test_sync_db(version)
 356         elif args[0] == 'simulate':
 357             manager.simulate=True
 358             manager.simulate_sync_db(version)
 359         elif args[0] == 'safesync':
 360             print 'Simluating migration'
 361             manager.simulate=True
 362             manager.simulate_sync_db(version)
 363             print 'Performing real migration'
 364             manager.simulate=False
 365             manager.do_sync_db(version)
 366         else:
 367             print USAGE
 368         return
 369
 370     print USAGE
 371
 372
 373 def get_migration_manager(db_name, debug, force):
 374     database = database_connection.DatabaseConnection(db_name)
 375     database.debug = debug
 376     database.reconnect_enabled = False
 377     database.connect()
 378     return MigrationManager(database, force=force)
 379
 380
 381 if __name__ == '__main__':
 382     main()