Files
ComfyAssets-ComfyUI_PromptM…/tests/test_library_move.py
T
Vito 8bf759d6f8 feat(system): keep the library in ComfyUI's user folder (#318)
* feat(system): resolve every library path in one place

utils/library_location.py decides where the library lives: the
PROMPT_MANAGER_DATA_DIR folder (validated: absolute, never inside a folder
ComfyUI or Prompt Manager serves), an explicit database.default_path as
before, a node-folder library that was not moved, else ComfyUI's
user/__promptmanager. The database, diagnostics, config.json, logs and the
LoRA image cache all ask it, so a new install never writes into the node
folder and an existing one keeps every path until the user moves it.
"Keep it here" is persisted as library.keep_in_node_folder.

* feat(system): move the library as a verified Activity job

POST /prompt_manager/system/library/move starts "Move library", an
exclusive, non-cancellable job (database/library_move.py). It checks the
target volume's free space, copies the database with SQLite's backup API
from one read snapshot, verifies the copy (integrity check, every workflow
document, the same rows in every table), copies the safety and
pre-migration backups and config.json (owner-only, checked by size and
SHA-256), then publishes under the old file's write lock and switches the
live process (every PromptModel, the review and history observers, the
API's plans and health report). The old prompts.db and config.json are
renamed *.moved-to-<folder> only after the switch. A failure removes
everything it wrote and leaves the old library in use; a marker in the new
folder makes a restart after any crash pick the right library.

GET /prompt_manager/system/library and the Overview's "location" say where
the library lives and why; POST .../keep persists "Keep it here".

* feat(system): show where the library lives and offer the move

D129. System > Overview's Library health gains a "Library location" row
(the database path, why it is there, the "Keep it here" choice and any
PROMPT_MANAGER_DATA_DIR problem) with Move library for a node-folder
library and Details otherwise. While a node-folder library is not kept
there, the workspace shows the concept's .notice above every page:
"Keep it here" saves the choice on the server and stops the notice;
"Move library" opens the shared review (From, To, Size, what moves, what
happens to the old files) and starts the exclusive job in Activity. The
notice and the Overview read the location again when the job ends.

The Overview comparison maps the inserted row; a new comparison holds the
notice to the concept's Sources notice (markup, tokens, alignment, both
themes and widths) with a mutation check.

* feat(system): close the browser server's log files before its root goes

The logs now live in the library's data folder, which for the browser test
server is <root>/user/__promptmanager; Windows could not remove the temp
root while the server still held prompt_manager.log open. The API test's
node folder is spelled in lower case because public paths are case-folded
on Windows.

* feat(system): show the library location on the legacy admin page

A line under the admin page's statistics names the database path and why it is there (ComfyUI's user folder, PROMPT_MANAGER_DATA_DIR, config.json, or the node folder with a pointer to Move library), read once from GET /prompt_manager/system/library.

* feat(system): prove a fresh host keeps its library in user/__promptmanager

The execution suite starts one pinned CPU ComfyUI with no configuration at all on an OS-chosen port: the library is created in user/__promptmanager (database and logs), nothing but byte code is added to the node folder, the location route reports the default folder, and ComfyUI's /userdata API serves none of it (traversal spellings, a listing, a system-user header).

* feat(system): document where the library lives and how it moves

README: where the library lives, the upgrade path (Move library / Keep it here), PROMPT_MANAGER_DATA_DIR, explicit paths unchanged, and the way back to an older version. TESTING.md: D129 and the library location contract (resolution order, move steps, crash windows, coded refusals, Windows handling).

* wip(system): library move fix round in progress

* fix(system): retry a fallback image link that races the library switch

The image monitor's fallback and recent-prompt links went through
PromptDatabase.link_image_to_prompt, which logs and returns 0 on any
error, so survive_switch never saw a "database is locked" or closed
connection during Move library and the link was lost in both files.
The monitor now links through link_image_to_prompt_or_raise (one
transaction per attempt); link_image_to_prompt keeps its
log-and-return-0 contract for other callers. A link that returns 0 is
no longer logged as "Successfully linked image 0".

