diff --git a/mnemosyne/core/beam.py b/mnemosyne/core/beam.py index 48bd2148..f02b2075 100644 --- a/mnemosyne/core/beam.py +++ b/mnemosyne/core/beam.py @@ -437,6 +437,27 @@ def _session_scope_params(session_id: str, extra_value=None) -> list: VEC_TYPE = "float32" +def _canonical_db_path(db_path: Path = None) -> tuple[str, Path]: + """Return a stable cache key and path for a database location.""" + path = Path(db_path) if db_path else _default_db_path() + path = path.expanduser().resolve() + return str(path), path + + +def _connections() -> dict[str, sqlite3.Connection]: + """Return this thread's BEAM connections, keyed by canonical path.""" + if not hasattr(_thread_local, "connections"): + _thread_local.connections = {} + return _thread_local.connections + + +def _owners() -> dict[str, int]: + """Return this thread's live BeamMemory-wrapper count per database.""" + if not hasattr(_thread_local, "owners"): + _thread_local.owners = {} + return _thread_local.owners + + def _get_connection(db_path: Path = None) -> sqlite3.Connection: """Get thread-local database connection with extensions loaded. @@ -445,20 +466,18 @@ def _get_connection(db_path: Path = None) -> sqlite3.Connection: `_deferred_commits`. Connection is otherwise identical to a plain sqlite3.Connection. """ - path = Path(db_path) if db_path else _default_db_path() - needs_reconnect = ( - not hasattr(_thread_local, 'conn') - or _thread_local.conn is None - or getattr(_thread_local, 'db_path', None) != str(path) - ) - if not needs_reconnect: + key, path = _canonical_db_path(db_path) + connections = _connections() + conn = connections.get(key) + if conn is not None: # Verify the cached connection is still alive try: - _thread_local.conn.execute("SELECT 1") + conn.execute("SELECT 1") except Exception: - needs_reconnect = True + connections.pop(key, None) + conn = None - if needs_reconnect: + if conn is None: path.parent.mkdir(parents=True, exist_ok=True) conn = sqlite3.connect( str(path), @@ -482,30 +501,59 @@ def _get_connection(db_path: Path = None) -> sqlite3.Connection: sqlite_vec.load(conn) except Exception: pass # Some environments don't support load_extension - _thread_local.conn = conn - _thread_local.db_path = str(path) - return _thread_local.conn + connections[key] = conn + return conn + + +def _acquire_beam_connection(db_path: Path = None) -> sqlite3.Connection: + """Get a connection and register one live BeamMemory wrapper for its path.""" + key, _ = _canonical_db_path(db_path) + conn = _get_connection(db_path) + owners = _owners() + owners[key] = owners.get(key, 0) + 1 + return conn + + +def _release_beam_connection(db_path: Path = None) -> None: + """Release one wrapper; close its connection only after its final owner.""" + key, path = _canonical_db_path(db_path) + owners = _owners() + remaining = owners.get(key, 0) - 1 + if remaining > 0: + owners[key] = remaining + return + owners.pop(key, None) + _close_beam_connection(path) -def _close_beam_connection() -> None: - """Close the thread-local BEAM connection and checkpoint the WAL. +def _close_beam_connection(db_path: Path = None) -> None: + """Close one cached BEAM connection, or all current-thread entries. - Thread-local connections under WAL mode can block checkpoints - after the owning thread exits. Explicitly closing the connection - and running ``PRAGMA wal_checkpoint(TRUNCATE)`` prevents the - ``database is locked`` cascade documented in #382. + Each close checkpoints WAL before releasing the handle, preventing the + ``database is locked`` cascade documented in #382. Explicit close is a + force-close operation; wrapper ownership uses ``_release_beam_connection``. """ - if hasattr(_thread_local, 'conn') and _thread_local.conn is not None: + connections = _connections() + owners = _owners() + if db_path is None: + keys = list(connections) + else: + key, _ = _canonical_db_path(db_path) + keys = [key] + + for key in keys: + conn = connections.pop(key, None) + owners.pop(key, None) + if conn is None: + continue try: - _thread_local.conn.execute("PRAGMA wal_checkpoint(TRUNCATE)") + conn.execute("PRAGMA wal_checkpoint(TRUNCATE)") except Exception: pass try: - _thread_local.conn.close() + conn.close() except Exception: pass - _thread_local.conn = None - _thread_local.db_path = None def _detect_vec_type(conn: sqlite3.Connection) -> str: @@ -2979,7 +3027,21 @@ def __init__(self, session_id: str = "default", db_path: Path = None, self._extraction_client = None # Lazy-loaded ExtractionClient self._extraction_buffer = [] # Buffer for batch extraction self._event_emitter = event_emitter # Streaming event callback - self.conn = _get_connection(self.db_path) + self._closed = False + self._owns_connection = False + try: + self.conn = _acquire_beam_connection(self.db_path) + self._owns_connection = True + self._initialize_stores() + except BaseException: + if self._owns_connection: + self._owns_connection = False + _release_beam_connection(self.db_path) + self._closed = True + raise + + def _initialize_stores(self) -> None: + """Initialize schemas and shared stores after acquiring ``self.conn``.""" init_beam(self.db_path) # E6: ensure schema split + auto-migrate legacy TripleStore rows @@ -3029,6 +3091,22 @@ def __init__(self, session_id: str = "default", db_path: Path = None, except Exception: logger.info("Regex extraction failed, skipping", exc_info=True) + def close(self) -> None: + """Release this wrapper's BEAM connection ownership.""" + if self._closed: + return + self._closed = True + if self._owns_connection: + self._owns_connection = False + _release_beam_connection(self.db_path) + + def __del__(self): + """Best-effort cleanup on garbage collection.""" + try: + self.close() + except Exception: + pass + # ------------------------------------------------------------------ # E6 schema split + auto-migration # ------------------------------------------------------------------ diff --git a/mnemosyne/core/memory.py b/mnemosyne/core/memory.py index 75e7dfd6..dc8f6f57 100644 --- a/mnemosyne/core/memory.py +++ b/mnemosyne/core/memory.py @@ -56,46 +56,106 @@ def _default_db_path() -> Path: return _default_data_dir() / "mnemosyne.db" -def _get_connection(db_path = None) -> sqlite3.Connection: - """Get thread-local database connection""" +def _canonical_db_path(db_path=None) -> tuple[str, Path]: + """Return a stable cache key and path for a database location.""" path = Path(db_path) if db_path else _default_db_path() - if not hasattr(_thread_local, 'conn') or _thread_local.conn is None or getattr(_thread_local, 'db_path', None) != str(path): + path = path.expanduser().resolve() + return str(path), path + + +def _connections() -> dict[str, sqlite3.Connection]: + """Return this thread's database connections, keyed by canonical path.""" + if not hasattr(_thread_local, "connections"): + _thread_local.connections = {} + return _thread_local.connections + + +def _owners() -> dict[str, int]: + """Return this thread's live Mnemosyne-wrapper count per database.""" + if not hasattr(_thread_local, "owners"): + _thread_local.owners = {} + return _thread_local.owners + + +def _get_connection(db_path=None) -> sqlite3.Connection: + """Get a live thread-local database connection for ``db_path``.""" + key, path = _canonical_db_path(db_path) + connections = _connections() + conn = connections.get(key) + if conn is not None: + try: + conn.execute("SELECT 1") + except Exception: + connections.pop(key, None) + conn = None + + if conn is None: path.parent.mkdir(parents=True, exist_ok=True) - _thread_local.conn = sqlite3.connect(str(path), check_same_thread=False) - _thread_local.conn.row_factory = sqlite3.Row - _thread_local.conn.execute("PRAGMA journal_mode=WAL") - _thread_local.conn.execute("PRAGMA busy_timeout=5000") - _thread_local.conn.execute("PRAGMA foreign_keys=ON") + conn = sqlite3.connect(str(path), check_same_thread=False) + conn.row_factory = sqlite3.Row + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("PRAGMA busy_timeout=5000") + conn.execute("PRAGMA foreign_keys=ON") # Load sqlite-vec extension for vector search (matches beam._get_connection) try: import sqlite_vec - _thread_local.conn.enable_load_extension(True) - sqlite_vec.load(_thread_local.conn) + conn.enable_load_extension(True) + sqlite_vec.load(conn) except Exception: pass - _thread_local.db_path = str(path) - return _thread_local.conn + connections[key] = conn + return conn -def _close_connection() -> None: - """Close the thread-local connection and checkpoint the WAL. +def _acquire_connection(db_path=None) -> sqlite3.Connection: + """Get a connection and register one live Mnemosyne wrapper for its path.""" + key, _ = _canonical_db_path(db_path) + conn = _get_connection(db_path) + owners = _owners() + owners[key] = owners.get(key, 0) + 1 + return conn + - Thread-local connections under WAL mode can block checkpoints - after the owning thread exits. Explicitly closing the connection - and running ``PRAGMA wal_checkpoint(TRUNCATE)`` prevents the - ``database is locked`` cascade documented in #382. +def _release_connection(db_path=None) -> None: + """Release one wrapper; close its connection only after its final owner.""" + key, path = _canonical_db_path(db_path) + owners = _owners() + remaining = owners.get(key, 0) - 1 + if remaining > 0: + owners[key] = remaining + return + owners.pop(key, None) + _close_connection(path) + + +def _close_connection(db_path=None) -> None: + """Close one cached connection, or every cached connection when omitted. + + Each close checkpoints WAL before releasing the handle, preventing the + ``database is locked`` cascade documented in #382. Explicit close is a + force-close operation; wrapper ownership uses ``_release_connection``. """ - if hasattr(_thread_local, 'conn') and _thread_local.conn is not None: + connections = _connections() + owners = _owners() + if db_path is None: + keys = list(connections) + else: + key, _ = _canonical_db_path(db_path) + keys = [key] + + for key in keys: + conn = connections.pop(key, None) + owners.pop(key, None) + if conn is None: + continue try: - _thread_local.conn.execute("PRAGMA wal_checkpoint(TRUNCATE)") + conn.execute("PRAGMA wal_checkpoint(TRUNCATE)") except Exception: pass try: - _thread_local.conn.close() + conn.close() except Exception: pass - _thread_local.conn = None - _thread_local.db_path = None def wal_checkpoint(db_path=None) -> dict: @@ -104,10 +164,7 @@ def wal_checkpoint(db_path=None) -> dict: Returns a dict with ``busy``, ``log``, ``checkpointed`` keys so callers can monitor whether the checkpoint succeeded. """ - conn = _get_connection(db_path) if db_path else ( - _thread_local.conn if hasattr(_thread_local, 'conn') and _thread_local.conn is not None - else _get_connection() - ) + conn = _get_connection(db_path) try: row = conn.execute("PRAGMA wal_checkpoint(TRUNCATE)").fetchone() return {"busy": row[0], "log": row[1], "checkpointed": row[2]} @@ -177,6 +234,8 @@ class Mnemosyne: def __init__(self, session_id: str = "default", db_path: Path = None, bank: str = None, author_id: str = None, author_type: str = None, channel_id: str = None): + self._closed = False + self._owns_connection = False # Auto-seed config.yaml on first Mnemosyne init from mnemosyne.core.config import get_config get_config() # triggers _seed() if config.yaml doesn't exist @@ -196,10 +255,29 @@ def __init__(self, session_id: str = "default", db_path: Path = None, bank: str else: self.db_path = _default_db_path() - self.conn = _get_connection(self.db_path) - init_db(self.db_path) - - self._closed = False + try: + self.conn = _acquire_connection(self.db_path) + self._owns_connection = True + init_db(self.db_path) + + # Phase 8: Streaming + Patterns + Plugins (lazy init) + self._stream = None + self._compressor = None + self._pattern_detector = None + self._delta_sync = None + self._plugin_manager = None + + # Create beam with streaming emitter wired + self.beam = BeamMemory(session_id=session_id, db_path=self.db_path, + author_id=author_id, author_type=author_type, + channel_id=channel_id, + event_emitter=self._stream_emit) + except BaseException: + if self._owns_connection: + self._owns_connection = False + _release_connection(self.db_path) + self._closed = True + raise def close(self) -> None: """Close the database connection and checkpoint the WAL. @@ -210,9 +288,11 @@ def close(self) -> None: if self._closed: return self._closed = True - _close_connection() - from mnemosyne.core.beam import _close_beam_connection - _close_beam_connection() + if self._owns_connection: + self._owns_connection = False + _release_connection(self.db_path) + if hasattr(self, "beam"): + self.beam.close() def __del__(self): """Best-effort cleanup on garbage collection.""" @@ -221,19 +301,6 @@ def __del__(self): except Exception: pass - # Phase 8: Streaming + Patterns + Plugins (lazy init) - self._stream = None - self._compressor = None - self._pattern_detector = None - self._delta_sync = None - self._plugin_manager = None - - # Create beam with streaming emitter wired - self.beam = BeamMemory(session_id=session_id, db_path=self.db_path, - author_id=author_id, author_type=author_type, - channel_id=channel_id, - event_emitter=self._stream_emit) - # ─── Phase 8: Streaming ───────────────────────────────────────── @property diff --git a/tests/conftest.py b/tests/conftest.py index 8cbf8c5e..28715b4e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -19,15 +19,14 @@ def _close_cached_connections(): try: import importlib mod = importlib.import_module(mod_path) - tl = getattr(mod, "_thread_local", None) - if tl is not None and hasattr(tl, "conn") and tl.conn is not None: - try: - tl.conn.close() - except Exception: - pass - tl.conn = None - if hasattr(tl, "db_path"): - tl.db_path = None + close_name = ( + "_close_beam_connection" + if mod_path.endswith(".beam") + else "_close_connection" + ) + close = getattr(mod, close_name, None) + if close is not None: + close() except Exception: pass @@ -78,9 +77,9 @@ def _reset_thread_local_connections(): tests that use different database paths. Both mnemosyne.core.beam and mnemosyne.core.memory maintain their own - thread-local caches (_thread_local.conn / _thread_local.db_path). - When tests create instances with different db_paths, the old connection - is never closed, leading to "database is locked" errors. + thread-local connection maps keyed by canonical database path. Closing all + entries around each test prevents per-path cache state from leaking into + later tests. """ _close_cached_connections() yield diff --git a/tests/test_beam.py b/tests/test_beam.py index c5a091d2..5cdb4644 100644 --- a/tests/test_beam.py +++ b/tests/test_beam.py @@ -952,6 +952,24 @@ def test_sleep_consolidated_content_is_recallable(self, temp_db, monkeypatch): class TestMnemosyneIntegration: + def test_close_isolated_by_database_path_and_live_owner(self, temp_db, tmp_path): + other_db = tmp_path / "other.db" + first = Mnemosyne(session_id="first", db_path=temp_db) + second = Mnemosyne(session_id="second", db_path=temp_db) + other = Mnemosyne(session_id="other", db_path=other_db) + + first.close() + assert second.conn.execute("SELECT 1").fetchone()[0] == 1 + assert second.beam.conn.execute("SELECT 1").fetchone()[0] == 1 + assert other.conn.execute("SELECT 1").fetchone()[0] == 1 + assert other.beam.conn.execute("SELECT 1").fetchone()[0] == 1 + + second.close() + with pytest.raises(sqlite3.ProgrammingError): + second.conn.execute("SELECT 1") + assert other.conn.execute("SELECT 1").fetchone()[0] == 1 + other.close() + def test_legacy_and_beam_dual_write(self, temp_db): mem = Mnemosyne(session_id="s2", db_path=temp_db) mid = mem.remember("Likes pizza", source="preference", importance=0.8)