Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 12 additions & 2 deletions mempalace/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -828,9 +828,18 @@ def cmd_repair(args):
# No prompt when source != dest AND dest does not exist (pure
# extract-into-fresh-dir case is non-destructive to existing
# palaces).
dry_run = getattr(args, "dry_run", False)
# A dry run only reads the source SQLite and prints a plan — it
# never archives, creates, or writes — so it must not trip the
# destructive-action confirmation, which would otherwise block the
# preview behind a y/N prompt (or require --yes just to look).
is_destructive_to_dest = source_path == palace_path or os.path.exists(palace_path)
if is_destructive_to_dest and not confirm_destructive_action(
"Rebuild from SQLite", palace_path, assume_yes=getattr(args, "yes", False)
if (
not dry_run
and is_destructive_to_dest
and not confirm_destructive_action(
"Rebuild from SQLite", palace_path, assume_yes=getattr(args, "yes", False)
)
):
return

Expand All @@ -839,6 +848,7 @@ def cmd_repair(args):
source_palace=source_path,
dest_palace=palace_path,
archive_existing_dest=archive_existing,
dry_run=dry_run,
)
except RebuildPartialError as exc:
# The error itself was already printed by rebuild_from_sqlite
Expand Down
36 changes: 36 additions & 0 deletions mempalace/repair.py
Original file line number Diff line number Diff line change
Expand Up @@ -1053,6 +1053,7 @@ def rebuild_from_sqlite(
*,
archive_existing_dest: bool = False,
batch_size: int = 1000,
dry_run: bool = False,
) -> dict[str, int]:
"""Rebuild a palace by reading drawers from ``source_palace``'s
``chroma.sqlite3`` and upserting them into a fresh palace at
Expand Down Expand Up @@ -1085,6 +1086,16 @@ def rebuild_from_sqlite(
instead. Used by the in-place CLI flow where ``--source`` defaults
to the same path as ``--palace``.

``dry_run`` (CLI: ``--dry-run``) previews the rebuild without making
any change: source validation runs as normal, then per-collection
row counts are read from the source SQLite and printed, and the
function returns those would-be counts *without* archiving the
existing palace, creating collections, or re-embedding. Useful before
a multi-hour rebuild on a large palace. A dry run returns a populated
dict (one key per recoverable collection) so CLI callers treat it as
success; a validation refusal still returns ``{}`` exactly as a real
run would, so the preview faithfully reflects what would happen.

Returns a ``{collection_name: row_count}`` dict so callers (CLI,
tests) can verify the per-collection rebuild count without parsing
stdout. A successful rebuild always returns a dict with one key per
Expand Down Expand Up @@ -1167,6 +1178,31 @@ def rebuild_from_sqlite(
)
return {}

# --dry-run: validation has passed, so report what a real run would
# do and stop before the first irreversible step (the in-place
# archive move). Counts come from ``sqlite_drawer_count`` — the same
# SQLite ground-truth helper repair uses elsewhere — so the preview
# matches the per-collection counts a real rebuild upserts. Reads the
# original ``source_palace`` (not yet archived at this point).
if dry_run:
print("\n DRY RUN — no changes will be made.")
if in_place:
print(
f" Would archive {dest_palace} → "
f"{dest_palace}.pre-rebuild-<timestamp>, then rebuild from the copy."
)
counts = {}
for cname in _recoverable_collections():
n = sqlite_drawer_count(source_palace, cname) or 0
counts[cname] = n
print(f" [{cname}] would re-embed and upsert {n} rows")
Comment on lines +1194 to +1198

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Using or 0 when sqlite_drawer_count returns None silently conflates an unreadable/corrupt database (or schema mismatch) with a genuinely empty collection. If the database is unreadable, the dry run will misleadingly report 0 rows and succeed, whereas a real run would fail.

We should explicitly check if n is None and return an empty dict {} (which indicates a validation/preflight failure to the CLI) to prevent silent failures during dry runs.

        counts = {}
        for cname in _recoverable_collections():
            n = sqlite_drawer_count(source_palace, cname)
            if n is None:
                print(f"  ERROR: [{cname}] count is unreadable (corrupt database or schema mismatch).")
                return {}
            counts[cname] = n
            print(f"  [{cname}] would re-embed and upsert {n} rows")

print(
f"\n Would rebuild {sum(counts.values())} total rows. "
"Re-run without --dry-run to execute."
)
print(f"{'=' * 55}\n")
return counts

archive_path: Optional[str] = None
if in_place:
ts = datetime.now().strftime("%Y%m%d-%H%M%S")
Expand Down
39 changes: 39 additions & 0 deletions tests/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -1384,3 +1384,42 @@ def test_cmd_repair_from_sqlite_success_does_not_exit(mock_config_cls, tmp_path)
with patch("mempalace.repair.rebuild_from_sqlite", return_value=fake_counts):
# Should return cleanly; no SystemExit raised.
cmd_repair(args)


@patch("mempalace.cli.MempalaceConfig")
def test_cmd_repair_from_sqlite_dry_run_passes_through_and_skips_confirm(
mock_config_cls, tmp_path
):
"""``repair --mode from-sqlite --dry-run`` must forward ``dry_run=True``
to ``rebuild_from_sqlite`` and must NOT call ``confirm_destructive_action``.
A preview is read-only, so it cannot be gated behind a y/N prompt or
require ``--yes`` just to look.

