diff --git a/hivemind_sqlite_database/__init__.py b/hivemind_sqlite_database/__init__.py index 25325a8..fa9e0cc 100644 --- a/hivemind_sqlite_database/__init__.py +++ b/hivemind_sqlite_database/__init__.py @@ -26,6 +26,13 @@ class SQLiteDB(AbstractDB): name: str = "clients" subfolder: str = "hivemind-core" password: Optional[str] = None + # Overrides the computed xdg path entirely, e.g. ":memory:" for tests. + # Every thread's connection is opened against this same target, so + # nothing ever silently falls back to the real client database. + # ":memory:" becomes a named shared in-memory database, one per + # SQLiteDB instance, so worker threads see the same tables as the + # thread that created them. + db_path: Optional[str] = None # How long SQLite waits for a file lock held by another connection # before giving up with SQLITE_BUSY, in milliseconds. @@ -51,9 +58,22 @@ def __post_init__(self): up as ``sqlite3.ProgrammingError: bad parameter or other API misuse`` and as writes that land with corrupted bindings. """ - self._db_path = os.path.join(xdg_data_home(), self.subfolder, self.name + ".db") + if self.db_path is not None: + self._db_path = self.db_path + if self._db_path == ":memory:": + # A plain ":memory:" database belongs to the one connection + # that opened it, so the next thread would open a second, + # empty database and find no tables. Give it a name and a + # shared cache instead. The name carries the instance id, so + # two in-memory databases in one process stay independent. + self._db_path = f"file:hivemind-{id(self):x}?mode=memory&cache=shared" + else: + self._db_path = os.path.join(xdg_data_home(), self.subfolder, self.name + ".db") + self._is_uri = self._db_path.startswith("file:") + parent = os.path.dirname(self._db_path) + if parent and not self._is_uri: + os.makedirs(parent, exist_ok=True) LOG.debug(f"sqlite database path: {self._db_path}") - os.makedirs(os.path.dirname(self._db_path), exist_ok=True) if self.password is not None and self.password == "": raise ValueError("password must be non-empty when encryption is enabled") @@ -61,6 +81,11 @@ def __post_init__(self): self._write_lock = threading.Lock() # opening the first connection also applies the WAL pragma self.conn.execute("PRAGMA journal_mode=WAL") + if "mode=memory" in self._db_path: + # A shared in-memory database lives only while a connection to + # it is open. Hold one for the lifetime of this object so the + # tables survive a thread closing its own connection. + self._keepalive = self._connect() self._initialize_database() self._maybe_migrate() @@ -75,12 +100,14 @@ def _connect(self): "Install the system library (e.g. 'apt install libsqlcipher-dev') " "then: pip install hivemind-sqlite-database[cipher]" ) - conn = _sqlcipher.connect(self._db_path, check_same_thread=False) + conn = _sqlcipher.connect(self._db_path, check_same_thread=False, + uri=self._is_uri) conn.row_factory = _sqlcipher.Row escaped_password = self.password.replace("'", "''") conn.execute(f"PRAGMA key='{escaped_password}'") else: - conn = sqlite3.connect(self._db_path, check_same_thread=False) + conn = sqlite3.connect(self._db_path, check_same_thread=False, + uri=self._is_uri) conn.row_factory = sqlite3.Row conn.execute(f"PRAGMA busy_timeout={int(self.BUSY_TIMEOUT_MS)}") return conn @@ -102,10 +129,15 @@ def conn(self): @conn.setter def conn(self, value) -> None: - """Adopt an already-open connection for the calling thread. - - Only the calling thread sees it; other threads still open their - own. Tests use this to inject an in-memory database. + """Adopt an already-open connection for the calling thread only. + + Any other thread that later touches ``.conn`` still opens its own + connection against ``self._db_path`` — if that is the real client + database, that thread silently writes to it. Pass ``db_path`` + (e.g. ``":memory:"``) to the constructor instead so every thread, + present and future, agrees on the same target; this setter exists + only for single-threaded call sites that already hold a connection + they want reused. """ self._thread_state().conn = value diff --git a/tests/test_sqlitedb.py b/tests/test_sqlitedb.py index 1b3d11a..027fd7f 100644 --- a/tests/test_sqlitedb.py +++ b/tests/test_sqlitedb.py @@ -422,6 +422,56 @@ def test_full_client_fields_preserved(self): self.assertEqual(r.intent_blacklist, ["a:b"]) +class TestSQLiteDBPathOverride(unittest.TestCase): + """A worker thread must never fall back to the real client database + just because the test only overrode `.conn` on the main thread.""" + + def test_db_path_override_keeps_worker_threads_off_disk(self): + import unittest.mock as mock + with tempfile.TemporaryDirectory() as tmpdir: + with mock.patch( + "hivemind_sqlite_database.xdg_data_home", return_value=tmpdir + ): + db = SQLiteDB(name="clients", subfolder="hivemind-core", + db_path=":memory:") + + # Read a row written on this thread. "SELECT 1" would answer + # the same against a private, empty database and so would + # never notice a per-thread database. + db.add_item(Client(client_id=1, api_key="key", + name="kitchen")) + + errors = [] + seen = [] + + def worker(): + try: + seen.extend(db.search_by_value("name", "kitchen")) + except Exception as e: # noqa: BLE001 + errors.append(e) + + t = threading.Thread(target=worker) + t.start() + t.join() + + self.assertEqual(errors, []) + self.assertEqual([c.client_id for c in seen], [1]) + real_db_file = os.path.join(tmpdir, "hivemind-core", "clients.db") + self.assertFalse(os.path.exists(real_db_file)) + + def test_two_in_memory_databases_stay_independent(self): + first = SQLiteDB(db_path=":memory:") + second = SQLiteDB(db_path=":memory:") + first.add_item(Client(client_id=1, api_key="key", name="kitchen")) + self.assertEqual(second.search_by_value("name", "kitchen"), []) + + def test_explicit_db_path_creates_missing_directories(self): + with tempfile.TemporaryDirectory() as tmpdir: + path = os.path.join(tmpdir, "newdir", "clients.db") + SQLiteDB(db_path=path) + self.assertTrue(os.path.isfile(path)) + + class TestSQLiteDBCommit(unittest.TestCase): def test_commit_returns_true(self): db = make_db()