From a6b7f445f854d9e05ac0d27c4098ed59ccb12102 Mon Sep 17 00:00:00 2001 From: phernandez Date: Sun, 30 Aug 2026 13:27:27 -0500 Subject: [PATCH] fix(core): batch relation search refresh cleanup Signed-off-by: phernandez --- .../repository/relation_repository.py | 25 +++++-- .../test_relation_search_refresh_batching.py | 69 +++++++++++++++++++ 2 files changed, 87 insertions(+), 7 deletions(-) create mode 100644 tests/repository/test_relation_search_refresh_batching.py diff --git a/src/basic_memory/repository/relation_repository.py b/src/basic_memory/repository/relation_repository.py index 6772eff50..8aa04458a 100644 --- a/src/basic_memory/repository/relation_repository.py +++ b/src/basic_memory/repository/relation_repository.py @@ -34,6 +34,7 @@ RESOLVED_RELATION_WRITE_STATEMENT_SIZE = 250 RELATION_GENERATION_WRITE_STATEMENT_SIZE = 250 +RELATION_SEARCH_REFRESH_DELETE_STATEMENT_SIZE = 500 LEGACY_RELATION_GENERATION = 0 @@ -676,12 +677,19 @@ async def clear_pending_search_refreshes( refresh_ids: Sequence[int], ) -> None: """Retire only refresh work completed by the caller's successful pass.""" - await session.execute( - delete(RelationSearchRefresh).where( - RelationSearchRefresh.project_id == self.project_id, - RelationSearchRefresh.id.in_(refresh_ids), + # Refresh backlogs can contain tens of thousands of durable markers. Keep each + # statement below database-driver parameter ceilings while preserving the caller's + # transaction as the atomic retirement boundary. + for refresh_id_batch in batched( + refresh_ids, + RELATION_SEARCH_REFRESH_DELETE_STATEMENT_SIZE, + ): + await session.execute( + delete(RelationSearchRefresh).where( + RelationSearchRefresh.project_id == self.project_id, + RelationSearchRefresh.id.in_(refresh_id_batch), + ) ) - ) async def complete_search_refresh_for_generation( self, @@ -704,11 +712,14 @@ async def complete_search_refresh_for_generation( entity_id=entity_id, generation=generation, ) - if refresh_ids: + for refresh_id_batch in batched( + refresh_ids, + RELATION_SEARCH_REFRESH_DELETE_STATEMENT_SIZE, + ): await session.execute( delete(RelationSearchRefresh).where( RelationSearchRefresh.project_id == self.project_id, - RelationSearchRefresh.id.in_(refresh_ids), + RelationSearchRefresh.id.in_(refresh_id_batch), generation_is_current, ) ) diff --git a/tests/repository/test_relation_search_refresh_batching.py b/tests/repository/test_relation_search_refresh_batching.py new file mode 100644 index 000000000..10a031442 --- /dev/null +++ b/tests/repository/test_relation_search_refresh_batching.py @@ -0,0 +1,69 @@ +"""Regression coverage for bounded relation-search refresh cleanup.""" + +from unittest.mock import AsyncMock + +import pytest +from sqlalchemy.dialects import postgresql +from sqlalchemy.ext.asyncio import AsyncSession + +from basic_memory.repository import relation_repository as relation_repository_module +from basic_memory.repository.relation_repository import RelationRepository + + +def _expanded_refresh_id_batches(session: AsyncMock) -> list[tuple[int, ...]]: + """Return the refresh IDs bound by each recorded DELETE statement.""" + batches: list[tuple[int, ...]] = [] + for call in session.execute.await_args_list: + statement = call.args[0] + compiled = statement.compile( + dialect=postgresql.dialect(), + compile_kwargs={"render_postcompile": True}, + ) + batches.append( + tuple( + value + for parameter, value in compiled.params.items() + if parameter.startswith("id_1_") + ) + ) + return batches + + +@pytest.mark.asyncio +async def test_clear_pending_search_refreshes_batches_delete_parameters(monkeypatch) -> None: + """A resolver backlog is retired without one unbounded asyncpg parameter list.""" + monkeypatch.setattr( + relation_repository_module, + "RELATION_SEARCH_REFRESH_DELETE_STATEMENT_SIZE", + 2, + ) + session = AsyncMock(spec=AsyncSession) + repository = RelationRepository(project_id=7) + + await repository.clear_pending_search_refreshes(session, [10, 11, 12, 13, 14]) + + assert _expanded_refresh_id_batches(session) == [(10, 11), (12, 13), (14,)] + + +@pytest.mark.asyncio +async def test_generation_refresh_completion_batches_delete_parameters(monkeypatch) -> None: + """Generation-fenced cleanup applies the same bound without losing its fence.""" + monkeypatch.setattr( + relation_repository_module, + "RELATION_SEARCH_REFRESH_DELETE_STATEMENT_SIZE", + 2, + ) + session = AsyncMock(spec=AsyncSession) + session.scalar.return_value = 1 + repository = RelationRepository(project_id=7) + + is_current = await repository.complete_search_refresh_for_generation( + session, + entity_id=23, + generation=5, + refresh_ids=[10, 11, 12, 13, 14], + ) + + assert is_current is True + assert _expanded_refresh_id_batches(session) == [(10, 11), (12, 13), (14,)] + session.add.assert_not_called()