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".
This commit is contained in:
Vito Sansevero
2026-10-24 07:31:49 -07:00
parent 2c3d3a58c4
commit edf7811454
11 changed files with 1820 additions and 12 deletions
+11
View File
@@ -196,3 +196,14 @@ class ExecutionHistoryOperations:
.fetchone()
)
return dict(row) if row is not None else None
def rekey_history_restore(old_path: str, new_path: str) -> None:
"""Move the observers of ``old_path`` to ``new_path`` (Move library)."""
old, new = _database_identity(old_path), _database_identity(new_path)
if old == new:
return
with _restore_lock:
observers = _restore_observers.pop(old, None)
if observers:
_restore_observers.setdefault(new, {}).update(observers)
+651
View File
@@ -0,0 +1,651 @@
"""Move library (#249 S6b): copy the node-folder library, verify it, switch.
``plan_move`` checks a ``utils.library_location`` location can move and
sizes the copy; ``move_library`` does it, in this order:
1. **Disk check** on the target's volume (the copy, the backups and
config.json, plus a margin). The source is never renamed across folders,
so another drive is simply a copy.
2. **Copy** with SQLite's backup API from a private connection holding one
read transaction: the per-table counts, the workflow documents' digest
and ``PRAGMA data_version`` describe exactly the copied snapshot. The copy
is a temp file in the target folder, made self-contained
(``journal_mode=DELETE``), fsynced and owner-only.
3. **Verify**: ``verify_database_file`` (integrity check, schema, every
workflow document decompressed against its hash, #315), the same row
count in every table and the same documents.
4. **Copy files**: every pooled safety backup and pre-migration backup,
checked by size and SHA-256 after the copy (their times are kept, so the
pool prunes in the same order).
5. **Publish** under the old database's write lock (``BEGIN IMMEDIATE``) and
the config.json lock: if any other connection wrote since the snapshot
(``data_version``), steps 2–4 run again (``attempts``, then
``library-busy``). Else config.json is copied (owner-only, checked), a
``pending`` marker is written, the temp files take their final names
(``os.replace`` inside the target folder), the marker turns ``moved``
and ``switch`` points the process at the new library while the old file
is still locked, so no write can land in it.
6. **Retire**: the old file's WAL is checkpointed and the old ``prompts.db``
and ``config.json`` are renamed ``*.moved-to-<folder>`` (never deleted;
the old backups stay where they are). A file still open elsewhere is left
in place and reported (``old-in-use``): the marker already makes a restart
use the new library.
Any failure before the marker says ``moved`` removes everything the move
wrote in the target folder and leaves the old library in use, untouched.
A crash at any point leaves either the old library in use (the marker is
absent or ``pending``; the next move replaces its own leftovers) or the new
one (``moved``). ``MoveFailed`` carries a code and a message without paths.
"""
import contextlib
import hashlib
import json
import os
import shutil
import sqlite3
import time
import uuid
from dataclasses import dataclass, field
from typing import Callable, Dict, List, Optional, Tuple
from .models import BUSY_TIMEOUT_SECONDS, PromptModel
try:
from ..utils import library_location as ll
from ..utils.logging_config import get_logger
except ImportError:
from utils import library_location as ll
from utils.logging_config import get_logger
logger = get_logger("prompt_manager.database.library_move")
TEMP_PREFIX = ".pm-move-"
MOVED_INFIX = ".moved-to-"
DISK_MARGIN_BYTES = 16 * 2**20
DEFAULT_ATTEMPTS = 3
_COPY_CHUNK = 4 * 2**20
# Backups that live beside the database and move with it.
BACKUP_INFIXES = (".backup_", ".pre-workflow-storage_")
MESSAGES = {
"not-movable": (
"This library is not in the node folder, so there is nothing to move."
),
"same-folder": "The new folder is the node folder itself.",
"target-invalid": "The new folder cannot hold the library.",
"target-has-library": (
"The new folder already holds a Prompt Manager library. Nothing was moved."
),
"target-conflict": (
"The new folder already holds files with the same names. Nothing was moved."
),
"insufficient-space": (
"There is not enough free space for the copy. Nothing was moved."
),
"copy-failed": "The copy could not be written or did not match. Nothing was moved.",
"verify-failed": "The copy did not pass verification. Nothing was moved.",
"library-busy": (
"Prompt Manager kept writing to the library while it was copied. "
"Nothing was moved; try again when no workflow is running."
),
"switch-failed": (
"The library was copied and verified, but this session could not "
"switch to it. Restart ComfyUI to use it."
),
}
class MoveFailed(Exception):
"""A move that did not happen (or did not finish), with a coded reason."""
def __init__(self, code: str, message: Optional[str] = None):
self.code = code
self.message = message or MESSAGES.get(code, "The library was not moved.")
super().__init__(self.message)
@dataclass(frozen=True)
class MovePlan:
location: "ll.LibraryLocation"
db_path: str
source_dir: str
target_dir: str
config_path: Optional[str]
backups: Tuple[str, ...]
db_bytes: int
files_bytes: int
@property
def required_bytes(self) -> int:
return self.db_bytes + self.files_bytes + DISK_MARGIN_BYTES
@property
def target_location(self) -> "ll.LibraryLocation":
"""Where the library is once moved (an explicit config.json stays)."""
source = self.location.target_source or ll.SOURCE_DEFAULT
config = (
self.location.config_path
if self.location.config_explicit
else os.path.join(self.target_dir, ll.CONFIG_NAME)
)
return ll.LibraryLocation(
source,
os.path.join(self.target_dir, ll.DATABASE_NAME),
config,
self.target_dir,
self.target_dir,
None,
source,
self.location.config_explicit,
)
@dataclass
class MoveResult:
location: "ll.LibraryLocation"
counts: Dict[str, int]
files: List[str]
renamed: List[str] = field(default_factory=list)
retire_problem: Optional[str] = None
def _size(path: str) -> int:
try:
return os.path.getsize(path)
except OSError:
return 0
def _free_bytes(path: str) -> int:
return shutil.disk_usage(path).free
def _existing_ancestor(path: str) -> str:
current = os.path.abspath(path)
while not os.path.isdir(current):
parent = os.path.dirname(current)
if parent == current:
break
current = parent
return current
def backup_names(db_path: str) -> Tuple[str, ...]:
"""The pooled and pre-migration backups beside ``db_path`` (names)."""
folder = os.path.dirname(os.path.abspath(db_path))
base = os.path.basename(db_path)
prefixes = tuple(base + infix for infix in BACKUP_INFIXES)
try:
names = os.listdir(folder)
except OSError:
return ()
return tuple(
sorted(
name
for name in names
if name.startswith(prefixes) and os.path.isfile(os.path.join(folder, name))
)
)
def plan_move(location) -> MovePlan:
"""Check ``location`` can move to its target; raises MoveFailed."""
if not location.movable:
raise MoveFailed("not-movable")
target = os.path.abspath(location.target_dir)
source_dir = os.path.dirname(os.path.abspath(location.db_path))
if ll.same_folder(target, source_dir) or os.path.normcase(
os.path.realpath(target)
) == os.path.normcase(os.path.realpath(source_dir)):
raise MoveFailed("same-folder")
if os.path.exists(target) and not os.path.isdir(target):
raise MoveFailed("target-invalid")
backups = backup_names(location.db_path)
# config.json moves with a node-folder library; one named by
# PROMPT_MANAGER_CONFIG_PATH is the user's and stays where it is.
config = None if location.config_explicit else location.config_path
config = config if config and os.path.isfile(config) else None
if ll.read_marker(target).get("state") != ll.MARKER_PENDING:
if os.path.exists(os.path.join(target, ll.DATABASE_NAME)):
raise MoveFailed("target-has-library")
names = [ll.CONFIG_NAME] if config else []
if any(os.path.exists(os.path.join(target, n)) for n in names + list(backups)):
raise MoveFailed("target-conflict")
db_bytes = _size(location.db_path) + _size(location.db_path + "-wal")
files_bytes = sum(_size(os.path.join(source_dir, n)) for n in backups)
files_bytes += _size(config) if config else 0
return MovePlan(
location,
os.path.abspath(location.db_path),
source_dir,
target,
config,
backups,
db_bytes,
files_bytes,
)
def space_problem(plan: MovePlan) -> Optional[Tuple[int, int]]:
"""``(required, free)`` when the target's volume is short, else None."""
free = _free_bytes(_existing_ancestor(plan.target_dir))
if free < plan.required_bytes:
return plan.required_bytes, free
return None
# ── copy and verify ───────────────────────────────────────────────────
def _tables(conn) -> List[str]:
return [
row[0]
for row in conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' "
"AND name NOT LIKE 'sqlite_%' ORDER BY name"
)
]
def _describe(conn) -> Tuple[Dict[str, int], str]:
"""Row counts per table and a digest of the workflow documents."""
counts = {}
for table in _tables(conn):
# A name read from sqlite_master, quoted as an identifier.
quoted = '"' + table.replace('"', '""') + '"'
sql = f"SELECT COUNT(*) FROM {quoted}" # nosec B608
counts[table] = conn.execute(sql).fetchone()[0]
digest = hashlib.sha256()
if "workflow_documents" in counts:
for row in conn.execute(
"SELECT id, sha256, byte_length FROM workflow_documents ORDER BY id"
):
digest.update(repr((row[0], bytes(row[1]), row[2])).encode())
return counts, digest.hexdigest()
def _fsync_file(path: str) -> None:
with open(path, "rb+") as handle:
os.fsync(handle.fileno())
def _fsync_dir(path: str) -> None:
if os.name != "posix":
return
fd = os.open(path, os.O_RDONLY)
try:
os.fsync(fd)
finally:
os.close(fd)
def _private(path: str) -> None:
if os.name == "posix":
os.chmod(path, 0o600)
def _copy_snapshot(source: sqlite3.Connection, temp: str):
"""Back up ``source``'s current snapshot to ``temp``; return what it holds."""
source.execute("BEGIN")
try:
version = source.execute("PRAGMA data_version").fetchone()[0]
counts, documents = _describe(source)
destination = sqlite3.connect(temp)
try:
source.backup(destination)
finally:
destination.close()
finally:
source.execute("COMMIT")
with contextlib.closing(sqlite3.connect(temp)) as conn:
conn.execute("PRAGMA journal_mode=DELETE").fetchone()
_fsync_file(temp)
_private(temp)
return version, counts, documents
def _verify_copy(temp: str, counts, documents) -> None:
ok, reason = PromptModel.verify_database_file(temp)
if not ok:
logger.error(f"Library move: the copy did not verify: {reason}")
detail = reason if reason.startswith("Workflow document") else None
raise MoveFailed(
"verify-failed",
f"{MESSAGES['verify-failed']} {detail}" if detail else None,
)
uri = "file:" + temp.replace(os.sep, "/") + "?mode=ro"
try:
with contextlib.closing(sqlite3.connect(uri, uri=True)) as conn:
copied = _describe(conn)
except sqlite3.Error as error:
logger.error(f"Library move: the copy could not be read: {error}")
raise MoveFailed("verify-failed") from error
if copied != (counts, documents):
logger.error("Library move: the copy's rows or documents differ")
raise MoveFailed("verify-failed")
def _hash_copy(source: str, destination: str) -> Tuple[int, str]:
"""Copy ``source`` to ``destination`` (fsynced, owner-only, times kept)."""
digest = hashlib.sha256()
size = 0
with open(source, "rb") as reader, open(destination, "wb") as writer:
while True:
chunk = reader.read(_COPY_CHUNK)
if not chunk:
break
digest.update(chunk)
size += len(chunk)
writer.write(chunk)
writer.flush()
os.fsync(writer.fileno())
stat = os.stat(source)
os.utime(destination, ns=(stat.st_atime_ns, stat.st_mtime_ns))
_private(destination)
return size, digest.hexdigest()
def _file_digest(path: str) -> Tuple[int, str]:
digest = hashlib.sha256()
size = 0
with open(path, "rb") as reader:
while True:
chunk = reader.read(_COPY_CHUNK)
if not chunk:
return size, digest.hexdigest()
digest.update(chunk)
size += len(chunk)
def _check_copies(copies) -> None:
for source, temp, expected in copies:
if _file_digest(temp) != expected or _file_digest(source) != expected:
logger.error(
f"Library move: {os.path.basename(source)} did not copy unchanged"
)
raise MoveFailed("copy-failed")
def write_marker(directory: str, state: str, source_db: str) -> str:
"""Atomically write the move marker; return its path."""
path = os.path.join(directory, ll.MARKER_NAME)
temp = os.path.join(directory, f"{TEMP_PREFIX}{uuid.uuid4().hex}.marker")
payload = json.dumps(
{
"version": 1,
"state": state,
"from": source_db,
"updated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
},
indent=2,
).encode("utf-8")
fd = os.open(temp, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
try:
with os.fdopen(fd, "wb") as handle:
handle.write(payload)
handle.flush()
os.fsync(handle.fileno())
_private(temp)
os.replace(temp, path)
except BaseException:
_remove(temp)
raise
_fsync_dir(directory)
return path
def _remove(*paths) -> None:
for path in paths:
try:
if path and os.path.lexists(path):
os.remove(path)
except OSError:
logger.warning(f"Library move: could not remove {os.path.basename(path)}")
def sweep_move_temps(directory: str) -> List[str]:
"""Remove the temp files an interrupted move left in ``directory``."""
removed = []
try:
names = os.listdir(directory)
except OSError:
return removed
for name in names:
if name.startswith(TEMP_PREFIX):
_remove(os.path.join(directory, name))
removed.append(name)
return removed
# ── the move ─────────────────────────────────────────────────────────
def _noop(*_args, **_kwargs):
return None
def switch_process(old_db: str, location) -> None:
"""Point this process at ``location``: resolver, every model, observers."""
from .execution_history import rekey_history_restore
from .review_changes import rekey_review_changes
ll.switch_to(location)
PromptModel.repoint_instances(old_db, location.db_path)
rekey_review_changes(old_db, location.db_path)
rekey_history_restore(old_db, location.db_path)
class _Move:
def __init__(self, plan, *, switch, config_lock, progress, checkpoint):
self.plan = plan
self.switch = switch
self.config_lock = config_lock
self.progress = progress
self.checkpoint = checkpoint
self.created_dir = False
self.temps: List[str] = []
self.published: List[str] = []
self.marker: Optional[str] = None
self.moved = False
def temp(self, suffix: str) -> str:
path = os.path.join(
self.plan.target_dir, f"{TEMP_PREFIX}{uuid.uuid4().hex}{suffix}"
)
self.temps.append(path)
return path
def prepare_target(self) -> None:
target = self.plan.target_dir
self.created_dir = not os.path.isdir(target)
ll.ensure_private_dir(target)
sweep_move_temps(target)
if ll.read_marker(target).get("state") == ll.MARKER_PENDING:
# An earlier move stopped before it switched: its files are ours.
for name in (ll.DATABASE_NAME, ll.CONFIG_NAME) + self.plan.backups:
_remove(os.path.join(target, name))
for suffix in ("-wal", "-shm", "-journal"):
_remove(os.path.join(target, ll.DATABASE_NAME + suffix))
def copy_round(self, source):
"""Steps 2–4 once: the verified temp copies and the snapshot version."""
self.progress("Copying the database")
temp_db = self.temp(".db")
version, counts, documents = _copy_snapshot(source, temp_db)
self.checkpoint("copied")
self.progress("Verifying the copy")
_verify_copy(temp_db, counts, documents)
self.checkpoint("verified")
copies = []
for name in self.plan.backups:
src = os.path.join(self.plan.source_dir, name)
temp = self.temp(".bak")
try:
expected = _hash_copy(src, temp)
except OSError as error:
logger.error(f"Library move: copying {name} failed: {error}")
raise MoveFailed("copy-failed") from error
copies.append((src, temp, expected, name))
if copies:
self.progress(f"Copied {len(copies)} backup file(s)")
self.checkpoint("files")
_check_copies([(s, t, e) for s, t, e, _ in copies])
return version, counts, temp_db, copies
def discard_round(self, temp_db, copies) -> None:
_remove(temp_db, *[temp for _s, temp, _e, _n in copies])
def publish(self, temp_db, copies) -> List[str]:
"""Under the locks: config.json, marker, final names, marker moved."""
target = self.plan.target_dir
files = []
if self.plan.config_path:
temp = self.temp(".config")
try:
expected = _hash_copy(self.plan.config_path, temp)
except OSError as error:
raise MoveFailed("copy-failed") from error
_check_copies([(self.plan.config_path, temp, expected)])
copies = copies + [(self.plan.config_path, temp, expected, ll.CONFIG_NAME)]
self.marker = write_marker(target, ll.MARKER_PENDING, self.plan.db_path)
self.checkpoint("pending")
for _src, temp, _expected, name in copies:
final = os.path.join(target, name)
os.replace(temp, final)
self.published.append(final)
files.append(name)
final_db = os.path.join(target, ll.DATABASE_NAME)
os.replace(temp_db, final_db)
self.published.append(final_db)
files.append(ll.DATABASE_NAME)
_fsync_dir(target)
self.checkpoint("published")
write_marker(target, ll.MARKER_MOVED, self.plan.db_path)
self.moved = True
self.checkpoint("marked")
return files
def cleanup(self) -> None:
if self.moved:
return
_remove(*self.temps, *self.published, self.marker)
if self.created_dir:
try:
os.rmdir(self.plan.target_dir)
except OSError:
pass
def _unique(path: str) -> str:
candidate, number = path, 1
while os.path.lexists(candidate):
candidate = f"{path}_{number}"
number += 1
return candidate
def retire_old_files(plan: MovePlan) -> Tuple[List[str], Optional[str]]:
"""Rename the old database and config.json; ``(renamed, problem code)``.
Every connection to the old file must be closed. The WAL is checkpointed
first; when it cannot be emptied (another process still has the file
open) nothing is renamed (``old-in-use``).
"""
old = plan.db_path
folder = os.path.basename(os.path.normpath(plan.target_dir))
try:
with contextlib.closing(sqlite3.connect(old, timeout=1)) as conn:
busy = conn.execute("PRAGMA wal_checkpoint(TRUNCATE)").fetchone()[0]
except sqlite3.Error as error:
logger.warning(
f"Library move: the old database could not be checkpointed: {error}"
)
return [], "old-in-use"
if busy or _size(old + "-wal") > 0:
return [], "old-in-use"
renamed = []
try:
for path in (old, plan.config_path):
if not path or not os.path.isfile(path):
continue
new_name = _unique(path + MOVED_INFIX + folder)
os.rename(path, new_name)
renamed.append(os.path.basename(new_name))
for suffix in ("-wal", "-shm"):
if os.path.isfile(old + suffix) and _size(old + suffix) == 0:
os.remove(old + suffix)
except OSError as error:
logger.warning(f"Library move: the old files could not be renamed: {error}")
return renamed, "old-locked"
return renamed, None
def move_library(
plan: MovePlan,
*,
switch: Callable = switch_process,
config_lock=None,
progress: Callable[[str], None] = _noop,
checkpoint: Callable[[str], None] = _noop,
attempts: int = DEFAULT_ATTEMPTS,
) -> MoveResult:
"""Move the library as planned (blocking; see the module docstring)."""
short = space_problem(plan)
if short is not None:
raise MoveFailed("insufficient-space")
lock = config_lock if config_lock is not None else contextlib.nullcontext()
move = _Move(
plan,
switch=switch,
config_lock=lock,
progress=progress,
checkpoint=checkpoint,
)
location = plan.target_location
try:
move.prepare_target()
files, counts = _copy_and_switch(move, location, attempts)
except Exception:
move.cleanup()
raise
from .execution_history import notify_history_restore
notify_history_restore(location.db_path)
progress("Renaming the old files")
renamed, problem = retire_old_files(plan)
checkpoint("retired")
return MoveResult(location, counts, files, renamed, problem)
def _copy_and_switch(move: _Move, location, attempts: int):
plan = move.plan
for attempt in range(1, attempts + 1):
source = sqlite3.connect(
plan.db_path, timeout=BUSY_TIMEOUT_SECONDS, isolation_level=None
)
try:
source.execute(f"PRAGMA busy_timeout = {BUSY_TIMEOUT_SECONDS * 1000}")
version, counts, temp_db, copies = move.copy_round(source)
source.execute("BEGIN IMMEDIATE")
try:
if source.execute("PRAGMA data_version").fetchone()[0] != version:
move.progress("The library changed during the copy; copying again")
move.discard_round(temp_db, copies)
continue
move.checkpoint("locked")
with move.config_lock:
files = move.publish(temp_db, copies)
try:
move.switch(plan.db_path, location)
except Exception as error:
logger.exception("Library move: switching failed")
raise MoveFailed("switch-failed") from error
move.checkpoint("switched")
finally:
if source.in_transaction:
source.execute("ROLLBACK")
return files, counts
finally:
source.close()
raise MoveFailed("library-busy")
+28
View File
@@ -630,13 +630,41 @@ class PromptModel:
for model in models:
model.close_all()
@classmethod
def repoint_instances(cls, old_path: str, new_path: str) -> int:
"""Point every model of ``old_path`` at ``new_path`` (Move library).
Each matching model's connections are closed and every thread opens
the new file on its next ``get_connection``. Closing a connection
another thread is using waits for its statement; a write waiting on
the mover's lock then fails instead of landing in the old file.
Returns how many models were repointed.
"""
old_key = os.path.normcase(os.path.realpath(old_path))
with cls._instances_lock:
models = list(cls._instances)
with cls._init_lock:
cls._initialized_paths.add(os.path.normcase(os.path.realpath(new_path)))
moved = 0
for model in models:
if model._schema_key() == old_key:
model._close_connections(new_path)
moved += 1
return moved
def close_all(self) -> None:
"""Close every thread's connection; they reopen lazily on next use.
Needed before deleting the database file (Windows keeps open files
locked) and by tests that ran queries through executor threads.
"""
self._close_connections()
def _close_connections(self, new_path: Optional[str] = None) -> None:
"""Close every connection (optionally pointing the model elsewhere)."""
with self._conn_lock:
if new_path is not None:
self.db_path = new_path
connections = list(self._connections)
self._connections.clear()
self._generation += 1
+11
View File
@@ -62,3 +62,14 @@ def notify_review_change(db_path: str) -> None:
callback()
except Exception:
_logger.warning("Needs review change observer failed", exc_info=True)
def rekey_review_changes(old_path: str, new_path: str) -> None:
"""Move the observers of ``old_path`` to ``new_path`` (Move library)."""
old, new = _identity(old_path), _identity(new_path)
if old == new:
return
with _lock:
observers = _observers.pop(old, None)
if observers:
_observers.setdefault(new, {}).update(observers)
+4
View File
@@ -13,6 +13,7 @@ Split into domain-specific mixins for maintainability:
- ReviewRecordingMixin: Needs review recording for the scan and the monitor
- ReviewEventsMixin: the Needs review change stream an open page re-reads on
- SystemRoutesMixin: System health check, Overview and usage statistics
- LibraryLocationRoutesMixin: where the library lives, Move library, Keep it here
- MaintenanceRoutesMixin: reviewed maintenance plans (preview, apply)
- WorkflowStorageRoutesMixin: store each workflow once, then Reclaim space
- DuplicateFilesMixin: duplicate image files as find and removal jobs
@@ -49,6 +50,7 @@ from .job_routes import JobRoutesMixin
from .image_tag_jobs import ImageTagJobsMixin
from .thumbnail_jobs import ThumbnailJobsMixin
from .system import SystemRoutesMixin
from .library_location import LibraryLocationRoutesMixin
from .maintenance import MaintenanceRoutesMixin
from .workflow_storage import WorkflowStorageRoutesMixin
from .duplicate_files import DuplicateFilesMixin
@@ -292,6 +294,7 @@ class PromptManagerAPI(
ImageTagJobsMixin,
ThumbnailJobsMixin,
SystemRoutesMixin,
LibraryLocationRoutesMixin,
MaintenanceRoutesMixin,
WorkflowStorageRoutesMixin,
DuplicateFilesMixin,
@@ -532,6 +535,7 @@ class PromptManagerAPI(
self._register_settings_routes(routes)
self._register_job_routes(routes)
self._register_system_routes(routes)
self._register_library_location_routes(routes)
self._register_maintenance_routes(routes)
self._register_duplicate_routes(routes)
self._register_logging_routes(routes)
+305
View File
@@ -0,0 +1,305 @@
"""Library location routes (#249 S6b): where the library lives, Move, Keep.
- ``GET /prompt_manager/system/library`` → ``{success, location}``: the
database's path (relative to ComfyUI, as the statistics show it), why it is
there (``source``: ``default``, ``environment``, ``config`` or
``node_folder``), whether it can move and where to (``target``), the
"Keep it here" choice (``keep``), whether to offer the move (``offer``:
movable and not kept), a ``problem`` code and the database's size. The
System Overview carries the same object as ``location``.
- ``POST /prompt_manager/system/library/move`` → 202 ``{job}``: the
``library_move`` job, ``exclusive`` (it runs alone and every other job,
legacy-page operations included, is refused while it runs), not
cancellable. It copies, verifies and switches the live library
(``database/library_move.py``) and ends with one truthful line; a failure
leaves the library in the node folder, unchanged. Refused before any job
starts: 400 ``not-movable``/``same-folder``/``target-invalid``, 409
``target-has-library``/``target-conflict``/``busy``, 507
``insufficient-space``.
- ``POST /prompt_manager/system/library/keep {keep: bool}`` persists
``library.keep_in_node_folder`` through the config store and answers the
new location.
Errors are ``{success:false, code, error}`` and never name a path.
"""
import asyncio
import functools
import os
from aiohttp import web
from .job_routes import busy_response
from .jobs import JobBusy
try:
from ...database import library_move as lm
from ...database.operations import resolve_library_location
from ...utils import library_location as ll
except ImportError:
from database import library_move as lm
from database.operations import resolve_library_location
from utils import library_location as ll
LIBRARY_MOVE_JOB_KIND = "library_move"
LIBRARY_MOVE_TITLE = "Move library"
LIBRARY_MOVE_GROUPS = ("files", "errors")
KEEP_BODY_MAX = 1024
_REFUSAL_STATUS = {
"not-movable": 400,
"same-folder": 400,
"target-invalid": 400,
"target-has-library": 409,
"target-conflict": 409,
"insufficient-space": 507,
}
STILL_THERE = "The library is still in the node folder, unchanged."
NOT_MOVED = "The library was not moved."
OLD_LEFT = {
"old-in-use": (
"The old database is still open elsewhere, so it was left in the node "
"folder; it is no longer used."
),
"old-locked": (
"The old files could not be renamed, so they were left in the node "
"folder; they are no longer used."
),
}
def _failure(message, status, code):
return web.json_response(
{"success": False, "code": code, "error": message}, status=status
)
def _plural(count, one, many=None):
return f"{count:,} {one if count == 1 else (many or one + 's')}"
def _same_file(first, second):
try:
return os.path.samefile(first, second)
except OSError:
return os.path.normcase(os.path.realpath(first)) == os.path.normcase(
os.path.realpath(second)
)
class LibraryLocationRoutesMixin:
"""Mixin for the /prompt_manager/system/library endpoints."""
# Replaced in tests with a resolver over a temp tree.
_library_resolver = staticmethod(resolve_library_location)
def _register_library_location_routes(self, routes):
@routes.get("/prompt_manager/system/library")
async def library_location_route(request):
return await self.get_library_location(request)
@routes.post("/prompt_manager/system/library/move")
async def library_move_route(request):
return await self.start_library_move(request)
@routes.post("/prompt_manager/system/library/keep")
async def library_keep_route(request):
return await self.keep_library(request)
# ── location ──────────────────────────────────────────────────────
def _live_location(self):
"""The resolved location, unless the live database is another file.
A database the API was handed (tests, the browser fixture server) or
one the resolver no longer names is reported as configured and is
never moved.
"""
location = self._library_resolver()
live = self.db.model.db_path
if _same_file(location.db_path, live):
return location
return ll.LibraryLocation(
ll.SOURCE_CONFIG,
live,
location.config_path,
os.path.dirname(live),
config_explicit=location.config_explicit,
)
def _location_payload(self, location):
"""The public view of ``location`` (blocking: stats the database)."""
from ..config import PromptManagerConfig
keep = bool(PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER)
size = 0
try:
size = os.path.getsize(location.db_path)
except OSError:
pass
return {
"path": self._public_path(location.db_path),
"source": location.source,
"movable": location.movable,
"keep": keep,
"offer": location.movable and not keep,
"target": (
self._public_path(location.target_dir) if location.movable else None
),
"problem": location.problem,
"size_bytes": size,
}
def _library_payload(self):
return self._location_payload(self._live_location())
async def get_library_location(self, request):
try:
payload = await self._run_in_executor(self._library_payload)
except Exception:
self.logger.exception("Library location could not be read")
return _failure(
"The library location could not be read.", 500, "location-failed"
)
return web.json_response(
{"success": True, "location": payload},
headers={"Cache-Control": "no-store"},
)
# ── move ──────────────────────────────────────────────────────────
def _plan_library_move(self):
"""The checked plan (blocking); raises MoveFailed."""
plan = lm.plan_move(self._live_location())
if lm.space_problem(plan) is not None:
raise lm.MoveFailed("insufficient-space")
return plan
async def start_library_move(self, request):
blocker = self._jobs.blocker(LIBRARY_MOVE_JOB_KIND, admission="exclusive")
if blocker is not None:
return busy_response(JobBusy(blocker))
try:
await self._run_in_executor(self._plan_library_move)
except lm.MoveFailed as refused:
status = _REFUSAL_STATUS.get(refused.code, 400)
return _failure(refused.message, status, refused.code)
except Exception:
self.logger.exception("Move library could not be planned")
return _failure(NOT_MOVED, 500, "move-failed")
try:
job = self._jobs.start(
LIBRARY_MOVE_JOB_KIND,
LIBRARY_MOVE_TITLE,
admission="exclusive",
cancellable=False,
scope={},
groups=LIBRARY_MOVE_GROUPS,
runner=self._run_library_move,
)
except JobBusy as busy:
return busy_response(busy)
return web.json_response({"success": True, "job": job.summary()}, status=202)
def _switch_library(self, old_db, location):
"""Point the process and this API at the moved library (blocking)."""
lm.switch_process(old_db, location)
# Plans, health reports and history reads described the old file.
self._database_replaced()
self._history_read_fence = None
async def _run_library_move(self, job):
from ..config_store import _COMMIT_LOCK
loop = asyncio.get_running_loop()
def progress(line):
def show():
job.log_line(line)
job.update(status_line=line)
loop.call_soon_threadsafe(show)
job.update(status_line="Checking the new folder")
try:
plan = await self._run_in_executor(self._plan_library_move)
target = self._public_path(plan.target_dir)
job.log_line(f"Moving the library to {target}")
result = await self._run_in_executor(
functools.partial(
lm.move_library,
plan,
switch=self._switch_library,
config_lock=_COMMIT_LOCK,
progress=progress,
)
)
except lm.MoveFailed as failed:
self.logger.error(f"Move library failed: {failed.code}")
tail = "" if failed.code == "switch-failed" else f" {STILL_THERE}"
return self._fail_library_move(job, failed.code, failed.message + tail)
except Exception:
self.logger.exception("Move library failed")
return self._fail_library_move(job, None, f"{NOT_MOVED} {STILL_THERE}")
await asyncio.sleep(0) # the progress lines come first
job.record_counted("files", len(result.files))
line = (
f"Moved the library to {target}: "
f"{_plural(result.counts.get('prompts', 0), 'prompt')} and "
f"{_plural(result.counts.get('generated_images', 0), 'image')} verified"
)
if result.retire_problem:
line += ". " + OLD_LEFT.get(result.retire_problem, OLD_LEFT["old-locked"])
elif result.renamed:
line += (
f". The old files were renamed {', '.join(result.renamed)}; "
"nothing was deleted"
)
line += ". Logs move at the next restart."
job.log_line(line)
job.update(progress=100, status_line=line)
@staticmethod
def _fail_library_move(job, code, message):
from .maintenance import _fail_job
_fail_job(job, code, message)
# ── keep it here ─────────────────────────────────────────────────
async def keep_library(self, request):
from ..config import PromptManagerConfig
from ..config_store import get_config_store
try:
if (request.content_length or 0) > KEEP_BODY_MAX:
raise ValueError("too large")
data = await request.json()
except Exception:
data = None
if not isinstance(data, dict) or not isinstance(data.get("keep"), bool):
return _failure("keep must be true or false.", 400, "invalid-request")
keep = data["keep"]
def mutate(config):
library = config.get("library")
library = dict(library) if isinstance(library, dict) else {}
library["keep_in_node_folder"] = keep
config["library"] = library
return config
def publish(_config):
PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER = keep
try:
await self._run_in_executor(
functools.partial(
get_config_store().commit,
mutate,
expected_revision=None,
publish=publish,
)
)
except Exception as exc:
return self._config_commit_failure(exc)
return await self.get_library_location(request)
+4 -1
View File
@@ -168,9 +168,10 @@ class SystemRoutesMixin:
from ..config import GalleryConfig
try:
(output_dirs, auto_detected), counts = await _gather_settled(
(output_dirs, auto_detected), counts, location = await _gather_settled(
self._run_in_executor(self._output_dirs_with_origin),
self._run_in_executor(self.db.library_counts),
self._run_in_executor(self._library_payload),
)
except Exception:
self.logger.exception("System overview failed")
@@ -194,6 +195,8 @@ class SystemRoutesMixin:
"effective": len(output_dirs),
"auto_detected": auto_detected,
},
# Where the library lives and why (py/api/library_location.py).
"location": location,
},
headers={"Cache-Control": "no-store"},
)
+279
View File
@@ -0,0 +1,279 @@
"""Library location routes (#249 S6b): where it lives, Move library, Keep it here.
The API's database is a node-folder library inside a fake ComfyUI tree; the
resolver is pointed at that tree (``api._library_resolver``) so nothing here
reads the real environment or the repository's node folder.
"""
import asyncio
import json
import os
import sys
import threading
from unittest.mock import MagicMock, patch
_mock_server = MagicMock()
_mock_server.PromptServer.instance.routes = MagicMock()
sys.modules.setdefault("server", _mock_server)
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.models import PromptModel # noqa: E402
from database.operations import PromptDatabase # noqa: E402
from py.config import PromptManagerConfig # noqa: E402
from utils import library_location as ll # noqa: E402
try:
from tests.test_api_admin import AdminAPITestCase # noqa: E402
except ImportError: # discovered with tests/ as the top-level directory
from test_api_admin import AdminAPITestCase # noqa: E402
WAIT = 10
URL = "/prompt_manager/system/library"
class LibraryAPITestCase(AdminAPITestCase):
async def get_application(self):
app = await super().get_application()
self.ext = os.path.join(str(self.comfy_dir), "custom_nodes", "PM")
os.makedirs(self.ext)
self.user = os.path.join(str(self.comfy_dir), "user")
os.makedirs(self.user)
self.target = os.path.join(self.user, "__promptmanager")
sys.modules["folder_paths"].get_user_directory = lambda: self.user
previous = self.api.db
self.api.db = PromptDatabase(os.path.join(self.ext, "prompts.db"))
previous.close_all()
self.api.db.save_prompt("a library prompt")
self.fp = sys.modules["folder_paths"]
self.environ = {}
self.api._library_resolver = lambda: ll.resolve_library(
None, environ=self.environ, extension_root=self.ext, folder_paths=self.fp
)
ll.clear_override()
self.addCleanup(ll.clear_override)
self.addCleanup(PromptModel.close_all_instances)
self._keep = PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER
self.addCleanup(
setattr, PromptManagerConfig, "KEEP_LIBRARY_IN_NODE_FOLDER", self._keep
)
PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER = False
return app
async def json(self, method, url, status=200, **kwargs):
response = await self.client.request(method, url, **kwargs)
text = await response.text()
self.assertEqual(response.status, status, text)
self.assertNotIn(self.tmpdir, text)
return json.loads(text)
async def finished(self, body):
job = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
return job
class TestLocationRoute(LibraryAPITestCase):
async def test_node_folder_library_offers_the_move(self):
body = await self.json("GET", URL)
self.assertEqual(
body["location"],
{
"path": "custom_nodes/PM/prompts.db",
"source": "node_folder",
"movable": True,
"keep": False,
"offer": True,
"target": "user/__promptmanager",
"problem": None,
"size_bytes": os.path.getsize(self.api.db.model.db_path),
},
)
async def test_overview_carries_the_location(self):
body = await self.json("GET", "/prompt_manager/system/overview")
self.assertEqual(body["location"]["source"], "node_folder")
self.assertTrue(body["location"]["offer"])
async def test_default_folder_library_offers_nothing(self):
self.api._library_resolver = lambda: ll.LibraryLocation(
ll.SOURCE_DEFAULT,
self.api.db.model.db_path,
os.path.join(self.target, "config.json"),
self.target,
self.target,
)
location = (await self.json("GET", URL))["location"]
self.assertEqual((location["source"], location["movable"]), ("default", False))
self.assertFalse(location["offer"])
async def test_a_resolved_file_that_is_not_the_live_database_is_never_moved(self):
self.api._library_resolver = lambda: ll.resolve_library(
None,
environ={},
extension_root=os.path.join(self.tmpdir, "elsewhere"),
folder_paths=self.fp,
)
location = (await self.json("GET", URL))["location"]
self.assertEqual(location["source"], "config")
self.assertFalse(location["movable"])
refused = await self.json("POST", URL + "/move", 400)
self.assertEqual(refused["code"], "not-movable")
class TestMoveRoute(LibraryAPITestCase):
async def test_move_runs_as_an_exclusive_job_and_switches_the_live_library(self):
body = await self.json("POST", URL + "/move", 202)
summary = body["job"]
self.assertEqual(summary["kind"], "library_move")
self.assertEqual(summary["title"], "Move library")
self.assertEqual(summary["admission"], "exclusive")
self.assertFalse(summary["cancellable"])
job = await self.finished(body)
self.assertEqual(job.state, "finished", job.status_line)
self.assertEqual(
self.api.db.model.db_path, os.path.join(self.target, "prompts.db")
)
self.assertIn("user/__promptmanager", job.status_line)
self.assertIn("1 prompt", job.status_line)
self.assertNotIn(self.tmpdir, json.dumps(job.report()))
location = (await self.json("GET", URL))["location"]
self.assertEqual(location["source"], "default")
self.assertEqual(location["path"], "user/__promptmanager/prompts.db")
self.assertFalse(location["offer"])
prompts = await self.json("GET", "/prompt_manager/stats")
self.assertEqual(prompts["stats"]["total_prompts"], 1)
self.assertTrue(
os.path.isfile(
os.path.join(self.ext, "prompts.db.moved-to-__promptmanager")
)
)
async def test_refusals_are_coded_without_paths(self):
with patch.object(lm, "_free_bytes", return_value=1024):
short = await self.json("POST", URL + "/move", 507)
self.assertEqual(short["code"], "insufficient-space")
os.makedirs(self.target)
with open(os.path.join(self.target, "prompts.db"), "wb") as handle:
handle.write(b"another library")
taken = await self.json("POST", URL + "/move", 409)
self.assertEqual(taken["code"], "target-has-library")
self.assertEqual(self.api._jobs.list(), [])
async def test_move_is_refused_while_another_job_runs_and_blocks_others(self):
gate = asyncio.Event()
async def scanning(job):
await gate.wait()
self.api._jobs.start(
"scan", "Metadata scan", admission="records", runner=scanning
)
busy = await self.json("POST", URL + "/move", 409)
self.assertEqual(busy["code"], "busy")
gate.set()
await asyncio.sleep(0)
for job in self.api._jobs.list():
if job.task is not None:
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
release = threading.Event()
entered = threading.Event()
real = lm.move_library
def gated(*args, **kwargs):
entered.set()
release.wait(WAIT)
return real(*args, **kwargs)
with patch.object(lm, "move_library", side_effect=gated):
body = await self.json("POST", URL + "/move", 202)
loop = asyncio.get_running_loop()
await loop.run_in_executor(None, entered.wait, WAIT)
refused = await self.json("POST", "/prompt_manager/system/health", 409)
self.assertEqual(refused["running"]["title"], "Move library")
release.set()
job = await self.finished(body)
self.assertEqual(job.state, "finished", job.status_line)
async def test_a_failed_move_ends_truthfully_and_keeps_the_old_library(self):
with patch.object(
lm, "move_library", side_effect=lm.MoveFailed("verify-failed")
):
body = await self.json("POST", URL + "/move", 202)
job = await self.finished(body)
self.assertEqual(job.state, "failed")
self.assertIn("did not pass verification", job.status_line)
self.assertIn("still in the node folder", job.status_line)
self.assertEqual(
self.api.db.model.db_path, os.path.join(self.ext, "prompts.db")
)
with patch.object(lm, "move_library", side_effect=OSError("disk gone")):
body = await self.json("POST", URL + "/move", 202)
job = await self.finished(body)
self.assertEqual(job.state, "failed")
self.assertNotIn("disk gone", job.status_line)
self.assertIn("still in the node folder", job.status_line)
async def test_old_file_left_in_place_is_reported(self):
real = lm.retire_old_files
with patch.object(lm, "retire_old_files", return_value=([], "old-in-use")):
body = await self.json("POST", URL + "/move", 202)
job = await self.finished(body)
self.assertEqual(job.state, "finished")
self.assertIn("left in the node folder", job.status_line)
self.assertTrue(callable(real))
class TestKeepRoute(LibraryAPITestCase):
async def test_keep_it_here_is_persisted_and_stops_the_offer(self):
body = await self.json("POST", URL + "/keep", 200, json={"keep": True})
self.assertTrue(body["location"]["keep"])
self.assertFalse(body["location"]["offer"])
self.assertTrue(body["location"]["movable"])
with open(self.config_path, encoding="utf-8") as handle:
saved = json.load(handle)
self.assertEqual(saved["library"], {"keep_in_node_folder": True})
self.assertTrue(PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER)
undone = await self.json("POST", URL + "/keep", 200, json={"keep": False})
self.assertTrue(undone["location"]["offer"])
async def test_keep_needs_a_boolean(self):
for payload in ({"keep": "yes"}, {}, [True]):
bad = await self.json("POST", URL + "/keep", 400, json=payload)
self.assertEqual(bad["code"], "invalid-request")
self.assertFalse(os.path.exists(self.config_path))
async def test_unwritable_config_reports_not_saved(self):
with patch("py.config_store._write_atomically", side_effect=OSError("ro")):
body = await self.json("POST", URL + "/keep", 500, json={"keep": True})
self.assertEqual(body["code"], "config-write-failed")
self.assertFalse(PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER)
class TestLegacyRoutesDuringAMove(LibraryAPITestCase):
async def test_legacy_restore_is_refused_while_the_move_runs(self):
from aiohttp import FormData
release = threading.Event()
entered = threading.Event()
real = lm.move_library
def gated(*args, **kwargs):
entered.set()
release.wait(WAIT)
return real(*args, **kwargs)
with patch.object(lm, "move_library", side_effect=gated):
body = await self.json("POST", URL + "/move", 202)
loop = asyncio.get_running_loop()
await loop.run_in_executor(None, entered.wait, WAIT)
form = FormData()
form.add_field("database_file", b"SQLite format 3\x00", filename="x.db")
response = await self.client.request(
"POST", "/prompt_manager/restore", data=form
)
self.assertEqual(response.status, 409, await response.text())
release.set()
await self.finished(body)
+497
View File
@@ -0,0 +1,497 @@
"""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()
+1
View File
@@ -277,6 +277,7 @@ class TestOverviewRoute(SystemAPITestCase):
"backup",
"needs_review",
"sources",
"location",
},
)
+29 -11
View File
@@ -66,9 +66,12 @@ _override: Optional["LibraryLocation"] = None
class LibraryLocation:
"""Where the library is and why.
``target_dir`` is where Move library would put it (None outside ComfyUI);
``target_dir`` is where Move library would put it (None outside ComfyUI)
and ``target_source`` why (the environment folder or the default one);
``problem`` is a code for something the user should know (an invalid or
empty ``PROMPT_MANAGER_DATA_DIR``, a move that did not finish).
``config_explicit``: config.json is named by ``PROMPT_MANAGER_CONFIG_PATH``
(it is the user's and never moves).
"""
source: str
@@ -77,6 +80,8 @@ class LibraryLocation:
data_dir: str
target_dir: Optional[str] = None
problem: Optional[str] = None
target_source: Optional[str] = None
config_explicit: bool = False
@property
def movable(self) -> bool:
@@ -287,6 +292,7 @@ class _Decision:
target_dir: Optional[str]
problem: Optional[str]
db_dir: Optional[str] = None # when the database is not in data_dir
target_source: Optional[str] = None
def _decide(environ, extension_root, folder_paths, gallery_roots) -> _Decision:
@@ -309,28 +315,35 @@ def _decide(environ, extension_root, folder_paths, gallery_roots) -> _Decision:
if problem is None:
env_dir = os.path.abspath(raw.strip())
target = env_dir or default
target_source = SOURCE_ENVIRONMENT if env_dir else SOURCE_DEFAULT
def decision(data_dir, source, problem=None, db_dir=None):
return _Decision(
data_dir, source, target, problem, db_dir, target_source if target else None
)
if env_dir is not None:
state = _marker_state(env_dir)
holds = os.path.isfile(os.path.join(env_dir, DATABASE_NAME))
if state != MARKER_PENDING and (holds or not legacy_db):
return _Decision(env_dir, SOURCE_ENVIRONMENT, env_dir, None)
return decision(env_dir, SOURCE_ENVIRONMENT)
problem = "move-unfinished" if state == MARKER_PENDING else "data-dir-empty"
return _Decision(extension_root, SOURCE_NODE_FOLDER, env_dir, problem)
return decision(extension_root, SOURCE_NODE_FOLDER, problem)
if default is None:
return _Decision(extension_root, SOURCE_NODE_FOLDER, None, problem)
return decision(extension_root, SOURCE_NODE_FOLDER, problem)
if _marker_state(default) == MARKER_MOVED and os.path.isfile(
os.path.join(default, DATABASE_NAME)
):
return _Decision(default, SOURCE_DEFAULT, target, problem)
return decision(default, SOURCE_DEFAULT, problem)
if legacy:
if _marker_state(default) == MARKER_PENDING:
problem = problem or "move-unfinished"
if legacy_db:
return _Decision(extension_root, SOURCE_NODE_FOLDER, target, problem)
return decision(extension_root, SOURCE_NODE_FOLDER, problem)
# Only a node-folder config.json: keep reading it, but a new database
# never goes into the node folder.
return _Decision(extension_root, SOURCE_DEFAULT, target, problem, default)
return _Decision(default, SOURCE_DEFAULT, target, problem)
return decision(extension_root, SOURCE_DEFAULT, problem, default)
return decision(default, SOURCE_DEFAULT, problem)
def config_location(
@@ -368,9 +381,12 @@ def resolve_library(
environ = os.environ if environ is None else environ
root = extension_root or EXTENSION_ROOT
decision = _decide(environ, root, folder_paths, gallery_roots)
config_path = environ.get(CONFIG_PATH_ENV, "").strip() or os.path.join(
decision.data_dir, CONFIG_NAME
)
explicit_config = environ.get(CONFIG_PATH_ENV, "").strip()
config_path = explicit_config or os.path.join(decision.data_dir, CONFIG_NAME)
common = {
"target_source": decision.target_source,
"config_explicit": bool(explicit_config),
}
if decision.source != SOURCE_ENVIRONMENT and is_explicit_db_setting(default_path):
db_path = default_path.strip()
if not os.path.isabs(db_path):
@@ -382,6 +398,7 @@ def resolve_library(
decision.data_dir,
decision.target_dir,
decision.problem,
**common,
)
return LibraryLocation(
decision.source,
@@ -390,6 +407,7 @@ def resolve_library(
decision.data_dir,
decision.target_dir,
decision.problem,
**common,
)