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, )