Catches the reported defect: ``cmd_repair`` ignored ``--dry-run`` for
from-sqlite mode (it never passed the flag through), so a preview
either ran the real destructive rebuild or aborted at the confirmation
prompt.
"""
palace_dir = tmp_path / "palace"
palace_dir.mkdir() # existing dir → would normally be destructive-to-dest and prompt
mock_config_cls.return_value.palace_path = str(palace_dir)

args = argparse.Namespace(
palace=str(palace_dir),
mode="from-sqlite",
source=None,
archive_existing=False,
yes=False, # no --yes: a real destructive run would block at the prompt
dry_run=True,
)
with (
patch(
"mempalace.repair.rebuild_from_sqlite",
return_value={"mempalace_drawers": 0, "mempalace_closets": 0},
) as mock_rebuild,
patch("mempalace.migrate.confirm_destructive_action") as mock_confirm,
):
cmd_repair(args)

mock_confirm.assert_not_called()
assert mock_rebuild.call_args.kwargs["dry_run"] is True
48 changes: 48 additions & 0 deletions tests/test_repair.py
Original file line number Diff line number Diff line change
Expand Up @@ -1750,6 +1750,54 @@ def test_rebuild_from_sqlite_source_missing_chroma_db(tmp_path):
assert not dest.exists()


def test_rebuild_from_sqlite_dry_run_cross_palace_writes_nothing(tmp_path):
"""``dry_run=True`` must report the per-collection row counts a real
rebuild would produce while creating nothing at dest.

Catches: a regression where ``--dry-run`` is ignored and the rebuild
runs anyway (dest gets written), or where the preview returns the
wrong counts. The counts asserted here match what
``test_rebuild_from_sqlite_roundtrips_via_real_chromadb`` proves a
real rebuild upserts, so the preview is verified against real
behavior, not just against itself.
"""
source = tmp_path / "source"
dest = tmp_path / "dest"
drawer_rows = [(f"d{i}", f"body {i}", {"wing": "w", "room": "r"}) for i in range(12)]
_seed_palace(source, "mempalace_drawers", drawer_rows)
_seed_palace(source, "mempalace_closets", [("c1", "abbrev", {"wing": "w"})])

counts = repair.rebuild_from_sqlite(str(source), str(dest), dry_run=True)

assert counts == {"mempalace_drawers": 12, "mempalace_closets": 1}
# Nothing written: dest never created, source untouched.
assert not dest.exists()
assert (source / "chroma.sqlite3").exists()


def test_rebuild_from_sqlite_dry_run_in_place_does_not_archive(tmp_path):
"""An in-place ``dry_run`` must NOT move the live palace aside. The
archive rename is the first irreversible step of a real in-place
rebuild; a dry run that performs it would strand the user's only
palace at a timestamped path and leave the expected location empty.

Catches: a regression where the dry-run early-return is placed after
the archive move instead of before it.
"""
palace = tmp_path / "palace"
_seed_palace(palace, "mempalace_drawers", [(f"d{i}", f"b{i}", {"wing": "w"}) for i in range(8)])
sqlite_before = (palace / "chroma.sqlite3").stat().st_size

counts = repair.rebuild_from_sqlite(
str(palace), str(palace), archive_existing_dest=True, dry_run=True
)

assert counts == {"mempalace_drawers": 8, "mempalace_closets": 0}
# Original palace untouched, no archive sibling created.
assert (palace / "chroma.sqlite3").stat().st_size == sqlite_before
assert [p for p in tmp_path.iterdir() if "pre-rebuild" in p.name] == []


def test_rebuild_from_sqlite_in_place_validates_source_before_archiving(tmp_path):
"""In-place + archive_existing_dest=True with a dir that lacks
chroma.sqlite3 must NOT rename the dir before bailing. An earlier
Expand Down
Loading