diff --git a/headroom/memory/adapters/fts5.py b/headroom/memory/adapters/fts5.py index 1b679e458..a0920a444 100644 --- a/headroom/memory/adapters/fts5.py +++ b/headroom/memory/adapters/fts5.py @@ -7,6 +7,7 @@ and Unicode tokenization for high-quality search results. from __future__ import annotations +import contextlib import re import sqlite3 from dataclasses import dataclass, field @@ -17,7 +18,7 @@ from ..models import Memory from ..ports import TextFilter, TextSearchResult if TYPE_CHECKING: - pass + from collections.abc import Iterator @dataclass @@ -67,15 +68,20 @@ class FTS5TextIndex: self.db_path = Path(db_path) self._init_db() - def _get_conn(self) -> sqlite3.Connection: + @contextlib.contextmanager + def _get_conn(self) -> Iterator[sqlite3.Connection]: """Get a new database connection (thread-safe pattern). - Returns: - A new SQLite connection with row factory configured. + Commits on clean exit, rolls back on exception, and always closes + the connection -- callers use ``with self._get_conn() as conn:``. """ conn = sqlite3.connect(str(self.db_path)) conn.row_factory = sqlite3.Row - return conn + try: + with conn: + yield conn + finally: + conn.close() def _init_db(self) -> None: """Initialize the FTS5 virtual table schema.""" diff --git a/headroom/memory/adapters/sqlite.py b/headroom/memory/adapters/sqlite.py index c5fcdf28c..c57d07132 100644 --- a/headroom/memory/adapters/sqlite.py +++ b/headroom/memory/adapters/sqlite.py @@ -9,6 +9,7 @@ Provides persistent storage for Memory objects with full support for: from __future__ import annotations +import contextlib import json import re import sqlite3 @@ -20,6 +21,8 @@ from ..models import Memory, ScopeLevel, normalize_entity_refs from ..ports import MemoryFilter if TYPE_CHECKING: + from collections.abc import Iterator + import numpy as np # Regex pattern for safe metadata keys: alphanumeric, underscores, hyphens only @@ -73,15 +76,20 @@ class SQLiteMemoryStore: self.db_path = Path(db_path) self._init_db() - def _get_conn(self) -> sqlite3.Connection: + @contextlib.contextmanager + def _get_conn(self) -> Iterator[sqlite3.Connection]: """Get a new database connection (thread-safe pattern). - Returns: - A new SQLite connection with row factory configured. + Commits on clean exit, rolls back on exception, and always closes + the connection -- callers use ``with self._get_conn() as conn:``. """ conn = sqlite3.connect(str(self.db_path)) conn.row_factory = sqlite3.Row - return conn + try: + with conn: + yield conn + finally: + conn.close() def _init_db(self) -> None: """Initialize the database schema with indexes.""" diff --git a/headroom/memory/adapters/sqlite_graph.py b/headroom/memory/adapters/sqlite_graph.py index 2c2ac432c..cda1fa6f7 100644 --- a/headroom/memory/adapters/sqlite_graph.py +++ b/headroom/memory/adapters/sqlite_graph.py @@ -13,6 +13,7 @@ This is a drop-in replacement for InMemoryGraphStore that: from __future__ import annotations +import contextlib import json import sqlite3 from collections import deque @@ -24,6 +25,8 @@ from typing import TYPE_CHECKING, Any from .graph_models import Entity, Relationship, RelationshipDirection, Subgraph if TYPE_CHECKING: + from collections.abc import Iterator + from ..tracker import ComponentStats @@ -71,11 +74,12 @@ class SQLiteGraphStore: self._lock = RLock() self._init_db() - def _get_conn(self) -> sqlite3.Connection: + @contextlib.contextmanager + def _get_conn(self) -> Iterator[sqlite3.Connection]: """Get a new database connection (thread-safe pattern). - Returns: - A new SQLite connection with row factory configured. + Commits on clean exit, rolls back on exception, and always closes + the connection -- callers use ``with self._get_conn() as conn:``. """ conn = sqlite3.connect(str(self.db_path)) conn.row_factory = sqlite3.Row @@ -87,7 +91,11 @@ class SQLiteGraphStore: # Enable foreign keys conn.execute("PRAGMA foreign_keys = ON") - return conn + try: + with conn: + yield conn + finally: + conn.close() def _init_db(self) -> None: """Initialize the database schema with indexes.""" diff --git a/tests/test_memory/test_adapters_connection_leak.py b/tests/test_memory/test_adapters_connection_leak.py new file mode 100644 index 000000000..06836196a --- /dev/null +++ b/tests/test_memory/test_adapters_connection_leak.py @@ -0,0 +1,74 @@ +"""Regression tests: SQLite adapters must close every connection they open. + +``_get_conn()`` in each adapter opens a brand-new ``sqlite3.Connection`` per +call. ``with conn:`` only commits/rolls back on exit -- it does not close the +connection -- so without an explicit ``.close()`` every operation leaks a +file descriptor. +""" + +from __future__ import annotations + +import sqlite3 +from pathlib import Path + +import pytest + +from headroom.memory.adapters.fts5 import FTS5TextIndex +from headroom.memory.adapters.graph_models import Entity +from headroom.memory.adapters.sqlite import SQLiteMemoryStore +from headroom.memory.adapters.sqlite_graph import SQLiteGraphStore +from headroom.memory.models import Memory + + +def _track_connections(monkeypatch: pytest.MonkeyPatch) -> list[sqlite3.Connection]: + """Wrap ``sqlite3.connect`` to record every connection it creates.""" + created: list[sqlite3.Connection] = [] + real_connect = sqlite3.connect + + def tracking_connect(*args: object, **kwargs: object) -> sqlite3.Connection: + conn = real_connect(*args, **kwargs) # type: ignore[arg-type] + created.append(conn) + return conn + + monkeypatch.setattr(sqlite3, "connect", tracking_connect) + return created + + +def _assert_all_closed(connections: list[sqlite3.Connection]) -> None: + assert connections, "expected at least one sqlite3.Connection to be created" + for conn in connections: + with pytest.raises(sqlite3.ProgrammingError, match="closed database"): + conn.execute("SELECT 1") + + +async def test_sqlite_memory_store_closes_connections( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + created = _track_connections(monkeypatch) + + store = SQLiteMemoryStore(str(tmp_path / "mem.db")) + await store.save(Memory(content="hello", user_id="alice")) + + _assert_all_closed(created) + + +def test_fts5_text_index_closes_connections( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + created = _track_connections(monkeypatch) + + index = FTS5TextIndex(str(tmp_path / "fts.db")) + index.index_raw("mem-1", "hello world", {"user_id": "alice"}) + + _assert_all_closed(created) + + +async def test_sqlite_graph_store_closes_connections( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + created = _track_connections(monkeypatch) + + store = SQLiteGraphStore(str(tmp_path / "graph.db")) + await store.add_entity(Entity(user_id="alice", name="Project X", entity_type="project")) + + _assert_all_closed(created)