perf(core): streamline entity write upserts

Signed-off-by: phernandez <paul@basicmachines.co>

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)

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:

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,

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

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",
         ],
     )