Merge pull request 'Run the Obsidian sync off the event loop and delete chunks by containment' (#84) from perf/obsidian-sync-off-loop into main

blueprints/rag/logic.py

@@ -1,4 +1,6 @@
+import asyncio
 import datetime
+import json
 import logging
 import os
 import re
@@ -325,6 +327,57 @@ def _load_obsidian_index_cache() -> dict[str, float]:
         return {row[0]: row[1] for row in result if row[0] is not None}
 
 
+def _walk_vault_mtimes(obsidian_service: ObsidianService) -> dict[str, float]:
+    """Walk the vault and stat every file. Blocking; run via asyncio.to_thread."""
+    return {
+        str(md_path): md_path.stat().st_mtime
+        for md_path in obsidian_service.walk_vault()
+    }
+
+
+def _read_obsidian_document(
+    obsidian_service: ObsidianService, filepath: str, now: float
+) -> Document | None:
+    """Read and parse one vault file into a Document, or None on error."""
+    try:
+        with open(filepath, "r", encoding="utf-8") as f:
+            content = f.read()
+        parsed = obsidian_service.parse_markdown(content, filepath)
+        metadata = {
+            "source": "obsidian",
+            "filepath": parsed["filepath"],
+            "folder": str(Path(parsed["filepath"]).parent)
+            if parsed["filepath"]
+            else "",
+            "tags": parsed["tags"],
+            "created_at": parsed["metadata"].get("created_at"),
+            "indexed_at": now,
+            **{
+                k: v
+                for k, v in parsed["metadata"].items()
+                if k not in ["created_at", "created_by"]
+            },
+        }
+        return Document(
+            page_content=parsed["content"], metadata=_make_serializable(metadata)
+        )
+    except Exception as e:
+        logger.warning(f"Error reading {filepath}: {e}")
+        return None
+
+
+def _read_obsidian_documents(
+    obsidian_service: ObsidianService, filepaths: list[str], now: float
+) -> list[Document]:
+    """Read and parse a batch of vault files. Blocking; run via asyncio.to_thread."""
+    documents = []
+    for filepath in filepaths:
+        document = _read_obsidian_document(obsidian_service, filepath, now)
+        if document is not None:
+            documents.append(document)
+    return documents
+
+
 async def sync_obsidian_documents() -> dict[str, int]:
     """Incrementally sync Obsidian documents to pgvector.
 
@@ -340,69 +393,52 @@ async def sync_obsidian_documents() -> dict[str, int]:
 
     # Load cache from DB on first run
     if not _obsidian_index_cache:
-        _obsidian_index_cache = _load_obsidian_index_cache()
+        _obsidian_index_cache = await asyncio.to_thread(_load_obsidian_index_cache)
 
     # Build map of current vault files -> mtime
-    vault_files: dict[str, float] = {}
-    for md_path in obsidian_service.walk_vault():
-        vault_files[str(md_path)] = md_path.stat().st_mtime
+    vault_files = await asyncio.to_thread(_walk_vault_mtimes, obsidian_service)
 
     added = 0
     updated = 0
-    deleted = 0
 
     # Find files to add or update
     files_to_index: list[str] = []
+    stale_filepaths: list[str] = []
     for filepath, mtime in vault_files.items():
         indexed_at = _obsidian_index_cache.get(filepath)
         if indexed_at is None:
             files_to_index.append(filepath)
             added += 1
         elif mtime > indexed_at:
-            # Delete old chunks first
-            delete_documents_by_metadata("filepath", filepath)
+            stale_filepaths.append(filepath)
             files_to_index.append(filepath)
             updated += 1
 
-    # Find deleted files (in cache but not on disk)
-    for filepath in list(_obsidian_index_cache):
-        if filepath not in vault_files:
-            delete_documents_by_metadata("filepath", filepath)
-            del _obsidian_index_cache[filepath]
-            deleted += 1
+    # Find deleted files (in cache but not on disk). Don't drop them from the
+    # cache yet: if the delete below raises, we want the next sync to still
+    # know about them and retry, rather than leaving their chunks searchable
+    # forever with no record that they need cleaning up.
+    removed_filepaths = [fp for fp in _obsidian_index_cache if fp not in vault_files]
+    stale_filepaths.extend(removed_filepaths)
+    deleted = len(removed_filepaths)
+
+    # Delete old chunks for updated/removed files before re-indexing, in one
+    # transaction instead of a round-trip per file. Left uncaught: a failure
+    # here must stop the sync so the cache isn't marked clean for work that
+    # didn't happen.
+    if stale_filepaths:
+        await asyncio.to_thread(delete_documents_by_filepaths, stale_filepaths)
+
+    # Only now that the delete has succeeded is it safe to forget these files.
+    for filepath in removed_filepaths:
+        del _obsidian_index_cache[filepath]
 
     # Index new/changed files
     if files_to_index:
         now = time.time()
-        documents = []
-        for filepath in files_to_index:
-            try:
-                with open(filepath, "r", encoding="utf-8") as f:
-                    content = f.read()
-                parsed = obsidian_service.parse_markdown(content, filepath)
-                metadata = {
-                    "source": "obsidian",
-                    "filepath": parsed["filepath"],
-                    "folder": str(Path(parsed["filepath"]).parent)
-                    if parsed["filepath"]
-                    else "",
-                    "tags": parsed["tags"],
-                    "created_at": parsed["metadata"].get("created_at"),
-                    "indexed_at": now,
-                    **{
-                        k: v
-                        for k, v in parsed["metadata"].items()
-                        if k not in ["created_at", "created_by"]
-                    },
-                }
-                document = Document(
-                    page_content=parsed["content"],
-                    metadata=_make_serializable(metadata),
-                )
-                documents.append(document)
-            except Exception as e:
-                logger.warning(f"Error reading {filepath}: {e}")
-                continue
+        documents = await asyncio.to_thread(
+            _read_obsidian_documents, obsidian_service, files_to_index, now
+        )
 
         if documents:
             splits = []
@@ -507,6 +543,27 @@ def delete_all_documents():
         conn.commit()
 
 
+def _delete_by_containment(conn, collection_id, filter_dict: dict) -> None:
+    """Delete rows whose cmetadata contains filter_dict.
+
+    Uses JSONB containment (`@>`) rather than `cmetadata->>key = value` so the
+    delete can use the GIN index langchain_postgres creates on cmetadata
+    (ix_cmetadata_gin, jsonb_path_ops); `->>` text extraction can't use it.
+
+    Unlike `->>`, which coerces to text before comparing, `@>` matches on the
+    stored JSON type: the string "123" does not match the number 123. Callers
+    must pass values with the same type as what was stored (today's only
+    caller, filepath/source strings, is fine).
+    """
+    conn.execute(
+        text(
+            "DELETE FROM langchain_pg_embedding "
+            "WHERE collection_id = :cid AND cmetadata @> CAST(:filter AS jsonb)"
+        ),
+        {"cid": collection_id, "filter": json.dumps(filter_dict)},
+    )
+
+
 def delete_documents_by_metadata(key: str, value: str):
     """Delete documents matching a metadata key/value pair."""
     collection_id = _get_collection_id()
@@ -514,13 +571,25 @@ def delete_documents_by_metadata(key: str, value: str):
         return
     engine = _get_engine()
     with engine.connect() as conn:
-        conn.execute(
-            text(
-                "DELETE FROM langchain_pg_embedding "
-                "WHERE collection_id = :cid AND cmetadata->>:key = :value"
-            ),
-            {"cid": collection_id, "key": key, "value": value},
-        )
+        _delete_by_containment(conn, collection_id, {key: value})
+        conn.commit()
+
+
+def delete_documents_by_filepaths(filepaths: list[str]) -> None:
+    """Delete all chunks for several obsidian filepaths in one transaction.
+
+    One containment delete per path, committed together, so a sync pass with
+    many stale files doesn't pay a commit round-trip per file.
+    """
+    if not filepaths:
+        return
+    collection_id = _get_collection_id()
+    if not collection_id:
+        return
+    engine = _get_engine()
+    with engine.connect() as conn:
+        for filepath in filepaths:
+            _delete_by_containment(conn, collection_id, {"filepath": filepath})
         conn.commit()
 
 

tests/unit/test_obsidian_sync.py

@@ -0,0 +1,188 @@
+"""Tests for the incremental Obsidian sync: off-loop execution and delete SQL."""
+
+import json
+import os
+import threading
+from unittest.mock import patch
+
+import pytest
+
+from blueprints.rag import logic
+
+
+class FakeConnection:
+    def __init__(self, store):
+        self.store = store
+
+    def execute(self, stmt, params=None):
+        self.store.append((str(stmt), params or {}))
+
+    def commit(self):
+        pass
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc_info):
+        return False
+
+
+class FakeEngine:
+    def __init__(self):
+        self.executed: list[tuple[str, dict]] = []
+
+    def connect(self):
+        return FakeConnection(self.executed)
+
+
+class FakeVectorStore:
+    def __init__(self):
+        self.calls: list[list[str]] = []
+
+    async def aadd_documents(self, documents):
+        self.calls.append([d.page_content for d in documents])
+
+
+@pytest.fixture
+def vault_dir(tmp_path):
+    (tmp_path / "new.md").write_text("# New\nNew note body")
+    (tmp_path / "updated.md").write_text("# Updated\nUpdated note body")
+    with patch.dict(os.environ, {"OBSIDIAN_VAULT_PATH": str(tmp_path)}):
+        yield tmp_path
+
+
+@pytest.fixture
+async def synced(monkeypatch, vault_dir):
+    """Run a sync with a fake vault: one new file, one updated file, one
+    file that's been deleted off disk since the last sync."""
+    updated_path = str(vault_dir / "updated.md")
+    missing_path = str(vault_dir / "gone.md")
+
+    monkeypatch.setattr(
+        logic, "_obsidian_index_cache", {updated_path: 1.0, missing_path: 1.0}
+    )
+
+    engine = FakeEngine()
+    monkeypatch.setattr(logic, "_get_engine", lambda: engine)
+    monkeypatch.setattr(logic, "_get_collection_id", lambda: "collection-uuid")
+
+    def _unexpected_cache_load():
+        raise AssertionError("cache load should be skipped when already populated")
+
+    monkeypatch.setattr(logic, "_load_obsidian_index_cache", _unexpected_cache_load)
+
+    store = FakeVectorStore()
+    monkeypatch.setattr(logic, "_get_vector_store", lambda: store)
+
+    threads: dict[str, threading.Thread] = {}
+
+    real_walk = logic._walk_vault_mtimes
+
+    def recording_walk(service):
+        threads["walk"] = threading.current_thread()
+        return real_walk(service)
+
+    monkeypatch.setattr(logic, "_walk_vault_mtimes", recording_walk)
+
+    real_read = logic._read_obsidian_documents
+
+    def recording_read(service, filepaths, now):
+        threads["read"] = threading.current_thread()
+        return real_read(service, filepaths, now)
+
+    monkeypatch.setattr(logic, "_read_obsidian_documents", recording_read)
+
+    real_delete = logic.delete_documents_by_filepaths
+
+    def recording_delete(filepaths):
+        threads["delete"] = threading.current_thread()
+        return real_delete(filepaths)
+
+    monkeypatch.setattr(logic, "delete_documents_by_filepaths", recording_delete)
+
+    result = await logic.sync_obsidian_documents()
+    return result, threads, engine, store, updated_path, missing_path
+
+
+class TestSyncOffLoop:
+    async def test_counts_match_vault_changes(self, synced):
+        result, *_ = synced
+        assert result == {"added": 1, "updated": 1, "deleted": 1}
+
+    async def test_walk_runs_off_the_main_thread(self, synced):
+        _, threads, *_ = synced
+        assert threads["walk"] is not threading.main_thread()
+
+    async def test_reads_run_off_the_main_thread(self, synced):
+        _, threads, *_ = synced
+        assert threads["read"] is not threading.main_thread()
+
+    async def test_deletes_run_off_the_main_thread(self, synced):
+        _, threads, *_ = synced
+        assert threads["delete"] is not threading.main_thread()
+
+    async def test_embeds_new_and_updated_content(self, synced):
+        _, _, _, store, *_ = synced
+        embedded = {text for call in store.calls for text in call}
+        assert any("New note body" in t for t in embedded)
+        assert any("Updated note body" in t for t in embedded)
+
+    async def test_cache_reflects_add_update_delete(self, synced):
+        result, _, _, _, updated_path, missing_path = synced
+        assert missing_path not in logic._obsidian_index_cache
+        assert updated_path in logic._obsidian_index_cache
+
+    async def test_deletes_use_jsonb_containment_with_bound_params(self, synced):
+        _, _, engine, *_ = synced
+        deletes = [(sql, params) for sql, params in engine.executed if "DELETE" in sql]
+        assert deletes, "expected at least one delete statement"
+        for sql, params in deletes:
+            assert "cmetadata @> CAST(:filter AS jsonb)" in sql
+            assert "->>:" not in sql  # not the old, unindexed text-extraction filter
+            assert "filter" in params
+            parsed_filter = json.loads(params["filter"])
+            assert "filepath" in parsed_filter
+
+
+async def test_delete_failure_keeps_deleted_path_in_cache_and_propagates(
+    monkeypatch, vault_dir
+):
+    """If the batched delete raises, the sync must not have already forgotten
+    the deleted file: otherwise the next sync never retries the delete and
+    its chunks stay searchable forever with no record that they're stale."""
+    updated_path = str(vault_dir / "updated.md")
+    missing_path = str(vault_dir / "gone.md")
+
+    monkeypatch.setattr(
+        logic, "_obsidian_index_cache", {updated_path: 1.0, missing_path: 1.0}
+    )
+
+    def _raise_on_delete(filepaths):
+        raise RuntimeError("db unavailable")
+
+    monkeypatch.setattr(logic, "delete_documents_by_filepaths", _raise_on_delete)
+
+    store = FakeVectorStore()
+    monkeypatch.setattr(logic, "_get_vector_store", lambda: store)
+
+    with pytest.raises(RuntimeError, match="db unavailable"):
+        await logic.sync_obsidian_documents()
+
+    assert missing_path in logic._obsidian_index_cache
+    # The updated file's old chunks were never actually deleted, so its cache
+    # entry must not be bumped to look freshly indexed.
+    assert logic._obsidian_index_cache[updated_path] == 1.0
+    assert store.calls == []
+
+
+async def test_delete_documents_by_metadata_uses_containment(monkeypatch):
+    engine = FakeEngine()
+    monkeypatch.setattr(logic, "_get_engine", lambda: engine)
+    monkeypatch.setattr(logic, "_get_collection_id", lambda: "collection-uuid")
+
+    logic.delete_documents_by_metadata("source", "obsidian")
+
+    assert len(engine.executed) == 1
+    sql, params = engine.executed[0]
+    assert "cmetadata @> CAST(:filter AS jsonb)" in sql
+    assert json.loads(params["filter"]) == {"source": "obsidian"}