admin/basicmachines-co-basic-memory
1
0
Fork 0
mirror of https://github.com/basicmachines-co/basic-memory synced 2026-06-21 13:47:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: use upsert to prevent IntegrityError during parallel search indexing

Browse Source 
Replace delete-then-insert pattern with INSERT ... ON CONFLICT for
PostgreSQL search index operations. This fixes race conditions where
parallel entity indexing could cause UniqueViolationError on the
uix_search_index_permalink_project constraint.

Changes:
- Add index_item() override in PostgresSearchRepository with upsert
- Update bulk_index_items() to use ON CONFLICT (permalink, project_id)
- Add CREATE_POSTGRES_SEARCH_INDEX_PERMALINK DDL for test fixtures
- Add tests for upsert behavior on duplicate permalinks

Technical note: Use column-based ON CONFLICT syntax instead of
ON CONSTRAINT (which only works for table constraints, not indexes).

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: phernandez <paul@basicmachines.co>
This commit is contained in:
phernandez
2025-12-30 21:55:43 -06:00
parent eb7fbaf0bf
commit 4ce21984a4
5 changed files with 200 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/basic_memory/models/search.py
+9
View File
@@ -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("""
src/basic_memory/repository/postgres_search_repository.py
+77 -11
View File
@@ -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,
test-int/conftest.py
+2
View File
@@ -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
tests/conftest.py
+2
View File
@@ -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
tests/repository/test_search_repository.py
+110
View File
@@ -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