import os import glob import imp import random from twisted.internet import reactor from twisted.enterprise import adbapi from twisted.internet import defer @defer.inlineCallbacks def get_schema_version(dbpool): schema_exists = yield dbpool.runQuery("SELECT name FROM sqlite_master WHERE type='table' AND name='schema_meta';") if not schema_exists: # If we get an empty list, we have no schema defer.returnValue(0) else: schema_version_result = yield dbpool.runQuery("SELECT version FROM schema_meta;") # There should only be one row in the meta table assert(len(schema_version_result) == 1) # Return the retrieved version version = schema_version_result[0][0] defer.returnValue(version) def add_schema_files(schemas): # Finds and imports all schema_*.py files into the list module_files = glob.glob(os.path.dirname(os.path.abspath(__file__)) + "/schema_*.py") for mod in module_files: module_name = os.path.basename(os.path.splitext(mod)[0]) newmod = imp.load_source('%s'%module_name, mod) schemas.append( (module_name, newmod) ) def copy_file(a, b): a_bytes = a.read(1024) while a_bytes: b.write(a_bytes) a_bytes = a.read(1024) def create_backup(filename): backup = filename + '.bak' while os.path.isfile(backup): backup = '%s.%d' % (backup, random.randint(0, 9999999999)) # Make sure backup file has secure permissions with os.fdopen(os.open(backup, os.O_CREAT, 0o0600), 'r') as f: pass # Copy the datafile with open(filename, 'r') as a: with open(backup, 'w') as b: copy_file(a, b) return backup @defer.inlineCallbacks def update_schema(dbpool, filename): # Update the database schema to the latest version schema_version = yield get_schema_version(dbpool) if schema_version == 0: verbose_update = False else: verbose_update = True schemas = [] add_schema_files(schemas) schemas = sorted(schemas, key=lambda tup: tup[0]) to_run = range(schema_version, len(schemas)) if len(to_run) > 0: # Back up data file if verbose_update: print 'Backing up data file' backup = create_backup(filename) if verbose_update: print 'Backed up to %s' % backup try: for i in to_run: # schemas[0] is v1, schemas[1] is v2, etc if verbose_update: print "Updating datafaile schema to version %d" % (i+1) yield schemas[i][1].update(dbpool) # Delete backup os.remove(backup) if verbose_update: print 'Update successful! Deleted backup' except Exception as e: # restore the backup print 'Update failed, restoring backup' with open(filename, 'w') as a: with open(backup, 'r') as b: copy_file(b, a) os.remove(backup) raise e @defer.inlineCallbacks def main(): dbpool = adbapi.ConnectionPool("sqlite3", "data.db", check_same_thread=False) yield update_schema(dbpool) reactor.stop() if __name__ == '__main__': reactor.callWhenRunning(main) reactor.run()