Skip to content

fix(sqlite): isolate thread-local connections by database #472

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
dplush wants to merge 1 commit into from
Closed

fix(sqlite): isolate thread-local connections by database #472

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
128 changes: 103 additions & 25 deletions mnemosyne/core/beam.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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
Comment on lines +469 to +478

@coderabbitai coderabbitai Bot Jul 15, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Arbitrary probe failures must not create a second connection generation.

Both caches evict without closing on any SELECT 1 exception, potentially leaking the original handle and splitting existing wrappers from future acquisitions.

  • mnemosyne/core/beam.py#L469-L478: recover only from a confirmed closed handle; propagate transient SQLite errors.
  • mnemosyne/core/memory.py#L85-L90: apply the same narrow recovery and explicit cleanup.
🧰 Tools
🪛 Ruff (0.15.21)

[warning] 476-476: Do not catch blind exception: Exception

(BLE001)

📍 Affects 2 files
  • mnemosyne/core/beam.py#L469-L478 (this comment)
  • mnemosyne/core/memory.py#L85-L90
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@mnemosyne/core/beam.py` around lines 469 - 478, In the cached connection
validation in mnemosyne/core/beam.py lines 469-478, recover only when the probe
confirms the handle is closed; propagate transient SQLite errors instead of
evicting the cache, and explicitly close the invalid connection before removal.
Apply the same narrow recovery and cleanup to the corresponding cache logic in
mnemosyne/core/memory.py lines 85-90, preserving a single connection generation
for arbitrary probe failures.

Source: Linters/SAST tools


if needs_reconnect:
if conn is None:
path.parent.mkdir(parents=True, exist_ok=True)
conn = sqlite3.connect(
str(path),
Expand All @@ -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)
Comment on lines +508 to +526

@coderabbitai coderabbitai Bot Jul 15, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift

Path-only ownership can close another live wrapper’s connection.

The release operation must identify the exact connection generation and acquisition thread, rather than consulting whichever thread-local path entry exists when close() runs.

  • mnemosyne/core/beam.py#L508-L526: return/store an exact BEAM lease token and validate it during release.
  • mnemosyne/core/beam.py#L3094-L3101: release the stored lease instead of only self.db_path.
  • mnemosyne/core/memory.py#L110-L128: make core ownership generation- and thread-aware.
  • mnemosyne/core/memory.py#L291-L293: release the wrapper’s stored core lease.
🧰 Tools
🪛 Ruff (0.15.21)

[warning] 508-508: PEP 484 prohibits implicit Optional

Convert to T | None

(RUF013)

[warning] 517-517: PEP 484 prohibits implicit Optional

Convert to T | None

(RUF013)

📍 Affects 2 files
  • mnemosyne/core/beam.py#L508-L526 (this comment)
  • mnemosyne/core/beam.py#L3094-L3101
  • mnemosyne/core/memory.py#L110-L128
  • mnemosyne/core/memory.py#L291-L293
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@mnemosyne/core/beam.py` around lines 508 - 526, Replace path-only ownership
tracking with exact lease tokens containing the connection generation and
acquisition thread. In mnemosyne/core/beam.py lines 508-526, have acquisition
return and release validate the stored BEAM lease; update BeamMemory.close
around lines 3094-3101 to release that lease instead of self.db_path. Apply the
same generation- and thread-aware lease tracking in mnemosyne/core/memory.py
lines 110-128, and make the wrapper release its stored core lease around lines
291-293.



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:
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
# ------------------------------------------------------------------
Expand Down
161 changes: 114 additions & 47 deletions mnemosyne/core/memory.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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]}
Expand Down Expand Up @@ -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
Expand All @@ -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
Comment on lines +258 to +280

@coderabbitai coderabbitai Bot Jul 15, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift

Avoid creating an unowned BEAM connection before guarded acquisition.

init_db() calls init_beam(), which can populate BEAM’s cache before BeamMemory acquires ownership. If BeamMemory then fails before _acquire_beam_connection()—for example in get_config()—this handler releases only the core connection, leaving the BEAM handle cached with zero owners.

Initialize BEAM under its wrapper lease, or add ownership-aware rollback for this path.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@mnemosyne/core/memory.py` around lines 258 - 280, Update BeamMemory
initialization around _acquire_connection and BeamMemory construction so BEAM is
initialized through its ownership-aware wrapper lease before any failure-prone
setup such as init_db or get_config. Ensure the exception path releases both
core and BEAM resources when acquired, preventing a cached BEAM handle from
remaining with zero owners; preserve existing cleanup for the core connection.


def close(self) -> None:
"""Close the database connection and checkpoint the WAL.
Expand All @@ -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."""
Expand All @@ -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
Expand Down
Loading
Loading