From 4e10e21a6c5e3296d40ae91e30b278eee9b853c0 Mon Sep 17 00:00:00 2001 From: phernandez Date: Sat, 4 Apr 2026 00:59:19 -0500 Subject: [PATCH] perf(core): streamline entity write upserts Signed-off-by: phernandez --- .../repository/entity_repository.py | 53 ++-- src/basic_memory/repository/repository.py | 19 +- src/basic_memory/services/entity_service.py | 230 ++++++++++++------ src/basic_memory/services/link_resolver.py | 51 +++- .../services/test_entity_service_telemetry.py | 11 +- 5 files changed, 257 insertions(+), 107 deletions(-) diff --git a/src/basic_memory/repository/entity_repository.py b/src/basic_memory/repository/entity_repository.py index 56271800..1568de3f 100644 --- a/src/basic_memory/repository/entity_repository.py +++ b/src/basic_memory/repository/entity_repository.py @@ -59,16 +59,27 @@ class EntityRepository(Repository[Entity]): ) return await self.find_one(query) - async def get_by_permalink(self, permalink: str) -> Optional[Entity]: + async def _find_one_by_query(self, query, *, load_relations: bool) -> Optional[Entity]: + """Return one entity row with optional eager loading.""" + if load_relations: + query = query.options(*self.get_load_options()) + return await self.find_one(query) + + result = await self.execute_query(query, use_query_options=False) + return result.scalars().one_or_none() + + async def get_by_permalink( + self, permalink: str, *, load_relations: bool = True + ) -> Optional[Entity]: """Get entity by permalink. Args: permalink: Unique identifier for the entity """ - query = self.select().where(Entity.permalink == permalink).options(*self.get_load_options()) - return await self.find_one(query) + query = self.select().where(Entity.permalink == permalink) + return await self._find_one_by_query(query, load_relations=load_relations) - async def get_by_title(self, title: str) -> Sequence[Entity]: + async def get_by_title(self, title: str, *, load_relations: bool = True) -> Sequence[Entity]: """Get entities by title, ordered by shortest path first. When multiple entities share the same title (in different folders), @@ -82,23 +93,20 @@ class EntityRepository(Repository[Entity]): self.select() .where(Entity.title == title) .order_by(func.length(Entity.file_path), Entity.file_path) - .options(*self.get_load_options()) ) - result = await self.execute_query(query) + result = await self.execute_query(query, use_query_options=load_relations) return list(result.scalars().all()) - async def get_by_file_path(self, file_path: Union[Path, str]) -> Optional[Entity]: + async def get_by_file_path( + self, file_path: Union[Path, str], *, load_relations: bool = True + ) -> Optional[Entity]: """Get entity by file_path. Args: file_path: Path to the entity file (will be converted to string internally) """ - query = ( - self.select() - .where(Entity.file_path == Path(file_path).as_posix()) - .options(*self.get_load_options()) - ) - return await self.find_one(query) + query = self.select().where(Entity.file_path == Path(file_path).as_posix()) + return await self._find_one_by_query(query, load_relations=load_relations) # ------------------------------------------------------------------------- # Lightweight methods for permalink resolution (no eager loading) @@ -306,7 +314,7 @@ class EntityRepository(Repository[Entity]): result = await self.execute_query(query) return list(result.scalars().all()) - async def upsert_entity(self, entity: Entity) -> Entity: + async def upsert_entity(self, entity: Entity, *, reload: bool = True) -> Entity: """Insert or update entity using simple try/catch with database-level conflict resolution. Handles file_path race conditions by checking for existing entity on IntegrityError. @@ -327,6 +335,9 @@ class EntityRepository(Repository[Entity]): session.add(entity) await session.flush() + if not reload: + return entity + # Return with relationships loaded query = ( self.select() @@ -363,13 +374,12 @@ class EntityRepository(Repository[Entity]): await session.rollback() # Re-query after rollback to get a fresh, attached entity - existing_result = await session.execute( - select(Entity) - .where( - Entity.file_path == entity.file_path, Entity.project_id == entity.project_id - ) - .options(*self.get_load_options()) + existing_query = select(Entity).where( + Entity.file_path == entity.file_path, Entity.project_id == entity.project_id ) + if reload: + existing_query = existing_query.options(*self.get_load_options()) + existing_result = await session.execute(existing_query) existing_entity = existing_result.scalar_one_or_none() if existing_entity: @@ -393,6 +403,9 @@ class EntityRepository(Repository[Entity]): await session.commit() + if not reload: + return merged_entity + # Re-query to get proper relationships loaded final_result = await session.execute( select(Entity) diff --git a/src/basic_memory/repository/repository.py b/src/basic_memory/repository/repository.py index 63cba2d2..7b99624f 100644 --- a/src/basic_memory/repository/repository.py +++ b/src/basic_memory/repository/repository.py @@ -268,8 +268,21 @@ class Repository[T: Base]: return await self.select_by_ids(session, [model.id for model in model_list]) # pyright: ignore [reportAttributeAccessIssue] - async def update(self, entity_id: int, entity_data: dict | T) -> Optional[T]: - """Update an entity with the given data.""" + async def update( + self, + entity_id: int, + entity_data: dict | T, + *, + reload: bool = True, + ) -> Optional[T]: + """Update an entity with the given data. + + Args: + entity_id: Primary key to update + entity_data: Column values or a model instance to copy from + reload: When True, re-select the entity with repository load options. + When False, return the attached row after flush/refresh. + """ logger.debug(f"Updating {self.Model.__name__} {entity_id} with data: {entity_data}") async with db.scoped_session(self.session_maker) as session: try: @@ -291,6 +304,8 @@ class Repository[T: Base]: await session.refresh(entity) # Refresh logger.debug(f"Updated {self.Model.__name__}: {entity_id}") + if not reload: + return entity return await self.select_by_id(session, entity.id) # pyright: ignore [reportAttributeAccessIssue] except NoResultFound: diff --git a/src/basic_memory/services/entity_service.py b/src/basic_memory/services/entity_service.py index 5aec0db2..f513256c 100644 --- a/src/basic_memory/services/entity_service.py +++ b/src/basic_memory/services/entity_service.py @@ -242,9 +242,17 @@ class EntityService(BaseService[EntityModel]): # Try to find existing entity using strict resolution (no fuzzy search) # This prevents incorrectly matching similar file paths like "Node A.md" and "Node C.md" - existing = await self.link_resolver.resolve_link(schema.file_path, strict=True) + existing = await self.link_resolver.resolve_link( + schema.file_path, + strict=True, + load_relations=False, + ) if not existing and schema.permalink: - existing = await self.link_resolver.resolve_link(schema.permalink, strict=True) + existing = await self.link_resolver.resolve_link( + schema.permalink, + strict=True, + load_relations=False, + ) if existing: logger.debug(f"Found existing entity: {existing.file_path}") @@ -324,17 +332,14 @@ class EntityService(BaseService[EntityModel]): action="create", phase="upsert_entity", ): - entity = await self.upsert_entity_from_markdown(file_path, entity_markdown, is_new=True) - - with telemetry.scope( - "entity_service.create.update_checksum", - domain="entity_service", - action="create", - phase="update_checksum", - ): - updated = await self.repository.update(entity.id, {"checksum": checksum}) + updated = await self.upsert_entity_from_markdown( + file_path, + entity_markdown, + is_new=True, + checksum=checksum, + ) if not updated: # pragma: no cover - raise ValueError(f"Failed to update entity checksum after create: {entity.id}") + raise ValueError(f"Failed to persist entity after create: {file_path}") return EntityWriteResult( entity=updated, content=final_content, @@ -451,18 +456,13 @@ class EntityService(BaseService[EntityModel]): phase="upsert_entity", ): entity = await self.upsert_entity_from_markdown( - file_path, entity_markdown, is_new=False + file_path, + entity_markdown, + is_new=False, + checksum=checksum, ) - - with telemetry.scope( - "entity_service.update.update_checksum", - domain="entity_service", - action="update", - phase="update_checksum", - ): - entity = await self.repository.update(entity.id, {"checksum": checksum}) if not entity: # pragma: no cover - raise ValueError(f"Failed to update entity checksum after update: {file_path}") + raise ValueError(f"Failed to persist entity after update: {file_path}") return EntityWriteResult( entity=entity, @@ -848,7 +848,7 @@ class EntityService(BaseService[EntityModel]): # Use UPSERT to handle conflicts cleanly try: - return await self.repository.upsert_entity(model) + return await self.repository.upsert_entity(model, reload=False) except Exception as e: logger.error(f"Failed to upsert entity for {file_path}: {e}") raise EntityCreationError(f"Failed to create entity: {str(e)}") from e @@ -863,7 +863,12 @@ class EntityService(BaseService[EntityModel]): """ logger.debug(f"Updating entity and observations: {file_path}") - db_entity = await self.repository.get_by_file_path(file_path.as_posix()) + db_entity = await self.repository.get_by_file_path( + file_path.as_posix(), + load_relations=False, + ) + if not db_entity: # pragma: no cover + raise EntityNotFoundError(f"Entity not found: {file_path}") # Clear observations for entity await self.observation_repository.delete_by_fields(entity_id=db_entity.id) @@ -880,23 +885,37 @@ class EntityService(BaseService[EntityModel]): ) for obs in markdown.observations ] - await self.observation_repository.add_all(observations) + if observations: + await self.observation_repository.add_all(observations) - # update values from markdown - db_entity = entity_model_from_markdown(file_path, markdown, db_entity) + # Trigger: the lightweight lookup above returns a detached row without loaded collections + # Why: assigning a new observation list onto that detached ORM object would trigger lazy loads + # Outcome: rebuild a fresh model from markdown, then copy over stable identity fields + db_entity_data = entity_model_from_markdown( + file_path, + markdown, + project_id=self.repository.project_id, + ) + db_entity_data.id = db_entity.id + db_entity_data.project_id = db_entity.project_id + db_entity_data.external_id = db_entity.external_id + db_entity_data.created_by = db_entity.created_by # checksum value is None == not finished with sync - db_entity.checksum = None + db_entity_data.checksum = None # Set last_updated_by for cloud usage (preserve existing created_by) user_id = self.get_user_id() if user_id is not None: - db_entity.last_updated_by = user_id + db_entity_data.last_updated_by = user_id + else: + db_entity_data.last_updated_by = db_entity.last_updated_by # update entity return await self.repository.update( db_entity.id, - db_entity, + db_entity_data, + reload=False, ) async def upsert_entity_from_markdown( @@ -905,26 +924,76 @@ class EntityService(BaseService[EntityModel]): markdown: EntityMarkdown, *, is_new: bool, + checksum: Optional[str] = None, ) -> EntityModel: """Create/update entity and relations from parsed markdown.""" - if is_new: - created = await self.create_entity_from_markdown(file_path, markdown) - else: - created = await self.update_entity_and_observations(file_path, markdown) - return await self.update_entity_relations(created.file_path, markdown) + # --- Base Entity Row --- + # Trigger: writes rebuild the entity row before touching relation edges + # Why: relations need a stable source entity ID, but not a fully hydrated graph + # Outcome: create/update the row with a lightweight return value + with telemetry.scope( + "entity_service.upsert.base_entity", + domain="entity_service", + action="upsert", + phase="base_entity", + ): + if is_new: + created = await self.create_entity_from_markdown(file_path, markdown) + else: + created = await self.update_entity_and_observations(file_path, markdown) + + # --- Relation Edges --- + with telemetry.scope( + "entity_service.upsert.relations", + domain="entity_service", + action="upsert", + phase="relations", + ): + await self.update_entity_relations(created, markdown) + + # --- Final Entity State --- + # Trigger: create/update/edit already computed the final file checksum + # Why: fold the checksum write into the upsert flow so callers do one hydrated read + # Outcome: the write path returns the final entity state without an extra checksum step + if checksum is not None: + with telemetry.scope( + "entity_service.upsert.persist_checksum", + domain="entity_service", + action="upsert", + phase="persist_checksum", + ): + updated = await self.repository.update(created.id, {"checksum": checksum}) + if not updated: # pragma: no cover + raise ValueError(f"Failed to update entity checksum after upsert: {file_path}") + return updated + + with telemetry.scope( + "entity_service.upsert.hydrate_entity", + domain="entity_service", + action="upsert", + phase="hydrate_entity", + ): + hydrated = await self.repository.get_by_file_path(created.file_path) + if not hydrated: # pragma: no cover + raise EntityNotFoundError(f"Entity not found after upsert: {created.file_path}") + return hydrated async def update_entity_relations( self, - path: str, + db_entity: EntityModel, markdown: EntityMarkdown, - ) -> EntityModel: + ) -> None: """Update relations for entity""" - logger.debug(f"Updating relations for entity: {path}") - - db_entity = await self.repository.get_by_file_path(path) + logger.debug(f"Updating relations for entity: {db_entity.file_path}") # Clear existing relations first - await self.relation_repository.delete_outgoing_relations_from_entity(db_entity.id) + with telemetry.scope( + "entity_service.upsert.delete_relations", + domain="entity_service", + action="upsert", + phase="delete_relations", + ): + await self.relation_repository.delete_outgoing_relations_from_entity(db_entity.id) # Batch resolve all relation targets in parallel if markdown.relations: @@ -933,13 +1002,23 @@ class EntityService(BaseService[EntityModel]): # Create tasks for all relation lookups # Use strict=True to disable fuzzy search - only exact matches should create resolved relations # This ensures forward references (links to non-existent entities) remain unresolved (to_id=NULL) - lookup_tasks = [ - self.link_resolver.resolve_link(rel.target, strict=True) - for rel in markdown.relations - ] + with telemetry.scope( + "entity_service.upsert.resolve_relation_targets", + domain="entity_service", + action="upsert", + phase="resolve_relation_targets", + ): + lookup_tasks = [ + self.link_resolver.resolve_link( + rel.target, + strict=True, + load_relations=False, + ) + for rel in markdown.relations + ] - # Execute all lookups in parallel - resolved_entities = await asyncio.gather(*lookup_tasks, return_exceptions=True) + # Execute all lookups in parallel + resolved_entities = await asyncio.gather(*lookup_tasks, return_exceptions=True) # Process results and create relation records relations_to_add = [] @@ -968,22 +1047,26 @@ class EntityService(BaseService[EntityModel]): # Batch insert all relations if relations_to_add: - try: - await self.relation_repository.add_all(relations_to_add) - except IntegrityError: - # Some relations might be duplicates - fall back to individual inserts - logger.debug("Batch relation insert failed, trying individual inserts") - for relation in relations_to_add: - try: - await self.relation_repository.add(relation) - except IntegrityError: - # Unique constraint violation - relation already exists - logger.debug( - f"Skipping duplicate relation {relation.relation_type} from {db_entity.permalink}" - ) - continue - - return await self.repository.get_by_file_path(path) + with telemetry.scope( + "entity_service.upsert.insert_relations", + domain="entity_service", + action="upsert", + phase="insert_relations", + ): + try: + await self.relation_repository.add_all(relations_to_add) + except IntegrityError: + # Some relations might be duplicates - fall back to individual inserts + logger.debug("Batch relation insert failed, trying individual inserts") + for relation in relations_to_add: + try: + await self.relation_repository.add(relation) + except IntegrityError: + # Unique constraint violation - relation already exists + logger.debug( + f"Skipping duplicate relation {relation.relation_type} from {db_entity.permalink}" + ) + continue async def edit_entity( self, @@ -1040,7 +1123,11 @@ class EntityService(BaseService[EntityModel]): action="edit", phase="resolve_entity", ): - entity = await self.link_resolver.resolve_link(identifier, strict=True) + entity = await self.link_resolver.resolve_link( + identifier, + strict=True, + load_relations=False, + ) if not entity: raise EntityNotFoundError(f"Entity not found: {identifier}") @@ -1089,18 +1176,13 @@ class EntityService(BaseService[EntityModel]): phase="upsert_entity", ): entity = await self.upsert_entity_from_markdown( - file_path, entity_markdown, is_new=False + file_path, + entity_markdown, + is_new=False, + checksum=checksum, ) - - with telemetry.scope( - "entity_service.edit.update_checksum", - domain="entity_service", - action="edit", - phase="update_checksum", - ): - entity = await self.repository.update(entity.id, {"checksum": checksum}) if not entity: # pragma: no cover - raise ValueError(f"Failed to update entity checksum after edit: {file_path}") + raise ValueError(f"Failed to persist entity after edit: {file_path}") return EntityWriteResult( entity=entity, diff --git a/src/basic_memory/services/link_resolver.py b/src/basic_memory/services/link_resolver.py index 739e0087..f0f95ad6 100644 --- a/src/basic_memory/services/link_resolver.py +++ b/src/basic_memory/services/link_resolver.py @@ -47,6 +47,7 @@ class LinkResolver: use_search: bool = True, strict: bool = False, source_path: Optional[str] = None, + load_relations: bool = True, ) -> Optional[Entity]: """Resolve a markdown link to a permalink. @@ -56,6 +57,7 @@ class LinkResolver: strict: If True, only exact matches are allowed (no fuzzy search fallback) source_path: Optional path of the source file containing the link. Used to prefer notes closer to the source (context-aware resolution). + load_relations: When False, skip eager loading and return a lightweight entity row. """ logger.trace(f"Resolving link: {link_text} (source: {source_path})") @@ -98,6 +100,7 @@ class LinkResolver: strict=strict, source_path=None, project_permalink=project.permalink, + load_relations=load_relations, ) current_project_permalink = await self._get_current_project_permalink() @@ -109,6 +112,7 @@ class LinkResolver: strict=strict, source_path=source_path, project_permalink=current_project_permalink, + load_relations=load_relations, ) if resolved: return resolved @@ -136,6 +140,7 @@ class LinkResolver: strict=strict, source_path=None, project_permalink=project.permalink, + load_relations=load_relations, ) def _normalize_link_text(self, link_text: str) -> Tuple[str, Optional[str]]: @@ -176,6 +181,7 @@ class LinkResolver: strict: bool, source_path: Optional[str], project_permalink: Optional[str], + load_relations: bool, ) -> Optional[Entity]: """Resolve a link within a specific project scope.""" clean_text = link_text @@ -223,12 +229,18 @@ class LinkResolver: # Try with .md extension if not relative_path.endswith(".md"): relative_path_md = f"{relative_path}.md" - entity = await entity_repository.get_by_file_path(relative_path_md) + entity = await entity_repository.get_by_file_path( + relative_path_md, + load_relations=load_relations, + ) if entity: return entity # Try as-is (already has extension or is a permalink) - entity = await entity_repository.get_by_file_path(relative_path) + entity = await entity_repository.get_by_file_path( + relative_path, + load_relations=load_relations, + ) if entity: return entity @@ -242,12 +254,18 @@ class LinkResolver: # Check permalink match for candidate_permalink in permalink_candidates: - permalink_entity = await entity_repository.get_by_permalink(candidate_permalink) + permalink_entity = await entity_repository.get_by_permalink( + candidate_permalink, + load_relations=load_relations, + ) if permalink_entity and permalink_entity.id not in [c.id for c in candidates]: candidates.append(permalink_entity) # Check title matches - title_entities = await entity_repository.get_by_title(clean_text) + title_entities = await entity_repository.get_by_title( + clean_text, + load_relations=load_relations, + ) for entity in title_entities: # Avoid duplicates (permalink match might also be in title matches) if entity.id not in [c.id for c in candidates]: @@ -263,13 +281,19 @@ class LinkResolver: # Standard resolution (no source context): permalink first, then title # 1. Try exact permalink match first (most efficient) for candidate_permalink in permalink_candidates: - entity = await entity_repository.get_by_permalink(candidate_permalink) + entity = await entity_repository.get_by_permalink( + candidate_permalink, + load_relations=load_relations, + ) if entity: logger.debug(f"Found exact permalink match: {entity.permalink}") return entity # 2. Try exact title match - found = await entity_repository.get_by_title(clean_text) + found = await entity_repository.get_by_title( + clean_text, + load_relations=load_relations, + ) if found: # Return first match (shortest path) if no source context entity = found[0] @@ -277,7 +301,10 @@ class LinkResolver: return entity # 3. Try file path - found_path = await entity_repository.get_by_file_path(clean_text) + found_path = await entity_repository.get_by_file_path( + clean_text, + load_relations=load_relations, + ) if found_path: logger.debug(f"Found entity with path: {found_path.file_path}") return found_path @@ -285,7 +312,10 @@ class LinkResolver: # 4. Try file path with .md extension if not already present if not clean_text.endswith(".md") and "/" in clean_text: file_path_with_md = f"{clean_text}.md" - found_path_md = await entity_repository.get_by_file_path(file_path_with_md) + found_path_md = await entity_repository.get_by_file_path( + file_path_with_md, + load_relations=load_relations, + ) if found_path_md: logger.debug(f"Found entity with path (with .md): {found_path_md.file_path}") return found_path_md @@ -309,7 +339,10 @@ class LinkResolver: f"Selected best match from {len(results)} results: {best_match.permalink}" ) if best_match.permalink: - return await entity_repository.get_by_permalink(best_match.permalink) + return await entity_repository.get_by_permalink( + best_match.permalink, + load_relations=load_relations, + ) # if we couldn't find anything then return None return None diff --git a/tests/services/test_entity_service_telemetry.py b/tests/services/test_entity_service_telemetry.py index e13dfa44..a2be315a 100644 --- a/tests/services/test_entity_service_telemetry.py +++ b/tests/services/test_entity_service_telemetry.py @@ -54,7 +54,9 @@ async def test_create_entity_emits_expected_phase_spans(entity_service, monkeypa "file_service.write", "entity_service.create.parse_markdown", "entity_service.create.upsert_entity", - "entity_service.create.update_checksum", + "entity_service.upsert.base_entity", + "entity_service.upsert.relations", + "entity_service.upsert.persist_checksum", ], ) @@ -93,7 +95,9 @@ async def test_edit_entity_emits_expected_phase_spans(entity_service, monkeypatc "file_service.write", "entity_service.edit.parse_markdown", "entity_service.edit.upsert_entity", - "entity_service.edit.update_checksum", + "entity_service.upsert.base_entity", + "entity_service.upsert.relations", + "entity_service.upsert.persist_checksum", ], ) @@ -124,6 +128,9 @@ async def test_reindex_entity_emits_expected_phase_spans(entity_service, monkeyp "file_service.read_content", "entity_service.reindex.parse_markdown", "entity_service.reindex.upsert_entity", + "entity_service.upsert.base_entity", + "entity_service.upsert.relations", + "entity_service.upsert.hydrate_entity", "entity_service.reindex.update_checksum", ], )