diff --git a/tests/test_repair_foreign_links.py b/tests/test_repair_foreign_links.py new file mode 100644 index 00000000..61a81910 --- /dev/null +++ b/tests/test_repair_foreign_links.py @@ -0,0 +1,951 @@ +"""tools/repair_foreign_links.py: the reviewed repair of the #247 dataset links. + +The fixture is a library whose foreign (no PromptManager node) dataset run +wrote shared-prefix files that the old scan linked to the wrong branch text. +Recorded history names each file's saver, so the ladder (``recorded-run``) +names the right branch. +""" + +import contextlib +import hashlib +import io +import json +import os +import shutil +import sqlite3 +import tempfile +import unittest +import unittest.mock +from pathlib import Path + +from PIL import Image +from PIL.PngImagePlugin import PngInfo + +try: + from tests.open_handles import assert_closed +except ImportError: # discovered with tests/ as the top-level directory + from open_handles import assert_closed + +from database.models import PromptModel +from database.operations import PromptDatabase +from tests import saver_anchor_graphs as shapes +from tests.test_run_attribution import RecordedRuns +from tools import repair_foreign_links as tool +from utils.file_prompt import Decision +from utils.hashing import generate_prompt_hash + +TEXTS = shapes.BRANCH_TEXTS +NAME = "IMG_GEN_{:05}_.png" + + +def png(path, graph): + info = PngInfo() + info.add_text("prompt", json.dumps(graph)) + Image.new("RGB", (4, 4)).save(path, pnginfo=info) + + +def sha(text): + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +def file_sha(path): + return hashlib.sha256(Path(path).read_bytes()).hexdigest() + + +def dump(path): + """The database's logical content (independent of WAL checkpoints).""" + conn = sqlite3.connect(path) + try: + return list(conn.iterdump()) + finally: + conn.close() + + +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 + (a wrong row beside a correct row), H relink from a rated prompt, G a + PromptManager run (not a candidate).""" + + def setUp(self): + self.tmp = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self.tmp, True) + self.addCleanup(assert_closed, self, self.tmp) + self.addCleanup(PromptModel.reset_schema_cache) + self.output = os.path.join(self.tmp, "output") + os.makedirs(os.path.join(self.output, "dataset")) + self.db_path = os.path.join(self.tmp, "prompts.db") + 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)) + png(path, shapes.shared_prefix()) + self.files[key] = path + db = PromptDatabase(self.db_path) + try: + self.build(db) + finally: + db.close_all() + os.remove(self.files["D"]) + + def build(self, db): + save = lambda text, **kw: db.save_prompt( # noqa: E731 + text=text, prompt_hash=generate_prompt_hash(text), **kw + ) + self.p0 = save(TEXTS[0], category="scanned", tags=["auto-scanned"]) + self.p2 = save(TEXTS[2], category="scanned", tags=["auto-scanned"]) + self.orphan = save( + "an old guess", + category="scanned", + tags=["auto-scanned"], + notes="Auto-scanned from IMG_GEN_00001_.png", + ) + self.rated = save("a rated guess", rating=5) + link = lambda prompt, key: db.link_image_to_prompt( # noqa: E731 + prompt, self.files[key], metadata={} + ) + self.rows = { + "A": link(self.orphan, "A"), + "B": link(self.p0, "B"), + "C": link(self.p0, "C"), + "D": link(self.p0, "D"), + "E": link(self.p0, "E"), + "F": link(self.p0, "F"), + "F_ok": link(self.p2, "F"), + "H": link(self.rated, "H"), + "G": link(self.p0, "G"), + } + conn = db.model.get_connection() + with conn: + pair = conn.execute( + "INSERT INTO prompt_pairs(positive_prompt_id) VALUES(?)", (self.p0,) + ).lastrowid + conn.execute( + "INSERT INTO image_prompt_pairs(image_id, pair_id, provenance) " + "VALUES(?,?,'{}')", + (self.rows["B"], pair), + ) + name = lambda key: ("dataset", os.path.basename(self.files[key])) # noqa + # Savers of the shared-prefix graph: "13", "23", "33", "43" (branches 1-4) + self.record_run( + db, + self.output, + [name(k) for k in "CDEAEBFH"], + [], + nodes=["13", "13", "13", "23", "23", "33", "33", "43"], + ) + self.record_run(db, self.output, [name("G")], [TEXTS[0]], item=1, nodes=["23"]) + + def dry_run(self, *args): + out = io.StringIO() + code = tool.main( + [ + self.db_path, + "--output-dir", + self.output, + "--plan", + self.plan_path, + *args, + ], + out=out, + ) + return code, out.getvalue() + + def apply(self, *args, probe=not_running, handles=no_handles): + out = io.StringIO() + code = tool.main( + [ + self.db_path, + "--apply", + self.plan_path, + "--comfyui-dir", + os.path.join(self.tmp, "comfy"), + "--backup-dir", + self.backups, + *args, + ], + out=out, + probe=probe, + handles=handles, + ) + return code, out.getvalue() + + def plan(self): + with open(self.plan_path, encoding="utf-8") as stream: + return json.load(stream) + + def write_plan(self, plan): + with open(self.plan_path, "w", encoding="utf-8") as stream: + json.dump(plan, stream) + + def query(self, sql, args=()): + conn = sqlite3.connect(self.db_path) + try: + return conn.execute(sql, args).fetchall() + 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")) + + +class DryRunTests(RepairFixture): + def test_plan_names_relink_keep_and_skip_rows(self): + code, out = self.dry_run() + self.assertEqual(code, 0, out) + items = {item["row_id"]: item for item in self.plan()["rows"]} + rows = self.rows + self.assertNotIn(rows["G"], items, "a PromptManager run is not foreign") + actions = {key: items[rows[key]]["action"] for key in rows if key != "G"} + self.assertEqual( + actions, + { + "A": "relink", + "B": "relink", + "C": "keep", + "D": "skip", + "E": "skip", + "F": "relink", + "F_ok": "keep", + "H": "relink", + }, + ) + a = items[rows["A"]] + self.assertEqual(a["current_prompt_id"], self.orphan) + self.assertEqual(a["current_text_sha256"], sha("an old guess")) + self.assertEqual(a["target_text"], TEXTS[1]) + self.assertEqual(a["target_hash"], generate_prompt_hash(TEXTS[1])) + self.assertEqual(a["target_prompt_id"], "new") + self.assertEqual(a["file"], self.files["A"]) + self.assertEqual(items[rows["B"]]["target_prompt_id"], self.p2) + self.assertEqual(items[rows["H"]]["target_prompt_id"], "new") + self.assertEqual(items[rows["D"]]["reason"], "file gone") + self.assertEqual( + items[rows["E"]]["reason"], "not decided by the recorded saver" + ) + self.assertIn("relink: 4", out) + self.assertIn("keep: 2", out) + self.assertIn("skip: 2", out) + + def test_orphan_audit_lists_the_prompts_losing_every_link(self): + self.dry_run() + plan = self.plan() + orphans = {o["id"]: o for o in plan["orphans"]} + self.assertEqual(set(orphans), {self.orphan, self.rated}) + self.assertEqual(orphans[self.orphan]["blockers"], []) + self.assertTrue(orphans[self.orphan]["safe"]) + self.assertEqual(orphans[self.rated]["blockers"], ["usage or rating"]) + self.assertFalse(orphans[self.rated]["safe"]) + self.assertFalse(plan["delete_orphans"], "off by default") + self.dry_run("--delete-orphans") + self.assertTrue(self.plan()["delete_orphans"]) + + def test_dry_run_writes_nothing_and_writes_the_markdown_list(self): + before = file_sha(self.db_path) + listing = sorted(os.listdir(self.tmp)) + code, out = self.dry_run() + self.assertEqual(code, 0) + self.assertEqual(file_sha(self.db_path), before) + self.assertIn("Dry run: nothing was changed.", out) + # A read-only reader of a WAL library may leave its empty side files + sidecars = {"prompts.db-shm", "prompts.db-wal"} + self.assertEqual( + sorted(set(os.listdir(self.tmp)) - 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.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") + self.assertIn(generate_prompt_hash(TEXTS[1])[:12], markdown) + self.assertIn("new", markdown) + 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() + code, out = self.dry_run() + self.assertEqual(code, 0, out) + + def test_apply_relinks_creates_targets_and_drops_pair_links(self): + code, out = self.apply() + self.assertEqual(code, 0, out) + links = self.links() + new = dict(self.query("SELECT hash, id FROM prompts")) + target_a = new[generate_prompt_hash(TEXTS[1])] + target_h = new[generate_prompt_hash(TEXTS[3])] + self.assertEqual(links[self.rows["A"]], target_a) + self.assertEqual(links[self.rows["B"]], self.p2) + self.assertEqual(links[self.rows["H"]], target_h) + self.assertEqual(links[self.rows["C"]], self.p0, "keep") + self.assertEqual(links[self.rows["E"]], self.p0, "skip") + self.assertEqual(links[self.rows["G"]], self.p0, "not foreign") + self.assertEqual( + self.query( + "SELECT COUNT(*) FROM image_prompt_pairs WHERE image_id=?", + (self.rows["B"],), + ), + [(0,)], + ) + category, notes = self.query( + "SELECT category, notes FROM prompts WHERE id=?", (target_a,) + )[0] + self.assertEqual(category, "scanned") + self.assertEqual( + notes, f"Auto-scanned from {os.path.basename(self.files['A'])}" + ) + tags = self.query( + "SELECT t.name FROM prompt_tags pt JOIN tags t ON t.id=pt.tag_id " + "WHERE pt.prompt_id=?", + (target_a,), + ) + self.assertEqual(tags, [("auto-scanned",)]) + self.assertIn("integrity_check", out) + self.assertIn("foreign_key_check", out) + self.assertRegex(out, r"integrity_check\s+ok\s+ok") + # Orphans are kept unless the plan chose to delete them + self.assertTrue(self.query("SELECT 1 FROM prompts WHERE id=?", (self.orphan,))) + + def test_unique_case_deletes_the_wrong_row(self): + code, out = self.apply() + self.assertEqual(code, 0, out) + links = self.links() + self.assertNotIn(self.rows["F"], links) + self.assertEqual(links[self.rows["F_ok"]], self.p2) + rows = self.query( + "SELECT prompt_id FROM generated_images WHERE image_path=?", + (self.files["F"],), + ) + self.assertEqual(rows, [(self.p2,)]) + + def test_backup_is_taken_before_writing(self): + before = self.links() + code, out = self.apply() + self.assertEqual(code, 0, out) + names = os.listdir(self.backups) + self.assertEqual(len(names), 1) + self.assertRegex(names[0], r"^prompts\.db\.pre-247-repair_\d{8}_\d{6}$") + backup = os.path.join(self.backups, names[0]) + conn = sqlite3.connect(backup) + try: + saved = dict(conn.execute("SELECT id, prompt_id FROM generated_images")) + finally: + conn.close() + self.assertEqual(saved, before) + self.assertIn(f"Backup: {backup}", out) + + def test_backup_beside_the_database_by_default(self): + out = io.StringIO() + code = tool.main( + [ + self.db_path, + "--apply", + self.plan_path, + "--comfyui-dir", + os.path.join(self.tmp, "comfy"), + ], + out=out, + probe=not_running, + handles=no_handles, + ) + self.assertEqual(code, 0, out.getvalue()) + self.assertTrue( + any( + n.startswith("prompts.db.pre-247-repair_") for n in os.listdir(self.tmp) + ) + ) + + def test_refuses_while_comfyui_runs(self): + before = file_sha(self.db_path) + code, out = self.apply(probe=lambda _d, _p: ["process 42 (main.py)"]) + self.assertEqual(code, 2) + self.assertIn("REFUSED", out) + self.assertIn("process 42", out) + self.assertEqual(os.listdir(self.backups), [], "no backup either") + self.assertEqual(file_sha(self.db_path), before) + + def test_rehearsal_skips_the_running_check_on_a_copy_only(self): + code, out = self.apply("--rehearsal", probe=lambda _d, _p: ["port 8188"]) + self.assertEqual(code, 0, out) + self.assertIn("rehearsal", out) + # The live library of that ComfyUI folder: --rehearsal is refused + live_dir = os.path.join(self.tmp, "comfy", "user", "__promptmanager") + os.makedirs(live_dir) + live = os.path.join(live_dir, "prompts.db") + shutil.copy(self.db_path, live) + out = io.StringIO() + code = tool.main( + [ + live, + "--apply", + self.plan_path, + "--comfyui-dir", + os.path.join(self.tmp, "comfy"), + "--rehearsal", + ], + out=out, + probe=lambda _d, _p: [], + handles=no_handles, + ) + self.assertEqual(code, 2) + self.assertIn("--rehearsal", out.getvalue()) + + def test_drift_refuses_and_changes_nothing(self): + cases = { + "row moved": lambda: self.query_write( + "UPDATE generated_images SET prompt_id=? WHERE id=?", + (self.p2, self.rows["H"]), + ), + "text changed": lambda: self.query_write( + "UPDATE prompts SET text='an older guess' WHERE id=?", (self.orphan,) + ), + "file changed": lambda: png(self.files["B"], shapes.branch_bench()), + "target stored since": lambda: self.query_write( + "INSERT INTO prompts(text, hash) VALUES(?,?)", + (TEXTS[1], generate_prompt_hash(TEXTS[1])), + ), + } + pristine = os.path.join(self.tmp, "pristine.db") + shutil.copy(self.db_path, pristine) + image = Path(self.files["B"]).read_bytes() + for name, drift in cases.items(): + 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() + self.assertEqual(code, 2, out) + self.assertIn("REFUSED (rolled back, nothing changed)", out) + self.assertEqual(dump(self.db_path), before) + + def test_late_failure_rolls_back_every_write(self): + plan = self.plan() + plan["delete_orphans"] = True + self.write_plan(plan) + # The safe orphan gained a user tag after the dry run + self.query_write("INSERT INTO tags(name) VALUES('mine')", ()) + self.query_write( + "INSERT INTO prompt_tags(prompt_id, tag_id) " + "SELECT ?, id FROM tags WHERE name='mine'", + (self.orphan,), + ) + before = self.links() + prompts = self.query("SELECT COUNT(*) FROM prompts") + code, out = self.apply() + self.assertEqual(code, 2, out) + self.assertIn(f"prompt {self.orphan}", out) + self.assertIn("user tags", out) + self.assertEqual(self.links(), before, "the relinks were rolled back") + self.assertEqual(self.query("SELECT COUNT(*) FROM prompts"), prompts) + + def test_orphans_are_deleted_only_when_chosen_and_unreferenced(self): + plan = self.plan() + plan["delete_orphans"] = True + self.write_plan(plan) + code, out = self.apply() + self.assertEqual(code, 0, out) + self.assertFalse(self.query("SELECT 1 FROM prompts WHERE id=?", (self.orphan,))) + self.assertTrue( + self.query("SELECT 1 FROM prompts WHERE id=?", (self.rated,)), + "a prompt with a rating is never deleted", + ) + self.assertIn("orphans_deleted", out) + + def test_an_applied_plan_refuses_a_second_time(self): + self.assertEqual(self.apply()[0], 0) + before = dump(self.db_path) + code, out = self.apply() + self.assertEqual(code, 2) + self.assertIn("REFUSED", out) + 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() + + +def _off_text_node(text, save=False): + return { + "class_type": "PromptManagerText", + "inputs": {"text": text, "save_to_library": save}, + } + + +def _encode_from(source): + return { + "class_type": "CLIPTextEncode", + "inputs": {"text": source, "clip": ["1", 1]}, + } + + +def off_behind_pass_through(branch=2): + """The dataset graph with ``branch``'s text coming from a switched-off + PromptManagerText through a pass-through the #90 walk does not follow.""" + graph = shapes.shared_prefix() + graph["50"] = _off_text_node(TEXTS[branch - 1]) + graph["51"] = {"class_type": "MyNode", "inputs": {"value": ["50", 0]}} + graph[str(branch * 10)] = _encode_from(["51", 0]) + return graph + + +def off_target(branch=2, save=False): + """``branch``'s text straight from a PromptManagerText whose switch is + ``save`` (False, or a link to a node the graph cannot read).""" + graph = shapes.shared_prefix() + graph["50"] = _off_text_node(TEXTS[branch - 1], save) + graph[str(branch * 10)] = _encode_from(["50", 0]) + if isinstance(save, list): + graph["52"] = {"class_type": "SomeBoolNode", "inputs": {}} + return graph + + +def off_negative_only(): + """Every branch's negative from a switched-off PromptManagerText.""" + graph = shapes.shared_prefix() + graph["50"] = _off_text_node(shapes.NEGATIVE) + graph["3"] = _encode_from(["50", 0]) + return graph + + +class SaveToLibraryTests(RecordedRuns, unittest.TestCase): + """#90: the repair never relinks a file to a switched-off node's text. + + The plan decides each file through ``decide_file_prompt``; every answer + but a positive from the recorded saver or the inspector is a skip, so a + ``not_saved`` or ``not_saved_wired`` file is never written. Each file is + reported by a run without contexts (as a cached or unrecorded run is), + so only the ladder can stop it. Saver "23" is branch 2. + """ + + def setUp(self): + self.tmp = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self.tmp, True) + self.addCleanup(assert_closed, self, self.tmp) + self.addCleanup(PromptModel.reset_schema_cache) + self.output = os.path.join(self.tmp, "output") + os.makedirs(os.path.join(self.output, "dataset")) + self.db_path = os.path.join(self.tmp, "prompts.db") + + def plan_for(self, graph): + path = os.path.join(self.output, "dataset", NAME.format(1)) + png(path, graph) + db = PromptDatabase(self.db_path) + try: + wrong = db.save_prompt( + text=TEXTS[0], prompt_hash=generate_prompt_hash(TEXTS[0]) + ) + db.link_image_to_prompt(wrong, path, metadata={}) + self.record_run( + db, self.output, [("dataset", NAME.format(1))], [], nodes=["23"] + ) + plan = tool.build_plan(db, self.output, self.db_path) + finally: + db.close_all() + (item,) = plan["rows"] + return item + + def test_an_off_node_behind_a_pass_through_is_skipped(self): + item = self.plan_for(off_behind_pass_through()) + self.assertEqual(item["action"], "skip") + self.assertEqual(item["decision"], "not_saved/none") + self.assertNotIn("target_text", item) + + def test_an_off_node_in_the_recorded_run_target_is_skipped(self): + cases = { + "off": (False, "not_saved/none"), + "wired": (["52", 0], "not_saved_wired/none"), + } + for name, (save, decision) in cases.items(): + with self.subTest(name): + if os.path.exists(self.db_path): + os.remove(self.db_path) + item = self.plan_for(off_target(save=save)) + self.assertEqual(item["action"], "skip") + self.assertEqual(item["decision"], decision) + + def test_an_off_node_on_another_branch_still_skips_the_file(self): + """The ladder judges the file's whole graph: any switched-off text + upstream of a positive makes the file not saved.""" + item = self.plan_for(off_target(branch=3)) + self.assertEqual(item["action"], "skip") + self.assertEqual(item["decision"], "not_saved/none") + + def test_an_off_node_wired_only_to_the_negative_still_repairs(self): + item = self.plan_for(off_negative_only()) + self.assertEqual(item["action"], "relink") + self.assertEqual(item["source"], "recorded-run") + self.assertEqual(item["target_text"], TEXTS[1]) + + def test_a_switched_on_node_in_the_target_still_repairs(self): + """Regression guard: a switch that is on is no reason to skip.""" + item = self.plan_for(off_target(save=True)) + self.assertEqual(item["action"], "relink") + self.assertEqual(item["target_text"], TEXTS[1]) + + +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.""" + + def test_action_by_decision(self): + linked = generate_prompt_hash("linked") + cases = [ + (Decision("positive", "other", source="recorded-run"), "relink"), + (Decision("positive", "linked", source="recorded-run"), "keep"), + (Decision("positive", "linked", source="inspector"), "keep"), + (Decision("positive", "other", source="inspector"), "skip"), + (Decision("positive", "linked", source="legacy-parser"), "skip"), + (Decision("pair", "linked", source="resolver"), "skip"), + (Decision("never"), "skip"), + (Decision("not_saved"), "skip"), + (Decision("not_saved_wired"), "skip"), + (Decision("positive", "other", source="opted-out-negative"), "skip"), + ] + for decision, action in cases: + with self.subTest(decision.source, text=decision.text): + self.assertEqual(tool.action_of(decision, linked), action) + + +class CommandLineTests(unittest.TestCase): + def test_running_probe_reports_a_listening_port(self): + import socket + + with socket.socket() as server: + server.bind(("127.0.0.1", 0)) + server.listen(1) + port = server.getsockname()[1] + found = tool.comfyui_running("/nonexistent/comfy", port) + self.assertTrue(any(str(port) in item for item in found)) + + def test_running_probe_is_quiet_when_nothing_runs(self): + import socket + + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + 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( + ["/nonexistent/prompts.db", "--output-dir", "/tmp"], out=io.StringIO() + ) + + def test_apply_needs_the_comfyui_folder(self): + with tempfile.TemporaryDirectory() as tmp: + db = os.path.join(tmp, "prompts.db") + Path(db).write_bytes(b"") + with ( + self.assertRaises(SystemExit), + contextlib.redirect_stderr(io.StringIO()), + ): + tool.main([db, "--apply", "plan.json"], out=io.StringIO()) + + +if __name__ == "__main__": + unittest.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 new file mode 100644 index 00000000..699faa96 --- /dev/null +++ b/tools/repair_foreign_links.py @@ -0,0 +1,776 @@ +"""Repair the links the old scan gave foreign dataset files (#247), on review. + +Before #247 the output scan linked each file of a foreign run (a workflow +with no PromptManager node) through the legacy parser, so every file of a +shared-prefix dataset run got the first branch's text. Recorded history +knows which save node wrote each file; the shared ladder +(``utils.file_prompt.decide_file_prompt``, scan policy) uses it +(``recorded-run``) to name the right branch. This tool judges every +``generated_images`` row of a file whose recording runs all have no +PromptManager context: + +``relink`` the ladder names a positive by the recorded saver, and its text + is not the linked prompt's +``keep`` the recorded saver or the inspector names the linked prompt +``skip`` anything else (file gone, no single recorded saver, ambiguous, + 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 +(``.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), only on the library the plan was +made from (its real path, recorded in the plan): + +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 + ``scanned``, tag ``auto-scanned``, note ``Auto-scanned from ``), + points each row at its target and drops its pair links (they named the old + positive). A file that already has a row for the target keeps that row and + loses the wrong one (UNIQUE(prompt_id, file_path)). Orphans are deleted only + when the plan chose it and each is still unreferenced. + +Any failed check rolls everything back and exits 2. Before/after table +counts, ``integrity_check`` and ``foreign_key_check`` are printed. + + python3 tools/repair_foreign_links.py --output-dir \\ + --plan repair_plan.json + python3 tools/repair_foreign_links.py --apply repair_plan.json \\ + --comfyui-dir + +Run from a checkout at the live code version (after the merge: main) with +its venv. +""" + +import argparse +import datetime +import hashlib +import json +import logging +import os +import sqlite3 +import sys +import tempfile +from pathlib import Path +from types import SimpleNamespace + +_ROOT = Path(__file__).resolve().parents[1] +if str(_ROOT) not in sys.path: + sys.path.insert(0, str(_ROOT)) + +from database.maintenance import ORPHANED_PROMPTS_SQL, _delete_ids # noqa: E402 +from database.operations import PromptDatabase # noqa: E402 +from database.pair_ingestion import ( # noqa: E402 + SCAN_MARKS, + file_link_rows, + mark_created, +) +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 +from utils.hashing import generate_prompt_hash # noqa: E402 + +PLAN_VERSION = 1 +COMFY_PORT = 8188 +LIVE_LIBRARY = ("user", "__promptmanager", "prompts.db") +EXCERPT = 80 +COUNTED = ( + "prompts", + "generated_images", + "prompt_tags", + "image_prompt_pairs", + "workflow_documents", + "prompt_pairs", + "prompt_library_state", + "execution_run_contexts", + "tags", +) +# The ladder steps whose answer is certain enough to keep a link. +KEEP_SOURCES = ("recorded-run", "inspector") +# What must be unchanged since the dry run for a row to be written. +DRIFT_KEYS = ( + "action", + "file", + "current_prompt_id", + "current_text_sha256", + "file_size", + "file_mtime_ns", + "target_text", + "target_hash", + "target_prompt_id", +) +# Locations reported only by runs without any PromptManager context. +_FOREIGN_LOCATIONS = ( + "SELECT o.subfolder, o.filename FROM execution_run_outputs o " + "WHERE o.folder_type = 'output' AND o.root_key = ? AND o.path_style = ? " + "GROUP BY o.subfolder, o.filename " + "HAVING SUM(EXISTS (SELECT 1 FROM execution_run_contexts c " + "WHERE c.run_id = o.run_id)) = 0 " + "ORDER BY o.subfolder, o.filename" +) + + +class Refused(Exception): + """A check failed: nothing is written.""" + + +def sha256(text): + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +def metadata_reader(): + """The scan's metadata extraction (PNG text chunks, ffprobe for videos).""" + api = SimpleNamespace(logger=logging.getLogger("prompt_manager.repair_247")) + return lambda path: PromptManagerAPI._extract_comfyui_metadata(api, path) or {} + + +def foreign_rows(db, output_dir): + """The link rows of every file only foreign runs reported, by id.""" + style = "windows" if os.name == "nt" else "posix" + root_key = snapshot_output_roots(style, output_dir, output_dir).output_key + conn = db.model.get_connection() + rows = {} + for subfolder, filename in conn.execute(_FOREIGN_LOCATIONS, (root_key, style)): + path = os.path.join(output_dir, subfolder, filename) + rows.update((row["id"], row) for row in file_link_rows(conn, path)) + return [rows[key] for key in sorted(rows)] + + +class Judge: + """One file's decision per path, for all of its rows.""" + + def __init__(self, db, output_dir, read=None): + self.db = db + self.output_dir = output_dir + self.read = read or metadata_reader() + self.decisions = {} + + def decide(self, path, size): + if path not in self.decisions: + self.decisions[path] = decide_file_prompt( + self.read(path), + file_name=os.path.basename(path), + file_size=size, + in_output=True, + lookup=lambda: recorded_output(self.db, path, self.output_dir), + policy=Policy.SCAN, + ) + return self.decisions[path] + + def row(self, row): + """The plan entry of one ``generated_images`` row.""" + conn = self.db.model.get_connection() + path = row["image_path"] + (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, + "current_prompt_id": row["prompt_id"], + "current_text_sha256": sha256(text), + } + try: + stat = os.stat(path) + except OSError: + return {**item, "action": "skip", "reason": "file gone"} + item.update(file_size=stat.st_size, file_mtime_ns=stat.st_mtime_ns) + decision = self.decide(path, stat.st_size) + action = action_of(decision, prompt_hash) + if action == "skip": + return { + **item, + "action": "skip", + "reason": "not decided by the recorded saver", + "decision": f"{decision.kind}/{decision.source}", + } + target_hash = generate_prompt_hash(decision.text) + stored = conn.execute( + "SELECT id FROM prompts WHERE hash = ?", (target_hash,) + ).fetchone() + item.update( + target_text=decision.text, + target_hash=target_hash, + target_prompt_id=stored[0] if stored else "new", + action=action, + source=decision.source, + ) + return item + + +def action_of(decision, linked_hash): + """``relink``, ``keep`` or ``skip`` for a row linked to ``linked_hash``. + + Only the recorded saver (``recorded-run``) names a new target. A certain + step (the recorded saver or the inspector) naming the linked prompt keeps + it; the legacy floor and every other answer are skipped. + """ + if decision.kind != POSITIVE or decision.source not in KEEP_SOURCES: + return "skip" + if generate_prompt_hash(decision.text) == linked_hash: + return "keep" + return "relink" if decision.source == "recorded-run" else "skip" + + +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": [], + "delete_orphans": False, + } + + +def blockers(conn, prompt_id, unused): + """Why ``prompt_id`` must not be deleted (empty: nothing references it).""" + found = [] + values = {"id": prompt_id, "tag": SCAN_MARKS.tags[0]} + checks = { + "linked images": "SELECT COUNT(*) FROM generated_images WHERE prompt_id=:id", + "saved pairs": "SELECT COUNT(*) FROM prompt_pairs " + "WHERE positive_prompt_id=:id OR negative_prompt_id=:id", + "recorded runs": "SELECT COUNT(*) FROM execution_run_contexts " + "WHERE positive_prompt_id=:id OR negative_prompt_id=:id", + "library state": "SELECT COUNT(*) FROM prompt_library_state " + "WHERE prompt_id=:id", + "user tags": "SELECT COUNT(*) FROM prompt_tags pt JOIN tags t " + "ON t.id=pt.tag_id WHERE pt.prompt_id=:id AND t.name != :tag", + } + for name, sql in checks.items(): + count = conn.execute(sql, values).fetchone()[0] + if count: + found.append(f"{name}: {count}") + run_count, last_used, rating, category, notes = conn.execute( + "SELECT run_count, last_used_at, rating, category, notes FROM prompts " + "WHERE id=?", + (prompt_id,), + ).fetchone() + if run_count or last_used or rating: + found.append("usage or rating") + if category not in (None, SCAN_MARKS.category): + found.append(f"category {category}") + if notes and not notes.startswith("Auto-scanned from "): + found.append("notes") + if prompt_id not in unused: + found.append("not unused by the app's rule") + return found + + +def _rejudge(conn, judge, plan): + for item in plan["rows"]: + if item["action"] == "skip": + continue + row = conn.execute( + "SELECT * FROM generated_images WHERE id = ?", (item["row_id"],) + ).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): + raise Refused( + f"row {item['row_id']}: {key} changed " + f"({item.get(key)!r} -> {fresh.get(key)!r})" + ) + + +def _targets(conn, db, relinks): + """``{target hash: prompt id}``, creating each "new" target once.""" + targets = {} + for item in relinks: + digest = item["target_hash"] + if digest in targets: + continue + if item["target_prompt_id"] != "new": + targets[digest] = item["target_prompt_id"] + continue + prompt_id, created = db._pair_component_outcome(conn, item["target_text"]) + if not created: + raise Refused(f"row {item['row_id']}: target prompt already stored") + mark_created(conn, db, prompt_id, SCAN_MARKS, os.path.basename(item["file"])) + targets[digest] = prompt_id + return targets + + +def _relink(conn, item, target): + """Point the row at ``target``; True when the wrong row was deleted instead.""" + row_id = item["row_id"] + correct = [ + row + for row in file_link_rows(conn, item["file"]) + if row["prompt_id"] == target and row["id"] != row_id + ] + if correct: # UNIQUE(prompt_id, file_path): the correct row stays alone + if _delete_ids(conn, "generated_images", [row_id]) != 1: + raise Refused(f"row {row_id}: delete touched no row") + return True + # A pair link records the old positive: it goes with the old link + conn.execute("DELETE FROM image_prompt_pairs WHERE image_id = ?", (row_id,)) + cursor = conn.execute( + "UPDATE generated_images SET prompt_id = ? WHERE id = ? AND prompt_id = ?", + (target, row_id, item["current_prompt_id"]), + ) + if cursor.rowcount != 1: + raise Refused(f"row {row_id}: relink touched {cursor.rowcount} rows") + return False + + +def _audit(conn, relinks): + """The prompts the relinks leave without any link, with their blockers.""" + sources = {} + for item in relinks: + sources.setdefault(item["current_prompt_id"], []).append(item) + unused = {row[0] for row in conn.execute(ORPHANED_PROMPTS_SQL)} + audit = [] + for prompt_id in sorted(sources): + if conn.execute( + "SELECT 1 FROM generated_images WHERE prompt_id = ?", (prompt_id,) + ).fetchone(): + continue + (text,) = conn.execute( + "SELECT text FROM prompts WHERE id = ?", (prompt_id,) + ).fetchone() + found = blockers(conn, prompt_id, unused) + audit.append( + { + "id": prompt_id, + "text_excerpt": " ".join(text.split())[:EXCERPT], + "text_sha256": sha256(text), + "images_before": len(sources[prompt_id]), + "blockers": found, + "safe": not found, + } + ) + return audit + + +def _delete_orphans(conn, plan): + chosen = [orphan for orphan in plan["orphans"] if orphan["safe"]] + unused = {row[0] for row in conn.execute(ORPHANED_PROMPTS_SQL)} + for orphan in chosen: + row = conn.execute( + "SELECT text FROM prompts WHERE id = ?", (orphan["id"],) + ).fetchone() + if row is None or sha256(row[0]) != orphan["text_sha256"]: + raise Refused(f"prompt {orphan['id']}: gone or text changed") + found = blockers(conn, orphan["id"], unused) + if found: + raise Refused(f"prompt {orphan['id']}: {', '.join(found)}") + ids = [orphan["id"] for orphan in chosen] + if _delete_ids(conn, "prompts", ids) != len(ids): + raise Refused("an orphan prompt was not deleted") + return len(ids) + + +def apply_plan(db, plan, read=None, before_write=None): + """Write ``plan`` in one transaction; returns ``(result, orphan audit)``. + + ``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) + dropped = sum( + _relink(conn, item, targets[item["target_hash"]]) for item in relinks + ) + audit = _audit(conn, relinks) + deleted = _delete_orphans(conn, plan) if plan.get("delete_orphans") else 0 + created = sum(1 for item in relinks if item["target_prompt_id"] == "new") + result = { + "relinked": len(relinks) - dropped, + "duplicates_dropped": dropped, + "prompts_created": len( + {i["target_hash"] for i in relinks if i["target_prompt_id"] == "new"} + ), + "rows_to_new_prompts": created, + "kept": sum(1 for item in plan["rows"] if item["action"] == "keep"), + "orphans_deleted": deleted, + } + return result, audit + + +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 = { + table: conn.execute(f'SELECT COUNT(*) FROM "{table}"').fetchone()[0] + for table in COUNTED # nosec B608 - fixed table names + } + found["integrity_check"] = conn.execute("PRAGMA integrity_check").fetchone()[0] + found["foreign_key_check"] = len( + conn.execute("PRAGMA foreign_key_check").fetchall() + ) + return found + finally: + conn.close() + + +def _print_counts(before, after, out): + out.write(f"{'':24}{'before':>12}{'after':>12}\n") + for key in before: + out.write(f"{key:24}{before[key]!s:>12}{after[key]!s:>12}\n") + + +def _summary(plan): + actions = [item["action"] for item in plan["rows"]] + return " ".join( + f"{name}: {actions.count(name)}" for name in ("relink", "keep", "skip") + ) + + +def _cell(text, limit=EXCERPT): + """One table cell: whitespace folded, cut to ``limit`` (None: whole).""" + folded = " ".join(str(text).split()) + 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: _relative(path, root) # noqa: E731 + rows = plan["rows"] + lines = [ + "# #247 repair plan (dry run)", + "", + f"Created {plan['created']}; output folder `{root}`. {_summary(plan)}.", + "", + f"## Relink ({sum(i['action'] == 'relink' for i in rows)})", + "", + "| row | file | linked prompt | linked text | target | target hash " + "| target text |", + "|---:|---|---:|---|---|---|---|", + ] + for item in (i for i in rows if i["action"] == "relink"): + lines.append( + f"| {item['row_id']} | {_cell(rel(item['file']), None)} | " + f"{item['current_prompt_id']} | " + f"{_cell(prompt_texts.get(item['current_prompt_id'], ''))} | " + f"{item['target_prompt_id']} | `{item['target_hash'][:12]}` | " + f"{_cell(item['target_text'])} |" + ) + lines += ["", f"## Keep ({sum(i['action'] == 'keep' for i in rows)})", ""] + lines += ["| row | file | prompt |", "|---:|---|---:|"] + for item in (i for i in rows if i["action"] == "keep"): + lines.append( + f"| {item['row_id']} | {_cell(rel(item['file']), None)} | " + f"{item['current_prompt_id']} |" + ) + lines += ["", f"## Skip ({sum(i['action'] == 'skip' for i in rows)})", ""] + lines += ["| row | file | prompt | reason |", "|---:|---|---:|---|"] + for item in (i for i in rows if i["action"] == "skip"): + reason = item["reason"] + ( + f" ({item['decision']})" if item.get("decision") else "" + ) + lines.append( + f"| {item['row_id']} | {_cell(rel(item['file']), None)} | " + f"{item['current_prompt_id']} | {reason} |" + ) + choice = "ON (--delete-orphans)" if plan["delete_orphans"] else "OFF" + lines += [ + "", + f"## Prompts that would lose every link ({len(plan['orphans'])})", + "", + f"Deleting the safe ones: {choice}.", + "", + "| id | text | images before | blockers | safe to delete |", + "|---:|---|---:|---|---|", + ] + for orphan in plan["orphans"]: + lines.append( + f"| {orphan['id']} | {_cell(orphan['text_excerpt'])} | " + f"{orphan['images_before']} | {', '.join(orphan['blockers']) or 'none'} | " + f"{'yes' if orphan['safe'] else 'no'} |" + ) + simulated = plan.get("simulated") + if simulated: + lines += ["", "## Rehearsed on the dry run's copy", ""] + lines += ["| table | before | after |", "|---|---:|---:|"] + for key, value in simulated["before"].items(): + lines.append(f"| {key} | {value} | {simulated['after'][key]} |") + 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, "judge.db") + backup_database(args.database, copy) + db = PromptDatabase(copy) + try: + plan = build_plan(db, os.path.abspath(args.output_dir), args.database) + texts = _texts(db) + finally: + db.close_all() + 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_name(plan_path.stem + ".list.md") + list_path.write_text(markdown(plan, texts), encoding="utf-8") + out.write(f"{_summary(plan)}\n") + out.write( + f"Prompts that would lose every link: {len(audit)} " + f"({sum(o['safe'] for o in audit)} safe to delete; deleting " + f"{'ON' if plan['delete_orphans'] else 'OFF'})\n" + ) + out.write(f"Plan: {plan_path}\nList: {list_path}\n") + out.write("Dry run: nothing was changed.\n") + return 0 + + +def _is_live(database, comfyui_dir): + live = os.path.join(comfyui_dir, *LIVE_LIBRARY) + return library_identity(database) == library_identity(live) + + +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) + refusal = _preflight(args, plan, probe, handles) + if refusal: + out.write(f"REFUSED: {refusal}\n") + return 2 + if args.rehearsal: + out.write("(rehearsal on a copy: ComfyUI's process and port not checked)\n") + before = counts(args.database) + try: + 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() + after = counts(args.database) + _print_counts(before, after, out) + out.write(f"Applied: {json.dumps(result, sort_keys=True)}\n") + if after["integrity_check"] != "ok" or after["foreign_key_check"]: + out.write(f"WARNING: integrity or foreign-key check failed; restore {backup}\n") + return 1 + return 0 + + +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)") + parser.add_argument( + "--plan", default="repair_foreign_links_plan.json", help="plan to write" + ) + parser.add_argument( + "--delete-orphans", + action="store_true", + help="record in the plan that the safe orphan prompts are deleted", + ) + parser.add_argument("--apply", metavar="PLAN", help="write this reviewed plan") + parser.add_argument("--comfyui-dir", help="the ComfyUI folder (apply)") + parser.add_argument("--backup-dir", help="where the backup goes (apply)") + parser.add_argument("--port", type=int, default=COMFY_PORT) + parser.add_argument( + "--rehearsal", + 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, 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) + + +if __name__ == "__main__": + sys.exit(main())