xiaoxue_admin
/
mempalace
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(mine): harden source adapter dispatch

This commit is contained in:
Grace Gettert 2026-08-06 16:33:30 +00:00
parent b6fad9ed3a
commit 141e16fb2b
2 changed files with 246 additions and 51 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

115
mempalace/cli.py
View File

@ -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):

182
tests/test_cli_source_adapters.py
View File

@ -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):