mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
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:
parent
0ec73faa28
commit
6da5cec7d1
4 changed files with 109 additions and 13 deletions
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
74
tests/test_memory/test_adapters_connection_leak.py
Normal file
74
tests/test_memory/test_adapters_connection_leak.py
Normal 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)
|
||||
Loading…
Add table
Add a link
Reference in a new issue