A switch that fails after repointing the models now closes the old
file's detached connections instead of leaving them to GC: switch()
adds them to a list the mover owns.

The mocked retry test is replaced by a real-lock race asserting the
image row lands in the new library only.

* fix(system): say which source is the default and where a symlinked library lives

Sources & scans: with no source configured, ComfyUI's output folder now
reads "ComfyUI/output · watching · default" on its own row instead of a
separate note line. GET /settings marks each effective source with
in_comfyui (its public path is relative to the ComfyUI base) so only
those get the "ComfyUI/" prefix; the scan still sends the server's
public path.

Library location: a library in the node folder reads through the node
folder's unresolved place in custom_nodes, so a node folder symlinked to
a checkout outside ComfyUI shows custom_nodes/promptmanager/prompts.db
instead of prompts.db (System library row, Move review, admin Library
line). Never absolute or "..", else the shared _public_path, which
sources and jobs keep using unchanged.
2026-10-24 12:54:09 -07:00

498 lines
20 KiB
Python

"""Move library (#249 S6b): copy, verify, switch, then retire the old files.
Each test builds a node-folder library (a real PromptDatabase with prompts,
images and workflow documents, a config.json holding a key, a pooled safety
backup and the pre-migration backup) and a ComfyUI user folder in a temp
directory, then moves it with ``database.library_move``. Failures and crash
windows are injected through the move's ``checkpoint`` hook: an ``Exception``
is a failure the move cleans up after, ``SimulatedCrash`` (a BaseException)
stops it dead like a killed process.
"""
import hashlib
import json
import os
import sqlite3
import sys
import tempfile
import threading
import types
import unittest
from unittest import mock
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from database import library_move as lm # noqa: E402
from database.execution_history import ( # noqa: E402
notify_history_restore,
subscribe_history_restore,
)
from database.models import PromptModel # noqa: E402
from database.operations import PromptDatabase # noqa: E402
from database.review_changes import ( # noqa: E402
notify_review_change,
subscribe_review_changes,
)
from tests.open_handles import assert_closed, open_handles_under # noqa: E402
from utils import library_location as ll # noqa: E402
KEY = "civitai-secret-key"
POOL = "prompts.db.backup_20261001_101010_000001"
PRE = "prompts.db.pre-workflow-storage_20261002_101010_000002"
class SimulatedCrash(BaseException):
"""The process died here: nothing after this point runs."""
def _digest(path):
with open(path, "rb") as handle:
return hashlib.sha256(handle.read()).hexdigest()
def _rows(path):
"""Every table's rows of ``path`` (read-only), to compare libraries."""
uri = "file:" + path.replace(os.sep, "/") + "?mode=ro"
conn = sqlite3.connect(uri, uri=True)
try:
tables = [
row[0]
for row in conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' "
"AND name NOT LIKE 'sqlite_%' ORDER BY name"
)
]
return {
table: sorted(
repr(tuple(row)) for row in conn.execute(f'SELECT * FROM "{table}"')
)
for table in tables
}
finally:
conn.close()
def _workflow(seed):
return {
"nodes": [
{"id": n, "type": f"Node{seed}", "widgets": [seed, n]} for n in range(30)
]
}
class _MoveCase(unittest.TestCase):
def setUp(self):
self._tmp = tempfile.TemporaryDirectory()
self.base = os.path.realpath(self._tmp.name)
self.comfy = os.path.join(self.base, "ComfyUI")
self.ext = os.path.join(self.comfy, "custom_nodes", "ComfyUI_PromptManager")
os.makedirs(self.ext)
self.fp = types.SimpleNamespace()
for name in ("output", "input", "temp", "user"):
folder = os.path.join(self.comfy, name)
os.makedirs(folder)
setattr(self.fp, f"get_{name}_directory", lambda f=folder: f)
self.target = os.path.join(self.comfy, "user", "__promptmanager")
self.old_db = os.path.join(self.ext, "prompts.db")
self.new_db = os.path.join(self.target, "prompts.db")
ll.clear_override()
self.db = PromptDatabase(self.old_db)
for index in range(3):
prompt_id = self.db.save_prompt(f"a cat number {index}", category="cats")
self.db.link_image_to_prompt(
str(prompt_id),
os.path.join(self.comfy, "output", f"cat{index}.png"),
{"workflow": _workflow(index % 2), "prompt": {"1": {"seed": index}}},
)
self.db.model.close_all()
self.config = os.path.join(self.ext, "config.json")
with open(self.config, "w", encoding="utf-8") as handle:
json.dump(
{"integrations": {"lora_manager": {"civitai_api_key": KEY}}}, handle
)
for name, stamp in ((POOL, 1_700_000_000), (PRE, 1_700_000_100)):
path = os.path.join(self.ext, name)
source = sqlite3.connect(self.old_db)
copy = sqlite3.connect(path)
try:
source.backup(copy)
finally:
copy.close()
source.close()
os.utime(path, (stamp, stamp))
self.before = {
"rows": _rows(self.old_db),
"config": _digest(self.config),
"pool": _digest(os.path.join(self.ext, POOL)),
"pre": _digest(os.path.join(self.ext, PRE)),
}
self.config_lock = threading.RLock()
def tearDown(self):
import gc
gc.collect()
PromptModel.close_all_instances()
ll.clear_override()
try:
assert_closed(self, self.ext, self.target)
finally:
self._tmp.cleanup()
def location(self, **kwargs):
return ll.resolve_library(
None, environ={}, extension_root=self.ext, folder_paths=self.fp, **kwargs
)
def move(self, **kwargs):
plan = lm.plan_move(self.location())
kwargs.setdefault("config_lock", self.config_lock)
return lm.move_library(plan, **kwargs)
def assert_old_library_untouched(self):
self.assertTrue(os.path.isfile(self.old_db))
self.assertEqual(_rows(self.old_db), self.before["rows"])
self.assertEqual(_digest(self.config), self.before["config"])
self.assertEqual(_digest(os.path.join(self.ext, POOL)), self.before["pool"])
self.assertEqual(_digest(os.path.join(self.ext, PRE)), self.before["pre"])
self.assertEqual(
os.path.normcase(self.db.model.db_path), os.path.normcase(self.old_db)
)
self.assertEqual(self.location().source, ll.SOURCE_NODE_FOLDER)
def assert_target_clean(self):
leftovers = os.listdir(self.target) if os.path.isdir(self.target) else []
self.assertEqual(leftovers, [])
class TestMove(_MoveCase):
def test_move_copies_verifies_switches_and_retires_the_old_files(self):
result = self.move()
self.assertEqual(_rows(self.new_db), self.before["rows"])
self.assertEqual(result.counts["prompts"], 3)
self.assertEqual(result.counts["generated_images"], 3)
self.assertEqual(
result.counts["workflow_documents"], 5
) # 2 workflows, 3 graphs
self.assertEqual(
sorted(result.files), sorted(["config.json", POOL, PRE, "prompts.db"])
)
self.assertIsNone(result.retire_problem)
# The live database object now reads and writes the new file.
self.assertEqual(self.db.model.db_path, self.new_db)
self.db.save_prompt("written after the move")
self.assertEqual(len(_rows(self.new_db)["prompts"]), 4)
image = self.db.get_prompt_images(
str(self.db.search_prompts(text="number 1")[0]["id"])
)[0]
self.assertEqual(image["workflow_data"], _workflow(1))
# config.json, the pooled backup and the pre-migration backup: same bytes.
self.assertEqual(
_digest(os.path.join(self.target, "config.json")), self.before["config"]
)
self.assertEqual(_digest(os.path.join(self.target, POOL)), self.before["pool"])
self.assertEqual(_digest(os.path.join(self.target, PRE)), self.before["pre"])
self.assertEqual(
int(os.path.getmtime(os.path.join(self.target, POOL))), 1_700_000_000
)
# The old files are renamed, never deleted; the backups stay.
renamed = sorted(os.listdir(self.ext))
self.assertIn("prompts.db.moved-to-__promptmanager", renamed)
self.assertIn("config.json.moved-to-__promptmanager", renamed)
self.assertNotIn("prompts.db", renamed)
self.assertNotIn("config.json", renamed)
self.assertFalse(any(name.endswith(("-wal", "-shm")) for name in renamed))
self.assertEqual(
_rows(os.path.join(self.ext, "prompts.db.moved-to-__promptmanager")),
self.before["rows"],
)
self.assertEqual(ll.read_marker(self.target)["state"], ll.MARKER_MOVED)
def test_new_files_are_private(self):
if os.name != "posix":
self.skipTest("POSIX permissions")
self.move()
self.assertEqual(os.stat(self.target).st_mode & 0o777, 0o700)
for name in ("prompts.db", "config.json", POOL, PRE, ll.MARKER_NAME):
mode = os.stat(os.path.join(self.target, name)).st_mode & 0o777
self.assertEqual(mode, 0o600, name)
def test_restart_uses_the_new_library_and_never_the_renamed_file(self):
self.move()
self.assertEqual(self.location().db_path, self.new_db)
ll.clear_override()
restarted = self.location()
self.assertEqual(restarted.source, ll.SOURCE_DEFAULT)
self.assertEqual(restarted.db_path, self.new_db)
self.assertEqual(
restarted.config_path, os.path.join(self.target, "config.json")
)
def test_every_open_database_object_is_repointed(self):
node = PromptDatabase(self.old_db)
monitor = PromptDatabase(self.old_db)
monitor.get_prompt_by_id(1)
other = PromptDatabase(os.path.join(self.base, "other.db"))
self.move()
self.assertEqual(node.model.db_path, self.new_db)
self.assertEqual(monitor.model.db_path, self.new_db)
self.assertEqual(monitor.get_prompt_by_id(1)["text"], "a cat number 0")
self.assertEqual(other.model.db_path, os.path.join(self.base, "other.db"))
assert_closed(
self, os.path.join(self.ext, "prompts.db.moved-to-__promptmanager")
)
def test_change_observers_follow_the_library(self):
reviews, restores = [], []
unsubscribe_review = subscribe_review_changes(
self.old_db, lambda: reviews.append(1)
)
unsubscribe_restore = subscribe_history_restore(
self.old_db, lambda: restores.append(1)
)
self.addCleanup(unsubscribe_review)
self.addCleanup(unsubscribe_restore)
self.move()
self.assertEqual(restores, [1]) # the history readers start a new epoch
notify_review_change(self.new_db)
notify_history_restore(self.new_db)
self.assertEqual(reviews, [1])
self.assertEqual(restores, [1, 1])
def test_files_are_only_renamed_inside_one_folder(self):
real_replace, real_rename = os.replace, os.rename
moves = []
def replace(src, dst, *args, **kwargs):
moves.append((src, dst))
return real_replace(src, dst, *args, **kwargs)
def rename(src, dst, *args, **kwargs):
moves.append((src, dst))
return real_rename(src, dst, *args, **kwargs)
with (
mock.patch.object(lm.os, "replace", replace),
mock.patch.object(lm.os, "rename", rename),
):
self.move()
self.assertTrue(moves)
for src, dst in moves:
self.assertEqual(os.path.dirname(src), os.path.dirname(dst), (src, dst))
def test_a_write_during_the_copy_is_copied_again(self):
writes = []
def checkpoint(step):
if step == "files" and not writes:
writes.append(PromptDatabase(self.old_db).save_prompt("late write"))
progress = []
self.move(checkpoint=checkpoint, progress=progress.append)
self.assertIn("late write", repr(_rows(self.new_db)["prompts"]))
self.assertTrue(any("changed" in line for line in progress), progress)
def test_a_library_that_keeps_changing_is_not_moved(self):
def checkpoint(step):
if step == "files":
PromptDatabase(self.old_db).save_prompt(f"write {os.urandom(4).hex()}")
with self.assertRaises(lm.MoveFailed) as caught:
self.move(checkpoint=checkpoint, attempts=2)
self.assertEqual(caught.exception.code, "library-busy")
self.assertEqual(self.location().source, ll.SOURCE_NODE_FOLDER)
self.assert_target_clean()
def test_old_file_still_open_elsewhere_is_left_in_place(self):
reader = sqlite3.connect(self.old_db)
self.addCleanup(reader.close)
reader.execute("BEGIN")
reader.execute("SELECT COUNT(*) FROM prompts").fetchone()
PromptDatabase(self.old_db).save_prompt("in the wal")
result = self.move()
self.assertEqual(result.retire_problem, "old-in-use")
self.assertTrue(os.path.isfile(self.old_db))
self.assertEqual(self.db.model.db_path, self.new_db)
reader.rollback()
reader.close()
ll.clear_override()
self.assertEqual(self.location().source, ll.SOURCE_DEFAULT)
class TestFailures(_MoveCase):
def _fail_at(self, step):
def checkpoint(name):
if name == step:
raise RuntimeError(f"failure at {step}")
return checkpoint
def test_a_failure_at_any_step_leaves_the_old_library_in_use_and_untouched(self):
for step in ("copied", "verified", "files", "locked", "pending", "published"):
with self.subTest(step=step):
with self.assertRaises(RuntimeError):
self.move(checkpoint=self._fail_at(step))
self.assert_old_library_untouched()
self.assert_target_clean()
assert_closed(self, self.target)
def test_a_copy_that_does_not_verify_is_refused(self):
def checkpoint(step):
if step == "copied":
(temp,) = [
name for name in os.listdir(self.target) if name.endswith(".db")
]
with open(os.path.join(self.target, temp), "r+b") as handle:
handle.truncate(4096)
with self.assertRaises(lm.MoveFailed) as caught:
self.move(checkpoint=checkpoint)
self.assertEqual(caught.exception.code, "verify-failed")
self.assertNotIn(self.base, caught.exception.message)
self.assert_old_library_untouched()
self.assert_target_clean()
def test_a_copied_file_that_differs_is_refused(self):
def checkpoint(step):
if step == "files":
for name in os.listdir(self.target):
if name.endswith(".bak"):
with open(os.path.join(self.target, name), "ab") as handle:
handle.write(b"x")
with self.assertRaises(lm.MoveFailed) as caught:
self.move(checkpoint=checkpoint)
self.assertEqual(caught.exception.code, "copy-failed")
self.assert_old_library_untouched()
self.assert_target_clean()
def test_not_movable_locations_are_refused(self):
os.makedirs(self.target)
with open(self.new_db, "wb") as handle:
handle.write(b"someone else's library")
with self.assertRaises(lm.MoveFailed) as caught:
lm.plan_move(self.location())
self.assertEqual(caught.exception.code, "target-has-library")
os.remove(self.new_db)
default = ll.LibraryLocation(
ll.SOURCE_DEFAULT, self.new_db, "c", self.target, self.target
)
with self.assertRaises(lm.MoveFailed) as caught:
lm.plan_move(default)
self.assertEqual(caught.exception.code, "not-movable")
def test_target_that_is_the_node_folder_is_refused(self):
link = os.path.join(self.base, "link-to-node-folder")
try:
os.symlink(self.ext, link)
except (OSError, NotImplementedError):
self.skipTest("symlinks are not available")
location = ll.LibraryLocation(
ll.SOURCE_NODE_FOLDER, self.old_db, self.config, self.ext, link
)
with self.assertRaises(lm.MoveFailed) as caught:
lm.plan_move(location)
self.assertEqual(caught.exception.code, "same-folder")
def test_disk_check_measures_the_target_volume(self):
measured = []
def free(path):
measured.append(path)
return 1024
with mock.patch.object(lm, "_free_bytes", free):
plan = lm.plan_move(self.location())
with self.assertRaises(lm.MoveFailed) as caught:
lm.move_library(plan, config_lock=self.config_lock)
self.assertEqual(caught.exception.code, "insufficient-space")
self.assertTrue(measured)
for path in measured:
self.assertTrue(
os.path.normcase(self.comfy) in os.path.normcase(path)
and "custom_nodes" not in path,
path,
)
self.assertGreater(plan.required_bytes, os.path.getsize(self.old_db))
self.assert_old_library_untouched()
class TestCrashWindows(_MoveCase):
"""A crash between steps never loses the library or picks a stale one."""
def _crash_at(self, step):
def checkpoint(name):
if name == step:
raise SimulatedCrash(step)
return checkpoint
def _restart(self):
PromptModel.close_all_instances()
ll.clear_override()
return self.location()
def test_crash_before_the_switch_keeps_the_old_library_and_a_retry_succeeds(self):
for step in ("copied", "files", "pending", "published"):
with self.subTest(step=step):
with self.assertRaises(SimulatedCrash):
self.move(checkpoint=self._crash_at(step))
restarted = self._restart()
self.assertEqual(restarted.source, ll.SOURCE_NODE_FOLDER)
self.assertEqual(restarted.db_path, self.old_db)
self.assertEqual(_rows(self.old_db), self.before["rows"])
result = self.move()
self.assertEqual(_rows(self.new_db), self.before["rows"])
self.assertIsNone(result.retire_problem)
self.assertFalse(
[name for name in os.listdir(self.target) if name.startswith(".pm-move-")]
)
def test_crash_after_the_marker_uses_the_new_library(self):
for step in ("marked", "switched"):
with self.subTest(step=step):
with self.assertRaises(SimulatedCrash):
self.move(checkpoint=self._crash_at(step))
restarted = self._restart()
self.assertEqual(restarted.source, ll.SOURCE_DEFAULT)
self.assertEqual(restarted.db_path, self.new_db)
self.assertEqual(_rows(self.new_db), self.before["rows"])
# The old file was not renamed yet and is never picked.
self.assertTrue(os.path.isfile(self.old_db))
# Reset for the next window: back to the node folder only.
for name in os.listdir(self.target):
os.remove(os.path.join(self.target, name))
def test_crash_after_the_rename(self):
with self.assertRaises(SimulatedCrash):
self.move(checkpoint=self._crash_at("retired"))
restarted = self._restart()
self.assertEqual(restarted.db_path, self.new_db)
self.assertFalse(os.path.exists(self.old_db))
class TestEnvironmentTarget(_MoveCase):
def test_move_into_the_environment_folder(self):
data = os.path.join(self.base, "pm-data")
environ = {ll.DATA_DIR_ENV: data}
location = ll.resolve_library(
None, environ=environ, extension_root=self.ext, folder_paths=self.fp
)
self.assertEqual(location.problem, "data-dir-empty")
result = lm.move_library(lm.plan_move(location), config_lock=self.config_lock)
self.assertEqual(result.location.source, ll.SOURCE_ENVIRONMENT)
self.assertEqual(_rows(os.path.join(data, "prompts.db")), self.before["rows"])
ll.clear_override()
self.assertEqual(
ll.resolve_library(
None, environ=environ, extension_root=self.ext, folder_paths=self.fp
).source,
ll.SOURCE_ENVIRONMENT,
)
self.assertEqual(open_handles_under(self.old_db), [])
if __name__ == "__main__":
unittest.main()