magic-wormhole/src/wormhole/server/database.py

80 lines
2.6 KiB
Python
Raw Normal View History

from __future__ import unicode_literals
2015-04-10 16:15:27 +00:00
import os, sys
import sqlite3
from pkg_resources import resource_string
class DBError(Exception):
pass
def get_schema(version):
schema_bytes = resource_string("wormhole.server",
"db-schemas/v%d.sql" % version)
2015-09-27 00:51:21 +00:00
return schema_bytes.decode("utf-8")
2015-04-10 16:15:27 +00:00
2016-06-21 00:29:34 +00:00
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 = 3
def dict_factory(cursor, row):
d = {}
for idx, col in enumerate(cursor.description):
d[col[0]] = row[idx]
return d
2016-06-21 00:29:34 +00:00
def get_db(dbfile, target_version=TARGET_VERSION, stderr=sys.stderr):
2015-04-10 16:15:27 +00:00
"""Open or create the given db file. The parent directory must exist.
Returns the db connection object, or raises DBError.
"""
2016-06-21 00:29:34 +00:00
must_create = (dbfile == ":memory:") or not os.path.exists(dbfile)
2015-04-10 16:15:27 +00:00
try:
db = sqlite3.connect(dbfile)
except (EnvironmentError, sqlite3.OperationalError) as e:
2015-04-10 16:15:27 +00:00
raise DBError("Unable to create/open db file %s: %s" % (dbfile, e))
db.row_factory = dict_factory
2016-06-23 23:26:36 +00:00
db.execute("PRAGMA foreign_keys = ON")
problems = db.execute("PRAGMA foreign_key_check").fetchall()
if problems:
raise DBError("failed foreign key check: %s" % (problems,))
2015-04-10 16:15:27 +00:00
if must_create:
2016-06-21 00:29:34 +00:00
schema = get_schema(target_version)
2015-04-10 16:15:27 +00:00
db.executescript(schema)
2016-06-21 00:29:34 +00:00
db.execute("INSERT INTO version (version) VALUES (?)",
(target_version,))
2015-04-10 16:15:27 +00:00
db.commit()
try:
version = db.execute("SELECT version FROM version").fetchone()["version"]
except sqlite3.DatabaseError as e:
2015-04-10 16:15:27 +00:00
# this indicates that the file is not a compatible database format.
# Perhaps it was created with an old version, or it might be junk.
raise DBError("db file is unusable: %s" % e)
2016-06-21 00:29:34 +00:00
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:
2015-04-10 16:15:27 +00:00
raise DBError("Unable to handle db version %s" % version)
return db
def dump_db(db):
# to let _iterdump work, we need to restore the original row factory
orig = db.row_factory
try:
db.row_factory = sqlite3.Row
return "".join(db.iterdump())
finally:
db.row_factory = orig