From 141e16fb2b6898c6068cef7ecbee9530798b419e Mon Sep 17 00:00:00 2001 From: Grace Gettert Date: Thu, 6 Aug 2026 16:33:30 +0000 Subject: [PATCH] fix(mine): harden source adapter dispatch --- mempalace/cli.py | 115 ++++++++++++------- tests/test_cli_source_adapters.py | 182 ++++++++++++++++++++++++++++-- 2 files changed, 246 insertions(+), 51 deletions(-) diff --git a/mempalace/cli.py b/mempalace/cli.py index 77c2781..b7eda8c 100644 --- a/mempalace/cli.py +++ b/mempalace/cli.py @@ -31,10 +31,12 @@ Examples: mempalace search "pricing discussion" --wing my_app --room costs """ -import os -import sys -import shlex import argparse +import contextlib +import os +import shlex +import sys +import warnings from pathlib import Path from .config import MempalaceConfig @@ -679,14 +681,14 @@ class UnsupportedSourceAdapterProtocolError(ValueError): class _DryRunCollectionProxy: - """Read-through collection facade that records, but never persists, writes. + """Read-only collection facade that records, but never persists, writes. Source adapters are deliberately allowed to access ``drawer_collection`` directly. Passing the live collection during a dry run would therefore make ``--dry-run`` advisory rather than safe. """ - def __init__(self, collection): + def __init__(self, collection=None): self._collection = collection self.operations = [] @@ -703,12 +705,18 @@ class _DryRunCollectionProxy: self.operations.append(("update", kwargs)) def query(self, **kwargs): + if self._collection is None: + return {"ids": [[]], "documents": [[]], "metadatas": [[]], "distances": [[]]} return self._collection.query(**kwargs) def get(self, **kwargs): + if self._collection is None: + return {"ids": [], "documents": [], "metadatas": []} return self._collection.get(**kwargs) def count(self): + if self._collection is None: + return 0 return self._collection.count() @@ -745,7 +753,7 @@ def mine_source_adapter( and ``--mode`` calls must retain their established dispatch paths. """ from .knowledge_graph import KnowledgeGraph - from .palace import get_collection + from .palace import get_collection, mine_palace_lock from .sources import ( DrawerRecord, PalaceContext, @@ -770,41 +778,68 @@ def mine_source_adapter( "mempalace mine does not support yet" ) - knowledge_graph = None - try: - drawer_collection = get_collection(palace_path) - if dry_run: - drawer_collection = _DryRunCollectionProxy(drawer_collection) - knowledge_graph = _DryRunKnowledgeGraphProxy() - else: - knowledge_graph = KnowledgeGraph( - db_path=os.path.join(palace_path, "knowledge_graph.sqlite3") - ) - context = PalaceContext( - drawer_collection=drawer_collection, - knowledge_graph=knowledge_graph, - palace_path=palace_path, - config=MempalaceConfig(palace_path=palace_path), - adapter_name=adapter.name, - adapter_version=adapter.adapter_version, - ) - drawers_written = 0 - for result in adapter.ingest( - source=SourceRef(local_path=source_path), - palace=context, - ): - if isinstance(result, SourceItemMetadata): - raise UnsupportedSourceAdapterProtocolError( - f"source adapter {adapter_name!r} yielded incremental item metadata, " - "which mempalace mine does not support yet" + # A dry run must never create a collection: opening a fresh Chroma palace + # can create storage and persist embedder identity before write proxies are + # installed. It may safely retain a read-only view of an existing + # collection so adapters can produce an accurate preview. Non-dry runs + # hold one writer lease from handle creation through adapter iteration, + # including direct KG mutations by adapters. + lock = mine_palace_lock(palace_path) if not dry_run else contextlib.nullcontext() + with lock: + knowledge_graph = None + try: + if dry_run: + try: + drawer_collection = _DryRunCollectionProxy( + get_collection(palace_path, create=False) + ) + except FileNotFoundError: + # No palace or collection has been initialized yet. Use + # an empty recording facade rather than materializing one. + drawer_collection = _DryRunCollectionProxy() + knowledge_graph = _DryRunKnowledgeGraphProxy() + else: + drawer_collection = get_collection(palace_path) + knowledge_graph = KnowledgeGraph( + db_path=os.path.join(palace_path, "knowledge_graph.sqlite3") ) - if isinstance(result, DrawerRecord): - drawers_written += 1 - context.upsert_drawer(result) - return drawers_written - finally: - if knowledge_graph is not None and hasattr(knowledge_graph, "close"): - knowledge_graph.close() + context = PalaceContext( + drawer_collection=drawer_collection, + knowledge_graph=knowledge_graph, + palace_path=palace_path, + config=MempalaceConfig(palace_path=palace_path), + adapter_name=adapter.name, + adapter_version=adapter.adapter_version, + ) + drawers_written = 0 + for result in adapter.ingest( + source=SourceRef(local_path=source_path), + palace=context, + ): + if isinstance(result, SourceItemMetadata): + # Non-incremental adapters may report a cursor or version + # while still doing a complete re-extract. Incremental + # adapters are rejected before ingest above, so accepting + # this avoids a late partial-ingest failure. + warnings.warn( + f"Source adapter {adapter_name!r} yielded non-incremental item " + "metadata; ignoring it during complete ingest", + RuntimeWarning, + stacklevel=2, + ) + continue + if isinstance(result, DrawerRecord): + drawers_written += 1 + context.upsert_drawer(result) + continue + raise TypeError( + f"source adapter {adapter_name!r} yielded unsupported result type " + f"{type(result).__name__}" + ) + return drawers_written + finally: + if knowledge_graph is not None and hasattr(knowledge_graph, "close"): + knowledge_graph.close() def cmd_sweep(args): diff --git a/tests/test_cli_source_adapters.py b/tests/test_cli_source_adapters.py index b0e8eaf..b5c3a4e 100644 --- a/tests/test_cli_source_adapters.py +++ b/tests/test_cli_source_adapters.py @@ -1,7 +1,11 @@ """CLI coverage for explicit RFC 002 source-adapter dispatch (#2062).""" import argparse +import contextlib +import multiprocessing +import os import sys +import time import pytest @@ -82,6 +86,18 @@ class _DirectMutationAdapter(BaseSourceAdapter): return AdapterSchema(version="1.0", fields={}) +class _ReadAwareAdapter(BaseSourceAdapter): + name = "read-aware" + observed_count = None + + def ingest(self, *, source, palace): + self.__class__.observed_count = palace.drawer_collection.count() + yield DrawerRecord(content="fixture content", source_file="fixture://record") + + def describe_schema(self): + return AdapterSchema(version="1.0", fields={}) + + class _IncrementalAdapter(BaseSourceAdapter): name = "incremental" capabilities = frozenset({"supports_incremental"}) @@ -97,16 +113,30 @@ class _MetadataAdapter(BaseSourceAdapter): name = "metadata" def ingest(self, *, source, palace): + yield DrawerRecord(content="before metadata", source_file="fixture://record") yield SourceItemMetadata(source_file="fixture://record", version="v1") def describe_schema(self): return AdapterSchema(version="1.0", fields={}) +def _hold_palace_lock(palace_path, ready_flag, release_flag): + """Hold a writer lease in a separate process for contention coverage.""" + from mempalace.palace import mine_palace_lock + + with mine_palace_lock(palace_path): + open(ready_flag, "w").close() + for _ in range(500): + if os.path.exists(release_flag): + return + time.sleep(0.01) + + @pytest.fixture(autouse=True) def _isolated_fixture_adapter(): _FixtureAdapter.instances.clear() _DirectMutationAdapter.instances.clear() + _ReadAwareAdapter.observed_count = None _FakeKnowledgeGraph.instances.clear() reset_adapters() try: @@ -170,10 +200,13 @@ def test_cmd_mine_source_rejects_unknown_adapter(capsys): def test_mine_source_dry_run_prevents_direct_collection_and_kg_mutations(monkeypatch): from mempalace import knowledge_graph, palace - collection = _FakeCollection() register("direct-mutation", _DirectMutationAdapter) monkeypatch.setattr(cli, "MempalaceConfig", _FakeConfig) - monkeypatch.setattr(palace, "get_collection", lambda palace_path: collection) + def read_only_collection(_palace_path, *, create): + assert create is False + raise FileNotFoundError + + monkeypatch.setattr(palace, "get_collection", read_only_collection) monkeypatch.setattr(knowledge_graph, "KnowledgeGraph", _FakeKnowledgeGraph) drawers_written = cli.mine_source_adapter( @@ -185,7 +218,6 @@ def test_mine_source_dry_run_prevents_direct_collection_and_kg_mutations(monkeyp adapter = _DirectMutationAdapter.instances[0] assert drawers_written == 1 - assert collection.upserts == [] assert _FakeKnowledgeGraph.instances == [] assert [operation[0] for operation in adapter.palace.drawer_collection.operations] == [ "upsert", @@ -196,6 +228,43 @@ def test_mine_source_dry_run_prevents_direct_collection_and_kg_mutations(monkeyp ] +def test_mine_source_dry_run_reads_existing_collection_without_writing(monkeypatch): + from mempalace import knowledge_graph, palace + + collection = _FakeCollection() + collection.count = lambda: 7 + register("read-aware", _ReadAwareAdapter) + monkeypatch.setattr(cli, "MempalaceConfig", _FakeConfig) + monkeypatch.setattr( + palace, + "get_collection", + lambda palace_path, *, create: collection if create is False else pytest.fail("must not create"), + ) + monkeypatch.setattr(knowledge_graph, "KnowledgeGraph", _FakeKnowledgeGraph) + + assert cli.mine_source_adapter( + source_name="read-aware", source_path="/source", palace_path="/fake/palace", dry_run=True + ) == 1 + + assert _ReadAwareAdapter.observed_count == 7 + assert collection.upserts == [] + + +def test_mine_source_dry_run_fresh_palace_creates_no_backend_artifacts(tmp_path): + """Dry-running a source adapter must not materialize a new palace.""" + register("fixture", _FixtureAdapter) + palace = tmp_path / "fresh-palace" + + assert cli.mine_source_adapter( + source_name="fixture", + source_path="/source", + palace_path=str(palace), + dry_run=True, + ) == 1 + + assert not palace.exists(), "dry run must not create backend storage" + + def test_mine_source_rejects_incremental_adapter_before_ingest(): register("incremental", _IncrementalAdapter) @@ -207,20 +276,111 @@ def test_mine_source_rejects_incremental_adapter_before_ingest(): ) -def test_mine_source_rejects_incremental_metadata(monkeypatch): +def test_mine_source_accepts_non_incremental_metadata(monkeypatch, recwarn): from mempalace import knowledge_graph, palace register("metadata", _MetadataAdapter) monkeypatch.setattr(cli, "MempalaceConfig", _FakeConfig) - monkeypatch.setattr(palace, "get_collection", lambda palace_path: _FakeCollection()) + collection = _FakeCollection() + monkeypatch.setattr(palace, "get_collection", lambda palace_path: collection) monkeypatch.setattr(knowledge_graph, "KnowledgeGraph", _FakeKnowledgeGraph) - with pytest.raises(cli.UnsupportedSourceAdapterProtocolError, match="item metadata"): - cli.mine_source_adapter( - source_name="metadata", - source_path="/source", - palace_path="/fake/palace", - ) + drawers_written = cli.mine_source_adapter( + source_name="metadata", + source_path="/source", + palace_path="/fake/palace", + ) + + assert drawers_written == 1 + assert collection.upserts[0]["documents"] == ["before metadata"] + assert "non-incremental item metadata" in str(recwarn.pop(RuntimeWarning).message) + + +def test_mine_source_holds_writer_lease_before_opening_handles(monkeypatch): + from mempalace import knowledge_graph, palace + + collection = _FakeCollection() + active = False + + @contextlib.contextmanager + def lock(_palace_path): + nonlocal active + active = True + try: + yield + finally: + active = False + + def get_collection(_palace_path): + assert active + return collection + + register("fixture", _FixtureAdapter) + monkeypatch.setattr(cli, "MempalaceConfig", _FakeConfig) + monkeypatch.setattr(palace, "mine_palace_lock", lock) + monkeypatch.setattr(palace, "get_collection", get_collection) + monkeypatch.setattr(knowledge_graph, "KnowledgeGraph", _FakeKnowledgeGraph) + + assert cli.mine_source_adapter( + source_name="fixture", source_path="/source", palace_path="/fake/palace" + ) == 1 + + +@pytest.mark.skipif(sys.platform == "win32", reason="cross-process lock semantics differ on Windows") +def test_mine_source_refuses_held_writer_lease_before_opening_handles(tmp_path, monkeypatch): + """A competing writer prevents adapter ingest and all handle creation.""" + from mempalace import knowledge_graph, palace + + home = tmp_path / "home" + home.mkdir() + monkeypatch.setenv("HOME", str(home)) + palace_path = str(tmp_path / "palace") + ready = str(tmp_path / "ready") + release = str(tmp_path / "release") + opened_collection = False + opened_kg = False + + def get_collection(_palace_path): + nonlocal opened_collection + opened_collection = True + return _FakeCollection() + + class TrackingKnowledgeGraph(_FakeKnowledgeGraph): + def __init__(self, db_path): + nonlocal opened_kg + opened_kg = True + super().__init__(db_path) + + register("fixture", _FixtureAdapter) + monkeypatch.setattr(cli, "MempalaceConfig", _FakeConfig) + monkeypatch.setattr(palace, "get_collection", get_collection) + monkeypatch.setattr(knowledge_graph, "KnowledgeGraph", TrackingKnowledgeGraph) + + ctx = multiprocessing.get_context("spawn") + holder = ctx.Process(target=_hold_palace_lock, args=(palace_path, ready, release)) + holder.start() + try: + for _ in range(500): + if os.path.exists(ready): + break + time.sleep(0.01) + assert os.path.exists(ready), "lock holder did not become ready" + + with pytest.raises(palace.MineAlreadyRunning): + cli.mine_source_adapter( + source_name="fixture", source_path="/source", palace_path=palace_path + ) + + assert len(_FixtureAdapter.instances) == 1 + assert _FixtureAdapter.instances[0].palace is None + assert not opened_collection + assert not opened_kg + finally: + open(release, "w").close() + holder.join(timeout=10) + if holder.is_alive(): + holder.terminate() + assert holder.exitcode == 0 def test_cmd_mine_without_mode_preserves_projects_legacy_path(monkeypatch):