fix(memory): close SQLite connections in memory/fts5/graph adapters

_get_conn() opened a fresh sqlite3.Connection per call and every call
site used `with self._get_conn() as conn:`. In Python's stdlib,
Connection.__enter__/__exit__ only commit/rollback the transaction on
exit — they never close the connection. With no close() anywhere in
these three files, every memory save/query/search/delete leaked one
OS file descriptor, eventually exhausting the FD ulimit in a
long-lived proxy process.

Turn _get_conn() into a contextmanager that wraps the same
commit/rollback semantics (`with conn:`) and closes the connection in
a finally block, so call sites are unchanged.
This commit is contained in:
Ashish Gaonker 2026-08-19 11:42:18 -07:00
parent 0ec73faa28
commit 6da5cec7d1
4 changed files with 109 additions and 13 deletions

View file

@ -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."""

View file

@ -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."""

View file

@ -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."""

View file

@ -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)