fix(tools): harden the #247 repair apply after review
- refuse while any other process holds the library, its -wal or -shm open (tools/comfyui_stopped.py), also in a rehearsal; without /proc only with --i-stopped-comfyui; require main.py in --comfyui-dir; check again just before BEGIN IMMEDIATE - the plan records its library (normcase realpath) and max row id; apply refuses on another library - refuse when opening the library changes its schema (run from a checkout at the live code version) - verify the backup (integrity_check and counts) and never overwrite one - target_text is a drift key and must match its hash; keep compares the linked text's hash, not the stored column - the dry run rehearses the orphan deletes when --delete-orphans is chosen - a .md plan path is refused; the list is <stem>.list.md - any error during apply rolls back, names the backup, exits 2 - relpath on another drive falls back to the full path
This commit is contained in:
@@ -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(
|
||||
|
||||
@@ -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
|
||||
+212
-90
@@ -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
|
||||
(``<plan stem>.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 <plan.json>`` (ComfyUI stopped):
|
||||
``--apply <plan.json>`` (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 ``<db>.pre-247-repair_<stamp>`` (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 ``<db>.pre-247-repair_<stamp>`` (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 <db> --apply repair_plan.json \\
|
||||
--comfyui-dir <ComfyUI folder>
|
||||
|
||||
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)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user