Add db_journal_mode trait (#61)

Author: kevin-bates

jupyter_server_fileid/manager.py

@@ -8,7 +8,7 @@
 from typing import Any, Callable, Dict, Optional
 
 from jupyter_core.paths import jupyter_data_dir
-from traitlets import TraitError, Unicode, validate
+from traitlets import TraitError, Unicode, default, validate
 from traitlets.config.configurable import LoggingConfigurable
 
 
@@ -75,6 +75,23 @@ def _validate_db_path(self, proposal):
             )
         return proposal["value"]
 
+    JOURNAL_MODES = ["DELETE", "TRUNCATE", "PERSIST", "MEMORY", "WAL", "OFF"]
+    db_journal_mode = Unicode(
+        help=(
+            f"The journal mode setting for the SQLite database.  Must be one of {JOURNAL_MODES}."
+        ),
+        config=True,
+    )
+
+    @validate("db_journal_mode")
+    def _validate_db_journal_mode(self, proposal):
+        candidate_value = proposal["value"]
+        if candidate_value is None or candidate_value.upper() not in self.JOURNAL_MODES:
+            raise TraitError(
+                f"db_journal_mode ('{candidate_value}') must be one of {self.JOURNAL_MODES}."
+            )
+        return candidate_value.upper()
+
     @staticmethod
     def _uuid() -> str:
         return str(uuid.uuid4())
@@ -229,6 +246,10 @@ def _validate_root_dir(self, proposal):
         normalized_content_root = self._normalize_separators(proposal["value"])
         return normalized_content_root
 
+    @default("db_journal_mode")
+    def _default_db_journal_mode(self):
+        return "DELETE"
+
     def __init__(self, *args, **kwargs):
         # pass args and kwargs to parent Configurable
         super().__init__(*args, **kwargs)
@@ -239,9 +260,11 @@ def __init__(self, *args, **kwargs):
         self.log.info(f"ArbitraryFileIdManager : Configured database path: {self.db_path}")
         self.con = sqlite3.connect(self.db_path)
         self.log.info("ArbitraryFileIdManager : Successfully connected to database file.")
-        self.log.info("ArbitraryFileIdManager : Creating File ID tables and indices.")
-        # do not allow reads to block writes. required when using multiple processes
-        self.con.execute("PRAGMA journal_mode = WAL")
+        self.log.info(
+            f"ArbitraryFileIdManager : Creating File ID tables and indices with "
+            f"journal_mode = {self.db_journal_mode}"
+        )
+        self.con.execute(f"PRAGMA journal_mode = {self.db_journal_mode}")
         self.con.execute(
             "CREATE TABLE IF NOT EXISTS Files("
             "id TEXT PRIMARY KEY NOT NULL, "
@@ -394,6 +417,10 @@ def _validate_root_dir(self, proposal):
             )
         return proposal["value"]
 
+    @default("db_journal_mode")
+    def _default_db_journal_mode(self):
+        return "WAL"
+
     def __init__(self, *args, **kwargs):
         # pass args and kwargs to parent Configurable
         super().__init__(*args, **kwargs)
@@ -405,9 +432,11 @@ def __init__(self, *args, **kwargs):
         self.log.info(f"LocalFileIdManager : Configured database path: {self.db_path}")
         self.con = sqlite3.connect(self.db_path)
         self.log.info("LocalFileIdManager : Successfully connected to database file.")
-        self.log.info("LocalFileIdManager : Creating File ID tables and indices.")
-        # do not allow reads to block writes. required when using multiple processes
-        self.con.execute("PRAGMA journal_mode = WAL")
+        self.log.info(
+            f"LocalFileIdManager : Creating File ID tables and indices with "
+            f"journal_mode = {self.db_journal_mode}"
+        )
+        self.con.execute(f"PRAGMA journal_mode = {self.db_journal_mode}")
         self.con.execute(
             "CREATE TABLE IF NOT EXISTS Files("
             "id TEXT PRIMARY KEY NOT NULL, "

tests/test_manager.py

@@ -565,3 +565,29 @@ def test_save(any_fid_manager, test_path, fs_helpers):
     any_fid_manager.save(test_path)
 
     assert any_fid_manager.get_id(test_path) == id
+
+
+@pytest.mark.parametrize(
+    "db_journal_mode", ["invalid", None, "DELETE", "TRUNCATE", "PERSIST", "MEMORY", "WAL", "OFF"]
+)
+def test_db_journal_mode(any_fid_manager_class, fid_db_path, jp_root_dir, db_journal_mode):
+    if db_journal_mode == "invalid":  # test invalid
+        with pytest.raises(TraitError, match=" must be one of "):
+            any_fid_manager_class(
+                db_path=fid_db_path, root_dir=str(jp_root_dir), db_journal_mode=db_journal_mode
+            )
+    else:
+        if not db_journal_mode:  # test correct defaults
+            expected_journal_mode = (
+                "WAL" if any_fid_manager_class.__name__ == "LocalFileIdManager" else "DELETE"
+            )
+            fid_manager = any_fid_manager_class(db_path=fid_db_path, root_dir=str(jp_root_dir))
+        else:  # test any valid value
+            expected_journal_mode = db_journal_mode
+            fid_manager = any_fid_manager_class(
+                db_path=fid_db_path, root_dir=str(jp_root_dir), db_journal_mode=db_journal_mode
+            )
+
+        cursor = fid_manager.con.execute("PRAGMA journal_mode")
+        actual_journal_mode = cursor.fetchone()
+        assert actual_journal_mode[0].upper() == expected_journal_mode