Compare commits
11
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
71108a6199 | ||
|
|
d5736423e8 | ||
|
|
55b49fb336 | ||
|
|
573396cac3 | ||
|
|
3f028c054b | ||
|
|
70fba1db6c | ||
|
|
16c42b27bf | ||
|
|
a865585220 | ||
|
|
ea23d5161d | ||
|
|
edf7811454 | ||
|
|
2c3d3a58c4 |
@@ -211,3 +211,6 @@ coverage/
|
||||
|
||||
# Disposable upstream checkout for real-host CI
|
||||
.ci-comfyui/
|
||||
|
||||
# Old library files a Move library renamed (never deleted; #249 S6b)
|
||||
*.moved-to-*
|
||||
|
||||
@@ -751,9 +751,34 @@ ComfyUI_PromptManager/
|
||||
|
||||
## Configuration
|
||||
|
||||
### Where your library lives
|
||||
|
||||
New installs keep the library (`prompts.db`, `config.json` with the CivitAI key, the
|
||||
database backups and the logs) in ComfyUI's system user folder
|
||||
`ComfyUI/user/__promptmanager/`. An update, a reinstall or ComfyUI-Manager replacing
|
||||
the node folder no longer touches it, and ComfyUI's `/userdata` web API does not
|
||||
serve `__` system folders. Thumbnails stay in `output/thumbnails`.
|
||||
|
||||
- **Upgrading from an earlier version:** your library stays where it is, in the node
|
||||
folder, until you move it; nothing is moved on its own. The workspace shows a notice
|
||||
above every page and System › Overview a "Library location" row. **Move library**
|
||||
copies the database with SQLite's backup API, verifies the copy (integrity check, every
|
||||
workflow document, the same rows in every table), copies `config.json` and the backups,
|
||||
switches to the new folder without a restart and only then renames the old files
|
||||
`prompts.db.moved-to-__promptmanager` and `config.json.moved-to-__promptmanager` (never
|
||||
deleted; the old backups stay too). If anything fails, the library stays in the node
|
||||
folder, unchanged. **Keep it here** stops the notice; the Overview row stays. Take a
|
||||
backup first (System › Backup & export or the dashboard's Backup) as with any change.
|
||||
- **Your own folder:** set `PROMPT_MANAGER_DATA_DIR` to an absolute folder outside
|
||||
ComfyUI's output, input and temp folders and your image sources. If it is empty and a
|
||||
node-folder library exists, that library stays in use until you move it there.
|
||||
- **An explicit `database.default_path`** in `config.json` and `PROMPT_MANAGER_CONFIG_PATH`
|
||||
keep working as before (described below); such a library is not offered a move.
|
||||
|
||||
### Database Settings
|
||||
|
||||
You can customize the database path by creating a `config.json` file in the extension root:
|
||||
You can customize the database path by creating a `config.json` file in the library's
|
||||
folder (the node folder for a library that was not moved):
|
||||
|
||||
```json
|
||||
{
|
||||
@@ -970,6 +995,10 @@ MIT License - see LICENSE file for details.
|
||||
|
||||
## Changelog
|
||||
|
||||
### Unreleased (Library location)
|
||||
|
||||
- **📁 Library outside the node folder**: new installs keep `prompts.db`, `config.json`, the database backups and the logs in `ComfyUI/user/__promptmanager/`, where an update or reinstall of the node cannot delete them and ComfyUI's `/userdata` API does not reach. An existing library stays in the node folder until you press **Move library** (a verified copy, then a live switch; the old files are renamed `*.moved-to-__promptmanager`, never deleted) or **Keep it here**. `PROMPT_MANAGER_DATA_DIR` names your own folder. After moving, restart ComfyUI once so the logs follow. Going back to an older version after a move: an older version does not look in `user/__promptmanager`, so stop ComfyUI and copy `user/__promptmanager/prompts.db` (with its `-wal` file, if one is left) and `config.json` into the node folder; the `*.moved-to-__promptmanager` files are the library as it was at the move, without later changes.
|
||||
|
||||
### Unreleased (Workflows stored once)
|
||||
|
||||
- **🗜️ Smaller database**: each distinct workflow and prompt graph is stored once, compressed, instead of again for every image. Images recorded by this version use it from the first start; the maintenance job "Store each workflow once" rewrites the existing records and then reclaims the space. A 20,000-image test library went from 1.4 GB to 60 MB. The job always takes a verified backup first: the first time, `prompts.db.pre-workflow-storage_<date time>` beside the database, kept until you delete it; on later runs, a normal pooled safety backup. No image or video file is opened or changed.
|
||||
|
||||
+134
-6
@@ -2115,10 +2115,11 @@ Sources & scans (`system-sources.mjs`, S4) is the concept's "sources" body on th
|
||||
server's sources. `GET /prompt_manager/settings` reports them (public paths only):
|
||||
`sources.configured` (each `gallery_root_paths` entry with `valid`, `code` and the
|
||||
server's `reason`), `sources.effective` (what is scanned and served, with its
|
||||
`root_index` and `watch`: `watching` only when the image monitor's observer is alive,
|
||||
`root_index`, `watch` (`watching` only when the image monitor's observer is alive,
|
||||
scheduled that folder and monitoring is on, `stopping` when it still watches it but
|
||||
monitoring was switched off, `manual` when monitoring is off or no monitor runs, else
|
||||
`restart_pending`), `sources.auto_detected`, `monitor` and `settings.monitoring_enabled`.
|
||||
`restart_pending`) and `in_comfyui` (its public path is relative to the ComfyUI base)),
|
||||
`sources.auto_detected`, `monitor` and `settings.monitoring_enabled`.
|
||||
`POST /prompt_manager/settings/validate-source {path}` checks one folder and writes
|
||||
nothing. Every caller (validate-source, the settings save, config.json at load) refuses
|
||||
a source through the same checks (`sources.source_problem`): its form first, then
|
||||
@@ -2310,7 +2311,7 @@ the connection's `total_changes` is 0 at close and the data reads back unchanged
|
||||
`-shm` bytes are not compared: every WAL reader rewrites its read marks there, so no
|
||||
reader can keep it byte-identical (`test_health_check_never_writes`).
|
||||
|
||||
Recorded deviations (plan D60–D67, D79–D81 renumbered D92–D99 and D111–D113; D121–D128 assigned):
|
||||
Recorded deviations (plan D60–D67, D79–D81 renumbered D92–D99 and D111–D113; D121–D129 assigned):
|
||||
|
||||
- **D92** (shell): no "Sample configuration" chip. The Save status is truthful: "No
|
||||
unsaved changes.", "Unsaved changes.", "Saved in this browser." / "Saved to the server
|
||||
@@ -2385,9 +2386,12 @@ Recorded deviations (plan D60–D67, D79–D81 renumbered D92–D99 and D111–D
|
||||
saved yet" with Remove only until Save; a configured source the server cannot use
|
||||
reads "<path> · <server reason>" with Remove only. A line under the list says when the
|
||||
list is not the whole story: "No source is configured. After saving, ComfyUI’s output
|
||||
folder is used." (the last one removed), "No source is configured; ComfyUI’s output
|
||||
folder is used." and "No configured source is available; ComfyUI’s output folder is
|
||||
used." (auto-detected rows, listed without Remove). "restart pending" is the third
|
||||
folder is used." (the last one removed) and "No configured source is available;
|
||||
ComfyUI’s output folder is used." (every configured source unusable). Auto-detected
|
||||
rows are listed without Remove and say they are the default on the row itself:
|
||||
"ComfyUI/output · watching · default" ("ComfyUI/" only when `in_comfyui`; their scan
|
||||
still sends the server's public path), with no line under the list when nothing is
|
||||
configured. "restart pending" is the third
|
||||
watch state. The operation "Refresh files" ("Read every source again to find files
|
||||
added or removed outside ComfyUI.", button "Refresh") closes the surface: it calls the
|
||||
legacy gallery's `/gallery/rescan` and toasts "Files refreshed. N files in M output
|
||||
@@ -2458,6 +2462,20 @@ Recorded deviations (plan D60–D67, D79–D81 renumbered D92–D99 and D111–D
|
||||
`Watch for new outputs saves monitoring_enabled; a restart applies it` and `a save
|
||||
refused for one unusable source names it and marks its row`
|
||||
(workspace-system-sources.spec) and the system-state and system-sources node tests.
|
||||
- **D129** (library location, S6b): the concept shows neither where the library lives
|
||||
nor a way to move it. System › Overview's "Library health" gains a "Library location"
|
||||
`.operation` row after Backup status, in that row's own markup: "<path> · <why>",
|
||||
plus "You chose to keep it here." and any `PROMPT_MANAGER_DATA_DIR` problem, with
|
||||
"Move library" for a node-folder library and "Details" (a `dl.meta-grid` dialog)
|
||||
otherwise. While a node-folder library is offered (`location.offer`), the shell shows
|
||||
the concept's `.notice` (the Sources notice's markup: `div > h2 + p`, then
|
||||
`.row-actions` with "Keep it here" and the `ca-primary` "Move library") in
|
||||
`#library-banner` above `main`, in main's own width and padding, with no bottom margin
|
||||
(main's padding spaces it); `workspace.css` carries the `/* D129 */` rules. Move
|
||||
library opens the shared review (From, To, Size, Moves with it, Old files). Pinned by
|
||||
`system shell and overview match the concept` (the mapping inserts the row) and
|
||||
`library notice matches the concept notice` with its mutation check
|
||||
(system-concept.spec), and the D129 tests of workspace-system.spec.
|
||||
- **D124** (inspector): with "Load metadata on selection" off, the Images panel and
|
||||
the viewer's details show "Metadata loads when you ask for it." and a "Load metadata"
|
||||
button instead of reading the file. The concept has no such state, so no concept
|
||||
@@ -2829,6 +2847,116 @@ Restoring a deduped backup on a version before this one is not supported: that
|
||||
version's restore accepts the file and then reads rewritten images without their
|
||||
workflow (README changelog).
|
||||
|
||||
### Library location contract (#249 System S6b)
|
||||
|
||||
`utils/library_location.py` is the one resolver of every data path: the database
|
||||
(`database/operations.py` `_resolve_db_path`, diagnostics), config.json
|
||||
(`PromptManagerConfig.get_config_path`, so the `ConfigStore`), the backups (beside the
|
||||
database), the logs (`utils/logging_config.py`) and the LoRA image cache
|
||||
(`py/lora_utils.py`). Order: a valid `PROMPT_MANAGER_DATA_DIR` (absolute, never inside
|
||||
ComfyUI's output, input or temp folder, a public user folder, an extra gallery root or
|
||||
an image source listed in any config.json the resolver may read; `config_location` and
|
||||
`resolve_library` read those sources the same way, so config.json and the database
|
||||
never split; an invalid one is ignored and reported as a problem code), used when
|
||||
it holds a library or no node-folder library exists; an explicit
|
||||
`database.default_path` other than `prompts.db` (relative to the node folder, as
|
||||
before) and `PROMPT_MANAGER_CONFIG_PATH`, honoured as before; ComfyUI's
|
||||
`user/__promptmanager` (`folder_paths.get_system_user_directory("promptmanager")`, or
|
||||
`get_user_directory()/__promptmanager`) once its marker `library-location.json` says
|
||||
`moved`; a node-folder `prompts.db` that was not moved; else `user/__promptmanager`
|
||||
(new installs; outside ComfyUI the node folder). An empty environment folder never
|
||||
replaces a node-folder library (`data-dir-empty`). Nothing is moved or copied at
|
||||
startup (`tests/test_library_location.py`, `tests/test_library_location_wiring.py`).
|
||||
|
||||
`POST /prompt_manager/system/library/move` (`py/api/library_location.py`,
|
||||
`database/library_move.py`) starts "Move library": kind `library_move`, admission
|
||||
`exclusive`, not cancellable. Refusals answer before any job: 400 `not-movable`,
|
||||
`same-folder`, `target-invalid`; 409 `target-has-library`, `target-conflict`, `busy`;
|
||||
507 `insufficient-space` (the free space of the target's volume, for the database, its
|
||||
backups and config.json plus 16 MiB). The job copies with SQLite's backup API from one
|
||||
read transaction (counts, document digest and `data_version` of that snapshot) into a
|
||||
`.pm-move-*` temp file in the target folder (`journal_mode=DELETE`, fsync, 0600),
|
||||
verifies it (`verify_database_file`, every workflow document, the same rows in every
|
||||
table, the same documents), copies every `prompts.db.backup_*` and
|
||||
`prompts.db.pre-workflow-storage_*` (times kept, size and SHA-256 checked), then, under
|
||||
the old file's `BEGIN IMMEDIATE` and the config lock, rechecks `data_version` (a write
|
||||
meanwhile copies again, up to three times, then `library-busy`), copies config.json
|
||||
(0600, checked), writes the marker `pending`, gives every file its final name with
|
||||
`os.replace` inside the target folder (a different drive is a copy, never a rename),
|
||||
writes the marker `moved` and switches the process (`switch_process`: the resolver's
|
||||
override, `PromptModel.repoint_instances` closes every connection of every model of the
|
||||
old file and points it at the new one, the review and history observers are re-keyed,
|
||||
the API's plans and health report are dropped). Only then are the old `prompts.db` and
|
||||
`config.json` renamed `*.moved-to-<folder>` after a `wal_checkpoint(TRUNCATE)`; a file
|
||||
still open elsewhere is left in place and named (`old-in-use`, `old-locked`); the old
|
||||
backups stay. A failure before the marker says `moved` removes everything it wrote and
|
||||
leaves the old library in use, untouched (`TestFailures`); a crash leaves the old
|
||||
library in use (no or a `pending` marker; the next move replaces its own leftovers) or
|
||||
the new one (`moved`) (`TestCrashWindows`). `GET /prompt_manager/system/library` and the
|
||||
Overview's `location` give `{path, source, movable, keep, offer, target, problem,
|
||||
size_bytes}` with paths relative to ComfyUI; a path in the node folder reads through
|
||||
the node folder's own unresolved place in custom_nodes (`custom_nodes/promptmanager/
|
||||
prompts.db` even when that folder is a symlink to a checkout outside ComfyUI; never
|
||||
absolute, never `..`, else the shared public path, which sources and jobs keep using
|
||||
unchanged); `POST .../keep {keep}` persists
|
||||
`library.keep_in_node_folder` through the `ConfigStore`. Errors are coded and name no
|
||||
path (`tests/test_api_library_location.py`). The real-host `library-location.spec.mjs`
|
||||
(execution suite, its own fresh host on an OS-chosen port, no configuration) proves a new
|
||||
install creates `user/__promptmanager/prompts.db` and its logs, adds nothing but byte
|
||||
code to the node folder, and that `/userdata` serves none of it.
|
||||
|
||||
The switch runs while the copy is write-locked too (a guard transaction on it), so a
|
||||
switch that fails at any point (the resolver, part of the models, the observers) is
|
||||
undone whole: every model points at the old file again, the copy and its marker are
|
||||
removed, the old file's connections the partial switch detached are closed (the models
|
||||
reopen it on next use), and the move fails "… Nothing was moved."
|
||||
(`TestFailedSwitchIsUndone`). A failed switch can stall up to the 30 s busy timeout when
|
||||
a writer is blocked on the copy's guard lock; this is not a deadlock, and the writer
|
||||
then retries on the old file. What
|
||||
follows a complete switch (dropping plans, the health report, the history read fence,
|
||||
which is closed first) is best effort: logged, never a failed move. A thread that
|
||||
opens a connection while the model is being repointed keeps it only if the model's
|
||||
generation did not change meanwhile (`PromptModel.get_connection`).
|
||||
|
||||
Writes racing the switch: a background write that was waiting on the old file's lock
|
||||
when the switch closed its connection fails, and run counting (`count_runs`,
|
||||
`count_pair_runs`) and the image monitor's link (`ImageGenerationHandler.
|
||||
link_image_to_prompt`: the resolved-pair ingest and the fallback or recent-prompt link)
|
||||
run it once more through `PromptDatabase.survive_library_switch` when the model's
|
||||
generation changed during the attempt; the retry opens the new library. The monitor
|
||||
links through `PromptDatabase.link_image_to_prompt_or_raise`, which raises database
|
||||
errors (`link_image_to_prompt` keeps its log-and-return-0 contract for other callers);
|
||||
each attempt is one transaction, so a failed attempt rolls back and a committed one
|
||||
does not retry, and the image is linked once. `TestBackgroundWritesSurviveTheSwitch`
|
||||
races a real UPDATE (run counting) and a real fallback image link against the lock and
|
||||
asserts each lands in the new file only. Any other failure is not retried. Known limit: that
|
||||
blocked writer, and the mover with it, waits out the writer's busy timeout (30 s)
|
||||
before its connection closes: SQLite holds the connection's mutex during the busy
|
||||
wait, so neither a close nor `Connection.interrupt()` ends it sooner (measured), and
|
||||
another thread cannot shorten the timeout of a connection in use. Other writes are
|
||||
unaffected: threads that ask for a connection after the switch get the new file.
|
||||
Writes from the HTTP routes (a prompt saved in the dashboard at that moment) are not
|
||||
retried and report their error.
|
||||
|
||||
A restart after a crash between the switch and the rename finishes the rename: the
|
||||
marker names the old `prompts.db` and `config.json` (`from`, `config_from`), and the API
|
||||
renames them `*.moved-to-<folder>` at startup when they are still there under their
|
||||
own names, never deleting them and leaving a file still in use alone
|
||||
(`TestRestartFinishesTheRename`). A second ComfyUI process using the same node folder
|
||||
(two installs sharing one custom_nodes folder) is not switched: it keeps writing to the
|
||||
old file, renamed or not, and those writes are not in the moved library; a rename while
|
||||
it has the file open is refused as `old-in-use`. The Overview flags a library that stays
|
||||
in the node folder through an explicit `database.default_path` or a node-folder
|
||||
config.json (`inside-node-folder`: a reinstall deletes it, CivitAI key included); it is
|
||||
not offered a move.
|
||||
|
||||
Windows: every connection to the old file is closed before its rename (the source
|
||||
snapshot connection, every model's connections, the checkpoint connection); temp files
|
||||
are closed before `os.replace`; renames stay inside one folder; folder identity compares
|
||||
`(st_dev, st_ino)` and `normcase(realpath)`, so case variants and links of a served
|
||||
folder are refused. The browser test server closes its own log handlers before its temp
|
||||
root (which now holds the logs) is removed.
|
||||
|
||||
## Workspace data fields for the concept pages (#249)
|
||||
|
||||
These fields are additive. The legacy `/prompt_manager/images/output` listing (no
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -0,0 +1,768 @@
|
||||
"""Move library (#249 S6b): copy the node-folder library, verify it, switch.
|
||||
|
||||
``plan_move`` checks a ``utils.library_location`` location can move and
|
||||
sizes the copy; ``move_library`` does it, in this order:
|
||||
|
||||
1. **Disk check** on the target's volume (the copy, the backups and
|
||||
config.json, plus a margin). The source is never renamed across folders,
|
||||
so another drive is simply a copy.
|
||||
2. **Copy** with SQLite's backup API from a private connection holding one
|
||||
read transaction: the per-table counts, the workflow documents' digest
|
||||
and ``PRAGMA data_version`` describe exactly the copied snapshot. The copy
|
||||
is a temp file in the target folder, made self-contained
|
||||
(``journal_mode=DELETE``), fsynced and owner-only.
|
||||
3. **Verify**: ``verify_database_file`` (integrity check, schema, every
|
||||
workflow document decompressed against its hash, #315), the same row
|
||||
count in every table and the same documents.
|
||||
4. **Copy files**: every pooled safety backup and pre-migration backup,
|
||||
checked by size and SHA-256 after the copy (their times are kept, so the
|
||||
pool prunes in the same order).
|
||||
5. **Publish** under the old database's write lock (``BEGIN IMMEDIATE``) and
|
||||
the config.json lock: if any other connection wrote since the snapshot
|
||||
(``data_version``), steps 2–4 run again (``attempts``, then
|
||||
``library-busy``). Else config.json is copied (owner-only, checked), a
|
||||
``pending`` marker is written, the temp files take their final names
|
||||
(``os.replace`` inside the target folder), the marker turns ``moved``
|
||||
and ``switch`` points the process at the new library while the old file
|
||||
is still locked, so no write can land in it.
|
||||
6. **Retire**: the old file's WAL is checkpointed and the old ``prompts.db``
|
||||
and ``config.json`` are renamed ``*.moved-to-<folder>`` (never deleted;
|
||||
the old backups stay where they are). A file still open elsewhere is left
|
||||
in place and reported (``old-in-use``): the marker already makes a restart
|
||||
use the new library.
|
||||
|
||||
Any failure before the marker says ``moved`` removes everything the move
|
||||
wrote in the target folder and leaves the old library in use, untouched.
|
||||
A crash at any point leaves either the old library in use (the marker is
|
||||
absent or ``pending``; the next move replaces its own leftovers) or the new
|
||||
one (``moved``). ``MoveFailed`` carries a code and a message without paths.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sqlite3
|
||||
import time
|
||||
import types
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
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": (
|
||||
"Prompt Manager could not switch to the verified copy, so the copy was "
|
||||
"removed. Nothing was moved."
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
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 = Path(os.path.abspath(temp)).as_uri() + "?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, source_config: Optional[str] = None
|
||||
) -> str:
|
||||
"""Atomically write the move marker; return its path.
|
||||
|
||||
``from`` and ``config_from`` name the old files, so a restart after a
|
||||
crash between the switch and the rename can finish the rename.
|
||||
"""
|
||||
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,
|
||||
"config_from": source_config,
|
||||
"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, detached: list) -> None:
|
||||
"""Point this process at ``location``: resolver, every model, observers.
|
||||
|
||||
Adds the old file's connections, detached and still open, to
|
||||
``detached`` as soon as they are detached, so the mover closes them
|
||||
(``PromptModel.close_detached``) whether or not the rest succeeds.
|
||||
"""
|
||||
from .execution_history import rekey_history_restore
|
||||
from .review_changes import rekey_review_changes
|
||||
|
||||
ll.switch_to(location)
|
||||
detached.extend(PromptModel.repoint_instances(old_db, location.db_path))
|
||||
rekey_review_changes(old_db, location.db_path)
|
||||
rekey_history_restore(old_db, location.db_path)
|
||||
|
||||
|
||||
def unswitch_process(old_db: str, new_db: str) -> list:
|
||||
"""Undo ``switch_process`` (any part of it); returns detached connections."""
|
||||
from .execution_history import rekey_history_restore
|
||||
from .review_changes import rekey_review_changes
|
||||
|
||||
ll.clear_override()
|
||||
detached = PromptModel.repoint_instances(new_db, old_db)
|
||||
rekey_review_changes(new_db, old_db)
|
||||
rekey_history_restore(new_db, old_db)
|
||||
return detached
|
||||
|
||||
|
||||
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.plan.config_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.plan.config_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 finish_interrupted_move(location) -> List[str]:
|
||||
"""Rename old files a crash left between the switch and the rename.
|
||||
|
||||
The marker of the folder in use says ``moved`` and names the old files;
|
||||
an old ``prompts.db`` (and config.json) still there under its own name is
|
||||
renamed ``*.moved-to-<folder>`` now, never deleted. Returns the new
|
||||
names (nothing when there is nothing to finish or the file is in use).
|
||||
"""
|
||||
if location.source not in (ll.SOURCE_DEFAULT, ll.SOURCE_ENVIRONMENT):
|
||||
return []
|
||||
marker = ll.read_marker(location.data_dir)
|
||||
old = marker.get("from")
|
||||
if marker.get("state") != ll.MARKER_MOVED or not isinstance(old, str):
|
||||
return []
|
||||
if not os.path.isfile(old) or os.path.basename(old) != ll.DATABASE_NAME:
|
||||
return []
|
||||
if ll.same_folder(os.path.dirname(old), location.data_dir):
|
||||
return []
|
||||
config = marker.get("config_from")
|
||||
if not (
|
||||
isinstance(config, str)
|
||||
and os.path.basename(config) == ll.CONFIG_NAME
|
||||
and os.path.isfile(config)
|
||||
and os.path.normcase(os.path.abspath(config))
|
||||
!= os.path.normcase(os.path.abspath(location.config_path))
|
||||
):
|
||||
config = None
|
||||
pending = types.SimpleNamespace(
|
||||
db_path=old, config_path=config, target_dir=location.data_dir
|
||||
)
|
||||
renamed, problem = retire_old_files(pending)
|
||||
if problem:
|
||||
logger.warning(f"Library move: the old files are still in use ({problem})")
|
||||
elif renamed:
|
||||
logger.info(f"Library move finished: renamed {', '.join(renamed)}")
|
||||
return renamed
|
||||
|
||||
|
||||
def move_library(
|
||||
plan: MovePlan,
|
||||
*,
|
||||
switch: Callable = switch_process,
|
||||
unswitch: Callable = unswitch_process,
|
||||
after_switch: Callable = _noop,
|
||||
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).
|
||||
|
||||
``switch(old_db, location, detached)`` points the process at the copy
|
||||
and adds the old connections it detached to ``detached``; if it raises,
|
||||
``unswitch(old_db, new_db)`` undoes it while both files are still
|
||||
locked, every detached connection is closed, the copy is removed and
|
||||
the move fails (``switch-failed``). ``after_switch()`` runs
|
||||
once the switch is final; it is best effort (logged, never fails the move).
|
||||
"""
|
||||
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,
|
||||
)
|
||||
move.unswitch = unswitch
|
||||
location = plan.target_location
|
||||
try:
|
||||
move.prepare_target()
|
||||
files, counts = _copy_and_switch(move, location, attempts)
|
||||
except Exception:
|
||||
move.cleanup()
|
||||
raise
|
||||
try:
|
||||
after_switch()
|
||||
except Exception:
|
||||
logger.exception("Library move: a follow-up after the switch failed")
|
||||
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 _switch_guarded(move: _Move, location) -> list:
|
||||
"""Switch while the copy is locked too; undo it all if any part fails.
|
||||
|
||||
The copy's write lock keeps every repointed thread from writing to it
|
||||
until the switch is whole, so undoing a partial switch loses nothing.
|
||||
Returns the old file's detached connections (still open).
|
||||
"""
|
||||
plan = move.plan
|
||||
guard = sqlite3.connect(
|
||||
location.db_path, timeout=BUSY_TIMEOUT_SECONDS, isolation_level=None
|
||||
)
|
||||
try:
|
||||
guard.execute("BEGIN IMMEDIATE")
|
||||
detached: list = []
|
||||
try:
|
||||
move.switch(plan.db_path, location, detached)
|
||||
except Exception as error:
|
||||
logger.exception("Library move: switching failed; undoing it")
|
||||
move.moved = False # the copy and its marker are removed
|
||||
try:
|
||||
PromptModel.close_detached(
|
||||
move.unswitch(plan.db_path, location.db_path), logger
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Library move: undoing the switch failed")
|
||||
# The old file's connections a partial switch detached: the
|
||||
# models reopen it on next use, so close them now, not at GC.
|
||||
PromptModel.close_detached(detached, logger)
|
||||
raise MoveFailed("switch-failed") from error
|
||||
move.checkpoint("switched")
|
||||
return detached
|
||||
finally:
|
||||
if guard.in_transaction:
|
||||
guard.execute("ROLLBACK")
|
||||
guard.close()
|
||||
|
||||
|
||||
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)
|
||||
detached = _switch_guarded(move, location)
|
||||
# Still holding the old file's lock: a write that was waiting
|
||||
# on it fails when its connection closes (and runs again on
|
||||
# the new library) instead of landing in the old file.
|
||||
PromptModel.close_detached(detached, logger)
|
||||
finally:
|
||||
if source.in_transaction:
|
||||
source.execute("ROLLBACK")
|
||||
return files, counts
|
||||
finally:
|
||||
source.close()
|
||||
raise MoveFailed("library-busy")
|
||||
+70
-15
@@ -609,19 +609,37 @@ class PromptModel:
|
||||
if conn is not None and self._local.generation == self._generation:
|
||||
return conn
|
||||
|
||||
# check_same_thread=False only so close_all can clean up other
|
||||
# threads' connections; each connection is still used by one thread.
|
||||
conn = sqlite3.connect(
|
||||
self.db_path, check_same_thread=False, timeout=BUSY_TIMEOUT_SECONDS
|
||||
)
|
||||
conn.row_factory = sqlite3.Row
|
||||
configure_connection(conn)
|
||||
while True:
|
||||
# The path and generation are read together, and the new
|
||||
# connection is kept only if no Move library repointed the model
|
||||
# meanwhile: otherwise it would write to the old file.
|
||||
with self._conn_lock:
|
||||
path, generation = self.db_path, self._generation
|
||||
# check_same_thread=False only so close_all can clean up other
|
||||
# threads' connections; each connection is still used by one thread.
|
||||
conn = sqlite3.connect(
|
||||
path, check_same_thread=False, timeout=BUSY_TIMEOUT_SECONDS
|
||||
)
|
||||
try:
|
||||
conn.row_factory = sqlite3.Row
|
||||
configure_connection(conn)
|
||||
except BaseException:
|
||||
conn.close()
|
||||
raise
|
||||
with self._conn_lock:
|
||||
if generation == self._generation:
|
||||
self._connections.add(conn)
|
||||
break
|
||||
conn.close()
|
||||
self._local.conn = conn
|
||||
self._local.generation = self._generation
|
||||
with self._conn_lock:
|
||||
self._connections.add(conn)
|
||||
self._local.generation = generation
|
||||
return conn
|
||||
|
||||
@property
|
||||
def generation(self) -> int:
|
||||
"""Bumped whenever every connection is closed or the model repointed."""
|
||||
return self._generation
|
||||
|
||||
@classmethod
|
||||
def close_all_instances(cls) -> None:
|
||||
"""Close every connection of every live model in this process."""
|
||||
@@ -630,22 +648,59 @@ class PromptModel:
|
||||
for model in models:
|
||||
model.close_all()
|
||||
|
||||
@classmethod
|
||||
def repoint_instances(cls, old_path: str, new_path: str) -> list:
|
||||
"""Point every model of ``old_path`` at ``new_path`` (Move library).
|
||||
|
||||
Every thread opens the new file on its next ``get_connection``. The
|
||||
old connections are detached, not closed, and returned: the caller
|
||||
closes them once nothing can block on the new file any more. Closing
|
||||
a connection another thread is using waits for its statement; a write
|
||||
waiting on the mover's lock then fails (and, through
|
||||
``PromptDatabase.survive_library_switch``, runs again on the new file).
|
||||
"""
|
||||
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)))
|
||||
detached = []
|
||||
for model in models:
|
||||
if model._schema_key() == old_key:
|
||||
detached.extend(model._detach_connections(new_path))
|
||||
return detached
|
||||
|
||||
@staticmethod
|
||||
def close_detached(connections, logger=None) -> None:
|
||||
"""Close connections ``repoint_instances`` detached."""
|
||||
for conn in connections:
|
||||
try:
|
||||
conn.close()
|
||||
except sqlite3.Error as e:
|
||||
if logger is not None:
|
||||
logger.warning(f"Error closing database connection: {e}")
|
||||
|
||||
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_detached(self._detach_connections(), self.logger)
|
||||
|
||||
def _detach_connections(self, new_path: Optional[str] = None) -> list:
|
||||
"""Forget every connection (optionally pointing the model elsewhere).
|
||||
|
||||
Returns them for the caller to close; threads reopen on next use.
|
||||
"""
|
||||
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
|
||||
self._local.conn = None
|
||||
for conn in connections:
|
||||
try:
|
||||
conn.close()
|
||||
except sqlite3.Error as e:
|
||||
self.logger.warning(f"Error closing database connection: {e}")
|
||||
return connections
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the calling thread's connection, if it has one.
|
||||
|
||||
+158
-107
@@ -46,6 +46,7 @@ try:
|
||||
from ..utils.logging_config import get_logger
|
||||
from ..utils.hashing import generate_prompt_hash
|
||||
from ..utils.json_strict import dumps_strict
|
||||
from ..utils import library_location
|
||||
except ImportError:
|
||||
import sys
|
||||
|
||||
@@ -54,6 +55,7 @@ except ImportError:
|
||||
from utils.logging_config import get_logger
|
||||
from utils.hashing import generate_prompt_hash
|
||||
from utils.json_strict import dumps_strict
|
||||
from utils import library_location
|
||||
|
||||
|
||||
# Subquery to fetch tags from junction table, embedded in SELECT statements
|
||||
@@ -99,40 +101,56 @@ def escape_like(term: str) -> str:
|
||||
)
|
||||
|
||||
|
||||
def _resolve_db_path(db_path: Optional[str] = None) -> str:
|
||||
"""Resolve the database path from config, falling back to defaults.
|
||||
def _configured_default_path():
|
||||
"""config.json's ``database.default_path``, or None.
|
||||
|
||||
When no explicit path is given, reads PromptManagerConfig.DEFAULT_DB_PATH
|
||||
(set via config.json). Relative paths are resolved against the extension
|
||||
root directory so the database location is predictable regardless of the
|
||||
process working directory (important for Docker deployments).
|
||||
The package-qualified config wins over a standalone ``py.config`` (a
|
||||
custom-node package uses its own configuration).
|
||||
"""
|
||||
try:
|
||||
try:
|
||||
from ..py.config import PromptManagerConfig
|
||||
except ImportError:
|
||||
from py.config import PromptManagerConfig
|
||||
except Exception:
|
||||
return None
|
||||
return getattr(PromptManagerConfig, "DEFAULT_DB_PATH", None)
|
||||
|
||||
|
||||
def resolve_library_location():
|
||||
"""Where the library lives now (``utils.library_location``, #249 S6b).
|
||||
|
||||
The image sources that rule out a ``PROMPT_MANAGER_DATA_DIR`` are read by
|
||||
the resolver itself, exactly as ``config_location`` reads them.
|
||||
"""
|
||||
return library_location.resolve_library(_configured_default_path())
|
||||
|
||||
|
||||
def _resolve_db_path(db_path: Optional[str] = None) -> str:
|
||||
"""Resolve the database path through the one library resolver.
|
||||
|
||||
With no explicit path, ``utils.library_location.resolve_library`` decides
|
||||
(``PROMPT_MANAGER_DATA_DIR``, an explicit ``database.default_path``, a
|
||||
node-folder library not moved yet, else ComfyUI's ``user/__promptmanager``).
|
||||
An explicit relative path is resolved against the extension root, as
|
||||
before, so the location does not depend on the working directory.
|
||||
|
||||
Args:
|
||||
db_path: Explicit database path, or None to use config/default.
|
||||
db_path: Explicit database path, or None to use the resolver.
|
||||
|
||||
Returns:
|
||||
Absolute path to the SQLite database file.
|
||||
Absolute path to the SQLite database file; its folder exists (a new
|
||||
data folder is created owner-only).
|
||||
"""
|
||||
extension_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
if db_path is None:
|
||||
try:
|
||||
try:
|
||||
from ..py.config import PromptManagerConfig
|
||||
except ImportError:
|
||||
from py.config import PromptManagerConfig
|
||||
|
||||
db_path = PromptManagerConfig.DEFAULT_DB_PATH
|
||||
except Exception:
|
||||
db_path = "prompts.db"
|
||||
|
||||
if not os.path.isabs(db_path):
|
||||
db_path = resolve_library_location().db_path
|
||||
elif not os.path.isabs(db_path):
|
||||
extension_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
db_path = os.path.normpath(os.path.join(extension_root, db_path))
|
||||
|
||||
# Ensure parent directory exists for custom paths
|
||||
parent = os.path.dirname(db_path)
|
||||
if parent:
|
||||
os.makedirs(parent, exist_ok=True)
|
||||
library_location.ensure_private_dir(parent)
|
||||
|
||||
return db_path
|
||||
|
||||
@@ -365,6 +383,25 @@ class PromptDatabase(
|
||||
self.model = PromptModel(db_path)
|
||||
self.logger.debug("Database operations initialized successfully")
|
||||
|
||||
def survive_library_switch(self, write):
|
||||
"""Run ``write()``; once more if it failed while the library moved.
|
||||
|
||||
A background write (run counting, the image monitor) that was waiting
|
||||
on the old file when Move library switched fails when its connection
|
||||
is closed. When the model's generation changed during the attempt,
|
||||
the write runs again on a fresh connection, which is the new library;
|
||||
any other failure is raised as it was. ``write`` must open its own
|
||||
transaction (every PromptDatabase method does).
|
||||
"""
|
||||
generation = self.model.generation
|
||||
try:
|
||||
return write()
|
||||
except Exception:
|
||||
if self.model.generation == generation:
|
||||
raise
|
||||
self.logger.info("Running a write again on the moved library")
|
||||
return write()
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the calling thread's database connection (see PromptModel.close)."""
|
||||
self.model.close()
|
||||
@@ -1580,103 +1617,117 @@ class PromptDatabase(
|
||||
metadata: Optional metadata about the image
|
||||
|
||||
Returns:
|
||||
int: The ID of the created image record
|
||||
int: The ID of the created image record, or 0 when nothing was
|
||||
linked (invalid or unknown prompt, already linked, or an error,
|
||||
which is logged)
|
||||
"""
|
||||
try:
|
||||
return self.link_image_to_prompt_or_raise(prompt_id, image_path, metadata)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error linking image to prompt {prompt_id}: {e}")
|
||||
return 0
|
||||
|
||||
def link_image_to_prompt_or_raise(
|
||||
self, prompt_id: str, image_path: str, metadata: Optional[Dict[str, Any]] = None
|
||||
) -> int:
|
||||
"""``link_image_to_prompt`` that raises database errors.
|
||||
|
||||
The image monitor links through this so a write that failed because
|
||||
Move library switched reaches ``survive_library_switch`` and runs
|
||||
again on the new library (#249 S6b). Each call is one transaction:
|
||||
a failed attempt rolls back, and a committed one does not raise.
|
||||
Returns 0 for an invalid or unknown prompt or an image already linked.
|
||||
"""
|
||||
# Validate that prompt_id is a valid integer and exists in prompts table
|
||||
try:
|
||||
# Convert prompt_id to integer if it's a string number
|
||||
if isinstance(prompt_id, str) and prompt_id.isdigit():
|
||||
prompt_id_int = int(prompt_id)
|
||||
elif isinstance(prompt_id, int):
|
||||
prompt_id_int = prompt_id
|
||||
# Convert prompt_id to integer if it's a string number
|
||||
if isinstance(prompt_id, str) and prompt_id.isdigit():
|
||||
prompt_id_int = int(prompt_id)
|
||||
elif isinstance(prompt_id, int):
|
||||
prompt_id_int = prompt_id
|
||||
else:
|
||||
# Handle temporary IDs like "temp_123456" - skip linking
|
||||
if isinstance(prompt_id, str) and prompt_id.startswith("temp_"):
|
||||
self.logger.debug(
|
||||
f"Skipping image linking for temporary prompt ID: {prompt_id}"
|
||||
)
|
||||
return 0
|
||||
else:
|
||||
# Handle temporary IDs like "temp_123456" - skip linking
|
||||
if isinstance(prompt_id, str) and prompt_id.startswith("temp_"):
|
||||
self.logger.debug(
|
||||
f"Skipping image linking for temporary prompt ID: {prompt_id}"
|
||||
)
|
||||
return 0
|
||||
else:
|
||||
self.logger.warning(
|
||||
f"Invalid prompt_id format: {prompt_id}, skipping image linking"
|
||||
)
|
||||
return 0
|
||||
|
||||
# Verify the prompt exists in the database
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT id FROM prompts WHERE id = ?", (prompt_id_int,)
|
||||
self.logger.warning(
|
||||
f"Invalid prompt_id format: {prompt_id}, skipping image linking"
|
||||
)
|
||||
if not cursor.fetchone():
|
||||
self.logger.warning(
|
||||
f"Prompt ID {prompt_id_int} not found in database, skipping "
|
||||
"image linking"
|
||||
)
|
||||
return 0
|
||||
return 0
|
||||
|
||||
# Proceed with linking
|
||||
if not isinstance(metadata, dict):
|
||||
metadata = {}
|
||||
filename = os.path.basename(image_path)
|
||||
file_path = normalize_image_path(image_path)
|
||||
file_size, width, height, image_format = _coerce_file_info(metadata)
|
||||
# Each workflow and prompt graph is stored once (new form).
|
||||
documents = stored_values(
|
||||
# Verify the prompt exists in the database
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT id FROM prompts WHERE id = ?", (prompt_id_int,)
|
||||
)
|
||||
if not cursor.fetchone():
|
||||
self.logger.warning(
|
||||
f"Prompt ID {prompt_id_int} not found in database, skipping "
|
||||
"image linking"
|
||||
)
|
||||
return 0
|
||||
|
||||
# Proceed with linking
|
||||
if not isinstance(metadata, dict):
|
||||
metadata = {}
|
||||
filename = os.path.basename(image_path)
|
||||
file_path = normalize_image_path(image_path)
|
||||
file_size, width, height, image_format = _coerce_file_info(metadata)
|
||||
# Each workflow and prompt graph is stored once (new form).
|
||||
documents = stored_values(
|
||||
conn,
|
||||
{
|
||||
"workflow_data": dumps_strict(metadata.get("workflow", {})),
|
||||
"prompt_metadata": dumps_strict(metadata.get("prompt", {})),
|
||||
},
|
||||
)
|
||||
|
||||
# INSERT OR IGNORE skips a file already linked to this prompt
|
||||
# (same prompt_id + normalised full path)
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
INSERT OR IGNORE INTO generated_images
|
||||
(prompt_id, image_path, filename, file_path, file_size, width,
|
||||
height, format, workflow_data, prompt_metadata, parameters,
|
||||
workflow_document_id, prompt_document_id)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
prompt_id_int,
|
||||
image_path,
|
||||
filename,
|
||||
file_path,
|
||||
file_size,
|
||||
width,
|
||||
height,
|
||||
image_format,
|
||||
documents["workflow_data"],
|
||||
documents["prompt_metadata"],
|
||||
dumps_strict(metadata.get("parameters", {})),
|
||||
documents["workflow_document_id"],
|
||||
documents["prompt_document_id"],
|
||||
),
|
||||
)
|
||||
if cursor.rowcount == 0:
|
||||
# Already linked: a document made only for this row goes.
|
||||
release_documents(
|
||||
conn,
|
||||
{
|
||||
"workflow_data": dumps_strict(metadata.get("workflow", {})),
|
||||
"prompt_metadata": dumps_strict(metadata.get("prompt", {})),
|
||||
},
|
||||
)
|
||||
|
||||
# INSERT OR IGNORE skips a file already linked to this prompt
|
||||
# (same prompt_id + normalised full path)
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
INSERT OR IGNORE INTO generated_images
|
||||
(prompt_id, image_path, filename, file_path, file_size, width,
|
||||
height, format, workflow_data, prompt_metadata, parameters,
|
||||
workflow_document_id, prompt_document_id)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
prompt_id_int,
|
||||
image_path,
|
||||
filename,
|
||||
file_path,
|
||||
file_size,
|
||||
width,
|
||||
height,
|
||||
image_format,
|
||||
documents["workflow_data"],
|
||||
documents["prompt_metadata"],
|
||||
dumps_strict(metadata.get("parameters", {})),
|
||||
documents["workflow_document_id"],
|
||||
documents["prompt_document_id"],
|
||||
),
|
||||
)
|
||||
if cursor.rowcount == 0:
|
||||
# Already linked: a document made only for this row goes.
|
||||
release_documents(
|
||||
conn,
|
||||
(
|
||||
documents["workflow_document_id"],
|
||||
documents["prompt_document_id"],
|
||||
),
|
||||
)
|
||||
conn.commit()
|
||||
conn.commit()
|
||||
|
||||
if cursor.rowcount == 0:
|
||||
self.logger.debug(
|
||||
f"Image {file_path} already linked to prompt {prompt_id_int}"
|
||||
)
|
||||
return 0
|
||||
if cursor.rowcount == 0:
|
||||
self.logger.debug(
|
||||
f"Image {file_path} already linked to prompt {prompt_id_int}"
|
||||
)
|
||||
return 0
|
||||
|
||||
return cursor.lastrowid
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error linking image to prompt {prompt_id}: {e}")
|
||||
return 0
|
||||
return cursor.lastrowid
|
||||
|
||||
def get_prompt_images(self, prompt_id: str) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -14,6 +14,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
|
||||
@@ -51,6 +52,7 @@ from .image_tag_jobs import ImageTagJobsMixin
|
||||
from .thumbnail_jobs import ThumbnailJobsMixin
|
||||
from .thumbnail_auto import AutoThumbnailMixin
|
||||
from .system import SystemRoutesMixin
|
||||
from .library_location import LibraryLocationRoutesMixin
|
||||
from .maintenance import MaintenanceRoutesMixin
|
||||
from .workflow_storage import WorkflowStorageRoutesMixin
|
||||
from .duplicate_files import DuplicateFilesMixin
|
||||
@@ -295,6 +297,7 @@ class PromptManagerAPI(
|
||||
ThumbnailJobsMixin,
|
||||
AutoThumbnailMixin,
|
||||
SystemRoutesMixin,
|
||||
LibraryLocationRoutesMixin,
|
||||
MaintenanceRoutesMixin,
|
||||
WorkflowStorageRoutesMixin,
|
||||
DuplicateFilesMixin,
|
||||
@@ -376,6 +379,10 @@ class PromptManagerAPI(
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Could not sweep interrupted backup copies: {e}")
|
||||
|
||||
# A crash between Move library's switch and its rename left the old
|
||||
# node-folder files under their own names (#249 S6b): rename them now.
|
||||
self.finish_interrupted_library_move()
|
||||
|
||||
# Run cleanup on initialization to remove any existing duplicates
|
||||
try:
|
||||
removed = self.db.cleanup_duplicates()
|
||||
@@ -535,6 +542,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)
|
||||
|
||||
@@ -0,0 +1,420 @@
|
||||
"""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)
|
||||
)
|
||||
|
||||
|
||||
CUSTOM_NODES = "custom_nodes"
|
||||
|
||||
|
||||
def _inside(path, folder):
|
||||
"""``path`` relative to ``folder`` by text (normalised spellings), or None."""
|
||||
child, parent = os.path.normcase(path), os.path.normcase(folder)
|
||||
if child == parent:
|
||||
return ""
|
||||
prefix = parent.rstrip(os.sep) + os.sep
|
||||
return path[len(prefix) :] if child.startswith(prefix) else None
|
||||
|
||||
|
||||
def _comfyui_custom_nodes():
|
||||
"""ComfyUI's custom_nodes folders as it lists them (unresolved), or []."""
|
||||
try:
|
||||
from ..config import _import_folder_paths
|
||||
|
||||
getter = getattr(_import_folder_paths(), "get_folder_paths", None)
|
||||
folders = getter(CUSTOM_NODES) if getter else []
|
||||
return [folder for folder in folders or [] if isinstance(folder, str)]
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
def node_folder_public_path(path, extension_root, custom_nodes):
|
||||
"""``custom_nodes/<node folder>/...`` for a path in the node folder, or None.
|
||||
|
||||
The node folder is matched at its own unresolved place in a custom_nodes
|
||||
folder (one of ``custom_nodes``, or the parent of ``extension_root`` when
|
||||
that is named custom_nodes), so a node folder that is a symlink to a
|
||||
checkout outside ComfyUI still reads relative to ComfyUI. ``path`` may be
|
||||
spelled through the link or resolved. Never absolute, never ``..``:
|
||||
anything else is None and the caller uses ``_public_path``.
|
||||
"""
|
||||
if not path or not extension_root:
|
||||
return None
|
||||
try:
|
||||
root = os.path.normpath(os.path.abspath(extension_root))
|
||||
name = os.path.basename(root)
|
||||
folders = [os.path.normpath(os.path.abspath(f)) for f in custom_nodes]
|
||||
parent = os.path.dirname(root)
|
||||
if os.path.normcase(os.path.basename(parent)) == CUSTOM_NODES:
|
||||
folders.append(parent)
|
||||
spellings = [os.path.normpath(os.path.abspath(path))]
|
||||
spellings.append(os.path.realpath(spellings[0]))
|
||||
except (OSError, ValueError, TypeError):
|
||||
return None
|
||||
for folder in folders:
|
||||
node = os.path.join(folder, name)
|
||||
for spelling, node_spelling in zip(spellings, (node, os.path.realpath(node))):
|
||||
rest = _inside(spelling, node_spelling)
|
||||
if rest is None:
|
||||
continue
|
||||
parts = [os.path.basename(folder), name] + (
|
||||
rest.split(os.sep) if rest else []
|
||||
)
|
||||
if all(part and part not in (".", "..") for part in parts):
|
||||
return "/".join(parts)
|
||||
return None
|
||||
|
||||
|
||||
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 _library_public_path(self, path):
|
||||
"""The library's path for display: through the node folder's place in
|
||||
custom_nodes when it is there (a symlinked node folder included), else
|
||||
``_public_path``. Only the library location uses this; sources and
|
||||
jobs keep ``_public_path``, which their requests round-trip."""
|
||||
shown = node_folder_public_path(
|
||||
path, self._extension_root, _comfyui_custom_nodes()
|
||||
)
|
||||
return shown or self._public_path(path)
|
||||
|
||||
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._library_public_path(location.db_path),
|
||||
# Where Move library would put it, and whether config.json goes too.
|
||||
"target_source": location.target_source if location.movable else None,
|
||||
"config_moves": location.movable and not location.config_explicit,
|
||||
"source": location.source,
|
||||
"movable": location.movable,
|
||||
"keep": keep,
|
||||
"offer": location.movable and not keep,
|
||||
"target": (
|
||||
self._library_public_path(location.target_dir)
|
||||
if location.movable
|
||||
else None
|
||||
),
|
||||
"problem": location.problem or self._node_folder_risk(location),
|
||||
"size_bytes": size,
|
||||
}
|
||||
|
||||
# The node folder (tests point it at a temp tree).
|
||||
_extension_root = ll.EXTENSION_ROOT
|
||||
|
||||
def _node_folder_risk(self, location):
|
||||
"""``inside-node-folder`` for a library that stays in the node folder.
|
||||
|
||||
A configured database path, or a config.json the resolver still reads
|
||||
from the node folder, is deleted with the node by a reinstall, and
|
||||
it is never offered a move.
|
||||
"""
|
||||
if location.source == ll.SOURCE_NODE_FOLDER:
|
||||
return None
|
||||
paths = [location.db_path] + (
|
||||
[] if location.config_explicit else [location.config_path]
|
||||
)
|
||||
root = self._extension_root
|
||||
for path in paths:
|
||||
parent = os.path.dirname(os.path.abspath(path))
|
||||
if ll.same_folder(parent, root) or os.path.normcase(
|
||||
os.path.realpath(parent)
|
||||
).startswith(os.path.normcase(os.path.realpath(root)) + os.sep):
|
||||
return "inside-node-folder"
|
||||
return None
|
||||
|
||||
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 _after_library_switch(self):
|
||||
"""Forget what described the old file (best effort, after the switch).
|
||||
|
||||
Plans, the health report and the history read fence named the old
|
||||
file; the fence is closed before it is dropped.
|
||||
"""
|
||||
self._database_replaced()
|
||||
fence, self._history_read_fence = self._history_read_fence, None
|
||||
if fence is not None:
|
||||
fence.close()
|
||||
|
||||
def finish_interrupted_library_move(self):
|
||||
"""Startup: rename old files a crash left after a move (blocking)."""
|
||||
try:
|
||||
location = self._library_resolver()
|
||||
if _same_file(location.db_path, self.db.model.db_path):
|
||||
return lm.finish_interrupted_move(location)
|
||||
except Exception:
|
||||
self.logger.exception("Could not finish an interrupted library move")
|
||||
return []
|
||||
|
||||
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._library_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,
|
||||
after_switch=self._after_library_switch,
|
||||
config_lock=_COMMIT_LOCK,
|
||||
progress=progress,
|
||||
)
|
||||
)
|
||||
except lm.MoveFailed as failed:
|
||||
self.logger.error(f"Move library failed: {failed.code}")
|
||||
return self._fail_library_move(
|
||||
job, failed.code, f"{failed.message} {STILL_THERE}"
|
||||
)
|
||||
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)
|
||||
+25
-1
@@ -13,7 +13,9 @@
|
||||
- ``describe_sources`` (blocking: it resolves paths) reports the configured
|
||||
sources with their validity, the effective sources with their listing
|
||||
index and watch state, and the image monitor's state. Every path in it is
|
||||
a public path (``_public_path``), never an absolute one.
|
||||
a public path (``_public_path``), never an absolute one; ``in_comfyui``
|
||||
says whether an effective source's public path is relative to the ComfyUI
|
||||
base (``in_comfyui_base``), so the page can read it as ``ComfyUI/<path>``.
|
||||
|
||||
A source's watch state is ``watching`` when the monitor's observer is alive,
|
||||
scheduled that directory and monitoring is on; ``stopping`` when it still
|
||||
@@ -209,6 +211,27 @@ def watch_state(path, *, present, alive, scheduled_keys, enabled):
|
||||
return "restart_pending"
|
||||
|
||||
|
||||
def in_comfyui_base(path):
|
||||
"""Whether ``_public_path(path)`` is relative to the ComfyUI base directory.
|
||||
|
||||
The base is the first anchor ``_public_path`` tries, so this agrees with
|
||||
it: the same canonical (realpath + normcase) spelling on both sides.
|
||||
"""
|
||||
from pathlib import Path
|
||||
|
||||
from ..config import _comfyui_base_directory
|
||||
|
||||
base = _comfyui_base_directory()
|
||||
if not path or not base:
|
||||
return False
|
||||
try:
|
||||
canonical = Path(os.path.normcase(os.path.realpath(str(path))))
|
||||
anchor = Path(os.path.normcase(os.path.realpath(base)))
|
||||
except (OSError, ValueError):
|
||||
return False
|
||||
return canonical.is_relative_to(anchor)
|
||||
|
||||
|
||||
def describe_sources(*, configured, output_dirs, auto_detected, monitor, enabled):
|
||||
"""``(sources, monitor)`` for GET /settings (blocking; public paths only).
|
||||
|
||||
@@ -233,6 +256,7 @@ def describe_sources(*, configured, output_dirs, auto_detected, monitor, enabled
|
||||
scheduled_keys=scheduled_keys,
|
||||
enabled=enabled,
|
||||
),
|
||||
"in_comfyui": in_comfyui_base(path),
|
||||
}
|
||||
for index, path in enumerate(output_dirs)
|
||||
]
|
||||
|
||||
+4
-1
@@ -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"},
|
||||
)
|
||||
|
||||
+18
-7
@@ -36,12 +36,14 @@ from typing import Any, Dict, List, Optional, Tuple # noqa: E402
|
||||
# Import logging system
|
||||
try:
|
||||
from ..utils.logging_config import get_logger
|
||||
from ..utils import library_location
|
||||
except ImportError:
|
||||
import sys
|
||||
|
||||
current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.insert(0, current_dir)
|
||||
from utils.logging_config import get_logger
|
||||
from utils import library_location
|
||||
|
||||
# Initialize logger for config operations
|
||||
config_logger = get_logger("prompt_manager.config")
|
||||
@@ -663,6 +665,10 @@ class PromptManagerConfig:
|
||||
# database first (System › Backup & export › Recovery policy, #249 S6).
|
||||
SAFETY_BACKUP = True
|
||||
|
||||
# "Keep it here" (#249 S6b): the user chose to keep a node-folder library
|
||||
# where it is, so the workspace stops offering to move it.
|
||||
KEEP_LIBRARY_IN_NODE_FOLDER = False
|
||||
|
||||
@classmethod
|
||||
def max_worker_threads(cls) -> int:
|
||||
"""Largest accepted worker-thread count: the detected core count."""
|
||||
@@ -718,6 +724,7 @@ class PromptManagerConfig:
|
||||
"worker_threads": cls.WORKER_THREADS,
|
||||
},
|
||||
"maintenance": {"safety_backup": cls.SAFETY_BACKUP},
|
||||
"library": {"keep_in_node_folder": cls.KEEP_LIBRARY_IN_NODE_FOLDER},
|
||||
"gallery": GalleryConfig.get_config(),
|
||||
"integrations": IntegrationConfig.get_config(),
|
||||
"autotag": AutotagConfig.get_config(),
|
||||
@@ -727,14 +734,12 @@ class PromptManagerConfig:
|
||||
def get_config_path(cls) -> str:
|
||||
"""Path of the persisted config.json.
|
||||
|
||||
Honours ``PROMPT_MANAGER_CONFIG_PATH`` when set; otherwise the file
|
||||
lives at the repository root next to ``pyproject.toml``.
|
||||
Honours ``PROMPT_MANAGER_CONFIG_PATH`` when set; otherwise it lives in
|
||||
the library's data folder (``utils.library_location``): the node
|
||||
folder for a library that was not moved, else ComfyUI's
|
||||
``user/__promptmanager``.
|
||||
"""
|
||||
override = os.environ.get(CONFIG_PATH_ENV, "").strip()
|
||||
if override:
|
||||
return override
|
||||
repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
return os.path.join(repo_root, "config.json")
|
||||
return library_location.config_location()[0]
|
||||
|
||||
@classmethod
|
||||
def load_from_file(cls, config_path: Optional[str] = None):
|
||||
@@ -859,6 +864,12 @@ class PromptManagerConfig:
|
||||
):
|
||||
cls.SAFETY_BACKUP = maintenance["safety_backup"]
|
||||
|
||||
library = new_config.get("library", {})
|
||||
if isinstance(library, dict) and isinstance(
|
||||
library.get("keep_in_node_folder"), bool
|
||||
):
|
||||
cls.KEEP_LIBRARY_IN_NODE_FOLDER = library["keep_in_node_folder"]
|
||||
|
||||
# Update gallery config
|
||||
if "gallery" in new_config:
|
||||
GalleryConfig.update_config(new_config["gallery"])
|
||||
|
||||
+14
-3
@@ -567,10 +567,21 @@ def download_civitai_images(
|
||||
|
||||
|
||||
def get_lora_image_cache_dir() -> Path:
|
||||
"""Get the directory used to cache downloaded LoRA example images."""
|
||||
# Store in the extension's own directory
|
||||
"""Get the directory used to cache downloaded LoRA example images.
|
||||
|
||||
``<data folder>/data/lora_images`` (``utils.library_location``): the node
|
||||
folder for a library that was not moved, else ComfyUI's
|
||||
``user/__promptmanager``.
|
||||
"""
|
||||
try:
|
||||
from ..utils import library_location
|
||||
except ImportError:
|
||||
from utils import library_location
|
||||
|
||||
ext_root = Path(__file__).resolve().parent.parent
|
||||
cache = ext_root / "data" / "lora_images"
|
||||
cache = Path(
|
||||
library_location.cache_dir("lora_images", extension_root=str(ext_root))
|
||||
)
|
||||
cache.mkdir(parents=True, exist_ok=True)
|
||||
return cache
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import socket
|
||||
@@ -656,6 +657,18 @@ def exit_when_stdin_closes(loop: asyncio.AbstractEventLoop) -> threading.Thread:
|
||||
return thread
|
||||
|
||||
|
||||
def close_owned_log_files() -> None:
|
||||
"""Close Prompt Manager's own log handlers before the temp root goes.
|
||||
|
||||
The logs live in the library's data folder (#249 S6b), which for this
|
||||
server is <root>/user/__promptmanager; Windows cannot remove a file
|
||||
this process still has open.
|
||||
"""
|
||||
for handler in logging.getLogger("prompt_manager").handlers[:]:
|
||||
if getattr(handler, "_prompt_manager_owned", False) is True:
|
||||
handler.close()
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
@@ -680,6 +693,7 @@ def main(argv: list[str] | None = None) -> None:
|
||||
finally:
|
||||
# create_app works inside the root; Windows cannot remove a process's cwd.
|
||||
os.chdir(original_cwd)
|
||||
close_owned_log_files()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -181,6 +181,17 @@ function mapOverviewConcept(reference, overview) {
|
||||
// D96: the dialog shows runs and ratings by category and library totals, not prompt, tag or date usage.
|
||||
Statistics: row => ({ ...row, description: 'Runs and ratings by category, with library totals.' }),
|
||||
};
|
||||
const surfaces = reference.columns.surfaces.map(surface => ({ ...surface,
|
||||
rows: surface.rows.map(row => (rows[row.title] ?? (value => value))(row)) }));
|
||||
// D129: "Library location" after Backup status, in the backup row's own markup; the fixture's database is configured
|
||||
// explicitly (config.json), so it does not move and its action is Details.
|
||||
const [health] = surfaces;
|
||||
const backup = health.rows.find(row => row.title === 'Backup status');
|
||||
const location = { ...backup, title: 'Library location',
|
||||
description: `${overview.location.path} · The database path set in config.json.`,
|
||||
action: { ...backup.action, text: 'Details' } };
|
||||
surfaces[0] = { ...health, tree: [...health.tree, 'div.operation'],
|
||||
rows: [...health.rows.slice(0, 2), location, ...health.rows.slice(2)] };
|
||||
return {
|
||||
...reference,
|
||||
// D92: no "Sample configuration" chip.
|
||||
@@ -188,8 +199,7 @@ function mapOverviewConcept(reference, overview) {
|
||||
// D122: the concept's literal label is its prototype action name ("system-tab sections").
|
||||
nav: { ...reference.nav, label: 'System sections' },
|
||||
metrics: { ...reference.metrics, items: reference.metrics.items.map(item => ({ ...item, value: relabel[item.label] })) },
|
||||
columns: { ...reference.columns, surfaces: reference.columns.surfaces.map(surface => ({ ...surface,
|
||||
rows: surface.rows.map(row => (rows[row.title] ?? (value => value))(row)) })) },
|
||||
columns: { ...reference.columns, surfaces },
|
||||
save: { ...reference.save,
|
||||
// D92: truthful status copy, and "Save preferences" is aria-disabled (dimmed) while nothing is unsaved.
|
||||
status: { ...reference.save.status, text: 'No unsaved changes.' },
|
||||
@@ -232,13 +242,21 @@ for (const theme of THEMES) for (const [width, height] of SIZES) {
|
||||
|
||||
// The recorded differences, by path: D93 the review row (its real count), the backup row, the
|
||||
// preferences row and the sources row, and metric 4; D96 the statistics row; D92 the chip, the status and the dimmed
|
||||
// Save; D122 the nav label. (The fixture's first three metrics are the concept's own counts.)
|
||||
// Save; D122 the nav label; D129 the inserted "Library location" row, which shifts the thumbnail row to the fourth place
|
||||
// (row 2 takes its title, text and action; row 3 and the surface's fifth child are new). (The fixture's first three
|
||||
// metrics are the concept's own counts.)
|
||||
const D129_ROW_ACTION = ['backgroundColor', 'borderTopColor', 'borderTopLeftRadius', 'borderTopWidth', 'className', 'color',
|
||||
'cursor', 'fontSize', 'opacity', 'paddingLeft', 'paddingTop', 'text'].map(key => `columns.surfaces.0.rows.3.action.${key}`);
|
||||
const UNMAPPED_PATHS = [
|
||||
'columns.surfaces.0.rows.0.title', 'columns.surfaces.0.rows.1.description',
|
||||
'columns.surfaces.0.rows.2.action.text', 'columns.surfaces.0.rows.2.description', 'columns.surfaces.0.rows.2.title',
|
||||
...D129_ROW_ACTION, 'columns.surfaces.0.rows.3.description', 'columns.surfaces.0.rows.3.title',
|
||||
'columns.surfaces.0.rows.3.tree.0', 'columns.surfaces.0.rows.3.tree.1', 'columns.surfaces.0.rows.3.tree.length',
|
||||
'columns.surfaces.0.rows.length', 'columns.surfaces.0.tree.4', 'columns.surfaces.0.tree.length',
|
||||
'columns.surfaces.1.rows.0.description', 'columns.surfaces.1.rows.1.title', 'columns.surfaces.1.rows.2.description',
|
||||
'heading.actions.0', 'heading.actions.length', 'metrics.items.3.value', 'nav.label',
|
||||
'save.button.cursor', 'save.button.opacity', 'save.status.text',
|
||||
];
|
||||
].sort();
|
||||
|
||||
test('system comparison detects mutations', async ({ page, request }) => {
|
||||
await page.setViewportSize({ width: 1600, height: 1100 });
|
||||
@@ -903,3 +921,85 @@ test('thumbnail dialog comparison detects mutations', async ({ page, request })
|
||||
await undoStyle();
|
||||
expect(diffInventory(expected, await page.evaluate(measureThumbnailReview))).toEqual([]);
|
||||
});
|
||||
|
||||
/**
|
||||
* Runs in the page: one `.notice` (the concept's markup: `div > h2 + p`, then `.row-actions` buttons) with its tokens and
|
||||
* geometry, and its left edge against `edgeSelector` (the content it lines up with).
|
||||
*/
|
||||
function measureNoticeOf([selector, edgeSelector]) {
|
||||
const notice = document.querySelector(selector);
|
||||
const css = (node, ...keys) => node ? Object.fromEntries(keys.map(key => [key, getComputedStyle(node)[key]])) : null;
|
||||
const name = node => [node.localName, ...[...node.classList].sort()].join('.');
|
||||
const text = node => node?.textContent.trim() ?? null;
|
||||
const button = node => ({ text: text(node), className: node.className,
|
||||
...css(node, 'backgroundColor', 'color', 'fontSize', 'paddingTop', 'paddingLeft', 'borderTopColor', 'borderTopLeftRadius') });
|
||||
const edge = document.querySelector(edgeSelector);
|
||||
return {
|
||||
tree: [...notice.children].map(name),
|
||||
text: [...notice.querySelectorAll(':scope > div:not(.row-actions) > *')].map(node => [node.localName, text(node)]),
|
||||
actions: [...notice.querySelectorAll(':scope > .row-actions > button')].map(button),
|
||||
frame: css(notice, 'display', 'alignItems', 'justifyContent', 'columnGap', 'borderTopWidth', 'borderTopColor',
|
||||
'borderTopLeftRadius', 'backgroundColor', 'paddingTop', 'paddingLeft', 'marginBottom', 'flexDirection'),
|
||||
h2: css(notice.querySelector('h2'), 'fontSize', 'fontWeight', 'color'),
|
||||
p: css(notice.querySelector('p'), 'fontSize', 'color', 'marginTop'),
|
||||
rowActions: css(notice.querySelector(':scope > .row-actions'), 'display', 'columnGap', 'flexShrink'),
|
||||
alignedWith: Math.round(notice.getBoundingClientRect().left - edge.getBoundingClientRect().left),
|
||||
pageScroll: document.documentElement.scrollWidth > document.documentElement.clientWidth,
|
||||
};
|
||||
}
|
||||
|
||||
/** D129: the concept's Sources notice as the library notice: its own copy and actions, no bottom margin (main's padding). */
|
||||
function mapLibraryNotice(reference) {
|
||||
return {
|
||||
...reference,
|
||||
text: [['h2', 'Your library is inside the node folder'], ['p', 'Your library is stored inside the node folder, where an '
|
||||
+ 'update or reinstall can delete it. Move it to ComfyUI\'s user folder: it is copied and verified first, and nothing '
|
||||
+ 'is deleted.']],
|
||||
actions: reference.actions.map((action, index) => ({ ...action, text: ['Keep it here', 'Move library'][index] })),
|
||||
frame: { ...reference.frame, marginBottom: '0px' },
|
||||
};
|
||||
}
|
||||
|
||||
async function conceptNotice(page, theme) {
|
||||
await openConcept(page, 'system', theme);
|
||||
await page.locator('#main .section-nav').getByRole('button', { name: 'Sources & scans', exact: true }).click();
|
||||
await expect(page.locator('#main .settings-layout > .notice')).toBeVisible();
|
||||
await page.mouse.move(0, 0);
|
||||
return page.evaluate(measureNoticeOf, ['#main .settings-layout > .notice', '#main .heading']);
|
||||
}
|
||||
|
||||
async function productionNotice(page, theme) {
|
||||
await openProduction(page, 'system/overview', theme);
|
||||
await expect(page.locator('#library-banner > .notice')).toBeVisible();
|
||||
await expect(page.locator('#workspace-main .summary-metrics')).toBeVisible();
|
||||
await page.mouse.move(0, 0);
|
||||
return page.evaluate(measureNoticeOf, ['#library-banner > .notice', '#workspace-main .heading']);
|
||||
}
|
||||
|
||||
for (const theme of THEMES) for (const [width, height] of SIZES) {
|
||||
test(`library notice matches the concept notice ${theme} ${width}x${height}`, async ({ page, request }, testInfo) => {
|
||||
await page.setViewportSize({ width, height });
|
||||
await expect(await request.post('/__test__/system/library', { data: { state: 'node_folder' } })).toBeOK();
|
||||
const concept = await conceptNotice(page, theme);
|
||||
const production = await productionNotice(page, theme);
|
||||
await page.screenshot({ path: testInfo.outputPath(`production-library-notice-${theme}-${width}.png`), fullPage: true });
|
||||
expect(diffInventory(mapLibraryNotice(concept), production)).toEqual([]);
|
||||
expect(production.alignedWith).toBe(0);
|
||||
expect(production.pageScroll).toBe(false);
|
||||
});
|
||||
}
|
||||
|
||||
test('library notice comparison detects mutations', async ({ page, request }) => {
|
||||
await page.setViewportSize({ width: 1600, height: 1100 });
|
||||
await expect(await request.post('/__test__/system/library', { data: { state: 'node_folder' } })).toBeOK();
|
||||
const expected = mapLibraryNotice(await conceptNotice(page, 'dark'));
|
||||
await productionNotice(page, 'dark');
|
||||
const paths = async () => diffInventory(expected,
|
||||
await page.evaluate(measureNoticeOf, ['#library-banner > .notice', '#workspace-main .heading'])).map(item => item.path);
|
||||
expect(await paths()).toEqual([]);
|
||||
const undoToken = await mutate(page, { css: '#library-banner { --soft-blue: rgb(255, 0, 0); }' });
|
||||
expect(await paths()).toEqual(['frame.backgroundColor']);
|
||||
await undoToken();
|
||||
await page.evaluate(() => { document.querySelector('#library-banner .notice h2').textContent = 'Library'; });
|
||||
expect(await paths()).toEqual(['text.0.1']);
|
||||
});
|
||||
|
||||
@@ -19,6 +19,13 @@ test.afterEach(async ({ page, browserName }) => {
|
||||
expect(page.workflowErrors).toEqual([]);
|
||||
});
|
||||
|
||||
// #249 S6b: the legacy admin page says where the library lives (the fixture's database is configured explicitly).
|
||||
test('the admin page names the library location', async ({ page }) => {
|
||||
await page.goto('/prompt_manager/admin');
|
||||
await expect(page.locator('#libraryLocation')).toHaveText('Library: browser.db (set in config.json)');
|
||||
await expect(page.locator('#totalPrompts')).not.toHaveText('-');
|
||||
});
|
||||
|
||||
test('actual gallery form saves a prompt that the admin page can edit and reload', async ({ page, request }) => {
|
||||
await page.goto('/prompt_manager/gallery');
|
||||
await page.locator('#addPromptBtn').click();
|
||||
|
||||
@@ -206,10 +206,9 @@ test('Remove marks the draft; removing the last entry states auto-detect', async
|
||||
expect(await rowsOf(page)).toEqual([['p', 'small', 'No source is configured. After saving, ComfyUI’s output folder is used.']]);
|
||||
await saveButton(page).click();
|
||||
await expect(statusOf(page)).toHaveText('Saved to the server configuration.');
|
||||
// Saved: ComfyUI's output folder is detected and listed without Remove.
|
||||
// Saved: ComfyUI's output folder is detected and listed without Remove, its row saying it is the default.
|
||||
await expect.poll(rowsOf.bind(null, page)).toEqual([
|
||||
['ComfyUI output', 'output · manual scan', ['Scan folder']],
|
||||
['p', 'small', 'No source is configured; ComfyUI’s output folder is used.'],
|
||||
['ComfyUI output', 'ComfyUI/output · manual scan · default', ['Scan folder']],
|
||||
]);
|
||||
const { sources } = await serverJson(request, SETTINGS_URL);
|
||||
expect([sources.configured, sources.auto_detected]).toEqual([[], true]);
|
||||
|
||||
@@ -68,6 +68,8 @@ test('overview shows real counts and Not checked', async ({ page, request }) =>
|
||||
expect(await rowsOf(layout, 'Library health')).toEqual([
|
||||
['0 ambiguous image links', 'Review metadata before assigning a positive / negative pair.', 'Review metadata'],
|
||||
['Backup status', 'No retained backup record.', 'View backups'],
|
||||
// D129: the fixture's database is configured explicitly, so it reads as config.json's choice and does not move.
|
||||
['Library location', 'browser.db · The database path set in config.json.', 'Details'],
|
||||
['No missing thumbnails', 'Generate only missing previews when necessary.', 'Review thumbnails'],
|
||||
]);
|
||||
expect(await rowsOf(layout, 'Workspace preferences')).toEqual([
|
||||
@@ -78,6 +80,95 @@ test('overview shows real counts and Not checked', async ({ page, request }) =>
|
||||
expect(writes).toEqual([]);
|
||||
});
|
||||
|
||||
// D129 (#249 S6b): where the library lives. The hook makes the fixture's database read as a node-folder library; its
|
||||
// Move library is a stand-in that touches no file (workspace_system_fixture.py).
|
||||
const banner = page => page.locator('#library-banner');
|
||||
const LIBRARY_URL = '/prompt_manager/system/library';
|
||||
|
||||
test('a node-folder library is offered a move above every page; Keep it here stops the offer', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/library', { data: { state: 'node_folder' } })).toBeOK();
|
||||
await page.goto(`${workspace}#/images`);
|
||||
await expect(banner(page).locator('.notice h2')).toHaveText('Your library is inside the node folder');
|
||||
await expect(banner(page).locator('.notice p')).toHaveText(/^Your library is stored inside the node folder, where an update or reinstall can delete it\./);
|
||||
await expect(banner(page).locator('.notice .row-actions button')).toHaveText(['Keep it here', 'Move library']);
|
||||
await expect(banner(page).getByRole('button', { name: 'Move library', exact: true })).toHaveClass(/ca-primary/);
|
||||
await page.goto(`${workspace}#/library`);
|
||||
await expect(banner(page).locator('.notice')).toBeVisible();
|
||||
await banner(page).getByRole('button', { name: 'Keep it here', exact: true }).click();
|
||||
await expect(banner(page)).toBeHidden();
|
||||
const { location } = await serverJson(request, LIBRARY_URL);
|
||||
expect([location.keep, location.offer, location.movable]).toEqual([true, false, true]);
|
||||
await page.reload();
|
||||
const layout = await openOverview(page);
|
||||
await expect(banner(page)).toBeHidden();
|
||||
const row = layout.locator('.operation').filter({ hasText: 'Library location' });
|
||||
await expect(row.locator('p')).toHaveText('browser.db · The node folder, where an update or reinstall can delete it. '
|
||||
+ 'You chose to keep it here.');
|
||||
// The row stays: the move is still one press away.
|
||||
await expect(row.locator('.row-actions button')).toHaveText('Move library');
|
||||
});
|
||||
|
||||
test('Move library is reviewed, runs as an Activity job and ends the offer', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/library', { data: { state: 'node_folder' } })).toBeOK();
|
||||
const layout = await openOverview(page);
|
||||
const row = layout.locator('.operation').filter({ hasText: 'Library location' });
|
||||
await expect(row.locator('p')).toHaveText('browser.db · The node folder, where an update or reinstall can delete it.');
|
||||
const move = banner(page).getByRole('button', { name: 'Move library', exact: true });
|
||||
await move.click();
|
||||
const dialog = page.locator('dialog#modal[open]');
|
||||
await expect(dialog.locator('#modal-title')).toHaveText('Move library');
|
||||
expect(await dialog.locator('dl.meta-grid > div').evaluateAll(items => items.map(item =>
|
||||
[item.querySelector('dt').textContent, item.querySelector('dd').textContent]))).toEqual([
|
||||
['From', 'browser.db'], ['To', 'user/__promptmanager/prompts.db'], ['Size', expect.stringMatching(/ (KB|MB)$/)],
|
||||
['Moves with it', 'config.json and the database backups'], ['Old files', 'Renamed *.moved-to-…, never deleted'],
|
||||
]);
|
||||
await expect(dialog.locator('.dialog-actions button')).toHaveText(['Cancel', 'Move library']);
|
||||
await dialog.getByRole('button', { name: 'Cancel', exact: true }).click();
|
||||
await expect(page.locator('dialog#modal[open]')).toHaveCount(0);
|
||||
await expect(move).toBeFocused();
|
||||
await row.getByRole('button', { name: 'Move library', exact: true }).click();
|
||||
await dialog.getByRole('button', { name: 'Move library', exact: true }).click();
|
||||
await expect(page).toHaveURL(/#\/activity\/[^/]+$/);
|
||||
await expect(page.locator('#workspace-main')).toContainText(
|
||||
/Moved the library to user\/__promptmanager: \d+ prompts? and 0 images verified\. The old files were renamed browser\.db\.moved-to-__promptmanager; nothing was deleted\./);
|
||||
// The offer ends with the job; the Overview says where the library lives now.
|
||||
await expect(banner(page)).toBeHidden();
|
||||
const after = await openOverview(page);
|
||||
await expect(after.locator('.operation').filter({ hasText: 'Library location' }).locator('p'))
|
||||
.toHaveText('browser.db · ComfyUI\'s user folder, which updates and reinstalls of the node leave alone.');
|
||||
});
|
||||
|
||||
test('a move that fails says so and keeps the offer; a refused start says why', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/library', { data: { state: 'failing' } })).toBeOK();
|
||||
await page.goto(`${workspace}#/images`);
|
||||
await banner(page).getByRole('button', { name: 'Move library', exact: true }).click();
|
||||
await page.locator('dialog#modal[open]').getByRole('button', { name: 'Move library', exact: true }).click();
|
||||
await expect(page).toHaveURL(/#\/activity\/[^/]+$/);
|
||||
await expect(page.locator('#workspace-main')).toContainText(
|
||||
'The copy did not pass verification. Nothing was moved. The library is still in the node folder, unchanged.');
|
||||
await expect(banner(page).locator('.notice')).toBeVisible();
|
||||
await expect(await request.post('/__test__/system/fail', { data: { path: `${LIBRARY_URL}/move` } })).toBeOK();
|
||||
await banner(page).getByRole('button', { name: 'Move library', exact: true }).click();
|
||||
await page.locator('dialog#modal[open]').getByRole('button', { name: 'Move library', exact: true }).click();
|
||||
await expect(page.locator('#workspace-toast')).toHaveText('Move library could not start. Try again.');
|
||||
});
|
||||
|
||||
test('a library that does not move shows its details', async ({ page }) => {
|
||||
const layout = await openOverview(page);
|
||||
await expect(banner(page)).toBeHidden();
|
||||
const details = layout.locator('.operation').filter({ hasText: 'Library location' }).getByRole('button', { name: 'Details', exact: true });
|
||||
await details.click();
|
||||
const dialog = page.locator('dialog#modal[open]');
|
||||
await expect(dialog.locator('#modal-title')).toHaveText('Library location');
|
||||
expect(await dialog.locator('dl.meta-grid > div').evaluateAll(items => items.map(item =>
|
||||
[item.querySelector('dt').textContent, item.querySelector('dd').textContent.replace(/[\d.,]+ (KB|MB)$/, 'N')]))).toEqual([
|
||||
['Database', 'browser.db'], ['Why here', 'The database path set in config.json'],
|
||||
['Beside it', 'config.json, the database backups and the logs'], ['Size', 'N'],
|
||||
]);
|
||||
await dialog.getByRole('button', { name: 'Close', exact: true }).click();
|
||||
await expect(details).toBeFocused();
|
||||
});
|
||||
|
||||
test('Run health checks opens Health report and updates metric 4', async ({ page, request }) => {
|
||||
await page.goto(`${workspace}#/system/maintenance`);
|
||||
const layout = layoutOf(page);
|
||||
|
||||
@@ -23,6 +23,12 @@ read. These per-server hooks only change what a System test needs on top:
|
||||
image monitor whose observer is ``alive`` (default true) and schedules the
|
||||
named folders of this server's root (``["output"]``), so Sources reads
|
||||
watching, manual scan and restart pending states (S4).
|
||||
- ``POST /__test__/system/library`` ``{state}`` (S6b): ``node_folder`` makes
|
||||
the live database read as a node-folder library that can move to
|
||||
``user/__promptmanager``; Move library then runs a stand-in move that
|
||||
touches no file and ends with the library in the default folder.
|
||||
``failing`` is the same library whose move fails verification. Reset
|
||||
restores the real resolver, the real move and "Keep it here".
|
||||
|
||||
Delays use the shared request holds (``request_hold.py``).
|
||||
"""
|
||||
@@ -54,7 +60,12 @@ def register_system_commands(app, routes, api, root: Path):
|
||||
from py.config import GalleryConfig, PromptManagerConfig
|
||||
from utils import image_monitor
|
||||
|
||||
from database import library_move
|
||||
from utils import library_location
|
||||
|
||||
configured = list(GalleryConfig.MONITORING_DIRECTORIES)
|
||||
real_move = library_move.move_library
|
||||
keep = PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER
|
||||
config_bytes = config_path.read_bytes() if config_path.exists() else None
|
||||
host_values = (
|
||||
PromptManagerConfig.RESULT_TIMEOUT,
|
||||
@@ -92,6 +103,9 @@ def register_system_commands(app, routes, api, root: Path):
|
||||
GalleryConfig.MONITORING_ENABLED,
|
||||
) = host_values
|
||||
image_monitor._monitor_instance = monitor
|
||||
api.__dict__.pop("_library_resolver", None)
|
||||
library_move.move_library = real_move
|
||||
PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER = keep
|
||||
if config_bytes is None:
|
||||
config_path.unlink(missing_ok=True)
|
||||
else:
|
||||
@@ -187,6 +201,43 @@ def register_system_commands(app, routes, api, root: Path):
|
||||
)
|
||||
return web.json_response({"success": True})
|
||||
|
||||
@routes.post("/__test__/system/library")
|
||||
async def library_state(request):
|
||||
data = await request.json()
|
||||
state = data.get("state")
|
||||
if state not in ("node_folder", "failing"):
|
||||
raise web.HTTPBadRequest()
|
||||
target = str(root / "user" / "__promptmanager")
|
||||
|
||||
def located(source):
|
||||
return library_location.LibraryLocation(
|
||||
source,
|
||||
api.db.model.db_path,
|
||||
str(config_path),
|
||||
str(root),
|
||||
target if source == library_location.SOURCE_NODE_FOLDER else None,
|
||||
target_source=library_location.SOURCE_DEFAULT,
|
||||
)
|
||||
|
||||
current = {"source": library_location.SOURCE_NODE_FOLDER}
|
||||
api._library_resolver = lambda: located(current["source"])
|
||||
|
||||
def stand_in_move(plan, *, progress, **_kwargs):
|
||||
if state == "failing":
|
||||
raise library_move.MoveFailed("verify-failed")
|
||||
progress("Copying the database")
|
||||
counts = api.db.library_counts()
|
||||
current["source"] = library_location.SOURCE_DEFAULT
|
||||
return library_move.MoveResult(
|
||||
located(library_location.SOURCE_DEFAULT),
|
||||
{"prompts": counts["positive_prompts"], "generated_images": 0},
|
||||
["prompts.db"],
|
||||
["browser.db.moved-to-__promptmanager"],
|
||||
)
|
||||
|
||||
library_move.move_library = stand_in_move
|
||||
return web.json_response({"success": True})
|
||||
|
||||
return reset
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
// #249 S6b: a fresh install keeps its library in ComfyUI's system user folder. A pinned CPU ComfyUI starts with Prompt
|
||||
// Manager and no configuration at all (no config.json, no PROMPT_MANAGER_* paths): the library must be created in
|
||||
// user/__promptmanager, nothing may be written into the node folder, and ComfyUI's /userdata API must not serve it.
|
||||
import { execFileSync } from 'node:child_process';
|
||||
import { access, readdir } from 'node:fs/promises';
|
||||
import { createServer } from 'node:net';
|
||||
import { join } from 'node:path';
|
||||
import { fileURLToPath } from 'node:url';
|
||||
import { test, expect } from '../../tools/vendor/node_modules/@playwright/test/index.mjs';
|
||||
import { launchSharedHost } from '../../tools/vendor/comfyui-host.mjs';
|
||||
import runtime from '../../tools/vendor/comfyui-runtime.json' with { type: 'json' };
|
||||
|
||||
const repositoryRoot = fileURLToPath(new URL('../../', import.meta.url));
|
||||
|
||||
/** An unused loopback port chosen by the OS (this host shares no port with the other suites). */
|
||||
function freePort() {
|
||||
return new Promise((resolve, reject) => {
|
||||
const probe = createServer();
|
||||
probe.once('error', reject);
|
||||
probe.listen({ host: '127.0.0.1', port: 0, exclusive: true }, () => {
|
||||
const { port } = probe.address();
|
||||
probe.close(error => error ? reject(error) : resolve(port));
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
test('a fresh host keeps the library in user/__promptmanager and serves none of it', async ({ playwright }, testInfo) => {
|
||||
test.skip(testInfo.project.name !== 'classic-start', 'one fresh host is enough; the renderer does not matter here');
|
||||
test.setTimeout(180000);
|
||||
const port = await freePort();
|
||||
const host = await launchSharedHost({ checkout: process.env.PM_COMFY_ROOT, repository: repositoryRoot,
|
||||
python: process.env.PM_TEST_PYTHON || 'python', port, pins: runtime, forwardOutput: false,
|
||||
logPath: join(repositoryRoot, 'tools/vendor/test-results/executions/fresh-host.log'), pluginName: 'promptmanager',
|
||||
// A fresh install: nothing written, and no inherited path variable may choose the folder.
|
||||
preparePlugin: async () => ({ PROMPT_MANAGER_CONFIG_PATH: '', PROMPT_MANAGER_DATA_DIR: '',
|
||||
PROMPT_MANAGER_EXTRA_GALLERY_ROOTS: '' }) });
|
||||
const request = await playwright.request.newContext({ baseURL: `http://127.0.0.1:${port}/` });
|
||||
try {
|
||||
await expect.poll(async () => {
|
||||
if (host.child.exitCode !== null) throw new Error('The fresh host exited before it answered');
|
||||
try { return (await request.get('/system_stats')).ok(); } catch { return false; }
|
||||
}, { timeout: 120000 }).toBe(true);
|
||||
expect((await (await request.get('/system_stats')).json()).system.comfyui_version).toBe(runtime.serverVersion);
|
||||
await expect(await request.get('/prompt_manager/stats')).toBeOK();
|
||||
|
||||
const dataFolder = join(host.stateDirectory, 'user', '__promptmanager');
|
||||
await access(join(dataFolder, 'prompts.db'));
|
||||
await access(join(dataFolder, 'logs'));
|
||||
const { location } = await (await request.get('/prompt_manager/system/library')).json();
|
||||
expect(location).toMatchObject({ path: 'user/__promptmanager/prompts.db', source: 'default', movable: false,
|
||||
offer: false, problem: null });
|
||||
|
||||
// Nothing but byte code was added to the node folder: no database, config.json, logs or caches.
|
||||
const installed = join(host.stateDirectory, 'custom_nodes', 'promptmanager');
|
||||
const tracked = new Set(execFileSync('git', ['-C', repositoryRoot, 'ls-files', '-z'], { encoding: 'utf8' })
|
||||
.split('\0').filter(Boolean).map(file => file.split('/')[0]));
|
||||
const added = (await readdir(installed)).filter(name => !tracked.has(name) && name !== '__pycache__');
|
||||
expect(added).toEqual([]);
|
||||
|
||||
// ComfyUI's /userdata serves public users only: the system user's folder is out of reach.
|
||||
for (const path of ['/userdata/..%2F__promptmanager%2Fprompts.db', '/api/userdata/..%2F__promptmanager%2Fprompts.db',
|
||||
'/userdata/%2E%2E%2F__promptmanager%2Fconfig.json', '/userdata/prompts.db']) {
|
||||
const response = await request.get(path);
|
||||
expect(response.status(), path).not.toBe(200);
|
||||
}
|
||||
const listing = await request.get('/userdata?dir=..%2F__promptmanager&recurse=true');
|
||||
expect(listing.status()).not.toBe(200);
|
||||
const asSystemUser = await request.get('/userdata/prompts.db', { headers: { 'comfy-user': '__promptmanager' } });
|
||||
expect(asSystemUser.status()).not.toBe(200);
|
||||
} finally {
|
||||
try { await request.dispose(); } finally { await host.stop(); }
|
||||
}
|
||||
});
|
||||
@@ -0,0 +1,159 @@
|
||||
// Where the library lives (#249 System S6b, D129): the Overview row's copy, the Move library review and its start,
|
||||
// "Keep it here" and the Move library job's end.
|
||||
const test = require('node:test');
|
||||
const assert = require('node:assert/strict');
|
||||
|
||||
const locationModule = import('../../web/js/workspace/library-location.mjs');
|
||||
|
||||
const NODE_FOLDER = Object.freeze({
|
||||
path: 'custom_nodes/ComfyUI_PromptManager/prompts.db', source: 'node_folder', movable: true, keep: false, offer: true,
|
||||
target: 'user/__promptmanager', problem: null, size_bytes: 1_700_000_000,
|
||||
target_source: 'default', config_moves: true,
|
||||
});
|
||||
|
||||
function scope() {
|
||||
const controller = new AbortController();
|
||||
return { signal: controller.signal, abort: () => controller.abort(), listen() {}, cleanup() {} };
|
||||
}
|
||||
|
||||
function activityStore({ running = [] } = {}) {
|
||||
const listeners = new Set();
|
||||
let jobs = [];
|
||||
return {
|
||||
tracked: [], loads: 0,
|
||||
track(job) { this.tracked.push(job); jobs = [...jobs, job]; return job; },
|
||||
async load() { this.loads += 1; },
|
||||
running() { return running; },
|
||||
jobs() { return jobs; },
|
||||
set(next) { jobs = next; for (const listener of listeners) listener(); },
|
||||
subscribe(listener) { listeners.add(listener); return () => listeners.delete(listener); },
|
||||
};
|
||||
}
|
||||
|
||||
class ApiError extends Error {
|
||||
constructor(status, message, code) { super(message); this.status = status; this.code = code; }
|
||||
}
|
||||
|
||||
test('the row says where the library is and why, the choice and any problem', async () => {
|
||||
const { locationDescription, SOURCE_TEXT } = await locationModule;
|
||||
assert.equal(locationDescription(NODE_FOLDER),
|
||||
'custom_nodes/ComfyUI_PromptManager/prompts.db · The node folder, where an update or reinstall can delete it.');
|
||||
assert.equal(locationDescription({ ...NODE_FOLDER, keep: true }),
|
||||
`${NODE_FOLDER.path} · ${SOURCE_TEXT.node_folder}. You chose to keep it here.`);
|
||||
assert.equal(locationDescription({ path: 'user/__promptmanager/prompts.db', source: 'default', movable: false, keep: true }),
|
||||
'user/__promptmanager/prompts.db · ComfyUI\'s user folder, which updates and reinstalls of the node leave alone.');
|
||||
assert.match(locationDescription({ ...NODE_FOLDER, problem: 'data-dir-empty' }), /PROMPT_MANAGER_DATA_DIR has no library yet/);
|
||||
assert.match(locationDescription({ path: 'prompts.db', source: 'environment', movable: false }), /PROMPT_MANAGER_DATA_DIR names/);
|
||||
assert.match(locationDescription({ path: 'prompts.db', source: 'config', movable: false }), /config\.json/);
|
||||
});
|
||||
|
||||
test('the notice and the review name the actual target and what moves', async () => {
|
||||
const { bannerText, reviewGrid, locationDescription } = await locationModule;
|
||||
assert.match(bannerText(NODE_FOLDER), /Move it to ComfyUI's user folder: it is copied/);
|
||||
assert.match(bannerText({ ...NODE_FOLDER, target_source: 'environment' }),
|
||||
/Move it to the folder PROMPT_MANAGER_DATA_DIR names: it is copied/);
|
||||
assert.match(bannerText(null), /ComfyUI's user folder/);
|
||||
assert.equal(reviewGrid({ ...NODE_FOLDER, config_moves: false })[3][1],
|
||||
'The database backups (config.json stays where PROMPT_MANAGER_CONFIG_PATH names it)');
|
||||
assert.equal(reviewGrid({ ...NODE_FOLDER, config_moves: true })[3][1], 'config.json and the database backups');
|
||||
assert.match(locationDescription({ path: 'custom_nodes/pm/data/my.db', source: 'config', movable: false,
|
||||
problem: 'inside-node-folder' }), /inside the node folder, where an update or reinstall deletes it/);
|
||||
});
|
||||
|
||||
test('only a well-formed location is shown', async () => {
|
||||
const { isLocation } = await locationModule;
|
||||
assert.equal(isLocation(NODE_FOLDER), true);
|
||||
for (const bad of [null, undefined, 'x', {}, { ...NODE_FOLDER, source: 'elsewhere' }, { ...NODE_FOLDER, path: 3 },
|
||||
{ ...NODE_FOLDER, movable: 'yes' }]) assert.equal(isLocation(bad), false);
|
||||
});
|
||||
|
||||
test('sizes read as the server would write them', async () => {
|
||||
const { sizeText } = await locationModule;
|
||||
assert.equal(sizeText(1_700_000_000), '1.6 GB');
|
||||
assert.equal(sizeText(120 * 2 ** 20), '120.0 MB');
|
||||
assert.equal(sizeText(1500), '2 KB');
|
||||
assert.equal(sizeText(-1), 'Unknown');
|
||||
assert.equal(sizeText(undefined), 'Unknown');
|
||||
});
|
||||
|
||||
test('the review names both folders, the size and what happens to the old files', async () => {
|
||||
const { reviewGrid } = await locationModule;
|
||||
assert.deepEqual(reviewGrid(NODE_FOLDER), [
|
||||
['From', 'custom_nodes/ComfyUI_PromptManager/prompts.db'],
|
||||
['To', 'user/__promptmanager/prompts.db'],
|
||||
['Size', '1.6 GB'],
|
||||
['Moves with it', 'config.json and the database backups'],
|
||||
['Old files', 'Renamed *.moved-to-…, never deleted'],
|
||||
]);
|
||||
});
|
||||
|
||||
test('Move library starts the job, tracks it and opens its view', async () => {
|
||||
const { startLibraryMove, MOVE_URL } = await locationModule;
|
||||
const posts = [];
|
||||
const api = { post: async (url, body) => { posts.push([url, body]); return { job: { id: 'j1', kind: 'library_move' } }; } };
|
||||
const activity = activityStore();
|
||||
const routes = [];
|
||||
const job = await startLibraryMove({ api, activity, scope: scope(), notify: () => assert.fail('no toast'),
|
||||
navigate: hash => routes.push(hash) });
|
||||
assert.equal(job.id, 'j1');
|
||||
assert.deepEqual(posts, [[MOVE_URL, {}]]);
|
||||
assert.deepEqual(activity.tracked.map(item => item.id), ['j1']);
|
||||
assert.deepEqual(routes, ['#/activity/j1']);
|
||||
});
|
||||
|
||||
test('a busy server names the running job; another refusal says the server reason', async () => {
|
||||
const { startLibraryMove, START_FAILED } = await locationModule;
|
||||
const running = { id: 'scan1', title: 'Metadata scan' };
|
||||
const activity = activityStore({ running: [running] });
|
||||
const toasts = [];
|
||||
const busy = { post: async () => { throw new ApiError(409, 'busy', 'busy'); } };
|
||||
assert.equal(await startLibraryMove({ api: busy, activity, scope: scope(), notify: t => toasts.push(t),
|
||||
navigate: () => assert.fail('nothing opens') }), null);
|
||||
assert.equal(activity.loads, 1);
|
||||
assert.deepEqual(toasts.at(-1), ['Another job is running: ', { text: 'Metadata scan', href: '#/activity/scan1' }, '.']);
|
||||
const short = { post: async () => { throw new ApiError(507, 'There is not enough free space for the copy.', 'insufficient-space'); } };
|
||||
await startLibraryMove({ api: short, activity, scope: scope(), notify: t => toasts.push(t), navigate: () => {} });
|
||||
assert.equal(toasts.at(-1), 'There is not enough free space for the copy.');
|
||||
const broken = { post: async () => { throw new TypeError('network'); } };
|
||||
await startLibraryMove({ api: broken, activity, scope: scope(), notify: t => toasts.push(t), navigate: () => {} });
|
||||
assert.equal(toasts.at(-1), START_FAILED);
|
||||
const gone = scope();
|
||||
gone.abort();
|
||||
const before = toasts.length;
|
||||
await startLibraryMove({ api: broken, activity, scope: gone, notify: t => toasts.push(t), navigate: () => {} });
|
||||
assert.equal(toasts.length, before);
|
||||
});
|
||||
|
||||
test('Keep it here posts the choice and resolves the new location, or null with a toast', async () => {
|
||||
const { keepLibrary, KEEP_URL, KEEP_FAILED } = await locationModule;
|
||||
const posts = [];
|
||||
const kept = { ...NODE_FOLDER, keep: true, offer: false };
|
||||
const api = { post: async (url, body) => { posts.push([url, body]); return { location: kept }; } };
|
||||
assert.deepEqual(await keepLibrary({ api, scope: scope() }), kept);
|
||||
assert.deepEqual(posts, [[KEEP_URL, { keep: true }]]);
|
||||
const toasts = [];
|
||||
const failing = { post: async () => { throw new ApiError(500, 'no', 'config-write-failed'); } };
|
||||
assert.equal(await keepLibrary({ api: failing, scope: scope(), notify: t => toasts.push(t) }), null);
|
||||
assert.deepEqual(toasts, [KEEP_FAILED]);
|
||||
const odd = { post: async () => ({ location: { path: 1 } }) };
|
||||
assert.equal(await keepLibrary({ api: odd, scope: scope() }), null);
|
||||
});
|
||||
|
||||
test('the end of a Move library job is reported once; other jobs are ignored', async () => {
|
||||
const { onMoveEnded } = await locationModule;
|
||||
const activity = activityStore();
|
||||
let ended = 0;
|
||||
const unsubscribe = onMoveEnded(activity, () => { ended += 1; });
|
||||
activity.set([{ id: 'm', kind: 'library_move', state: 'running' }, { id: 's', kind: 'scan', state: 'running' }]);
|
||||
assert.equal(ended, 0);
|
||||
activity.set([{ id: 'm', kind: 'library_move', state: 'running' }, { id: 's', kind: 'scan', state: 'finished' }]);
|
||||
assert.equal(ended, 0);
|
||||
activity.set([{ id: 'm', kind: 'library_move', state: 'failed' }]);
|
||||
assert.equal(ended, 1);
|
||||
activity.set([{ id: 'm', kind: 'library_move', state: 'failed' }]);
|
||||
assert.equal(ended, 1);
|
||||
unsubscribe();
|
||||
activity.set([{ id: 'n', kind: 'library_move', state: 'running' }]);
|
||||
activity.set([{ id: 'n', kind: 'library_move', state: 'finished' }]);
|
||||
assert.equal(ended, 1);
|
||||
});
|
||||
@@ -72,21 +72,35 @@ test('removing the last entry states auto-detect (D98)', async () => {
|
||||
assert.equal(NOTES.removedAll, 'No source is configured. After saving, ComfyUI’s output folder is used.');
|
||||
});
|
||||
|
||||
test('auto-detected sources are listed without Remove, with the line that says how they were chosen', async () => {
|
||||
test('auto-detected sources are listed without Remove and read as the default: "ComfyUI/<path> · <watch> · default"', async () => {
|
||||
const { sourceRows, NOTES } = await sourcesModule;
|
||||
const auto = { configured: [], effective: [{ label: 'ComfyUI output', path: 'output', root_index: 0, watch: 'manual' }], auto_detected: true };
|
||||
const auto = { configured: [], auto_detected: true,
|
||||
effective: [{ label: 'ComfyUI output', path: 'output', root_index: 0, watch: 'watching', in_comfyui: true }] };
|
||||
// The row says it is the default, so no line under the list repeats it; the scan keeps the server's public path.
|
||||
assert.deepEqual(sourceRows({ sources: auto, saved: [], current: [] }), { rows: [
|
||||
{ key: 'auto:0', title: 'ComfyUI output', description: 'output · manual scan', scan: { label: 'Scan folder', root: 0, path: 'output' },
|
||||
remove: null },
|
||||
], note: NOTES.autoDetected });
|
||||
// Every configured source unusable: ComfyUI's output folder is used instead.
|
||||
{ key: 'auto:0', title: 'ComfyUI output', description: 'ComfyUI/output · watching · default',
|
||||
scan: { label: 'Scan changes', root: 0, path: 'output' }, remove: null },
|
||||
], note: null });
|
||||
assert.equal('autoDetected' in NOTES, false);
|
||||
// An output folder outside the ComfyUI base (--output-directory) is shown by its own name.
|
||||
const outside = { ...auto, effective: [{ label: 'Renders', path: 'Renders', root_index: 0, watch: 'manual', in_comfyui: false }] };
|
||||
assert.equal(sourceRows({ sources: outside, saved: [], current: [] }).rows[0].description, 'Renders · manual scan · default');
|
||||
// Every configured source unusable: ComfyUI's output folder is used instead, and the line says why.
|
||||
const fallback = { ...auto, configured: [{ label: 'gone', path: 'gone', valid: false, code: 'not-found', reason: 'Gallery root does not exist' }] };
|
||||
const result = sourceRows({ sources: fallback, saved: ['gone'], current: ['gone'] });
|
||||
assert.deepEqual(result.rows.map(row => [row.title, row.remove]), [['gone', 'gone'], ['ComfyUI output', null]]);
|
||||
assert.deepEqual(result.rows.map(row => [row.title, row.description, row.remove]),
|
||||
[['gone', 'gone · Gallery root does not exist', 'gone'], ['ComfyUI output', 'ComfyUI/output · watching · default', null]]);
|
||||
assert.equal(result.note, NOTES.fallback);
|
||||
assert.equal(sourceRows({ sources: { configured: [], effective: [], auto_detected: false }, saved: [], current: [] }).note, NOTES.none);
|
||||
});
|
||||
|
||||
test('a configured source keeps its plain public path even inside ComfyUI', async () => {
|
||||
const { sourceRows } = await sourcesModule;
|
||||
const sources = { ...STATUS, effective: STATUS.effective.map(source => ({ ...source, in_comfyui: true })) };
|
||||
const { rows } = sourceRows({ sources, saved: ['output', 'input'], current: ['output', 'input'] });
|
||||
assert.deepEqual(rows.map(row => row.description), ['output · watching', 'input · manual scan']);
|
||||
});
|
||||
|
||||
test('a path listed twice in the draft is one row', async () => {
|
||||
const { sourceRows } = await sourcesModule;
|
||||
assert.equal(sourceRows({ sources: STATUS, saved: ['output', 'input'], current: ['output', 'output'] }).rows.length, 1);
|
||||
|
||||
@@ -0,0 +1,427 @@
|
||||
"""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()
|
||||
# Lower case: a public path is case-folded on Windows (normcase).
|
||||
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),
|
||||
"target_source": "default",
|
||||
"config_moves": True,
|
||||
},
|
||||
)
|
||||
|
||||
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 TestLocationDetails(LibraryAPITestCase):
|
||||
async def test_the_target_and_what_moves_follow_the_environment(self):
|
||||
data = os.path.join(self.tmpdir, "pm-data")
|
||||
self.environ = {
|
||||
ll.DATA_DIR_ENV: data,
|
||||
ll.CONFIG_PATH_ENV: os.path.join(self.tmpdir, "mine.json"),
|
||||
}
|
||||
location = (await self.json("GET", URL))["location"]
|
||||
self.assertEqual(location["target_source"], "environment")
|
||||
self.assertFalse(location["config_moves"])
|
||||
self.assertEqual(location["problem"], "data-dir-empty")
|
||||
|
||||
async def test_a_configured_library_inside_the_node_folder_is_flagged(self):
|
||||
self.api._library_resolver = lambda: ll.LibraryLocation(
|
||||
ll.SOURCE_CONFIG,
|
||||
self.api.db.model.db_path,
|
||||
os.path.join(self.ext, "config.json"),
|
||||
self.ext,
|
||||
)
|
||||
self.api._extension_root = self.ext
|
||||
location = (await self.json("GET", URL))["location"]
|
||||
self.assertEqual(location["problem"], "inside-node-folder")
|
||||
self.assertFalse(location["movable"])
|
||||
|
||||
async def test_a_startup_after_a_crash_finishes_the_rename(self):
|
||||
target_db = os.path.join(self.target, "prompts.db")
|
||||
os.makedirs(self.target)
|
||||
with open(target_db, "wb") as handle:
|
||||
handle.write(b"x")
|
||||
lm.write_marker(self.target, ll.MARKER_MOVED, self.api.db.model.db_path)
|
||||
old = self.api.db.model.db_path
|
||||
self.api.db.model.close_all()
|
||||
moved = ll.LibraryLocation(
|
||||
ll.SOURCE_DEFAULT,
|
||||
target_db,
|
||||
os.path.join(self.target, "config.json"),
|
||||
self.target,
|
||||
self.target,
|
||||
)
|
||||
self.api._library_resolver = lambda: moved
|
||||
with patch.object(self.api.db.model, "db_path", target_db):
|
||||
renamed = self.api.finish_interrupted_library_move()
|
||||
self.assertEqual(renamed, ["prompts.db.moved-to-__promptmanager"])
|
||||
self.assertFalse(os.path.exists(old))
|
||||
|
||||
|
||||
class TestLocationThroughASymlinkedNodeFolder(LibraryAPITestCase):
|
||||
"""custom_nodes/<name> is a symlink to a checkout outside ComfyUI.
|
||||
|
||||
The library location reads relative to ComfyUI through the node folder's
|
||||
own (unresolved) place in custom_nodes, not as a bare file name; the
|
||||
shared ``_public_path`` (sources, jobs) keeps resolving symlinks.
|
||||
"""
|
||||
|
||||
async def get_application(self):
|
||||
app = await super().get_application()
|
||||
self.real = os.path.join(self.tmpdir, "checkout", "ComfyUI_PromptManager")
|
||||
os.makedirs(self.real)
|
||||
self.link = os.path.join(str(self.comfy_dir), "custom_nodes", "promptmanager")
|
||||
try:
|
||||
os.symlink(self.real, self.link, target_is_directory=True)
|
||||
except (OSError, NotImplementedError) as exc:
|
||||
self.skipTest(f"symlinks not permitted here: {exc}")
|
||||
previous = self.api.db
|
||||
self.api.db = PromptDatabase(os.path.join(self.link, "prompts.db"))
|
||||
previous.close_all()
|
||||
self.api.db.save_prompt("a library prompt")
|
||||
self.api._library_resolver = lambda: ll.resolve_library(
|
||||
None, environ=self.environ, extension_root=self.link, folder_paths=self.fp
|
||||
)
|
||||
self.api._extension_root = self.link
|
||||
return app
|
||||
|
||||
async def test_location_names_the_node_folder_in_custom_nodes(self):
|
||||
location = (await self.json("GET", URL))["location"]
|
||||
self.assertEqual(location["source"], "node_folder")
|
||||
self.assertEqual(location["path"], "custom_nodes/promptmanager/prompts.db")
|
||||
self.assertEqual(location["target"], "user/__promptmanager")
|
||||
overview = await self.json("GET", "/prompt_manager/system/overview")
|
||||
self.assertEqual(
|
||||
overview["location"]["path"], "custom_nodes/promptmanager/prompts.db"
|
||||
)
|
||||
|
||||
async def test_a_resolved_spelling_of_the_database_reads_the_same(self):
|
||||
from py.api import library_location as routes
|
||||
|
||||
resolved = os.path.join(os.path.realpath(self.real), "prompts.db")
|
||||
self.assertEqual(
|
||||
routes.node_folder_public_path(resolved, self.link, []),
|
||||
"custom_nodes/promptmanager/prompts.db",
|
||||
)
|
||||
|
||||
async def test_the_shared_public_path_is_unchanged(self):
|
||||
self.assertEqual(
|
||||
self.api._public_path(os.path.join(self.link, "prompts.db")), "prompts.db"
|
||||
)
|
||||
|
||||
|
||||
class TestNodeFolderPublicPath(LibraryAPITestCase):
|
||||
def display(self, path, root=None, custom_nodes=()):
|
||||
from py.api import library_location as routes
|
||||
|
||||
return routes.node_folder_public_path(path, root or self.ext, custom_nodes)
|
||||
|
||||
async def test_a_plain_node_folder_reads_relative_to_custom_nodes(self):
|
||||
self.assertEqual(
|
||||
self.display(os.path.join(self.ext, "prompts.db")),
|
||||
"custom_nodes/pm/prompts.db",
|
||||
)
|
||||
self.assertEqual(
|
||||
self.display(os.path.join(self.ext, "data", "prompts.db")),
|
||||
"custom_nodes/pm/data/prompts.db",
|
||||
)
|
||||
|
||||
async def test_a_custom_nodes_entry_from_comfyui_names_the_folder(self):
|
||||
nodes = os.path.join(self.tmpdir, "more_nodes")
|
||||
root = os.path.join(nodes, "pm")
|
||||
self.assertEqual(
|
||||
self.display(os.path.join(root, "prompts.db"), root, [nodes]),
|
||||
"more_nodes/pm/prompts.db",
|
||||
)
|
||||
|
||||
async def test_anything_else_is_not_a_node_folder_path(self):
|
||||
# Outside the node folder, a parent that is no custom_nodes folder, or
|
||||
# a spelling that climbs out: the caller falls back to _public_path.
|
||||
self.assertIsNone(self.display(os.path.join(self.user, "prompts.db")))
|
||||
loose = os.path.join(self.tmpdir, "somewhere", "pm")
|
||||
self.assertIsNone(self.display(os.path.join(loose, "prompts.db"), loose))
|
||||
self.assertIsNone(self.display(os.path.join(self.ext, "..", "prompts.db")))
|
||||
self.assertIsNone(self.display(""))
|
||||
self.assertIsNone(self.display(None))
|
||||
|
||||
async def test_a_library_outside_the_node_folder_keeps_the_public_path(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,
|
||||
)
|
||||
self.api._extension_root = os.path.join(self.tmpdir, "elsewhere")
|
||||
location = (await self.json("GET", URL))["location"]
|
||||
self.assertEqual(location["path"], "custom_nodes/pm/prompts.db")
|
||||
|
||||
|
||||
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)
|
||||
@@ -66,7 +66,7 @@ class TestIngestionKeepsDefaultProfile(unittest.TestCase):
|
||||
handler.link_image_to_prompt("/out/f03.png", {"id": 1}, metadata)
|
||||
self.assertEqual(len(spy.calls), 2)
|
||||
spy.assert_default_and_gated(self)
|
||||
handler.db_manager.link_image_to_prompt.assert_called_once()
|
||||
handler.db_manager.link_image_to_prompt_or_raise.assert_called_once()
|
||||
|
||||
def test_admin_scan_ingestion_uses_the_default_profile(self):
|
||||
from py.api import admin
|
||||
|
||||
@@ -0,0 +1,369 @@
|
||||
"""Where the library lives (#249 S6b): the one resolver, ``utils/library_location``.
|
||||
|
||||
Every case builds its own extension root (the node folder) and ComfyUI user
|
||||
folder in a temp directory and hands the resolver a fake ``folder_paths``, so
|
||||
no test reads the real environment, the repository's own files or ComfyUI.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from utils import library_location as ll # noqa: E402
|
||||
|
||||
|
||||
def _touch(path, data=b"x"):
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
with open(path, "wb") as handle:
|
||||
handle.write(data)
|
||||
|
||||
|
||||
def _folder_paths(base, *, helper=True):
|
||||
"""A ComfyUI ``folder_paths`` stand-in rooted at ``base``."""
|
||||
module = types.SimpleNamespace(base_path=base)
|
||||
for name in ("output", "input", "temp", "user"):
|
||||
directory = os.path.join(base, name)
|
||||
os.makedirs(directory, exist_ok=True)
|
||||
setattr(module, f"get_{name}_directory", lambda d=directory: d)
|
||||
if helper:
|
||||
module.get_system_user_directory = lambda name="system": os.path.join(
|
||||
base, "user", "__" + name
|
||||
)
|
||||
return module
|
||||
|
||||
|
||||
class _Case(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self._tmp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self._tmp.cleanup)
|
||||
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 = _folder_paths(self.comfy)
|
||||
self.default_dir = os.path.join(self.comfy, "user", "__promptmanager")
|
||||
ll.clear_override()
|
||||
self.addCleanup(ll.clear_override)
|
||||
|
||||
def resolve(self, default_path=None, environ=None, **kwargs):
|
||||
kwargs.setdefault("folder_paths", self.fp)
|
||||
return ll.resolve_library(
|
||||
default_path, environ=environ or {}, extension_root=self.ext, **kwargs
|
||||
)
|
||||
|
||||
def config_location(self, environ=None, **kwargs):
|
||||
kwargs.setdefault("folder_paths", self.fp)
|
||||
return ll.config_location(
|
||||
environ=environ or {}, extension_root=self.ext, **kwargs
|
||||
)
|
||||
|
||||
|
||||
class TestDefaultDataFolder(_Case):
|
||||
def test_uses_comfyui_system_user_helper(self):
|
||||
self.assertEqual(ll.default_data_dir(self.fp), self.default_dir)
|
||||
|
||||
def test_falls_back_to_user_directory_without_helper(self):
|
||||
fp = _folder_paths(self.comfy, helper=False)
|
||||
self.assertEqual(ll.default_data_dir(fp), self.default_dir)
|
||||
|
||||
def test_none_outside_comfyui(self):
|
||||
self.assertIsNone(ll.default_data_dir(None))
|
||||
self.assertIsNone(ll.default_data_dir(types.SimpleNamespace()))
|
||||
|
||||
|
||||
class TestResolutionOrder(_Case):
|
||||
def test_new_install_uses_default_folder_and_creates_nothing_in_node_folder(self):
|
||||
location = self.resolve()
|
||||
self.assertEqual(location.source, ll.SOURCE_DEFAULT)
|
||||
self.assertEqual(location.db_path, os.path.join(self.default_dir, "prompts.db"))
|
||||
self.assertEqual(
|
||||
location.config_path, os.path.join(self.default_dir, "config.json")
|
||||
)
|
||||
self.assertFalse(location.movable)
|
||||
self.assertEqual(os.listdir(self.ext), [])
|
||||
|
||||
def test_legacy_node_folder_library_stays_in_use(self):
|
||||
_touch(os.path.join(self.ext, "prompts.db"))
|
||||
location = self.resolve()
|
||||
self.assertEqual(location.source, ll.SOURCE_NODE_FOLDER)
|
||||
self.assertEqual(location.db_path, os.path.join(self.ext, "prompts.db"))
|
||||
self.assertEqual(location.config_path, os.path.join(self.ext, "config.json"))
|
||||
self.assertEqual(location.data_dir, self.ext)
|
||||
self.assertTrue(location.movable)
|
||||
self.assertEqual(location.target_dir, self.default_dir)
|
||||
|
||||
def test_legacy_wins_over_an_unmarked_default_folder_library(self):
|
||||
_touch(os.path.join(self.ext, "prompts.db"))
|
||||
_touch(os.path.join(self.default_dir, "prompts.db"))
|
||||
self.assertEqual(self.resolve().source, ll.SOURCE_NODE_FOLDER)
|
||||
|
||||
def test_explicit_default_path_is_honoured_as_today(self):
|
||||
_touch(os.path.join(self.ext, "prompts.db"))
|
||||
elsewhere = os.path.join(self.base, "elsewhere", "pm.db")
|
||||
location = self.resolve(elsewhere)
|
||||
self.assertEqual(location.source, ll.SOURCE_CONFIG)
|
||||
self.assertEqual(location.db_path, elsewhere)
|
||||
self.assertFalse(location.movable)
|
||||
relative = self.resolve("data/my.db")
|
||||
self.assertEqual(relative.db_path, os.path.join(self.ext, "data", "my.db"))
|
||||
self.assertEqual(relative.source, ll.SOURCE_CONFIG)
|
||||
|
||||
def test_default_setting_value_is_not_explicit(self):
|
||||
for value in (None, "", "prompts.db"):
|
||||
self.assertEqual(self.resolve(value).source, ll.SOURCE_DEFAULT, value)
|
||||
|
||||
def test_environment_folder_wins(self):
|
||||
data = os.path.join(self.base, "pm-data")
|
||||
location = self.resolve(environ={ll.DATA_DIR_ENV: data})
|
||||
self.assertEqual(location.source, ll.SOURCE_ENVIRONMENT)
|
||||
self.assertEqual(location.db_path, os.path.join(data, "prompts.db"))
|
||||
self.assertEqual(location.config_path, os.path.join(data, "config.json"))
|
||||
|
||||
def test_environment_folder_with_a_library_wins_over_legacy(self):
|
||||
data = os.path.join(self.base, "pm-data")
|
||||
_touch(os.path.join(data, "prompts.db"))
|
||||
_touch(os.path.join(self.ext, "prompts.db"))
|
||||
location = self.resolve(environ={ll.DATA_DIR_ENV: data})
|
||||
self.assertEqual(location.source, ll.SOURCE_ENVIRONMENT)
|
||||
|
||||
def test_empty_environment_folder_never_replaces_a_legacy_library(self):
|
||||
data = os.path.join(self.base, "pm-data")
|
||||
_touch(os.path.join(self.ext, "prompts.db"))
|
||||
location = self.resolve(environ={ll.DATA_DIR_ENV: data})
|
||||
self.assertEqual(location.source, ll.SOURCE_NODE_FOLDER)
|
||||
self.assertEqual(location.db_path, os.path.join(self.ext, "prompts.db"))
|
||||
self.assertEqual(location.problem, "data-dir-empty")
|
||||
self.assertEqual(location.target_dir, data)
|
||||
self.assertTrue(location.movable)
|
||||
|
||||
def test_config_path_environment_is_honoured(self):
|
||||
custom = os.path.join(self.base, "cfg", "custom.json")
|
||||
location = self.resolve(environ={ll.CONFIG_PATH_ENV: custom})
|
||||
self.assertEqual(location.config_path, custom)
|
||||
self.assertEqual(location.source, ll.SOURCE_DEFAULT)
|
||||
|
||||
def test_without_comfyui_the_node_folder_is_used(self):
|
||||
location = self.resolve(folder_paths=None)
|
||||
self.assertEqual(location.source, ll.SOURCE_NODE_FOLDER)
|
||||
self.assertEqual(location.db_path, os.path.join(self.ext, "prompts.db"))
|
||||
self.assertFalse(location.movable)
|
||||
|
||||
def test_override_after_a_live_move_wins(self):
|
||||
_touch(os.path.join(self.ext, "prompts.db"))
|
||||
moved = ll.LibraryLocation(
|
||||
source=ll.SOURCE_DEFAULT,
|
||||
db_path=os.path.join(self.default_dir, "prompts.db"),
|
||||
config_path=os.path.join(self.default_dir, "config.json"),
|
||||
data_dir=self.default_dir,
|
||||
target_dir=self.default_dir,
|
||||
)
|
||||
ll.switch_to(moved)
|
||||
self.assertEqual(self.resolve(), moved)
|
||||
self.assertEqual(self.config_location(), (moved.config_path, moved.data_dir))
|
||||
ll.clear_override()
|
||||
self.assertEqual(self.resolve().source, ll.SOURCE_NODE_FOLDER)
|
||||
|
||||
|
||||
class TestMarker(_Case):
|
||||
def _marker(self, directory, state):
|
||||
os.makedirs(directory, exist_ok=True)
|
||||
with open(os.path.join(directory, ll.MARKER_NAME), "w") as handle:
|
||||
json.dump({"version": 1, "state": state}, handle)
|
||||
|
||||
def test_moved_marker_wins_while_the_old_file_was_not_renamed_yet(self):
|
||||
_touch(os.path.join(self.ext, "prompts.db"))
|
||||
_touch(os.path.join(self.default_dir, "prompts.db"))
|
||||
self._marker(self.default_dir, "moved")
|
||||
location = self.resolve()
|
||||
self.assertEqual(location.source, ll.SOURCE_DEFAULT)
|
||||
self.assertEqual(location.db_path, os.path.join(self.default_dir, "prompts.db"))
|
||||
self.assertEqual(
|
||||
location.config_path, os.path.join(self.default_dir, "config.json")
|
||||
)
|
||||
|
||||
def test_pending_marker_keeps_the_old_library(self):
|
||||
_touch(os.path.join(self.ext, "prompts.db"))
|
||||
_touch(os.path.join(self.default_dir, "prompts.db"))
|
||||
self._marker(self.default_dir, "pending")
|
||||
self.assertEqual(self.resolve().source, ll.SOURCE_NODE_FOLDER)
|
||||
data = os.path.join(self.base, "pm-data")
|
||||
_touch(os.path.join(data, "prompts.db"))
|
||||
self._marker(data, "pending")
|
||||
location = self.resolve(environ={ll.DATA_DIR_ENV: data})
|
||||
self.assertEqual(location.source, ll.SOURCE_NODE_FOLDER)
|
||||
self.assertEqual(location.problem, "move-unfinished")
|
||||
|
||||
def test_unreadable_marker_is_not_trusted(self):
|
||||
_touch(os.path.join(self.ext, "prompts.db"))
|
||||
_touch(os.path.join(self.default_dir, "prompts.db"))
|
||||
_touch(os.path.join(self.default_dir, ll.MARKER_NAME), b"{not json")
|
||||
self.assertEqual(self.resolve().source, ll.SOURCE_NODE_FOLDER)
|
||||
self.assertEqual(ll.read_marker(self.default_dir), {})
|
||||
|
||||
def test_renamed_old_file_is_never_picked(self):
|
||||
_touch(os.path.join(self.ext, "prompts.db.moved-to-__promptmanager"))
|
||||
_touch(os.path.join(self.default_dir, "prompts.db"))
|
||||
self.assertEqual(self.resolve().source, ll.SOURCE_DEFAULT)
|
||||
|
||||
|
||||
class TestConfigLocation(_Case):
|
||||
def test_legacy_config_without_a_database_keeps_being_read(self):
|
||||
_touch(os.path.join(self.ext, "config.json"), b"{}")
|
||||
path, data_dir = self.config_location()
|
||||
self.assertEqual(path, os.path.join(self.ext, "config.json"))
|
||||
self.assertEqual(data_dir, self.ext)
|
||||
|
||||
def test_legacy_config_alone_never_puts_a_new_database_in_the_node_folder(self):
|
||||
_touch(os.path.join(self.ext, "config.json"), b"{}")
|
||||
location = self.resolve()
|
||||
self.assertEqual(location.db_path, os.path.join(self.default_dir, "prompts.db"))
|
||||
self.assertEqual(location.config_path, os.path.join(self.ext, "config.json"))
|
||||
self.assertFalse(location.movable)
|
||||
|
||||
def test_new_install_reads_config_from_the_data_folder(self):
|
||||
path, data_dir = self.config_location()
|
||||
self.assertEqual(path, os.path.join(self.default_dir, "config.json"))
|
||||
self.assertEqual(data_dir, self.default_dir)
|
||||
|
||||
def test_logs_and_cache_follow_the_data_folder(self):
|
||||
self.assertEqual(
|
||||
ll.logs_dir(environ={}, extension_root=self.ext, folder_paths=self.fp),
|
||||
os.path.join(self.default_dir, "logs"),
|
||||
)
|
||||
_touch(os.path.join(self.ext, "prompts.db"))
|
||||
self.assertEqual(
|
||||
ll.logs_dir(environ={}, extension_root=self.ext, folder_paths=self.fp),
|
||||
os.path.join(self.ext, "logs"),
|
||||
)
|
||||
self.assertEqual(
|
||||
ll.cache_dir(
|
||||
"lora_images", environ={}, extension_root=self.ext, folder_paths=self.fp
|
||||
),
|
||||
os.path.join(self.ext, "data", "lora_images"),
|
||||
)
|
||||
|
||||
|
||||
class TestDataDirValidation(_Case):
|
||||
def problem(self, path, **kwargs):
|
||||
kwargs.setdefault("folder_paths", self.fp)
|
||||
return ll.data_dir_problem(path, **kwargs)
|
||||
|
||||
def test_valid_folder(self):
|
||||
self.assertIsNone(self.problem(os.path.join(self.base, "pm-data")))
|
||||
self.assertIsNone(self.problem(self.default_dir))
|
||||
|
||||
def test_relative_and_empty_paths_are_refused(self):
|
||||
self.assertEqual(self.problem("pm-data"), "data-dir-not-absolute")
|
||||
self.assertEqual(self.problem(""), "data-dir-not-absolute")
|
||||
|
||||
def test_served_folders_are_refused(self):
|
||||
for name in ("output", "input", "temp"):
|
||||
inside = os.path.join(self.comfy, name, "pm")
|
||||
self.assertEqual(self.problem(inside), "data-dir-served", name)
|
||||
self.assertEqual(
|
||||
self.problem(os.path.join(self.comfy, "output")), "data-dir-served"
|
||||
)
|
||||
|
||||
def test_public_user_folder_is_refused_but_system_users_are_fine(self):
|
||||
public = os.path.join(self.comfy, "user", "default", "pm")
|
||||
self.assertEqual(self.problem(public), "data-dir-served")
|
||||
self.assertIsNone(self.problem(os.path.join(self.comfy, "user", "__pm")))
|
||||
|
||||
def test_gallery_source_is_refused(self):
|
||||
source = os.path.join(self.base, "gallery")
|
||||
os.makedirs(source)
|
||||
self.assertEqual(
|
||||
self.problem(os.path.join(source, "pm"), gallery_roots=[source]),
|
||||
"data-dir-served",
|
||||
)
|
||||
self.assertEqual(
|
||||
self.problem(
|
||||
os.path.join(source, "pm"),
|
||||
environ={ll.EXTRA_GALLERY_ROOTS_ENV: source},
|
||||
),
|
||||
"data-dir-served",
|
||||
)
|
||||
|
||||
def test_symlink_into_output_is_refused(self):
|
||||
link = os.path.join(self.base, "link-to-output")
|
||||
try:
|
||||
os.symlink(os.path.join(self.comfy, "output"), link)
|
||||
except (OSError, NotImplementedError):
|
||||
self.skipTest("symlinks are not available")
|
||||
self.assertEqual(self.problem(os.path.join(link, "pm")), "data-dir-served")
|
||||
|
||||
def test_case_variant_spelling_of_output_is_refused_where_the_volume_folds_case(
|
||||
self,
|
||||
):
|
||||
upper = os.path.join(self.comfy, "OUTPUT")
|
||||
if not os.path.isdir(upper):
|
||||
self.skipTest("this volume is case-sensitive")
|
||||
self.assertEqual(self.problem(os.path.join(upper, "pm")), "data-dir-served")
|
||||
|
||||
def test_invalid_environment_folder_is_ignored_with_a_problem(self):
|
||||
inside = os.path.join(self.comfy, "output", "pm")
|
||||
location = self.resolve(environ={ll.DATA_DIR_ENV: inside})
|
||||
self.assertEqual(location.source, ll.SOURCE_DEFAULT)
|
||||
self.assertEqual(location.problem, "data-dir-served")
|
||||
self.assertFalse(os.path.exists(inside))
|
||||
relative = self.resolve(environ={ll.DATA_DIR_ENV: "relative/dir"})
|
||||
self.assertEqual(relative.problem, "data-dir-not-absolute")
|
||||
|
||||
|
||||
class TestOneRuleForConfigAndDatabase(_Case):
|
||||
"""M1: config.json and the database never split over PROMPT_MANAGER_DATA_DIR."""
|
||||
|
||||
def test_a_source_listed_only_in_a_config_file_rules_out_the_folder_for_both(self):
|
||||
source = os.path.join(self.base, "gallery")
|
||||
os.makedirs(source)
|
||||
_touch(
|
||||
os.path.join(self.ext, "config.json"),
|
||||
json.dumps({"gallery": {"monitoring": {"directories": [source]}}}).encode(),
|
||||
)
|
||||
environ = {ll.DATA_DIR_ENV: os.path.join(source, "pm")}
|
||||
location = self.resolve(environ=environ)
|
||||
config_path, data_dir = self.config_location(environ=environ)
|
||||
self.assertEqual(location.problem, "data-dir-served")
|
||||
self.assertEqual(
|
||||
(location.config_path, location.data_dir), (config_path, data_dir)
|
||||
)
|
||||
self.assertNotEqual(location.source, ll.SOURCE_ENVIRONMENT)
|
||||
|
||||
def test_a_relative_source_is_taken_from_comfyui(self):
|
||||
self.fp.base_path = self.comfy
|
||||
os.makedirs(os.path.join(self.comfy, "pictures"))
|
||||
_touch(
|
||||
os.path.join(self.ext, "config.json"),
|
||||
json.dumps(
|
||||
{"gallery": {"monitoring": {"directories": ["pictures"]}}}
|
||||
).encode(),
|
||||
)
|
||||
environ = {ll.DATA_DIR_ENV: os.path.join(self.comfy, "pictures", "pm")}
|
||||
self.assertEqual(self.resolve(environ=environ).problem, "data-dir-served")
|
||||
self.assertEqual(self.config_location(environ=environ)[1], self.ext)
|
||||
|
||||
def test_unreadable_or_odd_config_files_add_no_source(self):
|
||||
for payload in (
|
||||
b"{not json",
|
||||
b"[]",
|
||||
json.dumps({"gallery": {"monitoring": {"directories": "x"}}}).encode(),
|
||||
):
|
||||
_touch(os.path.join(self.ext, "config.json"), payload)
|
||||
environ = {ll.DATA_DIR_ENV: os.path.join(self.base, "pm-data")}
|
||||
self.assertIsNone(self.resolve(environ=environ).problem)
|
||||
|
||||
|
||||
class TestSameFolder(_Case):
|
||||
def test_same_folder_follows_links_and_case(self):
|
||||
self.assertTrue(ll.same_folder(self.ext, self.ext + os.sep))
|
||||
self.assertFalse(ll.same_folder(self.ext, self.default_dir))
|
||||
self.assertFalse(ll.same_folder(self.ext, os.path.join(self.base, "missing")))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,169 @@
|
||||
"""Every data path asks the one resolver (#249 S6b).
|
||||
|
||||
The database, the diagnostics database, config.json, the logs and the LoRA
|
||||
image cache all follow ``utils.library_location``: a folder named by
|
||||
``PROMPT_MANAGER_DATA_DIR`` (and, after a live move, the switched location).
|
||||
"""
|
||||
|
||||
import os
|
||||
import pathlib
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
import unittest.mock as mock
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
_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 utils import library_location as ll # noqa: E402
|
||||
|
||||
|
||||
class _EnvCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self._tmp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self._tmp.cleanup)
|
||||
base = os.path.realpath(self._tmp.name)
|
||||
self.data = os.path.join(base, "pm-data")
|
||||
patcher = mock.patch.dict(
|
||||
os.environ,
|
||||
{ll.DATA_DIR_ENV: self.data, ll.CONFIG_PATH_ENV: ""},
|
||||
)
|
||||
patcher.start()
|
||||
self.addCleanup(patcher.stop)
|
||||
# This test's own ComfyUI folders (the shared stub serves the system
|
||||
# temp folder as "output", which would rightly refuse self.data).
|
||||
comfy = os.path.join(base, "ComfyUI")
|
||||
fp = types.ModuleType("folder_paths")
|
||||
for name in ("output", "input", "temp", "user"):
|
||||
setattr(fp, f"get_{name}_directory", lambda n=name: os.path.join(comfy, n))
|
||||
modules = mock.patch.dict(sys.modules, {"folder_paths": fp})
|
||||
modules.start()
|
||||
self.addCleanup(modules.stop)
|
||||
ll.clear_override()
|
||||
self.addCleanup(ll.clear_override)
|
||||
|
||||
|
||||
class TestDatabasePath(_EnvCase):
|
||||
def test_database_follows_the_data_folder(self):
|
||||
from database.operations import _resolve_db_path
|
||||
|
||||
with mock.patch.dict(sys.modules, {"py.config": None}):
|
||||
path = _resolve_db_path(None)
|
||||
self.assertEqual(path, os.path.join(self.data, "prompts.db"))
|
||||
self.assertTrue(os.path.isdir(self.data))
|
||||
if os.name == "posix":
|
||||
self.assertEqual(os.stat(self.data).st_mode & 0o777, 0o700)
|
||||
|
||||
def test_diagnostics_use_the_same_resolver(self):
|
||||
from database.operations import _resolve_db_path
|
||||
from utils.diagnostics import GalleryDiagnostics
|
||||
|
||||
with mock.patch.dict(sys.modules, {"py.config": None}):
|
||||
self.assertEqual(GalleryDiagnostics._resolve_db_path(), _resolve_db_path())
|
||||
|
||||
def test_switched_location_wins(self):
|
||||
from database.operations import _resolve_db_path
|
||||
|
||||
moved = ll.LibraryLocation(
|
||||
ll.SOURCE_DEFAULT,
|
||||
os.path.join(self.data, "moved", "prompts.db"),
|
||||
os.path.join(self.data, "moved", "config.json"),
|
||||
os.path.join(self.data, "moved"),
|
||||
)
|
||||
ll.switch_to(moved)
|
||||
self.assertEqual(_resolve_db_path(None), moved.db_path)
|
||||
|
||||
|
||||
class TestConfigPath(_EnvCase):
|
||||
def test_config_json_follows_the_data_folder(self):
|
||||
from py.config import PromptManagerConfig
|
||||
|
||||
self.assertEqual(
|
||||
PromptManagerConfig.get_config_path(),
|
||||
os.path.join(self.data, "config.json"),
|
||||
)
|
||||
|
||||
def test_explicit_config_path_still_wins(self):
|
||||
from py.config import PromptManagerConfig
|
||||
|
||||
custom = os.path.join(self.data, "elsewhere.json")
|
||||
with mock.patch.dict(os.environ, {ll.CONFIG_PATH_ENV: custom}):
|
||||
self.assertEqual(PromptManagerConfig.get_config_path(), custom)
|
||||
|
||||
def test_keep_choice_round_trips(self):
|
||||
from py.config import PromptManagerConfig
|
||||
|
||||
original = PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER
|
||||
self.addCleanup(
|
||||
setattr, PromptManagerConfig, "KEEP_LIBRARY_IN_NODE_FOLDER", original
|
||||
)
|
||||
PromptManagerConfig.update_config({"library": {"keep_in_node_folder": True}})
|
||||
self.assertTrue(PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER)
|
||||
self.assertEqual(
|
||||
PromptManagerConfig.get_config()["library"], {"keep_in_node_folder": True}
|
||||
)
|
||||
PromptManagerConfig.update_config({"library": {"keep_in_node_folder": "yes"}})
|
||||
self.assertTrue(PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER)
|
||||
|
||||
|
||||
class TestLogsAndCache(_EnvCase):
|
||||
def test_logs_go_to_the_data_folder(self):
|
||||
from utils import logging_config
|
||||
from utils.logging_config import PromptManagerLogger
|
||||
|
||||
manager = object.__new__(PromptManagerLogger)
|
||||
self.addCleanup(self._close, manager)
|
||||
manager.__init__()
|
||||
self.assertEqual(manager.log_dir, pathlib.Path(self.data, "logs"))
|
||||
self.assertTrue(manager.log_dir.is_dir())
|
||||
self.assertIs(logging_config.library_location, ll)
|
||||
|
||||
@staticmethod
|
||||
def _close(manager):
|
||||
import logging
|
||||
|
||||
for handler in logging.getLogger("prompt_manager").handlers[:]:
|
||||
if getattr(handler, "_prompt_manager_owned", False):
|
||||
handler.close()
|
||||
try:
|
||||
from utils.logging_config import get_logger_manager
|
||||
|
||||
get_logger_manager()._setup_loggers()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def test_lora_image_cache_follows_the_data_folder(self):
|
||||
from py import lora_utils
|
||||
|
||||
cache = lora_utils.get_lora_image_cache_dir()
|
||||
self.assertEqual(cache, pathlib.Path(self.data, "data", "lora_images"))
|
||||
|
||||
|
||||
class TestNodeFolderUnchangedForLegacyInstalls(unittest.TestCase):
|
||||
"""A node folder that holds the library keeps every path where it was."""
|
||||
|
||||
def test_legacy_paths(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
ext = os.path.join(tmp, "ext")
|
||||
os.makedirs(ext)
|
||||
open(os.path.join(ext, "prompts.db"), "wb").close()
|
||||
fp = types.SimpleNamespace(
|
||||
get_user_directory=lambda: os.path.join(tmp, "u")
|
||||
)
|
||||
location = ll.resolve_library(
|
||||
None, environ={}, extension_root=ext, folder_paths=fp
|
||||
)
|
||||
self.assertEqual(location.db_path, os.path.join(ext, "prompts.db"))
|
||||
self.assertEqual(
|
||||
ll.logs_dir(environ={}, extension_root=ext, folder_paths=fp),
|
||||
os.path.join(ext, "logs"),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -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()
|
||||
@@ -0,0 +1,352 @@
|
||||
"""Move library, final review round (#249 S6b PR #318).
|
||||
|
||||
- I1: a switch that fails at any point is undone while both files are still
|
||||
locked: the copy is removed, the old library stays in use and a restart
|
||||
keeps it. A follow-up after a complete switch only logs.
|
||||
- I2: a background write (run counting, the image monitor) that fails
|
||||
because Move library switched runs once more, on the new library; a write
|
||||
racing the switch on a real connection (a run count, an image link) lands
|
||||
in the new library only.
|
||||
- M2: the copy is verified through a file URI that survives ``#`` and ``%``.
|
||||
- M6: a restart after a crash between the switch and the rename renames the
|
||||
old files, never deletes them.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import threading
|
||||
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 import models # noqa: E402
|
||||
from database.operations import PromptDatabase # noqa: E402
|
||||
from tests.open_handles import assert_closed # noqa: E402
|
||||
from utils import library_location as ll # noqa: E402
|
||||
from utils import usage_tracking # noqa: E402
|
||||
from utils.hashing import generate_prompt_hash # noqa: E402
|
||||
|
||||
try:
|
||||
from tests.test_library_move import KEY, SimulatedCrash, _MoveCase, _rows
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from test_library_move import KEY, SimulatedCrash, _MoveCase, _rows
|
||||
|
||||
WAIT = 30
|
||||
|
||||
|
||||
def _run_count(path, text):
|
||||
uri = "file:" + path.replace(os.sep, "/") + "?mode=ro"
|
||||
with contextlib.closing(sqlite3.connect(uri, uri=True)) as conn:
|
||||
row = conn.execute("SELECT run_count FROM prompts WHERE text=?", (text,))
|
||||
return row.fetchone()[0]
|
||||
|
||||
|
||||
class TestFailedSwitchIsUndone(_MoveCase):
|
||||
"""I1: whatever part of the switch fails, nothing was moved."""
|
||||
|
||||
def _assert_not_moved(self, caught):
|
||||
self.assertEqual(caught.exception.code, "switch-failed")
|
||||
self.assertIn("Nothing was moved", caught.exception.message)
|
||||
self.assertIsNone(ll.current_override())
|
||||
self.assertEqual(self.db.model.db_path, self.old_db)
|
||||
# Writes keep landing in the old library, and a restart keeps it.
|
||||
self.db.save_prompt("written after the failed switch")
|
||||
self.assertIn("written after the failed switch", repr(_rows(self.old_db)))
|
||||
self.assertEqual(self.location().source, ll.SOURCE_NODE_FOLDER)
|
||||
self.assertNotEqual(ll.read_marker(self.target).get("state"), "moved")
|
||||
self.assert_target_clean()
|
||||
self.assertTrue(os.path.isfile(self.config))
|
||||
assert_closed(self, self.target)
|
||||
|
||||
def _failing(self, point):
|
||||
real_switch_to = ll.switch_to
|
||||
real_repoint = models.PromptModel.repoint_instances
|
||||
|
||||
def switch(old_db, location, detached):
|
||||
if point == "before":
|
||||
raise RuntimeError("switch failed before anything")
|
||||
if point == "resolver":
|
||||
real_switch_to(location)
|
||||
raise RuntimeError("switch failed after the resolver")
|
||||
if point == "repoint":
|
||||
real_switch_to(location)
|
||||
detached.extend(real_repoint(old_db, location.db_path))
|
||||
raise RuntimeError("switch failed after the models")
|
||||
lm.switch_process(old_db, location, detached)
|
||||
raise RuntimeError("switch failed at its very end")
|
||||
|
||||
return switch
|
||||
|
||||
def test_a_switch_failing_at_any_point_is_undone(self):
|
||||
node = PromptDatabase(self.old_db)
|
||||
node.get_prompt_by_id(1)
|
||||
for point in ("before", "resolver", "repoint", "complete"):
|
||||
with self.subTest(point=point):
|
||||
with self.assertRaises(lm.MoveFailed) as caught:
|
||||
self.move(switch=self._failing(point))
|
||||
self._assert_not_moved(caught)
|
||||
self.assertEqual(node.model.db_path, self.old_db)
|
||||
self.assertEqual(node.get_prompt_by_id(1)["text"], "a cat number 0")
|
||||
|
||||
def test_a_failed_switch_closes_the_old_connections_it_detached(self):
|
||||
node = PromptDatabase(self.old_db)
|
||||
for point in ("repoint", "complete"):
|
||||
with self.subTest(point=point):
|
||||
before = node.model.get_connection()
|
||||
with self.assertRaises(lm.MoveFailed):
|
||||
self.move(switch=self._failing(point))
|
||||
with self.assertRaises(sqlite3.ProgrammingError):
|
||||
before.execute("SELECT 1")
|
||||
# The model reopens the old library on next use.
|
||||
self.assertEqual(node.get_prompt_by_id(1)["text"], "a cat number 0")
|
||||
|
||||
def test_a_failing_follow_up_after_the_switch_only_logs(self):
|
||||
def follow_up():
|
||||
raise RuntimeError("the health report could not be dropped")
|
||||
|
||||
with self.assertLogs("prompt_manager.database.library_move", "ERROR"):
|
||||
result = self.move(after_switch=follow_up)
|
||||
self.assertIsNone(result.retire_problem)
|
||||
self.assertEqual(self.db.model.db_path, self.new_db)
|
||||
self.assertTrue(
|
||||
os.path.isfile(
|
||||
os.path.join(self.ext, "prompts.db.moved-to-__promptmanager")
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class _Flaky:
|
||||
"""A write that fails once as if Move library switched while it waited."""
|
||||
|
||||
def __init__(self, db, name):
|
||||
self.db, self.name, self.calls = db, name, 0
|
||||
self.real = getattr(db, name)
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
self.calls += 1
|
||||
if self.calls == 1:
|
||||
self.db.model.close_all() # what the switch does: a new generation
|
||||
raise sqlite3.ProgrammingError("Cannot operate on a closed database.")
|
||||
return self.real(*args, **kwargs)
|
||||
|
||||
|
||||
class TestBackgroundWritesSurviveTheSwitch(_MoveCase):
|
||||
"""I2: run counting and the image monitor write again on a fresh connection."""
|
||||
|
||||
def test_run_counting_retries_once_after_a_switch(self):
|
||||
db = PromptDatabase(self.old_db)
|
||||
text = "a counted prompt"
|
||||
db.save_prompt(text, prompt_hash=generate_prompt_hash(text))
|
||||
flaky = _Flaky(db, "record_prompt_use")
|
||||
with mock.patch.object(db, "record_prompt_use", flaky):
|
||||
pending = usage_tracking.PendingFirstUse()
|
||||
usage_tracking.count_runs([(generate_prompt_hash(text), text)], db, pending)
|
||||
self.assertEqual(flaky.calls, 2)
|
||||
self.assertEqual(_run_count(self.old_db, text), 1)
|
||||
|
||||
def test_a_failure_without_a_switch_is_not_retried(self):
|
||||
db = PromptDatabase(self.old_db)
|
||||
calls = []
|
||||
|
||||
def broken():
|
||||
calls.append(1)
|
||||
raise sqlite3.OperationalError("disk I/O error")
|
||||
|
||||
with self.assertRaises(sqlite3.OperationalError):
|
||||
db.survive_library_switch(broken)
|
||||
self.assertEqual(calls, [1])
|
||||
|
||||
def test_an_image_link_racing_the_switch_lands_in_the_new_library(self):
|
||||
from utils.image_monitor import ImageGenerationHandler
|
||||
|
||||
self.db.model.close_all()
|
||||
image = os.path.join(self.comfy, "output", "late.png")
|
||||
# A short busy timeout so the blocked writer gives up in seconds.
|
||||
with mock.patch.object(models, "BUSY_TIMEOUT_SECONDS", 2):
|
||||
writer_db = PromptDatabase(self.old_db)
|
||||
handler = ImageGenerationHandler(writer_db, mock.MagicMock())
|
||||
started = threading.Event()
|
||||
outcome = {}
|
||||
|
||||
def write():
|
||||
conn = writer_db.model.get_connection()
|
||||
conn.set_trace_callback(
|
||||
lambda sql: (
|
||||
started.set()
|
||||
if sql.lstrip().upper().startswith(("BEGIN", "INSERT"))
|
||||
else None
|
||||
)
|
||||
)
|
||||
try:
|
||||
# No metadata: the fallback link, not the resolved pair.
|
||||
handler.link_image_to_prompt(image, {"id": 1}, {})
|
||||
outcome["ok"] = True
|
||||
except Exception as error: # pragma: no cover - reported below
|
||||
outcome["error"] = error
|
||||
finally:
|
||||
writer_db.model.close()
|
||||
|
||||
thread = threading.Thread(target=write)
|
||||
|
||||
def checkpoint(step):
|
||||
if step == "locked":
|
||||
# The old file is write-locked now: start the writer and
|
||||
# wait until its first write is running (it then waits on us).
|
||||
thread.start()
|
||||
self.assertTrue(started.wait(WAIT))
|
||||
|
||||
self.move(checkpoint=checkpoint)
|
||||
thread.join(WAIT)
|
||||
self.assertFalse(thread.is_alive())
|
||||
self.assertEqual(outcome, {"ok": True})
|
||||
moved = [
|
||||
row
|
||||
for row in _rows(self.new_db)["generated_images"]
|
||||
if "late.png" in repr(row)
|
||||
]
|
||||
self.assertEqual(len(moved), 1)
|
||||
renamed = os.path.join(self.ext, "prompts.db.moved-to-__promptmanager")
|
||||
self.assertNotIn("late.png", repr(_rows(renamed).get("generated_images")))
|
||||
|
||||
def test_an_image_that_was_not_linked_is_not_reported_as_linked(self):
|
||||
from utils.image_monitor import ImageGenerationHandler
|
||||
|
||||
db = PromptDatabase(self.old_db)
|
||||
handler = ImageGenerationHandler(db, mock.MagicMock())
|
||||
image = os.path.join(self.comfy, "output", "orphan.png")
|
||||
with mock.patch.object(handler, "logger") as log:
|
||||
handler.link_image_to_prompt(image, {"id": 99999}, {})
|
||||
messages = [str(call) for call in log.debug.call_args_list]
|
||||
self.assertFalse([m for m in messages if "Successfully linked" in m])
|
||||
self.assertNotIn("orphan.png", repr(_rows(self.old_db)["generated_images"]))
|
||||
|
||||
def test_a_write_racing_the_switch_lands_in_the_new_library(self):
|
||||
text = "a racing prompt"
|
||||
self.db.save_prompt(text, prompt_hash=generate_prompt_hash(text))
|
||||
self.db.model.close_all()
|
||||
# A short busy timeout so the blocked writer gives up in seconds.
|
||||
with mock.patch.object(models, "BUSY_TIMEOUT_SECONDS", 2):
|
||||
writer_db = PromptDatabase(self.old_db)
|
||||
started = threading.Event()
|
||||
outcome = {}
|
||||
|
||||
def write():
|
||||
conn = writer_db.model.get_connection()
|
||||
conn.set_trace_callback(
|
||||
lambda sql: started.set() if sql.startswith("UPDATE") else None
|
||||
)
|
||||
try:
|
||||
usage_tracking.count_runs(
|
||||
[(generate_prompt_hash(text), text)],
|
||||
writer_db,
|
||||
usage_tracking.PendingFirstUse(),
|
||||
)
|
||||
outcome["ok"] = True
|
||||
except Exception as error: # pragma: no cover - reported below
|
||||
outcome["error"] = error
|
||||
finally:
|
||||
writer_db.model.close()
|
||||
|
||||
thread = threading.Thread(target=write)
|
||||
|
||||
def checkpoint(step):
|
||||
if step == "locked":
|
||||
# The old file is write-locked now: start the writer and
|
||||
# wait until its UPDATE is running (it then waits on us).
|
||||
thread.start()
|
||||
self.assertTrue(started.wait(WAIT))
|
||||
|
||||
self.move(checkpoint=checkpoint)
|
||||
thread.join(WAIT)
|
||||
self.assertFalse(thread.is_alive())
|
||||
self.assertEqual(outcome, {"ok": True})
|
||||
self.assertEqual(_run_count(self.new_db, text), 1)
|
||||
renamed = os.path.join(self.ext, "prompts.db.moved-to-__promptmanager")
|
||||
self.assertEqual(_run_count(renamed, text), 0)
|
||||
|
||||
|
||||
class TestVerifyUri(unittest.TestCase):
|
||||
"""M2: a copy whose path holds ``#`` and ``%`` verifies."""
|
||||
|
||||
def test_copy_path_with_uri_characters_verifies(self):
|
||||
import tempfile
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
folder = os.path.join(tmp, "lib #1 %20")
|
||||
os.makedirs(folder)
|
||||
path = os.path.join(folder, "copy.db")
|
||||
db = PromptDatabase(path)
|
||||
db.save_prompt("a prompt")
|
||||
db.model.close_all()
|
||||
with contextlib.closing(sqlite3.connect(path)) as conn:
|
||||
conn.execute("PRAGMA journal_mode=DELETE")
|
||||
counts, documents = lm._describe(conn)
|
||||
lm._verify_copy(path, counts, documents)
|
||||
models.PromptModel.close_all_instances()
|
||||
|
||||
|
||||
class TestRestartFinishesTheRename(_MoveCase):
|
||||
"""M6: the marker names the old files; a restart renames them."""
|
||||
|
||||
def _crash_after_switch(self):
|
||||
def checkpoint(step):
|
||||
if step == "switched":
|
||||
raise SimulatedCrash(step)
|
||||
|
||||
with self.assertRaises(SimulatedCrash):
|
||||
self.move(checkpoint=checkpoint)
|
||||
models.PromptModel.close_all_instances()
|
||||
ll.clear_override()
|
||||
|
||||
def test_a_restart_renames_the_old_files_never_deletes(self):
|
||||
self._crash_after_switch()
|
||||
self.assertTrue(os.path.isfile(self.old_db))
|
||||
restarted = self.location()
|
||||
self.assertEqual(restarted.db_path, self.new_db)
|
||||
marker = ll.read_marker(self.target)
|
||||
self.assertEqual(
|
||||
(marker["from"], marker["config_from"]), (self.old_db, self.config)
|
||||
)
|
||||
renamed = lm.finish_interrupted_move(restarted)
|
||||
self.assertEqual(
|
||||
sorted(renamed),
|
||||
[
|
||||
"config.json.moved-to-__promptmanager",
|
||||
"prompts.db.moved-to-__promptmanager",
|
||||
],
|
||||
)
|
||||
self.assertFalse(os.path.exists(self.old_db))
|
||||
moved_config = os.path.join(self.ext, "config.json.moved-to-__promptmanager")
|
||||
with open(moved_config, encoding="utf-8") as handle:
|
||||
self.assertEqual(
|
||||
json.load(handle)["integrations"]["lora_manager"]["civitai_api_key"],
|
||||
KEY,
|
||||
)
|
||||
self.assertEqual(lm.finish_interrupted_move(self.location()), [])
|
||||
|
||||
def test_nothing_is_renamed_while_the_old_file_is_in_use(self):
|
||||
self._crash_after_switch()
|
||||
reader = sqlite3.connect(self.old_db)
|
||||
try:
|
||||
reader.execute("BEGIN")
|
||||
reader.execute("SELECT COUNT(*) FROM prompts").fetchone()
|
||||
PromptDatabase(self.old_db).save_prompt("in the wal")
|
||||
self.assertEqual(lm.finish_interrupted_move(self.location()), [])
|
||||
self.assertTrue(os.path.isfile(self.old_db))
|
||||
finally:
|
||||
reader.rollback()
|
||||
reader.close()
|
||||
models.PromptModel.close_all_instances()
|
||||
|
||||
def test_a_node_folder_library_has_nothing_to_finish(self):
|
||||
self.assertEqual(lm.finish_interrupted_move(self.location()), [])
|
||||
self.assertTrue(os.path.isfile(self.old_db))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -150,7 +150,15 @@ class TestSourceStatus(SourcesTestCase):
|
||||
)
|
||||
self.assertEqual(
|
||||
data["sources"]["effective"],
|
||||
[{"label": "a", "path": "output/a", "root_index": 0, "watch": "manual"}],
|
||||
[
|
||||
{
|
||||
"label": "a",
|
||||
"path": "output/a",
|
||||
"root_index": 0,
|
||||
"watch": "manual",
|
||||
"in_comfyui": True,
|
||||
}
|
||||
],
|
||||
)
|
||||
self.assertFalse(data["sources"]["auto_detected"])
|
||||
# The legacy keys are unchanged.
|
||||
@@ -183,10 +191,21 @@ class TestSourceStatus(SourcesTestCase):
|
||||
"path": "output",
|
||||
"root_index": 0,
|
||||
"watch": "manual",
|
||||
"in_comfyui": True,
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
def test_a_source_outside_comfyui_is_not_marked_in_comfyui(self):
|
||||
# Sources & scans prefixes "ComfyUI/" only to a path relative to the
|
||||
# ComfyUI base; a folder outside it is shown by its own name.
|
||||
from py.api import sources
|
||||
|
||||
self.assertTrue(sources.in_comfyui_base(str(self.a)))
|
||||
self.assertTrue(sources.in_comfyui_base(str(self.comfy_dir)))
|
||||
self.assertFalse(sources.in_comfyui_base(self.tmpdir))
|
||||
self.assertFalse(sources.in_comfyui_base(os.path.join(self.tmpdir, "Renders")))
|
||||
|
||||
async def test_watching_requires_live_observer_with_root(self):
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.a), self._root(self.b)]
|
||||
|
||||
|
||||
@@ -277,6 +277,7 @@ class TestOverviewRoute(SystemAPITestCase):
|
||||
"backup",
|
||||
"needs_review",
|
||||
"sources",
|
||||
"location",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
Vendored
+3
-1
@@ -25,14 +25,16 @@ test('execution script discovers only the new matrix and preserves all 78 origin
|
||||
assert.match(stock, /Total: 82 tests in 8 files/); // + local Inspect host case (inspect-sender.spec.mjs, 2 projects)
|
||||
assert.doesNotMatch(stock, /execution-history/);
|
||||
assert.doesNotMatch(stock, /execution-outputs/);
|
||||
assert.doesNotMatch(stock, /library-location/); // it starts its own fresh host (#249 S6b)
|
||||
const executions = list('tools/vendor/playwright.execution.config.mjs');
|
||||
assert.match(executions, /\[classic-start\]/);
|
||||
assert.match(executions, /\[nodes2-start\]/);
|
||||
assert.match(executions, /Total: 46 tests in 2 files/);
|
||||
assert.match(executions, /Total: 48 tests in 3 files/); // + the fresh-host library location case
|
||||
for (const project of ['classic-start', 'nodes2-start']) {
|
||||
const cases = executions.split('\n').filter(line => line.includes(`[${project}]`));
|
||||
assert.equal(cases.filter(line => line.includes('execution-history.spec.mjs')).length, 15);
|
||||
assert.equal(cases.filter(line => line.includes('execution-outputs.spec.mjs')).length, 8);
|
||||
assert.equal(cases.filter(line => line.includes('library-location.spec.mjs')).length, 1);
|
||||
}
|
||||
assert.doesNotMatch(executions, /prompt-manager.spec/);
|
||||
});
|
||||
|
||||
Vendored
+1
-1
@@ -19,7 +19,7 @@ const servers = [{
|
||||
|
||||
export default defineConfig({
|
||||
testDir: '../../tests/comfyui',
|
||||
testIgnore: ['execution-history.spec.mjs', 'execution-outputs.spec.mjs'],
|
||||
testIgnore: ['execution-history.spec.mjs', 'execution-outputs.spec.mjs', 'library-location.spec.mjs'],
|
||||
outputDir: './test-results/comfyui',
|
||||
timeout: 60000,
|
||||
workers: 1,
|
||||
|
||||
@@ -335,6 +335,12 @@
|
||||
"functions": 98,
|
||||
"statements": 98
|
||||
},
|
||||
"workspace/library-location.mjs": {
|
||||
"lines": 94,
|
||||
"branches": 91,
|
||||
"functions": 92,
|
||||
"statements": 94
|
||||
},
|
||||
"workspace/library-state.mjs": {
|
||||
"lines": 98,
|
||||
"branches": 95,
|
||||
|
||||
+1
-1
@@ -6,7 +6,7 @@ if (host.href !== 'http://127.0.0.1:8196/') throw new Error('Execution tests req
|
||||
|
||||
export default defineConfig({
|
||||
testDir: '../../tests/comfyui',
|
||||
testMatch: ['execution-history.spec.mjs', 'execution-outputs.spec.mjs'],
|
||||
testMatch: ['execution-history.spec.mjs', 'execution-outputs.spec.mjs', 'library-location.spec.mjs'],
|
||||
outputDir: './test-results/executions/cases',
|
||||
timeout: 60000,
|
||||
workers: 1,
|
||||
|
||||
+7
-13
@@ -69,8 +69,8 @@ class GalleryDiagnostics:
|
||||
|
||||
Args:
|
||||
db_path: Path to the SQLite database file to diagnose.
|
||||
If None, resolves from config.json or defaults to
|
||||
prompts.db in the extension root.
|
||||
If None, the library's database (see
|
||||
``utils.library_location``).
|
||||
"""
|
||||
if db_path is None:
|
||||
db_path = self._resolve_db_path()
|
||||
@@ -79,18 +79,12 @@ class GalleryDiagnostics:
|
||||
|
||||
@staticmethod
|
||||
def _resolve_db_path() -> str:
|
||||
"""Resolve database path from config, matching PromptDatabase logic."""
|
||||
extension_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
"""The library's database: the same resolver as PromptDatabase."""
|
||||
try:
|
||||
from py.config import PromptManagerConfig
|
||||
|
||||
db_path = PromptManagerConfig.DEFAULT_DB_PATH
|
||||
except Exception:
|
||||
db_path = "prompts.db"
|
||||
|
||||
if not os.path.isabs(db_path):
|
||||
db_path = os.path.normpath(os.path.join(extension_root, db_path))
|
||||
return db_path
|
||||
from ..database.operations import resolve_library_location
|
||||
except ImportError:
|
||||
from database.operations import resolve_library_location
|
||||
return resolve_library_location().db_path
|
||||
|
||||
def run_full_diagnostic(self) -> Dict[str, Any]:
|
||||
"""Run a complete diagnostic check.
|
||||
|
||||
+32
-17
@@ -36,6 +36,7 @@ from .metadata_extractor import ComfyUIMetadataExtractor
|
||||
from .logging_config import get_logger
|
||||
from .prompt_graph import resolve_text, run_prompt_nodes
|
||||
from .prompt_pair_resolver import resolve_prompt_pair
|
||||
from .usage_tracking import survive_switch
|
||||
|
||||
try:
|
||||
from ..database.pair_ingestion import ingest_resolved_pair
|
||||
@@ -587,24 +588,38 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
"""
|
||||
try:
|
||||
resolution = resolve_prompt_pair(metadata)
|
||||
if resolution["state"] == "resolved":
|
||||
positive = resolution["positive"]
|
||||
expected = self.db_manager.get_prompt_by_hash(
|
||||
generate_prompt_hash(positive.get("storage_text", positive["text"]))
|
||||
)
|
||||
if expected and expected["id"] == prompt_context["id"]:
|
||||
ingest_resolved_pair(
|
||||
self.db_manager, image_path, metadata, resolution
|
||||
|
||||
def link():
|
||||
# One write; run again on the moved library when a Move
|
||||
# library switched while it waited (#249 S6b).
|
||||
if resolution["state"] == "resolved":
|
||||
positive = resolution["positive"]
|
||||
expected = self.db_manager.get_prompt_by_hash(
|
||||
generate_prompt_hash(
|
||||
positive.get("storage_text", positive["text"])
|
||||
)
|
||||
)
|
||||
return
|
||||
image_id = self.db_manager.link_image_to_prompt(
|
||||
prompt_id=prompt_context["id"], image_path=image_path, metadata=metadata
|
||||
)
|
||||
fallback_note = " (fallback)" if prompt_context.get("fallback") else ""
|
||||
self.logger.debug(
|
||||
f"Successfully linked image {image_id} to prompt "
|
||||
f"{prompt_context['id']}{fallback_note}"
|
||||
)
|
||||
if expected and expected["id"] == prompt_context["id"]:
|
||||
ingest_resolved_pair(
|
||||
self.db_manager, image_path, metadata, resolution
|
||||
)
|
||||
return None
|
||||
# The raising form: a failed write reaches survive_switch.
|
||||
return self.db_manager.link_image_to_prompt_or_raise(
|
||||
prompt_id=prompt_context["id"],
|
||||
image_path=image_path,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
image_id = survive_switch(self.db_manager, link)
|
||||
if image_id is None:
|
||||
return
|
||||
if image_id: # 0: nothing new was linked (logged by the database)
|
||||
fallback_note = " (fallback)" if prompt_context.get("fallback") else ""
|
||||
self.logger.debug(
|
||||
f"Successfully linked image {image_id} to prompt "
|
||||
f"{prompt_context['id']}{fallback_note}"
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to link image to prompt: {e}")
|
||||
return
|
||||
|
||||
@@ -0,0 +1,480 @@
|
||||
"""Where the library lives (#249 S6b): the one resolver for every data path.
|
||||
|
||||
The library is ``prompts.db``, ``config.json`` (it can hold the CivitAI key),
|
||||
the safety backups beside the database, the logs and the LoRA image cache.
|
||||
Inside the node folder an uninstall, a reinstall or a folder-replacing update
|
||||
deletes it, so new installs keep it in ComfyUI's system user folder
|
||||
``user/__promptmanager`` (``folder_paths.get_system_user_directory``), which
|
||||
ComfyUI's ``/userdata`` HTTP API refuses to serve.
|
||||
|
||||
Resolution order (``resolve_library``):
|
||||
|
||||
1. ``PROMPT_MANAGER_DATA_DIR``: an absolute folder outside every folder
|
||||
ComfyUI or Prompt Manager serves. It is used when it holds a library or
|
||||
when no node-folder library exists; an empty one never replaces a
|
||||
node-folder library (that one stays in use, and Move library targets the
|
||||
environment folder). An invalid value is ignored and reported.
|
||||
2. An explicit ``database.default_path`` in config.json, honoured as before
|
||||
(a relative path is relative to the node folder). ``PROMPT_MANAGER_CONFIG_PATH``
|
||||
still names config.json itself.
|
||||
3. The default data folder once a move finished there (its marker says
|
||||
``moved``), even before the old file was renamed.
|
||||
4. A node-folder ``prompts.db`` that was not moved: it stays in use until the
|
||||
user moves it.
|
||||
5. The default data folder (new installs). Outside ComfyUI (no
|
||||
``folder_paths``) the node folder, as before.
|
||||
|
||||
Nothing here moves, copies or creates data: it only reads the file system.
|
||||
After a live move the process uses ``switch_to``'s location until restart.
|
||||
This module imports only the standard library, so the logger can use it.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
from typing import Iterable, List, Mapping, Optional, Tuple
|
||||
|
||||
DATA_DIR_ENV = "PROMPT_MANAGER_DATA_DIR"
|
||||
CONFIG_PATH_ENV = "PROMPT_MANAGER_CONFIG_PATH"
|
||||
EXTRA_GALLERY_ROOTS_ENV = "PROMPT_MANAGER_EXTRA_GALLERY_ROOTS"
|
||||
SYSTEM_USER_NAME = "promptmanager"
|
||||
SYSTEM_USER_PREFIX = "__"
|
||||
DATABASE_NAME = "prompts.db"
|
||||
CONFIG_NAME = "config.json"
|
||||
MARKER_NAME = "library-location.json"
|
||||
# config.json has always written its default; that value is not a choice.
|
||||
DEFAULT_DB_SETTING = "prompts.db"
|
||||
|
||||
SOURCE_ENVIRONMENT = "environment"
|
||||
SOURCE_CONFIG = "config"
|
||||
SOURCE_NODE_FOLDER = "node_folder"
|
||||
SOURCE_DEFAULT = "default"
|
||||
SOURCES = (SOURCE_ENVIRONMENT, SOURCE_CONFIG, SOURCE_NODE_FOLDER, SOURCE_DEFAULT)
|
||||
|
||||
MARKER_PENDING = "pending"
|
||||
MARKER_MOVED = "moved"
|
||||
|
||||
EXTENSION_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
_UNSET = object()
|
||||
_override_lock = threading.Lock()
|
||||
_override: Optional["LibraryLocation"] = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LibraryLocation:
|
||||
"""Where the library is and why.
|
||||
|
||||
``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
|
||||
db_path: str
|
||||
config_path: str
|
||||
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:
|
||||
return self.source == SOURCE_NODE_FOLDER and self.target_dir is not None
|
||||
|
||||
|
||||
def import_folder_paths():
|
||||
"""ComfyUI's ``folder_paths`` module, or None outside ComfyUI."""
|
||||
try:
|
||||
import folder_paths
|
||||
except ImportError:
|
||||
return None
|
||||
return folder_paths
|
||||
|
||||
|
||||
def _folder_paths(value):
|
||||
return import_folder_paths() if value is _UNSET else value
|
||||
|
||||
|
||||
def _call(module, name, *args) -> Optional[str]:
|
||||
getter = getattr(module, name, None)
|
||||
if not callable(getter):
|
||||
return None
|
||||
try:
|
||||
value = getter(*args)
|
||||
except Exception: # an older or broken host: treat the folder as unknown
|
||||
return None
|
||||
return str(value) if value else None
|
||||
|
||||
|
||||
def default_data_dir(folder_paths=_UNSET) -> Optional[str]:
|
||||
"""``user/__promptmanager``, or None outside ComfyUI.
|
||||
|
||||
ComfyUI builds without ``get_system_user_directory`` get the same folder
|
||||
joined by hand.
|
||||
"""
|
||||
module = _folder_paths(folder_paths)
|
||||
if module is None:
|
||||
return None
|
||||
folder = _call(module, "get_system_user_directory", SYSTEM_USER_NAME)
|
||||
if folder:
|
||||
return os.path.abspath(folder)
|
||||
user = _call(module, "get_user_directory")
|
||||
if not user:
|
||||
return None
|
||||
return os.path.abspath(os.path.join(user, SYSTEM_USER_PREFIX + SYSTEM_USER_NAME))
|
||||
|
||||
|
||||
def _identity(path: str) -> Optional[Tuple[int, int]]:
|
||||
try:
|
||||
stat = os.stat(path)
|
||||
except (OSError, ValueError):
|
||||
return None
|
||||
return (stat.st_dev, stat.st_ino)
|
||||
|
||||
|
||||
def same_folder(first: str, second: str) -> bool:
|
||||
"""Whether two spellings name one existing folder (links and case folded)."""
|
||||
one, two = _identity(first), _identity(second)
|
||||
if one is None or two is None:
|
||||
return False
|
||||
if one[1] == 0: # no file id on this volume: compare the spelling
|
||||
return _spelling(first) == _spelling(second)
|
||||
return one == two
|
||||
|
||||
|
||||
def _spelling(path: str) -> str:
|
||||
return os.path.normcase(os.path.realpath(path))
|
||||
|
||||
|
||||
def _existing_ancestors(path: str) -> List[str]:
|
||||
"""``path`` (resolved) and its parents that exist, nearest first."""
|
||||
current = os.path.realpath(path)
|
||||
found = []
|
||||
while True:
|
||||
if os.path.isdir(current):
|
||||
found.append(current)
|
||||
parent = os.path.dirname(current)
|
||||
if parent == current:
|
||||
return found
|
||||
current = parent
|
||||
|
||||
|
||||
def _inside_any(path: str, folders: Iterable[str]) -> bool:
|
||||
"""Whether ``path`` is one of ``folders`` or below one.
|
||||
|
||||
Compared by spelling (realpath + normcase) and by file identity of every
|
||||
existing ancestor, so symlinks and case-insensitive volumes (Windows,
|
||||
macOS) cannot hide a served folder.
|
||||
"""
|
||||
folders = [folder for folder in folders if folder]
|
||||
if not folders:
|
||||
return False
|
||||
spelling = _spelling(path)
|
||||
sep = os.path.sep
|
||||
for folder in folders:
|
||||
root = _spelling(folder)
|
||||
if spelling == root or spelling.startswith(root.rstrip(sep) + sep):
|
||||
return True
|
||||
identities = {_identity(folder) for folder in folders}
|
||||
identities.discard(None)
|
||||
identities = {item for item in identities if item[1] != 0}
|
||||
return any(_identity(a) in identities for a in _existing_ancestors(path))
|
||||
|
||||
|
||||
def _served_folders(module, environ, gallery_roots) -> List[str]:
|
||||
"""Folders whose files ComfyUI or Prompt Manager hand out over HTTP."""
|
||||
folders = []
|
||||
if module is not None:
|
||||
for name in (
|
||||
"get_output_directory",
|
||||
"get_input_directory",
|
||||
"get_temp_directory",
|
||||
):
|
||||
folder = _call(module, name)
|
||||
if folder:
|
||||
folders.append(folder)
|
||||
extra = (environ or {}).get(EXTRA_GALLERY_ROOTS_ENV, "")
|
||||
folders.extend(part for part in extra.split(os.pathsep) if part.strip())
|
||||
folders.extend(str(root) for root in gallery_roots or () if root)
|
||||
return folders
|
||||
|
||||
|
||||
def _in_public_user_folder(path: str, module) -> bool:
|
||||
"""Whether ``path`` lies in a ``/userdata``-served user folder (not ``__``)."""
|
||||
user = _call(module, "get_user_directory") if module is not None else None
|
||||
if not user:
|
||||
return False
|
||||
if not _inside_any(path, [user]) or same_folder(path, user):
|
||||
return False
|
||||
real_user = _spelling(user)
|
||||
relative = os.path.relpath(_spelling(path), real_user)
|
||||
first = relative.split(os.sep, 1)[0]
|
||||
if first.startswith(os.pardir):
|
||||
# Reached through a link or another spelling: judge by identity.
|
||||
for ancestor in _existing_ancestors(path):
|
||||
if same_folder(os.path.dirname(ancestor), user):
|
||||
first = os.path.basename(ancestor)
|
||||
break
|
||||
else:
|
||||
return False
|
||||
return not first.startswith(SYSTEM_USER_PREFIX)
|
||||
|
||||
|
||||
def data_dir_problem(
|
||||
path, *, folder_paths=_UNSET, environ=None, gallery_roots=()
|
||||
) -> Optional[str]:
|
||||
"""Why ``path`` cannot hold the library, as a code, or None.
|
||||
|
||||
``data-dir-not-absolute``: empty or relative. ``data-dir-served``: inside
|
||||
ComfyUI's output, input or temp folder, a public user folder, an extra
|
||||
gallery root or a gallery source (their files are served over HTTP).
|
||||
"""
|
||||
if not isinstance(path, str) or not path.strip() or not os.path.isabs(path):
|
||||
return "data-dir-not-absolute"
|
||||
module = _folder_paths(folder_paths)
|
||||
if _inside_any(path, _served_folders(module, environ, gallery_roots)):
|
||||
return "data-dir-served"
|
||||
if _in_public_user_folder(path, module):
|
||||
return "data-dir-served"
|
||||
return None
|
||||
|
||||
|
||||
def read_marker(directory: Optional[str]) -> dict:
|
||||
"""The move marker in ``directory``, or ``{}`` when absent or unreadable."""
|
||||
if not directory:
|
||||
return {}
|
||||
try:
|
||||
with open(os.path.join(directory, MARKER_NAME), "r", encoding="utf-8") as fh:
|
||||
data = json.load(fh)
|
||||
except (OSError, ValueError):
|
||||
return {}
|
||||
if not isinstance(data, dict) or data.get("state") not in (
|
||||
MARKER_PENDING,
|
||||
MARKER_MOVED,
|
||||
):
|
||||
return {}
|
||||
return data
|
||||
|
||||
|
||||
def _marker_state(directory) -> Optional[str]:
|
||||
return read_marker(directory).get("state")
|
||||
|
||||
|
||||
def switch_to(location: LibraryLocation) -> None:
|
||||
"""Use ``location`` for the rest of this process (after a live move)."""
|
||||
global _override
|
||||
with _override_lock:
|
||||
_override = location
|
||||
|
||||
|
||||
def clear_override() -> None:
|
||||
"""Forget a live move's location (tests)."""
|
||||
global _override
|
||||
with _override_lock:
|
||||
_override = None
|
||||
|
||||
|
||||
def current_override() -> Optional[LibraryLocation]:
|
||||
with _override_lock:
|
||||
return _override
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Decision:
|
||||
data_dir: str
|
||||
source: str
|
||||
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 _configured_sources(module, environ, extension_root, default, raw_env):
|
||||
"""The image sources every config.json this resolver might read lists.
|
||||
|
||||
``config_location`` runs before any config.json is loaded, so it cannot be
|
||||
handed the configured sources; both it and ``resolve_library`` read them
|
||||
here, from each candidate file (the explicit one, the environment
|
||||
folder's, the node folder's and the default folder's), so the two always
|
||||
judge ``PROMPT_MANAGER_DATA_DIR`` alike. Relative sources are taken from
|
||||
ComfyUI's base folder, as the gallery resolves them.
|
||||
"""
|
||||
candidates = [environ.get(CONFIG_PATH_ENV, "").strip()]
|
||||
if raw_env and os.path.isabs(raw_env):
|
||||
candidates.append(os.path.join(raw_env, CONFIG_NAME))
|
||||
candidates.append(os.path.join(extension_root, CONFIG_NAME))
|
||||
if default:
|
||||
candidates.append(os.path.join(default, CONFIG_NAME))
|
||||
base = getattr(module, "base_path", None) if module is not None else None
|
||||
sources = []
|
||||
for candidate in candidates:
|
||||
if not candidate or not os.path.isfile(candidate):
|
||||
continue
|
||||
try:
|
||||
with open(candidate, "r", encoding="utf-8") as handle:
|
||||
data = json.load(handle)
|
||||
directories = data["gallery"]["monitoring"]["directories"]
|
||||
except (OSError, ValueError, KeyError, TypeError):
|
||||
continue
|
||||
if not isinstance(directories, list):
|
||||
continue
|
||||
for directory in directories:
|
||||
if not isinstance(directory, str) or not directory.strip():
|
||||
continue
|
||||
directory = directory.strip()
|
||||
if not os.path.isabs(directory):
|
||||
if not base:
|
||||
continue
|
||||
directory = os.path.join(str(base), directory)
|
||||
sources.append(directory)
|
||||
return sources
|
||||
|
||||
|
||||
def _decide(environ, extension_root, folder_paths, gallery_roots) -> _Decision:
|
||||
"""Which folder holds the library (before an explicit database path)."""
|
||||
module = _folder_paths(folder_paths)
|
||||
environ = os.environ if environ is None else environ
|
||||
legacy_db = os.path.isfile(os.path.join(extension_root, DATABASE_NAME))
|
||||
legacy = legacy_db or os.path.isfile(os.path.join(extension_root, CONFIG_NAME))
|
||||
default = default_data_dir(module)
|
||||
problem = None
|
||||
env_dir = None
|
||||
raw = environ.get(DATA_DIR_ENV, "")
|
||||
if raw.strip():
|
||||
sources = list(gallery_roots or ()) + _configured_sources(
|
||||
module, environ, extension_root, default, raw.strip()
|
||||
)
|
||||
problem = data_dir_problem(
|
||||
raw.strip(),
|
||||
folder_paths=module,
|
||||
environ=environ,
|
||||
gallery_roots=sources,
|
||||
)
|
||||
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)
|
||||
problem = "move-unfinished" if state == MARKER_PENDING else "data-dir-empty"
|
||||
return decision(extension_root, SOURCE_NODE_FOLDER, problem)
|
||||
if default is None:
|
||||
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, problem)
|
||||
if legacy:
|
||||
if _marker_state(default) == MARKER_PENDING:
|
||||
problem = problem or "move-unfinished"
|
||||
if legacy_db:
|
||||
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, problem, default)
|
||||
return decision(default, SOURCE_DEFAULT, problem)
|
||||
|
||||
|
||||
def config_location(
|
||||
environ=None, extension_root=None, folder_paths=_UNSET, gallery_roots=()
|
||||
) -> Tuple[str, str]:
|
||||
"""``(config.json path, data folder)``; needs no config value (no cycle)."""
|
||||
override = current_override()
|
||||
if override is not None:
|
||||
return override.config_path, override.data_dir
|
||||
environ = os.environ if environ is None else environ
|
||||
root = extension_root or EXTENSION_ROOT
|
||||
decision = _decide(environ, root, folder_paths, gallery_roots)
|
||||
explicit = environ.get(CONFIG_PATH_ENV, "").strip()
|
||||
config_path = explicit or os.path.join(decision.data_dir, CONFIG_NAME)
|
||||
return config_path, decision.data_dir
|
||||
|
||||
|
||||
def is_explicit_db_setting(value) -> bool:
|
||||
"""Whether config.json's ``database.default_path`` names a chosen file."""
|
||||
return isinstance(value, str) and value.strip() not in ("", DEFAULT_DB_SETTING)
|
||||
|
||||
|
||||
def resolve_library(
|
||||
default_path=None,
|
||||
*,
|
||||
environ: Optional[Mapping[str, str]] = None,
|
||||
extension_root: Optional[str] = None,
|
||||
folder_paths=_UNSET,
|
||||
gallery_roots=(),
|
||||
) -> LibraryLocation:
|
||||
"""The library's location now (the module docstring has the order)."""
|
||||
override = current_override()
|
||||
if override is not None:
|
||||
return override
|
||||
environ = os.environ if environ is None else environ
|
||||
root = extension_root or EXTENSION_ROOT
|
||||
decision = _decide(environ, root, folder_paths, gallery_roots)
|
||||
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):
|
||||
db_path = os.path.normpath(os.path.join(root, db_path))
|
||||
return LibraryLocation(
|
||||
SOURCE_CONFIG,
|
||||
db_path,
|
||||
config_path,
|
||||
decision.data_dir,
|
||||
decision.target_dir,
|
||||
decision.problem,
|
||||
**common,
|
||||
)
|
||||
return LibraryLocation(
|
||||
decision.source,
|
||||
os.path.join(decision.db_dir or decision.data_dir, DATABASE_NAME),
|
||||
config_path,
|
||||
decision.data_dir,
|
||||
decision.target_dir,
|
||||
decision.problem,
|
||||
**common,
|
||||
)
|
||||
|
||||
|
||||
def ensure_private_dir(path: str) -> str:
|
||||
"""Create ``path`` (and parents) if missing; a new folder is owner-only.
|
||||
|
||||
An existing folder is left as it is (its permissions are the user's).
|
||||
"""
|
||||
if not os.path.isdir(path):
|
||||
os.makedirs(path, mode=0o700, exist_ok=True)
|
||||
if os.name == "posix":
|
||||
try:
|
||||
os.chmod(path, 0o700)
|
||||
except OSError:
|
||||
pass
|
||||
return path
|
||||
|
||||
|
||||
def logs_dir(**kwargs) -> str:
|
||||
"""The folder for Prompt Manager's log files."""
|
||||
return os.path.join(config_location(**kwargs)[1], "logs")
|
||||
|
||||
|
||||
def cache_dir(name: str, **kwargs) -> str:
|
||||
"""A re-creatable cache folder, ``<data folder>/data/<name>``."""
|
||||
return os.path.join(config_location(**kwargs)[1], "data", name)
|
||||
@@ -18,6 +18,8 @@ from datetime import datetime
|
||||
from typing import Dict, Any, Optional, List
|
||||
from pathlib import Path
|
||||
|
||||
from . import library_location
|
||||
|
||||
|
||||
class _OwnedRotatingFileHandler(logging.handlers.RotatingFileHandler):
|
||||
"""Keep a retired owned sink closed when an in-flight record resumes."""
|
||||
@@ -65,8 +67,12 @@ class PromptManagerLogger:
|
||||
return
|
||||
|
||||
self._initialized = True
|
||||
self.log_dir = Path(__file__).parent.parent / "logs"
|
||||
self.log_dir.mkdir(exist_ok=True)
|
||||
# The library's data folder (#249 S6b): the node folder for a library
|
||||
# that was not moved, else ComfyUI's user/__promptmanager.
|
||||
self.log_dir = Path(
|
||||
library_location.logs_dir(extension_root=str(Path(__file__).parent.parent))
|
||||
)
|
||||
library_location.ensure_private_dir(str(self.log_dir))
|
||||
|
||||
# Configuration
|
||||
self.config = {
|
||||
|
||||
+29
-8
@@ -187,6 +187,19 @@ def find_pair_id(db: Any, positive: str, negative: str) -> Optional[int]:
|
||||
return row["id"] if row else None
|
||||
|
||||
|
||||
def survive_switch(db: Any, write: Callable[[], Any]) -> Any:
|
||||
"""``write()`` through ``db.survive_library_switch`` when the database has it.
|
||||
|
||||
A write that was waiting on the old file when Move library switched runs
|
||||
again on the new library instead of being lost (#249 S6b). Stand-in
|
||||
databases without the method (tests) just run ``write()``.
|
||||
"""
|
||||
survive = getattr(type(db), "survive_library_switch", None)
|
||||
if survive is None:
|
||||
return write()
|
||||
return survive(db, write)
|
||||
|
||||
|
||||
def count_pair_runs(
|
||||
runs: List[Tuple[str, str, str]],
|
||||
db: Any,
|
||||
@@ -194,20 +207,28 @@ def count_pair_runs(
|
||||
) -> None:
|
||||
"""Count existing recipes at dequeue and defer new recipes until first save."""
|
||||
for key, positive, negative in runs:
|
||||
pair_id = find_pair_id(db, positive, negative)
|
||||
if pair_id:
|
||||
db.record_prompt_pair_use(pair_id)
|
||||
else:
|
||||
|
||||
def count(positive=positive, negative=negative):
|
||||
pair_id = find_pair_id(db, positive, negative)
|
||||
if pair_id:
|
||||
db.record_prompt_pair_use(pair_id)
|
||||
return bool(pair_id)
|
||||
|
||||
if not survive_switch(db, count):
|
||||
pending.add(key)
|
||||
|
||||
|
||||
def count_runs(runs: List[Run], db: Any, pending: PendingFirstUse) -> None:
|
||||
"""Record one use per prompt in the database; remember the rest for first save."""
|
||||
for prompt_hash, _text in runs:
|
||||
existing = db.get_prompt_by_hash(prompt_hash)
|
||||
if existing:
|
||||
db.record_prompt_use(existing["id"])
|
||||
else:
|
||||
|
||||
def count(prompt_hash=prompt_hash):
|
||||
existing = db.get_prompt_by_hash(prompt_hash)
|
||||
if existing:
|
||||
db.record_prompt_use(existing["id"])
|
||||
return bool(existing)
|
||||
|
||||
if not survive_switch(db, count):
|
||||
pending.add(prompt_hash)
|
||||
|
||||
|
||||
|
||||
@@ -79,6 +79,8 @@
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
<!-- #249 S6b: where the library (database, config.json, backups) lives and why. -->
|
||||
<p class="text-pm-secondary text-xs mt-2" id="libraryLocation" hidden></p>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
|
||||
@@ -297,6 +297,7 @@
|
||||
this.loadTags(),
|
||||
this.loadRecentPrompts(),
|
||||
this.loadSettings(),
|
||||
this.loadLibraryLocation(),
|
||||
]);
|
||||
} catch (error) {
|
||||
this.showNotification("Failed to load initial data", "error");
|
||||
@@ -344,6 +345,29 @@
|
||||
this.listRequest = ApiClient.latestOnly({ lifetime: this.pageLifetime.signal });
|
||||
}
|
||||
|
||||
/** #249 S6b: one line under the statistics saying where the library lives and why. */
|
||||
async loadLibraryLocation() {
|
||||
const { signal } = this.pageLifetime;
|
||||
if (signal.aborted) return;
|
||||
const line = document.getElementById("libraryLocation");
|
||||
const reasons = {
|
||||
default: "ComfyUI's user folder",
|
||||
environment: "set by PROMPT_MANAGER_DATA_DIR",
|
||||
config: "set in config.json",
|
||||
node_folder: "the node folder, where an update or reinstall can delete it; move it from the workspace's System page",
|
||||
};
|
||||
try {
|
||||
const response = await ApiClient.fetchUrl("/prompt_manager/system/library", { signal });
|
||||
const data = response.ok ? await response.json() : null;
|
||||
const location = data && data.location;
|
||||
if (!line || !location || typeof location.path !== "string") return;
|
||||
line.textContent = `Library: ${location.path} (${reasons[location.source] || location.source})`;
|
||||
line.hidden = false;
|
||||
} catch (error) {
|
||||
if (!signal.aborted) console.error("Library location error:", error);
|
||||
}
|
||||
}
|
||||
|
||||
async loadStatistics() {
|
||||
// A leaving page starts no read; one in flight is aborted with the page and not reported.
|
||||
const { signal } = this.pageLifetime;
|
||||
|
||||
@@ -18,6 +18,7 @@ import { mountInspect } from './inspect.mjs';
|
||||
import { mountReview } from './review.mjs';
|
||||
import { createPromptSender } from './prompt-sender.mjs';
|
||||
import { mountSystem } from './system.mjs';
|
||||
import { mountLibraryBanner } from './library-location.mjs';
|
||||
|
||||
const { createApiClient, resolveUrl } = window.ApiClient;
|
||||
|
||||
@@ -75,6 +76,8 @@ function mountWorkspace() {
|
||||
const activity = createActivityStore({ api, scope, connect: url => new EventSource(resolveUrl(url)) });
|
||||
mountActivityTopbar({ store: activity, scope });
|
||||
activity.load();
|
||||
// D129: a node-folder library is offered a move (or "Keep it here") above every page.
|
||||
mountLibraryBanner(document.getElementById('library-banner'), { api, activity, scope });
|
||||
const sidebar = document.querySelector('aside.sidebar');
|
||||
// The route key a nav link carries in `data-route`: a saved view's is `saved/<id>`.
|
||||
const routeKey = route => route.id === 'saved' ? `saved/${route.savedId}` : route.id;
|
||||
|
||||
@@ -0,0 +1,255 @@
|
||||
/**
|
||||
* Where the library lives (#249 System S6b, D129): the workspace banner, System › Overview's "Library location" row,
|
||||
* the Move library review and its start.
|
||||
*
|
||||
* - The server says where the database is and why (`GET /prompt_manager/system/library` and the Overview's
|
||||
* `location`): `default` (ComfyUI's user folder), `environment` (PROMPT_MANAGER_DATA_DIR), `config`
|
||||
* (config.json's database path) or `node_folder`, where an update or reinstall can delete it.
|
||||
* - D129: while a node-folder library is not kept there, the shell shows the concept's `.notice` (`h2` + `p` +
|
||||
* `.row-actions`) above every page: "Keep it here" saves the choice on the server (the banner stops; the Overview row
|
||||
* stays), "Move library" opens the review. The concept has no such notice; its markup and classes are the Sources
|
||||
* notice's (app.js 680).
|
||||
* - The review is the shared record-changing review (`openReviewDialog`): From, To, Size, what moves with it and what
|
||||
* happens to the old files. Move library starts the server's exclusive "Move library" job, tracks it in Activity and
|
||||
* opens its view; a refusal (another job running, no room, a library already there) says why and nothing starts.
|
||||
* - When a Move library job ends, the banner and the Overview read the location again.
|
||||
*/
|
||||
import { createButton } from '../shared/toolkit/core.js';
|
||||
import { busyToast } from './activity-scan.mjs';
|
||||
import { closeWorkspaceDialog, openContentDialog, openReviewDialog } from './workspace-dialogs.mjs';
|
||||
import { navigateTo } from './state.mjs';
|
||||
import { showWorkspaceToast } from './toast.mjs';
|
||||
|
||||
export const LIBRARY_URL = '/prompt_manager/system/library';
|
||||
export const MOVE_URL = `${LIBRARY_URL}/move`;
|
||||
export const KEEP_URL = `${LIBRARY_URL}/keep`;
|
||||
export const MOVE_KIND = 'library_move';
|
||||
|
||||
export const BANNER_TITLE = 'Your library is inside the node folder';
|
||||
const BANNER_LEAD = 'Your library is stored inside the node folder, where an update or reinstall can delete it.';
|
||||
/** Where Move library puts it, by the server's `target_source`. */
|
||||
export const TARGET_TEXT = Object.freeze({
|
||||
default: 'ComfyUI\'s user folder',
|
||||
environment: 'the folder PROMPT_MANAGER_DATA_DIR names',
|
||||
});
|
||||
/** The notice's sentence: it names the actual target. */
|
||||
export function bannerText(location) {
|
||||
const target = TARGET_TEXT[location?.target_source] ?? TARGET_TEXT.default;
|
||||
return `${BANNER_LEAD} Move it to ${target}: it is copied and verified first, and nothing is deleted.`;
|
||||
}
|
||||
export const BANNER_TEXT = bannerText({ target_source: 'default' });
|
||||
export const KEEP_LABEL = 'Keep it here';
|
||||
export const MOVE_LABEL = 'Move library';
|
||||
export const ROW_TITLE = 'Library location';
|
||||
export const START_FAILED = 'Move library could not start. Try again.';
|
||||
export const KEEP_FAILED = 'Your choice could not be saved. Try again.';
|
||||
|
||||
/** Why the library is where it is, by `source`. */
|
||||
export const SOURCE_TEXT = Object.freeze({
|
||||
default: 'ComfyUI\'s user folder, which updates and reinstalls of the node leave alone',
|
||||
environment: 'The folder PROMPT_MANAGER_DATA_DIR names',
|
||||
config: 'The database path set in config.json',
|
||||
node_folder: 'The node folder, where an update or reinstall can delete it',
|
||||
});
|
||||
|
||||
/** A `problem` code as a sentence. */
|
||||
export const PROBLEM_TEXT = Object.freeze({
|
||||
'data-dir-not-absolute': 'PROMPT_MANAGER_DATA_DIR is not an absolute folder, so it was ignored.',
|
||||
'data-dir-served': 'PROMPT_MANAGER_DATA_DIR names a folder whose files are served over the web, so it was ignored.',
|
||||
'data-dir-empty': 'PROMPT_MANAGER_DATA_DIR has no library yet; Move library puts this one there.',
|
||||
'move-unfinished': 'An earlier move did not finish; this library is still the one in use.',
|
||||
'inside-node-folder': 'It is inside the node folder, where an update or reinstall deletes it, config.json and its '
|
||||
+ 'CivitAI key included. Set database.default_path or PROMPT_MANAGER_DATA_DIR to a folder outside it.',
|
||||
});
|
||||
|
||||
const number = value => Number(value).toLocaleString('en-US');
|
||||
|
||||
/** "1.6 GB" or "120.0 MB" or "512 KB". */
|
||||
export function sizeText(bytes) {
|
||||
if (!Number.isFinite(bytes) || bytes < 0) return 'Unknown';
|
||||
if (bytes >= 2 ** 30) return `${(bytes / 2 ** 30).toLocaleString('en-US', { maximumFractionDigits: 1, minimumFractionDigits: 1 })} GB`;
|
||||
if (bytes >= 2 ** 20) return `${(bytes / 2 ** 20).toLocaleString('en-US', { maximumFractionDigits: 1, minimumFractionDigits: 1 })} MB`;
|
||||
return `${number(Math.ceil(bytes / 1024))} KB`;
|
||||
}
|
||||
|
||||
/** Whether `location` is a well-formed server answer. */
|
||||
export function isLocation(location) {
|
||||
return Boolean(location) && typeof location === 'object' && typeof location.path === 'string'
|
||||
&& Object.hasOwn(SOURCE_TEXT, location.source) && typeof location.movable === 'boolean';
|
||||
}
|
||||
|
||||
/** The Overview row's description: the path, why it is there, the user's choice and any problem. */
|
||||
export function locationDescription(location) {
|
||||
const parts = [`${location.path} · ${SOURCE_TEXT[location.source]}.`];
|
||||
if (location.movable && location.keep) parts.push('You chose to keep it here.');
|
||||
const problem = PROBLEM_TEXT[location.problem];
|
||||
if (problem) parts.push(problem);
|
||||
return parts.join(' ');
|
||||
}
|
||||
|
||||
/** The review's `[term, value]` rows. */
|
||||
export function reviewGrid(location) {
|
||||
return [
|
||||
['From', location.path],
|
||||
['To', `${location.target}/prompts.db`],
|
||||
['Size', sizeText(location.size_bytes)],
|
||||
['Moves with it', location.config_moves === false
|
||||
? 'The database backups (config.json stays where PROMPT_MANAGER_CONFIG_PATH names it)'
|
||||
: 'config.json and the database backups'],
|
||||
['Old files', 'Renamed *.moved-to-…, never deleted'],
|
||||
];
|
||||
}
|
||||
|
||||
export const REVIEW_SENTENCE = 'The copy is verified before Prompt Manager switches to it. No other job can run until it '
|
||||
+ 'finishes, and closing this page does not stop it.';
|
||||
|
||||
function el(tag, text, className) {
|
||||
const node = document.createElement(tag);
|
||||
if (text !== undefined) node.textContent = text;
|
||||
if (className) node.className = className;
|
||||
return node;
|
||||
}
|
||||
|
||||
/**
|
||||
* Start the Move library job, follow it in Activity and open its view. A busy server toasts the D57 link to the running
|
||||
* job; any other refusal says the server's reason. Nothing is said once `scope` ended.
|
||||
*/
|
||||
export async function startLibraryMove({ api, activity, scope, notify = showWorkspaceToast, navigate = navigateTo }) {
|
||||
try {
|
||||
const { job } = await api.post(MOVE_URL, {}, { signal: scope.signal });
|
||||
if (scope.signal.aborted) return null;
|
||||
const tracked = activity.track(job);
|
||||
navigate(`#/activity/${tracked.id}`);
|
||||
return tracked;
|
||||
} catch (error) {
|
||||
if (scope.signal.aborted) return null;
|
||||
if (error?.status === 409 && error?.code === 'busy') {
|
||||
await activity.load();
|
||||
if (!scope.signal.aborted) notify(busyToast(activity.running()[0] ?? null));
|
||||
return null;
|
||||
}
|
||||
notify(typeof error?.code === 'string' && error.message ? error.message : START_FAILED);
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
/** The Move library review over `location`; Move library starts the job. */
|
||||
export function openMoveReview({ api, activity, scope, location, trigger = null, notify = showWorkspaceToast,
|
||||
navigate = navigateTo }) {
|
||||
return openReviewDialog({ scope, trigger, title: 'Move library', grid: reviewGrid(location), sentence: REVIEW_SENTENCE,
|
||||
cancelLabel: 'Cancel', primary: { label: MOVE_LABEL },
|
||||
onStart: () => { startLibraryMove({ api, activity, scope, notify, navigate }); } });
|
||||
}
|
||||
|
||||
/** Save "Keep it here" (`keep` true) or take it back; resolves the new location, or null when it was not saved. */
|
||||
export async function keepLibrary({ api, scope, keep = true, notify = showWorkspaceToast }) {
|
||||
try {
|
||||
const { location } = await api.post(KEEP_URL, { keep }, { signal: scope.signal });
|
||||
return isLocation(location) ? location : null;
|
||||
} catch (error) {
|
||||
if (!scope.signal.aborted) notify(KEEP_FAILED);
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
/** The details of a library that does not move: where it is, why, and what lives beside it. */
|
||||
export function openLocationDetails({ scope, location, trigger = null }) {
|
||||
const list = el('dl', undefined, 'meta-grid');
|
||||
for (const [term, value] of [['Database', location.path], ['Why here', SOURCE_TEXT[location.source]],
|
||||
['Beside it', 'config.json, the database backups and the logs'], ['Size', sizeText(location.size_bytes)]]) {
|
||||
const entry = el('div');
|
||||
entry.append(el('dt', term), el('dd', value));
|
||||
list.append(entry);
|
||||
}
|
||||
const close = createButton('Close');
|
||||
scope.listen(close, 'click', () => closeWorkspaceDialog());
|
||||
const actions = el('div', undefined, 'dialog-actions');
|
||||
actions.append(close);
|
||||
const nodes = [list];
|
||||
const problem = PROBLEM_TEXT[location.problem];
|
||||
if (problem) nodes.push(el('p', problem, 'small'));
|
||||
nodes.push(actions);
|
||||
openContentDialog({ scope, title: ROW_TITLE, trigger, nodes });
|
||||
}
|
||||
|
||||
/** D129: the Overview's "Library location" `.operation` row; Move library when it can move, else Details. */
|
||||
export function locationRow(location, { api, activity, scope, notify = showWorkspaceToast, navigate = navigateTo }) {
|
||||
const row = el('div', undefined, 'operation');
|
||||
const text = el('div');
|
||||
text.append(el('h3', ROW_TITLE), el('p', locationDescription(location)));
|
||||
const actions = el('div', undefined, 'row-actions');
|
||||
const action = createButton(location.movable ? MOVE_LABEL : 'Details');
|
||||
scope.listen(action, 'click', () => {
|
||||
if (location.movable) openMoveReview({ api, activity, scope, location, trigger: action, notify, navigate });
|
||||
else openLocationDetails({ scope, location, trigger: action });
|
||||
});
|
||||
actions.append(action);
|
||||
row.append(text, actions);
|
||||
return row;
|
||||
}
|
||||
|
||||
/** D129: the concept's `.notice` for a node-folder library: Keep it here and Move library. */
|
||||
export function libraryNotice({ scope, location = null, onKeep, onMove }) {
|
||||
const notice = el('div', undefined, 'notice');
|
||||
const text = el('div');
|
||||
text.append(el('h2', BANNER_TITLE), el('p', bannerText(location)));
|
||||
const keep = createButton(KEEP_LABEL);
|
||||
const move = createButton(MOVE_LABEL, 'ca-primary');
|
||||
scope.listen(keep, 'click', () => onKeep(keep));
|
||||
scope.listen(move, 'click', () => onMove(move));
|
||||
const actions = el('div', undefined, 'row-actions');
|
||||
actions.append(keep, move);
|
||||
notice.append(text, actions);
|
||||
return notice;
|
||||
}
|
||||
|
||||
/** Calls `onEnded()` whenever a Move library job this page sees stops running. Returns the unsubscribe. */
|
||||
export function onMoveEnded(activity, onEnded) {
|
||||
const running = new Set();
|
||||
return activity.subscribe(() => {
|
||||
let ended = false;
|
||||
for (const job of activity.jobs()) {
|
||||
if (job.kind !== MOVE_KIND) continue;
|
||||
if (job.state === 'running') running.add(job.id);
|
||||
else if (running.delete(job.id)) ended = true;
|
||||
}
|
||||
if (ended) onEnded();
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Mount the banner in `container` (hidden until the server offers the move). It reads the location once, again when a
|
||||
* Move library job ends, and hides for good this page once the user keeps the library where it is.
|
||||
*/
|
||||
export function mountLibraryBanner(container, { api, activity, scope, notify = showWorkspaceToast, navigate = navigateTo,
|
||||
latestOnly = window.ApiClient.latestOnly, STALE = window.ApiClient.STALE }) {
|
||||
const read = latestOnly({ lifetime: scope.signal });
|
||||
const show = location => {
|
||||
container.replaceChildren();
|
||||
container.hidden = !(isLocation(location) && location.offer === true);
|
||||
if (container.hidden) return;
|
||||
container.append(libraryNotice({ scope, location,
|
||||
onKeep: async button => {
|
||||
button.setAttribute('aria-disabled', 'true');
|
||||
const saved = await keepLibrary({ api, scope, notify });
|
||||
if (scope.signal.aborted) return;
|
||||
button.removeAttribute('aria-disabled');
|
||||
if (saved) show(saved);
|
||||
},
|
||||
onMove: button => openMoveReview({ api, activity, scope, location, trigger: button, notify, navigate }) }));
|
||||
};
|
||||
async function load() {
|
||||
try {
|
||||
const answer = await read(signal => api.get(LIBRARY_URL, { signal, cache: 'no-store' }));
|
||||
if (answer === STALE || scope.signal.aborted) return;
|
||||
show(answer.location);
|
||||
} catch {
|
||||
// The banner is advisory: an unreadable location shows nothing (the Overview row reports it).
|
||||
if (!scope.signal.aborted) show(null);
|
||||
}
|
||||
}
|
||||
container.hidden = true;
|
||||
scope.cleanup(onMoveEnded(activity, load));
|
||||
load();
|
||||
return { reload: load };
|
||||
}
|
||||
@@ -22,6 +22,7 @@ import { showWorkspaceToast } from './toast.mjs';
|
||||
import { noun } from './plural.mjs';
|
||||
import { operationRow } from './system-controls.mjs';
|
||||
import { openThumbnailReview } from './thumbnail-review.mjs';
|
||||
import { isLocation, locationRow, onMoveEnded } from './library-location.mjs';
|
||||
|
||||
export const OVERVIEW_URL = '/prompt_manager/system/overview';
|
||||
export const STATISTICS_URL = '/prompt_manager/system/statistics';
|
||||
@@ -321,7 +322,7 @@ function surface(heading, rows) {
|
||||
}
|
||||
|
||||
/** The Overview's nodes, and its thumbnail row (painted by its own read). */
|
||||
function overviewNodes(overview, { scope, navigate, openStatistics, openThumbnails }) {
|
||||
function overviewNodes(overview, { api, activity, notify, scope, navigate, openStatistics, openThumbnails }) {
|
||||
const go = section => () => navigate(`#/system/${section}`);
|
||||
const review = operationRow({ scope, title: ambiguousTitle(overview.needs_review),
|
||||
description: Number.isSafeInteger(overview.needs_review?.pending)
|
||||
@@ -334,6 +335,9 @@ function overviewNodes(overview, { scope, navigate, openStatistics, openThumbnai
|
||||
review,
|
||||
operationRow({ scope, title: 'Backup status', description: backupText(overview.backup), label: 'View backups',
|
||||
onClick: go('backups') }),
|
||||
// D129: where the library lives and why; Move library when it is in the node folder.
|
||||
...(isLocation(overview.location)
|
||||
? [locationRow(overview.location, { api, activity, scope, notify, navigate })] : []),
|
||||
// D99: Review thumbnails opens the thumbnail dialog (the concept's row starts a sample job directly).
|
||||
thumbnails,
|
||||
]);
|
||||
@@ -396,7 +400,7 @@ export function mountOverview(container, { api, activity, scope, navigate = navi
|
||||
try {
|
||||
const overview = await overviewRead(signal => api.get(OVERVIEW_URL, { signal, cache: 'no-store' }));
|
||||
if (overview === STALE || scope.signal.aborted) return;
|
||||
const { nodes, thumbnails } = overviewNodes(overview, { scope, navigate, openStatistics,
|
||||
const { nodes, thumbnails } = overviewNodes(overview, { api, activity, notify, scope, navigate, openStatistics,
|
||||
openThumbnails: trigger => openThumbnailReview({ api, store: activity, scope, trigger, navigate }) });
|
||||
thumbnailRow = thumbnails;
|
||||
sources = overview.sources;
|
||||
@@ -429,6 +433,8 @@ export function mountOverview(container, { api, activity, scope, navigate = navi
|
||||
paintThumbnails();
|
||||
}
|
||||
const first = swap(loading());
|
||||
// D129: a finished Move library changes where the library lives.
|
||||
scope.cleanup(onMoveEnded(activity, load));
|
||||
load();
|
||||
loadThumbnails();
|
||||
return first;
|
||||
|
||||
@@ -18,7 +18,9 @@
|
||||
* starts the thumbnail Activity job.
|
||||
*
|
||||
* D98: per-source Remove; an added source reads "not saved yet" until Save; the line under the list says when
|
||||
* ComfyUI's output folder is used instead; "restart pending" for a source the running monitor does not watch yet; the
|
||||
* ComfyUI's output folder is used instead of configured sources or after an unsaved removal of the last one; with no
|
||||
* source configured, ComfyUI's output folder is the row "ComfyUI/output · watching · default" (`in_comfyui`: its public
|
||||
* path is relative to the ComfyUI base) and its scan keeps the server's public path; "restart pending" for a source the running monitor does not watch yet; the
|
||||
* "Refresh files" operation (`/gallery/rescan`) closes the surface. D127: the add dialog's field is "Folder path",
|
||||
* starts empty and says the server's reason in place; its action is "Add", and Enter adds too.
|
||||
* D128: "watching until restart" for a source the running monitor still watches after monitoring was switched off;
|
||||
@@ -46,11 +48,13 @@ export const WATCH_COPY = Object.freeze({ watching: 'watching', stopping: 'watch
|
||||
// The states whose observer runs now: their scan only has to catch up on changes.
|
||||
const WATCHED = new Set(['watching', 'stopping']);
|
||||
const NOT_SAVED = 'not saved yet';
|
||||
// An auto-detected source (ComfyUI's output folder) says it is the default on its own row.
|
||||
const DEFAULT_SOURCE = 'default';
|
||||
const COMFYUI_PREFIX = 'ComfyUI/';
|
||||
export const DUPLICATE_SOURCE = 'same folder as another source';
|
||||
/** D98: the line under the source list when it does not say everything by itself. */
|
||||
export const NOTES = Object.freeze({
|
||||
removedAll: 'No source is configured. After saving, ComfyUI’s output folder is used.',
|
||||
autoDetected: 'No source is configured; ComfyUI’s output folder is used.',
|
||||
fallback: 'No configured source is available; ComfyUI’s output folder is used.',
|
||||
none: 'No source is available.',
|
||||
});
|
||||
@@ -95,8 +99,11 @@ const folderName = path => path.split('/').filter(Boolean).at(-1) ?? path;
|
||||
|
||||
function effectiveRow(source, remove) {
|
||||
const watched = WATCHED.has(source.watch);
|
||||
return { key: remove ?? `auto:${source.root_index}`, title: source.label,
|
||||
description: `${source.path} · ${WATCH_COPY[source.watch] ?? WATCH_COPY.manual}`,
|
||||
const auto = remove === null;
|
||||
// The default row reads from ComfyUI's own folder ("ComfyUI/output") when its public path is relative to the base.
|
||||
const shown = auto && source.in_comfyui === true ? `${COMFYUI_PREFIX}${source.path}` : source.path;
|
||||
const parts = [shown, WATCH_COPY[source.watch] ?? WATCH_COPY.manual, ...(auto ? [DEFAULT_SOURCE] : [])];
|
||||
return { key: remove ?? `auto:${source.root_index}`, title: source.label, description: parts.join(' · '),
|
||||
scan: { label: watched ? 'Scan changes' : 'Scan folder', root: source.root_index, path: source.path }, remove };
|
||||
}
|
||||
|
||||
@@ -124,7 +131,7 @@ export function sourceRows({ sources, saved, current, labels = addedLabels }) {
|
||||
if (dirty) return { rows, note: roots.length ? null : NOTES.removedAll };
|
||||
if (sources.auto_detected) {
|
||||
rows.push(...effective.map(source => effectiveRow(source, null)));
|
||||
return { rows, note: configured.length ? NOTES.fallback : NOTES.autoDetected };
|
||||
return { rows, note: configured.length ? NOTES.fallback : null };
|
||||
}
|
||||
return { rows, note: rows.length ? null : NOTES.none };
|
||||
}
|
||||
|
||||
@@ -414,3 +414,9 @@ pre.prompt-text { font: inherit; font-size: 12px; line-height: 1.75; margin-bloc
|
||||
field) beside their label, as the concept's own radio rows (s1160 `.compare-item label`). */
|
||||
.field[role="radiogroup"] > label { display: flex; align-items: center; gap: 8px; /* D90 */ }
|
||||
.field[role="radiogroup"] input[type="radio"] { width: auto; margin: 0; /* D90 */ }
|
||||
/* D129: the library-location notice sits above every page, in main's own padding and width, so its `.notice` lines
|
||||
up with the page below; main's top padding then spaces it from the page. */
|
||||
.library-banner { max-width: 1900px; margin: auto; padding: 30px 30px 0; /* D129 */ }
|
||||
.library-banner > .notice { margin-bottom: 0; /* D129 */ }
|
||||
@media (max-width: 1250px) { .library-banner { padding: 24px 24px 0; /* D129 */ } }
|
||||
@media (max-width: 700px) { .library-banner { padding: 20px 14px 0; /* D129 */ } }
|
||||
|
||||
@@ -40,6 +40,8 @@
|
||||
<option value="system" selected>System</option>
|
||||
</select></label><a class="canvas-link" href="../">ComfyUI <span aria-hidden="true">↗</span></a></div>
|
||||
</header>
|
||||
<!-- D129: the library-location notice, shown above every page while a node-folder library can move. -->
|
||||
<div id="library-banner" class="library-banner" hidden></div>
|
||||
<main id="workspace-main" tabindex="-1"></main>
|
||||
<footer class="app-footer"><span>ComfyAssets · Local tools. Your library.</span></footer>
|
||||
</div>
|
||||
|
||||
Reference in New Issue
Block a user