From edf7811454d2e6d348e753c0dea0b141a30043a3 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Sat, 24 Oct 2026 07:31:49 -0700 Subject: [PATCH] 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- 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". --- database/execution_history.py | 11 + database/library_move.py | 651 +++++++++++++++++++++++++++++ database/models.py | 28 ++ database/review_changes.py | 11 + py/api/__init__.py | 4 + py/api/library_location.py | 305 ++++++++++++++ py/api/system.py | 5 +- tests/test_api_library_location.py | 279 +++++++++++++ tests/test_library_move.py | 497 ++++++++++++++++++++++ tests/test_system_api.py | 1 + utils/library_location.py | 40 +- 11 files changed, 1820 insertions(+), 12 deletions(-) create mode 100644 database/library_move.py create mode 100644 py/api/library_location.py create mode 100644 tests/test_api_library_location.py create mode 100644 tests/test_library_move.py diff --git a/database/execution_history.py b/database/execution_history.py index d9c6cb69..6a72ae6a 100644 --- a/database/execution_history.py +++ b/database/execution_history.py @@ -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) diff --git a/database/library_move.py b/database/library_move.py new file mode 100644 index 00000000..cf4e3a12 --- /dev/null +++ b/database/library_move.py @@ -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-`` (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") diff --git a/database/models.py b/database/models.py index f36fdbc9..20bd94ec 100644 --- a/database/models.py +++ b/database/models.py @@ -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 diff --git a/database/review_changes.py b/database/review_changes.py index 31c54b69..d46a3f3c 100644 --- a/database/review_changes.py +++ b/database/review_changes.py @@ -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) diff --git a/py/api/__init__.py b/py/api/__init__.py index 62551a93..a8b56c44 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -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) diff --git a/py/api/library_location.py b/py/api/library_location.py new file mode 100644 index 00000000..fde87076 --- /dev/null +++ b/py/api/library_location.py @@ -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) diff --git a/py/api/system.py b/py/api/system.py index 092042e2..45d67da6 100644 --- a/py/api/system.py +++ b/py/api/system.py @@ -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"}, ) diff --git a/tests/test_api_library_location.py b/tests/test_api_library_location.py new file mode 100644 index 00000000..68eae643 --- /dev/null +++ b/tests/test_api_library_location.py @@ -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) diff --git a/tests/test_library_move.py b/tests/test_library_move.py new file mode 100644 index 00000000..b838af3a --- /dev/null +++ b/tests/test_library_move.py @@ -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() diff --git a/tests/test_system_api.py b/tests/test_system_api.py index 99891bdc..df7ff935 100644 --- a/tests/test_system_api.py +++ b/tests/test_system_api.py @@ -277,6 +277,7 @@ class TestOverviewRoute(SystemAPITestCase): "backup", "needs_review", "sources", + "location", }, ) diff --git a/utils/library_location.py b/utils/library_location.py index f5b1797b..7e37e8ed 100644 --- a/utils/library_location.py +++ b/utils/library_location.py @@ -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, )