Modifications to SqliteStorage to allow for custom filepaths

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

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):