Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
32 changes: 24 additions & 8 deletions src/basic_memory/repository/repository.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down Expand Up @@ -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:
"""
Expand Down
25 changes: 25 additions & 0 deletions tests/repository/test_repository.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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."""
Expand Down
Loading