diff --git a/src/basic_memory/repository/repository.py b/src/basic_memory/repository/repository.py index 377d14051..4947e2e4f 100644 --- a/src/basic_memory/repository/repository.py +++ b/src/basic_memory/repository/repository.py @@ -23,6 +23,12 @@ T = TypeVar("T", bound=Base) +# SQLite caps the number of bound parameters per statement (999 on builds older +# than 3.32.0), so id-list queries must stay below that floor. 500 leaves +# headroom for additional binds (e.g. the project_id filter) while keeping the +# number of round-trips low. +SELECT_BY_IDS_CHUNK_SIZE = 500 + class Repository[T: Base]: """Base repository implementation with explicit caller-owned sessions.""" @@ -83,15 +89,25 @@ async def select_by_id(self, session: AsyncSession, entity_id: int) -> Optional[ return result.scalars().one_or_none() async def select_by_ids(self, session: AsyncSession, ids: List[int]) -> Sequence[T]: - """Select multiple entities by IDs using an existing session.""" - query = ( - select(self.Model).where(self.primary_key.in_(ids)).options(*self.get_load_options()) - ) - # Add project filter if applicable - query = self._add_project_filter(query) + """Select multiple entities by IDs using an existing session. - result = await session.execute(query) - return result.scalars().all() + Queries in chunks so callers can pass arbitrarily large id lists + without hitting SQLite's bound-parameter limit (issue #1045). + """ + results: list[T] = [] + for start in range(0, len(ids), SELECT_BY_IDS_CHUNK_SIZE): + chunk = ids[start : start + SELECT_BY_IDS_CHUNK_SIZE] + query = ( + select(self.Model) + .where(self.primary_key.in_(chunk)) + .options(*self.get_load_options()) + ) + # Add project filter if applicable + query = self._add_project_filter(query) + + result = await session.execute(query) + results.extend(result.scalars().all()) + return results async def add(self, session: AsyncSession, model: T) -> T: """ diff --git a/tests/repository/test_repository.py b/tests/repository/test_repository.py index bb900c7f8..d974db881 100644 --- a/tests/repository/test_repository.py +++ b/tests/repository/test_repository.py @@ -1,6 +1,8 @@ """Test repository implementation.""" from datetime import datetime, UTC +from unittest.mock import AsyncMock + import pytest from sqlalchemy import String, DateTime from sqlalchemy.orm import Mapped, mapped_column @@ -133,6 +135,29 @@ async def test_find_by_ids(repository, session_maker): assert len(not_found) == 0 +@pytest.mark.asyncio +async def test_find_by_ids_chunks_large_requests(repository, session_maker, monkeypatch): + """Regression test for #1045: bulk id hydration must issue bounded queries.""" + instances = [ModelTest(id=f"test_{i}", name=f"Test {i}") for i in range(5)] + async with db.scoped_session(session_maker) as session: + await repository.add_all_no_return(session, instances) + + # Use a tiny deterministic chunk size so this test proves the query is + # split even on SQLite builds whose real parameter cap exceeds 1,100. + monkeypatch.setattr( + "basic_memory.repository.repository.SELECT_BY_IDS_CHUNK_SIZE", + 2, + ) + execute = AsyncMock(wraps=session.execute) + monkeypatch.setattr(session, "execute", execute) + + ids_to_find = [instance.id for instance in instances] + found = await repository.find_by_ids(session, ids_to_find) + + assert execute.await_count == 3 + assert sorted(e.id for e in found) == sorted(ids_to_find) + + @pytest.mark.asyncio async def test_delete_by_ids(repository, session_maker): """Test finding multiple entities by IDs."""