diff --git a/tests/test_repair_foreign_links.py b/tests/test_repair_foreign_links.py index 3d4e47e5..23628183 100644 --- a/tests/test_repair_foreign_links.py +++ b/tests/test_repair_foreign_links.py @@ -15,6 +15,7 @@ import shutil import sqlite3 import tempfile import unittest +import unittest.mock from pathlib import Path from PIL import Image @@ -64,6 +65,10 @@ def not_running(_comfyui_dir, _port): return [] +def no_handles(_database): + return [] + + class RepairFixture(RecordedRuns, unittest.TestCase): """Rows: A relink to a new prompt, B relink to a stored prompt (with a pair link), C keep, D file gone, E two savers (skip), F the UNIQUE case @@ -81,6 +86,9 @@ class RepairFixture(RecordedRuns, unittest.TestCase): self.plan_path = os.path.join(self.tmp, "plan.json") self.backups = os.path.join(self.tmp, "backups") os.makedirs(self.backups) + self.comfy = os.path.join(self.tmp, "comfy") + os.makedirs(self.comfy) + Path(self.comfy, "main.py").write_text("# ComfyUI\n", encoding="utf-8") self.files = {} for key, number in zip("ABCDEFHG", range(1, 9)): path = os.path.join(self.output, "dataset", NAME.format(number)) @@ -156,7 +164,7 @@ class RepairFixture(RecordedRuns, unittest.TestCase): ) return code, out.getvalue() - def apply(self, *args, probe=not_running): + def apply(self, *args, probe=not_running, handles=no_handles): out = io.StringIO() code = tool.main( [ @@ -171,6 +179,7 @@ class RepairFixture(RecordedRuns, unittest.TestCase): ], out=out, probe=probe, + handles=handles, ) return code, out.getvalue() @@ -189,6 +198,14 @@ class RepairFixture(RecordedRuns, unittest.TestCase): finally: conn.close() + def query_write(self, sql, args): + conn = sqlite3.connect(self.db_path) + try: + with conn: + conn.execute(sql, args) + finally: + conn.close() + def links(self): return dict(self.query("SELECT id, prompt_id FROM generated_images")) @@ -255,11 +272,11 @@ class DryRunTests(RepairFixture): sidecars = {"prompts.db-shm", "prompts.db-wal"} self.assertEqual( sorted(set(os.listdir(self.tmp)) - sidecars), - sorted(set(listing + ["plan.json", "plan.md"]) - sidecars), + sorted(set(listing + ["plan.json", "plan.list.md"]) - sidecars), ) wal = Path(self.tmp, "prompts.db-wal") self.assertEqual(wal.stat().st_size if wal.exists() else 0, 0) - markdown = Path(self.tmp, "plan.md").read_text(encoding="utf-8") + markdown = Path(self.tmp, "plan.list.md").read_text(encoding="utf-8") self.assertIn(f"| {self.rows['A']} |", markdown) long_path = os.path.join("dataset", NAME.format(1)) self.assertIn(f"| {long_path} |", markdown, "file paths are never cut") @@ -268,6 +285,58 @@ class DryRunTests(RepairFixture): self.assertIn("would lose every link", markdown) +class DryRunReviewTests(RepairFixture): + def test_plan_records_its_library_and_fingerprint(self): + self.dry_run() + plan = self.plan() + self.assertEqual( + plan["source_db"], os.path.normcase(os.path.realpath(self.db_path)) + ) + self.assertEqual(plan["max_row_id"], max(self.links())) + + def test_a_markdown_plan_path_is_refused(self): + out = io.StringIO() + with ( + self.assertRaises(SystemExit), + contextlib.redirect_stderr(io.StringIO()), + ): + tool.main( + [ + self.db_path, + "--output-dir", + self.output, + "--plan", + os.path.join(self.tmp, "plan.md"), + ], + out=out, + ) + self.assertFalse(os.path.exists(os.path.join(self.tmp, "plan.md"))) + + def test_keep_compares_the_linked_text_not_the_stored_hash(self): + self.query_write("UPDATE prompts SET hash='stale' WHERE id=?", (self.p0,)) + self.dry_run() + items = {item["row_id"]: item for item in self.plan()["rows"]} + self.assertEqual(items[self.rows["C"]]["action"], "keep") + + def test_delete_orphans_is_rehearsed_with_the_deletes(self): + code, out = self.dry_run("--delete-orphans") + self.assertEqual(code, 0, out) + simulated = self.plan()["simulated"] + self.assertEqual(simulated["result"]["orphans_deleted"], 1) + self.assertEqual( + simulated["after"]["prompts"], simulated["before"]["prompts"] + 2 - 1 + ) + + def test_a_file_on_another_drive_is_listed_by_its_full_path(self): + self.dry_run() + plan = self.plan() + with unittest.mock.patch.object( + tool.os.path, "relpath", side_effect=ValueError("other drive") + ): + listing = tool.markdown(plan, {}) + self.assertIn(self.files["A"], listing) + + class ApplyTests(RepairFixture): def setUp(self): super().setUp() @@ -353,6 +422,7 @@ class ApplyTests(RepairFixture): ], out=out, probe=not_running, + handles=no_handles, ) self.assertEqual(code, 0, out.getvalue()) self.assertTrue( @@ -391,6 +461,7 @@ class ApplyTests(RepairFixture): ], out=out, probe=lambda _d, _p: [], + handles=no_handles, ) self.assertEqual(code, 2) self.assertIn("--rehearsal", out.getvalue()) @@ -417,6 +488,8 @@ class ApplyTests(RepairFixture): with self.subTest(name): shutil.copy(pristine, self.db_path) Path(self.files["B"]).write_bytes(image) + shutil.rmtree(self.backups) + os.makedirs(self.backups) drift() before = dump(self.db_path) code, out = self.apply() @@ -424,14 +497,6 @@ class ApplyTests(RepairFixture): self.assertIn("REFUSED (rolled back, nothing changed)", out) self.assertEqual(dump(self.db_path), before) - def query_write(self, sql, args): - conn = sqlite3.connect(self.db_path) - try: - with conn: - conn.execute(sql, args) - finally: - conn.close() - def test_late_failure_rolls_back_every_write(self): plan = self.plan() plan["delete_orphans"] = True @@ -474,6 +539,201 @@ class ApplyTests(RepairFixture): self.assertEqual(dump(self.db_path), before) +class ApplySafetyTests(RepairFixture): + def setUp(self): + super().setUp() + code, out = self.dry_run() + self.assertEqual(code, 0, out) + + def assert_unchanged(self, before, code, out): + self.assertEqual(code, 2, out) + self.assertEqual(dump(self.db_path), before) + + def test_refuses_when_another_process_has_the_library_open(self): + before = dump(self.db_path) + held = lambda _db: ["process 7 has prompts.db-wal open"] # noqa: E731 + for args in ((), ("--rehearsal",)): + with self.subTest(args): + code, out = self.apply(*args, handles=held) + self.assert_unchanged(before, code, out) + self.assertIn("process 7", out) + self.assertEqual(os.listdir(self.backups), []) + + def test_the_running_check_runs_again_just_before_writing(self): + calls = [] + + def late(_d, _p): + calls.append(1) + return [] if len(calls) == 1 else ["something listens on port 8188"] + + before = dump(self.db_path) + code, out = self.apply(probe=late) + self.assert_unchanged(before, code, out) + self.assertEqual(len(calls), 2) + self.assertIn("rolled back", out) + self.assertIn("Backup:", out) + + def test_without_proc_only_an_explicit_stop_lets_it_through(self): + blind = lambda _db: [tool.NO_PROC] # noqa: E731 + before = dump(self.db_path) + code, out = self.apply(handles=blind) + self.assert_unchanged(before, code, out) + self.assertIn("--i-stopped-comfyui", out) + code, out = self.apply("--i-stopped-comfyui", handles=blind) + self.assertEqual(code, 0, out) + + def test_a_folder_without_main_py_is_refused(self): + os.remove(os.path.join(self.comfy, "main.py")) + before = dump(self.db_path) + code, out = self.apply() + self.assert_unchanged(before, code, out) + self.assertIn("main.py", out) + + def test_a_plan_for_another_library_is_refused(self): + other = os.path.join(self.tmp, "other.db") + shutil.copy(self.db_path, other) + out = io.StringIO() + code = tool.main( + [other, "--apply", self.plan_path, "--comfyui-dir", self.comfy], + out=out, + probe=not_running, + handles=no_handles, + ) + self.assertEqual(code, 2) + self.assertIn("made for", out.getvalue()) + + def test_opening_the_library_must_not_change_its_schema(self): + self.query_write("DROP INDEX idx_prompt_pairs_positive_only", ()) + PromptModel.reset_schema_cache() + code, out = self.apply() + self.assertEqual(code, 2, out) + self.assertIn("schema", out) + self.assertIn("live code version", out) + + def test_an_existing_backup_file_is_refused(self): + stamp = "20261028_000000" + Path(self.backups, f"prompts.db.pre-247-repair_{stamp}").write_bytes(b"x") + before = dump(self.db_path) + with unittest.mock.patch.object(tool, "_stamp", return_value=stamp): + code, out = self.apply() + self.assert_unchanged(before, code, out) + self.assertIn("already exists", out) + + def test_a_backup_that_does_not_match_is_refused(self): + real = tool.backup_database + + def short(source, target): + real(source, target) + conn = sqlite3.connect(target) + try: + with conn: + conn.execute("DELETE FROM image_prompt_pairs") + finally: + conn.close() + + before = dump(self.db_path) + with unittest.mock.patch.object(tool, "backup_database", short): + code, out = self.apply() + self.assert_unchanged(before, code, out) + self.assertIn("backup", out) + + def test_a_changed_target_text_is_refused(self): + plan = self.plan() + for item in plan["rows"]: + if item["row_id"] == self.rows["A"]: + item["target_text"] = "something else" + self.write_plan(plan) + before = dump(self.db_path) + code, out = self.apply() + self.assert_unchanged(before, code, out) + self.assertIn("target_text", out) + + def test_an_orphan_gone_or_changed_at_apply_is_refused(self): + plan = self.plan() + plan["delete_orphans"] = True + plan["orphans"][0]["text_sha256"] = "0" * 64 + self.write_plan(plan) + before = dump(self.db_path) + code, out = self.apply() + self.assert_unchanged(before, code, out) + self.assertIn("gone or text changed", out) + + def test_any_error_rolls_back_and_names_the_backup(self): + before = dump(self.db_path) + with unittest.mock.patch.object( + tool, "_relink", side_effect=RuntimeError("disk on fire") + ): + code, out = self.apply() + self.assert_unchanged(before, code, out) + self.assertIn("rolled back", out) + self.assertIn("disk on fire", out) + self.assertIn(self.backups, out) + + def test_a_failed_integrity_check_after_writing_exits_1(self): + real = tool.counts + seen = [] + + def counts(path, readonly=True): + found = real(path, readonly) + if path == self.db_path: + seen.append(path) + if len(seen) > 1: + return {**found, "integrity_check": "corrupt"} + return found + + with unittest.mock.patch.object(tool, "counts", counts): + code, out = self.apply() + self.assertEqual(code, 1, out) + self.assertIn("restore", out) + + +class BlockerTests(RepairFixture): + """Each reference that keeps a prompt from being deleted.""" + + def test_every_blocker(self): + pid = self.orphan + cases = { + "linked images": None, + "saved pairs": "INSERT INTO prompt_pairs(positive_prompt_id) VALUES(:id)", + "recorded runs": "UPDATE execution_run_contexts " + "SET positive_prompt_id=:id", + "library state": "INSERT INTO prompt_library_state(prompt_id) " + "VALUES(:id)", + "user tags": "INSERT INTO prompt_tags(prompt_id, tag_id) " + "SELECT :id, id FROM tags WHERE name='mine'", + "usage or rating": "UPDATE prompts SET run_count=3 WHERE id=:id", + "category": "UPDATE prompts SET category='portraits' WHERE id=:id", + "notes": "UPDATE prompts SET notes='mine' WHERE id=:id", + } + conn = sqlite3.connect(self.db_path) + try: + conn.execute("INSERT INTO tags(name) VALUES('mine')") + conn.commit() + for name, sql in cases.items(): + with self.subTest(name): + conn.execute("SAVEPOINT blocker") + if sql is None: + unused = {pid} + else: + conn.execute( + "DELETE FROM generated_images WHERE prompt_id=?", (pid,) + ) + conn.execute(sql, {"id": pid}) + unused = {pid} + found = tool.blockers(conn, pid, unused) + self.assertTrue(any(item.startswith(name) for item in found), found) + conn.execute("ROLLBACK TO blocker") + conn.execute("RELEASE blocker") + conn.execute("DELETE FROM generated_images WHERE prompt_id=?", (pid,)) + self.assertEqual(tool.blockers(conn, pid, {pid}), []) + self.assertEqual( + tool.blockers(conn, pid, set()), ["not unused by the app's rule"] + ) + conn.rollback() + finally: + conn.close() + + class ActionTests(unittest.TestCase): """Only the recorded saver relinks; a certain step naming the linked prompt keeps it; the legacy floor and the rest are skipped.""" @@ -513,6 +773,41 @@ class CommandLineTests(unittest.TestCase): port = probe.getsockname()[1] self.assertEqual(tool.comfyui_running("/nonexistent/comfy", port), []) + def fake_proc(self, tmp): + proc = Path(tmp, "proc") + db = Path(tmp, "prompts.db") + db.write_bytes(b"") + Path(tmp, "prompts.db-wal").write_bytes(b"") + comfy = Path(tmp, "comfy") + comfy.mkdir() + holder = proc / "125" / "fd" + holder.mkdir(parents=True) + os.symlink(Path(tmp, "prompts.db-wal"), holder / "3") + (proc / "125" / "cmdline").write_bytes(b"python\0main.py\0") + os.symlink(comfy, proc / "125" / "cwd") + (proc / "123").mkdir() # unreadable: no fd, cmdline or cwd + (proc / "self").mkdir() # not a pid + return proc, db, comfy + + @unittest.skipIf(os.name == "nt", "symlinks need privileges on Windows") + def test_proc_scan_finds_holders_and_skips_unreadable_entries(self): + with tempfile.TemporaryDirectory() as tmp: + proc, db, comfy = self.fake_proc(tmp) + found = tool.open_by_others(str(db), proc=str(proc), own_pid=1) + self.assertEqual(len(found), 1) + self.assertIn("125", found[0]) + self.assertEqual( + tool.open_by_others(str(db), proc=str(proc), own_pid=125), [] + ) + running = tool.comfyui_running(str(comfy), 1, proc=str(proc)) + self.assertTrue(any("125" in item for item in running), running) + + def test_without_proc_the_scan_says_it_cannot_tell(self): + self.assertEqual( + tool.open_by_others("/x/prompts.db", proc="/nonexistent/proc"), + [tool.NO_PROC], + ) + def test_missing_database_is_an_argument_error(self): with self.assertRaises(SystemExit), contextlib.redirect_stderr(io.StringIO()): tool.main( diff --git a/tools/comfyui_stopped.py b/tools/comfyui_stopped.py new file mode 100644 index 00000000..c21bfec1 --- /dev/null +++ b/tools/comfyui_stopped.py @@ -0,0 +1,81 @@ +"""Is ComfyUI stopped? The checks a library repair runs before writing (#247). + +``comfyui_running`` looks for ``main.py`` processes of one ComfyUI folder +and a listener on its port. ``open_by_others`` lists the processes that hold +the library file, its ``-wal`` or its ``-shm`` open. Both read ``/proc`` +(Linux); without it ``open_by_others`` answers ``[NO_PROC]``: it cannot +tell, and the caller refuses unless the user says ComfyUI is stopped. +Entries that cannot be read (another user's process, one that just exited) +are skipped. +""" + +import os +import socket + +NO_PROC = ( + "cannot list the processes that hold the library open (no /proc): stop " + "ComfyUI and pass --i-stopped-comfyui" +) + + +def _same(path): + return os.path.normcase(os.path.realpath(path)) + + +def _pids(proc, own_pid=None): + try: + names = os.listdir(proc) + except OSError: + return [] + return [name for name in names if name.isdigit() and name != str(own_pid)] + + +def comfyui_running(comfyui_dir, port, proc="/proc"): + """``main.py`` processes running from ``comfyui_dir`` and a listener on + ``port`` (127.0.0.1), each as one line; empty when none is found.""" + found = [] + folder = _same(comfyui_dir) + for pid in _pids(proc): + try: + with open(os.path.join(proc, pid, "cmdline"), "rb") as stream: + parts = stream.read().split(b"\0") + cwd = _same(os.readlink(os.path.join(proc, pid, "cwd"))) + except OSError: + continue + if any(part.endswith(b"main.py") for part in parts) and ( + cwd == folder or any(folder.encode() in part for part in parts) + ): + found.append(f"process {pid} (main.py in {cwd})") + with socket.socket() as probe: + probe.settimeout(1) + if probe.connect_ex(("127.0.0.1", port)) == 0: + found.append(f"something listens on port {port}") + return found + + +def open_by_others(database, proc="/proc", own_pid=None): + """Processes other than ``own_pid`` (default: this one) holding + ``database``, its ``-wal`` or ``-shm`` open; ``[NO_PROC]`` without /proc.""" + if not os.path.isdir(proc): + return [NO_PROC] + own_pid = os.getpid() if own_pid is None else own_pid + targets = {_same(database + suffix): suffix for suffix in ("", "-wal", "-shm")} + found = [] + for pid in _pids(proc, own_pid): + fd_dir = os.path.join(proc, pid, "fd") + try: + fds = os.listdir(fd_dir) + except OSError: + continue + held = set() + for fd in fds: + try: + target = os.readlink(os.path.join(fd_dir, fd)) + except OSError: + continue + target = target.replace(" (deleted)", "") + if os.path.isabs(target) and _same(target) in targets: + held.add(os.path.basename(database) + targets[_same(target)]) + if held: + found.append(f"process {pid} has {', '.join(sorted(held))} open") + return found diff --git a/tools/repair_foreign_links.py b/tools/repair_foreign_links.py index 8b4162a4..699faa96 100644 --- a/tools/repair_foreign_links.py +++ b/tools/repair_foreign_links.py @@ -16,18 +16,28 @@ PromptManager context: the legacy floor or another step decided) Dry run (the default) judges a copy of the database (SQLite backup from a -read-only connection) and writes a JSON plan and a Markdown list beside it. -It also rehearses the apply on that copy, to list the prompts that would -lose every link with their remaining references. Deleting those is OFF -unless the dry run is given ``--delete-orphans``; the choice is recorded in -the plan. The library is never written. +read-only connection) and writes a JSON plan and a Markdown list beside it +(``.list.md``). It also rehearses the apply on a second copy, to +list the prompts that would lose every link with their remaining +references. Deleting those is OFF unless the dry run is given +``--delete-orphans``; the choice is recorded in the plan and rehearsed on a +third copy. The library is never written. -``--apply `` (ComfyUI stopped): +``--apply `` (ComfyUI stopped), only on the library the plan was +made from (its real path, recorded in the plan): -1. refuses while a ``main.py`` runs from ``--comfyui-dir`` or anything listens - on ``--port`` (``--rehearsal`` skips this, never for the library of that - ComfyUI folder); -2. saves ``.pre-247-repair_`` (or into ``--backup-dir``); +1. refuses unless ``--comfyui-dir`` holds ``main.py``; refuses while a + ``main.py`` runs from it, anything listens on ``--port``, or any other + process holds the library, its ``-wal`` or ``-shm`` open + (``tools/comfyui_stopped.py``; without /proc only with + ``--i-stopped-comfyui``). ``--rehearsal`` skips the process and port + check, never the open-file check, and never for that folder's library; +2. saves ``.pre-247-repair_`` (or into ``--backup-dir``), never + over an existing file, and refuses unless the backup passes + ``integrity_check`` with the library's counts. Opening the library must + not change its schema (a newer checkout's migration): run the tool from a + checkout at the live code version. The running check runs once more just + before writing; 3. in ONE ``BEGIN IMMEDIATE`` transaction: re-judges every planned row and refuses on any drift (row moved, prompt text changed, file changed, target stored since), creates the target prompts that do not exist (category @@ -45,8 +55,8 @@ counts, ``integrity_check`` and ``foreign_key_check`` are printed. python3 tools/repair_foreign_links.py --apply repair_plan.json \\ --comfyui-dir -Run from this repository with its venv. The process check reads /proc -(Linux); elsewhere only the port is checked, so stop ComfyUI first. +Run from a checkout at the live code version (after the merge: main) with +its venv. """ import argparse @@ -55,7 +65,6 @@ import hashlib import json import logging import os -import socket import sqlite3 import sys import tempfile @@ -75,6 +84,11 @@ from database.pair_ingestion import ( # noqa: E402 ) from database.run_attribution import recorded_output # noqa: E402 from py.api import PromptManagerAPI # noqa: E402 +from tools.comfyui_stopped import ( # noqa: E402 + NO_PROC, + comfyui_running, + open_by_others, +) from tools.review_image_links import backup_database # noqa: E402 from utils.execution_outputs import snapshot_output_roots # noqa: E402 from utils.file_prompt import POSITIVE, Policy, decide_file_prompt # noqa: E402 @@ -105,6 +119,7 @@ DRIFT_KEYS = ( "current_text_sha256", "file_size", "file_mtime_ns", + "target_text", "target_hash", "target_prompt_id", ) @@ -127,29 +142,6 @@ def sha256(text): return hashlib.sha256(text.encode("utf-8")).hexdigest() -def comfyui_running(comfyui_dir, port): - """What shows ComfyUI is running: ``main.py`` processes of that folder - (Linux /proc) and a listener on ``port``.""" - found = [] - folder = os.path.realpath(comfyui_dir) - proc = Path("/proc") - for pid in filter(str.isdigit, os.listdir(proc) if proc.is_dir() else ()): - try: - parts = (proc / pid / "cmdline").read_bytes().split(b"\0") - cwd = os.readlink(proc / pid / "cwd") - except OSError: - continue - if any(part.endswith(b"main.py") for part in parts) and ( - cwd == folder or any(folder.encode() in part for part in parts) - ): - found.append(f"process {pid} (main.py in {cwd})") - with socket.socket() as probe: - probe.settimeout(1) - if probe.connect_ex(("127.0.0.1", port)) == 0: - found.append(f"something listens on port {port}") - return found - - def metadata_reader(): """The scan's metadata extraction (PNG text chunks, ffprobe for videos).""" api = SimpleNamespace(logger=logging.getLogger("prompt_manager.repair_247")) @@ -193,9 +185,10 @@ class Judge: """The plan entry of one ``generated_images`` row.""" conn = self.db.model.get_connection() path = row["image_path"] - text, prompt_hash = conn.execute( - "SELECT text, hash FROM prompts WHERE id = ?", (row["prompt_id"],) + (text,) = conn.execute( + "SELECT text FROM prompts WHERE id = ?", (row["prompt_id"],) ).fetchone() + prompt_hash = generate_prompt_hash(text) item = { "row_id": row["id"], "file": path, @@ -244,11 +237,21 @@ def action_of(decision, linked_hash): return "relink" if decision.source == "recorded-run" else "skip" -def build_plan(db, output_dir, read=None): +def library_identity(path): + """How a plan names its library: the case-normalized real path.""" + return os.path.normcase(os.path.realpath(path)) + + +def build_plan(db, output_dir, source_db, read=None): + """The plan for ``db`` (a copy of ``source_db``).""" judge = Judge(db, output_dir, read) + conn = db.model.get_connection() + (max_row_id,) = conn.execute("SELECT MAX(id) FROM generated_images").fetchone() return { "version": PLAN_VERSION, "created": datetime.datetime.now().isoformat(timespec="seconds"), + "source_db": library_identity(source_db), + "max_row_id": max_row_id, "output_dir": output_dir, "rows": [judge.row(row) for row in foreign_rows(db, output_dir)], "orphans": [], @@ -300,6 +303,10 @@ def _rejudge(conn, judge, plan): ).fetchone() if row is None: raise Refused(f"row {item['row_id']}: gone") + if generate_prompt_hash(item.get("target_text") or "") != item.get( + "target_hash" + ): + raise Refused(f"row {item['row_id']}: target_text does not match its hash") fresh = judge.row(row) for key in DRIFT_KEYS: if fresh.get(key) != item.get(key): @@ -397,13 +404,16 @@ def _delete_orphans(conn, plan): return len(ids) -def apply_plan(db, plan, read=None): +def apply_plan(db, plan, read=None, before_write=None): """Write ``plan`` in one transaction; returns ``(result, orphan audit)``. - Raises ``Refused`` (and rolls back) on any drift or failed check. + ``before_write()`` runs just before ``BEGIN IMMEDIATE`` (the running + check). Raises ``Refused`` (and rolls back) on any drift or failed check. """ judge = Judge(db, plan["output_dir"], read) relinks = [item for item in plan["rows"] if item["action"] == "relink"] + if before_write is not None: + before_write() with db._pair_transaction() as conn: _rejudge(conn, judge, plan) targets = _targets(conn, db, relinks) @@ -426,9 +436,13 @@ def apply_plan(db, plan, read=None): return result, audit -def counts(path): - """Table counts plus the integrity and foreign-key checks (read-only).""" - uri = Path(os.path.abspath(path)).as_uri() + "?mode=ro" +def counts(path, readonly=True): + """Table counts plus the integrity and foreign-key checks. + + Read-only for a library; the tool's own backup is opened read-write so + a WAL file leaves no side files behind once it closes. + """ + uri = Path(os.path.abspath(path)).as_uri() + ("?mode=ro" if readonly else "") conn = sqlite3.connect(uri, uri=True) try: found = { @@ -463,10 +477,17 @@ def _cell(text, limit=EXCERPT): return (folded if limit is None else folded[:limit]).replace("|", "\\|") +def _relative(path, root): + try: + return os.path.relpath(path, root) + except ValueError: # another drive on Windows + return path + + def markdown(plan, prompt_texts): """The maintainer's list of the plan.""" root = plan["output_dir"] - rel = lambda path: os.path.relpath(path, root) # noqa: E731 + rel = lambda path: _relative(path, root) # noqa: E731 rows = plan["rows"] lines = [ "# #247 repair plan (dry run)", @@ -529,34 +550,49 @@ def markdown(plan, prompt_texts): return "\n".join(lines) + "\n" +def _texts(db): + return dict( + db.model.get_connection().execute( + "SELECT id, text FROM prompts WHERE id IN (SELECT prompt_id " + "FROM generated_images)" + ) + ) + + +def _rehearse(source, scratch, name, plan): + """``apply_plan`` on a fresh backup of ``source``: (counts, result, audit).""" + copy = os.path.join(scratch, name) + backup_database(source, copy) + db = PromptDatabase(copy) + try: + before = counts(copy) + result, audit = apply_plan(db, plan) + finally: + db.close_all() + return {"before": before, "after": counts(copy), "result": result}, audit + + def dry_run(args, out): with tempfile.TemporaryDirectory() as scratch: - copy = os.path.join(scratch, "repair.db") + copy = os.path.join(scratch, "judge.db") backup_database(args.database, copy) db = PromptDatabase(copy) try: - plan = build_plan(db, os.path.abspath(args.output_dir)) - texts = dict( - db.model.get_connection().execute( - "SELECT id, text FROM prompts WHERE id IN (SELECT prompt_id " - "FROM generated_images)" - ) - ) - before = counts(copy) - result, audit = apply_plan(db, plan) + plan = build_plan(db, os.path.abspath(args.output_dir), args.database) + texts = _texts(db) finally: db.close_all() - after = counts(copy) - plan.update( - orphans=audit, - delete_orphans=bool(args.delete_orphans), - simulated={"before": before, "after": after, "result": result}, - ) + simulated, audit = _rehearse(args.database, scratch, "relinks.db", plan) + plan.update(orphans=audit, delete_orphans=bool(args.delete_orphans)) + if plan["delete_orphans"]: + # The plan as the apply will write it, deletes included + simulated, _audit = _rehearse(args.database, scratch, "deletes.db", plan) + plan["simulated"] = simulated plan_path = Path(args.plan) plan_path.write_text( json.dumps(plan, indent=1, ensure_ascii=False) + "\n", encoding="utf-8" ) - list_path = plan_path.with_suffix(".md") + list_path = plan_path.with_name(plan_path.stem + ".list.md") list_path.write_text(markdown(plan, texts), encoding="utf-8") out.write(f"{_summary(plan)}\n") out.write( @@ -571,40 +607,119 @@ def dry_run(args, out): def _is_live(database, comfyui_dir): live = os.path.join(comfyui_dir, *LIVE_LIBRARY) - return os.path.realpath(database) == os.path.realpath(live) + return library_identity(database) == library_identity(live) -def apply(args, out, probe): +def _stamp(): + return datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + + +def schema_of(path): + """``sqlite_master`` and ``user_version`` from a read-only connection.""" + uri = Path(os.path.abspath(path)).as_uri() + "?mode=ro" + conn = sqlite3.connect(uri, uri=True) + try: + master = conn.execute( + "SELECT type, name, tbl_name, sql FROM sqlite_master ORDER BY type, name" + ).fetchall() + return master, conn.execute("PRAGMA user_version").fetchone()[0] + finally: + conn.close() + + +def running_reasons(args, probe, handles): + """Why writing must not start now (empty: go ahead). + + Another process holding the library open always counts, also in a + rehearsal; ComfyUI's processes and port count unless rehearsing. + """ + held = [ + reason + for reason in handles(args.database) + if not (reason == NO_PROC and args.i_stopped_comfyui) + ] + return held + ([] if args.rehearsal else probe(args.comfyui_dir, args.port)) + + +def _preflight(args, plan, probe, handles): + """The refusals before any backup or write, as a message (None: go on).""" + if plan.get("version") != PLAN_VERSION: + return "not a plan of this tool" + if args.rehearsal and _is_live(args.database, args.comfyui_dir): + return f"--rehearsal is for a copy; this is the library of {args.comfyui_dir}" + if plan.get("source_db") != library_identity(args.database): + return f"this plan was made for {plan.get('source_db')}" + if not os.path.isfile(os.path.join(args.comfyui_dir, "main.py")): + return f"no main.py in {args.comfyui_dir}: not a ComfyUI folder" + reasons = running_reasons(args, probe, handles) + if reasons: + return "ComfyUI may be running: " + "; ".join(reasons) + return None + + +def _backup(args, before): + """The verified backup's path; raises ``Refused``.""" + folder = args.backup_dir or os.path.dirname(os.path.abspath(args.database)) + name = f"{os.path.basename(args.database)}.pre-247-repair_{_stamp()}" + backup = os.path.join(folder, name) + if os.path.exists(backup): + raise Refused(f"the backup {backup} already exists") + backup_database(args.database, backup) + saved = counts(backup, readonly=False) + if saved != before or saved["integrity_check"] != "ok": + raise Refused(f"the backup {backup} does not match the library") + return backup + + +def _open_unchanged(path): + """``PromptDatabase(path)``; raises ``Refused`` when opening it changed + its schema (a migration of a newer checkout, committed on open).""" + schema = schema_of(path) + db = PromptDatabase(path) + db.model.get_connection() + if schema_of(path) != schema: + db.close_all() + raise Refused( + "opening the library with this checkout changed its schema (that " + "change is committed); run the tool from a checkout at the live " + "code version" + ) + return db + + +def apply(args, out, probe, handles): with open(args.apply, encoding="utf-8") as stream: plan = json.load(stream) - if plan.get("version") != PLAN_VERSION: - out.write("REFUSED: not a plan of this tool\n") - return 2 - if args.rehearsal and _is_live(args.database, args.comfyui_dir): - out.write( - "REFUSED: --rehearsal is for a copy; this is the library of " - f"{args.comfyui_dir}\n" - ) - return 2 - running = probe(args.comfyui_dir, args.port) - if running and not args.rehearsal: - out.write("REFUSED: ComfyUI is running: " + "; ".join(running) + "\n") + refusal = _preflight(args, plan, probe, handles) + if refusal: + out.write(f"REFUSED: {refusal}\n") return 2 if args.rehearsal: - out.write(f"(rehearsal on a copy; running check: {running or 'nothing'})\n") - stamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") - folder = args.backup_dir or os.path.dirname(os.path.abspath(args.database)) - backup = os.path.join( - folder, f"{os.path.basename(args.database)}.pre-247-repair_{stamp}" - ) - backup_database(args.database, backup) - out.write(f"Backup: {backup}\n") + out.write("(rehearsal on a copy: ComfyUI's process and port not checked)\n") before = counts(args.database) - db = PromptDatabase(args.database) try: - result, _audit = apply_plan(db, plan) - except (Refused, sqlite3.Error) as error: - out.write(f"REFUSED (rolled back, nothing changed): {error}\n") + backup = _backup(args, before) + except Refused as error: + out.write(f"REFUSED: {error}\n") + return 2 + out.write(f"Backup: {backup}\n") + + def still_stopped(): + reasons = running_reasons(args, probe, handles) + if reasons: + raise Refused("ComfyUI may be running: " + "; ".join(reasons)) + + try: + db = _open_unchanged(args.database) + except Refused as error: + out.write(f"REFUSED: {error}. Backup: {backup}\n") + return 2 + try: + result, _audit = apply_plan(db, plan, before_write=still_stopped) + except Exception as error: + out.write( + f"REFUSED (rolled back, nothing changed): {error}. Backup: {backup}\n" + ) return 2 finally: db.close_all() @@ -617,7 +732,7 @@ def apply(args, out, probe): return 0 -def main(argv=None, out=sys.stdout, probe=comfyui_running): +def main(argv=None, out=sys.stdout, probe=comfyui_running, handles=open_by_others): parser = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0]) parser.add_argument("database", help="the library database (prompts.db)") parser.add_argument("--output-dir", help="ComfyUI's output folder (dry run)") @@ -638,15 +753,22 @@ def main(argv=None, out=sys.stdout, probe=comfyui_running): action="store_true", help="apply to a copy while ComfyUI runs (never its own library)", ) + parser.add_argument( + "--i-stopped-comfyui", + action="store_true", + help="without /proc: say ComfyUI is stopped (the port is still checked)", + ) args = parser.parse_args(argv) if not os.path.isfile(args.database): parser.error(f"no database at {args.database}") if args.apply: if not args.comfyui_dir: parser.error("--apply needs --comfyui-dir (the running check)") - return apply(args, out, probe) + return apply(args, out, probe, handles) if not args.output_dir: parser.error("the dry run needs --output-dir") + if Path(args.plan).suffix.lower() == ".md": + parser.error("--plan is the JSON plan; the list is written beside it") return dry_run(args, out)