diff --git a/aw_datastore/migration.py b/aw_datastore/migration.py index 468eb46..54f0c0e 100644 --- a/aw_datastore/migration.py +++ b/aw_datastore/migration.py @@ -8,6 +8,7 @@ logger = logging.getLogger(__name__) + def detect_db_files(data_dir: str, datastore_name: str = None, version=None) -> List[str]: db_files = [filename for filename in os.listdir(data_dir)] if datastore_name: @@ -16,17 +17,19 @@ def detect_db_files(data_dir: str, datastore_name: str = None, version=None) -> db_files = [filename for filename in db_files if filename.split(".")[1] == "v{}".format(version)] return db_files -def check_for_migration(datastore: AbstractStorage, datastore_name: str, version: int): + +def check_for_migration(datastore: AbstractStorage): data_dir = get_data_dir("aw-server") if datastore.sid == "sqlite": peewee_type = "peewee-sqlite" - peewee_name = peewee_type + "-testing" if datastore.testing else "" + peewee_name = peewee_type + ("-testing" if datastore.testing else "") # Migrate from peewee v2 peewee_db_v2 = detect_db_files(data_dir, peewee_name, 2) if len(peewee_db_v2) > 0: peewee_v2_to_sqlite_v1(datastore) + def peewee_v2_to_sqlite_v1(datastore): logger.info("Migrating database from peewee v2 to sqlite v1") from .storages import PeeweeStorage diff --git a/aw_datastore/storages/sqlite.py b/aw_datastore/storages/sqlite.py index ea53757..0f7b40e 100644 --- a/aw_datastore/storages/sqlite.py +++ b/aw_datastore/storages/sqlite.py @@ -57,13 +57,19 @@ class SqliteStorage(AbstractStorage): sid = "sqlite" - def __init__(self, testing): + def __init__(self, testing, filepath: str = None, enable_lazy_commit=True) -> None: self.testing = testing - data_dir = get_data_dir("aw-server") + self.enable_lazy_commit = enable_lazy_commit + + # Ignore the migration check if custom filepath is set + ignore_migration_check = filepath is not None ds_name = self.sid + ('-testing' if testing else '') - filename = ds_name + ".v{}".format(LATEST_VERSION) + '.db' - filepath = os.path.join(data_dir, filename) + if not filepath: + data_dir = get_data_dir("aw-server") + filename = ds_name + ".v{}".format(LATEST_VERSION) + '.db' + filepath = os.path.join(data_dir, filename) + new_db_file = not os.path.exists(filepath) self.conn = sqlite3.connect(filepath) logger.info("Using database file: {}".format(filepath)) @@ -77,10 +83,10 @@ def __init__(self, testing): self.conn.execute("PRAGMA journal_mode=WAL;") self.commit() - if new_db_file: + if new_db_file and not ignore_migration_check: logger.info("Created new SQlite db file") from aw_datastore import check_for_migration - check_for_migration(self, ds_name, LATEST_VERSION) + check_for_migration(self) self.last_commit = datetime.now() self.num_uncommited_statements = 0 @@ -102,10 +108,13 @@ def conditional_commit(self, num_statements): This is because sqlite is very slow with small inserts, this is a way to batch them together and lower CPU+disk usage """ - self.num_uncommited_statements += num_statements - if self.num_uncommited_statements > 50: - self.commit() - if (self.last_commit - datetime.now()) > timedelta(seconds=10): + if self.enable_lazy_commit: + self.num_uncommited_statements += num_statements + if self.num_uncommited_statements > 50: + self.commit() + if (self.last_commit - datetime.now()) > timedelta(seconds=10): + self.commit() + else: self.commit() def buckets(self):