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

Signed-off-by: phernandez <paul@basicmachines.co>
This commit is contained in:
Paul Hernandez
2026-04-17 07:08:21 -05:00
committed by GitHub
parent 1c343bed66
commit c4cf0aff1e
9 changed files with 446 additions and 59 deletions
+23 -5
View File
@@ -186,6 +186,7 @@ class BatchIndexer:
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 @@ class BatchIndexer:
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 @@ class BatchIndexer:
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 @@ class BatchIndexer:
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 @@ class BatchIndexer:
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"
@@ -33,7 +33,7 @@ class EntityRepository(Repository[Entity]):
"""
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 @@ class EntityRepository(Repository[Entity]):
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]:
+34
View File
@@ -13,6 +13,7 @@ from sqlalchemy import (
Result,
and_,
delete,
update as sqlalchemy_update,
)
from sqlalchemy.engine import CursorResult
from sqlalchemy.exc import NoResultFound
@@ -140,6 +141,20 @@ class Repository[T: Base]:
# 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 @@ class Repository[T: Base]:
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}")
+122 -53
View File
@@ -1,5 +1,6 @@
"""Service for managing entities in the database."""
import asyncio
from collections.abc import Callable
from copy import deepcopy
from dataclasses import dataclass
@@ -9,8 +10,6 @@ from typing import Any, List, Optional, Sequence, Tuple, Union
import frontmatter
import yaml
from loguru import logger
from sqlalchemy import inspect as sa_inspect
from sqlalchemy.exc import IntegrityError
from basic_memory.config import ProjectConfig, BasicMemoryConfig
from basic_memory.file_utils import (
@@ -883,24 +882,23 @@ class EntityService(BaseService[EntityModel]):
"""
logger.debug(f"Updating entity and observations: {file_path}")
db_entity = existing_entity
if db_entity is not None:
state = sa_inspect(db_entity)
# Trigger: update flows can hand us an entity loaded in a different session.
# Why: clearing and rebuilding observations touches relationship state that cannot
# lazy-load from a detached ORM instance.
# Outcome: reload the canonical row before mutating observations.
if state.detached or "observations" in state.unloaded:
db_entity = await self.repository.get_by_id(db_entity.id)
if existing_entity is not None:
db_entity = await self.repository.get_by_id(
existing_entity.id,
load_relations=False,
)
else:
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 db_entity is None: # pragma: no cover
raise EntityNotFoundError(f"Entity not found for file path: {file_path}")
# Clear observations for entity
# Observations are owned by the markdown file, so re-indexing replaces the old set.
# We only need the entity id here; loading the old relationship collection is wasted work.
await self.observation_repository.delete_by_fields(entity_id=db_entity.id)
# add new observations
observations = [
Observation(
project_id=self.observation_repository.project_id,
@@ -912,10 +910,9 @@ class EntityService(BaseService[EntityModel]):
)
for obs in markdown.observations
]
await self.observation_repository.add_all(observations)
await self.observation_repository.add_all_no_return(observations)
# update values from markdown
db_entity = entity_model_from_markdown(file_path, markdown, db_entity)
self._apply_markdown_entity_fields(db_entity, file_path, markdown)
# checksum value is None == not finished with sync
db_entity.checksum = None
@@ -925,11 +922,49 @@ class EntityService(BaseService[EntityModel]):
if user_id is not None:
db_entity.last_updated_by = user_id
# update entity
return await self.repository.update(
entity_updates = {
"title": db_entity.title,
"note_type": db_entity.note_type,
"permalink": db_entity.permalink,
"file_path": db_entity.file_path,
"content_type": db_entity.content_type,
"created_at": db_entity.created_at,
"updated_at": db_entity.updated_at,
"entity_metadata": db_entity.entity_metadata,
"checksum": db_entity.checksum,
"last_updated_by": db_entity.last_updated_by,
}
updated = await self.repository.update_fields(
db_entity.id,
db_entity,
entity_updates,
)
if not updated: # pragma: no cover
raise EntityNotFoundError(f"Entity not found for file path: {file_path}")
return db_entity
def _apply_markdown_entity_fields(
self,
entity: EntityModel,
file_path: Path,
markdown: EntityMarkdown,
) -> None:
"""Apply parsed markdown scalar fields without touching ORM relationships."""
if not markdown.created or not markdown.modified: # pragma: no cover
raise ValueError("Both created and modified dates are required in markdown")
entity.title = markdown.frontmatter.title
entity.note_type = markdown.frontmatter.type
if markdown.frontmatter.permalink is not None:
entity.permalink = markdown.frontmatter.permalink
entity.file_path = file_path.as_posix()
entity.content_type = "text/markdown"
entity.created_at = markdown.created
entity.updated_at = markdown.modified
normalized_metadata = normalize_frontmatter_metadata(markdown.frontmatter.metadata or {})
entity.entity_metadata = {
key: value for key, value in normalized_metadata.items() if value is not None
}
async def upsert_entity_from_markdown(
self,
@@ -938,6 +973,8 @@ class EntityService(BaseService[EntityModel]):
*,
is_new: bool,
existing_entity: EntityModel | None = None,
resolve_relations: bool = True,
reload_entity: bool = True,
) -> EntityModel:
"""Create/update entity and relations from parsed markdown."""
if is_new:
@@ -948,13 +985,21 @@ class EntityService(BaseService[EntityModel]):
markdown,
existing_entity=existing_entity,
)
# Pass entity directly — avoids redundant get_by_file_path inside update_entity_relations
return await self.update_entity_relations(created, markdown)
# Pass the entity through so relation work does not have to rediscover the source row.
return await self.update_entity_relations(
created,
markdown,
resolve_targets=resolve_relations,
reload_entity=reload_entity,
)
async def update_entity_relations(
self,
entity: EntityModel,
markdown: EntityMarkdown,
*,
resolve_targets: bool = True,
reload_entity: bool = True,
) -> EntityModel:
"""Update relations for entity.
@@ -967,24 +1012,24 @@ class EntityService(BaseService[EntityModel]):
# Clear existing relations first
await self.relation_repository.delete_outgoing_relations_from_entity(entity_id)
# Batch resolve all relation targets in parallel
if markdown.relations:
import asyncio
# 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,
load_relations=False,
if resolve_targets:
# Exact target resolution is useful for local sync, but expensive for cloud
# one-file jobs. Cloud can write unresolved rows and let a relation repair pass
# fill in to_id later.
resolved_entities = await asyncio.gather(
*(
self.link_resolver.resolve_link(
rel.target,
strict=True,
load_relations=False,
)
for rel in markdown.relations
),
return_exceptions=True,
)
for rel in markdown.relations
]
# Execute all lookups in parallel
resolved_entities = await asyncio.gather(*lookup_tasks, return_exceptions=True)
else:
resolved_entities = [None] * len(markdown.relations)
# Process results and create relation records
relations_to_add = []
@@ -995,6 +1040,9 @@ class EntityService(BaseService[EntityModel]):
# Type narrowing: resolved is Optional[Entity] here, not Exception
target_entity = resolved # pyright: ignore [reportAssignmentType]
if target_entity is None and not resolve_targets:
target_entity = await self._resolve_deferred_self_relation(rel.target, entity)
# if the target is found, store the id
target_id = target_entity.id if target_entity else None
# if the target is found, store the title, otherwise add the target for a "forward link"
@@ -1013,25 +1061,46 @@ 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 {entity.permalink}"
)
continue
await self.relation_repository.add_all_ignore_duplicates(relations_to_add)
# Reload entity with relations via PK lookup (faster than get_by_file_path string match)
if not reload_entity:
return entity
# Reload entity with relations via PK lookup (faster than get_by_file_path string match).
reloaded = await self.repository.find_by_ids([entity_id])
return reloaded[0]
async def _resolve_deferred_self_relation(
self, target: str, entity: EntityModel
) -> EntityModel | None:
"""Resolve only self-relations that are safe to identify in deferred mode."""
clean_target = target.strip()
if clean_target.startswith("[[") and clean_target.endswith("]]"):
clean_target = clean_target[2:-2].strip()
if "|" in clean_target:
clean_target = clean_target.split("|", 1)[0].strip()
candidates = {entity.file_path}
if entity.permalink:
candidates.add(entity.permalink)
if entity.file_path.endswith(".md"):
candidates.add(entity.file_path[:-3])
if clean_target in candidates:
return entity
if clean_target != entity.title:
return None
# Title-only links are ambiguous because Basic Memory allows duplicate titles.
# Collapse them to self only when the title lookup proves this source is the sole candidate;
# otherwise leave the relation unresolved so we do not create a wrong permanent edge.
title_matches = await self.repository.get_by_title(clean_target, load_relations=False)
if len(title_matches) == 1 and title_matches[0].id == entity.id:
return entity
return None
async def edit_entity(
self,
identifier: str,
+2
View File
@@ -1067,6 +1067,7 @@ class SyncService:
*,
new: bool = True,
index_search: bool = True,
resolve_relations: bool = True,
) -> SyncedMarkdownFile:
"""Sync one markdown file and return the final canonical file state.
@@ -1115,6 +1116,7 @@ class SyncService:
),
new=new,
index_search=False,
resolve_relations=resolve_relations,
)
final_markdown_content = (
indexed.markdown_content
+60
View File
@@ -628,6 +628,66 @@ async def test_batch_indexer_index_markdown_file_rewrites_permalink_after_reposi
assert indexed.markdown_content == persisted_content
@pytest.mark.asyncio
async def test_batch_indexer_index_markdown_file_can_defer_relation_resolution(
app_config,
entity_service,
entity_repository,
relation_repository,
search_service,
file_service,
project_config,
monkeypatch,
):
await entity_service.create_entity_with_content(
EntitySchema(
title="Deferred Target",
directory="notes",
content="# Deferred Target\n",
)
)
path = "notes/deferred-source.md"
await _create_file(
project_config.home / path,
dedent(
"""
---
title: Deferred Source
type: note
---
# Deferred Source
- links_to [[Deferred Target]]
"""
),
)
resolve_link = AsyncMock(side_effect=AssertionError("relation lookup should be deferred"))
monkeypatch.setattr(entity_service.link_resolver, "resolve_link", resolve_link)
batch_indexer = _make_batch_indexer(
app_config,
entity_service,
entity_repository,
relation_repository,
search_service,
file_service,
)
await batch_indexer.index_markdown_file(
await _load_input(file_service, path),
index_search=False,
resolve_relations=False,
)
resolve_link.assert_not_awaited()
source = await entity_repository.get_by_file_path(path)
assert source is not None
assert len(source.outgoing_relations) == 1
assert source.outgoing_relations[0].to_id is None
assert source.outgoing_relations[0].to_name == "Deferred Target"
@pytest.mark.asyncio
async def test_batch_indexer_strips_frontmatter_from_search_content_when_body_is_empty(
app_config,
+35
View File
@@ -59,6 +59,19 @@ async def test_add_all(repository):
assert found.name == "Test 0"
@pytest.mark.asyncio
async def test_add_all_no_return(repository):
"""Bulk inserts can skip the follow-up reload when callers do not need rows back."""
instances = [ModelTest(id=f"test_{i}", name=f"Test {i}") for i in range(3)]
inserted = await repository.add_all_no_return(instances)
assert inserted == 3
found = await repository.find_by_id("test_0")
assert found is not None
assert found.name == "Test 0"
@pytest.mark.asyncio
async def test_bulk_create(repository):
"""Test bulk creation of entities."""
@@ -151,6 +164,28 @@ async def test_update(repository):
assert modified.name == "Updated"
@pytest.mark.asyncio
async def test_update_fields(repository):
"""Column-only updates can skip eager reloads on write-heavy paths."""
instance = ModelTest(id="test_add", name="Test Add")
await repository.add(instance)
updated = await repository.update_fields(instance.id, {"name": "Updated"})
assert updated is True
found = await repository.find_by_id(instance.id)
assert found is not None
assert found.name == "Updated"
@pytest.mark.asyncio
async def test_update_fields_not_found(repository):
"""Column-only updates still report when no row matched."""
updated = await repository.update_fields("missing", {"name": "Updated"})
assert updated is False
@pytest.mark.asyncio
async def test_update_model(repository):
"""Test finding entities modified since a timestamp."""
@@ -191,6 +191,124 @@ async def test_upsert_with_relations_uses_lightweight_exact_resolution(
]
@pytest.mark.asyncio
async def test_upsert_can_defer_relation_target_resolution(
entity_service: EntityService, monkeypatch
):
"""Cloud one-file indexing can store unresolved relation rows for later repair."""
await entity_service.create_entity(
EntitySchema(
title="Deferred Target",
directory="notes",
note_type="note",
content="# Deferred Target",
)
)
source = await entity_service.create_entity(
EntitySchema(
title="Deferred Source",
directory="notes",
note_type="note",
content="# Deferred Source",
)
)
resolve_link = AsyncMock(side_effect=AssertionError("relation lookup should be deferred"))
monkeypatch.setattr(entity_service.link_resolver, "resolve_link", resolve_link)
markdown = _make_markdown(
title="Deferred Source",
relations=[MarkdownRelation(type="links_to", target="Deferred Target")],
)
updated = await entity_service.upsert_entity_from_markdown(
Path(source.file_path),
markdown,
is_new=False,
resolve_relations=False,
)
resolve_link.assert_not_awaited()
outgoing = updated.outgoing_relations
assert len(outgoing) == 1
assert outgoing[0].to_id is None
assert outgoing[0].to_name == "Deferred Target"
@pytest.mark.asyncio
async def test_upsert_deferred_relation_resolution_keeps_self_links_resolved(
entity_service: EntityService, monkeypatch
):
"""Deferred relation mode should still resolve self-links without a target lookup."""
source = await entity_service.create_entity(
EntitySchema(
title="Deferred Self",
directory="notes",
note_type="note",
content="# Deferred Self",
)
)
resolve_link = AsyncMock(side_effect=AssertionError("relation lookup should be deferred"))
monkeypatch.setattr(entity_service.link_resolver, "resolve_link", resolve_link)
markdown = _make_markdown(
title="Deferred Self",
relations=[MarkdownRelation(type="links_to", target="Deferred Self")],
)
updated = await entity_service.upsert_entity_from_markdown(
Path(source.file_path),
markdown,
is_new=False,
resolve_relations=False,
)
resolve_link.assert_not_awaited()
outgoing = updated.outgoing_relations
assert len(outgoing) == 1
assert outgoing[0].to_id == source.id
assert outgoing[0].to_name == "Deferred Self"
@pytest.mark.asyncio
async def test_upsert_deferred_relation_resolution_does_not_guess_duplicate_titles(
entity_service: EntityService, monkeypatch
):
"""Deferred relation mode should not treat duplicate titles as guaranteed self-links."""
source = await entity_service.create_entity(
EntitySchema(
title="Duplicate Title",
directory="notes/source",
note_type="note",
content="# Duplicate Title",
)
)
await entity_service.create_entity(
EntitySchema(
title="Duplicate Title",
directory="notes/other",
note_type="note",
content="# Duplicate Title",
)
)
resolve_link = AsyncMock(side_effect=AssertionError("relation lookup should be deferred"))
monkeypatch.setattr(entity_service.link_resolver, "resolve_link", resolve_link)
markdown = _make_markdown(
title="Duplicate Title",
relations=[MarkdownRelation(type="links_to", target="Duplicate Title")],
)
updated = await entity_service.upsert_entity_from_markdown(
Path(source.file_path),
markdown,
is_new=False,
resolve_relations=False,
)
resolve_link.assert_not_awaited()
outgoing = updated.outgoing_relations
assert len(outgoing) == 1
assert outgoing[0].to_id is None
assert outgoing[0].to_name == "Duplicate Title"
# --- Correctness: full round-trip ---
+47
View File
@@ -291,6 +291,53 @@ async def test_sync_one_markdown_file_indexes_when_checksum_differs(
assert result.checksum != initial.checksum
@pytest.mark.asyncio
async def test_sync_one_markdown_file_can_defer_relation_resolution(
sync_service,
entity_service,
test_project,
monkeypatch,
):
"""Cloud callers can keep one-file sync cheap and repair relations later."""
await entity_service.create_entity_with_content(
EntitySchema(
title="Deferred Target",
directory="notes",
content="# Deferred Target\n",
)
)
_write_markdown(
Path(test_project.path),
"notes/deferred-source.md",
dedent(
"""
---
title: Deferred Source
type: note
---
# Deferred Source
- links_to [[Deferred Target]]
"""
),
)
resolve_link = AsyncMock(side_effect=AssertionError("relation lookup should be deferred"))
monkeypatch.setattr(sync_service.entity_service.link_resolver, "resolve_link", resolve_link)
result = await sync_service.sync_one_markdown_file(
"notes/deferred-source.md",
index_search=False,
resolve_relations=False,
)
resolve_link.assert_not_awaited()
assert len(result.entity.outgoing_relations) == 1
assert result.entity.outgoing_relations[0].to_id is None
assert result.entity.outgoing_relations[0].to_name == "Deferred Target"
@pytest.mark.asyncio
async def test_sync_markdown_file_remains_tuple_compatible(sync_service, test_project):
"""The legacy tuple-returning API still works for existing callers."""