perf(sync): speed up single markdown file indexing (#751)

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

Author: phernandez

src/basic_memory/indexing/batch_indexer.py

@@ -186,6 +186,7 @@ async def index_markdown_file(
         new: bool | None = None,
         existing_permalink_by_path: dict[str, str | None] | None = None,
         index_search: bool = True,
+        resolve_relations: bool = True,
     ) -> IndexedEntity:
         """Index one markdown file using the same normalization and upsert path as batches."""
         if not self._is_markdown(file):
@@ -212,7 +213,12 @@ async def index_markdown_file(
         existing_permalink_by_path[file.path] = prepared.markdown.frontmatter.permalink
 
         with telemetry.span("index.markdown_file.persist", path=file.path, is_new=new):
-            persisted = await self._persist_markdown_file(prepared, is_new=new)
+            persisted = await self._persist_markdown_file(
+                prepared,
+                is_new=new,
+                resolve_relations=resolve_relations,
+                reload_entity=False,
+            )
         existing_permalink_by_path[file.path] = persisted.entity.permalink
 
         with telemetry.span(
@@ -550,6 +556,8 @@ async def _persist_markdown_file(
         prepared: _PreparedMarkdownFile,
         *,
         is_new: bool | None = None,
+        resolve_relations: bool = True,
+        reload_entity: bool = True,
     ) -> _PersistedMarkdownFile:
         existing = await self.entity_repository.get_by_file_path(
             prepared.file.path,
@@ -561,15 +569,20 @@ async def _persist_markdown_file(
             Path(prepared.file.path),
             prepared.markdown,
             is_new=is_new,
+            existing_entity=existing,
+            resolve_relations=resolve_relations,
+            reload_entity=reload_entity,
         )
         prepared = await self._reconcile_persisted_permalink(prepared, entity)
-        updated = await self.entity_repository.update(
+        metadata_updates = self._entity_metadata_updates(prepared.file, prepared.final_checksum)
+        updated = await self.entity_repository.update_fields(
             entity.id,
-            self._entity_metadata_updates(prepared.file, prepared.final_checksum),
+            metadata_updates,
         )
-        if updated is None:
+        if not updated:
             raise ValueError(f"Failed to update markdown entity metadata for {prepared.file.path}")
-        return _PersistedMarkdownFile(prepared=prepared, entity=updated)
+        self._apply_entity_metadata_updates(entity, metadata_updates)
+        return _PersistedMarkdownFile(prepared=prepared, entity=entity)
 
     async def _reconcile_persisted_permalink(
         self,
@@ -659,6 +672,11 @@ def _entity_metadata_updates(
             updates["content_type"] = file.content_type
         return updates
 
+    def _apply_entity_metadata_updates(self, entity: Entity, updates: dict[str, object]) -> None:
+        """Keep the returned entity aligned with metadata written without reload."""
+        for key, value in updates.items():
+            setattr(entity, key, value)
+
     def _is_markdown(self, file: IndexInputFile) -> bool:
         if file.content_type is not None:
             return file.content_type == "text/markdown"

src/basic_memory/repository/entity_repository.py

@@ -33,7 +33,7 @@ def __init__(self, session_maker: async_sessionmaker[AsyncSession], project_id:
         """
         super().__init__(session_maker, Entity, project_id=project_id)
 
-    async def get_by_id(self, entity_id: int) -> Optional[Entity]:  # pragma: no cover
+    async def get_by_id(self, entity_id: int, *, load_relations: bool = True) -> Optional[Entity]:
         """Get entity by numeric ID.
 
         Args:
@@ -43,6 +43,10 @@ async def get_by_id(self, entity_id: int) -> Optional[Entity]:  # pragma: no cov
             Entity if found, None otherwise
         """
         async with db.scoped_session(self.session_maker) as session:
+            if not load_relations:
+                result = await session.execute(self.select().where(Entity.id == entity_id))
+                return result.scalars().one_or_none()
+
             return await self.select_by_id(session, entity_id)
 
     async def _find_one_by_query(self, query, *, load_relations: bool) -> Optional[Entity]:

src/basic_memory/repository/repository.py

@@ -13,6 +13,7 @@
     Result,
     and_,
     delete,
+    update as sqlalchemy_update,
 )
 from sqlalchemy.engine import CursorResult
 from sqlalchemy.exc import NoResultFound
@@ -140,6 +141,20 @@ async def add_all(self, models: List[T]) -> Sequence[T]:
             # Query within same session
             return await self.select_by_ids(session, [m.id for m in models])  # pyright: ignore [reportAttributeAccessIssue]
 
+    async def add_all_no_return(self, models: List[T]) -> int:
+        """Insert models without reloading them afterward."""
+        if not models:
+            return 0
+
+        async with db.scoped_session(self.session_maker) as session:
+            for model in models:
+                self._set_project_id_if_needed(model)
+
+            session.add_all(models)
+            await session.flush()
+            logger.debug(f"Added {len(models)} {self.Model.__name__} records")
+            return len(models)
+
     def select(self, *entities: Any) -> Select:
         """Create a new SELECT statement.
 
@@ -298,6 +313,25 @@ async def update(self, entity_id: int, entity_data: dict[str, Any] | T) -> Optio
                 logger.debug(f"No {self.Model.__name__} found to update: {entity_id}")
                 return None
 
+    async def update_fields(self, entity_id: Any, entity_data: dict[str, Any]) -> bool:
+        """Update columns without reloading the model graph afterward."""
+        update_data = {k: v for k, v in entity_data.items() if k in self.valid_columns}
+        if not update_data:
+            return True
+
+        async with db.scoped_session(self.session_maker) as session:
+            conditions = [self.primary_key == entity_id]
+            if self.has_project_id and self.project_id is not None:
+                conditions.append(getattr(self.Model, "project_id") == self.project_id)
+
+            result = cast(
+                CursorResult[Any],
+                await session.execute(
+                    sqlalchemy_update(self.Model).where(and_(*conditions)).values(**update_data)
+                ),
+            )
+            return result.rowcount > 0
+
     async def delete(self, entity_id: int) -> bool:
         """Delete an entity from the database."""
         logger.debug(f"Deleting {self.Model.__name__}: {entity_id}")