diff --git a/src/wormhole/server/database.py b/src/wormhole/server/database.py index 253b536..7cd09f9 100644 --- a/src/wormhole/server/database.py +++ b/src/wormhole/server/database.py @@ -11,6 +11,12 @@ def get_schema(version): "db-schemas/v%d.sql" % version) return schema_bytes.decode("utf-8") +def get_upgrader(new_version): + schema_bytes = resource_string("wormhole.server", + "db-schemas/upgrade-to-v%d.sql" % new_version) + return schema_bytes.decode("utf-8") + +TARGET_VERSION = 2 def dict_factory(cursor, row): d = {} @@ -18,23 +24,23 @@ def dict_factory(cursor, row): d[col[0]] = row[idx] return d -def get_db(dbfile, stderr=sys.stderr): +def get_db(dbfile, target_version=TARGET_VERSION, stderr=sys.stderr): """Open or create the given db file. The parent directory must exist. Returns the db connection object, or raises DBError. """ - must_create = not os.path.exists(dbfile) + must_create = (dbfile == ":memory:") or not os.path.exists(dbfile) try: db = sqlite3.connect(dbfile) except (EnvironmentError, sqlite3.OperationalError) as e: raise DBError("Unable to create/open db file %s: %s" % (dbfile, e)) db.row_factory = dict_factory - VERSION = 2 if must_create: - schema = get_schema(VERSION) + schema = get_schema(target_version) db.executescript(schema) - db.execute("INSERT INTO version (version) VALUES (?)", (VERSION,)) + db.execute("INSERT INTO version (version) VALUES (?)", + (target_version,)) db.commit() try: @@ -44,7 +50,17 @@ def get_db(dbfile, stderr=sys.stderr): # Perhaps it was created with an old version, or it might be junk. raise DBError("db file is unusable: %s" % e) - if version != VERSION: + while version < target_version: + try: + upgrader = get_upgrader(version+1) + except ValueError: # ResourceError?? + raise DBError("Unable to upgrade %s to version %s, left at %s" + % (dbfile, version+1, version)) + db.executescript(upgrader) + db.commit() + version = version+1 + + if version != target_version: raise DBError("Unable to handle db version %s" % version) return db