perf(core): streamline entity write upserts

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

Author: phernandez

src/basic_memory/repository/entity_repository.py

@@ -59,16 +59,27 @@ async def get_by_external_id(self, external_id: str) -> Optional[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 @@ async def get_by_title(self, title: str) -> Sequence[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 @@ async def find_by_permalinks(self, permalinks: List[str]) -> Sequence[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 @@ async def upsert_entity(self, entity: Entity) -> Entity:
                 session.add(entity)
                 await session.flush()
 
+                if not reload:
+                    return entity
+
                 # Return with relationships loaded
                 query = (
                     self.select()
@@ -363,13 +374,12 @@ async def upsert_entity(self, entity: Entity) -> 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 @@ async def upsert_entity(self, entity: Entity) -> 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 @@ async def create_all(self, data_list: List[dict]) -> Sequence[T]:
 
             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 @@ async def update(self, entity_id: int, entity_data: dict | T) -> Optional[T]:
                 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: