diff --git a/src/basic_memory/models/search.py b/src/basic_memory/models/search.py index 292c0ac6..c409ed86 100644 --- a/src/basic_memory/models/search.py +++ b/src/basic_memory/models/search.py @@ -48,6 +48,15 @@ CREATE_POSTGRES_SEARCH_INDEX_METADATA = DDL(""" CREATE INDEX IF NOT EXISTS idx_search_index_metadata_gin ON search_index USING gin(metadata jsonb_path_ops) """) +# Partial unique index on (permalink, project_id) for non-null permalinks +# This prevents duplicate permalinks per project and is used by upsert operations +# in PostgresSearchRepository to handle race conditions during parallel indexing +CREATE_POSTGRES_SEARCH_INDEX_PERMALINK = DDL(""" +CREATE UNIQUE INDEX IF NOT EXISTS uix_search_index_permalink_project +ON search_index (permalink, project_id) +WHERE permalink IS NOT NULL +""") + # Define FTS5 virtual table creation for SQLite only # This DDL is executed separately for SQLite databases CREATE_SEARCH_INDEX = DDL(""" diff --git a/src/basic_memory/repository/postgres_search_repository.py b/src/basic_memory/repository/postgres_search_repository.py index a4c1673c..f0bc0f28 100644 --- a/src/basic_memory/repository/postgres_search_repository.py +++ b/src/basic_memory/repository/postgres_search_repository.py @@ -24,6 +24,11 @@ class PostgresSearchRepository(SearchRepositoryBase): - GIN indexes for performance - ts_rank() function for relevance scoring - JSONB containment operators for metadata search + + Note: This implementation uses UPSERT patterns (INSERT ... ON CONFLICT) instead of + delete-then-insert to handle race conditions during parallel entity indexing. + The partial unique index uix_search_index_permalink_project prevents duplicate + permalinks per project. """ async def init_search_index(self): @@ -41,6 +46,63 @@ class PostgresSearchRepository(SearchRepositoryBase): # - CREATE INDEX USING GIN on metadata jsonb_path_ops pass + async def index_item(self, search_index_row: SearchIndexRow) -> None: + """Index or update a single item using UPSERT. + + Uses INSERT ... ON CONFLICT to handle race conditions during parallel + entity indexing. The partial unique index uix_search_index_permalink_project + on (permalink, project_id) WHERE permalink IS NOT NULL prevents duplicate + permalinks. + + For rows with non-null permalinks (entities), conflicts are resolved by + updating the existing row. For rows with null permalinks, no conflict + occurs on this index. + """ + async with db.scoped_session(self.session_maker) as session: + # Serialize JSON for raw SQL + insert_data = search_index_row.to_insert(serialize_json=True) + insert_data["project_id"] = self.project_id + + # Use upsert to handle race conditions during parallel indexing + # ON CONFLICT (permalink, project_id) matches the partial unique index + # uix_search_index_permalink_project WHERE permalink IS NOT NULL + # For rows with NULL permalinks, no conflict occurs (partial index doesn't apply) + await session.execute( + text(""" + INSERT INTO search_index ( + id, title, content_stems, content_snippet, permalink, file_path, type, metadata, + from_id, to_id, relation_type, + entity_id, category, + created_at, updated_at, + project_id + ) VALUES ( + :id, :title, :content_stems, :content_snippet, :permalink, :file_path, :type, :metadata, + :from_id, :to_id, :relation_type, + :entity_id, :category, + :created_at, :updated_at, + :project_id + ) + ON CONFLICT (permalink, project_id) WHERE permalink IS NOT NULL DO UPDATE SET + id = EXCLUDED.id, + title = EXCLUDED.title, + content_stems = EXCLUDED.content_stems, + content_snippet = EXCLUDED.content_snippet, + file_path = EXCLUDED.file_path, + type = EXCLUDED.type, + metadata = EXCLUDED.metadata, + from_id = EXCLUDED.from_id, + to_id = EXCLUDED.to_id, + relation_type = EXCLUDED.relation_type, + entity_id = EXCLUDED.entity_id, + category = EXCLUDED.category, + created_at = EXCLUDED.created_at, + updated_at = EXCLUDED.updated_at + """), + insert_data, + ) + logger.debug(f"indexed row {search_index_row}") + await session.commit() + def _prepare_search_term(self, term: str, is_prefix: bool = True) -> str: """Prepare a search term for tsquery format. @@ -316,10 +378,14 @@ class PostgresSearchRepository(SearchRepositoryBase): async def bulk_index_items(self, search_index_rows: List[SearchIndexRow]) -> None: """Index multiple items in a single batch operation using UPSERT. - Uses INSERT ... ON CONFLICT DO UPDATE to handle re-indexing of existing - entities (e.g., during forward reference resolution) without requiring - a separate delete operation. This eliminates race conditions between - delete and insert operations in separate transactions. + Uses INSERT ... ON CONFLICT to handle race conditions during parallel + entity indexing. The partial unique index uix_search_index_permalink_project + on (permalink, project_id) WHERE permalink IS NOT NULL prevents duplicate + permalinks. + + For rows with non-null permalinks (entities), conflicts are resolved by + updating the existing row. For rows with null permalinks (observations, + relations), the partial index doesn't apply and they are inserted directly. Args: search_index_rows: List of SearchIndexRow objects to index @@ -338,11 +404,10 @@ class PostgresSearchRepository(SearchRepositoryBase): insert_data["project_id"] = self.project_id insert_data_list.append(insert_data) - # Use UPSERT (INSERT ... ON CONFLICT) to handle re-indexing - # Primary key is (id, type, project_id) - # This handles race conditions during forward reference resolution - # where an entity might be re-indexed before the delete commits - # Syntax works for both SQLite 3.24+ and PostgreSQL + # Use upsert to handle race conditions during parallel indexing + # ON CONFLICT (permalink, project_id) matches the partial unique index + # uix_search_index_permalink_project WHERE permalink IS NOT NULL + # For rows with NULL permalinks (observations, relations), no conflict occurs await session.execute( text(""" INSERT INTO search_index ( @@ -358,12 +423,13 @@ class PostgresSearchRepository(SearchRepositoryBase): :created_at, :updated_at, :project_id ) - ON CONFLICT (id, type, project_id) DO UPDATE SET + ON CONFLICT (permalink, project_id) WHERE permalink IS NOT NULL DO UPDATE SET + id = EXCLUDED.id, title = EXCLUDED.title, content_stems = EXCLUDED.content_stems, content_snippet = EXCLUDED.content_snippet, - permalink = EXCLUDED.permalink, file_path = EXCLUDED.file_path, + type = EXCLUDED.type, metadata = EXCLUDED.metadata, from_id = EXCLUDED.from_id, to_id = EXCLUDED.to_id, diff --git a/test-int/conftest.py b/test-int/conftest.py index 97b2fb41..0783ce21 100644 --- a/test-int/conftest.py +++ b/test-int/conftest.py @@ -125,6 +125,7 @@ async def engine_factory( CREATE_POSTGRES_SEARCH_INDEX_TABLE, CREATE_POSTGRES_SEARCH_INDEX_FTS, CREATE_POSTGRES_SEARCH_INDEX_METADATA, + CREATE_POSTGRES_SEARCH_INDEX_PERMALINK, ) from basic_memory import db @@ -160,6 +161,7 @@ async def engine_factory( await conn.execute(CREATE_POSTGRES_SEARCH_INDEX_TABLE) await conn.execute(CREATE_POSTGRES_SEARCH_INDEX_FTS) await conn.execute(CREATE_POSTGRES_SEARCH_INDEX_METADATA) + await conn.execute(CREATE_POSTGRES_SEARCH_INDEX_PERMALINK) yield engine, session_maker diff --git a/tests/conftest.py b/tests/conftest.py index fb7bb208..55bfbfbb 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -210,6 +210,7 @@ async def engine_factory( CREATE_POSTGRES_SEARCH_INDEX_TABLE, CREATE_POSTGRES_SEARCH_INDEX_FTS, CREATE_POSTGRES_SEARCH_INDEX_METADATA, + CREATE_POSTGRES_SEARCH_INDEX_PERMALINK, ) # Drop and recreate all tables for test isolation @@ -223,6 +224,7 @@ async def engine_factory( await conn.execute(CREATE_POSTGRES_SEARCH_INDEX_TABLE) await conn.execute(CREATE_POSTGRES_SEARCH_INDEX_FTS) await conn.execute(CREATE_POSTGRES_SEARCH_INDEX_METADATA) + await conn.execute(CREATE_POSTGRES_SEARCH_INDEX_PERMALINK) yield engine, session_maker diff --git a/tests/repository/test_search_repository.py b/tests/repository/test_search_repository.py index 3cb7e65f..c8a99982 100644 --- a/tests/repository/test_search_repository.py +++ b/tests/repository/test_search_repository.py @@ -138,6 +138,116 @@ async def test_index_item(search_repository, search_entity): assert results[0].project_id == search_repository.project_id +@pytest.mark.asyncio +async def test_index_item_upsert_on_duplicate_permalink(search_repository, search_entity): + """Test that indexing the same permalink twice uses upsert instead of failing. + + This tests the fix for the race condition where parallel entity indexing + could cause IntegrityError on the unique permalink constraint. + """ + # First insert + search_row1 = SearchIndexRow( + id=search_entity.id, + type=SearchItemType.ENTITY.value, + title="Original Title", + content_stems="original content", + content_snippet="Original content snippet", + permalink=search_entity.permalink, + file_path=search_entity.file_path, + entity_id=search_entity.id, + metadata={"entity_type": search_entity.entity_type}, + created_at=search_entity.created_at, + updated_at=search_entity.updated_at, + project_id=search_repository.project_id, + ) + await search_repository.index_item(search_row1) + + # Verify first insert worked + results = await search_repository.search(search_text="original") + assert len(results) == 1 + assert results[0].title == "Original Title" + + # Second insert with same permalink but different content (simulates race condition) + # This should NOT raise IntegrityError - it should upsert (update) instead + search_row2 = SearchIndexRow( + id=search_entity.id, + type=SearchItemType.ENTITY.value, + title="Updated Title", + content_stems="updated content", + content_snippet="Updated content snippet", + permalink=search_entity.permalink, # Same permalink! + file_path=search_entity.file_path, + entity_id=search_entity.id, + metadata={"entity_type": search_entity.entity_type}, + created_at=search_entity.created_at, + updated_at=search_entity.updated_at, + project_id=search_repository.project_id, + ) + # This should succeed without raising IntegrityError + await search_repository.index_item(search_row2) + + # Verify the row was updated, not duplicated + results_after = await search_repository.search(search_text="updated") + assert len(results_after) == 1 + assert results_after[0].title == "Updated Title" + + # Verify old content is gone (was replaced) + results_old = await search_repository.search(search_text="original") + assert len(results_old) == 0 + + +@pytest.mark.asyncio +async def test_bulk_index_items_upsert_on_duplicate_permalink(search_repository, search_entity): + """Test that bulk_index_items uses upsert for duplicate permalinks. + + This tests the fix for race conditions during bulk entity indexing. + """ + # First bulk insert + search_row1 = SearchIndexRow( + id=search_entity.id, + type=SearchItemType.ENTITY.value, + title="Bulk Original Title", + content_stems="bulk original content", + content_snippet="Bulk original content snippet", + permalink=search_entity.permalink, + file_path=search_entity.file_path, + entity_id=search_entity.id, + metadata={"entity_type": search_entity.entity_type}, + created_at=search_entity.created_at, + updated_at=search_entity.updated_at, + project_id=search_repository.project_id, + ) + await search_repository.bulk_index_items([search_row1]) + + # Verify first insert worked + results = await search_repository.search(search_text="bulk original") + assert len(results) == 1 + assert results[0].title == "Bulk Original Title" + + # Second bulk insert with same permalink (simulates race condition) + search_row2 = SearchIndexRow( + id=search_entity.id, + type=SearchItemType.ENTITY.value, + title="Bulk Updated Title", + content_stems="bulk updated content", + content_snippet="Bulk updated content snippet", + permalink=search_entity.permalink, # Same permalink! + file_path=search_entity.file_path, + entity_id=search_entity.id, + metadata={"entity_type": search_entity.entity_type}, + created_at=search_entity.created_at, + updated_at=search_entity.updated_at, + project_id=search_repository.project_id, + ) + # This should succeed without raising IntegrityError + await search_repository.bulk_index_items([search_row2]) + + # Verify the row was updated + results_after = await search_repository.search(search_text="bulk updated") + assert len(results_after) == 1 + assert results_after[0].title == "Bulk Updated Title" + + @pytest.mark.asyncio async def test_project_isolation( search_repository, second_project_repository, search_entity, second_entity