diff --git a/src/basic_memory/repository/entity_repository.py b/src/basic_memory/repository/entity_repository.py index 1159e280..9fbfb72f 100644 --- a/src/basic_memory/repository/entity_repository.py +++ b/src/basic_memory/repository/entity_repository.py @@ -8,6 +8,7 @@ from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from sqlalchemy.orm import selectinload from sqlalchemy.orm.interfaces import LoaderOption +from sqlalchemy.dialects.sqlite import insert as sqlite_insert from basic_memory import db from basic_memory.models.knowledge import Entity, Observation, Relation @@ -252,3 +253,159 @@ class EntityRepository(Repository[Entity]): if not found: # pragma: no cover raise RuntimeError(f"Failed to retrieve entity after insert: {entity.file_path}") return found + + async def upsert_entity_optimistic( + self, entity_data: dict, conflict_columns: Optional[List[str]] = None + ) -> Entity: + """Insert or update entity using SQLite's ON CONFLICT clause for optimistic concurrency. + + This method uses SQLite's native ON CONFLICT semantics to handle race conditions + efficiently without manual exception handling. + + Args: + entity_data: Dictionary of entity data to insert/update + conflict_columns: Columns to use for conflict detection (defaults to ['file_path']) + + Returns: + The inserted or updated entity + """ + if conflict_columns is None: + conflict_columns = ["file_path"] + + async with db.scoped_session(self.session_maker) as session: + # Ensure project_id is set + if ( + self.project_id is not None + and "project_id" not in entity_data + and "project_id" in self.valid_columns + ): + entity_data["project_id"] = self.project_id + + # Create the upsert statement using SQLite's ON CONFLICT + stmt = sqlite_insert(Entity).values(entity_data) + stmt = stmt.on_conflict_do_update( + index_elements=conflict_columns, + set_={ + key: stmt.excluded[key] + for key in entity_data.keys() + if key not in conflict_columns and key != "id" + }, + ) + + await session.execute(stmt) + await session.flush() + + # Retrieve the entity with relationships loaded + file_path = entity_data.get("file_path") + if not file_path: + raise ValueError("file_path is required for upsert_entity_optimistic") + + query = ( + self.select() + .where(Entity.file_path == file_path) + .options(*self.get_load_options()) + ) + result = await session.execute(query) + found = result.scalar_one_or_none() + + if not found: # pragma: no cover + raise RuntimeError(f"Failed to retrieve entity after optimistic upsert: {file_path}") + + return found + + async def update_entity_optimistic( + self, file_path: str, updates: dict, conflict_columns: Optional[List[str]] = None + ) -> Optional[Entity]: + """Update an entity by file_path using SQLite's ON CONFLICT clause. + + This method is designed to replace direct update() calls in sync operations + to eliminate race conditions. + + Args: + file_path: The file_path of the entity to update + updates: Dictionary of fields to update + conflict_columns: Columns to use for conflict detection (defaults to ['file_path']) + + Returns: + The updated entity or None if no entity exists with the given file_path + """ + if conflict_columns is None: + conflict_columns = ["file_path"] + + async with db.scoped_session(self.session_maker) as session: + # First check if entity exists + existing_entity = await session.execute( + select(Entity).where( + Entity.file_path == file_path, + Entity.project_id == self.project_id + ) + ) + entity = existing_entity.scalar_one_or_none() + + if not entity: + # Entity doesn't exist, cannot update + return None + + # Prepare the data for upsert (include the file_path and any other required fields) + entity_data = { + "file_path": file_path, + "project_id": self.project_id, + **updates + } + + # For updates, we need to include the current values that shouldn't change + # This ensures we don't lose data during the upsert + for field in ["title", "entity_type", "content_type", "permalink", "checksum", "created_at", "updated_at"]: + if field not in entity_data and hasattr(entity, field): + current_value = getattr(entity, field) + if current_value is not None: + entity_data[field] = current_value + + # Use the optimistic upsert method + return await self.upsert_entity_optimistic(entity_data, conflict_columns) + + async def update_entity_by_id_optimistic( + self, entity_id: int, updates: dict + ) -> Optional[Entity]: + """Update an entity by ID using current values + updates, then upsert optimistically. + + This method is useful when you have an entity ID and want to update specific fields + while preserving all other current values. + + Args: + entity_id: The ID of the entity to update + updates: Dictionary of fields to update + + Returns: + The updated entity or None if no entity exists with the given ID + """ + async with db.scoped_session(self.session_maker) as session: + # First get the current entity + existing_entity = await session.execute( + select(Entity).where(Entity.id == entity_id) + ) + entity = existing_entity.scalar_one_or_none() + + if not entity: + # Entity doesn't exist, cannot update + return None + + # Prepare the data for upsert with current values + updates + entity_data = { + "file_path": entity.file_path, + "project_id": entity.project_id, + "title": entity.title, + "entity_type": entity.entity_type, + "content_type": entity.content_type, + "permalink": entity.permalink, + "checksum": entity.checksum, + "created_at": entity.created_at, + "updated_at": entity.updated_at, + **updates # Override with new values + } + + # For moves, we need to handle the file_path conflict properly + conflict_columns = ["file_path"] if "file_path" in updates else ["file_path"] + + # Use the optimistic upsert method + return await self.upsert_entity_optimistic(entity_data, conflict_columns) diff --git a/src/basic_memory/sync/sync_service.py b/src/basic_memory/sync/sync_service.py index 96e6f792..77bdae7b 100644 --- a/src/basic_memory/sync/sync_service.py +++ b/src/basic_memory/sync/sync_service.py @@ -389,8 +389,8 @@ class SyncService: logger.error(f"Entity not found after constraint violation, path={path}") raise ValueError(f"Entity not found after constraint violation: {path}") - updated = await self.entity_repository.update( - entity.id, {"file_path": path, "checksum": checksum} + updated = await self.entity_repository.update_entity_optimistic( + path, {"checksum": checksum} ) if updated is None: # pragma: no cover @@ -407,8 +407,8 @@ class SyncService: logger.error(f"Entity not found for existing file, path={path}") raise ValueError(f"Entity not found for existing file: {path}") - updated = await self.entity_repository.update( - entity.id, {"file_path": path, "checksum": checksum} + updated = await self.entity_repository.update_entity_optimistic( + path, {"checksum": checksum} ) if updated is None: # pragma: no cover @@ -477,7 +477,7 @@ class SyncService: f"new_checksum={new_checksum}" ) - updated = await self.entity_repository.update(entity.id, updates) + updated = await self.entity_repository.update_entity_by_id_optimistic(entity.id, updates) if updated is None: # pragma: no cover logger.error(