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:
Vito Sansevero
2026-10-28 07:16:50 -07:00
parent b7cf8ef221
commit 1a90d3b0ae
3 changed files with 599 additions and 101 deletions
+306 -11
View File
@@ -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(
+81
View File
@@ -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
View File
@@ -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)