Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Vito Sansevero
|
f6e8977c79 | ||
|
Vito Sansevero
|
7f75533a63 | ||
|
Vito Sansevero
|
6f78b4f46f | ||
|
Vito Sansevero
|
2b4b415638 | ||
|
Vito Sansevero
|
c0ef4b6090 | ||
|
Vito Sansevero
|
762fda1448 | ||
|
Vito Sansevero
|
cb1de9de88 | ||
|
Vito Sansevero
|
d5bcae301b |
@@ -211,6 +211,3 @@ coverage/
|
||||
|
||||
# Disposable upstream checkout for real-host CI
|
||||
.ci-comfyui/
|
||||
|
||||
# Old library files a Move library renamed (never deleted; #249 S6b)
|
||||
*.moved-to-*
|
||||
|
||||
@@ -751,34 +751,9 @@ 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 library's
|
||||
folder (the node folder for a library that was not moved):
|
||||
You can customize the database path by creating a `config.json` file in the extension root:
|
||||
|
||||
```json
|
||||
{
|
||||
@@ -995,15 +970,6 @@ 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.
|
||||
- **⚠️ Downgrading**: images recorded by this version (and, after the job, all images) are stored in the new form. An older Prompt Manager version still opens and even restores such a database, but shows those images without workflows or prompt graphs until you update again. To go back to an older version, stop ComfyUI, delete `prompts.db-wal` and `prompts.db-shm`, and copy the **earliest** `prompts.db.pre-workflow-storage_<date time>` over `prompts.db`, or, if the job never ran, a backup taken before this update (records added since that backup are lost). Otherwise stay on this version.
|
||||
|
||||
### v3.3.1 (Proxy Routing Fix)
|
||||
|
||||
- **🌐 Proxy support**: the dashboard, gallery, tag page and metadata viewer now work when ComfyUI is reached through a proxy mount such as SwarmUI's `/ComfyBackendDirect/`. Scripts, images, API calls, progress streams and gallery folder filters stay under the mount, and the canvas node's buttons use ComfyUI's `fileURL()`. Direct ComfyUI URLs are unchanged. Thanks to @ghosttigerllc-bit ([#164](https://github.com/ComfyAssets/ComfyUI_PromptManager/pull/164))
|
||||
|
||||
+16
-507
@@ -1932,8 +1932,8 @@ pending count (`GET /prompt_manager/review/items?limit=1`, its `total`).
|
||||
unavailable file's empty media button keeps the concept's 1.33 media area, so its
|
||||
card lines up with the others. A long file name in the panel's `p.small` wraps inside
|
||||
the panel (`overflow-wrap: anywhere`), where the concept's short sample name never
|
||||
overflows. A queue card shows the Images listing's thumbnail (`thumbnail_url`, the
|
||||
same URL the Images page uses; a preview not made yet is made on first view); a re-read
|
||||
overflows. A queue card shows the Images listing's thumbnail when the server knows one
|
||||
(`thumbnail_url`, the same URL the Images page uses), else the original; a re-read
|
||||
rebuilds a card whose media, target or availability changed (a root came or went), so
|
||||
Open and Tag follow it. When an unanswered decision's item leaves the queue in a
|
||||
re-read, the decision was recorded: the page finishes as a recorded decision does (its
|
||||
@@ -2115,11 +2115,10 @@ 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`, `watch` (`watching` only when the image monitor's observer is alive,
|
||||
`root_index` and `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`) and `in_comfyui` (its public path is relative to the ComfyUI base)),
|
||||
`sources.auto_detected`, `monitor` and `settings.monitoring_enabled`.
|
||||
`restart_pending`), `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
|
||||
@@ -2231,59 +2230,6 @@ through the Activity store, which toasts "Thumbnail generation started."; the jo
|
||||
in Activity. A start refused because another job holds the slot toasts the D57 busy
|
||||
link toast and stays. The legacy gallery's thumbnail tools are unchanged.
|
||||
|
||||
Automatic previews (#249, `py/api/thumbnail_auto.py`). The approved mock has no
|
||||
repeated global thumbnail prompt (`docs/concepts/workspace/README.md`), so the
|
||||
workspace never asks; previews are a cache the server keeps:
|
||||
|
||||
- Every listing gives an image or video a thumbnail URL (`media_thumbnail_url`: Images,
|
||||
saved views and favorites, Library cards and galleries, Needs review, duplicate
|
||||
groups). An existing preview's URL keeps its `v=<mtime>`; a missing one is listed
|
||||
without `v=`.
|
||||
- `GET /images/serve/thumbnails/<rel>_thumb<ext>` for a missing preview is refused as
|
||||
text first (traversal, a backslash, a drive or stream colon, empty or dot segments, so
|
||||
`D:/x`, `\\server\share` and a leading `/` never reach a joined path), then admitted
|
||||
on the event loop (`_PreviewGate`): 3 previews are written at once on the API's own
|
||||
`pm-preview` pool (never the default executor), one flight runs per file (the history
|
||||
route shares the flight key and the admission), and at most 64 requests wait; one
|
||||
more gets its fallback at once without listing a folder. Inside the admission the
|
||||
preview maps back to its media file under the same root (exact names, then one folder
|
||||
listing for a mixed-case extension; a regular file inside the root, never through a
|
||||
link) and is made with the job's defaults (512 px, medium). When it cannot be made an
|
||||
image's original is served and a video's placeholder frame (a video is never sent to
|
||||
an `<img>`), both `Cache-Control: no-store`; cards and the inspector cover also swap a
|
||||
video back to its own element when its preview image fails. Only a file that cannot
|
||||
be decoded is remembered for the session (path, size, mtime; retried once it changes)
|
||||
and counted as `failed` in `/images/thumbnails/status`, which Overview's row adds
|
||||
("N missing thumbnails · M failed"); I/O errors and a preview another program holds
|
||||
open (Windows) are not. A preview folder (or any link segment of it) resolving out of
|
||||
the root is never written. Media inside a nested `thumbnails` folder is listed with no
|
||||
thumbnail and keeps its original. The listing's `thumbnail_ready` says whether the
|
||||
preview file exists (the legacy gallery's dot and "Fast").
|
||||
- The video decoder falls through OpenCV, ffmpeg (a tracked process, killed on
|
||||
shutdown) and the placeholder frame: an OpenCV that cannot open or read a clip (AV1)
|
||||
no longer ends the run.
|
||||
- The image monitor hands every settled output to `queue_capture_thumbnail`, which
|
||||
queues it for one background worker: each wake takes at most Worker threads files and
|
||||
writes them on that many threads; a queued file is not queued twice; the queue holds
|
||||
10,000 and a file left out is made on view.
|
||||
- A finished scan that added files counts the missing previews of the sources it
|
||||
walked: 200 or more start the thumbnail job (missing only, visible and cancellable in
|
||||
Activity, never silent bulk work); fewer go through the worker.
|
||||
- Holds: automatic previews write only `thumbnails/` files and no records, so they run
|
||||
beside a scan, image tagging, the thumbnail job and the Needs review recording (S5
|
||||
M7). They wait (a view gets the original, a queued file is left for its view) while a
|
||||
restore, a duplicate removal, a sources change, the gallery's thumbnail clear or an
|
||||
`exclusive` job runs (`_auto_thumbnail_blocker`), checked when a write starts (after
|
||||
its slot) and again just before its rename.
|
||||
- Shutdown is bounded and logged: the worker stops at a file boundary, ffmpeg is
|
||||
killed, an idle worker leaves (the next output starts a fresh one), the view pool
|
||||
shuts down, and after the app's cleanup a late monitor event
|
||||
starts nothing. A package reload stops the workers the previous import left.
|
||||
- Cards and strips load previews only: Images in every layout and size, Library cards,
|
||||
Needs review cards and the inspector's cover (a video's cover is its preview frame).
|
||||
The media viewer's main media and downloads load the original
|
||||
(`workspace-thumbnails.spec.mjs` records every media request; `tests/test_thumbnail_auto.py`).
|
||||
|
||||
Overview (`system-overview.mjs`) reads `GET /prompt_manager/system/overview` once per
|
||||
visit through `latestOnly()` on the section scope; a late answer after the section left
|
||||
paints nothing. Until it answers, the section shows "Loading the overview…"; a failure
|
||||
@@ -2311,7 +2257,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–D129 assigned):
|
||||
Recorded deviations (plan D60–D67, D79–D81 renumbered D92–D99 and D111–D113; D121–D128 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
|
||||
@@ -2331,8 +2277,7 @@ Recorded deviations (plan D60–D67, D79–D81 renumbered D92–D99 and D111–D
|
||||
"Backup status" reads "No retained backup record." until verified backups exist (S8).
|
||||
The thumbnail row is "Checking thumbnails…" until its own read answers, then "No
|
||||
missing thumbnails", "N missing thumbnails" (the pending count of
|
||||
`/images/thumbnails/status`, followed by " · M failed" when automatic previews
|
||||
failed this session) or "Thumbnail status unavailable"; Review thumbnails
|
||||
`/images/thumbnails/status`) or "Thumbnail status unavailable"; Review thumbnails
|
||||
opens the thumbnail dialog (D99). "One display configuration" says "Grid
|
||||
and density settings are shared across views." (no paging). The sources row is
|
||||
"N output sources" with "Monitoring and enrichment share one source list.", prefixed
|
||||
@@ -2386,12 +2331,9 @@ 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) 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
|
||||
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
|
||||
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
|
||||
@@ -2462,20 +2404,6 @@ 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
|
||||
@@ -2571,392 +2499,11 @@ node --test tests/js/workspace-system-state.test.js tests/js/workspace-system-co
|
||||
tests/js/workspace-inspection-state.test.js tests/js/workspace-images-state.test.js \
|
||||
tests/js/workspace-system-sources.test.js tests/js/workspace-thumbnail-review.test.js
|
||||
python3 -m unittest tests.test_settings_sources tests.test_scan_scope -v
|
||||
python3 -m unittest tests.test_maintenance_plans tests.test_api_maintenance -v
|
||||
npm run --prefix tools/vendor test:workflows -- system-concept.spec.mjs workspace-system.spec.mjs \
|
||||
workspace-system-sources.spec.mjs \
|
||||
--project chromium --project touch --project firefox --project webkit
|
||||
```
|
||||
|
||||
### Maintenance contract (#249 System S6)
|
||||
|
||||
`POST /prompt_manager/maintenance/preview {operation, options:{}}` reviews one
|
||||
operation and answers a `Plan` (`py/api/maintenance.py`, `database/maintenance.py`):
|
||||
`duplicate_prompts` (legacy `cleanup_duplicates`: equal text ignoring case and outer
|
||||
spaces, the oldest prompt of each group kept; a group with a member in a saved pair is
|
||||
kept whole), `optimize` (VACUUM, with the database size, its free pages and the free
|
||||
space it needs), `unused_components` (legacy `prune_orphaned_prompts`; the #239
|
||||
protections are the predicates `ORPHANED_PROMPTS_SQL` is built from, so the plan's
|
||||
protected reasons are exactly the prune's rules) and `missing_links` (legacy
|
||||
`cleanup_orphaned_images`, judged as the health check's `missing_files` through
|
||||
`health.link_states`: only a link under a source that can be read now counts; one under
|
||||
an unavailable source or outside every source is protected). The operation enum is
|
||||
strict, `options` must be `{}` and the body is capped at 4 KiB (512 KiB for
|
||||
`duplicate_files`: 1 MiB, below). A preview meets the
|
||||
read-only standard above on its own `mode=ro` snapshot (`test_preview_writes_nothing`,
|
||||
per operation, plus the media bytes) and stats files after the snapshot ended. The
|
||||
plan lists at most 50 detail rows and no server path. Its fingerprint (the ids, text
|
||||
hashes, metadata, tags, Library state and pair refs of every prompt involved, or each
|
||||
link's row and source), its targets and the sources it judged stay in the server's
|
||||
`PlanStore` under a single-use `secrets.token_urlsafe(32)` token: 15 minutes, at most 8
|
||||
plans and 16 MiB of payload. A plan over that budget is refused with `too-large`, never
|
||||
truncated; a new plan drops the oldest ones, which then read as expired; a restart
|
||||
forgets them all. A plan is bound to the database it read: a restore during the preview
|
||||
answers `storage-changed`. Errors are coded (`invalid-request`, `too-large`,
|
||||
`storage-changed`, or a null code with "Failed to review maintenance") and carry no
|
||||
path or exception text.
|
||||
|
||||
`POST /prompt_manager/maintenance/apply {token}` starts the reviewed plan as an
|
||||
Activity job (202 `{job}`): kind `maintenance`, titled after the concept's row
|
||||
("Duplicate prompt text", "Optimize database", "Unused prompt components", "Missing-file
|
||||
links"), admission `records` (it never runs beside a scan, restore, LoRA import or other
|
||||
record-changing job, legacy-page holds included: `test_apply_refused_while_restore_runs`),
|
||||
not cancellable, with the groups `changed`, `protected` and `errors`. The token is consumed
|
||||
on the loop turn the job starts, so a second apply answers 409 `plan-used` with that job;
|
||||
a refused start (`busy`) keeps it. Other errors: `invalid-request`, `plan-expired` (410:
|
||||
unknown, expired, or forgotten after a restart) and `storage-changed` (409: a restore
|
||||
replaced the database since the review; the plan is dropped). The job:
|
||||
|
||||
1. takes a safety backup when the plan was reviewed with the policy on or the policy is
|
||||
on now (`maintenance.safety_backup` in `config.json`, default `true`; GET/POST
|
||||
`/prompt_manager/settings` field `safety_backup` with the revision). The copy is made
|
||||
with SQLite's backup API into a temp file beside the database, switched to a
|
||||
rollback journal (so a later read-only open leaves no `-wal`/`-shm` beside it),
|
||||
verified with `verify_database_file`, fsynced, made `0o600` on POSIX and only then
|
||||
renamed with `os.replace` onto `<db>.backup_<YYYYmmdd_HHMMSS_ffffff>`, reserved with
|
||||
`O_EXCL` (`_1`, `_2`… when a coarse clock repeats a stamp:
|
||||
`test_backup_names_unique_within_one_second`). A backup file on disk is therefore
|
||||
always a verified one. Any failure removes what it wrote and ends the job
|
||||
`backup-failed` before anything is changed (`test_backup_failure_aborts_before_write`);
|
||||
2. applies the plan: inside `BEGIN IMMEDIATE` the fingerprint is computed again on the
|
||||
write connection and a difference ends the job `plan-changed` with nothing written
|
||||
(`test_apply_rejects_changed_target`; a missing link whose file is back or whose source
|
||||
went away is never removed). The targets are deleted in batches of 400 ids in that one
|
||||
transaction. Missing-file links remove database rows only; every media file's bytes
|
||||
are unchanged (`test_missing_links_never_deletes_files`). VACUUM runs outside a
|
||||
transaction and binds to its database only;
|
||||
3. counts the changed and protected records (outcome rows hold root-relative paths or
|
||||
prompt ids only), logs the backup by its file name, and prunes safety backups to the
|
||||
newest 10. With no backup the log and the status line say "No safety backup".
|
||||
|
||||
The legacy restore (`POST /prompt_manager/restore`) takes a backup of any size the disks
|
||||
have room for, so a real library's safety backup restores: the upload, its staged copy
|
||||
and the swap (the upload twice in the temporary folder, once more plus the current
|
||||
database beside the database, added up on one disk, each with the 16 MiB margin) must
|
||||
fit, under a 64 GiB ceiling (`restore_max_bytes`). A request whose Content-Length does
|
||||
not fit is refused before anything is written, one without is refused as soon as its
|
||||
running size does not, both with 507 `insufficient-space` and "Not enough free disk
|
||||
space to restore a 1.6 GB backup; there is room for about …. Nothing was changed."
|
||||
(no path); past the ceiling it is the legacy 400 "File too large" with code `too-large`.
|
||||
The upload streams with `read_chunk`, which aiohttp's `client_max_size` (ComfyUI's
|
||||
`--max-upload-size`, 100 MB by default) does not cap, and its temp file is closed and
|
||||
removed on every path (`TestLargeRestores`, a 150 MB backup under a 100 MB app limit).
|
||||
|
||||
Safety backups and restore backups are one pool: every `<db>.backup_*` file beside the
|
||||
database counts toward the same `MAX_SAFETY_BACKUPS` (10), whichever wrote it, and the
|
||||
pool is pruned after every backup a job took, also when the apply then failed
|
||||
(`test_safety_backups_share_the_restore_pool`,
|
||||
`test_the_pool_is_pruned_when_the_apply_fails_after_the_backup`). The legacy
|
||||
`GET /backup` download writes no file beside the database, so it adds nothing to the
|
||||
pool. Each plan reports `backup.expected_bytes` (the backup's size, `page_count *
|
||||
page_size`, or null with the policy off) and `disk: {free_bytes, required_bytes,
|
||||
sufficient}` for the database folder: the backup needs its size there, and Optimize
|
||||
needs the database's size twice (the WAL beside it and SQLite's temporary copy, added
|
||||
up when the temporary folder is on the same disk), each plus a 16 MiB margin. The
|
||||
temporary folder is the one SQLite really uses (`sqlite_temp_dir`): on Unix
|
||||
`SQLITE_TMPDIR`, `TMPDIR`, then the first writable of `/var/tmp`, `/usr/tmp` and
|
||||
`/tmp` (not Python's `tempfile` order, which picks a RAM-backed `/tmp` on many Linux
|
||||
systems); on Windows `TMP`, `TEMP`, then `USERPROFILE`. Whether a backup is taken is
|
||||
decided once, when Apply measured the disk, and handed to the job. Apply looks at the
|
||||
token again with no await before it starts the job, so a token used or expired during
|
||||
the measurement answers its code, never a 500. A backup is put in place with
|
||||
`os.link` (atomic, refuses an existing name; reserve-and-replace where a filesystem
|
||||
has no hard links), and startup removes the `.pm-safety-*.tmp` copies an interrupted
|
||||
backup left. Optimize runs alone: its job has the `exclusive` admission class (refused
|
||||
beside any job, and every job is refused while it runs), and its plan says the database
|
||||
is locked while VACUUM runs and that run counts or captures waiting longer than the
|
||||
30-second busy timeout are not recorded. The duplicate prompt merge is strict inside
|
||||
the apply: a group whose copies link one image twice, or whose removed copy has its own
|
||||
Library title, is protected in the plan, and any merge error rolls the transaction back
|
||||
(`MergeConflict`); the legacy cleanup keeps its lenient behaviour. Apply
|
||||
measures again before it starts a job and refuses a short disk up front with 507
|
||||
`insufficient-space` and a truthful message ("Not enough free disk space: the safety
|
||||
backup needs about … in the database folder and … is free. Nothing was changed."); no
|
||||
job starts and the token stays usable
|
||||
(`test_apply_refuses_up_front_when_the_backup_does_not_fit`,
|
||||
`test_apply_refuses_up_front_when_vacuum_does_not_fit`).
|
||||
|
||||
**Duplicate image files** (`py/api/duplicate_files.py`, `py/api/duplicate_review.py`;
|
||||
`tests/test_api_duplicate_files.py`). `POST /prompt_manager/jobs/duplicates` starts the
|
||||
find: kind `duplicates`, "Find duplicate image files", admission `read`, cancellable,
|
||||
group `groups_found`, over every effective source (the legacy scan's
|
||||
`_hash_media_groups`: no symlinks, no `thumbnails`, bounded, each file once). The
|
||||
find lists a physical file once, by `(st_dev, st_ino)`: a source listed twice under two
|
||||
spellings, a symlinked source inside another one and a hardlink are one file, never a
|
||||
copy of themselves (`test_a_folder_listed_twice_never_loses_its_only_copy`,
|
||||
`test_a_symlinked_source_inside_another_is_one_file`,
|
||||
`test_a_hardlink_is_never_a_copy_to_remove`). One file reached through two mounts is
|
||||
protected, never guessed at: a removal on another `st_dev` than its kept copy, with the
|
||||
same size and mtimes within 2 s, is kept as "may be the same file reached through another
|
||||
mount or share" when its `st_ino` is equal or its path within its source is the same, or
|
||||
one ends with the other on a whole folder name (sources listed at different depths:
|
||||
`test_a_deeper_source_on_another_mount_is_protected`)
|
||||
(NFS and SMB, a loopback mount, a FUSE pool: `TestAnotherMount`). Some true copies on
|
||||
another disk are therefore never offered. A volume that reports no file id protects
|
||||
every pair ("Can't confirm this is a separate file on this drive"). The job checks both
|
||||
before each file's records move, holds the kept copy open, and right before the unlink
|
||||
checks again that the held kept copy is still the file at its path (`fstat` of the held
|
||||
handle against `os.stat` of the path) and that the guard still holds; otherwise the file
|
||||
is kept, its records move back to it, and it counts as not removed
|
||||
(`test_a_kept_copy_deleted_during_the_move_keeps_the_file`). On Windows the hold also
|
||||
makes deleting the kept copy itself through another path fail. The server holds the groups (content
|
||||
hash, root-relative paths, source index, size, mtime and inode) for the last three
|
||||
finds; `GET /prompt_manager/maintenance/duplicates/{job_id}?offset&limit` pages
|
||||
them ten at a time (`not-found` 404, `job-running` or `no-groups` 409). The preview
|
||||
`{operation:"duplicate_files", options:{job_id, keep:[{group_id, root_index, path}]}}`
|
||||
(body cap 1 MiB, at most 2,000 keep entries; `test_the_keep_limit_fits_the_body_cap`) removes, in each group named, every copy
|
||||
not kept; a keeper that is not in its group is `invalid-request`, and a source no longer
|
||||
listed at its index is `sources-changed`. Every keeper and removal is checked with
|
||||
`file_problem`: containment as text first, with no filesystem access (inside its
|
||||
source, not under `thumbnails`, a media extension), then `lstat` (a link or any
|
||||
non-regular file is refused), the resolved path still inside the resolved source, not
|
||||
the kept copy itself (`os.path.samestat`), and the size, mtime and inode the find saw.
|
||||
The preview reads no file; the removal job hashes each file and its kept copy again
|
||||
right before the unlink (`test_preview_reads_no_file_and_spots_a_touched_one`). A removal that fails, or whose kept copies all fail, is
|
||||
protected with its reason. The plan names each removal, its kept copy and how many
|
||||
database links move, says the safety backup covers the database and not the files,
|
||||
and writes nothing. Apply runs `duplicate_cleanup` (records, not cancellable, groups
|
||||
`removed`, `kept`, `errors`): sources unchanged, the promised safety backup, then per
|
||||
file the kept copy and the file are checked again (hashes included, and never the same
|
||||
file), the file's records move to the kept copy in one short `BEGIN IMMEDIATE`
|
||||
transaction, and only then is the file removed (every hash handle closed first:
|
||||
`test_handles_closed_before_unlink`), with its thumbnail (a video's `.jpg` too). A crash
|
||||
or a failed unlink therefore leaves the records on the identical kept copy, and a failed
|
||||
move leaves the file and its records alone: no row ever names a deleted file
|
||||
(`test_a_failure_between_two_files_leaves_no_row_on_a_deleted_file`). The status line
|
||||
counts the files removed and not removed. A changed file is kept (`test_changed_content_skips_file`); the
|
||||
first failed unlink stops the job, a locked or read-only file with "The file is open in
|
||||
another program or read-only; it was not deleted, and its links now name the kept copy."
|
||||
and never the OS text, and the
|
||||
files after it are listed as not attempted (`test_partial_failure_reports_each_path`).
|
||||
Records move for the reviewed files only (`database/duplicate_copies.py`): each image
|
||||
record, looked up under every spelling of the file by containment
|
||||
(`py/api/duplicate_spellings.py`: each configured, monitored, ComfyUI output or listed
|
||||
folder whose `os.stat` identity is a folder above the file, plus the path below it, so a
|
||||
parent spelling of a nested source and a case variant match:
|
||||
`test_links_under_a_symlinked_source_spelling_move`,
|
||||
`test_a_parent_spelling_of_a_nested_source`), moves to the kept copy under the same
|
||||
spelling, or, when its prompt already links the kept copy, hands its saved-pair
|
||||
links to that record and is deleted; a favorite under any listed root that holds the
|
||||
file moves with it (`test_nested_sources_subfolder_first_with_a_parent_favorite`)
|
||||
(`test_db_links_reconciled_only_for_removed_files`); the copy's image tags join the kept
|
||||
copy's as their union; and its pending Needs review items (still `pending` at revision
|
||||
0 with no decision row) are settled (`test_tags_move_as_a_union_and_pending_items_settle`).
|
||||
The plan states each of these. The find refuses while a removal
|
||||
runs and the removal while a find runs.
|
||||
|
||||
The legacy gallery routes keep their shapes on the same kinds: `/scan_duplicates`
|
||||
holds `duplicates` (read, "Find duplicate image files (gallery page)"),
|
||||
`/scan_duplicates/stream` starts the find over the gallery's output folder or attaches
|
||||
to a running one that covers it (409 otherwise), `/scan_duplicates/status` reports the
|
||||
latest find, and `/delete_duplicate_images` holds `duplicate_cleanup` (records) and is
|
||||
refused while a find runs (`test_legacy_delete_duplicates_unchanged`). The link
|
||||
statements are index searches (`test_link_queries_use_the_path_indexes`; the fixture
|
||||
check is in the PR). Windows: the hash reads close before `os.remove`; a sharing
|
||||
violation is a `PermissionError` answered with the coded message (native on the
|
||||
Windows runners, `test_locked_file_gives_a_coded_message`).
|
||||
|
||||
A failure the job cannot name ends it "Maintenance failed; nothing was changed." with a
|
||||
null code and no exception text. Windows: every connection the backup and the previews
|
||||
open is closed before the temp copy is renamed or removed, and the suites' teardown
|
||||
asserts no handle under the temp root is still open.
|
||||
|
||||
The legacy `POST /prompt_manager/maintenance` keeps its shape. Its `vacuum` result is
|
||||
`success:false` with the error when VACUUM fails (`vacuum_database` now raises after
|
||||
logging), and `cleanup_orphaned_images` removes only links under a source that can be
|
||||
read now (the health check's sources), never every link of an unmounted drive.
|
||||
`cleanup_missing_images()` without sources keeps the old rule for its other callers.
|
||||
Query plans on a fixture of 5,100 prompts, 8,000 images, 300 pairs and 200 recorded
|
||||
runs: every per-prompt protection and fingerprint read is an index search, the
|
||||
duplicate groups and the link rows are one scan each, and a delete batch (400 ids)
|
||||
searches by rowid with its cascades by index.
|
||||
|
||||
**Workflow storage** (`database/workflow_documents.py`, `database/workflow_storage.py`,
|
||||
`py/api/workflow_storage.py`; `tests/test_workflow_documents.py`,
|
||||
`tests/test_workflow_migration.py`, `tests/test_api_workflow_storage.py`). Each distinct
|
||||
workflow and prompt graph text is one `workflow_documents` row: `sha256` of its exact
|
||||
UTF-8 bytes (unique), `byte_length` and `data` = `zlib.compress` of those bytes. Image rows
|
||||
reference it through `workflow_document_id` / `prompt_document_id` (foreign keys, one
|
||||
index each). The text is stored and hashed exactly as written, never parsed and
|
||||
re-serialized. Two forms are readable forever: old (text inline, no id) and new (inline
|
||||
NULL, id set); inline text wins when both are present. Every reader goes through
|
||||
`read_stored_text` / `resolve_image_row` (image lists, image by id, prompt info for the
|
||||
inspector and metadata page, diagnostics, pair ingestion's comparison), which also drops
|
||||
the id keys so payloads keep their shape; a document that fails its length or hash check
|
||||
reads as unreadable (None), never as other text. Round trips are proven byte for byte
|
||||
over edge texts (spacing, key order, duplicate keys, number spellings, `NaN`/`Infinity`
|
||||
tokens, NUL, BOM, astral and combining characters, 500-deep nesting, a 1.2 MB+
|
||||
workflow, `""` and NULL) and seeded random text and workflows
|
||||
(`TestRoundTrip`). New writes (`link_image_to_prompt`, pair ingestion for scan, monitor
|
||||
and Needs review) store the new form. No trigger is installed (restore validators refuse
|
||||
them): a document goes in the transaction that drops its last reference (image and
|
||||
prompt deletes, missing-link cleanup, duplicate merges, duplicate-file link merges,
|
||||
ingestion updates: `TestGarbage`, `TestMoreGarbagePaths`), and startup, the migration's
|
||||
end and Reclaim space sweep anything a path missed. `verify_database_file` checks the
|
||||
table definition, the id columns, their foreign keys and indexes, the stored values and
|
||||
dangling ids (`TestSchemaErrors`); a pre-dedupe file is accepted and migrated on open
|
||||
and on restore. The `workflow_storage` maintenance plan counts the old-form records
|
||||
(one pass over record headers) and its apply starts "Store each workflow once"
|
||||
(`records`, cancellable): first, always and whatever `maintenance.safety_backup` says,
|
||||
a verified pre-migration backup named `<db>.pre-workflow-storage_<YYYYmmdd_HHMMSS_ffffff>`,
|
||||
outside the pruned `<db>.backup_*` pool and kept until the user removes it (the preview
|
||||
reports `backup.policy: "required"`; `test_the_backup_is_taken_even_with_the_policy_off`,
|
||||
`test_the_backup_is_outside_the_pool_never_pruned_and_restores`); then batches of at most
|
||||
200 rows, 16 MB of text or 0.5 s, each one `BEGIN IMMEDIATE` transaction found by a rowid
|
||||
range search and read one row at a time, followed by 0.15 s with the write lock free so
|
||||
a writer on another connection gets in (`test_a_writer_on_another_connection_gets_in_between_batches`).
|
||||
Each text is read back through the accessor before its row changes. Each row has its own
|
||||
savepoint: a row that cannot be rewritten (text SQLite cannot decode, a text that would
|
||||
not come back) stays inline, is listed in `errors` and named in the job's last line, and
|
||||
the rest go on, Reclaim space included (`TestRowErrors`); any other failure (a full disk,
|
||||
an I/O error, a malformed database, a bug) rolls the batch back whole and fails the job
|
||||
(`TestSystemErrorsFailTheBatch`, `TestSystemErrorsInTheJob`). With nothing to rewrite the
|
||||
plan says "Nothing to rewrite" and Apply answers 409 `nothing-to-do` with no backup; only
|
||||
the first apply keeps a `pre-workflow-storage` backup, and a later one (after a Cancel, a
|
||||
rerun) takes a pooled safety backup (`TestOneKeptBackup`). The job loop's pause is tested
|
||||
with a second-connection writer (`TestLiveLoopLetsWritersIn`). A reused document is decompressed and checked first and repaired from the text
|
||||
when damaged, and `verify_database_file` (restore and backup verification) reads every
|
||||
document and refuses a damaged one (`TestDamagedDocuments`). The legacy table rebuilds
|
||||
refuse a file that already has the document id columns (`TestRebuildGuards`).
|
||||
`tools/workflow_text_digests.py <db> <out>` writes the sha256 of every record's stored
|
||||
texts, so a `diff` of its output before and after the rewrite proves byte identity. It
|
||||
loads only `database/workflow_documents.py` (no package start-up, no log handler) and
|
||||
opens the database `mode=ro`, with `immutable=1` when no `-wal`/`-shm` exists, so it
|
||||
creates none (`TestDigestToolSideEffects`). A damaged document's verify reason names it
|
||||
and says what to do, and a backup that fails on it says so in the job line. A crash or Cancel leaves
|
||||
committed batches new form and the rest old form, and the next apply continues
|
||||
(`test_a_crash_mid_migration_leaves_every_row_readable_and_resumes`,
|
||||
`test_cancel_stops_after_the_current_batch_and_a_rerun_continues`). Values that are not
|
||||
text stay inline (protected). "Reclaim space" follows when every row was looked at
|
||||
(`exclusive`, not cancellable): a disk check for VACUUM's copies of the live pages,
|
||||
then sweep, VACUUM and `wal_checkpoint(TRUNCATE)`; when another job holds the registry it
|
||||
does not start and the finished migration names Optimize database. No image or video
|
||||
file is opened (`test_the_jobs_never_open_a_media_file` tracks every Python open under
|
||||
the output folder and compares bytes and mtimes). Memory stays bounded by one row plus
|
||||
the batch's distinct documents (`test_memory_stays_bounded_by_one_row`, tracemalloc).
|
||||
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
|
||||
@@ -3178,52 +2725,15 @@ still have made the file, so that text is not certain either. Both report
|
||||
`unresolved` ("could not be traced"), are never linked, and their pending,
|
||||
undecided items whose own evidence offers no choice (a limited item counts as
|
||||
a choice when it holds at least two candidates) leave the queue on the next
|
||||
scan (`settle_untraced_review_items`). The maintainer's
|
||||
second live check found 39 items whose 11 candidates were all "unresolved".
|
||||
Inspection-profile choices (`tests/test_review_inspection.py`, ruling a): when
|
||||
the default profile is ambiguous without a choice and nothing links the file,
|
||||
the scan (never the monitor) resolves it again with `INSPECTION_PROFILE`, which
|
||||
knows `TextEncodeQwenImageEdit`, `StringConstantMultiline`, `JoinStringMulti`
|
||||
and the rest of `web/js/inspection-semantics.json`. With at least two choosable
|
||||
candidates there, the file is `ambiguous` and is queued from that profile's
|
||||
evidence, which adds `profile: "inspection"` and `semantics`: the SHA-256 of
|
||||
`class_projection`, the table restricted to the graph's node classes plus its
|
||||
class-independent `schema` and `stage_inputs` (the resolver looks the table up
|
||||
only by a graph node's class). The item's `resolver_version` is
|
||||
`inspection_review_version(metadata)`, 2^40 plus 40 bits of SHA-256 over the
|
||||
profile with its budgets, `RESOLVER_VERSION`, `INSPECTION_RESOLVER_VERSION`
|
||||
(bump it when the inspection profile would resolve the same metadata otherwise)
|
||||
and that fingerprint. A table edit about the
|
||||
graph's classes behaves like a resolver change (I2); an edit about other
|
||||
classes changes nothing. One file never holds two pending items for the same
|
||||
bytes: a new version's item supersedes the byte record's other items when all
|
||||
are pending and undecided (revision 0, no decision row, the self-heal's safe
|
||||
delete); beside a decided or advanced item nothing new is stored, and the
|
||||
rescan's stat shortcut then skips those bytes under any version. A pending item
|
||||
stays applicable while the bytes still offer its chosen candidate.
|
||||
Apply re-resolves the capture with the profile the item recorded and re-checks
|
||||
the chosen hash; the link goes through the same compatible-only ingestion. The
|
||||
inspection profile never links on its own, and ingestion keeps the default
|
||||
profile (D1). The heavier profile runs only for files already ambiguous without
|
||||
a choice, and once per file state: the capture reuses the scan's result when
|
||||
it holds the same records; a rescan skips files already queued or decided at
|
||||
their current size and `mtime_ns` under their version, and files this process
|
||||
already found untraced at the same state (an in-memory map of at most 65,536
|
||||
files, oldest forgotten first, keyed by 16-byte digests: 8.6 MiB measured at
|
||||
the cap). A file whose inspection raises is logged, counted in
|
||||
`review_errors` and reported `unresolved`; the scan goes on. An older version's
|
||||
pending, undecided item of bytes now untraced leaves the queue
|
||||
(`settle_stale_review_items`), so an item never stays stuck at 409. The 16-candidate limited
|
||||
rule and the 8 MiB bound apply unchanged. The queue heals itself: after each batch, the scan removes the pending
|
||||
scan or monitor pass (`settle_untraced_review_items`). The maintainer's
|
||||
second live check found 39 items whose 11 candidates were all "unresolved". The queue heals itself: after each batch, the scan removes the pending
|
||||
items of every listed file (read or skipped as known) that a saved row now
|
||||
links, and the monitor does the same for each ambiguous output it links
|
||||
(`settle_linked_review_items`; the review index revision moves). Only an
|
||||
undecided item goes (`pending`, revision 0, no decision row): a decision is
|
||||
never deleted or changed, so no schema state was added. The monitor never
|
||||
records an item (only scans add items, maintainer ruling 2026-10-23): an
|
||||
unlinked file passed to `record_review_file` is `skipped` with no capture,
|
||||
and only a linked one's pending items are settled
|
||||
(`test_monitor_never_records_an_unlinked_ambiguous_output`). The
|
||||
queues an output it linked; an unlinked file passed to `record_review_file`
|
||||
follows the scan's rule (a file the fallback would link is `skipped`). The
|
||||
maintainer's live check of #284/#288 found the queue at 10,232 against a review
|
||||
group of 4,677: every fallback-linked ambiguous file had been queued. It captures
|
||||
the file (scope `needs-review`), reads its metadata again from the capture,
|
||||
@@ -3255,10 +2765,9 @@ Apply, and waits up to 2 s for room before counting the file as not added
|
||||
pass runs after its own link, best effort: a failure is logged and the link
|
||||
stays. Recording follows the registry's `records`
|
||||
admission (#289): the scan's recording runs under its own `records` job, and
|
||||
the monitor's pass (one file, not a job, which only settles) is refused
|
||||
while a `records` or `guarded` job or hold runs, except thumbnail work, which
|
||||
writes no records (S5 M7). It then writes nothing, logs that the file waits
|
||||
for the next scan, and answers `deferred`; the next "new" scan reads the ambiguous file
|
||||
the monitor's (one file, not a job) is refused while a `records` or `guarded`
|
||||
job or hold runs, except thumbnail work, which writes no records (S5 M7). It then records nothing, logs that the file waits for the
|
||||
next scan, and answers `deferred`; the next "new" scan reads the ambiguous file
|
||||
again (`test_monitor_recording_is_refused_beside_record_changing_jobs`). A new item moves `review_index_revision`. A
|
||||
rescan does not copy or hash an ambiguous file whose current size and
|
||||
`mtime_ns` equal the verified fingerprint of bytes an item already cites (a
|
||||
|
||||
+1
-8
@@ -132,18 +132,11 @@ try:
|
||||
_review_api = globals().get("api")
|
||||
if _review_api is not None:
|
||||
_global_image_monitor.set_review_recorder(_review_api.record_review_file)
|
||||
# Every settled output's preview is queued for the API's background
|
||||
# worker (#249): the workspace never asks to generate thumbnails.
|
||||
_global_image_monitor.set_thumbnail_queue(_review_api.queue_capture_thumbnail)
|
||||
_review_host = getattr(config.server_instance, "app", None)
|
||||
if _review_host is not None and hasattr(_review_host, "on_cleanup"):
|
||||
|
||||
async def _stop_review_recording(
|
||||
app, monitor=_global_image_monitor, api=_review_api
|
||||
):
|
||||
async def _stop_review_recording(app, monitor=_global_image_monitor):
|
||||
monitor.set_review_recorder(None)
|
||||
monitor.set_thumbnail_queue(None)
|
||||
api.close_auto_thumbnails(final=True)
|
||||
|
||||
_review_host.on_cleanup.insert(0, _stop_review_recording)
|
||||
|
||||
|
||||
@@ -1,215 +0,0 @@
|
||||
"""Records of duplicate image files: count them and move them (#249 S6, T17).
|
||||
|
||||
A removed copy's records move to its verified kept copy before the copy is
|
||||
unlinked (``py/api/duplicate_files.py``), one file per short ``BEGIN
|
||||
IMMEDIATE`` transaction, so no row ever names a deleted file:
|
||||
|
||||
- its image records (``generated_images``), found under every spelling the
|
||||
file is linked by. A record moves to the kept copy, unless its prompt
|
||||
already links the kept copy: then its saved-pair links join that record
|
||||
and it is deleted;
|
||||
- its favorite;
|
||||
- its image tags (``media_image_tags`` of its byte record), joined onto the
|
||||
kept copy's byte record as their union, in order;
|
||||
- its pending Needs review items, settled: only an item still ``pending`` at
|
||||
revision 0 with no decision row is removed (``review._SETTLE_ITEM``), so a
|
||||
decided item is never touched.
|
||||
|
||||
Every statement is an index search (see ``FILE_LINK_SQL``).
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
from .annotation_schema import BYTES_TABLE, FILES_TABLE
|
||||
from .annotations import _BUMP_INDEX
|
||||
from .models import normalize_image_path, utc_now_iso
|
||||
from .review import _BUMP_REVIEW_INDEX, _PENDING_OF_LOCATOR, _SETTLE_ITEM
|
||||
from .review_changes import notify_review_change
|
||||
from .workflow_documents import documents_of_rows, release_documents
|
||||
|
||||
LINKS_BY_FILE_SQL = "SELECT id, prompt_id FROM generated_images WHERE file_path = ?"
|
||||
LEGACY_LINKS_SQL = (
|
||||
"SELECT id, prompt_id FROM generated_images "
|
||||
"WHERE image_path = ? AND file_path IS NULL"
|
||||
)
|
||||
KEEPER_LINK_SQL = (
|
||||
"SELECT id FROM generated_images WHERE prompt_id = ? AND file_path = ?"
|
||||
)
|
||||
MOVE_PAIR_LINKS_SQL = (
|
||||
"UPDATE OR IGNORE image_prompt_pairs SET image_id = ? WHERE image_id = ?"
|
||||
)
|
||||
MOVE_LINK_SQL = (
|
||||
"UPDATE generated_images SET image_path = ?, filename = ?, file_path = ? "
|
||||
"WHERE id = ?"
|
||||
)
|
||||
DELETE_LINK_SQL = "DELETE FROM generated_images WHERE id = ?"
|
||||
FAVORITE_SQL = (
|
||||
"SELECT created_at FROM image_favorites WHERE root_key = ? AND relative_path = ?"
|
||||
)
|
||||
INSERT_FAVORITE_SQL = "INSERT OR IGNORE INTO image_favorites VALUES (?, ?, ?)"
|
||||
DELETE_FAVORITE_SQL = (
|
||||
"DELETE FROM image_favorites WHERE root_key = ? AND relative_path = ?"
|
||||
)
|
||||
COPY_BYTE_SQL = (
|
||||
f"SELECT b.id FROM {FILES_TABLE} f JOIN {BYTES_TABLE} b " # nosec B608
|
||||
"ON b.file_id=f.id WHERE f.locator=? AND b.sha256=?"
|
||||
)
|
||||
COPY_TAG_COUNT_SQL = (
|
||||
"SELECT COUNT(*) FROM media_image_tags WHERE byte_id=?" # nosec B608
|
||||
)
|
||||
PENDING_COUNT_SQL = (
|
||||
f"SELECT COUNT(*) FROM {FILES_TABLE} f " # nosec B608
|
||||
f"JOIN {BYTES_TABLE} b ON b.file_id=f.id "
|
||||
"JOIN media_review_items i ON i.byte_id=b.id "
|
||||
"WHERE f.locator=? AND +i.state='pending' AND i.revision=0 "
|
||||
"AND NOT EXISTS (SELECT 1 FROM media_review_decisions d WHERE d.item_id=i.id)"
|
||||
)
|
||||
FILE_LINK_SQL = (
|
||||
LINKS_BY_FILE_SQL,
|
||||
LEGACY_LINKS_SQL,
|
||||
KEEPER_LINK_SQL,
|
||||
MOVE_PAIR_LINKS_SQL,
|
||||
MOVE_LINK_SQL,
|
||||
DELETE_LINK_SQL,
|
||||
FAVORITE_SQL,
|
||||
DELETE_FAVORITE_SQL,
|
||||
COPY_BYTE_SQL,
|
||||
COPY_TAG_COUNT_SQL,
|
||||
PENDING_COUNT_SQL,
|
||||
)
|
||||
|
||||
|
||||
def _file_link_rows(conn, spelled):
|
||||
"""``(id, prompt_id)`` of the image records linking one spelling."""
|
||||
rows = list(conn.execute(LINKS_BY_FILE_SQL, (normalize_image_path(spelled),)))
|
||||
rows += list(conn.execute(LEGACY_LINKS_SQL, (spelled,)))
|
||||
return [tuple(row) for row in rows]
|
||||
|
||||
|
||||
def _copy_byte(conn, locator, sha256):
|
||||
row = conn.execute(COPY_BYTE_SQL, (locator, sha256)).fetchone()
|
||||
return row[0] if row else None
|
||||
|
||||
|
||||
class DuplicateCopyOperations:
|
||||
"""PromptDatabase mixin: the records of duplicate copies."""
|
||||
|
||||
def count_copy_records(self, copies) -> dict:
|
||||
"""``{path: {links, tags, review}}`` for each copy (one read snapshot).
|
||||
|
||||
``copies`` maps each path to ``{spellings, locator, sha256}``:
|
||||
``links`` counts the image records under any spelling, ``tags`` the
|
||||
image tags of its bytes and ``review`` its pending, undecided Needs
|
||||
review items.
|
||||
"""
|
||||
from .maintenance import _read_snapshot
|
||||
|
||||
def read(conn):
|
||||
found = {}
|
||||
for path, copy in copies.items():
|
||||
ids = {
|
||||
row[0]
|
||||
for spelled in copy["spellings"]
|
||||
for row in _file_link_rows(conn, spelled)
|
||||
}
|
||||
byte_id = _copy_byte(conn, copy["locator"], copy["sha256"])
|
||||
tags = (
|
||||
conn.execute(COPY_TAG_COUNT_SQL, (byte_id,)).fetchone()[0]
|
||||
if byte_id
|
||||
else 0
|
||||
)
|
||||
review = conn.execute(PENDING_COUNT_SQL, (copy["locator"],))
|
||||
found[path] = {
|
||||
"links": len(ids),
|
||||
"tags": tags,
|
||||
"review": review.fetchone()[0],
|
||||
}
|
||||
return found
|
||||
|
||||
return _read_snapshot(self.model.db_path, read)
|
||||
|
||||
def move_copy_records(self, move) -> dict:
|
||||
"""Move one removed copy's records to its kept copy (one transaction).
|
||||
|
||||
``move``: ``removed`` and ``kept`` (the spellings each file is linked
|
||||
by), ``targets`` (``{removed spelling: kept spelling}``, the spelling
|
||||
a moved record takes; else ``kept_path``), ``favorites`` (``(root key,
|
||||
removed key, kept root key, kept key)`` per listed root holding the
|
||||
file), ``removed_locator``,
|
||||
``kept_locator``, ``sha256`` and the kept copy's ``kept_size`` and
|
||||
``kept_mtime_ns``. Returns ``{moved, merged, favorites, tags,
|
||||
review}``.
|
||||
"""
|
||||
counts = dict.fromkeys(("moved", "merged", "favorites", "tags", "review"), 0)
|
||||
kept_keys = [normalize_image_path(path) for path in move["kept"]]
|
||||
with self._pair_transaction() as conn:
|
||||
for spelled in dict.fromkeys(move["removed"]):
|
||||
target = move["targets"].get(spelled, move["kept_path"])
|
||||
keys = list(dict.fromkeys(kept_keys + [normalize_image_path(target)]))
|
||||
for image_id, prompt_id in _file_link_rows(conn, spelled):
|
||||
existing = None
|
||||
for key in keys:
|
||||
existing = conn.execute(
|
||||
KEEPER_LINK_SQL, (prompt_id, key)
|
||||
).fetchone()
|
||||
if existing is not None:
|
||||
break
|
||||
if existing is not None:
|
||||
conn.execute(MOVE_PAIR_LINKS_SQL, (existing[0], image_id))
|
||||
documents = documents_of_rows(conn, "id", [image_id])
|
||||
conn.execute(DELETE_LINK_SQL, (image_id,))
|
||||
release_documents(conn, documents)
|
||||
counts["merged"] += 1
|
||||
continue
|
||||
conn.execute(
|
||||
MOVE_LINK_SQL,
|
||||
(
|
||||
target,
|
||||
os.path.basename(target),
|
||||
normalize_image_path(target),
|
||||
image_id,
|
||||
),
|
||||
)
|
||||
counts["moved"] += 1
|
||||
for favorite in move.get("favorites", ()):
|
||||
counts["favorites"] += _move_favorite(conn, *favorite)
|
||||
counts["tags"] = self._move_copy_tags(conn, move)
|
||||
for (item_id,) in conn.execute(
|
||||
_PENDING_OF_LOCATOR, (move["removed_locator"],)
|
||||
).fetchall():
|
||||
counts["review"] += conn.execute(_SETTLE_ITEM, (item_id,)).rowcount
|
||||
if counts["review"]:
|
||||
conn.execute(_BUMP_REVIEW_INDEX)
|
||||
if counts["review"]:
|
||||
notify_review_change(self.model.db_path)
|
||||
return counts
|
||||
|
||||
def _move_copy_tags(self, conn, move):
|
||||
"""Join the removed copy's image tags onto the kept copy's; the count added."""
|
||||
removed = _copy_byte(conn, move["removed_locator"], move["sha256"])
|
||||
if removed is None:
|
||||
return 0
|
||||
tags = self._byte_tags(conn, removed)
|
||||
if not tags:
|
||||
return 0
|
||||
now = utc_now_iso()
|
||||
file_id = self._annotation_owner(conn, move["kept_locator"], now)
|
||||
byte_id, _revision, created, moved = self._annotation_byte(
|
||||
conn, file_id, move["sha256"], move["kept_size"], move["kept_mtime_ns"], now
|
||||
)
|
||||
current = self._byte_tags(conn, byte_id)
|
||||
union = current + [tag for tag in tags if tag not in current]
|
||||
if union != current:
|
||||
self._write_image_tags(conn, byte_id, union)
|
||||
elif created or moved:
|
||||
conn.execute(_BUMP_INDEX)
|
||||
return len(union) - len(current)
|
||||
|
||||
|
||||
def _move_favorite(conn, root_key, removed_key, kept_root, kept_key):
|
||||
row = conn.execute(FAVORITE_SQL, (root_key, removed_key)).fetchone()
|
||||
if row is None:
|
||||
return 0
|
||||
conn.execute(INSERT_FAVORITE_SQL, (kept_root, kept_key, row[0]))
|
||||
conn.execute(DELETE_FAVORITE_SQL, (root_key, removed_key))
|
||||
return 1
|
||||
@@ -196,14 +196,3 @@ 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)
|
||||
|
||||
+6
-32
@@ -189,33 +189,6 @@ def _image_owner(image_path, file_path, prefixes):
|
||||
return owner
|
||||
|
||||
|
||||
def readable_roots(roots):
|
||||
"""Whether each source can be read now (an unmounted drive cannot)."""
|
||||
return [_readable(root) for root in roots]
|
||||
|
||||
|
||||
def link_states(rows, roots, readable):
|
||||
"""``(row, owner, state)`` for each image row ``(id, image_path, filename,
|
||||
file_path)``: ``owner`` is the index of the source it lies under (None
|
||||
outside every source) and ``state`` is ``missing``, ``present``,
|
||||
``unavailable`` (its source cannot be read, so the file is not looked at)
|
||||
or ``outside``. ``readable`` is ``readable_roots(roots)``. The health
|
||||
check's ``missing_files`` and the maintenance missing-link plan both
|
||||
judge links this way, so their counts agree.
|
||||
"""
|
||||
prefixes = [_root_prefix(root) for root in roots]
|
||||
for row in rows:
|
||||
owner = _image_owner(row[1], row[3], prefixes)
|
||||
if owner is None:
|
||||
yield row, None, "outside"
|
||||
elif not readable[owner]:
|
||||
yield row, owner, "unavailable"
|
||||
elif os.path.exists(row[1]):
|
||||
yield row, owner, "present"
|
||||
else:
|
||||
yield row, owner, "missing"
|
||||
|
||||
|
||||
def _unavailable_line(unavailable, *, partial):
|
||||
noun = "source is" if unavailable == 1 else "sources are"
|
||||
effect = "its files were" if unavailable == 1 else "their files were"
|
||||
@@ -251,7 +224,7 @@ def _missing_files(conn, roots):
|
||||
|
||||
def _check_files(rows, roots):
|
||||
"""The ``missing_files`` result of the image *rows*, from the file system."""
|
||||
readable = readable_roots(roots)
|
||||
readable = [_readable(root) for root in roots]
|
||||
sources = [
|
||||
{
|
||||
"name": _source_name(root, index),
|
||||
@@ -266,12 +239,13 @@ def _check_files(rows, roots):
|
||||
None, [_unavailable_line(unavailable, partial=False)], status="unavailable"
|
||||
)
|
||||
return {**result, "partial": False, "sources": sources}
|
||||
prefixes = [_root_prefix(root) for root in roots]
|
||||
outside, missing = 0, []
|
||||
for row, owner, state in link_states(rows, roots, readable):
|
||||
if state == "outside":
|
||||
for image_id, image_path, filename, file_path in rows:
|
||||
owner = _image_owner(image_path, file_path, prefixes)
|
||||
if owner is None:
|
||||
outside += 1
|
||||
elif state == "missing":
|
||||
image_id, image_path, filename, _file_path = row
|
||||
elif readable[owner] and not os.path.exists(image_path):
|
||||
sources[owner]["count"] += 1
|
||||
if len(missing) < MAX_DETAIL:
|
||||
name = filename or os.path.basename(image_path)
|
||||
|
||||
@@ -1,805 +0,0 @@
|
||||
"""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 retires them
|
||||
(``PromptModel.retire_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
|
||||
# Retired connections still in use; awaited once the locks are free.
|
||||
self.retiring: list = []
|
||||
|
||||
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). On failure
|
||||
every detached connection is retired; those still in use are added to
|
||||
``move.retiring`` for the caller to await once the locks are released.
|
||||
"""
|
||||
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
|
||||
# Retire, never wait, here: a connection still in use may be
|
||||
# waiting on the very locks this thread holds.
|
||||
try:
|
||||
undone = move.unswitch(plan.db_path, location.db_path)
|
||||
except Exception:
|
||||
logger.exception("Library move: undoing the switch failed")
|
||||
else:
|
||||
_retire(move, undone)
|
||||
# The old file's connections a partial switch detached: the
|
||||
# models reopen it on next use, so close them now, not at GC.
|
||||
_retire(move, detached)
|
||||
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):
|
||||
"""Copy, then switch under both locks; retired connections are awaited.
|
||||
|
||||
A connection that was inside a statement when it was retired is closed
|
||||
by its own thread once the statement ends, which may need the locks
|
||||
released first; the old files are renamed (and a failed copy removed)
|
||||
only after that, since Windows refuses to rename or delete an open file.
|
||||
"""
|
||||
try:
|
||||
return _copy_and_switch_locked(move, location, attempts)
|
||||
finally:
|
||||
try:
|
||||
PromptModel.await_detached(move.retiring, logger)
|
||||
except Exception:
|
||||
logger.exception("Library move: waiting for old connections failed")
|
||||
|
||||
|
||||
def _retire(move: _Move, connections: list) -> None:
|
||||
"""Retire ``connections``, adding those still in use to ``move.retiring``.
|
||||
|
||||
Never raises: a switch that completed must not be reported as failed
|
||||
because one of the old connections could not be retired (logged; the
|
||||
rename then reports the old files as still in use).
|
||||
"""
|
||||
try:
|
||||
move.retiring.extend(PromptModel.retire_detached(connections, logger))
|
||||
except Exception:
|
||||
logger.exception("Library move: retiring the old connections failed")
|
||||
|
||||
|
||||
def _copy_and_switch_locked(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 is interrupted and cannot commit, so it fails (and
|
||||
# runs again on the new library) instead of landing in the
|
||||
# old file. Its own thread closes the connection.
|
||||
_retire(move, detached)
|
||||
finally:
|
||||
if source.in_transaction:
|
||||
source.execute("ROLLBACK")
|
||||
return files, counts
|
||||
finally:
|
||||
source.close()
|
||||
raise MoveFailed("library-busy")
|
||||
@@ -1,740 +0,0 @@
|
||||
"""Reviewed maintenance plans: read-only previews and their exact apply (#249 S6).
|
||||
|
||||
Four operations, each the native owner of a legacy ``/maintenance`` one and
|
||||
built on the same queries, so a preview counts what the legacy cleanup would
|
||||
change:
|
||||
|
||||
- ``duplicate_prompts``: prompts whose text is equal ignoring case and outer
|
||||
spaces (``DUPLICATE_PROMPTS_SQL``, legacy ``cleanup_duplicates``); each
|
||||
group keeps its oldest prompt, which takes over the others' images and
|
||||
metadata. A group with a member in a saved pair is kept whole.
|
||||
- ``optimize``: VACUUM (legacy ``vacuum``).
|
||||
- ``unused_components``: prompts nothing references (``ORPHANED_PROMPTS_SQL``,
|
||||
legacy ``prune_orphaned_prompts``, with the #239 protections).
|
||||
- ``missing_links``: image rows whose file is gone (legacy
|
||||
``cleanup_orphaned_images``), counted only under a source that can be read
|
||||
now, as the health check's ``missing_files`` judges them
|
||||
(``health.link_states``). It deletes database rows only, never a file.
|
||||
|
||||
``preview_maintenance`` reads one snapshot on its own ``mode=ro`` connection
|
||||
(``health.readonly_uri``) and closes it in ``finally``; files are looked at
|
||||
after the snapshot ended. It returns a ``MaintenanceDraft``: the public plan
|
||||
fields, the targets an apply needs and the fingerprint of exactly what was
|
||||
reviewed (ids, text hashes, metadata, tags, Library state and pair refs of
|
||||
every prompt involved, or each link's row and source). The fingerprint never
|
||||
leaves the server.
|
||||
|
||||
``apply_maintenance`` recomputes that fingerprint on the write connection
|
||||
inside ``BEGIN IMMEDIATE`` and raises ``PlanChanged`` before writing anything
|
||||
when it differs; otherwise it applies the reviewed targets in batches of
|
||||
``ID_BATCH`` in that one transaction. VACUUM cannot run in a transaction, so
|
||||
``optimize`` binds to its database only (the caller checks that).
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sqlite3
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
from .execution_contexts import LIVE_HINT_CTES
|
||||
from .health import BATCH_ROWS, link_states, readable_roots, readonly_uri
|
||||
from .pairs import PairDependencyError
|
||||
from .workflow_documents import documents_of_rows, release_documents
|
||||
|
||||
MAINTENANCE_OPERATIONS = (
|
||||
"duplicate_prompts",
|
||||
"optimize",
|
||||
"unused_components",
|
||||
"missing_links",
|
||||
# Applied as its own batched job (database/workflow_storage.py).
|
||||
"workflow_storage",
|
||||
)
|
||||
MAX_DETAILS = 50 # detail rows a plan lists
|
||||
MAX_RESULT_ROWS = 200 # changed records an apply describes (the job lists 200)
|
||||
EXCERPT_CHARS = 80
|
||||
# Ids bound per statement: the pair-reference read binds each id twice, which
|
||||
# stays under SQLite's historical 999-variable limit.
|
||||
ID_BATCH = 400
|
||||
|
||||
# Why a prompt no image links to is kept (#239): predicates on ``p``, in the
|
||||
# order the plan names them. ORPHANED_PROMPTS_SQL is built from them, so the
|
||||
# reasons a plan lists are exactly the rules the prune applies.
|
||||
PROTECTIONS = (
|
||||
(
|
||||
"EXISTS (SELECT 1 FROM prompt_pairs pp "
|
||||
"WHERE pp.positive_prompt_id = p.id OR pp.negative_prompt_id = p.id)",
|
||||
"in saved pairs",
|
||||
"in a saved pair",
|
||||
),
|
||||
(
|
||||
"EXISTS (SELECT 1 FROM prompt_tags pt JOIN tags t ON pt.tag_id = t.id "
|
||||
"WHERE pt.prompt_id = p.id AND t.name = '__protected__')",
|
||||
"tagged __protected__",
|
||||
"tagged __protected__",
|
||||
),
|
||||
(
|
||||
"EXISTS (SELECT 1 FROM prompt_library_state ls WHERE ls.prompt_id = p.id)",
|
||||
"with Library state (a title or an explicit parent)",
|
||||
"with Library state (a title or an explicit parent)",
|
||||
),
|
||||
(
|
||||
"p.id IN (SELECT prompt_id FROM live_hint_prompts)",
|
||||
"used by recorded runs",
|
||||
"used by a recorded run",
|
||||
),
|
||||
)
|
||||
|
||||
# Exempt references: images, pairs, the __protected__ tag, Library parent
|
||||
# state, and live recorded-run hints (computed once, see LIVE_HINT_CTES).
|
||||
ORPHANED_PROMPTS_SQL = (
|
||||
"WITH "
|
||||
+ LIVE_HINT_CTES
|
||||
+ " SELECT p.id FROM prompts p"
|
||||
+ " LEFT JOIN generated_images gi ON p.id = gi.prompt_id"
|
||||
+ " WHERE gi.prompt_id IS NULL"
|
||||
+ "".join(f" AND NOT ({predicate})" for predicate, _many, _one in PROTECTIONS)
|
||||
)
|
||||
# The prompts no image links to, with which protections keep each one.
|
||||
_UNLINKED_PROTECTIONS_SQL = (
|
||||
"WITH "
|
||||
+ LIVE_HINT_CTES
|
||||
+ " SELECT "
|
||||
+ ", ".join(
|
||||
f"COALESCE(SUM({predicate}), 0)" for predicate, _many, _one in PROTECTIONS
|
||||
)
|
||||
+ ", COALESCE(SUM("
|
||||
+ " OR ".join(f"({predicate})" for predicate, _many, _one in PROTECTIONS)
|
||||
+ "), 0) FROM prompts p WHERE NOT EXISTS "
|
||||
"(SELECT 1 FROM generated_images gi WHERE gi.prompt_id = p.id)"
|
||||
)
|
||||
|
||||
# Duplicate prompt groups, as the legacy cleanup finds them (case-insensitive
|
||||
# text; GROUP_CONCAT order is sorted by created_at afterwards).
|
||||
DUPLICATE_PROMPTS_SQL = """
|
||||
SELECT LOWER(TRIM(text)) as normalized_text, COUNT(*) as count,
|
||||
GROUP_CONCAT(id) as ids,
|
||||
GROUP_CONCAT(created_at) as created_dates
|
||||
FROM prompts
|
||||
GROUP BY LOWER(TRIM(text))
|
||||
HAVING COUNT(*) > 1
|
||||
"""
|
||||
|
||||
_LINK_ROWS_SQL = (
|
||||
"SELECT id, image_path, filename, file_path FROM generated_images ORDER BY id"
|
||||
)
|
||||
|
||||
# What a prompt's fingerprint covers, per batch of ids.
|
||||
_PROMPT_STATE_SQL = (
|
||||
"SELECT id, hash, text, category, rating, notes, created_at, updated_at "
|
||||
"FROM prompts WHERE id IN ({marks}) ORDER BY id",
|
||||
"SELECT prompt_id, tag_id FROM prompt_tags WHERE prompt_id IN ({marks}) "
|
||||
"ORDER BY prompt_id, tag_id",
|
||||
"SELECT * FROM prompt_library_state WHERE prompt_id IN ({marks}) "
|
||||
"ORDER BY prompt_id",
|
||||
"SELECT id, positive_prompt_id, negative_prompt_id FROM prompt_pairs "
|
||||
"WHERE positive_prompt_id IN ({marks}) OR negative_prompt_id IN ({marks}) "
|
||||
"ORDER BY id",
|
||||
)
|
||||
|
||||
NO_FILES_DELETED = "No image file is deleted."
|
||||
LINKS_ONLY = "Database links only; no image file is deleted."
|
||||
UNUSED_DETAIL = "No image, saved pair, Library state or recorded run uses it"
|
||||
|
||||
|
||||
# Why a duplicate group is kept whole (``_duplicate_conflict``).
|
||||
CONFLICT_REASONS = {
|
||||
"collision": ("group links one image twice", "groups link one image twice"),
|
||||
"title": (
|
||||
"group has its own Library title",
|
||||
"groups have their own Library titles",
|
||||
),
|
||||
}
|
||||
CONFLICT_DETAILS = {
|
||||
"collision": "Kept: two copies link the same image",
|
||||
"title": "Kept: a copy has its own Library title",
|
||||
}
|
||||
OPTIMIZE_LOCK_NOTE = (
|
||||
"Runs alone: the database is locked while it runs, and on a large library "
|
||||
"this takes a while. Prompt runs counted and images captured meanwhile "
|
||||
"wait up to 30 seconds each; anything that waits longer is not recorded."
|
||||
)
|
||||
|
||||
|
||||
class PlanChanged(Exception):
|
||||
"""What a plan was reviewed on changed before it was applied."""
|
||||
|
||||
|
||||
class MergeConflict(Exception):
|
||||
"""A strict duplicate merge would drop an image link or a Library title."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class MaintenanceDraft:
|
||||
"""A preview: the public plan fields plus what only the server keeps."""
|
||||
|
||||
operation: str
|
||||
fingerprint: str
|
||||
targets: list
|
||||
scope: dict
|
||||
protected: dict
|
||||
media: dict
|
||||
details: list = field(default_factory=list)
|
||||
|
||||
def public(self):
|
||||
"""The plan fields a response may carry (no fingerprint, no targets)."""
|
||||
return {
|
||||
"operation": self.operation,
|
||||
"scope": dict(self.scope),
|
||||
"protected": {
|
||||
"count": self.protected["count"],
|
||||
"reasons": list(self.protected["reasons"]),
|
||||
},
|
||||
"media": dict(self.media),
|
||||
"details": [dict(row) for row in self.details],
|
||||
}
|
||||
|
||||
|
||||
def _plural(count, one, many=None):
|
||||
return f"{count:,} {one if count == 1 else (many or one + 's')}"
|
||||
|
||||
|
||||
def _excerpt(text):
|
||||
text = " ".join(str(text or "").split())
|
||||
if len(text) <= EXCERPT_CHARS:
|
||||
return text
|
||||
return text[: EXCERPT_CHARS - 1].rstrip() + "…"
|
||||
|
||||
|
||||
def _megabytes(size):
|
||||
return f"{size / 2**20:,.1f} MB"
|
||||
|
||||
|
||||
class _Fingerprint:
|
||||
"""A streaming sha256 over JSON-encoded parts."""
|
||||
|
||||
def __init__(self, operation):
|
||||
self._hash = hashlib.sha256()
|
||||
self.add(operation)
|
||||
|
||||
def add(self, part):
|
||||
encoded = json.dumps(part, separators=(",", ":"), default=str)
|
||||
self._hash.update(encoded.encode("utf-8") + b"\n")
|
||||
|
||||
def hexdigest(self):
|
||||
return "sha256:" + self._hash.hexdigest()
|
||||
|
||||
|
||||
def _batches(ids):
|
||||
ids = list(ids)
|
||||
for start in range(0, len(ids), ID_BATCH):
|
||||
yield ids[start : start + ID_BATCH]
|
||||
|
||||
|
||||
def _marks(batch):
|
||||
return ",".join("?" for _ in batch)
|
||||
|
||||
|
||||
def _add_prompt_state(fingerprint, conn, ids):
|
||||
"""Fold every reviewed fact about the prompts ``ids`` into ``fingerprint``."""
|
||||
for batch in _batches(sorted(ids)):
|
||||
marks = _marks(batch)
|
||||
for sql in _PROMPT_STATE_SQL:
|
||||
params = batch * sql.count("{marks}")
|
||||
for row in conn.execute(sql.format(marks=marks), params):
|
||||
fingerprint.add(list(row))
|
||||
|
||||
|
||||
def _texts(conn, ids, limit=MAX_DETAILS):
|
||||
"""``{id: text}`` of the first ``limit`` prompts, for detail rows."""
|
||||
ids = list(ids)[:limit]
|
||||
if not ids:
|
||||
return {}
|
||||
rows = conn.execute(
|
||||
f"SELECT id, text FROM prompts WHERE id IN ({_marks(ids)})", ids
|
||||
)
|
||||
return {row[0]: row[1] for row in rows}
|
||||
|
||||
|
||||
def _orphaned_ids(conn):
|
||||
return sorted(row[0] for row in conn.execute(ORPHANED_PROMPTS_SQL))
|
||||
|
||||
|
||||
def _delete_ids(conn, table, ids):
|
||||
"""Delete ``table`` rows by id in batches; returns the rows removed.
|
||||
|
||||
The workflow documents only those rows referenced go with them (a
|
||||
prompt's images go by cascade), in the same transaction.
|
||||
"""
|
||||
if table not in ("prompts", "generated_images"):
|
||||
raise ValueError(f"Unsupported table {table!r}")
|
||||
removed = 0
|
||||
for batch in _batches(ids):
|
||||
documents = documents_of_rows(
|
||||
conn, "id" if table == "generated_images" else "prompt_id", batch
|
||||
)
|
||||
cursor = conn.execute(
|
||||
f"DELETE FROM {table} WHERE id IN ({_marks(batch)})", # nosec B608
|
||||
batch,
|
||||
)
|
||||
removed += cursor.rowcount
|
||||
release_documents(conn, documents)
|
||||
return removed
|
||||
|
||||
|
||||
def _read_snapshot(db_path, read):
|
||||
"""``read(conn)`` on one ``mode=ro`` snapshot; the connection always closes."""
|
||||
conn = sqlite3.connect(readonly_uri(db_path), uri=True)
|
||||
try:
|
||||
conn.row_factory = sqlite3.Row
|
||||
conn.execute("BEGIN") # one read snapshot for the whole plan
|
||||
return read(conn)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _relative(image_path, root):
|
||||
try:
|
||||
return Path(os.path.relpath(image_path, root)).as_posix()
|
||||
except ValueError: # another drive (Windows)
|
||||
return os.path.basename(image_path)
|
||||
|
||||
|
||||
def _source_name(root, index):
|
||||
return os.path.basename(str(root).rstrip("\\/")) or f"Source {index + 1}"
|
||||
|
||||
|
||||
class MaintenanceOperations:
|
||||
"""Mixin for PromptDatabase: maintenance previews and their apply."""
|
||||
|
||||
# ── shared with the legacy cleanup ───────────────────────────────
|
||||
|
||||
def _duplicate_groups(self, conn):
|
||||
"""``(groups, paired)``: duplicate id groups, oldest first.
|
||||
|
||||
``paired`` are the groups with a member in a saved pair, which the
|
||||
cleanup keeps whole.
|
||||
"""
|
||||
groups, paired = [], []
|
||||
for duplicate in conn.execute(DUPLICATE_PROMPTS_SQL).fetchall():
|
||||
ids = duplicate["ids"].split(",")
|
||||
created_dates = duplicate["created_dates"].split(",")
|
||||
# Sort IDs by created_at date to keep the oldest
|
||||
id_date_pairs = list(zip(ids, created_dates))
|
||||
id_date_pairs.sort(key=lambda x: x[1])
|
||||
sorted_ids = [int(pair[0]) for pair in id_date_pairs]
|
||||
try:
|
||||
self._ensure_prompts_unpaired(conn, sorted_ids)
|
||||
except PairDependencyError:
|
||||
paired.append(sorted_ids)
|
||||
continue
|
||||
groups.append(sorted_ids)
|
||||
return groups, paired
|
||||
|
||||
@staticmethod
|
||||
def _duplicate_conflict(conn, sorted_ids):
|
||||
"""Why merging the group would lose something, or None.
|
||||
|
||||
``collision``: two members link the same file, so moving the images
|
||||
to the oldest prompt would hit ``UNIQUE(prompt_id, file_path)``.
|
||||
``title``: a member that would be deleted has its own Library title.
|
||||
"""
|
||||
marks = _marks(sorted_ids)
|
||||
collision = conn.execute(
|
||||
"SELECT 1 FROM generated_images WHERE prompt_id IN " # nosec B608
|
||||
f"({marks}) AND file_path IS NOT NULL GROUP BY file_path "
|
||||
"HAVING COUNT(DISTINCT prompt_id) > 1 LIMIT 1",
|
||||
sorted_ids,
|
||||
).fetchone()
|
||||
if collision:
|
||||
return "collision"
|
||||
duplicates = sorted_ids[1:]
|
||||
titled = conn.execute(
|
||||
"SELECT 1 FROM prompt_library_state WHERE prompt_id IN " # nosec B608
|
||||
f"({_marks(duplicates)}) AND trim(coalesce(title, '')) != '' LIMIT 1",
|
||||
duplicates,
|
||||
).fetchone()
|
||||
return "title" if titled else None
|
||||
|
||||
def _merge_duplicate_group(self, conn, sorted_ids, strict=False):
|
||||
"""Keep the oldest prompt of a group; returns (removed, images moved).
|
||||
|
||||
``strict`` (the reviewed apply): a merge that would drop an image
|
||||
link or a Library title raises MergeConflict, and every helper error
|
||||
propagates, so the transaction rolls back. The legacy cleanup keeps
|
||||
its old, lenient behaviour.
|
||||
"""
|
||||
primary_id, duplicate_ids = sorted_ids[0], sorted_ids[1:]
|
||||
if strict:
|
||||
conflict = self._duplicate_conflict(conn, sorted_ids)
|
||||
if conflict:
|
||||
raise MergeConflict(conflict)
|
||||
self.logger.debug(
|
||||
f"Merging duplicates: keeping {primary_id}, removing {duplicate_ids}"
|
||||
)
|
||||
primary = conn.execute(
|
||||
"SELECT * FROM prompts WHERE id = ?", (primary_id,)
|
||||
).fetchone()
|
||||
if not primary:
|
||||
return 0, 0
|
||||
merged_metadata = self._merge_duplicate_metadata(
|
||||
conn, primary_id, duplicate_ids, strict=strict
|
||||
)
|
||||
moved = self._transfer_images_to_primary(
|
||||
conn, primary_id, duplicate_ids, strict=strict
|
||||
)
|
||||
if merged_metadata:
|
||||
self._update_primary_with_merged_metadata(
|
||||
conn, primary_id, merged_metadata, strict=strict
|
||||
)
|
||||
removed = 0
|
||||
# Images the lenient merge could not move go with their prompt.
|
||||
documents = documents_of_rows(conn, "prompt_id", duplicate_ids)
|
||||
for duplicate_id in duplicate_ids:
|
||||
conn.execute("DELETE FROM prompts WHERE id = ?", (duplicate_id,))
|
||||
removed += 1
|
||||
self.logger.debug(f"Removed duplicate prompt {duplicate_id}")
|
||||
release_documents(conn, documents)
|
||||
return removed, moved
|
||||
|
||||
@staticmethod
|
||||
def _missing_link_ids(rows, roots):
|
||||
"""Ids of the ``rows`` whose file is gone under a readable source."""
|
||||
readable = readable_roots(roots)
|
||||
return [
|
||||
row[0]
|
||||
for row, _owner, state in link_states(rows, roots, readable)
|
||||
if state == "missing"
|
||||
]
|
||||
|
||||
# ── previews ──────────────────────────────────────────────────────
|
||||
|
||||
def preview_maintenance(self, operation, roots=()) -> MaintenanceDraft:
|
||||
"""The reviewed plan of ``operation``; writes nothing (blocking).
|
||||
|
||||
``roots`` are the sources missing links are judged under (the health
|
||||
check's: the effective sources plus configured ones that are gone).
|
||||
"""
|
||||
if operation not in MAINTENANCE_OPERATIONS:
|
||||
raise ValueError(f"Unknown maintenance operation {operation!r}")
|
||||
roots = [str(root) for root in roots]
|
||||
if operation == "optimize":
|
||||
return self._plan_optimize()
|
||||
if operation == "workflow_storage":
|
||||
return self._plan_workflow_storage()
|
||||
if operation == "missing_links":
|
||||
rows = _read_snapshot(self.model.db_path, self._link_rows)
|
||||
return self._plan_missing_links(rows, roots)
|
||||
reader = (
|
||||
self._plan_duplicate_prompts
|
||||
if operation == "duplicate_prompts"
|
||||
else self._plan_unused_components
|
||||
)
|
||||
return _read_snapshot(self.model.db_path, reader)
|
||||
|
||||
@staticmethod
|
||||
def _link_rows(conn):
|
||||
rows = []
|
||||
cursor = conn.execute(_LINK_ROWS_SQL)
|
||||
while batch := cursor.fetchmany(BATCH_ROWS):
|
||||
rows.extend(tuple(row) for row in batch)
|
||||
return rows
|
||||
|
||||
def _duplicates_fingerprint(self, conn):
|
||||
"""``(groups, paired, conflicted, fingerprint)`` of the duplicate merge.
|
||||
|
||||
``conflicted`` holds ``(ids, reason)`` of the groups a merge would
|
||||
lose an image link or a Library title in (``_duplicate_conflict``);
|
||||
the plan protects them.
|
||||
"""
|
||||
candidates, paired = self._duplicate_groups(conn)
|
||||
groups, conflicted = [], []
|
||||
for ids in candidates:
|
||||
reason = self._duplicate_conflict(conn, ids)
|
||||
if reason:
|
||||
conflicted.append((ids, reason))
|
||||
else:
|
||||
groups.append(ids)
|
||||
fingerprint = _Fingerprint("duplicate_prompts")
|
||||
fingerprint.add({"groups": groups, "paired": paired, "conflicted": conflicted})
|
||||
_add_prompt_state(fingerprint, conn, [i for group in groups for i in group])
|
||||
return groups, paired, conflicted, fingerprint.hexdigest()
|
||||
|
||||
def _plan_duplicate_prompts(self, conn):
|
||||
groups, paired, conflicted, fingerprint = self._duplicates_fingerprint(conn)
|
||||
removed = sum(len(group) - 1 for group in groups)
|
||||
held = [ids for ids, _reason in conflicted]
|
||||
texts = _texts(conn, [group[0] for group in (groups + paired + held)])
|
||||
details = (
|
||||
[
|
||||
{
|
||||
"label": _excerpt(texts.get(group[0])),
|
||||
"detail": f"Keeps the oldest of {len(group)} copies",
|
||||
}
|
||||
for group in groups
|
||||
]
|
||||
+ [
|
||||
{
|
||||
"label": _excerpt(texts.get(group[0])),
|
||||
"detail": "Kept: belongs to a saved pair",
|
||||
}
|
||||
for group in paired
|
||||
]
|
||||
+ [
|
||||
{
|
||||
"label": _excerpt(texts.get(ids[0])),
|
||||
"detail": CONFLICT_DETAILS[reason],
|
||||
}
|
||||
for ids, reason in conflicted
|
||||
]
|
||||
)
|
||||
if groups:
|
||||
label = (
|
||||
f"{_plural(removed, 'duplicate prompt')} in "
|
||||
f"{_plural(len(groups), 'group')}"
|
||||
)
|
||||
else:
|
||||
label = "No duplicate prompt text"
|
||||
reasons = []
|
||||
if len(paired) == 1:
|
||||
reasons.append("1 group belongs to a saved pair")
|
||||
elif paired:
|
||||
reasons.append(f"{len(paired):,} groups belong to saved pairs")
|
||||
for reason, (one, many) in CONFLICT_REASONS.items():
|
||||
count = sum(1 for _ids, why in conflicted if why == reason)
|
||||
if count:
|
||||
reasons.append(f"{count:,} {one if count == 1 else many}")
|
||||
return MaintenanceDraft(
|
||||
operation="duplicate_prompts",
|
||||
fingerprint=fingerprint,
|
||||
targets=groups,
|
||||
scope={"count": removed, "label": label},
|
||||
protected={
|
||||
"count": sum(len(group) for group in paired + held),
|
||||
"reasons": reasons,
|
||||
},
|
||||
media={"deletes_files": False, "label": NO_FILES_DELETED},
|
||||
details=details[:MAX_DETAILS],
|
||||
)
|
||||
|
||||
def _unused_fingerprint(self, conn):
|
||||
ids = _orphaned_ids(conn)
|
||||
fingerprint = _Fingerprint("unused_components")
|
||||
fingerprint.add(ids)
|
||||
_add_prompt_state(fingerprint, conn, ids)
|
||||
return ids, fingerprint.hexdigest()
|
||||
|
||||
def _plan_unused_components(self, conn):
|
||||
ids, fingerprint = self._unused_fingerprint(conn)
|
||||
*by_reason, kept = conn.execute(_UNLINKED_PROTECTIONS_SQL).fetchone()
|
||||
reasons = [
|
||||
f"{count:,} {many if count != 1 else one}"
|
||||
for count, (_predicate, many, one) in zip(by_reason, PROTECTIONS)
|
||||
if count
|
||||
]
|
||||
texts = _texts(conn, ids)
|
||||
count = len(ids)
|
||||
return MaintenanceDraft(
|
||||
operation="unused_components",
|
||||
fingerprint=fingerprint,
|
||||
targets=ids,
|
||||
scope={
|
||||
"count": count,
|
||||
"label": (
|
||||
_plural(count, "unused prompt component")
|
||||
if count
|
||||
else "No unused prompt components"
|
||||
),
|
||||
},
|
||||
protected={"count": kept, "reasons": reasons},
|
||||
media={"deletes_files": False, "label": NO_FILES_DELETED},
|
||||
details=[
|
||||
{"label": _excerpt(texts.get(prompt_id)), "detail": UNUSED_DETAIL}
|
||||
for prompt_id in ids[:MAX_DETAILS]
|
||||
],
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _links_fingerprint(candidates):
|
||||
fingerprint = _Fingerprint("missing_links")
|
||||
for row, owner in candidates:
|
||||
fingerprint.add([row[0], row[1], row[3], owner])
|
||||
return fingerprint.hexdigest()
|
||||
|
||||
def _plan_missing_links(self, rows, roots):
|
||||
readable = readable_roots(roots)
|
||||
candidates, unavailable, outside = [], 0, 0
|
||||
for row, owner, state in link_states(rows, roots, readable):
|
||||
if state == "missing":
|
||||
candidates.append((row, owner))
|
||||
elif state == "unavailable":
|
||||
unavailable += 1
|
||||
elif state == "outside":
|
||||
outside += 1
|
||||
reasons = []
|
||||
if unavailable:
|
||||
reasons.append(
|
||||
_plural(
|
||||
unavailable,
|
||||
"link under an unavailable source",
|
||||
"links under unavailable sources",
|
||||
)
|
||||
)
|
||||
if outside:
|
||||
reasons.append(
|
||||
_plural(
|
||||
outside,
|
||||
"link outside the image sources",
|
||||
"links outside the image sources",
|
||||
)
|
||||
)
|
||||
count = len(candidates)
|
||||
return MaintenanceDraft(
|
||||
operation="missing_links",
|
||||
fingerprint=self._links_fingerprint(candidates),
|
||||
targets=[row[0] for row, _owner in candidates],
|
||||
scope={
|
||||
"count": count,
|
||||
"label": (
|
||||
_plural(count, "link to a missing file", "links to missing files")
|
||||
if count
|
||||
else "No links to missing files"
|
||||
),
|
||||
},
|
||||
protected={"count": unavailable + outside, "reasons": reasons},
|
||||
media={"deletes_files": False, "label": LINKS_ONLY},
|
||||
details=[
|
||||
{
|
||||
"label": _relative(row[1], roots[owner]),
|
||||
"detail": (
|
||||
f"{_source_name(roots[owner], owner)} · the file is missing"
|
||||
),
|
||||
}
|
||||
for row, owner in candidates[:MAX_DETAILS]
|
||||
],
|
||||
)
|
||||
|
||||
def _plan_optimize(self):
|
||||
def read(conn):
|
||||
return [
|
||||
conn.execute(f"PRAGMA {name}").fetchone()[0]
|
||||
for name in ("page_size", "page_count", "freelist_count")
|
||||
]
|
||||
|
||||
page_size, page_count, free_pages = _read_snapshot(self.model.db_path, read)
|
||||
db_path = self.model.db_path
|
||||
wal = db_path + "-wal"
|
||||
size = page_size * page_count
|
||||
wal_size = os.path.getsize(wal) if os.path.exists(wal) else 0
|
||||
free = shutil.disk_usage(os.path.dirname(os.path.abspath(db_path))).free
|
||||
# VACUUM writes the rebuilt database through the WAL beside it.
|
||||
return MaintenanceDraft(
|
||||
operation="optimize",
|
||||
fingerprint=_Fingerprint("optimize").hexdigest(),
|
||||
targets=[],
|
||||
scope={"count": 1, "label": "The whole database"},
|
||||
protected={"count": 0, "reasons": []},
|
||||
media={"deletes_files": False, "label": NO_FILES_DELETED},
|
||||
details=[
|
||||
{
|
||||
"label": "Database file",
|
||||
"detail": f"{_megabytes(size)}, plus {_megabytes(wal_size)} "
|
||||
"not yet checkpointed",
|
||||
},
|
||||
{
|
||||
"label": "Free pages",
|
||||
"detail": f"About {_megabytes(free_pages * page_size)} "
|
||||
"can be reclaimed",
|
||||
},
|
||||
{
|
||||
"label": "Lock",
|
||||
"detail": OPTIMIZE_LOCK_NOTE,
|
||||
},
|
||||
{
|
||||
"label": "Free space needed",
|
||||
"detail": f"About {_megabytes(size)} beside the database and "
|
||||
f"{_megabytes(size)} in SQLite's temporary folder; "
|
||||
f"{_megabytes(free)} available beside the database",
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
# ── apply ─────────────────────────────────────────────────────────
|
||||
|
||||
def apply_maintenance(self, operation, fingerprint, targets, roots=()) -> dict:
|
||||
"""Apply a reviewed plan exactly, or raise PlanChanged (blocking).
|
||||
|
||||
Inside ``BEGIN IMMEDIATE`` the fingerprint is computed again on the
|
||||
write connection; when it differs nothing is written. ``targets`` and
|
||||
``roots`` are the plan's. Returns ``{changed, rows}``: how many
|
||||
records changed and up to ``MAX_RESULT_ROWS`` descriptions of them
|
||||
(``{prompt_id[, kept_id]}`` or ``{image_path}``).
|
||||
"""
|
||||
if operation not in MAINTENANCE_OPERATIONS or operation == "workflow_storage":
|
||||
raise ValueError(f"Unknown maintenance operation {operation!r}")
|
||||
if operation == "optimize":
|
||||
# VACUUM cannot run in a transaction; the plan binds the database.
|
||||
if fingerprint != _Fingerprint("optimize").hexdigest():
|
||||
raise PlanChanged("optimize")
|
||||
self.model.vacuum_database()
|
||||
return {"changed": 1, "rows": []}
|
||||
appliers = {
|
||||
"duplicate_prompts": self._apply_duplicate_prompts,
|
||||
"unused_components": self._apply_unused_components,
|
||||
"missing_links": self._apply_missing_links,
|
||||
}
|
||||
with self._pair_transaction() as conn:
|
||||
return appliers[operation](conn, fingerprint, targets, roots)
|
||||
|
||||
def _apply_duplicate_prompts(self, conn, fingerprint, targets, _roots):
|
||||
groups, _paired, _conflicted, current = self._duplicates_fingerprint(conn)
|
||||
if current != fingerprint or groups != targets:
|
||||
raise PlanChanged("duplicate_prompts")
|
||||
changed, rows = 0, []
|
||||
for sorted_ids in groups:
|
||||
removed, _moved = self._merge_duplicate_group(conn, sorted_ids, strict=True)
|
||||
changed += removed
|
||||
rows.extend(
|
||||
{"prompt_id": duplicate_id, "kept_id": sorted_ids[0]}
|
||||
for duplicate_id in sorted_ids[1:]
|
||||
)
|
||||
return {"changed": changed, "rows": rows[:MAX_RESULT_ROWS]}
|
||||
|
||||
def _apply_unused_components(self, conn, fingerprint, targets, _roots):
|
||||
ids, current = self._unused_fingerprint(conn)
|
||||
if current != fingerprint or ids != targets:
|
||||
raise PlanChanged("unused_components")
|
||||
texts = _texts(conn, ids, limit=MAX_RESULT_ROWS)
|
||||
changed = _delete_ids(conn, "prompts", ids)
|
||||
rows = [
|
||||
{"prompt_id": prompt_id, "text": _excerpt(texts.get(prompt_id))}
|
||||
for prompt_id in ids[:MAX_RESULT_ROWS]
|
||||
]
|
||||
return {"changed": changed, "rows": rows}
|
||||
|
||||
def _apply_missing_links(self, conn, fingerprint, targets, roots):
|
||||
rows = []
|
||||
for batch in _batches(sorted(targets)):
|
||||
rows.extend(
|
||||
tuple(row)
|
||||
for row in conn.execute(
|
||||
"SELECT id, image_path, filename, file_path FROM generated_images "
|
||||
f"WHERE id IN ({_marks(batch)}) ORDER BY id", # nosec B608
|
||||
batch,
|
||||
)
|
||||
)
|
||||
readable = readable_roots(roots)
|
||||
candidates = []
|
||||
for row, owner, state in link_states(rows, roots, readable):
|
||||
if state != "missing": # back, or its source went away
|
||||
raise PlanChanged("missing_links")
|
||||
candidates.append((row, owner))
|
||||
if self._links_fingerprint(candidates) != fingerprint:
|
||||
raise PlanChanged("missing_links")
|
||||
changed = _delete_ids(conn, "generated_images", [row[0] for row in rows])
|
||||
return {
|
||||
"changed": changed,
|
||||
"rows": [{"image_path": row[1]} for row in rows[:MAX_RESULT_ROWS]],
|
||||
}
|
||||
@@ -1,353 +0,0 @@
|
||||
"""A SQLite connection that is never closed under the thread using it.
|
||||
|
||||
``PromptModel`` gives each thread its own connection, but Move library and
|
||||
``close_all`` must retire every thread's connection from one thread. Closing
|
||||
a ``sqlite3.Connection`` while its owner is inside a statement (Python runs
|
||||
``sqlite3_step`` without the GIL) frees memory the owner then reads: on
|
||||
Windows the process dies with "access violation", elsewhere the owner sees
|
||||
garbage such as a ``MemoryError``.
|
||||
|
||||
``ManagedConnection`` counts the calls its owner is inside. ``retire()`` is
|
||||
safe from any thread and never raises: an idle connection is closed at once;
|
||||
one in use is closed by its owner as the call ends. A retired connection
|
||||
refuses new work (cursors, statements, commits) with the error a closed
|
||||
connection raises.
|
||||
|
||||
What retiring a connection in use does and does not do:
|
||||
|
||||
- It calls ``sqlite3_interrupt`` (thread-safe). That ends a running
|
||||
statement at its next step, but it does not wake a statement that is
|
||||
waiting for another connection's lock: that one waits until the lock is
|
||||
released (or the busy timeout ends it), and only then is it interrupted.
|
||||
An interrupt that arrives between two statements, for example between the
|
||||
implicit ``BEGIN`` and the DML or between ``executemany`` rows, is lost.
|
||||
- Writes are kept out of a file the process has switched away from because
|
||||
a retired connection never commits: legacy-mode DML runs inside a
|
||||
transaction that only ``commit()`` (refused) or ``close()`` (rolls back)
|
||||
ends. A statement that commits by itself is not covered: DDL or a PRAGMA
|
||||
write outside a transaction, or an explicit ``COMMIT`` passed to
|
||||
``execute``/``executescript`` (``executescript`` also commits first, in C,
|
||||
on Python before 3.12). Managed connections must never be used that way
|
||||
for library writes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
import weakref
|
||||
from typing import Iterable, Iterator, List, Optional
|
||||
|
||||
import sqlite3
|
||||
|
||||
CLOSED_MESSAGE = "Cannot operate on a closed database."
|
||||
# Rows fetched per step when a cursor is iterated: the connection is marked
|
||||
# in use once per chunk instead of once per row (iteration costs about what
|
||||
# a plain cursor costs). A loop that stops early leaves up to this many rows
|
||||
# fetched; ``fetchone``/``fetchmany`` after a partial loop skip them.
|
||||
ITER_CHUNK = 128
|
||||
|
||||
|
||||
class ManagedCursor(sqlite3.Cursor):
|
||||
"""A cursor whose calls into SQLite count as its connection being in use."""
|
||||
|
||||
def execute(self, sql, parameters=(), /):
|
||||
conn = self.connection
|
||||
conn._enter()
|
||||
try:
|
||||
return super().execute(sql, parameters)
|
||||
finally:
|
||||
conn._leave()
|
||||
|
||||
def executemany(self, sql, seq_of_parameters, /):
|
||||
conn = self.connection
|
||||
conn._enter()
|
||||
try:
|
||||
return super().executemany(sql, seq_of_parameters)
|
||||
finally:
|
||||
conn._leave()
|
||||
|
||||
def executescript(self, sql_script, /):
|
||||
conn = self.connection
|
||||
conn._enter()
|
||||
try:
|
||||
return super().executescript(sql_script)
|
||||
finally:
|
||||
conn._leave()
|
||||
|
||||
def fetchone(self):
|
||||
conn = self.connection
|
||||
conn._enter()
|
||||
try:
|
||||
return super().fetchone()
|
||||
finally:
|
||||
conn._leave()
|
||||
|
||||
def fetchmany(self, *args, **kwargs):
|
||||
conn = self.connection
|
||||
conn._enter()
|
||||
try:
|
||||
return super().fetchmany(*args, **kwargs)
|
||||
finally:
|
||||
conn._leave()
|
||||
|
||||
def fetchall(self):
|
||||
conn = self.connection
|
||||
conn._enter()
|
||||
try:
|
||||
return super().fetchall()
|
||||
finally:
|
||||
conn._leave()
|
||||
|
||||
def __iter__(self) -> Iterator:
|
||||
return self._chunks()
|
||||
|
||||
def _chunks(self) -> Iterator:
|
||||
while True:
|
||||
rows = self.fetchmany(ITER_CHUNK)
|
||||
if not rows:
|
||||
return
|
||||
yield from rows
|
||||
|
||||
def __next__(self):
|
||||
conn = self.connection
|
||||
conn._enter()
|
||||
try:
|
||||
return super().__next__()
|
||||
finally:
|
||||
conn._leave()
|
||||
|
||||
def close(self):
|
||||
conn = self.connection
|
||||
try:
|
||||
conn._enter()
|
||||
except sqlite3.ProgrammingError:
|
||||
return super().close() # the connection is closed or closing
|
||||
try:
|
||||
return super().close()
|
||||
finally:
|
||||
conn._leave()
|
||||
|
||||
|
||||
class ManagedConnection(sqlite3.Connection):
|
||||
"""A connection only its owner closes while a call into it is running."""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
# Guards _depth, _retired, _closing and _cursors. Re-entrant so that
|
||||
# code running in the closing thread while it holds the lock (a
|
||||
# cursor() during the cursor-set copy) is refused instead of hanging.
|
||||
self._use_lock = threading.RLock()
|
||||
self._depth = 0
|
||||
self._retired = False
|
||||
self._closing = False
|
||||
self._released = threading.Event()
|
||||
# Python before 3.12 keeps the file open after close() while a cursor
|
||||
# still holds a statement; these are closed first (see _close_now).
|
||||
self._cursors: "weakref.WeakSet[sqlite3.Cursor]" = weakref.WeakSet()
|
||||
self.close_logger: Optional[logging.Logger] = None
|
||||
|
||||
# ── in use ───────────────────────────────────────────────────────
|
||||
|
||||
def _enter(self) -> None:
|
||||
with self._use_lock:
|
||||
if self._retired:
|
||||
raise sqlite3.ProgrammingError(CLOSED_MESSAGE)
|
||||
self._depth += 1
|
||||
|
||||
def _leave(self) -> None:
|
||||
with self._use_lock:
|
||||
self._depth -= 1
|
||||
close_now = self._retired and self._depth == 0
|
||||
if close_now:
|
||||
self._close_now()
|
||||
|
||||
@contextlib.contextmanager
|
||||
def using(self):
|
||||
"""Mark the connection in use for a call this class cannot wrap.
|
||||
|
||||
For example as the *target* of another connection's ``backup()``.
|
||||
"""
|
||||
self._enter()
|
||||
try:
|
||||
yield self
|
||||
finally:
|
||||
self._leave()
|
||||
|
||||
def cursor(self, factory=None):
|
||||
with self._use_lock:
|
||||
if self._retired:
|
||||
raise sqlite3.ProgrammingError(CLOSED_MESSAGE)
|
||||
cursor = super().cursor(factory or ManagedCursor)
|
||||
self._cursors.add(cursor)
|
||||
return cursor
|
||||
|
||||
def execute(self, sql, parameters=(), /):
|
||||
return self.cursor().execute(sql, parameters)
|
||||
|
||||
def executemany(self, sql, seq_of_parameters, /):
|
||||
return self.cursor().executemany(sql, seq_of_parameters)
|
||||
|
||||
def executescript(self, sql_script, /):
|
||||
return self.cursor().executescript(sql_script)
|
||||
|
||||
def commit(self):
|
||||
self._enter()
|
||||
try:
|
||||
return super().commit()
|
||||
finally:
|
||||
self._leave()
|
||||
|
||||
def rollback(self):
|
||||
self._enter() # a closed connection raises, as sqlite3's does
|
||||
try:
|
||||
return super().rollback()
|
||||
finally:
|
||||
self._leave()
|
||||
|
||||
def backup(self, target, **kwargs):
|
||||
self._enter()
|
||||
try:
|
||||
return super().backup(target, **kwargs)
|
||||
finally:
|
||||
self._leave()
|
||||
|
||||
def create_function(self, *args, **kwargs):
|
||||
self._enter()
|
||||
try:
|
||||
return super().create_function(*args, **kwargs)
|
||||
finally:
|
||||
self._leave()
|
||||
|
||||
@property
|
||||
def in_transaction(self) -> bool:
|
||||
self._enter()
|
||||
try:
|
||||
return super().in_transaction
|
||||
finally:
|
||||
self._leave()
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback):
|
||||
# The C implementation calls its own commit/rollback, not these.
|
||||
if exc_type is None:
|
||||
try:
|
||||
self.commit()
|
||||
except BaseException:
|
||||
with contextlib.suppress(sqlite3.Error):
|
||||
self.rollback()
|
||||
raise
|
||||
else:
|
||||
self.rollback()
|
||||
return False
|
||||
|
||||
# ── closing ──────────────────────────────────────────────────────
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close now if idle, else as soon as the running call ends.
|
||||
|
||||
Unlike ``sqlite3.Connection.close``, a call from another thread while
|
||||
the owner is inside a call returns before the file is closed. A caller
|
||||
that must know the file is closed (before renaming or deleting it)
|
||||
waits for ``wait_closed`` (``PromptModel.close_detached`` does).
|
||||
"""
|
||||
self.retire()
|
||||
|
||||
def retire(self) -> bool:
|
||||
"""Refuse new work and close; True when the connection is closed already.
|
||||
|
||||
Safe from any thread, and never raises: failures are logged.
|
||||
"""
|
||||
try:
|
||||
with self._use_lock:
|
||||
if not self._retired:
|
||||
self._retired = True
|
||||
if self._depth:
|
||||
# Ends the running statement (see the module docstring
|
||||
# for what it does not end). The owner cannot start
|
||||
# closing the connection while this lock is held.
|
||||
self.interrupt()
|
||||
return False
|
||||
elif self._depth:
|
||||
return False
|
||||
else:
|
||||
return self._released.wait(0)
|
||||
except Exception:
|
||||
self._log_exception("Error retiring database connection")
|
||||
return self._released.wait(0)
|
||||
self._close_now()
|
||||
return True
|
||||
|
||||
def wait_closed(self, timeout: float) -> bool:
|
||||
return self._released.wait(timeout)
|
||||
|
||||
def _close_now(self) -> None:
|
||||
# Only ever called with no call into the connection running, and a
|
||||
# retired connection refuses new calls and cursors: its cursors are
|
||||
# idle, and the set is copied under the lock that guards it.
|
||||
with self._use_lock:
|
||||
if self._closing:
|
||||
return # another thread closes it and sets _released
|
||||
self._closing = True
|
||||
cursors = list(self._cursors)
|
||||
try:
|
||||
for cursor in cursors:
|
||||
with contextlib.suppress(sqlite3.Error):
|
||||
sqlite3.Cursor.close(cursor)
|
||||
super().close()
|
||||
except Exception:
|
||||
self._log_exception("Error closing database connection")
|
||||
finally:
|
||||
self._released.set()
|
||||
|
||||
def _log_exception(self, message: str) -> None:
|
||||
if self.close_logger is not None:
|
||||
self.close_logger.exception(message)
|
||||
|
||||
|
||||
def retire_connections(
|
||||
connections: Iterable[sqlite3.Connection], logger=None
|
||||
) -> List[ManagedConnection]:
|
||||
"""Retire every connection; return those whose owner has yet to close them.
|
||||
|
||||
Safe while holding a lock the owners wait on: nothing here waits for them.
|
||||
One connection failing to retire is logged and never skips the rest.
|
||||
"""
|
||||
pending = []
|
||||
for conn in connections:
|
||||
try:
|
||||
if isinstance(conn, ManagedConnection):
|
||||
if logger is not None:
|
||||
conn.close_logger = logger
|
||||
if not conn.retire():
|
||||
pending.append(conn)
|
||||
else:
|
||||
conn.close()
|
||||
except Exception:
|
||||
if logger is not None:
|
||||
logger.exception("Error closing database connection")
|
||||
return pending
|
||||
|
||||
|
||||
def await_closed(
|
||||
connections: Iterable[ManagedConnection], timeout: float, logger=None
|
||||
) -> bool:
|
||||
"""Wait until the owners closed ``connections`` (retired); False on timeout.
|
||||
|
||||
A connection still open after ``timeout`` (its owner is blocked, for
|
||||
example on another process's lock after a lost interrupt) is logged and
|
||||
left to close when its call ends; a Move library then reports the old
|
||||
files as still in use instead of renaming them.
|
||||
"""
|
||||
deadline = time.monotonic() + timeout
|
||||
still_open = 0
|
||||
for conn in connections:
|
||||
if not conn.wait_closed(max(0.0, deadline - time.monotonic())):
|
||||
still_open += 1
|
||||
if still_open and logger is not None:
|
||||
logger.warning(
|
||||
f"{still_open} database connection(s) still in use after "
|
||||
f"{timeout:g}s; they close when their statement ends"
|
||||
)
|
||||
return not still_open
|
||||
+20
-137
@@ -15,11 +15,6 @@ from contextlib import closing
|
||||
from pathlib import Path
|
||||
from typing import Optional, Set, Tuple
|
||||
|
||||
from .managed_connection import (
|
||||
ManagedConnection,
|
||||
await_closed,
|
||||
retire_connections,
|
||||
)
|
||||
from .library_schema import library_schema_error, migrate_library_state
|
||||
from .image_favorites import image_favorites_schema_error, migrate_image_favorites
|
||||
from .annotation_schema import annotation_schema_error, migrate_annotation_state
|
||||
@@ -29,13 +24,6 @@ from .review_schema import (
|
||||
review_without_annotations_error,
|
||||
)
|
||||
from .saved_views import migrate_saved_views, saved_views_schema_error
|
||||
from .workflow_documents import (
|
||||
DOCUMENT_ID_COLUMNS,
|
||||
migrate_workflow_documents,
|
||||
workflow_documents_schema_error,
|
||||
)
|
||||
|
||||
|
||||
from .execution_history_schema import (
|
||||
history_schema_error,
|
||||
migrate_execution_history,
|
||||
@@ -62,19 +50,6 @@ MAX_PAIR_PROVENANCE_LENGTH = 4_000_000
|
||||
MAX_PAIR_PROVENANCE_DEPTH = 64
|
||||
|
||||
|
||||
def _refuse_rebuild_with_documents(columns) -> None:
|
||||
"""The legacy generated_images rebuilds copy a fixed column list.
|
||||
|
||||
A file that needs one (``prompt_id`` TEXT, no ``file_path``) but already
|
||||
has the workflow document ids would lose them, so it is refused.
|
||||
"""
|
||||
if any(column in columns for column in DOCUMENT_ID_COLUMNS):
|
||||
raise sqlite3.DatabaseError(
|
||||
"This database needs an old image table rebuild but already has "
|
||||
"workflow document references; the rebuild would drop them"
|
||||
)
|
||||
|
||||
|
||||
def pair_provenance_error(encoded: object) -> Optional[str]:
|
||||
"""Validate the same bounded JSON object for live writes and restores."""
|
||||
if not isinstance(encoded, str):
|
||||
@@ -341,8 +316,6 @@ class PromptModel:
|
||||
raise sqlite3.DatabaseError(review_error)
|
||||
migrate_annotation_state(conn)
|
||||
migrate_review_state(conn)
|
||||
# Last: the older generated_images rebuilds run only before file_path.
|
||||
migrate_workflow_documents(conn)
|
||||
|
||||
def _migrate_prompt_pairs(self, conn: sqlite3.Connection) -> None:
|
||||
conn.execute("SAVEPOINT prompt_pair_schema")
|
||||
@@ -614,42 +587,19 @@ class PromptModel:
|
||||
if conn is not None and self._local.generation == self._generation:
|
||||
return 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 and Move library can
|
||||
# retire other threads' connections; each connection is still
|
||||
# used by one thread, and ManagedConnection never closes it under
|
||||
# a running statement (see managed_connection).
|
||||
conn = sqlite3.connect(
|
||||
path,
|
||||
check_same_thread=False,
|
||||
timeout=BUSY_TIMEOUT_SECONDS,
|
||||
factory=ManagedConnection,
|
||||
)
|
||||
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()
|
||||
# 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)
|
||||
self._local.conn = conn
|
||||
self._local.generation = generation
|
||||
self._local.generation = self._generation
|
||||
with self._conn_lock:
|
||||
self._connections.add(conn)
|
||||
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."""
|
||||
@@ -658,76 +608,22 @@ 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
|
||||
retires them (``retire_detached``) while it still holds the old
|
||||
file's write lock. A write waiting on that lock is then interrupted
|
||||
and can never commit, so it 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 retire_detached(connections, logger=None) -> list:
|
||||
"""Retire detached connections without waiting for their owners.
|
||||
|
||||
Idle ones close now. One whose owner thread is inside a statement is
|
||||
never closed from here (that frees memory the owner is using: an
|
||||
access violation on Windows); it is interrupted, refuses to commit and
|
||||
is closed by its owner as the statement ends. Returns those still
|
||||
open, for ``await_detached`` once no lock the owners wait on is held.
|
||||
"""
|
||||
return retire_connections(connections, logger)
|
||||
|
||||
@staticmethod
|
||||
def await_detached(connections, logger=None) -> bool:
|
||||
"""Wait until retired connections are closed; False if some are not."""
|
||||
return await_closed(connections, BUSY_TIMEOUT_SECONDS, logger)
|
||||
|
||||
@staticmethod
|
||||
def close_detached(connections, logger=None) -> bool:
|
||||
"""Retire connections ``repoint_instances`` detached and await them.
|
||||
|
||||
Do not call it while holding a lock their owners may wait on; retire
|
||||
them there (``retire_detached``) and await them after releasing it.
|
||||
"""
|
||||
pending = retire_connections(connections, logger)
|
||||
return await_closed(pending, BUSY_TIMEOUT_SECONDS, logger)
|
||||
|
||||
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
|
||||
return connections
|
||||
for conn in connections:
|
||||
try:
|
||||
conn.close()
|
||||
except sqlite3.Error as e:
|
||||
self.logger.warning(f"Error closing database connection: {e}")
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the calling thread's connection, if it has one.
|
||||
@@ -741,7 +637,10 @@ class PromptModel:
|
||||
self._local.conn = None
|
||||
with self._conn_lock:
|
||||
self._connections.discard(conn)
|
||||
self.close_detached([conn], self.logger)
|
||||
try:
|
||||
conn.close()
|
||||
except sqlite3.Error as e:
|
||||
self.logger.warning(f"Error closing database connection: {e}")
|
||||
|
||||
def _migrate_workflow_name_removal(self, conn: sqlite3.Connection) -> None:
|
||||
"""
|
||||
@@ -812,7 +711,6 @@ class PromptModel:
|
||||
columns = {column[1]: column[2] for column in cursor.fetchall()}
|
||||
if columns.get("prompt_id") != "TEXT":
|
||||
return
|
||||
_refuse_rebuild_with_documents(columns)
|
||||
|
||||
self.logger.info("Migrating foreign key types: prompt_id TEXT -> INTEGER")
|
||||
with self._table_rebuild(conn):
|
||||
@@ -978,7 +876,6 @@ class PromptModel:
|
||||
}
|
||||
if "file_path" in columns:
|
||||
return
|
||||
_refuse_rebuild_with_documents(columns)
|
||||
|
||||
self.logger.info("Migrating database: image uniqueness by full path")
|
||||
with self._table_rebuild(conn):
|
||||
@@ -1044,10 +941,6 @@ class PromptModel:
|
||||
|
||||
Reclaims unused space and defragments the database file,
|
||||
improving query performance and reducing file size.
|
||||
|
||||
Raises:
|
||||
Exception: whatever stopped VACUUM (logged first); the database is
|
||||
then unchanged.
|
||||
"""
|
||||
try:
|
||||
with closing(sqlite3.connect(self.db_path)) as conn:
|
||||
@@ -1055,7 +948,6 @@ class PromptModel:
|
||||
conn.commit()
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error vacuuming database: {e}")
|
||||
raise
|
||||
|
||||
def get_database_info(self) -> dict:
|
||||
"""
|
||||
@@ -1220,11 +1112,6 @@ class PromptModel:
|
||||
)
|
||||
if review_error:
|
||||
return False, review_error
|
||||
documents_error = workflow_documents_schema_error(
|
||||
conn, allow_missing=not require_library_state, documents=True
|
||||
)
|
||||
if documents_error:
|
||||
return False, documents_error
|
||||
return True, "ok"
|
||||
except sqlite3.Error as e:
|
||||
return False, f"Invalid SQLite database: {e}"
|
||||
@@ -1247,11 +1134,7 @@ class PromptModel:
|
||||
Path(src_path).resolve().as_uri() + "?mode=ro", uri=True
|
||||
)
|
||||
try:
|
||||
target = self.get_connection()
|
||||
# The target is written under backup's step too: mark it in use
|
||||
# so a concurrent close_all retires it instead of closing it.
|
||||
with target.using():
|
||||
source.backup(target)
|
||||
source.backup(self.get_connection())
|
||||
finally:
|
||||
source.close()
|
||||
|
||||
|
||||
+219
-280
@@ -10,7 +10,7 @@ import os
|
||||
from typing import Optional, List, Dict, Any, Set, Tuple
|
||||
|
||||
from .models import PromptModel, normalize_image_path, utc_now_iso
|
||||
from .pairs import PairDependencyError, PairOperations # noqa: F401 (re-exported)
|
||||
from .pairs import PairDependencyError, PairOperations
|
||||
from .library import LibraryOperations, titled
|
||||
from .image_favorites import ImageFavoriteOperations
|
||||
from .annotations import AnnotationOperations
|
||||
@@ -20,33 +20,16 @@ from .saved_views import SavedViewOperations
|
||||
from .library_commands import LibraryCommandOperations
|
||||
from .vocabulary import VocabularyOperations
|
||||
from .execution_history import ExecutionHistoryOperations
|
||||
from .execution_contexts import ExecutionContextOperations
|
||||
from .execution_contexts import ExecutionContextOperations, LIVE_HINT_CTES
|
||||
from .execution_outputs import ExecutionOutputOperations
|
||||
from .history_reads import HistoryReadOperations
|
||||
from .health import HealthOperations, iter_consistency_issues
|
||||
from .duplicate_copies import DuplicateCopyOperations
|
||||
from .workflow_storage import WorkflowStorageOperations
|
||||
from .workflow_documents import (
|
||||
DocumentReader,
|
||||
documents_of_rows,
|
||||
release_documents,
|
||||
resolve_image_row,
|
||||
stored_values,
|
||||
)
|
||||
from .maintenance import ( # noqa: F401 (ORPHANED_PROMPTS_SQL is re-exported)
|
||||
ORPHANED_PROMPTS_SQL,
|
||||
_LINK_ROWS_SQL,
|
||||
MaintenanceOperations,
|
||||
_delete_ids,
|
||||
_orphaned_ids,
|
||||
)
|
||||
|
||||
# Import logging system
|
||||
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
|
||||
|
||||
@@ -55,7 +38,6 @@ 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
|
||||
@@ -101,56 +83,40 @@ def escape_like(term: str) -> str:
|
||||
)
|
||||
|
||||
|
||||
def _configured_default_path():
|
||||
"""config.json's ``database.default_path``, or None.
|
||||
|
||||
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.
|
||||
"""Resolve the database path from config, falling back to defaults.
|
||||
|
||||
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.
|
||||
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).
|
||||
|
||||
Args:
|
||||
db_path: Explicit database path, or None to use the resolver.
|
||||
db_path: Explicit database path, or None to use config/default.
|
||||
|
||||
Returns:
|
||||
Absolute path to the SQLite database file; its folder exists (a new
|
||||
data folder is created owner-only).
|
||||
Absolute path to the SQLite database file.
|
||||
"""
|
||||
extension_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
if db_path is None:
|
||||
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__)))
|
||||
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 = 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:
|
||||
library_location.ensure_private_dir(parent)
|
||||
os.makedirs(parent, exist_ok=True)
|
||||
|
||||
return db_path
|
||||
|
||||
@@ -200,6 +166,28 @@ def _coerce_file_info(metadata: Any) -> tuple:
|
||||
)
|
||||
|
||||
|
||||
# Exempt references: images, pairs, the __protected__ tag, Library parent
|
||||
# state, and live recorded-run hints (computed once, see LIVE_HINT_CTES).
|
||||
ORPHANED_PROMPTS_SQL = "WITH " + LIVE_HINT_CTES + """
|
||||
SELECT p.id FROM prompts p
|
||||
LEFT JOIN generated_images gi ON p.id = gi.prompt_id
|
||||
WHERE gi.prompt_id IS NULL
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM prompt_pairs pp
|
||||
WHERE pp.positive_prompt_id = p.id OR pp.negative_prompt_id = p.id
|
||||
)
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM prompt_tags pt
|
||||
JOIN tags t ON pt.tag_id = t.id
|
||||
WHERE pt.prompt_id = p.id AND t.name = '__protected__'
|
||||
)
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM prompt_library_state ls WHERE ls.prompt_id = p.id
|
||||
)
|
||||
AND p.id NOT IN (SELECT prompt_id FROM live_hint_prompts)
|
||||
"""
|
||||
|
||||
|
||||
# The scan's known set (#249 Activity-2): every image row, with its stored file
|
||||
# size and whether its pair link is resolved: exactly one distinct pair, the rule
|
||||
# _LINK_SUMMARY_SQL and get_image_prompt_info use. A file is known and resolved
|
||||
@@ -363,9 +351,6 @@ class PromptDatabase(
|
||||
ReviewDecisionOperations,
|
||||
SavedViewOperations,
|
||||
HealthOperations,
|
||||
MaintenanceOperations,
|
||||
DuplicateCopyOperations,
|
||||
WorkflowStorageOperations,
|
||||
):
|
||||
"""Database operations class for managing prompts."""
|
||||
|
||||
@@ -383,25 +368,6 @@ 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()
|
||||
@@ -822,12 +788,10 @@ class PromptDatabase(
|
||||
"""
|
||||
with self._pair_transaction() as conn:
|
||||
self._ensure_prompts_unpaired(conn, [prompt_id])
|
||||
documents = documents_of_rows(conn, "prompt_id", [prompt_id])
|
||||
# First delete related images to avoid foreign key constraint
|
||||
conn.execute(
|
||||
"DELETE FROM generated_images WHERE prompt_id = ?", (prompt_id,)
|
||||
)
|
||||
release_documents(conn, documents)
|
||||
# Then delete the prompt
|
||||
cursor = conn.execute("DELETE FROM prompts WHERE id = ?", (prompt_id,))
|
||||
conn.commit()
|
||||
@@ -851,12 +815,10 @@ class PromptDatabase(
|
||||
return 0
|
||||
self._ensure_prompts_unpaired(conn, ids)
|
||||
placeholders = ",".join("?" * len(ids))
|
||||
documents = documents_of_rows(conn, "prompt_id", ids)
|
||||
conn.execute(
|
||||
f"DELETE FROM generated_images WHERE prompt_id IN ({placeholders})",
|
||||
ids,
|
||||
)
|
||||
release_documents(conn, documents)
|
||||
conn.execute(
|
||||
f"DELETE FROM prompt_tags WHERE prompt_id IN ({placeholders})",
|
||||
ids,
|
||||
@@ -1579,14 +1541,80 @@ class PromptDatabase(
|
||||
self.logger.info("Starting duplicate cleanup process with image preservation")
|
||||
try:
|
||||
with self._pair_transaction() as conn:
|
||||
groups, _paired = self._duplicate_groups(conn)
|
||||
self.logger.debug(f"Found {len(groups)} groups of duplicate prompts")
|
||||
# Find duplicates by text content (case-insensitive)
|
||||
# Note: Removed ORDER BY from GROUP_CONCAT for SQLite compatibility
|
||||
# We'll sort the IDs manually after fetching
|
||||
cursor = conn.execute("""
|
||||
SELECT LOWER(TRIM(text)) as normalized_text, COUNT(*) as count,
|
||||
GROUP_CONCAT(id) as ids,
|
||||
GROUP_CONCAT(created_at) as created_dates
|
||||
FROM prompts
|
||||
GROUP BY LOWER(TRIM(text))
|
||||
HAVING COUNT(*) > 1
|
||||
""")
|
||||
|
||||
duplicates = cursor.fetchall()
|
||||
self.logger.debug(
|
||||
f"Found {len(duplicates)} groups of duplicate prompts"
|
||||
)
|
||||
total_removed = 0
|
||||
total_images_transferred = 0
|
||||
for sorted_ids in groups:
|
||||
removed, moved = self._merge_duplicate_group(conn, sorted_ids)
|
||||
total_removed += removed
|
||||
total_images_transferred += moved
|
||||
|
||||
for duplicate in duplicates:
|
||||
ids = duplicate["ids"].split(",")
|
||||
created_dates = duplicate["created_dates"].split(",")
|
||||
|
||||
# Sort IDs by created_at date to keep the oldest
|
||||
id_date_pairs = list(zip(ids, created_dates))
|
||||
id_date_pairs.sort(key=lambda x: x[1]) # Sort by date
|
||||
sorted_ids = [int(pair[0]) for pair in id_date_pairs]
|
||||
|
||||
try:
|
||||
self._ensure_prompts_unpaired(conn, sorted_ids)
|
||||
except PairDependencyError:
|
||||
continue
|
||||
|
||||
# Keep the oldest one (first), merge and delete the rest
|
||||
primary_id = sorted_ids[0] # Keep the oldest
|
||||
duplicate_ids = sorted_ids[1:]
|
||||
|
||||
self.logger.debug(
|
||||
f"Merging duplicates: keeping {primary_id}, removing "
|
||||
f"{duplicate_ids}"
|
||||
)
|
||||
|
||||
# Get primary prompt details
|
||||
cursor = conn.execute(
|
||||
"SELECT * FROM prompts WHERE id = ?", (primary_id,)
|
||||
)
|
||||
primary_prompt = cursor.fetchone()
|
||||
if not primary_prompt:
|
||||
continue
|
||||
|
||||
# Collect and merge metadata from all duplicates
|
||||
merged_metadata = self._merge_duplicate_metadata(
|
||||
conn, primary_id, duplicate_ids
|
||||
)
|
||||
|
||||
# Transfer all images from duplicates to primary prompt
|
||||
images_transferred = self._transfer_images_to_primary(
|
||||
conn, primary_id, duplicate_ids
|
||||
)
|
||||
total_images_transferred += images_transferred
|
||||
|
||||
# Update primary prompt with merged metadata
|
||||
if merged_metadata:
|
||||
self._update_primary_with_merged_metadata(
|
||||
conn, primary_id, merged_metadata
|
||||
)
|
||||
|
||||
# Delete duplicate prompts (images already transferred)
|
||||
for duplicate_id in duplicate_ids:
|
||||
conn.execute(
|
||||
"DELETE FROM prompts WHERE id = ?", (duplicate_id,)
|
||||
)
|
||||
total_removed += 1
|
||||
self.logger.debug(f"Removed duplicate prompt {duplicate_id}")
|
||||
|
||||
conn.commit()
|
||||
|
||||
@@ -1617,118 +1645,84 @@ class PromptDatabase(
|
||||
metadata: Optional metadata about the image
|
||||
|
||||
Returns:
|
||||
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)
|
||||
int: The ID of the created image record
|
||||
"""
|
||||
# Validate that prompt_id is a valid integer and exists in prompts table
|
||||
try:
|
||||
return self.link_image_to_prompt_or_raise(prompt_id, image_path, metadata)
|
||||
# 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:
|
||||
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,)
|
||||
)
|
||||
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)
|
||||
|
||||
# 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)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
prompt_id_int,
|
||||
image_path,
|
||||
filename,
|
||||
file_path,
|
||||
file_size,
|
||||
width,
|
||||
height,
|
||||
image_format,
|
||||
dumps_strict(metadata.get("workflow", {})),
|
||||
dumps_strict(metadata.get("prompt", {})),
|
||||
dumps_strict(metadata.get("parameters", {})),
|
||||
),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
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
|
||||
|
||||
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
|
||||
# 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:
|
||||
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,)
|
||||
)
|
||||
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,
|
||||
(
|
||||
documents["workflow_document_id"],
|
||||
documents["prompt_document_id"],
|
||||
),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
if cursor.rowcount == 0:
|
||||
self.logger.debug(
|
||||
f"Image {file_path} already linked to prompt {prompt_id_int}"
|
||||
)
|
||||
return 0
|
||||
|
||||
return cursor.lastrowid
|
||||
|
||||
def get_prompt_images(self, prompt_id: str) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Get all images associated with a prompt.
|
||||
@@ -1748,8 +1742,7 @@ class PromptDatabase(
|
||||
""",
|
||||
(prompt_id,),
|
||||
)
|
||||
reader = DocumentReader(conn)
|
||||
return [self._image_row_to_dict(row, reader) for row in cursor]
|
||||
return [self._image_row_to_dict(row) for row in cursor.fetchall()]
|
||||
|
||||
def get_recent_images(
|
||||
self, limit: int = 50, offset: int = 0
|
||||
@@ -1775,8 +1768,7 @@ class PromptDatabase(
|
||||
""",
|
||||
(int(limit), max(0, int(offset))),
|
||||
)
|
||||
reader = DocumentReader(conn)
|
||||
return [self._image_row_to_dict(row, reader) for row in cursor]
|
||||
return [self._image_row_to_dict(row) for row in cursor.fetchall()]
|
||||
|
||||
def get_all_images(
|
||||
self,
|
||||
@@ -1810,10 +1802,10 @@ class PromptDatabase(
|
||||
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(sql, params)
|
||||
reader = DocumentReader(conn)
|
||||
rows = cursor.fetchall()
|
||||
result = []
|
||||
for row in cursor:
|
||||
data = self._image_row_to_dict(row, reader)
|
||||
for row in rows:
|
||||
data = self._image_row_to_dict(row)
|
||||
if row["_prompt_tags_list"]:
|
||||
data["prompt_tags"] = row["_prompt_tags_list"].split("|||")
|
||||
else:
|
||||
@@ -1851,8 +1843,7 @@ class PromptDatabase(
|
||||
max(0, int(offset)),
|
||||
),
|
||||
)
|
||||
reader = DocumentReader(conn)
|
||||
return [self._image_row_to_dict(row, reader) for row in cursor]
|
||||
return [self._image_row_to_dict(row) for row in cursor.fetchall()]
|
||||
|
||||
def get_image_by_id(self, image_id: int) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
@@ -1869,7 +1860,7 @@ class PromptDatabase(
|
||||
"SELECT * FROM generated_images WHERE id = ?", (image_id,)
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
return self._image_row_to_dict(row, DocumentReader(conn)) if row else None
|
||||
return self._image_row_to_dict(row) if row else None
|
||||
|
||||
def delete_image(self, image_id: int) -> bool:
|
||||
"""
|
||||
@@ -1882,54 +1873,31 @@ class PromptDatabase(
|
||||
bool: True if deletion was successful
|
||||
"""
|
||||
with self.model.get_connection() as conn:
|
||||
documents = documents_of_rows(conn, "id", [image_id])
|
||||
cursor = conn.execute(
|
||||
"DELETE FROM generated_images WHERE id = ?", (image_id,)
|
||||
)
|
||||
release_documents(conn, documents)
|
||||
conn.commit()
|
||||
return cursor.rowcount > 0
|
||||
|
||||
def cleanup_missing_images(self, roots=None) -> int:
|
||||
def cleanup_missing_images(self) -> int:
|
||||
"""
|
||||
Remove image records where the actual file no longer exists.
|
||||
|
||||
Args:
|
||||
roots: The image sources links are judged under (the health
|
||||
check's). When given, only a link under a source that can be
|
||||
read now is removed: a link under an unavailable source (an
|
||||
unmounted drive) or outside every source is kept. None keeps
|
||||
the old rule: every link whose file does not exist.
|
||||
|
||||
Returns:
|
||||
int: Number of orphaned records removed
|
||||
"""
|
||||
if roots is not None:
|
||||
with self.model.get_connection() as conn:
|
||||
rows = [tuple(row) for row in conn.execute(_LINK_ROWS_SQL)]
|
||||
# Files are looked at before the write transaction starts.
|
||||
ids = self._missing_link_ids(rows, [str(root) for root in roots])
|
||||
if not ids:
|
||||
return 0
|
||||
with self._pair_transaction() as conn:
|
||||
return _delete_ids(conn, "generated_images", ids)
|
||||
|
||||
removed_count = 0
|
||||
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute("SELECT id, image_path FROM generated_images")
|
||||
images = cursor.fetchall()
|
||||
|
||||
removed = [
|
||||
image["id"]
|
||||
for image in images
|
||||
if not os.path.exists(image["image_path"])
|
||||
]
|
||||
documents = documents_of_rows(conn, "id", removed)
|
||||
for image_id in removed:
|
||||
conn.execute("DELETE FROM generated_images WHERE id = ?", (image_id,))
|
||||
removed_count += 1
|
||||
release_documents(conn, documents)
|
||||
for image in images:
|
||||
if not os.path.exists(image["image_path"]):
|
||||
conn.execute(
|
||||
"DELETE FROM generated_images WHERE id = ?", (image["id"],)
|
||||
)
|
||||
removed_count += 1
|
||||
|
||||
conn.commit()
|
||||
|
||||
@@ -1955,11 +1923,7 @@ class PromptDatabase(
|
||||
return obj
|
||||
|
||||
def _merge_duplicate_metadata(
|
||||
self,
|
||||
conn: sqlite3.Connection,
|
||||
primary_id: int,
|
||||
duplicate_ids: List[int],
|
||||
strict: bool = False,
|
||||
self, conn: sqlite3.Connection, primary_id: int, duplicate_ids: List[int]
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Merge metadata from duplicate prompts, prioritizing non-empty values.
|
||||
@@ -2019,17 +1983,11 @@ class PromptDatabase(
|
||||
return merged
|
||||
|
||||
except Exception as e:
|
||||
if strict: # a reviewed apply rolls back instead
|
||||
raise
|
||||
self.logger.error(f"Error merging metadata: {e}", exc_info=True)
|
||||
return {}
|
||||
|
||||
def _transfer_images_to_primary(
|
||||
self,
|
||||
conn: sqlite3.Connection,
|
||||
primary_id: int,
|
||||
duplicate_ids: List[int],
|
||||
strict: bool = False,
|
||||
self, conn: sqlite3.Connection, primary_id: int, duplicate_ids: List[int]
|
||||
) -> int:
|
||||
"""
|
||||
Transfer all images from duplicate prompts to the primary prompt.
|
||||
@@ -2061,17 +2019,11 @@ class PromptDatabase(
|
||||
return transferred_count
|
||||
|
||||
except Exception as e:
|
||||
if strict: # a reviewed apply rolls back instead
|
||||
raise
|
||||
self.logger.error(f"Error transferring images: {e}", exc_info=True)
|
||||
return 0
|
||||
|
||||
def _update_primary_with_merged_metadata(
|
||||
self,
|
||||
conn: sqlite3.Connection,
|
||||
primary_id: int,
|
||||
merged_metadata: Dict[str, Any],
|
||||
strict: bool = False,
|
||||
self, conn: sqlite3.Connection, primary_id: int, merged_metadata: Dict[str, Any]
|
||||
) -> None:
|
||||
"""
|
||||
Update the primary prompt with merged metadata.
|
||||
@@ -2102,8 +2054,6 @@ class PromptDatabase(
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
if strict: # a reviewed apply rolls back instead
|
||||
raise
|
||||
self.logger.error(
|
||||
f"Error updating primary prompt metadata: {e}", exc_info=True
|
||||
)
|
||||
@@ -2220,13 +2170,11 @@ class PromptDatabase(
|
||||
count = 0
|
||||
with self._pair_transaction() as conn:
|
||||
self._ensure_prompts_unpaired(conn, prompt_ids)
|
||||
documents = documents_of_rows(conn, "prompt_id", prompt_ids)
|
||||
for pid in prompt_ids:
|
||||
conn.execute("DELETE FROM generated_images WHERE prompt_id = ?", (pid,))
|
||||
cursor = conn.execute("DELETE FROM prompts WHERE id = ?", (pid,))
|
||||
if cursor.rowcount > 0:
|
||||
count += 1
|
||||
release_documents(conn, documents)
|
||||
conn.commit()
|
||||
return count
|
||||
|
||||
@@ -2300,8 +2248,7 @@ class PromptDatabase(
|
||||
cursor = conn.execute(
|
||||
"""SELECT gi.prompt_id, p.text, p.category, p.rating, p.notes,
|
||||
gi.workflow_data, gi.prompt_metadata, gi.generation_time,
|
||||
gi.parameters, gi.id, gi.image_path, gi.file_path,
|
||||
gi.workflow_document_id, gi.prompt_document_id
|
||||
gi.parameters, gi.id, gi.image_path, gi.file_path
|
||||
FROM generated_images gi
|
||||
JOIN prompts p ON gi.prompt_id = p.id
|
||||
WHERE gi.image_path = ? OR gi.file_path = ?
|
||||
@@ -2381,21 +2328,12 @@ class PromptDatabase(
|
||||
"text_source": source,
|
||||
}
|
||||
tags = self._get_prompt_tags(conn, row[0])
|
||||
texts = resolve_image_row(DocumentReader(conn), row)
|
||||
try:
|
||||
workflow = (
|
||||
json.loads(texts["workflow_data"])
|
||||
if texts["workflow_data"]
|
||||
else None
|
||||
)
|
||||
workflow = json.loads(row[5]) if row[5] else None
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
workflow = None
|
||||
try:
|
||||
prompt_meta = (
|
||||
json.loads(texts["prompt_metadata"])
|
||||
if texts["prompt_metadata"]
|
||||
else None
|
||||
)
|
||||
prompt_meta = json.loads(row[6]) if row[6] else None
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
prompt_meta = None
|
||||
return {
|
||||
@@ -2629,12 +2567,16 @@ class PromptDatabase(
|
||||
Returns number of prompts removed.
|
||||
"""
|
||||
with self._pair_transaction() as conn:
|
||||
orphaned = _orphaned_ids(conn)
|
||||
cursor = conn.execute(ORPHANED_PROMPTS_SQL)
|
||||
orphaned = [row["id"] for row in cursor.fetchall()]
|
||||
if not orphaned:
|
||||
return 0
|
||||
removed = _delete_ids(conn, "prompts", orphaned)
|
||||
placeholders = ",".join(["?"] * len(orphaned))
|
||||
cursor = conn.execute(
|
||||
f"DELETE FROM prompts WHERE id IN ({placeholders})", orphaned
|
||||
)
|
||||
conn.commit()
|
||||
return removed
|
||||
return cursor.rowcount
|
||||
|
||||
def check_consistency(self) -> List[str]:
|
||||
"""Run consistency checks and return a list of issue descriptions.
|
||||
@@ -2645,20 +2587,17 @@ class PromptDatabase(
|
||||
with self.model.get_connection() as conn:
|
||||
return list(iter_consistency_issues(conn))
|
||||
|
||||
def _image_row_to_dict(
|
||||
self, row: sqlite3.Row, reader: DocumentReader
|
||||
) -> Dict[str, Any]:
|
||||
def _image_row_to_dict(self, row: sqlite3.Row) -> Dict[str, Any]:
|
||||
"""
|
||||
Convert an image database row to a dictionary with parsed JSON fields.
|
||||
|
||||
Args:
|
||||
row: SQLite row object
|
||||
reader: The workflow document reader of the row's connection
|
||||
|
||||
Returns:
|
||||
Dictionary representation of the row
|
||||
"""
|
||||
data = resolve_image_row(reader, row)
|
||||
data = dict(row)
|
||||
|
||||
# Parse JSON fields
|
||||
for field in ["workflow_data", "prompt_metadata", "parameters"]:
|
||||
|
||||
@@ -5,13 +5,6 @@ from typing import Any, Dict
|
||||
|
||||
from .models import normalize_image_path
|
||||
from .operations import _coerce_file_info
|
||||
from .workflow_documents import (
|
||||
TEXT_COLUMNS,
|
||||
DocumentReader,
|
||||
release_documents,
|
||||
resolve_image_row,
|
||||
stored_values,
|
||||
)
|
||||
|
||||
try:
|
||||
from ..utils.hashing import generate_prompt_hash
|
||||
@@ -123,35 +116,25 @@ def _ingest_resolved_pair_in(
|
||||
changed = False
|
||||
if existing:
|
||||
image_id = existing["id"]
|
||||
# Compared as stored text, whichever form the row is in.
|
||||
current = resolve_image_row(DocumentReader(conn), existing)
|
||||
updates = {
|
||||
key: value for key, value in snapshots.items() if current[key] != value
|
||||
key: value for key, value in snapshots.items() if existing[key] != value
|
||||
}
|
||||
if updates:
|
||||
replaced = [
|
||||
existing[reference]
|
||||
for column, reference in TEXT_COLUMNS
|
||||
if column in updates
|
||||
]
|
||||
updates = _document_form(conn, updates)
|
||||
conn.execute(
|
||||
"UPDATE generated_images SET "
|
||||
+ ",".join(f"{key}=?" for key in updates)
|
||||
+ " WHERE id=?",
|
||||
[*updates.values(), image_id],
|
||||
)
|
||||
release_documents(conn, replaced)
|
||||
changed = True
|
||||
else:
|
||||
stored = _document_form(conn, snapshots)
|
||||
columns = ["prompt_id", "image_path", "filename", "file_path", *stored]
|
||||
columns = ["prompt_id", "image_path", "filename", "file_path", *snapshots]
|
||||
values = [
|
||||
pair["positive_prompt_id"],
|
||||
image_path,
|
||||
os.path.basename(image_path),
|
||||
normalized,
|
||||
*stored.values(),
|
||||
*snapshots.values(),
|
||||
]
|
||||
image_id = conn.execute(
|
||||
"INSERT INTO generated_images("
|
||||
@@ -179,18 +162,3 @@ def _ingest_resolved_pair_in(
|
||||
"pair_id": pair_id,
|
||||
"image_id": image_id,
|
||||
}
|
||||
|
||||
|
||||
def _document_form(conn: Any, values: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""``values`` with the workflow and prompt graph texts stored once.
|
||||
|
||||
Each of those columns becomes NULL beside its document id (new form);
|
||||
every other column is unchanged.
|
||||
"""
|
||||
texts = {
|
||||
column: values[column]
|
||||
for column, _reference in TEXT_COLUMNS
|
||||
if column in values
|
||||
}
|
||||
rest = {key: value for key, value in values.items() if key not in texts}
|
||||
return {**rest, **stored_values(conn, texts)}
|
||||
|
||||
+2
-89
@@ -90,30 +90,6 @@ _PENDING_AT = (
|
||||
"WHERE f.locator=? AND b.size=? AND b.mtime_ns=? AND i.resolver_version=? "
|
||||
"AND i.state='pending' LIMIT 1"
|
||||
)
|
||||
# Bytes verified at this stat whose item carries a decision (any version).
|
||||
_DECIDED_AT = (
|
||||
f"SELECT 1 FROM {FILES_TABLE} f " # nosec B608
|
||||
f"JOIN {BYTES_TABLE} b ON b.file_id=f.id "
|
||||
f"JOIN {ITEMS_TABLE} i ON i.byte_id=b.id "
|
||||
"WHERE f.locator=? AND b.size=? AND b.mtime_ns=? "
|
||||
"AND (i.state<>'pending' OR i.revision<>0 OR EXISTS (SELECT 1 FROM "
|
||||
f"{DECISIONS_TABLE} d WHERE d.item_id=i.id)) LIMIT 1"
|
||||
)
|
||||
# The other items of one byte record (another version or evidence), with
|
||||
# whether each may be superseded: pending, revision 0, no decision row.
|
||||
_OTHER_ITEMS_OF_BYTES = (
|
||||
"SELECT i.id, (i.state='pending' AND i.revision=0 AND NOT EXISTS " # nosec B608
|
||||
f"(SELECT 1 FROM {DECISIONS_TABLE} d WHERE d.item_id=i.id)) "
|
||||
f"FROM {ITEMS_TABLE} i WHERE i.byte_id=? "
|
||||
"AND NOT (i.resolver_version=? AND i.evidence_hash=?)"
|
||||
)
|
||||
# Pending items of the bytes verified at one stat, with their versions.
|
||||
_PENDING_AT_STAT = (
|
||||
f"SELECT i.id,i.resolver_version FROM {FILES_TABLE} f " # nosec B608
|
||||
f"JOIN {BYTES_TABLE} b ON b.file_id=f.id "
|
||||
f"JOIN {ITEMS_TABLE} i ON i.byte_id=b.id "
|
||||
"WHERE f.locator=? AND b.size=? AND b.mtime_ns=? AND i.state='pending'"
|
||||
)
|
||||
# Pending items of one file: the candidates a link settles.
|
||||
_PENDING_OF_LOCATOR = (
|
||||
f"SELECT i.id FROM {FILES_TABLE} f " # nosec B608
|
||||
@@ -298,12 +274,6 @@ class ReviewOperations:
|
||||
``LibraryCommandError("storage-changed")`` when the storage epoch is
|
||||
no longer ``epoch`` (a restore since the capture) and ``ValueError``
|
||||
for invalid facts or evidence.
|
||||
|
||||
One file never holds two pending items for the same bytes: the new
|
||||
item supersedes the byte record's other items (an older resolver or
|
||||
semantics version) when every one is pending and undecided (the safe
|
||||
delete of the self-heal); when one carries a decision or a later
|
||||
revision, it is kept and nothing new is stored (False).
|
||||
"""
|
||||
_validate_target(locator, sha256, size, mtime_ns)
|
||||
_validate_spelling(spelled_path, locator)
|
||||
@@ -327,18 +297,6 @@ class ReviewOperations:
|
||||
)
|
||||
if created or moved:
|
||||
conn.execute(_BUMP_IMAGE_INDEX)
|
||||
# The byte record is created or refreshed even when the item is
|
||||
# refused below (an observe would do the same); that commits
|
||||
# with the image index bump, and the review queue is unchanged.
|
||||
others = conn.execute(
|
||||
_OTHER_ITEMS_OF_BYTES, (byte_id, resolver_version, evidence_hash)
|
||||
).fetchall()
|
||||
if any(not removable for _id, removable in others):
|
||||
return False
|
||||
superseded = sum(
|
||||
conn.execute(_SETTLE_ITEM, (item_id,)).rowcount
|
||||
for item_id, _removable in others
|
||||
)
|
||||
cursor = conn.execute(
|
||||
_INSERT_ITEM,
|
||||
(
|
||||
@@ -356,9 +314,9 @@ class ReviewOperations:
|
||||
)
|
||||
inserted = cursor.rowcount > 0
|
||||
refreshed = _refresh_spelling(conn, locator, spelled_path)
|
||||
if inserted or refreshed or superseded:
|
||||
if inserted or refreshed:
|
||||
conn.execute(_BUMP_REVIEW_INDEX)
|
||||
if inserted or refreshed or superseded:
|
||||
if inserted or refreshed:
|
||||
notify_review_change(self.model.db_path)
|
||||
return inserted
|
||||
|
||||
@@ -431,19 +389,6 @@ class ReviewOperations:
|
||||
)
|
||||
return row is not None
|
||||
|
||||
def review_decided_at(self, locator: str, size: int, mtime_ns: int) -> bool:
|
||||
"""True when bytes verified at this stat carry a decided item.
|
||||
|
||||
Any version: a newer resolver or table never adds an item beside a
|
||||
decision, so the rescan need not copy those bytes again.
|
||||
"""
|
||||
row = (
|
||||
self.model.get_connection()
|
||||
.execute(_DECIDED_AT, (locator, size, mtime_ns))
|
||||
.fetchone()
|
||||
)
|
||||
return row is not None
|
||||
|
||||
def review_pending_at(
|
||||
self, locator: str, size: int, mtime_ns: int, resolver_version: int
|
||||
) -> bool:
|
||||
@@ -500,38 +445,6 @@ class ReviewOperations:
|
||||
notify_review_change(self.model.db_path)
|
||||
return removed
|
||||
|
||||
def settle_stale_review_items(
|
||||
self, locator: str, size: int, mtime_ns: int, versions
|
||||
) -> int:
|
||||
"""Remove pending, undecided items of the bytes verified at this stat
|
||||
whose ``resolver_version`` is not one of ``versions``.
|
||||
|
||||
The caller found these bytes untraced under the current versions, so
|
||||
an older version's item could only refuse Apply. The same safe delete
|
||||
as the other settles; one indexed read when nothing is pending.
|
||||
"""
|
||||
_validate_locator(locator)
|
||||
current = set(versions)
|
||||
conn = self.model.get_connection()
|
||||
stale = [
|
||||
item_id
|
||||
for item_id, version in conn.execute(
|
||||
_PENDING_AT_STAT, (locator, size, mtime_ns)
|
||||
)
|
||||
if version not in current
|
||||
]
|
||||
if not stale:
|
||||
return 0
|
||||
removed = 0
|
||||
with self._pair_transaction() as conn:
|
||||
for item_id in stale:
|
||||
removed += conn.execute(_SETTLE_ITEM, (item_id,)).rowcount
|
||||
if removed:
|
||||
conn.execute(_BUMP_REVIEW_INDEX)
|
||||
if removed:
|
||||
notify_review_change(self.model.db_path)
|
||||
return removed
|
||||
|
||||
def settle_untraced_review_items(self, locators) -> int:
|
||||
"""Remove the pending, undecided items that offer nothing to choose.
|
||||
|
||||
|
||||
@@ -62,14 +62,3 @@ 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)
|
||||
|
||||
@@ -1,477 +0,0 @@
|
||||
"""Each workflow and prompt graph stored once, compressed (#249 S6 carry-over).
|
||||
|
||||
``generated_images.workflow_data`` (the UI workflow) and ``prompt_metadata``
|
||||
(the API prompt graph) used to hold the same large JSON text again for every
|
||||
image of a run. Now each distinct text is one ``workflow_documents`` row:
|
||||
|
||||
- ``sha256``: the SHA-256 of the text's exact UTF-8 bytes (unique key);
|
||||
- ``byte_length``: the length of those bytes;
|
||||
- ``data``: ``zlib.compress`` of those bytes.
|
||||
|
||||
An image row references its texts through ``workflow_document_id`` and
|
||||
``prompt_document_id``. Two forms coexist until the migration has rewritten
|
||||
every row, and both stay readable forever:
|
||||
|
||||
- old form: the text inline in ``workflow_data`` / ``prompt_metadata``, no id;
|
||||
- new form: the inline column NULL and the document id set.
|
||||
|
||||
The text is stored and hashed exactly as written, never parsed and
|
||||
re-serialized, because a JSON round trip can change bytes (key order,
|
||||
spacing, number spelling, duplicate keys, ComfyUI's ``NaN`` tokens).
|
||||
|
||||
``read_stored_text`` is the one decoder. Inline text, when present, wins (a
|
||||
row an older app version rewrote inline stays right); otherwise the document
|
||||
is read by its rowid, decompressed within its recorded length and checked
|
||||
against its length and hash. A document that fails a check reads as
|
||||
unreadable (None, logged), never as other text. ``resolve_image_row`` puts
|
||||
the texts back into a ``generated_images`` row and drops the id keys, so
|
||||
every payload keeps its shape.
|
||||
|
||||
No trigger: every restore validator refuses triggers in a restored file. A
|
||||
document is removed in code: ``release_documents`` in the same transaction
|
||||
that drops a row's last reference, and ``sweep_documents`` (at startup, after
|
||||
the migration and before Reclaim space) for anything a path missed.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import re
|
||||
import sqlite3
|
||||
import zlib
|
||||
from collections import OrderedDict
|
||||
from typing import Any, Dict, Iterable, Optional, Set
|
||||
|
||||
logger = logging.getLogger("prompt_manager.database.workflow_documents")
|
||||
|
||||
TABLE = "workflow_documents"
|
||||
# (inline text column, document id column) on generated_images.
|
||||
TEXT_COLUMNS = (
|
||||
("workflow_data", "workflow_document_id"),
|
||||
("prompt_metadata", "prompt_document_id"),
|
||||
)
|
||||
DOCUMENT_ID_COLUMNS = tuple(reference for _column, reference in TEXT_COLUMNS)
|
||||
INDEXES = {
|
||||
"workflow_document_id": "idx_image_workflow_document",
|
||||
"prompt_document_id": "idx_image_prompt_document",
|
||||
}
|
||||
COMPRESSION_LEVEL = 6
|
||||
# A text larger than this stays inline: a document is decompressed in one
|
||||
# piece, and its recorded length bounds that (a restored file cannot make a
|
||||
# small row expand without limit).
|
||||
MAX_DOCUMENT_BYTES = 256 * 2**20
|
||||
ID_BATCH = 400 # ids bound per statement (under SQLite's old 999 limit)
|
||||
|
||||
_DDL = (
|
||||
f"CREATE TABLE {TABLE} (id INTEGER PRIMARY KEY, sha256 BLOB NOT NULL UNIQUE, "
|
||||
"byte_length INTEGER NOT NULL, data BLOB NOT NULL)"
|
||||
)
|
||||
_AUTOINDEX = f"sqlite_autoindex_{TABLE}_1"
|
||||
_REFERENCE_TYPE = f"INTEGER REFERENCES {TABLE}(id)"
|
||||
|
||||
_UNREFERENCED = (
|
||||
"NOT EXISTS (SELECT 1 FROM generated_images g "
|
||||
"WHERE g.workflow_document_id=workflow_documents.id) "
|
||||
"AND NOT EXISTS (SELECT 1 FROM generated_images g "
|
||||
"WHERE g.prompt_document_id=workflow_documents.id)"
|
||||
)
|
||||
LOOKUP_SQL = (
|
||||
"SELECT id, sha256, byte_length, data FROM workflow_documents WHERE sha256=?"
|
||||
)
|
||||
REPAIR_SQL = "UPDATE workflow_documents SET byte_length=?, data=? WHERE id=?"
|
||||
ALL_DOCUMENTS_SQL = "SELECT sha256, byte_length, data, id FROM workflow_documents"
|
||||
READ_SQL = "SELECT sha256, byte_length, data FROM workflow_documents WHERE id=?"
|
||||
INSERT_SQL = (
|
||||
"INSERT INTO workflow_documents (sha256, byte_length, data) VALUES (?, ?, ?)"
|
||||
)
|
||||
# Built from constants only; the ids are bound as "?" marks.
|
||||
RELEASE_SQL = "DELETE FROM workflow_documents WHERE id IN ({marks}) AND "
|
||||
RELEASE_SQL += _UNREFERENCED
|
||||
# The sweep walks the ids through the small unique index, not the blobs.
|
||||
SWEEP_SQL = (
|
||||
"DELETE FROM workflow_documents WHERE id IN (SELECT id FROM workflow_documents "
|
||||
"INDEXED BY sqlite_autoindex_workflow_documents_1 WHERE "
|
||||
)
|
||||
SWEEP_SQL += _UNREFERENCED + ")"
|
||||
IMAGE_DOCUMENTS_SQL = (
|
||||
"SELECT workflow_document_id, prompt_document_id FROM generated_images WHERE "
|
||||
)
|
||||
|
||||
# Statements whose query plans tests/test_workflow_documents.py records.
|
||||
QUERY_PLAN_STATEMENTS = {
|
||||
"lookup by hash": (LOOKUP_SQL, (b"\0" * 32,)),
|
||||
"read by id": (READ_SQL, (1,)),
|
||||
"release": (RELEASE_SQL.format(marks="?,?"), (1, 2)),
|
||||
"sweep": (SWEEP_SQL, ()),
|
||||
"documents of images": (IMAGE_DOCUMENTS_SQL + "id IN (?,?)", (1, 2)),
|
||||
"documents of prompts": (IMAGE_DOCUMENTS_SQL + "prompt_id IN (?,?)", (1, 2)),
|
||||
}
|
||||
|
||||
|
||||
class DocumentMismatch(Exception):
|
||||
"""A document would not give back exactly the text it was made from."""
|
||||
|
||||
|
||||
def _compact(sql: str) -> str:
|
||||
return re.sub(r"\s+", "", sql or "").lower()
|
||||
|
||||
|
||||
# ── schema ────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _reference_columns(conn: sqlite3.Connection) -> Dict[str, tuple]:
|
||||
return {
|
||||
row[1]: tuple(row)
|
||||
for row in conn.execute("PRAGMA table_info(generated_images)")
|
||||
if row[1] in DOCUMENT_ID_COLUMNS
|
||||
}
|
||||
|
||||
|
||||
def _reference_error(conn: sqlite3.Connection, columns: Dict[str, tuple]):
|
||||
for name, info in columns.items():
|
||||
# (cid, name, type, notnull, default, pk)
|
||||
if (str(info[2]).upper(), info[3], info[4], info[5]) != ("INTEGER", 0, None, 0):
|
||||
return "Workflow document references have an invalid column"
|
||||
keys = {
|
||||
(row[3], row[2], row[4])
|
||||
for row in conn.execute("PRAGMA foreign_key_list(generated_images)")
|
||||
if row[3] in DOCUMENT_ID_COLUMNS
|
||||
}
|
||||
if keys != {(name, TABLE, "id") for name in DOCUMENT_ID_COLUMNS}:
|
||||
return "Workflow document references have invalid foreign keys"
|
||||
for column, index in INDEXES.items():
|
||||
info = conn.execute(f'PRAGMA index_info("{index}")').fetchall()
|
||||
owner = conn.execute(
|
||||
"SELECT tbl_name FROM sqlite_master WHERE type='index' AND name=?",
|
||||
(index,),
|
||||
).fetchone()
|
||||
if [row[2] for row in info] != [column] or not owner:
|
||||
return "Workflow document references are missing their index"
|
||||
if owner[0] != "generated_images":
|
||||
return "Workflow document references are missing their index"
|
||||
return None
|
||||
|
||||
|
||||
def _value_error(conn: sqlite3.Connection):
|
||||
if conn.execute(
|
||||
"SELECT 1 FROM workflow_documents WHERE typeof(sha256)!='blob' "
|
||||
"OR length(sha256)!=32 OR typeof(byte_length)!='integer' OR byte_length<0 "
|
||||
"OR byte_length>? OR typeof(data)!='blob' LIMIT 1",
|
||||
(MAX_DOCUMENT_BYTES,),
|
||||
).fetchone():
|
||||
return "Workflow documents contain invalid values"
|
||||
for column in DOCUMENT_ID_COLUMNS:
|
||||
# typeof() reads no overflow page; a dangling id is a rowid probe.
|
||||
# ``column`` is one of DOCUMENT_ID_COLUMNS, never caller data.
|
||||
sql = (
|
||||
f"SELECT 1 FROM generated_images g WHERE typeof(g.{column}) " # nosec B608
|
||||
f"NOT IN ('null','integer') OR (g.{column} IS NOT NULL AND NOT EXISTS "
|
||||
f"(SELECT 1 FROM workflow_documents d WHERE d.id=g.{column})) LIMIT 1"
|
||||
)
|
||||
if conn.execute(sql).fetchone():
|
||||
return "Workflow document references point at missing documents"
|
||||
return None
|
||||
|
||||
|
||||
def damaged_message(document_id: int) -> str:
|
||||
"""The verify reason for a damaged document: which one, and what to do."""
|
||||
return (
|
||||
f"Workflow document {document_id} is damaged: its text no longer matches "
|
||||
"its checksum, and the images that use it show no workflow. To fix it, "
|
||||
"restore a backup from before the damage, or store a new image with the "
|
||||
"same workflow (that repairs it)."
|
||||
)
|
||||
|
||||
|
||||
def _damaged_document(conn: sqlite3.Connection) -> Optional[int]:
|
||||
"""The first document failing its length or hash check, or None."""
|
||||
for row in conn.execute(ALL_DOCUMENTS_SQL):
|
||||
if _decode(row[:3], row[3], quiet=True) is None:
|
||||
return row[3]
|
||||
return None
|
||||
|
||||
|
||||
def workflow_documents_schema_error(
|
||||
conn: sqlite3.Connection, *, allow_missing: bool = True, documents: bool = False
|
||||
) -> Optional[str]:
|
||||
"""Reject an unsupported document schema or invalid stored values (reads only).
|
||||
|
||||
A database from before the dedupe has neither the table nor the id
|
||||
columns; that is acceptable unless ``allow_missing`` is false. Half of
|
||||
the schema never is.
|
||||
"""
|
||||
objects = conn.execute(
|
||||
"SELECT type,name,sql FROM sqlite_master WHERE name=? "
|
||||
"OR (type IN ('index','trigger') AND tbl_name=?)",
|
||||
(TABLE, TABLE),
|
||||
).fetchall()
|
||||
has_images = conn.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE type='table' AND name='generated_images'"
|
||||
).fetchone()
|
||||
columns = _reference_columns(conn) if has_images else {}
|
||||
if not objects and not columns:
|
||||
if allow_missing or not has_images:
|
||||
return None
|
||||
return "Missing workflow document schema"
|
||||
if not objects or set(columns) != set(DOCUMENT_ID_COLUMNS):
|
||||
return "Incomplete workflow document schema"
|
||||
tables = [sql for kind, _name, sql in objects if kind == "table"]
|
||||
others = {(kind, name, sql) for kind, name, sql in objects if kind != "table"}
|
||||
if len(tables) != 1 or others != {("index", _AUTOINDEX, None)}:
|
||||
return "Workflow documents name is reserved by an unsupported schema object"
|
||||
# The exact definition fixes the columns, their types and constraints.
|
||||
if _compact(tables[0]) != _compact(_DDL):
|
||||
return "Workflow documents have an invalid table definition"
|
||||
error = _reference_error(conn, columns) or _value_error(conn)
|
||||
if error is None and documents:
|
||||
# Restore and backup verification only: every document is read.
|
||||
damaged = _damaged_document(conn)
|
||||
if damaged is not None:
|
||||
return damaged_message(damaged)
|
||||
return error
|
||||
|
||||
|
||||
def migrate_workflow_documents(conn: sqlite3.Connection) -> None:
|
||||
"""Install the table, the id columns and their indexes atomically.
|
||||
|
||||
Runs after every older migration (the table rebuilds only run on
|
||||
databases without ``file_path``, before this). Validates before and
|
||||
after, then sweeps documents no row references.
|
||||
"""
|
||||
conn.execute("SAVEPOINT workflow_documents_schema")
|
||||
try:
|
||||
error = workflow_documents_schema_error(conn)
|
||||
if error:
|
||||
raise sqlite3.DatabaseError(error)
|
||||
conn.execute(_DDL.replace("CREATE TABLE", "CREATE TABLE IF NOT EXISTS", 1))
|
||||
present = _reference_columns(conn)
|
||||
for column in DOCUMENT_ID_COLUMNS:
|
||||
if column not in present:
|
||||
conn.execute(
|
||||
f"ALTER TABLE generated_images ADD COLUMN {column} "
|
||||
f"{_REFERENCE_TYPE}"
|
||||
)
|
||||
for column, index in INDEXES.items():
|
||||
conn.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS {index} ON generated_images({column})"
|
||||
)
|
||||
error = workflow_documents_schema_error(conn, allow_missing=False)
|
||||
if error:
|
||||
raise sqlite3.DatabaseError(error)
|
||||
sweep_documents(conn)
|
||||
conn.execute("RELEASE workflow_documents_schema")
|
||||
except BaseException:
|
||||
conn.execute("ROLLBACK TO workflow_documents_schema")
|
||||
conn.execute("RELEASE workflow_documents_schema")
|
||||
raise
|
||||
|
||||
|
||||
# ── write ─────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def storable(text: Any) -> bool:
|
||||
"""Whether ``text`` can be a document (text within MAX_DOCUMENT_BYTES)."""
|
||||
if not isinstance(text, str):
|
||||
return False
|
||||
# UTF-8 takes at most 4 bytes a character: most texts need no encoding.
|
||||
if len(text) * 4 <= MAX_DOCUMENT_BYTES:
|
||||
return True
|
||||
return len(text.encode("utf-8")) <= MAX_DOCUMENT_BYTES
|
||||
|
||||
|
||||
def store_text(
|
||||
conn: sqlite3.Connection, text: Optional[str], verified: Optional[set] = None
|
||||
) -> Optional[int]:
|
||||
"""The id of the document holding exactly ``text``; None for None.
|
||||
|
||||
Reuses the document with the same hash, or inserts a new one whose
|
||||
compressed bytes were checked to decompress to ``text``'s bytes. A reused
|
||||
document is decompressed and checked first; a damaged one (a restored
|
||||
file's corruption under a valid hash) is repaired from ``text``, whose
|
||||
hash is its key, so every row that references it reads right again.
|
||||
``verified``, when given, holds the ids already checked (a migration
|
||||
batch), which are then not read again. Opens a write transaction first
|
||||
when none is open, so a concurrent ``release_documents`` cannot remove a
|
||||
reused document before the caller references it. Raises TypeError for a
|
||||
non-text value and ValueError for a text over MAX_DOCUMENT_BYTES.
|
||||
"""
|
||||
if text is None:
|
||||
return None
|
||||
if not isinstance(text, str):
|
||||
raise TypeError("Only text is stored as a workflow document")
|
||||
raw = text.encode("utf-8")
|
||||
if len(raw) > MAX_DOCUMENT_BYTES:
|
||||
raise ValueError("Text is too large for a workflow document")
|
||||
digest = hashlib.sha256(raw).digest()
|
||||
if not conn.in_transaction:
|
||||
conn.execute("BEGIN IMMEDIATE")
|
||||
row = conn.execute(LOOKUP_SQL, (digest,)).fetchone()
|
||||
if row is not None and verified is not None and row[0] in verified:
|
||||
return row[0]
|
||||
data = zlib.compress(raw, COMPRESSION_LEVEL)
|
||||
if zlib.decompress(data) != raw: # pragma: no cover - zlib is exact
|
||||
raise DocumentMismatch("Compression did not round-trip")
|
||||
if row is None:
|
||||
document_id = conn.execute(INSERT_SQL, (digest, len(raw), data)).lastrowid
|
||||
else:
|
||||
document_id = row[0]
|
||||
if _decode(row[1:], document_id, quiet=True) != text:
|
||||
logger.warning(f"Workflow document {document_id} was damaged; repaired")
|
||||
conn.execute(REPAIR_SQL, (len(raw), data, document_id))
|
||||
if verified is not None:
|
||||
verified.add(document_id)
|
||||
return document_id
|
||||
|
||||
|
||||
def stored_values(
|
||||
conn: sqlite3.Connection, texts: Dict[str, Optional[str]]
|
||||
) -> Dict[str, Any]:
|
||||
"""New-form column values for ``{text column: text}`` (writers).
|
||||
|
||||
Each text becomes ``{column: None, id column: document id}``; None
|
||||
stays ``{column: None, id column: None}``; a value that cannot be a
|
||||
document (not text, or too large) stays inline with no id.
|
||||
"""
|
||||
references = dict(TEXT_COLUMNS)
|
||||
values: Dict[str, Any] = {}
|
||||
for column, text in texts.items():
|
||||
reference = references[column]
|
||||
if text is None or not storable(text):
|
||||
values[column], values[reference] = text, None
|
||||
else:
|
||||
values[column], values[reference] = None, store_text(conn, text)
|
||||
return values
|
||||
|
||||
|
||||
# ── read ──────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _decode(row, document_id, quiet=False) -> Optional[str]:
|
||||
"""``(sha256, byte_length, data)`` as checked text, or None (logged)."""
|
||||
if row is None:
|
||||
logger.warning(f"Workflow document {document_id} is missing")
|
||||
return None
|
||||
digest, length, data = bytes(row[0]), row[1], row[2]
|
||||
try:
|
||||
if not isinstance(length, int) or not 0 <= length <= MAX_DOCUMENT_BYTES:
|
||||
raise ValueError("invalid length")
|
||||
inflater = zlib.decompressobj()
|
||||
raw = inflater.decompress(bytes(data), length + 1)
|
||||
if len(raw) != length or inflater.unconsumed_tail or not inflater.eof:
|
||||
raise ValueError("length differs")
|
||||
if hashlib.sha256(raw).digest() != digest:
|
||||
raise ValueError("hash differs")
|
||||
return raw.decode("utf-8")
|
||||
except (zlib.error, ValueError, TypeError) as error:
|
||||
if not quiet:
|
||||
logger.warning(f"Workflow document {document_id} is unreadable: {error}")
|
||||
return None
|
||||
|
||||
|
||||
def read_stored_text(
|
||||
conn: sqlite3.Connection, inline: Any, document_id: Optional[int]
|
||||
) -> Any:
|
||||
"""The stored text of one column: inline when present, else the document.
|
||||
|
||||
Returns None when there is neither, or when the document is missing or
|
||||
fails its checks (logged). An inline value is returned unchanged.
|
||||
"""
|
||||
if inline is not None:
|
||||
return inline
|
||||
if document_id is None:
|
||||
return None
|
||||
return _decode(conn.execute(READ_SQL, (document_id,)).fetchone(), document_id)
|
||||
|
||||
|
||||
class DocumentReader:
|
||||
"""``read_stored_text`` with a small memo for list reads on one connection."""
|
||||
|
||||
MEMO = 32
|
||||
|
||||
def __init__(self, conn: sqlite3.Connection):
|
||||
self.conn = conn
|
||||
self._memo: "OrderedDict[int, Optional[str]]" = OrderedDict()
|
||||
|
||||
def text(self, inline: Any, document_id: Optional[int]) -> Any:
|
||||
if inline is not None or document_id is None:
|
||||
return read_stored_text(self.conn, inline, document_id)
|
||||
if document_id in self._memo:
|
||||
self._memo.move_to_end(document_id)
|
||||
return self._memo[document_id]
|
||||
text = read_stored_text(self.conn, None, document_id)
|
||||
self._memo[document_id] = text
|
||||
if len(self._memo) > self.MEMO:
|
||||
self._memo.popitem(last=False)
|
||||
return text
|
||||
|
||||
|
||||
def resolve_image_row(reader: DocumentReader, row) -> Dict[str, Any]:
|
||||
"""A ``generated_images`` row as a dict with its texts read back.
|
||||
|
||||
The id columns are dropped, so a payload built from it has the shape it
|
||||
had before the dedupe. Columns the row did not select stay absent.
|
||||
"""
|
||||
data = dict(row)
|
||||
for column, reference in TEXT_COLUMNS:
|
||||
has_reference = reference in data
|
||||
document_id = data.pop(reference, None)
|
||||
if column in data or has_reference:
|
||||
data[column] = reader.text(data.get(column), document_id)
|
||||
return data
|
||||
|
||||
|
||||
# ── garbage ───────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _chunks(values):
|
||||
values = list(values)
|
||||
for start in range(0, len(values), ID_BATCH):
|
||||
yield values[start : start + ID_BATCH]
|
||||
|
||||
|
||||
def image_documents(conn: sqlite3.Connection, where: str, params) -> Set[int]:
|
||||
"""The document ids the image rows matching ``where`` reference.
|
||||
|
||||
``where`` is a fixed condition written in this package (``id IN (...)``,
|
||||
``prompt_id IN (...)``), never caller data; values go in ``params``.
|
||||
"""
|
||||
ids: Set[int] = set()
|
||||
sql = IMAGE_DOCUMENTS_SQL + where
|
||||
for row in conn.execute(sql, list(params)): # nosec B608
|
||||
ids.update(value for value in row if value is not None)
|
||||
return ids
|
||||
|
||||
|
||||
def documents_of_rows(
|
||||
conn: sqlite3.Connection, column: str, values: Iterable[Any]
|
||||
) -> Set[int]:
|
||||
"""Document ids of the image rows whose ``column`` is in ``values``.
|
||||
|
||||
``column`` is ``id`` or ``prompt_id`` (the rows a prompt delete cascades to).
|
||||
"""
|
||||
if column not in ("id", "prompt_id"):
|
||||
raise ValueError(f"Unsupported column {column!r}")
|
||||
ids: Set[int] = set()
|
||||
for chunk in _chunks(values):
|
||||
marks = ",".join("?" for _ in chunk)
|
||||
ids |= image_documents(conn, f"{column} IN ({marks})", chunk)
|
||||
return ids
|
||||
|
||||
|
||||
def release_documents(conn: sqlite3.Connection, ids: Iterable[Optional[int]]) -> int:
|
||||
"""Delete the documents in ``ids`` that no image row references now.
|
||||
|
||||
Runs in the caller's transaction, right after the delete or update that
|
||||
dropped the references. Returns how many documents were removed.
|
||||
"""
|
||||
removed = 0
|
||||
for chunk in _chunks(sorted({i for i in ids if i is not None})):
|
||||
marks = ",".join("?" for _ in chunk)
|
||||
# Only "?" marks are interpolated.
|
||||
sql = RELEASE_SQL.format(marks=marks)
|
||||
removed += conn.execute(sql, chunk).rowcount # nosec B608
|
||||
return removed
|
||||
|
||||
|
||||
def sweep_documents(conn: sqlite3.Connection) -> int:
|
||||
"""Delete every document no image row references; returns the count."""
|
||||
return conn.execute(SWEEP_SQL).rowcount
|
||||
@@ -1,455 +0,0 @@
|
||||
"""The workflow storage migration and Reclaim space (#249 S6 carry-over).
|
||||
|
||||
``migrate_workflow_storage`` rewrites old-form image rows (workflow and
|
||||
prompt graph text inline) to the new form (a ``workflow_documents`` id,
|
||||
``database/workflow_documents.py``) in batches:
|
||||
|
||||
- each batch is one ``BEGIN IMMEDIATE`` transaction over at most
|
||||
``BATCH_ROWS`` rows or about ``BATCH_BYTES`` of text, found by a rowid
|
||||
range search (``NEXT_IDS_SQL``) and read one row at a time, so memory is
|
||||
bounded by one row plus the batch's distinct documents;
|
||||
- each text is stored (or its document reused) and read back through the
|
||||
one accessor before its row changes; a text that would not come back
|
||||
identical fails the batch (``DocumentMismatch``) and rolls it back whole;
|
||||
- a value that is not text (or is over ``MAX_DOCUMENT_BYTES``) stays inline
|
||||
and is counted as protected;
|
||||
- each row runs in its own savepoint: one that cannot be rewritten (text
|
||||
SQLite cannot decode, a text that would not come back) stays inline and
|
||||
is listed, and the rest of the batch goes on;
|
||||
- a batch ends after ``BATCH_SECONDS``, and the write lock is then left
|
||||
free for ``BATCH_PAUSE``, so other writers are never starved.
|
||||
|
||||
A crash or a stop between batches leaves committed batches in the new form
|
||||
and the rest in the old form; both read the same, and a rerun continues
|
||||
from the first old-form row. Only database rows change: no image or video
|
||||
file is opened.
|
||||
|
||||
``reclaim_space`` is the VACUUM afterwards: it sweeps unreferenced
|
||||
documents, rebuilds the file and checkpoints the WAL with TRUNCATE, so the
|
||||
freed pages go back to the disk.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sqlite3
|
||||
import time
|
||||
from contextlib import closing
|
||||
|
||||
from .workflow_documents import (
|
||||
TEXT_COLUMNS,
|
||||
DocumentMismatch,
|
||||
read_stored_text,
|
||||
release_documents,
|
||||
storable,
|
||||
store_text,
|
||||
sweep_documents,
|
||||
)
|
||||
|
||||
OPERATION = "workflow_storage"
|
||||
# The first apply's backup: ``<db>.pre-workflow-storage_<stamp>``, outside the
|
||||
# pruned ``<db>.backup_*`` pool. Only one is ever made; later applies take a
|
||||
# pooled safety backup.
|
||||
PRE_MIGRATION_LABEL = "pre-workflow-storage"
|
||||
BATCH_ROWS = 200
|
||||
BATCH_BYTES = 16 * 2**20 # characters of text rewritten per transaction
|
||||
# A batch holds the write lock: it ends after this long, and the loop then
|
||||
# leaves the lock free for BATCH_PAUSE, longer than SQLite's longest busy
|
||||
# poll (100 ms), so the image monitor and the run counting get in between.
|
||||
BATCH_SECONDS = 0.5
|
||||
BATCH_PAUSE = 0.15
|
||||
|
||||
|
||||
def _row_error(error) -> bool:
|
||||
"""Whether ``error`` is one row's own cause, which keeps only that row inline.
|
||||
|
||||
Text SQLite cannot decode for Python, and a text that would not come
|
||||
back identical. Anything else (a full disk, an I/O error, a malformed
|
||||
database, a bug) fails the whole batch. The decode error is matched by
|
||||
its message: ``sqlite_errorcode`` needs Python 3.11.
|
||||
"""
|
||||
if isinstance(error, (DocumentMismatch, UnicodeError)):
|
||||
return True
|
||||
return isinstance(error, sqlite3.OperationalError) and str(error).startswith(
|
||||
"Could not decode"
|
||||
)
|
||||
|
||||
|
||||
logger = logging.getLogger("prompt_manager.database.workflow_storage")
|
||||
_clock = time.monotonic
|
||||
|
||||
# typeof() reads a column's type from the record header: no overflow page.
|
||||
NEXT_IDS_SQL = (
|
||||
"SELECT id FROM generated_images WHERE id>? AND "
|
||||
"(typeof(workflow_data)!='null' OR typeof(prompt_metadata)!='null') "
|
||||
"ORDER BY id LIMIT ?"
|
||||
)
|
||||
ROW_SQL = (
|
||||
"SELECT workflow_data, prompt_metadata, workflow_document_id, "
|
||||
"prompt_document_id FROM generated_images WHERE id=?"
|
||||
)
|
||||
UPDATE_SQL = (
|
||||
"UPDATE generated_images SET {assignments} WHERE id=?" # columns are constants
|
||||
)
|
||||
# One pass over the record headers: rows with inline text, and inline
|
||||
# values that are not text (they stay inline).
|
||||
INLINE_COUNTS_SQL = (
|
||||
"SELECT "
|
||||
"SUM(typeof(workflow_data)='text' OR typeof(prompt_metadata)='text'), "
|
||||
"SUM(typeof(workflow_data) NOT IN ('null','text')) "
|
||||
"+ SUM(typeof(prompt_metadata) NOT IN ('null','text')) "
|
||||
"FROM generated_images"
|
||||
)
|
||||
# octet_length() (SQLite 3.43+) reads no overflow page either.
|
||||
INLINE_BYTES_SQL = (
|
||||
"SELECT SUM(octet_length(workflow_data)), SUM(octet_length(prompt_metadata)) "
|
||||
"FROM generated_images"
|
||||
)
|
||||
|
||||
QUERY_PLAN_STATEMENTS = {
|
||||
"next old-form rows": (NEXT_IDS_SQL, (0, BATCH_ROWS)),
|
||||
"one row": (ROW_SQL, (1,)),
|
||||
}
|
||||
|
||||
|
||||
def _rewrite_row(conn, image_id, row, verified, checked):
|
||||
"""Store the row's inline texts as documents; returns (changed, protected, size).
|
||||
|
||||
``verified`` maps document ids already read back in this batch to their
|
||||
text, and ``checked`` holds the reused ids ``store_text`` already
|
||||
checked, so a document shared by many rows is decompressed once.
|
||||
"""
|
||||
changes, replaced = {}, []
|
||||
protected = size = 0
|
||||
for index, (column, reference) in enumerate(TEXT_COLUMNS):
|
||||
value, current = row[index], row[index + len(TEXT_COLUMNS)]
|
||||
if value is None:
|
||||
continue
|
||||
if not storable(value):
|
||||
protected += 1
|
||||
continue
|
||||
document_id = store_text(conn, value, checked)
|
||||
if document_id not in verified:
|
||||
verified[document_id] = read_stored_text(conn, None, document_id)
|
||||
if verified[document_id] != value:
|
||||
raise DocumentMismatch(f"Image {image_id}: {column} would change")
|
||||
changes[column], changes[reference] = None, document_id
|
||||
if current is not None and current != document_id:
|
||||
replaced.append(current)
|
||||
size += len(value)
|
||||
if changes:
|
||||
assignments = ",".join(f"{key}=?" for key in changes)
|
||||
conn.execute(
|
||||
UPDATE_SQL.format(assignments=assignments),
|
||||
[*changes.values(), image_id],
|
||||
)
|
||||
release_documents(conn, replaced)
|
||||
return bool(changes), protected, size
|
||||
|
||||
|
||||
def migrate_batch(
|
||||
conn,
|
||||
after_id,
|
||||
max_rows=BATCH_ROWS,
|
||||
max_bytes=BATCH_BYTES,
|
||||
max_seconds=BATCH_SECONDS,
|
||||
):
|
||||
"""Rewrite the next batch after ``after_id`` in one transaction.
|
||||
|
||||
The batch ends after ``max_rows`` rows, about ``max_bytes`` of text or
|
||||
``max_seconds``, whichever comes first. Each row runs in its own
|
||||
savepoint: a row that cannot be rewritten (``_row_error``) is rolled
|
||||
back alone, stays inline and is listed in ``errors``; anything else rolls
|
||||
the whole batch back and propagates. Returns ``{last_id, rows,
|
||||
protected, errors, done}``: the last row looked at, rows rewritten,
|
||||
inline values kept (not text), the ids kept by an error and whether
|
||||
there was nothing left to look at.
|
||||
"""
|
||||
if conn.in_transaction:
|
||||
raise RuntimeError("The migration needs its own transaction")
|
||||
started = _clock()
|
||||
conn.execute("BEGIN IMMEDIATE")
|
||||
try:
|
||||
ids = [row[0] for row in conn.execute(NEXT_IDS_SQL, (after_id, max_rows))]
|
||||
result = {
|
||||
"last_id": after_id,
|
||||
"rows": 0,
|
||||
"protected": 0,
|
||||
"errors": [],
|
||||
"done": not ids,
|
||||
}
|
||||
verified, checked, size = {}, set(), 0
|
||||
for image_id in ids:
|
||||
conn.execute("SAVEPOINT workflow_row")
|
||||
try:
|
||||
row = conn.execute(ROW_SQL, (image_id,)).fetchone()
|
||||
changed, protected, row_size = _rewrite_row(
|
||||
conn, image_id, row, verified, checked
|
||||
)
|
||||
except Exception as error:
|
||||
if not _row_error(error):
|
||||
raise
|
||||
conn.execute("ROLLBACK TO workflow_row")
|
||||
conn.execute("RELEASE workflow_row")
|
||||
verified.clear() # the row's documents were rolled back too
|
||||
checked.clear()
|
||||
logger.warning(f"Image record {image_id} kept inline: {error}")
|
||||
result["errors"].append(image_id)
|
||||
changed, protected, row_size = False, 0, 0
|
||||
else:
|
||||
conn.execute("RELEASE workflow_row")
|
||||
result["rows"] += int(changed)
|
||||
result["protected"] += protected
|
||||
result["last_id"] = image_id
|
||||
size += row_size
|
||||
if size >= max_bytes or _clock() - started >= max_seconds:
|
||||
break
|
||||
conn.commit()
|
||||
return result
|
||||
except BaseException:
|
||||
conn.rollback()
|
||||
raise
|
||||
|
||||
|
||||
def pre_migration_backups(db_path):
|
||||
"""The kept pre-migration backups beside ``db_path``, earliest first."""
|
||||
db_path = os.path.abspath(db_path)
|
||||
prefix = os.path.basename(db_path) + "." + PRE_MIGRATION_LABEL + "_"
|
||||
folder = os.path.dirname(db_path)
|
||||
try:
|
||||
names = os.listdir(folder)
|
||||
except OSError:
|
||||
return []
|
||||
return sorted(
|
||||
name
|
||||
for name in names
|
||||
if name.startswith(prefix) and os.path.isfile(os.path.join(folder, name))
|
||||
)
|
||||
|
||||
|
||||
def _file_bytes(db_path):
|
||||
return sum(
|
||||
os.path.getsize(path)
|
||||
for path in (db_path, db_path + "-wal")
|
||||
if os.path.exists(path)
|
||||
)
|
||||
|
||||
|
||||
class WorkflowStorageOperations:
|
||||
"""PromptDatabase mixin: the workflow storage migration and Reclaim space."""
|
||||
|
||||
def migrate_workflow_storage(
|
||||
self,
|
||||
*,
|
||||
on_batch=None,
|
||||
should_stop=None,
|
||||
max_rows=BATCH_ROWS,
|
||||
max_bytes=BATCH_BYTES,
|
||||
) -> dict:
|
||||
"""Rewrite every old-form row, batch by batch (blocking).
|
||||
|
||||
``on_batch(progress)`` runs after each committed batch with the
|
||||
running totals ``{rows, protected, batches, last_id}``;
|
||||
``should_stop()`` is asked before each batch. Each committed batch is
|
||||
followed by ``BATCH_PAUSE`` with the write lock free. Returns
|
||||
``{rows, protected, batches, stopped, errors, documents_removed}``
|
||||
(``errors``: the ids of rows kept inline by an error). Unreferenced
|
||||
documents are swept once every row was looked at.
|
||||
"""
|
||||
totals = {"rows": 0, "protected": 0, "batches": 0, "last_id": 0}
|
||||
errors = []
|
||||
stopped = False
|
||||
while True:
|
||||
if should_stop is not None and should_stop():
|
||||
stopped = True
|
||||
break
|
||||
batch = self.migrate_workflow_batch(totals["last_id"], max_rows, max_bytes)
|
||||
if batch["done"]:
|
||||
break
|
||||
time.sleep(BATCH_PAUSE) # leave the write lock to other writers
|
||||
errors += batch["errors"]
|
||||
totals["rows"] += batch["rows"]
|
||||
totals["protected"] += batch["protected"]
|
||||
totals["batches"] += 1
|
||||
totals["last_id"] = batch["last_id"]
|
||||
if on_batch is not None:
|
||||
on_batch(dict(totals))
|
||||
removed = 0 if stopped else self.sweep_workflow_documents()
|
||||
return {
|
||||
"rows": totals["rows"],
|
||||
"protected": totals["protected"],
|
||||
"batches": totals["batches"],
|
||||
"stopped": stopped,
|
||||
"errors": errors,
|
||||
"documents_removed": removed,
|
||||
}
|
||||
|
||||
def migrate_workflow_batch(
|
||||
self, after_id, max_rows=BATCH_ROWS, max_bytes=BATCH_BYTES
|
||||
) -> dict:
|
||||
"""One batch on the calling thread's connection (``migrate_batch``)."""
|
||||
return migrate_batch(self.model.get_connection(), after_id, max_rows, max_bytes)
|
||||
|
||||
def sweep_workflow_documents(self) -> int:
|
||||
"""Remove every document no image row references (one transaction)."""
|
||||
conn = self.model.get_connection()
|
||||
with conn:
|
||||
return sweep_documents(conn)
|
||||
|
||||
def workflow_storage_status(self) -> dict:
|
||||
"""``workflow_storage_counts`` on one read-only snapshot (blocking)."""
|
||||
from .maintenance import _read_snapshot
|
||||
|
||||
return _read_snapshot(self.model.db_path, self.workflow_storage_counts)
|
||||
|
||||
def workflow_storage_counts(self, conn) -> dict:
|
||||
"""``{rows, protected, documents, inline_bytes}`` on ``conn`` (reads only).
|
||||
|
||||
``inline_bytes`` is None where SQLite has no ``octet_length``.
|
||||
"""
|
||||
rows, protected = conn.execute(INLINE_COUNTS_SQL).fetchone()
|
||||
documents = conn.execute("SELECT COUNT(*) FROM workflow_documents").fetchone()
|
||||
inline_bytes = None
|
||||
if sqlite3.sqlite_version_info >= (3, 43, 0):
|
||||
row = conn.execute(INLINE_BYTES_SQL).fetchone()
|
||||
inline_bytes = sum(value or 0 for value in row)
|
||||
return {
|
||||
"rows": rows or 0,
|
||||
"protected": protected or 0,
|
||||
"documents": documents[0],
|
||||
"inline_bytes": inline_bytes,
|
||||
}
|
||||
|
||||
def reclaim_space(self) -> dict:
|
||||
"""Sweep, VACUUM and truncate the WAL (blocking); file bytes before/after.
|
||||
|
||||
The caller holds the registry's ``exclusive`` slot. Raises whatever
|
||||
stopped VACUUM; the database is then unchanged.
|
||||
"""
|
||||
db_path = self.model.db_path
|
||||
removed = self.sweep_workflow_documents()
|
||||
before = _file_bytes(db_path)
|
||||
self.model.vacuum_database()
|
||||
with closing(sqlite3.connect(db_path)) as checkpoint:
|
||||
busy = checkpoint.execute("PRAGMA wal_checkpoint(TRUNCATE)").fetchone()[0]
|
||||
return {
|
||||
"before_bytes": before,
|
||||
"after_bytes": _file_bytes(db_path),
|
||||
"documents_removed": removed,
|
||||
"checkpoint_busy": bool(busy),
|
||||
}
|
||||
|
||||
# ── preview ───────────────────────────────────────────────────────
|
||||
|
||||
def _plan_workflow_storage(self):
|
||||
"""The reviewed plan of the migration (one read-only snapshot)."""
|
||||
from .maintenance import (
|
||||
MaintenanceDraft,
|
||||
_Fingerprint,
|
||||
_megabytes,
|
||||
_plural,
|
||||
_read_snapshot,
|
||||
)
|
||||
|
||||
def read(conn):
|
||||
sizes = [
|
||||
conn.execute(f"PRAGMA {name}").fetchone()[0]
|
||||
for name in ("page_size", "page_count", "freelist_count")
|
||||
]
|
||||
return sizes, self.workflow_storage_counts(conn)
|
||||
|
||||
(page_size, page_count, free_pages), counts = _read_snapshot(
|
||||
self.model.db_path, read
|
||||
)
|
||||
size = page_size * page_count
|
||||
details = [
|
||||
{
|
||||
"label": "Database file",
|
||||
"detail": f"{_megabytes(size)}, of which "
|
||||
f"{_megabytes(free_pages * page_size)} is free pages",
|
||||
},
|
||||
{
|
||||
"label": "Image records",
|
||||
"detail": f"{_plural(counts['rows'], 'record')} keep their workflow "
|
||||
"or prompt graph inline; each distinct text is stored once, "
|
||||
"compressed",
|
||||
},
|
||||
]
|
||||
if not counts["rows"]:
|
||||
details[1]["detail"] = NOTHING_TO_DO
|
||||
if counts["inline_bytes"] is not None:
|
||||
details.append(
|
||||
{
|
||||
"label": "Inline text",
|
||||
"detail": f"{_megabytes(counts['inline_bytes'])} now",
|
||||
}
|
||||
)
|
||||
if counts["documents"]:
|
||||
details.append(
|
||||
{
|
||||
"label": "Already stored once",
|
||||
"detail": _plural(counts["documents"], "workflow document"),
|
||||
}
|
||||
)
|
||||
name = os.path.basename(self.model.db_path)
|
||||
kept = pre_migration_backups(self.model.db_path)
|
||||
backup = (
|
||||
POOLED_BACKUP_NOTE.format(
|
||||
name=f"{name}.backup_<date and time>", kept=kept[0]
|
||||
)
|
||||
if kept
|
||||
else BACKUP_NOTE.format(
|
||||
name=f"{name}.{PRE_MIGRATION_LABEL}_<date and time>"
|
||||
)
|
||||
)
|
||||
details += [
|
||||
{"label": "Backup", "detail": backup},
|
||||
{"label": "Reclaim space", "detail": RECLAIM_NOTE},
|
||||
{"label": "Older versions", "detail": OLDER_VERSIONS_NOTE},
|
||||
]
|
||||
protected = counts["protected"]
|
||||
return MaintenanceDraft(
|
||||
operation=OPERATION,
|
||||
fingerprint=_Fingerprint(OPERATION).hexdigest(),
|
||||
targets=[],
|
||||
scope={
|
||||
"count": counts["rows"],
|
||||
"label": (
|
||||
_plural(counts["rows"], "image record") + " to rewrite"
|
||||
if counts["rows"]
|
||||
else "Nothing to rewrite"
|
||||
),
|
||||
},
|
||||
protected={
|
||||
"count": protected,
|
||||
"reasons": (
|
||||
[f"{_plural(protected, 'stored value')} not text; kept as is"]
|
||||
if protected
|
||||
else []
|
||||
),
|
||||
},
|
||||
media={"deletes_files": False, "label": MEDIA_NOTE},
|
||||
details=details,
|
||||
)
|
||||
|
||||
|
||||
MEDIA_NOTE = "No image or video file is opened or changed; only the database."
|
||||
RECLAIM_NOTE = (
|
||||
"When every record is rewritten, Reclaim space runs alone to give the freed "
|
||||
"space back to the disk (VACUUM). The database is locked while it runs."
|
||||
)
|
||||
BACKUP_NOTE = (
|
||||
"Always taken first, whatever the safety backup setting: {name} beside "
|
||||
"the database, verified, and kept until you remove it (it is never pruned)."
|
||||
)
|
||||
POOLED_BACKUP_NOTE = (
|
||||
"A safety backup is taken first, whatever the safety backup setting: {name} "
|
||||
"among the pooled backups; the pre-migration backup {kept} is already kept."
|
||||
)
|
||||
NOTHING_TO_DO = "Every image record already stores its workflow once."
|
||||
OLDER_VERSIONS_NOTE = (
|
||||
"This version stores the workflows of images it records in the new form, "
|
||||
"whether or not this job runs, and this job rewrites the rest. An older "
|
||||
"Prompt Manager version shows those images without their workflow until you "
|
||||
"update again. The way back to an older version is the earliest "
|
||||
"pre-workflow-storage backup."
|
||||
)
|
||||
+2
-39
@@ -9,15 +9,10 @@ Split into domain-specific mixins for maintainability:
|
||||
- JobRoutesMixin: the session job registry behind Activity (py/api/jobs.py)
|
||||
- ImageTagJobsMixin: bulk Add tags / Auto-tag as image-tags jobs
|
||||
- ThumbnailJobsMixin: thumbnail plan, thumbnail job and the gallery stream on it
|
||||
- AutoThumbnailMixin: previews made on first view, on capture and after a scan
|
||||
- ReviewRoutesMixin: the Needs review queue, item detail and decisions
|
||||
- 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
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
@@ -50,13 +45,7 @@ from .settings_routes import SettingsRoutesMixin
|
||||
from .job_routes import JobRoutesMixin
|
||||
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
|
||||
from .maintenance_plans import PlanStore
|
||||
from .jobs import JobRegistry
|
||||
from .logging_routes import LoggingRoutesMixin
|
||||
from .autotag_routes import AutotagRoutesMixin
|
||||
@@ -295,12 +284,7 @@ class PromptManagerAPI(
|
||||
JobRoutesMixin,
|
||||
ImageTagJobsMixin,
|
||||
ThumbnailJobsMixin,
|
||||
AutoThumbnailMixin,
|
||||
SystemRoutesMixin,
|
||||
LibraryLocationRoutesMixin,
|
||||
MaintenanceRoutesMixin,
|
||||
WorkflowStorageRoutesMixin,
|
||||
DuplicateFilesMixin,
|
||||
LoggingRoutesMixin,
|
||||
AutotagRoutesMixin,
|
||||
LoraIntegrationMixin,
|
||||
@@ -357,31 +341,13 @@ class PromptManagerAPI(
|
||||
self._jobs = JobRegistry(self.logger)
|
||||
# The latest metadata scan job (also in self._jobs), if one was started.
|
||||
self._scan_job = None
|
||||
# Duplicate-file finds by job id, with the groups each one holds for
|
||||
# review (py/api/duplicate_files.py); the last few only.
|
||||
self._duplicate_finds = {}
|
||||
# The current duplicate scan (a py/api/jobs.Job outside the registry).
|
||||
self._duplicate_job = None
|
||||
# The last System health report this session and the database
|
||||
# generation it describes, bumped when a restore replaces self.db
|
||||
# (py/api/system.py).
|
||||
self._health_report = None
|
||||
self._db_generation = 0
|
||||
# Reviewed maintenance plans under single-use tokens; a restart
|
||||
# forgets them (py/api/maintenance.py).
|
||||
self._maintenance_plans = PlanStore()
|
||||
|
||||
# A safety backup interrupted by a crash leaves its temp copy behind.
|
||||
try:
|
||||
from .maintenance import sweep_safety_temps
|
||||
|
||||
swept = sweep_safety_temps(self.db.model.db_path)
|
||||
if swept:
|
||||
self.logger.info(f"Removed {len(swept)} interrupted backup copies")
|
||||
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:
|
||||
@@ -542,9 +508,6 @@ 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)
|
||||
self._register_autotag_routes(routes)
|
||||
self._register_lora_routes(routes)
|
||||
|
||||
+145
-259
@@ -24,9 +24,9 @@ from .jobs import (
|
||||
SCAN_GROUPS,
|
||||
SCAN_SCOPES,
|
||||
DUPLICATE_CLEANUP_JOB_KIND,
|
||||
DUPLICATES_JOB_KIND,
|
||||
MAINTENANCE_JOB_KIND,
|
||||
RESTORE_JOB_KIND,
|
||||
Job,
|
||||
JobBusy,
|
||||
)
|
||||
|
||||
@@ -69,9 +69,7 @@ except ImportError:
|
||||
IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp", ".tiff")
|
||||
VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv", ".webm", ".m4v", ".wmv")
|
||||
SCAN_BATCH_SIZE = 50 # Files per executor call during the output scan
|
||||
# '<db>.backup_*' files kept next to the live database: one pool for the
|
||||
# restore backups and the maintenance safety backups (py/api/maintenance.py).
|
||||
MAX_SAFETY_BACKUPS = 10
|
||||
MAX_SAFETY_BACKUPS = 10 # '<db>.backup_*' files kept next to the live database
|
||||
MAX_SCAN_FILES = 50_000 # Upper bound on media files one scan will look at
|
||||
SCAN_COUNT_KEYS = (
|
||||
"processed",
|
||||
@@ -92,7 +90,6 @@ MAX_SCAN_DEPTH = 12 # Directory levels below an output root that are scanned
|
||||
# the common single-f TIFF spelling. Anything else inside the output
|
||||
# directory (databases, scripts, notes) is never a "duplicate image".
|
||||
DELETABLE_MEDIA_EXTENSIONS = frozenset(IMAGE_EXTENSIONS + VIDEO_EXTENSIONS + (".tif",))
|
||||
DUPLICATE_FIND_LEGACY_TITLE = "Find duplicate image files (gallery page)"
|
||||
|
||||
|
||||
def _collect_media_files(output_dirs, extensions):
|
||||
@@ -135,52 +132,11 @@ def _collect_media_files(output_dirs, extensions):
|
||||
return found
|
||||
|
||||
|
||||
def _physical_file_key(path, stat):
|
||||
"""What tells two physical files apart: ``(st_dev, st_ino)``.
|
||||
|
||||
A filesystem that reports no inode (``st_ino`` 0, some network or FAT
|
||||
volumes on Windows) falls back to the resolved, case-folded path.
|
||||
"""
|
||||
if stat.st_ino:
|
||||
return ("inode", stat.st_dev, stat.st_ino)
|
||||
return ("path", os.path.normcase(os.path.realpath(str(path))))
|
||||
|
||||
|
||||
def _thumbnail_rel_path(rel_path, thumbnail_ext):
|
||||
"""Relative path of the thumbnail the gallery would generate for ``rel_path``."""
|
||||
return f"thumbnails/{rel_path.with_suffix('').as_posix()}_thumb{thumbnail_ext}"
|
||||
|
||||
|
||||
def _media_thumbnail_rel_path(rel_path):
|
||||
"""The thumbnail of a media file: a video's preview is a ``.jpg``."""
|
||||
extension = rel_path.suffix.lower()
|
||||
return _thumbnail_rel_path(
|
||||
rel_path, ".jpg" if extension in VIDEO_EXTENSIONS else rel_path.suffix
|
||||
)
|
||||
|
||||
|
||||
def _remove_thumbnail(output_path, rel_path, logger=None):
|
||||
"""Remove the thumbnail of ``rel_path`` under ``output_path``, if any.
|
||||
|
||||
Shared by the legacy duplicate delete and the reviewed removal job. A
|
||||
thumbnail that is a link, or one that cannot be removed, is left alone
|
||||
(a stale preview never fails a removal). Returns whether one was removed.
|
||||
"""
|
||||
output_path = Path(output_path)
|
||||
thumbnail_path = output_path / _media_thumbnail_rel_path(Path(rel_path))
|
||||
try:
|
||||
if thumbnail_path.is_symlink() or not thumbnail_path.is_file():
|
||||
return False
|
||||
os.remove(thumbnail_path)
|
||||
except OSError as e:
|
||||
if logger is not None:
|
||||
logger.warning(f"Could not delete thumbnail for {rel_path}: {e}")
|
||||
return False
|
||||
if logger is not None:
|
||||
logger.debug(f"Deleted associated thumbnail: {thumbnail_path}")
|
||||
return True
|
||||
|
||||
|
||||
async def _read_json(request):
|
||||
"""Parse a JSON object body through the package's capped reader.
|
||||
|
||||
@@ -605,21 +561,12 @@ class AdminRoutesMixin:
|
||||
LONG_JOB_DUPLICATES = "duplicates"
|
||||
|
||||
async def scan_duplicates_endpoint(self, request):
|
||||
"""Scan for duplicate images without removing them.
|
||||
|
||||
Holds the ``duplicates`` kind on the registry (a ``read`` hold, so it
|
||||
is listed in Activity), so it is single flight with every duplicate
|
||||
find, and it is refused while a duplicate removal runs.
|
||||
"""
|
||||
cleanup = self._jobs.active(DUPLICATE_CLEANUP_JOB_KIND)
|
||||
if cleanup is not None:
|
||||
return self._long_job_busy_response(cleanup.title)
|
||||
return await self._run_held(
|
||||
DUPLICATES_JOB_KIND,
|
||||
DUPLICATE_FIND_LEGACY_TITLE,
|
||||
"read",
|
||||
functools.partial(self._scan_duplicates, request),
|
||||
)
|
||||
"""Scan for duplicate images without removing them (single flight)."""
|
||||
job = self._acquire_long_job(self.LONG_JOB_DUPLICATES)
|
||||
if job is None:
|
||||
return self._long_job_busy_response("Duplicate scan")
|
||||
async with job:
|
||||
return await self._scan_duplicates(request)
|
||||
|
||||
async def _scan_duplicates(self, request):
|
||||
try:
|
||||
@@ -641,29 +588,24 @@ class AdminRoutesMixin:
|
||||
)
|
||||
|
||||
async def scan_duplicates_stream(self, request):
|
||||
"""Stream a duplicate find to the legacy gallery (its event shapes).
|
||||
"""Run the duplicate scan as a background job and stream its progress.
|
||||
|
||||
Starts the registry's ``duplicates`` job over the gallery's output
|
||||
folder, or attaches to the running one when it covers that folder
|
||||
(one started in the workspace, for instance). A held legacy scan, a
|
||||
find that does not cover the folder or a running removal answers the
|
||||
legacy 409. The job outlives the request.
|
||||
A request while a scan is running attaches to it. The job holds the
|
||||
duplicates lock, so a delete is refused (409) until it finishes.
|
||||
"""
|
||||
job = self._jobs.active(DUPLICATES_JOB_KIND)
|
||||
if job is None:
|
||||
try:
|
||||
job = self._start_duplicate_find(origin="gallery")
|
||||
except JobBusy as busy:
|
||||
return self._long_job_busy_response(busy.running.title)
|
||||
else:
|
||||
legacy_dir = await self._run_in_executor(self._find_comfyui_output_dir)
|
||||
if not self._find_covers(job, legacy_dir):
|
||||
return self._long_job_busy_response(job.title)
|
||||
job = self._duplicate_job
|
||||
if job is None or job.finished:
|
||||
guard = self._acquire_long_job(self.LONG_JOB_DUPLICATES)
|
||||
if guard is None:
|
||||
return self._long_job_busy_response("Duplicate scan or delete")
|
||||
job = Job("duplicates", "Duplicate scan")
|
||||
self._duplicate_job = job
|
||||
job.task = asyncio.ensure_future(self._run_duplicate_job(job, guard))
|
||||
return await self._stream_scan_job(request, job)
|
||||
|
||||
async def scan_duplicates_status(self, request):
|
||||
"""Whether a duplicate find is running, and its latest legacy event."""
|
||||
job = self._latest_duplicate_find()
|
||||
"""Whether a duplicate scan is running, and its latest event."""
|
||||
job = self._duplicate_job
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
@@ -672,6 +614,62 @@ class AdminRoutesMixin:
|
||||
}
|
||||
)
|
||||
|
||||
async def _run_duplicate_job(self, job, guard):
|
||||
loop = asyncio.get_running_loop()
|
||||
counts = {"processed": 0}
|
||||
|
||||
def report(done, total):
|
||||
# Called from worker threads; hand the event to the loop.
|
||||
counts["processed"] = done
|
||||
pct = int(done / total * 100) if total else 100
|
||||
loop.call_soon_threadsafe(
|
||||
job.publish,
|
||||
{
|
||||
"type": "progress",
|
||||
"progress": pct,
|
||||
"status": f"Hashing file {done}/{total}...",
|
||||
"processed": done,
|
||||
"found": 0,
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
async with guard:
|
||||
job.publish(
|
||||
{
|
||||
"type": "progress",
|
||||
"progress": 0,
|
||||
"status": "Collecting media files...",
|
||||
"processed": 0,
|
||||
"found": 0,
|
||||
}
|
||||
)
|
||||
output_dir = await self._run_in_executor(self._find_comfyui_output_dir)
|
||||
groups = await self._run_in_executor(
|
||||
self._find_duplicate_images_sync, output_dir, report
|
||||
)
|
||||
job.publish(
|
||||
{
|
||||
"type": "complete",
|
||||
"processed": counts["processed"],
|
||||
"found": len(groups),
|
||||
"duplicates": groups,
|
||||
"message": f"Found {len(groups)} groups of duplicate images",
|
||||
}
|
||||
)
|
||||
except Exception:
|
||||
self.logger.exception("Duplicate scan error")
|
||||
job.publish(
|
||||
{
|
||||
"type": "error",
|
||||
"message": (
|
||||
"An internal error occurred. Check server logs for details."
|
||||
),
|
||||
}
|
||||
)
|
||||
finally:
|
||||
job.finish()
|
||||
|
||||
async def cleanup_duplicates_endpoint(self, request):
|
||||
"""Cleanup duplicate prompts endpoint (holds the registry: records)."""
|
||||
return await self._run_held(
|
||||
@@ -708,109 +706,31 @@ class AdminRoutesMixin:
|
||||
def _find_duplicate_images_sync(self, output_dir, progress=None):
|
||||
"""Group media files under ``output_dir`` by content hash (blocking).
|
||||
|
||||
The legacy response shape over ``_hash_media_groups``, which the
|
||||
workspace's duplicate find shares.
|
||||
"""
|
||||
if not output_dir or not os.path.isdir(output_dir):
|
||||
self.logger.warning("ComfyUI output directory not found")
|
||||
return []
|
||||
groups = self._hash_media_groups([os.path.abspath(output_dir)], progress)
|
||||
duplicates = self._legacy_duplicate_groups(groups, output_dir)
|
||||
self.logger.info(f"Found {len(duplicates)} groups of duplicate images")
|
||||
return duplicates
|
||||
|
||||
def _legacy_duplicate_groups(self, groups, output_dir):
|
||||
"""``_hash_media_groups`` groups as the legacy gallery reads them.
|
||||
|
||||
Only the files under ``output_dir`` count (paths are relative to it);
|
||||
a group needs two of them. A file that vanished since it was hashed
|
||||
is left out.
|
||||
"""
|
||||
output_path = Path(os.path.abspath(output_dir))
|
||||
base = os.path.normcase(str(output_path)).rstrip("\\/") + os.sep
|
||||
duplicates = []
|
||||
for group in groups:
|
||||
images = []
|
||||
for record in group["files"]:
|
||||
if not os.path.normcase(record["abs"]).startswith(base):
|
||||
continue
|
||||
try:
|
||||
images.append(
|
||||
self._describe_media_file(
|
||||
Path(record["abs"]), output_path, file_hash=group["hash"]
|
||||
)
|
||||
)
|
||||
except (OSError, ValueError) as e:
|
||||
self.logger.error(
|
||||
f"Error describing {Path(record['abs']).name}: {e}"
|
||||
)
|
||||
if len(images) > 1:
|
||||
duplicates.append(
|
||||
{"hash": group["hash"], "images": images, "count": len(images)}
|
||||
)
|
||||
return duplicates
|
||||
|
||||
def _hash_media_groups(self, roots, progress=None, stop=None):
|
||||
"""Content-identical media files under ``roots`` (blocking).
|
||||
|
||||
Each root is walked as ``_collect_media_files`` walks (no symlinks,
|
||||
no ``thumbnails``, bounded). A physical file is listed once, under
|
||||
the first root that holds it: files are told apart by ``(st_dev,
|
||||
st_ino)``, so a source listed twice under two spellings (a symlink, a
|
||||
bind mount, a case variant), a source inside another one or a
|
||||
hardlink is one file, never a copy of itself. At most MAX_SCAN_FILES
|
||||
files are read. Files
|
||||
are hashed on ``PromptManagerConfig.WORKER_THREADS`` threads;
|
||||
Files are hashed on ``PromptManagerConfig.WORKER_THREADS`` threads.
|
||||
``progress(done, total)`` is called from those threads about every
|
||||
percent and once at the end, and ``stop()`` true skips the files not
|
||||
yet read. Returns ``[{hash, size, files:[{slot, abs, rel, size,
|
||||
mtime, mtime_ns, inode}]}]`` for every hash shared by two or more
|
||||
files (``inode`` is ``[st_dev, st_ino]``), ``slot``
|
||||
being the root's place in ``roots`` and the files oldest first.
|
||||
percent and once at the end.
|
||||
"""
|
||||
from ..config import PromptManagerConfig
|
||||
|
||||
if not output_dir or not os.path.isdir(output_dir):
|
||||
self.logger.warning("ComfyUI output directory not found")
|
||||
return []
|
||||
|
||||
output_path = Path(output_dir)
|
||||
extensions = IMAGE_EXTENSIONS + VIDEO_EXTENSIONS
|
||||
media_files = []
|
||||
seen = set()
|
||||
for slot, root in enumerate(roots):
|
||||
for path in _collect_media_files([root], extensions):
|
||||
try:
|
||||
# os.stat, not Path.stat: one stat call for every check
|
||||
# (Python 3.10's Path.stat bypasses a patched os.stat).
|
||||
stat = os.stat(path)
|
||||
except OSError:
|
||||
continue
|
||||
key = _physical_file_key(path, stat)
|
||||
if key not in seen:
|
||||
seen.add(key)
|
||||
media_files.append((slot, root, path, stat))
|
||||
if len(media_files) >= MAX_SCAN_FILES:
|
||||
break
|
||||
media_files = media_files[:MAX_SCAN_FILES]
|
||||
media_files = _collect_media_files([output_path], extensions)
|
||||
total = len(media_files)
|
||||
self.logger.info(f"Found {total} media files to analyze")
|
||||
step = max(1, total // 100)
|
||||
counter = {"done": 0}
|
||||
lock = threading.Lock()
|
||||
|
||||
def describe(item):
|
||||
slot, root, media_path, stat = item
|
||||
info = None
|
||||
if stop is None or not stop():
|
||||
try:
|
||||
info = {
|
||||
"slot": slot,
|
||||
"abs": str(media_path),
|
||||
"rel": Path(os.path.relpath(str(media_path), root)).as_posix(),
|
||||
"size": stat.st_size,
|
||||
"mtime": stat.st_mtime,
|
||||
"mtime_ns": stat.st_mtime_ns,
|
||||
"inode": [stat.st_dev, stat.st_ino],
|
||||
"hash": self._calculate_file_hash(media_path),
|
||||
}
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error processing file {media_path.name}: {e}")
|
||||
def describe(media_path):
|
||||
try:
|
||||
info = self._describe_media_file(media_path, output_path)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error processing file {media_path.name}: {e}")
|
||||
info = None
|
||||
with lock:
|
||||
counter["done"] += 1
|
||||
done = counter["done"]
|
||||
@@ -824,22 +744,24 @@ class AdminRoutesMixin:
|
||||
|
||||
infos = map_parallel(describe, media_files, PromptManagerConfig.WORKER_THREADS)
|
||||
|
||||
by_hash = {}
|
||||
file_hashes = {}
|
||||
for info in infos:
|
||||
if info is not None:
|
||||
by_hash.setdefault(info.pop("hash"), []).append(info)
|
||||
groups = []
|
||||
for file_hash, files in by_hash.items():
|
||||
if len(files) > 1:
|
||||
files.sort(key=lambda record: (record["mtime"], record["abs"]))
|
||||
groups.append(
|
||||
{"hash": file_hash, "size": files[0]["size"], "files": files}
|
||||
)
|
||||
return groups
|
||||
file_hashes.setdefault(info["hash"], []).append(info)
|
||||
|
||||
def _describe_media_file(self, media_path, output_path, file_hash=None):
|
||||
"""Describe one media file for the duplicates response (hashing it
|
||||
unless ``file_hash`` is given)."""
|
||||
duplicates = []
|
||||
for file_hash, images in file_hashes.items():
|
||||
if len(images) > 1:
|
||||
images.sort(key=lambda x: x["modified_time"])
|
||||
duplicates.append(
|
||||
{"hash": file_hash, "images": images, "count": len(images)}
|
||||
)
|
||||
|
||||
self.logger.info(f"Found {len(duplicates)} groups of duplicate images")
|
||||
return duplicates
|
||||
|
||||
def _describe_media_file(self, media_path, output_path):
|
||||
"""Hash one media file and describe it for the duplicates response."""
|
||||
stat = media_path.stat()
|
||||
rel_path = media_path.relative_to(output_path)
|
||||
extension = media_path.suffix.lower()
|
||||
@@ -860,7 +782,7 @@ class AdminRoutesMixin:
|
||||
"extension": extension,
|
||||
"media_type": "video" if is_video else "image",
|
||||
"is_video": is_video,
|
||||
"hash": file_hash or self._calculate_file_hash(media_path),
|
||||
"hash": self._calculate_file_hash(media_path),
|
||||
}
|
||||
|
||||
def _calculate_file_hash(self, file_path):
|
||||
@@ -874,9 +796,9 @@ class AdminRoutesMixin:
|
||||
async def delete_duplicate_images_endpoint(self, request):
|
||||
"""Delete duplicate image files from disk.
|
||||
|
||||
Refused while a duplicate find runs, and a records hold on the
|
||||
registry, so no scan, restore, removal or tagging job runs meanwhile.
|
||||
A body without paths is refused before the hold (it is no job).
|
||||
Single flight with the duplicate scan, and a records hold on the
|
||||
registry, so no scan, restore or tagging job runs meanwhile. A body
|
||||
without paths is refused before the hold (it is no job).
|
||||
"""
|
||||
data, error_response = await _read_json(request)
|
||||
if error_response is not None:
|
||||
@@ -887,14 +809,20 @@ class AdminRoutesMixin:
|
||||
{"success": False, "error": "No image paths provided"},
|
||||
status=400,
|
||||
)
|
||||
find = self._jobs.active(DUPLICATES_JOB_KIND)
|
||||
if find is not None:
|
||||
return self._long_job_busy_response(find.title)
|
||||
guard = self._acquire_long_job(self.LONG_JOB_DUPLICATES)
|
||||
if guard is None:
|
||||
return self._long_job_busy_response("Duplicate scan or delete")
|
||||
|
||||
async def delete():
|
||||
# No await since the check above: the lock is still free.
|
||||
async with guard:
|
||||
return await self._delete_duplicate_images(image_paths)
|
||||
|
||||
return await self._run_held(
|
||||
DUPLICATE_CLEANUP_JOB_KIND,
|
||||
"Duplicate image delete (gallery page)",
|
||||
"records",
|
||||
functools.partial(self._delete_duplicate_images, image_paths),
|
||||
delete,
|
||||
)
|
||||
|
||||
async def _delete_duplicate_images(self, image_paths):
|
||||
@@ -974,7 +902,14 @@ class AdminRoutesMixin:
|
||||
|
||||
os.remove(file_path)
|
||||
self.logger.info(f"Deleted duplicate image: {image_path}")
|
||||
_remove_thumbnail(output_path.resolve(), rel_path, self.logger)
|
||||
|
||||
thumbnail_path = output_path / _thumbnail_rel_path(rel_path, file_path.suffix)
|
||||
try:
|
||||
if thumbnail_path.exists():
|
||||
os.remove(thumbnail_path)
|
||||
self.logger.debug(f"Deleted associated thumbnail: {thumbnail_path}")
|
||||
except OSError as e:
|
||||
self.logger.warning(f"Could not delete thumbnail for {image_path}: {e}")
|
||||
return None
|
||||
|
||||
async def get_statistics(self, request):
|
||||
@@ -1177,10 +1112,8 @@ class AdminRoutesMixin:
|
||||
db.cleanup_duplicates(), "duplicate prompts"
|
||||
),
|
||||
"vacuum": vacuum,
|
||||
# Only under sources that can be read now (never an unmounted drive).
|
||||
"cleanup_orphaned_images": lambda: count_result(
|
||||
db.cleanup_missing_images(self._health_roots()),
|
||||
"orphaned image records",
|
||||
db.cleanup_missing_images(), "orphaned image records"
|
||||
),
|
||||
"check_hash_duplicates": lambda: self._hash_duplicates_result(
|
||||
db.check_hash_duplicates()
|
||||
@@ -1274,11 +1207,7 @@ class AdminRoutesMixin:
|
||||
async def restore_database(self, request):
|
||||
"""Restore the prompts database from an uploaded SQLite file.
|
||||
|
||||
The upload is streamed to a temp file. Free disk space decides how
|
||||
large it may be (``maintenance.restore_room``), under a 64 GiB ceiling
|
||||
(``restore_max_bytes``): a request whose Content-Length does not fit
|
||||
is refused before anything is written, and one without is refused as
|
||||
soon as its running size does not. The file is
|
||||
The upload is streamed to a temp file (capped at restore_max_bytes),
|
||||
verified with PromptModel.verify_database_file and only then copied
|
||||
over the live database, which is backed up first. The whole request
|
||||
holds the registry (records): it is refused while a scan or another
|
||||
@@ -1291,44 +1220,10 @@ class AdminRoutesMixin:
|
||||
lambda: self._restore_database(request),
|
||||
)
|
||||
|
||||
def _restore_refusal(self, size, ceiling, room, *, known):
|
||||
"""``400 too-large`` past the ceiling, else ``507 insufficient-space``."""
|
||||
from .maintenance import restore_space_message
|
||||
|
||||
if size > ceiling:
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"code": "too-large",
|
||||
"error": (
|
||||
"File too large. Maximum size is "
|
||||
f"{ceiling // (1024 * 1024)}MB"
|
||||
),
|
||||
},
|
||||
status=400,
|
||||
)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"code": "insufficient-space",
|
||||
"error": restore_space_message(size, room, known=known),
|
||||
},
|
||||
status=507,
|
||||
)
|
||||
|
||||
async def _restore_database(self, request):
|
||||
from .maintenance import RESTORE_MAX_BYTES, restore_room
|
||||
|
||||
ceiling = getattr(self, "restore_max_bytes", RESTORE_MAX_BYTES)
|
||||
max_bytes = getattr(self, "restore_max_bytes", 100 * 1024 * 1024)
|
||||
temp_path = None
|
||||
try:
|
||||
# Free disk space decides how large a backup may be (#313): the
|
||||
# upload, its staged copy and the swap must all fit.
|
||||
room = await self._run_in_executor(restore_room, self.db.model.db_path)
|
||||
max_bytes = min(ceiling, room)
|
||||
length = request.content_length
|
||||
if length is not None and length > max_bytes:
|
||||
return self._restore_refusal(length, ceiling, room, known=True)
|
||||
reader = await request.multipart()
|
||||
field = await reader.next()
|
||||
|
||||
@@ -1354,7 +1249,16 @@ class AdminRoutesMixin:
|
||||
break
|
||||
size += len(chunk)
|
||||
if size > max_bytes:
|
||||
return self._restore_refusal(size, ceiling, room, known=False)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"error": (
|
||||
"File too large. Maximum size is "
|
||||
f"{max_bytes // (1024 * 1024)}MB"
|
||||
),
|
||||
},
|
||||
status=400,
|
||||
)
|
||||
out.write(chunk)
|
||||
|
||||
if size == 0:
|
||||
@@ -1504,9 +1408,6 @@ class AdminRoutesMixin:
|
||||
verb = "cancelled" if state == "cancelled" else "finished"
|
||||
job.log_line(f"Scan {verb} · {_plural(job.processed, 'file')} processed")
|
||||
job.finish(state)
|
||||
if state == "finished":
|
||||
# After the scan, so the guarded thumbnail job is admitted (#249).
|
||||
await self._hand_off_scan_thumbnails(job, output_dirs)
|
||||
|
||||
async def _stream_scan_job(self, request, job):
|
||||
"""Stream ``job`` as SSE; stop quietly if the client goes away."""
|
||||
@@ -1745,7 +1646,7 @@ class AdminRoutesMixin:
|
||||
if self._scan_job and self._scan_job.cancelled:
|
||||
break
|
||||
counts["processed"] += 1
|
||||
ambiguous, untraced, records = [], [], {}
|
||||
ambiguous, untraced = [], []
|
||||
try:
|
||||
outcome = self._ingest_scanned_file(
|
||||
media_file,
|
||||
@@ -1753,7 +1654,6 @@ class AdminRoutesMixin:
|
||||
detailed=True,
|
||||
ambiguous=ambiguous,
|
||||
untraced=untraced,
|
||||
records=records,
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error processing {media_file.name}: {e}")
|
||||
@@ -1766,11 +1666,7 @@ class AdminRoutesMixin:
|
||||
else:
|
||||
self._record_scanned_review(media_file, roots, review)
|
||||
elif untraced and outcome == "unresolved":
|
||||
# Untraced by the default profile: the inspection profile may
|
||||
# still offer a choice (queued for review, never linked).
|
||||
outcome = self._untraced_scan_outcome(
|
||||
media_file, roots, records.get("canonical"), review
|
||||
)
|
||||
self._settle_untraced_reviews(untraced, roots, review)
|
||||
if outcome is None:
|
||||
note(media_file, None)
|
||||
continue
|
||||
@@ -1805,13 +1701,7 @@ class AdminRoutesMixin:
|
||||
return positive.strip() or None
|
||||
|
||||
def _ingest_scanned_file(
|
||||
self,
|
||||
media_file,
|
||||
metadata,
|
||||
detailed=False,
|
||||
ambiguous=None,
|
||||
untraced=None,
|
||||
records=None,
|
||||
self, media_file, metadata, detailed=False, ambiguous=None, untraced=None
|
||||
):
|
||||
"""Enrich exact pairs; retain legacy positive-only import for absent pairs.
|
||||
|
||||
@@ -1824,9 +1714,7 @@ class AdminRoutesMixin:
|
||||
``"ambiguous"`` (nothing linked it) is a Needs review case. An
|
||||
ambiguous result with nothing to choose (``reviewable`` is False: fewer
|
||||
than two candidates carry a prompt) that nothing links returns
|
||||
``"unresolved"`` instead, and ``media_file`` goes to ``untraced``;
|
||||
its parsed records go to ``records["canonical"]`` when ``records`` is
|
||||
a dict.
|
||||
``"unresolved"`` instead, and ``media_file`` goes to ``untraced``.
|
||||
"""
|
||||
if not metadata:
|
||||
return None
|
||||
@@ -1847,8 +1735,6 @@ class AdminRoutesMixin:
|
||||
if unchoosable:
|
||||
if untraced is not None:
|
||||
untraced.append(media_file)
|
||||
if records is not None:
|
||||
records["canonical"] = canonical
|
||||
elif resolution["state"] == "ambiguous" and ambiguous is not None:
|
||||
ambiguous.append(media_file)
|
||||
if resolution["state"] == "resolved":
|
||||
|
||||
@@ -1,779 +0,0 @@
|
||||
"""Duplicate image files on the job registry (#249 System S6, Task 17).
|
||||
|
||||
- ``POST /prompt_manager/jobs/duplicates`` starts the find (202 ``{job}``):
|
||||
kind ``duplicates``, "Find duplicate image files", admission ``read``,
|
||||
cancellable, outcome group ``groups_found``. It hashes the media files of
|
||||
every effective source (``_hash_media_groups``, shared with the legacy
|
||||
scan) and the server keeps the groups it found under the job id: content
|
||||
hash, root-relative paths and each file's source index. A find is refused
|
||||
while a duplicate removal runs, and one find runs at a time.
|
||||
- ``GET /prompt_manager/maintenance/duplicates/{job_id}?offset&limit`` pages
|
||||
those groups, at most ``MAX_GROUP_PAGE`` at a time: ``{success, job:{id,
|
||||
state}, total, offset, limit, groups:[{group_id, size, count, files}]}``.
|
||||
Each file names its root-relative ``path``, its source's public label and
|
||||
index and the Images serve URLs; never an absolute path.
|
||||
- ``POST /prompt_manager/maintenance/preview {operation:"duplicate_files",
|
||||
options:{job_id, keep:[{group_id, root_index, path}]}}`` reviews a removal:
|
||||
in each group named, the files kept are the ones listed and every other
|
||||
copy is removed. Each keeper and each removal is checked against the held
|
||||
group, with containment decided as text before any filesystem access
|
||||
(``file_problem``): inside its source, a media file outside
|
||||
``thumbnails``, then a regular file (not a link) whose resolved path is
|
||||
still inside the source, with the same size and content hash. A removal
|
||||
that fails, or whose kept copy fails, is protected and listed with the
|
||||
reason. The plan carries the files' database links, which move to the
|
||||
kept copy.
|
||||
- ``POST /prompt_manager/maintenance/apply {token}`` runs the reviewed plan
|
||||
as the ``duplicate_cleanup`` job: admission ``records``, not cancellable,
|
||||
groups ``removed``, ``kept`` and ``errors``. It checks the sources are
|
||||
unchanged, takes the safety backup the policy promises (the database only;
|
||||
removed files are not in it), then for each removal rechecks its kept copy
|
||||
and the file itself (``file_problem``) and removes it with every handle
|
||||
closed (``_unlink``), together with its thumbnail. A file that changed is
|
||||
kept; the first unlink that fails stops the job, a locked file with the
|
||||
coded ``LOCKED_FILE`` message (never the OS text), and the files after it
|
||||
are not attempted. The database links of the files actually removed then
|
||||
move to their kept copy (``PromptDatabase.move_file_links``).
|
||||
|
||||
The legacy gallery routes (``/scan_duplicates``, its stream and status, and
|
||||
``/delete_duplicate_images``) keep their shapes on the same registry kinds;
|
||||
see ``py/api/admin.py``.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
from .admin import DUPLICATE_FIND_LEGACY_TITLE, _remove_thumbnail
|
||||
from .duplicate_review import ( # noqa: F401 (CHANGED, OUTSIDE_SOURCE: public)
|
||||
CHANGED,
|
||||
DELETE_FAILED,
|
||||
FIND_CANCELLED,
|
||||
FIND_ERROR_EVENT,
|
||||
FIND_FAILED,
|
||||
GONE,
|
||||
KEEPER,
|
||||
KEEPER_CHANGED,
|
||||
KEEPER_GONE,
|
||||
KEEPER_WENT,
|
||||
KEPT_COPY,
|
||||
LOCKED_FILE,
|
||||
MAYBE_SAME_FILE,
|
||||
NO_FILE_ID_ROW,
|
||||
NOT_ATTEMPTED,
|
||||
OUTSIDE_SOURCE,
|
||||
RECORDS_FAILED,
|
||||
ROW_DETAILS,
|
||||
SAME_FILE,
|
||||
SOURCES_CHANGED,
|
||||
SOURCES_CHANGED_REVIEW,
|
||||
HeldFind,
|
||||
SourcesChanged,
|
||||
_fingerprint,
|
||||
_group_order,
|
||||
_identity,
|
||||
_parse_options,
|
||||
_plural,
|
||||
_public_plan,
|
||||
_select,
|
||||
file_problem,
|
||||
human_bytes,
|
||||
still_kept,
|
||||
)
|
||||
from .duplicate_spellings import Spellings, copy_move
|
||||
from .images import _favorite_root_key, _output_image_entry
|
||||
from .job_routes import busy_response
|
||||
from .jobs import (
|
||||
DUPLICATE_CLEANUP_GROUP_NAMES,
|
||||
DUPLICATE_CLEANUP_JOB_KIND,
|
||||
DUPLICATE_GROUP_NAMES,
|
||||
DUPLICATES_JOB_KIND,
|
||||
JobBusy,
|
||||
)
|
||||
from .maintenance import DUPLICATE_FILES, _fail_job, _failure, _invalid, backup_note
|
||||
|
||||
try:
|
||||
from ...database.duplicate_copies import FILE_LINK_SQL
|
||||
except ImportError:
|
||||
from database.duplicate_copies import FILE_LINK_SQL
|
||||
|
||||
DUPLICATE_FIND_TITLE = "Find duplicate image files"
|
||||
DUPLICATE_CLEANUP_TITLE = "Remove duplicate image files"
|
||||
MAX_GROUP_PAGE = 10 # groups a page lists (D69: ten at a time)
|
||||
MAX_HELD_FINDS = 3 # finds whose groups the server keeps for review
|
||||
favorite_root_key = _favorite_root_key
|
||||
|
||||
|
||||
def query_plan_sql():
|
||||
"""The link reconciliation's statements (EXPLAIN QUERY PLAN checks)."""
|
||||
return list(FILE_LINK_SQL)
|
||||
|
||||
|
||||
def _unlink(path):
|
||||
"""Remove one file (the caller closed every handle on it)."""
|
||||
os.remove(path)
|
||||
|
||||
|
||||
class DuplicateFilesMixin:
|
||||
"""Mixin providing the duplicate-file find, its groups and its removal."""
|
||||
|
||||
def _register_duplicate_routes(self, routes):
|
||||
@routes.post("/prompt_manager/jobs/duplicates")
|
||||
async def start_duplicate_find_route(request):
|
||||
return await self.start_duplicate_find(request)
|
||||
|
||||
@routes.get("/prompt_manager/maintenance/duplicates/{job_id}")
|
||||
async def duplicate_groups_route(request):
|
||||
return await self.duplicate_groups(request)
|
||||
|
||||
# ── find ──────────────────────────────────────────────────────────
|
||||
|
||||
async def start_duplicate_find(self, request):
|
||||
"""Start the duplicate find (202), or 409 naming the job in the way."""
|
||||
try:
|
||||
job = self._start_duplicate_find(origin="workspace")
|
||||
except JobBusy as busy:
|
||||
return busy_response(busy)
|
||||
return web.json_response({"success": True, "job": job.summary()}, status=202)
|
||||
|
||||
def _start_duplicate_find(self, origin):
|
||||
"""Register a find job (on the loop); raises JobBusy.
|
||||
|
||||
A workspace find walks every effective source; a gallery one the
|
||||
gallery's output folder, as the legacy scan did. Both resolve their
|
||||
folders when the job begins, so a source change later cannot move it.
|
||||
"""
|
||||
cleanup = self._jobs.active(DUPLICATE_CLEANUP_JOB_KIND)
|
||||
if cleanup is not None:
|
||||
raise JobBusy(cleanup)
|
||||
held = HeldFind(origin)
|
||||
if origin == "workspace":
|
||||
title = DUPLICATE_FIND_TITLE
|
||||
else:
|
||||
title = DUPLICATE_FIND_LEGACY_TITLE
|
||||
job = self._jobs.start(
|
||||
DUPLICATES_JOB_KIND,
|
||||
title,
|
||||
admission="read",
|
||||
cancellable=True,
|
||||
groups=DUPLICATE_GROUP_NAMES,
|
||||
runner=lambda job: self._run_duplicate_find(job, held),
|
||||
)
|
||||
held.job_id = job.id
|
||||
self._remember_find(held)
|
||||
return job
|
||||
|
||||
def _remember_find(self, held):
|
||||
finds = self._duplicate_finds
|
||||
finds[held.job_id] = held
|
||||
for job_id in list(finds):
|
||||
if self._jobs.get(job_id) is None:
|
||||
del finds[job_id]
|
||||
while len(finds) > MAX_HELD_FINDS:
|
||||
del finds[next(iter(finds))]
|
||||
|
||||
def _find_covers(self, job, legacy_dir):
|
||||
held = self._duplicate_finds.get(job.id)
|
||||
return held is not None and held.covers(legacy_dir)
|
||||
|
||||
def _latest_duplicate_find(self):
|
||||
"""The running find, else the newest finished one this session."""
|
||||
for job in self._jobs.list():
|
||||
if job.kind == DUPLICATES_JOB_KIND and job.id in self._duplicate_finds:
|
||||
return job
|
||||
return None
|
||||
|
||||
def _resolve_find_roots(self, held):
|
||||
"""Fill in where a find looks (blocking): sources and the gallery folder."""
|
||||
listed = [os.path.abspath(str(d)) for d in self._get_all_output_dirs()]
|
||||
held.legacy_dir = self._find_comfyui_output_dir()
|
||||
if held.origin == "workspace":
|
||||
held.roots = list(enumerate(listed))
|
||||
return
|
||||
if not held.legacy_dir or not os.path.isdir(held.legacy_dir):
|
||||
self.logger.warning("ComfyUI output directory not found")
|
||||
held.roots = []
|
||||
return
|
||||
legacy = os.path.abspath(held.legacy_dir)
|
||||
index = next(
|
||||
(i for i, d in enumerate(listed) if _identity(d) == _identity(legacy)),
|
||||
None,
|
||||
)
|
||||
held.roots = [(index, legacy)]
|
||||
|
||||
def _find_progress(self, job, done, total):
|
||||
if done < job.processed:
|
||||
return # hashing threads report out of order; keep the count rising
|
||||
pct = int(done / total * 100) if total else 100
|
||||
job.update(
|
||||
processed=done,
|
||||
total=total,
|
||||
status_line=f"Hashing file {done:,} of {total:,}",
|
||||
)
|
||||
job.publish(
|
||||
{
|
||||
"type": "progress",
|
||||
"progress": pct,
|
||||
"status": f"Hashing file {done}/{total}...",
|
||||
"processed": done,
|
||||
"found": 0,
|
||||
}
|
||||
)
|
||||
|
||||
async def _run_duplicate_find(self, job, held):
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
def report(done, total): # from the hashing threads
|
||||
loop.call_soon_threadsafe(self._find_progress, job, done, total)
|
||||
|
||||
try:
|
||||
job.update(status_line="Looking for media files")
|
||||
job.publish(
|
||||
{
|
||||
"type": "progress",
|
||||
"progress": 0,
|
||||
"status": "Collecting media files...",
|
||||
"processed": 0,
|
||||
"found": 0,
|
||||
}
|
||||
)
|
||||
await self._run_in_executor(self._resolve_find_roots, held)
|
||||
job.scope["output_folders"] = len(held.roots)
|
||||
job.log_line(f"Searching {_plural(len(held.roots), 'source')}")
|
||||
groups = await self._run_in_executor(
|
||||
self._hash_media_groups,
|
||||
[root for _index, root in held.roots],
|
||||
report,
|
||||
lambda: job.cancelled,
|
||||
)
|
||||
if job.cancelled:
|
||||
job.update(status_line="Cancelled; nothing was changed.")
|
||||
job.publish({"type": "error", "message": FIND_CANCELLED})
|
||||
job.finish("cancelled")
|
||||
return
|
||||
await self._finish_find(job, held, groups)
|
||||
except Exception:
|
||||
self.logger.exception("Duplicate search failed")
|
||||
job.update(status_line=FIND_FAILED)
|
||||
job.publish({"type": "error", "message": FIND_ERROR_EVENT})
|
||||
job.finish("failed")
|
||||
|
||||
async def _finish_find(self, job, held, groups):
|
||||
for group in groups:
|
||||
group["group_id"] = group["hash"]
|
||||
for record in group["files"]:
|
||||
index, root = held.roots[record.pop("slot")]
|
||||
record["root_index"] = index
|
||||
record["root"] = root
|
||||
groups.sort(key=_group_order)
|
||||
held.groups = groups
|
||||
labels = {root: self._public_path(root) for _index, root in held.roots}
|
||||
rows = [
|
||||
{
|
||||
"path": group["files"][0]["rel"],
|
||||
"root": labels[group["files"][0]["root"]],
|
||||
"root_index": group["files"][0]["root_index"],
|
||||
"detail": f"{len(group['files'])} identical files · "
|
||||
f"{human_bytes(group['size'])} each",
|
||||
}
|
||||
for group in groups
|
||||
]
|
||||
job.record_counted("groups_found", len(groups), rows)
|
||||
line = f"Found {_plural(len(groups), 'group')} of duplicate files"
|
||||
job.log_line(line)
|
||||
legacy = []
|
||||
if held.covers(held.legacy_dir) and held.legacy_dir:
|
||||
legacy = await self._run_in_executor(
|
||||
self._legacy_duplicate_groups, groups, held.legacy_dir
|
||||
)
|
||||
# Every file was read by now (the last reports may still be queued).
|
||||
job.update(progress=100, status_line=line, processed=job.total)
|
||||
job.publish(
|
||||
{
|
||||
"type": "complete",
|
||||
"processed": job.processed,
|
||||
"found": len(legacy),
|
||||
"duplicates": legacy,
|
||||
"message": f"Found {len(legacy)} groups of duplicate images",
|
||||
}
|
||||
)
|
||||
|
||||
# ── groups ────────────────────────────────────────────────────────
|
||||
|
||||
def _held_result(self, job_id):
|
||||
"""``(held, None)`` for a finished find's groups, else an error reply."""
|
||||
held = self._duplicate_finds.get(job_id)
|
||||
job = self._jobs.get(job_id)
|
||||
if held is None or job is None or job.kind != DUPLICATES_JOB_KIND:
|
||||
return None, _failure(
|
||||
"No duplicate search with this id. Search again.", 404, "not-found"
|
||||
)
|
||||
if job.running:
|
||||
return None, _failure(
|
||||
"The duplicate search is still running.", 409, "job-running"
|
||||
)
|
||||
if held.groups is None:
|
||||
return None, _failure(
|
||||
"This search did not finish. Search again.", 409, "no-groups"
|
||||
)
|
||||
return held, None
|
||||
|
||||
async def duplicate_groups(self, request):
|
||||
"""One page of a finished find's groups (at most MAX_GROUP_PAGE)."""
|
||||
try:
|
||||
offset = int(request.query.get("offset", "0"))
|
||||
limit = int(request.query.get("limit", str(MAX_GROUP_PAGE)))
|
||||
except ValueError:
|
||||
return _invalid("offset and limit must be whole numbers")
|
||||
if offset < 0 or not 0 < limit <= MAX_GROUP_PAGE:
|
||||
return _invalid(f"limit must be 1 to {MAX_GROUP_PAGE}; offset 0 or more")
|
||||
job_id = request.match_info["job_id"]
|
||||
held, error = self._held_result(job_id)
|
||||
if error is not None:
|
||||
return error
|
||||
page = held.groups[offset : offset + limit]
|
||||
groups = await self._run_in_executor(self._describe_groups, page)
|
||||
job = self._jobs.get(job_id)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"job": {"id": job.id, "state": job.state},
|
||||
"total": len(held.groups),
|
||||
"offset": offset,
|
||||
"limit": limit,
|
||||
"groups": groups,
|
||||
},
|
||||
headers={"Cache-Control": "no-store"},
|
||||
)
|
||||
|
||||
def _describe_groups(self, groups):
|
||||
"""Public group rows (blocking: one stat per file)."""
|
||||
return [
|
||||
{
|
||||
"group_id": group["group_id"],
|
||||
"size": group["size"],
|
||||
"count": len(group["files"]),
|
||||
"files": [self._describe_copy(record) for record in group["files"]],
|
||||
}
|
||||
for group in groups
|
||||
]
|
||||
|
||||
def _describe_copy(self, record):
|
||||
label = self._public_path(record["root"])
|
||||
base = {
|
||||
"filename": os.path.basename(record["rel"]),
|
||||
"path": record["rel"],
|
||||
"root": label,
|
||||
"root_index": record["root_index"],
|
||||
"url": None,
|
||||
"thumbnail_url": None,
|
||||
"size": record["size"],
|
||||
"modified_time": record["mtime"],
|
||||
"media_type": "image",
|
||||
"missing": False,
|
||||
}
|
||||
if record["root_index"] is None:
|
||||
return base
|
||||
try:
|
||||
entry = _output_image_entry(
|
||||
Path(record["abs"]), Path(record["root"]), record["root_index"], label
|
||||
)
|
||||
except (OSError, ValueError):
|
||||
return {**base, "missing": True}
|
||||
return {
|
||||
**base,
|
||||
"url": entry["url"],
|
||||
"thumbnail_url": entry["thumbnail_url"],
|
||||
"size": entry["size"],
|
||||
"modified_time": entry["modified_time"],
|
||||
"media_type": entry["media_type"],
|
||||
}
|
||||
|
||||
# ── review ────────────────────────────────────────────────────────
|
||||
|
||||
def _listed_sources(self):
|
||||
"""``{identity: index}`` of the effective sources now (blocking)."""
|
||||
return {
|
||||
_identity(d): index for index, d in enumerate(self._get_all_output_dirs())
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _source_listed(listed, index, root):
|
||||
return index is not None and listed.get(_identity(root)) == index
|
||||
|
||||
async def preview_duplicate_files(self, options):
|
||||
"""The ``duplicate_files`` plan for one finished find and its keepers."""
|
||||
job_id, keep, message = _parse_options(options)
|
||||
if message:
|
||||
return _invalid(message)
|
||||
held, error = self._held_result(job_id)
|
||||
if error is not None:
|
||||
return error
|
||||
selection, message = _select(held, keep)
|
||||
if message:
|
||||
return _invalid(message)
|
||||
db = self.db
|
||||
identity = self._database_identity(db)
|
||||
try:
|
||||
review = await self._run_in_executor(
|
||||
self._review_duplicate_files, db, selection
|
||||
)
|
||||
except SourcesChanged:
|
||||
return _failure(SOURCES_CHANGED_REVIEW, 409, "sources-changed")
|
||||
except Exception:
|
||||
self.logger.exception("Duplicate file review failed")
|
||||
return _failure("Failed to review maintenance", 500)
|
||||
return await self._plan_response(
|
||||
db,
|
||||
identity,
|
||||
DUPLICATE_FILES,
|
||||
review["fingerprint"],
|
||||
review["public"],
|
||||
{
|
||||
"targets": review["targets"],
|
||||
"kept": review["kept"],
|
||||
"protected": review["public"]["protected"]["count"],
|
||||
},
|
||||
)
|
||||
|
||||
def _review_duplicate_files(self, db, selection):
|
||||
"""Check every keeper and removal now (blocking); the plan's parts.
|
||||
|
||||
No file is read: each copy is bound to what the search saw by its
|
||||
size, mtime and inode, and never to its kept copy itself. The removal
|
||||
job hashes each file again right before it unlinks it.
|
||||
"""
|
||||
listed = self._listed_sources()
|
||||
targets, kept, protected = [], [], {}
|
||||
for group, keepers in selection.values():
|
||||
removals = [r for r in group["files"] if r not in keepers]
|
||||
for record in removals:
|
||||
if record["root_index"] is not None and not self._source_listed(
|
||||
listed, record["root_index"], record["root"]
|
||||
):
|
||||
raise SourcesChanged(record["root"])
|
||||
keeper = next(
|
||||
(k for k in keepers if _bound_problem(k, group) is None), None
|
||||
)
|
||||
reviewed = []
|
||||
for record in removals:
|
||||
if keeper is None:
|
||||
problem = KEEPER
|
||||
elif record["root_index"] is None:
|
||||
problem = OUTSIDE_SOURCE
|
||||
else:
|
||||
problem = _bound_problem(record, group, keeper)
|
||||
if problem:
|
||||
protected[problem] = protected.get(problem, 0) + 1
|
||||
continue
|
||||
reviewed.append(self._target(group, record, keeper))
|
||||
if reviewed:
|
||||
targets.extend(reviewed)
|
||||
kept.append(self._copy_row(keeper, KEPT_COPY))
|
||||
records = {}
|
||||
if targets:
|
||||
spellings = self._source_spellings()
|
||||
records = db.count_copy_records(
|
||||
{
|
||||
t["abs"]: {
|
||||
"spellings": list(spellings.of(t["abs"]).values()),
|
||||
"locator": t["locator"],
|
||||
"sha256": t["hash"],
|
||||
}
|
||||
for t in targets
|
||||
}
|
||||
)
|
||||
return {
|
||||
"targets": targets,
|
||||
"kept": kept,
|
||||
"fingerprint": _fingerprint(targets),
|
||||
"public": _public_plan(targets, protected, records),
|
||||
}
|
||||
|
||||
def _source_spellings(self):
|
||||
"""The folder spellings records may name files by (blocking).
|
||||
|
||||
The search walks each source's real path, while a record names a file
|
||||
as it was seen when written: under the folder the image monitor
|
||||
watches, the configured spelling, ComfyUI's output folder, a symlink
|
||||
to it, or a parent of a source added later. ``Spellings`` matches
|
||||
them to each file by containment (``py/api/duplicate_spellings.py``).
|
||||
"""
|
||||
from ..config import GalleryConfig
|
||||
|
||||
listed = [os.path.abspath(str(d)) for d in self._get_all_output_dirs()]
|
||||
candidates = list(self._monitored_directories())
|
||||
candidates += list(GalleryConfig.MONITORING_DIRECTORIES)
|
||||
try:
|
||||
import folder_paths
|
||||
|
||||
candidates.append(folder_paths.get_output_directory())
|
||||
except Exception: # nosec B110 - outside ComfyUI there is none
|
||||
pass
|
||||
candidates.append(self._find_comfyui_output_dir())
|
||||
folders = [
|
||||
os.path.abspath(str(c)) for c in candidates if c and os.path.isabs(str(c))
|
||||
]
|
||||
return Spellings(folders + listed, listed)
|
||||
|
||||
def _target(self, group, record, keeper):
|
||||
return {
|
||||
"group_id": group["group_id"],
|
||||
"hash": group["hash"],
|
||||
"size": record["size"],
|
||||
"root_index": record["root_index"],
|
||||
"root": record["root"],
|
||||
"root_label": self._public_path(record["root"]),
|
||||
"abs": record["abs"],
|
||||
"rel": record["rel"],
|
||||
"keeper_abs": keeper["abs"],
|
||||
"keeper_rel": keeper["rel"],
|
||||
"keeper_root": keeper["root"],
|
||||
"keeper_size": keeper["size"],
|
||||
"locator": os.path.normcase(record["abs"]),
|
||||
"keeper_locator": os.path.normcase(keeper["abs"]),
|
||||
}
|
||||
|
||||
def _copy_row(self, record, detail):
|
||||
return {
|
||||
"path": record["rel"],
|
||||
"root": self._public_path(record["root"]),
|
||||
"root_index": record["root_index"],
|
||||
"detail": detail,
|
||||
}
|
||||
|
||||
# ── removal ───────────────────────────────────────────────────────
|
||||
|
||||
def _start_duplicate_cleanup(self, db, plan, backup_on):
|
||||
"""Register the removal job for a reviewed plan (on the loop)."""
|
||||
find = self._jobs.active(DUPLICATES_JOB_KIND)
|
||||
if find is not None:
|
||||
raise JobBusy(find)
|
||||
return self._jobs.start(
|
||||
DUPLICATE_CLEANUP_JOB_KIND,
|
||||
DUPLICATE_CLEANUP_TITLE,
|
||||
admission="records",
|
||||
cancellable=False,
|
||||
scope={"operation": DUPLICATE_FILES},
|
||||
groups=DUPLICATE_CLEANUP_GROUP_NAMES,
|
||||
runner=lambda job: self._run_duplicate_cleanup(job, db, plan, backup_on),
|
||||
)
|
||||
|
||||
def _sources_unchanged(self, targets):
|
||||
listed = self._listed_sources()
|
||||
return all(
|
||||
self._source_listed(listed, t["root_index"], t["root"]) for t in targets
|
||||
)
|
||||
|
||||
async def _run_duplicate_cleanup(self, job, db, plan, backup_on):
|
||||
targets = plan.payload["targets"]
|
||||
unchanged = await self._run_in_executor(self._sources_unchanged, targets)
|
||||
if not unchanged:
|
||||
return _fail_job(job, "sources-changed", SOURCES_CHANGED)
|
||||
ok, backup = await self._safety_backup_step(job, db, backup_on)
|
||||
if not ok:
|
||||
return
|
||||
try:
|
||||
await self._remove_duplicates(job, db, plan, backup)
|
||||
finally:
|
||||
if backup:
|
||||
await self._prune_backup_pool(job, db)
|
||||
|
||||
async def _remove_duplicates(self, job, db, plan, backup):
|
||||
targets = plan.payload["targets"]
|
||||
for row in plan.payload["kept"]:
|
||||
job.record("kept", row, f"= {row['root']}/{row['path']} · {KEPT_COPY}")
|
||||
job.update(total=len(targets), status_line="Removing duplicate files")
|
||||
spellings = await self._run_in_executor(self._source_spellings)
|
||||
removed, failure, keepers = 0, None, {}
|
||||
totals = dict.fromkeys(("moved", "merged", "favorites", "tags", "review"), 0)
|
||||
for done, target in enumerate(targets, start=1):
|
||||
if failure is not None:
|
||||
self._record_target(job, "kept", target, NOT_ATTEMPTED)
|
||||
continue
|
||||
outcome, detail, counts = await self._run_in_executor(
|
||||
self._remove_one_duplicate, db, target, keepers, spellings
|
||||
)
|
||||
for key, value in (counts or {}).items():
|
||||
totals[key] += value
|
||||
if outcome == "removed":
|
||||
removed += 1
|
||||
text = f"A copy of {target['keeper_rel']}"
|
||||
self._record_target(job, "removed", target, text, "-")
|
||||
elif outcome == "kept":
|
||||
self._record_target(job, "kept", target, ROW_DETAILS[detail])
|
||||
else:
|
||||
failure = detail
|
||||
self._record_target(job, "errors", target, detail, "!")
|
||||
job.update(processed=done, progress=int(done / len(targets) * 100))
|
||||
if removed:
|
||||
self.invalidate_gallery_cache()
|
||||
self._log_moved_records(job, totals)
|
||||
line = f"Removed {_plural(removed, 'duplicate file')}"
|
||||
if removed < len(targets):
|
||||
line += f", {len(targets) - removed:,} not removed"
|
||||
job.log_line(line)
|
||||
if failure is not None:
|
||||
line = f"Stopped after an error · {line}"
|
||||
job.update(progress=100, status_line=f"{line} · {backup_note(backup)}")
|
||||
if failure is not None:
|
||||
job.publish({"type": "error", "code": "delete-failed", "message": line})
|
||||
job.finish("failed")
|
||||
|
||||
@staticmethod
|
||||
def _log_moved_records(job, totals):
|
||||
links = totals["moved"] + totals["merged"]
|
||||
if links:
|
||||
job.log_line(f"Moved {_plural(links, 'database link')} to the kept copies")
|
||||
if totals["favorites"]:
|
||||
job.log_line(
|
||||
f"Moved {_plural(totals['favorites'], 'favorite')} to the kept copies"
|
||||
)
|
||||
if totals["tags"]:
|
||||
job.log_line(
|
||||
f"Added {_plural(totals['tags'], 'image tag')} to the kept copies"
|
||||
)
|
||||
if totals["review"]:
|
||||
job.log_line(
|
||||
f"Settled {_plural(totals['review'], 'Needs review item')} "
|
||||
"of removed copies"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _record_target(job, group, target, detail, mark="="):
|
||||
row = {
|
||||
"path": target["rel"],
|
||||
"root": target["root_label"],
|
||||
"root_index": target["root_index"],
|
||||
"detail": detail,
|
||||
}
|
||||
job.record(
|
||||
group, row, f"{mark} {target['root_label']}/{target['rel']} · {detail}"
|
||||
)
|
||||
|
||||
def _remove_one_duplicate(self, db, target, keepers, spellings):
|
||||
"""``(outcome, detail, counts)`` for one reviewed removal (blocking).
|
||||
|
||||
The kept copy (once per group; cached in ``keepers``) and the file
|
||||
are checked again, hashes included, and the file must not be its kept
|
||||
copy. Then the file's records move to the kept copy in one short
|
||||
transaction, and only then is the file unlinked: a crash or a failed
|
||||
unlink leaves its records on the identical kept copy, never on a
|
||||
deleted file. A failed move leaves the file and its records alone.
|
||||
"""
|
||||
hasher = self._calculate_file_hash
|
||||
group_id = target["group_id"]
|
||||
if group_id not in keepers:
|
||||
keepers[group_id] = _kept_copy_state(target, hasher)
|
||||
if keepers[group_id] is None:
|
||||
return "kept", KEEPER, None
|
||||
problem = file_problem(
|
||||
target["abs"],
|
||||
target["root"],
|
||||
target["hash"],
|
||||
target["size"],
|
||||
hasher,
|
||||
keeper=target["keeper_abs"],
|
||||
keeper_root=target["keeper_root"],
|
||||
)
|
||||
if problem:
|
||||
return "kept", problem, None
|
||||
try:
|
||||
# Held open until the unlink: on Windows, deleting the kept copy
|
||||
# itself through another path then fails instead of removing it.
|
||||
hold = open(target["keeper_abs"], "rb")
|
||||
except OSError:
|
||||
return "kept", KEEPER, None
|
||||
with hold:
|
||||
return self._move_then_unlink(
|
||||
db, target, keepers[group_id], spellings, hold
|
||||
)
|
||||
|
||||
def _move_then_unlink(self, db, target, kept_stat, spellings, hold):
|
||||
try:
|
||||
counts = db.move_copy_records(copy_move(target, kept_stat, spellings))
|
||||
except Exception:
|
||||
self.logger.exception(f"Could not move the records of {target['rel']}")
|
||||
return "error", RECORDS_FAILED, None
|
||||
if not still_kept(target, hold):
|
||||
# The kept copy went away or changed while the records moved:
|
||||
# keep this file and give its records back to it.
|
||||
return "kept", KEEPER_GONE, self._move_back(db, target, spellings)
|
||||
try:
|
||||
_unlink(target["abs"])
|
||||
except FileNotFoundError:
|
||||
return "removed", None, counts
|
||||
except PermissionError:
|
||||
self.logger.warning(
|
||||
f"Duplicate file is locked: {target['rel']}; "
|
||||
"its records now name the kept copy"
|
||||
)
|
||||
return "error", LOCKED_FILE, counts
|
||||
except OSError:
|
||||
self.logger.warning(
|
||||
f"Could not delete duplicate file {target['rel']}; "
|
||||
"its records now name the kept copy",
|
||||
exc_info=True,
|
||||
)
|
||||
return "error", DELETE_FAILED, counts
|
||||
self.logger.info(f"Deleted duplicate file: {target['rel']}")
|
||||
_remove_thumbnail(Path(target["root"]), Path(target["rel"]), self.logger)
|
||||
return "removed", None, counts
|
||||
|
||||
def _move_back(self, db, target, spellings):
|
||||
"""Give the kept copy's records to the file that stays (blocking)."""
|
||||
swapped = {
|
||||
**target,
|
||||
"abs": target["keeper_abs"],
|
||||
"root": target["keeper_root"],
|
||||
"rel": target["keeper_rel"],
|
||||
"locator": target["keeper_locator"],
|
||||
"keeper_abs": target["abs"],
|
||||
"keeper_root": target["root"],
|
||||
"keeper_rel": target["rel"],
|
||||
"keeper_locator": target["locator"],
|
||||
}
|
||||
try:
|
||||
return db.move_copy_records(
|
||||
copy_move(swapped, os.stat(target["abs"]), spellings)
|
||||
)
|
||||
except Exception:
|
||||
self.logger.exception(f"Could not give the records back to {target['rel']}")
|
||||
return None
|
||||
|
||||
|
||||
def _bound_problem(record, group, keeper=None):
|
||||
"""``file_problem`` without reading the file: size, mtime and inode.
|
||||
|
||||
``keeper`` is the kept copy's record, never offered as a copy of itself.
|
||||
"""
|
||||
return file_problem(
|
||||
record["abs"],
|
||||
record["root"],
|
||||
group["hash"],
|
||||
record["size"],
|
||||
hash_check=False,
|
||||
mtime_ns=record.get("mtime_ns"),
|
||||
inode=record.get("inode"),
|
||||
keeper=keeper["abs"] if keeper else None,
|
||||
keeper_root=keeper["root"] if keeper else None,
|
||||
)
|
||||
|
||||
|
||||
def _kept_copy_state(target, hasher):
|
||||
"""The kept copy's ``os.stat`` when it is still the reviewed bytes, else None."""
|
||||
problem = file_problem(
|
||||
target["keeper_abs"],
|
||||
target["keeper_root"],
|
||||
target["hash"],
|
||||
target["keeper_size"],
|
||||
hasher,
|
||||
)
|
||||
if problem:
|
||||
return None
|
||||
try:
|
||||
return os.stat(target["keeper_abs"])
|
||||
except OSError:
|
||||
return None
|
||||
@@ -1,451 +0,0 @@
|
||||
"""Reviewing duplicate image files: the checks and the plan (#249 S6, T17).
|
||||
|
||||
Pure helpers of ``py/api/duplicate_files.py``: the keep-list parsing, the
|
||||
per-file check ``file_problem`` (containment decided as text before any
|
||||
filesystem access, then a regular file of the found size and hash) and the
|
||||
public ``duplicate_files`` plan.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import stat as stat_module
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
from .admin import DELETABLE_MEDIA_EXTENSIONS
|
||||
from .maintenance import DUPLICATE_FILES
|
||||
|
||||
# Keep entries one review takes; with paths of a few hundred characters the
|
||||
# list stays inside the preview body cap (DUPLICATE_PREVIEW_BODY_MAX, 1 MiB).
|
||||
MAX_KEEP = 2000
|
||||
MAX_DETAILS = 50
|
||||
MAX_JOB_ID = 64
|
||||
HASH_CHUNK = 1024 * 1024
|
||||
|
||||
# What ``file_problem`` finds, and how a plan counts it as protected.
|
||||
OUTSIDE_SOURCE = "outside"
|
||||
LINK = "link"
|
||||
CHANGED_CONTENT = "changed"
|
||||
GONE = "gone"
|
||||
UNREADABLE = "unreadable"
|
||||
KEEPER = "keeper"
|
||||
SAME = "same"
|
||||
MAYBE_SAME = "maybe-same"
|
||||
NO_FILE_ID = "no-file-id"
|
||||
# Two mounts of one share may round mtimes differently (SMB keeps 100 ns,
|
||||
# FAT 2 s): within this, equal sizes on another device may be one file.
|
||||
MOUNT_MTIME_SLACK_NS = 2_000_000_000
|
||||
PROTECTED_REASONS = {
|
||||
MAYBE_SAME: (
|
||||
"file may be the same file reached through another mount or share",
|
||||
"files may be the same file reached through another mount or share",
|
||||
),
|
||||
NO_FILE_ID: (
|
||||
"file: can't confirm this is a separate file on this drive",
|
||||
"files: can't confirm these are separate files on this drive",
|
||||
),
|
||||
SAME: (
|
||||
"file is the same file as its kept copy",
|
||||
"files are the same file as their kept copy",
|
||||
),
|
||||
OUTSIDE_SOURCE: (
|
||||
"file is not inside a configured source",
|
||||
"files are not inside a configured source",
|
||||
),
|
||||
LINK: ("file is a link, not a regular file", "files are links, not regular files"),
|
||||
CHANGED_CONTENT: (
|
||||
"file changed since the search",
|
||||
"files changed since the search",
|
||||
),
|
||||
GONE: ("file no longer exists", "files no longer exist"),
|
||||
UNREADABLE: ("file could not be read", "files could not be read"),
|
||||
KEEPER: (
|
||||
"file kept: its kept copy changed or is gone",
|
||||
"files kept: their kept copy changed or is gone",
|
||||
),
|
||||
}
|
||||
# Job outcome rows (never an OS error text).
|
||||
CHANGED = "The file changed since the review; it was kept."
|
||||
KEEPER_CHANGED = "The kept copy changed or is gone; this file was kept."
|
||||
SAME_FILE = "This is the same file as its kept copy; it was kept."
|
||||
MAYBE_SAME_FILE = (
|
||||
"May be the same file reached through another mount or share; it was kept."
|
||||
)
|
||||
NO_FILE_ID_ROW = "Can't confirm this is a separate file on this drive; it was kept."
|
||||
KEEPER_GONE = "keeper-gone"
|
||||
KEEPER_WENT = (
|
||||
"The kept copy changed or went away during the removal; this file was kept "
|
||||
"and its links stay with it."
|
||||
)
|
||||
ROW_DETAILS = {
|
||||
KEEPER_GONE: None, # filled below
|
||||
MAYBE_SAME: MAYBE_SAME_FILE,
|
||||
NO_FILE_ID: NO_FILE_ID_ROW,
|
||||
SAME: SAME_FILE,
|
||||
OUTSIDE_SOURCE: "The file is no longer inside its source; it was kept.",
|
||||
LINK: "The file is now a link, not a regular file; it was kept.",
|
||||
CHANGED_CONTENT: CHANGED,
|
||||
GONE: "The file was already gone.",
|
||||
UNREADABLE: "The file could not be read; it was kept.",
|
||||
KEEPER: KEEPER_CHANGED,
|
||||
}
|
||||
# An unlink fails only after the file's records moved to its kept copy.
|
||||
LOCKED_FILE = (
|
||||
"The file is open in another program or read-only; it was not deleted, "
|
||||
"and its links now name the kept copy."
|
||||
)
|
||||
DELETE_FAILED = "The file could not be deleted; its links now name the kept copy."
|
||||
NOT_ATTEMPTED = "Not attempted: the job stopped after an error."
|
||||
KEPT_COPY = "Kept copy"
|
||||
SOURCES_CHANGED = (
|
||||
"The sources changed since the review; nothing was deleted. Review again."
|
||||
)
|
||||
RECORDS_FAILED = (
|
||||
"The database could not be updated, so the file was not deleted. "
|
||||
"Its links and tags are unchanged."
|
||||
)
|
||||
FIND_FAILED = "The duplicate search failed."
|
||||
FIND_ERROR_EVENT = "An internal error occurred. Check server logs for details."
|
||||
FIND_CANCELLED = "Duplicate scan cancelled"
|
||||
|
||||
|
||||
class SourcesChanged(Exception):
|
||||
"""A source a review or removal names is no longer listed at its index."""
|
||||
|
||||
|
||||
ROW_DETAILS[KEEPER_GONE] = KEEPER_WENT
|
||||
|
||||
|
||||
def still_kept(target, hold):
|
||||
"""Right before an unlink: the held kept copy is still the file at its
|
||||
path, and the file is still not the kept copy itself (``keeper_problem``).
|
||||
"""
|
||||
try:
|
||||
if not os.path.samestat(os.fstat(hold.fileno()), os.stat(target["keeper_abs"])):
|
||||
return False
|
||||
info = os.lstat(target["abs"])
|
||||
problem = keeper_problem(
|
||||
target["abs"],
|
||||
target["root"],
|
||||
info,
|
||||
target["keeper_abs"],
|
||||
target["keeper_root"],
|
||||
)
|
||||
return problem is None
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
|
||||
def _plural(count, one, many=None):
|
||||
return f"{count:,} {one if count == 1 else (many or one + 's')}"
|
||||
|
||||
|
||||
def human_bytes(size):
|
||||
"""``10 B``, ``1.5 KB``, ``2.0 MB``…"""
|
||||
if size < 1024:
|
||||
return f"{size:,} B"
|
||||
for unit in ("KB", "MB", "GB", "TB"):
|
||||
size /= 1024
|
||||
if size < 1024 or unit == "TB":
|
||||
return f"{size:,.1f} {unit}"
|
||||
return f"{size:,.1f} TB" # pragma: no cover - unreachable
|
||||
|
||||
|
||||
def _hash_file(path):
|
||||
digest = hashlib.sha256()
|
||||
with open(path, "rb") as handle:
|
||||
for chunk in iter(lambda: handle.read(HASH_CHUNK), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def file_problem(
|
||||
path,
|
||||
root,
|
||||
digest,
|
||||
size,
|
||||
hasher=None,
|
||||
*,
|
||||
hash_check=True,
|
||||
mtime_ns=None,
|
||||
inode=None,
|
||||
keeper=None,
|
||||
keeper_root=None,
|
||||
):
|
||||
"""Why ``path`` may not be treated as a copy in ``root``, or None (blocking).
|
||||
|
||||
Containment is decided as text first, with no filesystem access: the
|
||||
normalized path must lie below the normalized root, outside
|
||||
``thumbnails``, with a media extension. Only then is the file looked at:
|
||||
a regular file, not a link, whose resolved path is still inside the
|
||||
resolved root and that is not ``keeper`` itself (``keeper_problem``: a
|
||||
hardlink, one file reached through two spellings or two mounts, or a
|
||||
volume that cannot tell files apart), of ``size`` bytes.
|
||||
``mtime_ns`` and ``inode`` (``[st_dev, st_ino]``), when given, bind it to
|
||||
what the search saw; ``hash_check`` compares its content hash with
|
||||
``digest`` (the removal job, right before an unlink). Every handle is
|
||||
closed when this returns.
|
||||
"""
|
||||
# py.config imports ComfyUI's server module: only when it is in use.
|
||||
from ..config import _canonical_path, _lexical_path, _lexically_inside
|
||||
|
||||
lexical = _lexical_path(str(path))
|
||||
base = _lexical_path(str(root))
|
||||
if lexical == base or not _lexically_inside(lexical, [base]):
|
||||
return OUTSIDE_SOURCE
|
||||
parts = Path(os.path.relpath(lexical, base)).parts
|
||||
if "thumbnails" in parts or ".." in parts:
|
||||
return OUTSIDE_SOURCE
|
||||
if os.path.splitext(lexical)[1].lower() not in DELETABLE_MEDIA_EXTENSIONS:
|
||||
return OUTSIDE_SOURCE
|
||||
try:
|
||||
info = os.lstat(path)
|
||||
except FileNotFoundError:
|
||||
return GONE
|
||||
except OSError:
|
||||
return UNREADABLE
|
||||
if stat_module.S_ISLNK(info.st_mode):
|
||||
return LINK
|
||||
if not stat_module.S_ISREG(info.st_mode):
|
||||
return LINK
|
||||
try:
|
||||
if not _lexically_inside(_canonical_path(path), [_canonical_path(root)]):
|
||||
return OUTSIDE_SOURCE # a linked folder on the way
|
||||
if keeper is not None:
|
||||
same = keeper_problem(path, root, info, keeper, keeper_root)
|
||||
if same:
|
||||
return same
|
||||
if info.st_size != size:
|
||||
return CHANGED_CONTENT
|
||||
if mtime_ns is not None and info.st_mtime_ns != mtime_ns:
|
||||
return CHANGED_CONTENT
|
||||
if inode is not None and [info.st_dev, info.st_ino] != list(inode):
|
||||
return CHANGED_CONTENT
|
||||
if hash_check and (hasher or _hash_file)(path) != digest:
|
||||
return CHANGED_CONTENT
|
||||
except FileNotFoundError:
|
||||
return GONE
|
||||
except OSError:
|
||||
return UNREADABLE
|
||||
return None
|
||||
|
||||
|
||||
def keeper_problem(path, root, info, keeper, keeper_root=None):
|
||||
"""Why ``path`` (its ``lstat`` is ``info``) may be ``keeper`` itself, or None.
|
||||
|
||||
- ``NO_FILE_ID``: a volume reports no file id (``st_ino`` 0), so the two
|
||||
cannot be told apart;
|
||||
- ``SAME``: ``os.path.samestat`` (a hardlink, or one file reached through
|
||||
two spellings of one mount);
|
||||
- ``MAYBE_SAME``: another device, the same size, mtimes within
|
||||
``MOUNT_MTIME_SLACK_NS``, and the same ``st_ino`` (NFS and SMB report
|
||||
the server's file id) or the same path within their sources, or one
|
||||
ending with the other on a folder name (one share mounted twice, a
|
||||
FUSE pool over a listed disk, sources listed at different depths).
|
||||
Protect, never guess: a true copy on another disk at the same relative
|
||||
path is kept too.
|
||||
|
||||
Raises OSError when the kept copy cannot be read (the caller's GONE or
|
||||
UNREADABLE).
|
||||
"""
|
||||
kept = os.stat(keeper)
|
||||
if not info.st_ino or not kept.st_ino:
|
||||
return NO_FILE_ID
|
||||
if os.path.samestat(info, kept):
|
||||
return SAME
|
||||
if info.st_dev == kept.st_dev or info.st_size != kept.st_size:
|
||||
return None
|
||||
if abs(info.st_mtime_ns - kept.st_mtime_ns) > MOUNT_MTIME_SLACK_NS:
|
||||
return None
|
||||
if info.st_ino == kept.st_ino:
|
||||
return MAYBE_SAME
|
||||
if keeper_root is not None and _same_tail(
|
||||
_relative_key(path, root), _relative_key(keeper, keeper_root)
|
||||
):
|
||||
return MAYBE_SAME
|
||||
return None
|
||||
|
||||
|
||||
def _same_tail(one, other):
|
||||
"""Whether one relative path is the other, or ends with it on a folder name.
|
||||
|
||||
Two mounts of one tree may be listed at different depths (``2024/x.png``
|
||||
in one source, ``x.png`` in the other); ``ax.png`` never matches ``x.png``.
|
||||
"""
|
||||
if one is None or other is None:
|
||||
return False
|
||||
return one == other or one.endswith("/" + other) or other.endswith("/" + one)
|
||||
|
||||
|
||||
def _relative_key(path, root):
|
||||
try:
|
||||
relative = os.path.relpath(str(path), str(root))
|
||||
except ValueError: # another drive (Windows)
|
||||
return None
|
||||
return os.path.normcase(relative).replace("\\", "/")
|
||||
|
||||
|
||||
@dataclass
|
||||
class HeldFind:
|
||||
"""One find: where it looked and, once it finished, what it found."""
|
||||
|
||||
origin: str # "workspace" or "gallery"
|
||||
job_id: str = ""
|
||||
roots: list = field(default_factory=list) # [(source index or None, root)]
|
||||
legacy_dir: str = None
|
||||
groups: list = None # [{group_id, hash, size, files:[record]}]
|
||||
|
||||
def covers(self, legacy_dir):
|
||||
"""Whether its events answer the legacy gallery for ``legacy_dir``."""
|
||||
if self.origin == "gallery":
|
||||
return True
|
||||
if not legacy_dir:
|
||||
return False
|
||||
wanted = _identity(legacy_dir)
|
||||
return any(_identity(root) == wanted for _index, root in self.roots)
|
||||
|
||||
|
||||
def _identity(path):
|
||||
return os.path.normcase(os.path.realpath(str(path)))
|
||||
|
||||
|
||||
def _group_order(group):
|
||||
return (-(len(group["files"]) - 1) * group["size"], group["hash"])
|
||||
|
||||
|
||||
def _keep_entry(entry):
|
||||
"""``((group_id, root_index, path), None)`` or ``(None, message)``."""
|
||||
if not isinstance(entry, dict) or set(entry) != {"group_id", "root_index", "path"}:
|
||||
return None, "Each keep entry is {group_id, root_index, path}"
|
||||
group_id, index, path = entry["group_id"], entry["root_index"], entry["path"]
|
||||
if not isinstance(group_id, str) or not 0 < len(group_id) <= 128:
|
||||
return None, "group_id must be a group's id"
|
||||
if index is not None and (
|
||||
isinstance(index, bool) or not isinstance(index, int) or index < 0
|
||||
):
|
||||
return None, "root_index must be a source index"
|
||||
if not isinstance(path, str) or not 0 < len(path) <= 4096:
|
||||
return None, "path must be a file's path in its group"
|
||||
return (group_id, index, path), None
|
||||
|
||||
|
||||
def _parse_options(options):
|
||||
"""``(job_id, keep, None)`` for duplicate_files options, else a message."""
|
||||
if not isinstance(options, dict) or set(options) != {"job_id", "keep"}:
|
||||
return None, None, "duplicate_files takes options {job_id, keep}"
|
||||
job_id, keep = options["job_id"], options["keep"]
|
||||
if not isinstance(job_id, str) or not 0 < len(job_id) <= MAX_JOB_ID:
|
||||
return None, None, "job_id must be a duplicate search's job id"
|
||||
if not isinstance(keep, list) or not 0 < len(keep) <= MAX_KEEP:
|
||||
return None, None, f"keep must list 1 to {MAX_KEEP:,} files to keep"
|
||||
parsed = []
|
||||
for entry in keep:
|
||||
key, message = _keep_entry(entry)
|
||||
if message:
|
||||
return None, None, message
|
||||
parsed.append(key)
|
||||
return job_id, parsed, None
|
||||
|
||||
|
||||
def _select(held, keep):
|
||||
"""``({group_id: (group, [keeper records])}, None)`` or ``(None, message)``."""
|
||||
groups = {group["group_id"]: group for group in held.groups}
|
||||
selection = {}
|
||||
for group_id, index, path in keep:
|
||||
group = groups.get(group_id)
|
||||
if group is None:
|
||||
return None, "A kept file names a group this search did not find"
|
||||
match = [
|
||||
record
|
||||
for record in group["files"]
|
||||
if record["root_index"] == index and record["rel"] == path
|
||||
]
|
||||
if not match:
|
||||
return None, "A kept file is not in its group"
|
||||
_group, keepers = selection.setdefault(group_id, (group, []))
|
||||
if match[0] not in keepers:
|
||||
keepers.append(match[0])
|
||||
return selection, None
|
||||
|
||||
|
||||
SOURCES_CHANGED_REVIEW = (
|
||||
"The sources changed since the search. Search for duplicates again."
|
||||
)
|
||||
|
||||
|
||||
def _fingerprint(targets):
|
||||
digest = hashlib.sha256()
|
||||
for target in sorted(targets, key=lambda t: _identity(t["abs"])):
|
||||
part = [
|
||||
os.path.normcase(target["abs"]),
|
||||
target["hash"],
|
||||
target["size"],
|
||||
os.path.normcase(target["keeper_abs"]),
|
||||
]
|
||||
digest.update(json.dumps(part).encode("utf-8") + b"\n")
|
||||
return "sha256:" + digest.hexdigest()
|
||||
|
||||
|
||||
def _record_parts(records, target):
|
||||
"""What moves with one removal: its links, tags and Needs review items."""
|
||||
found = records.get(target["abs"], {})
|
||||
links, tags, review = (found.get(key, 0) for key in ("links", "tags", "review"))
|
||||
parts = [
|
||||
(
|
||||
(
|
||||
f"{_plural(links, 'database link')} "
|
||||
f"{'moves' if links == 1 else 'move'} to the kept copy"
|
||||
)
|
||||
if links
|
||||
else "no database links"
|
||||
)
|
||||
]
|
||||
if tags:
|
||||
parts.append("its image tags join the kept copy's")
|
||||
if review:
|
||||
item = "Needs review item is" if review == 1 else "Needs review items are"
|
||||
parts.append(f"{review:,} {item} settled")
|
||||
return parts
|
||||
|
||||
|
||||
def _public_plan(targets, protected, records):
|
||||
count = len(targets)
|
||||
groups = len({t["group_id"] for t in targets})
|
||||
size = sum(t["size"] for t in targets)
|
||||
reasons = [
|
||||
f"{n:,} {PROTECTED_REASONS[code][0 if n == 1 else 1]}"
|
||||
for code, n in sorted(protected.items())
|
||||
]
|
||||
details = [
|
||||
{
|
||||
"label": target["rel"],
|
||||
"detail": f"{target['root_label']} · a copy of {target['keeper_rel']}"
|
||||
+ "".join(f" · {part}" for part in _record_parts(records, target)),
|
||||
}
|
||||
for target in targets[:MAX_DETAILS]
|
||||
]
|
||||
label = f"{_plural(count, 'duplicate file')} in {_plural(groups, 'group')}"
|
||||
media = (
|
||||
f"Deletes {_plural(count, 'file')} ({human_bytes(size)}) from disk; the "
|
||||
"kept copies stay. The safety backup covers the database, not these files."
|
||||
)
|
||||
tagged = sum(1 for t in targets if records.get(t["abs"], {}).get("tags"))
|
||||
review = sum(records.get(t["abs"], {}).get("review", 0) for t in targets)
|
||||
if tagged:
|
||||
media += (
|
||||
f" The image tags of {_plural(tagged, 'removed copy', 'removed copies')}"
|
||||
" join their kept copy's."
|
||||
)
|
||||
if review:
|
||||
media += (
|
||||
f" {_plural(review, 'pending Needs review item')} of the removed "
|
||||
f"{'copy is' if review == 1 else 'copies are'} settled."
|
||||
)
|
||||
return {
|
||||
"operation": DUPLICATE_FILES,
|
||||
"scope": {"count": count, "label": label},
|
||||
"protected": {"count": sum(protected.values()), "reasons": reasons},
|
||||
"media": {"deletes_files": True, "label": media},
|
||||
"details": details,
|
||||
}
|
||||
@@ -1,133 +0,0 @@
|
||||
"""Every path a duplicate copy is recorded under (#249 System S6, Task 17).
|
||||
|
||||
The search walks each source's real path, but a record names a file the way
|
||||
it was seen when it was written: under the folder the image monitor watches,
|
||||
the configured spelling, ComfyUI's output folder, a symlink to it, or a
|
||||
parent of a source listed later. ``Spellings`` finds them by containment:
|
||||
for each known folder spelling ``S`` (resolving to the folder ``R``) and a
|
||||
file ``P`` inside ``R``, the file is also recorded as ``S + relpath(P, R)``.
|
||||
|
||||
Folders are compared by ``os.stat`` identity (``st_dev``, ``st_ino``), not
|
||||
by text, so a case variant on a case-insensitive filesystem (macOS, Windows)
|
||||
or a symlinked spelling matches. A volume that reports no file id falls
|
||||
back to ``normcase(realpath)``. Favorites are keyed by a listed root and the
|
||||
path below it, so ``favorites`` lists every listed root that holds a file.
|
||||
Blocking: each lookup stats the file's folders once (cached).
|
||||
"""
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from .images import _favorite_relative_key, _favorite_root_key
|
||||
|
||||
|
||||
def _identity(path):
|
||||
try:
|
||||
info = os.stat(path)
|
||||
except OSError:
|
||||
return None
|
||||
if info.st_ino:
|
||||
return ("inode", info.st_dev, info.st_ino)
|
||||
return ("path", os.path.normcase(os.path.realpath(str(path))))
|
||||
|
||||
|
||||
class Spellings:
|
||||
"""Folder spellings and listed roots, matched to files by containment."""
|
||||
|
||||
def __init__(self, folders, listed_roots):
|
||||
self._cache = {}
|
||||
self._folders = []
|
||||
for spelled in dict.fromkeys(str(f) for f in folders if f):
|
||||
identity = self._identity(spelled)
|
||||
if identity is not None:
|
||||
self._folders.append((spelled, identity))
|
||||
self._listed = []
|
||||
for root in dict.fromkeys(str(r) for r in listed_roots if r):
|
||||
identity = self._identity(root)
|
||||
if identity is not None:
|
||||
self._listed.append((root, identity))
|
||||
|
||||
def _identity(self, path):
|
||||
if path not in self._cache:
|
||||
self._cache[path] = _identity(path)
|
||||
return self._cache[path]
|
||||
|
||||
def _ancestors(self, path):
|
||||
"""``{identity: folder}`` of every folder above ``path``, nearest first."""
|
||||
found = {}
|
||||
folder = os.path.dirname(os.path.abspath(str(path)))
|
||||
while True:
|
||||
identity = self._identity(folder)
|
||||
if identity is not None:
|
||||
found.setdefault(identity, folder)
|
||||
parent = os.path.dirname(folder)
|
||||
if parent == folder:
|
||||
return found
|
||||
folder = parent
|
||||
|
||||
def of(self, path):
|
||||
"""``{folder spelling: path spelled under it}``; ``""`` is ``path`` itself."""
|
||||
path = os.path.abspath(str(path))
|
||||
spelled = {"": path}
|
||||
ancestors = self._ancestors(path)
|
||||
for folder, identity in self._folders:
|
||||
inside = ancestors.get(identity)
|
||||
if inside is not None:
|
||||
spelled.setdefault(
|
||||
folder, os.path.join(folder, os.path.relpath(path, inside))
|
||||
)
|
||||
return spelled
|
||||
|
||||
def favorites(self, path):
|
||||
"""``{listed root: path below it (POSIX)}`` for each listed root holding it."""
|
||||
path = os.path.abspath(str(path))
|
||||
ancestors = self._ancestors(path)
|
||||
held = {}
|
||||
for root, identity in self._listed:
|
||||
inside = ancestors.get(identity)
|
||||
if inside is not None:
|
||||
held[root] = Path(os.path.relpath(path, inside)).as_posix()
|
||||
return held
|
||||
|
||||
|
||||
def copy_move(target, kept_stat, spellings):
|
||||
"""What ``move_copy_records`` moves for one removal.
|
||||
|
||||
A record found under a folder spelling moves to the kept copy under
|
||||
the same spelling when that folder holds it too, else to its real
|
||||
path. Its favorite moves under every listed root holding the file.
|
||||
"""
|
||||
removed = spellings.of(target["abs"])
|
||||
kept = spellings.of(target["keeper_abs"])
|
||||
targets = {
|
||||
spelled: kept.get(folder, target["keeper_abs"])
|
||||
for folder, spelled in removed.items()
|
||||
}
|
||||
kept_favorites = spellings.favorites(target["keeper_abs"])
|
||||
favorites = []
|
||||
for root, relative in spellings.favorites(target["abs"]).items():
|
||||
kept_root, kept_rel = (
|
||||
(root, kept_favorites[root])
|
||||
if root in kept_favorites
|
||||
else (target["keeper_root"], target["keeper_rel"])
|
||||
)
|
||||
favorites.append(
|
||||
(
|
||||
_favorite_root_key(root),
|
||||
_favorite_relative_key(relative),
|
||||
_favorite_root_key(kept_root),
|
||||
_favorite_relative_key(kept_rel),
|
||||
)
|
||||
)
|
||||
return {
|
||||
"removed": list(removed.values()),
|
||||
"kept": list(kept.values()),
|
||||
"targets": targets,
|
||||
"kept_path": target["keeper_abs"],
|
||||
"favorites": favorites,
|
||||
"removed_locator": target["locator"],
|
||||
"kept_locator": target["keeper_locator"],
|
||||
"sha256": target["hash"],
|
||||
"kept_size": kept_stat.st_size,
|
||||
"kept_mtime_ns": kept_stat.st_mtime_ns,
|
||||
}
|
||||
+20
-42
@@ -734,10 +734,7 @@ class HistoryRoutesMixin:
|
||||
raise HistoryReadError(422, "response-limit") from None
|
||||
else:
|
||||
prepared = _prepare_current_media(
|
||||
resolved,
|
||||
scope,
|
||||
thumbnail=kind == "thumbnail",
|
||||
**media_options,
|
||||
resolved, scope, thumbnail=kind == "thumbnail"
|
||||
)
|
||||
_check_current(scope, self, fence, epoch, generation)
|
||||
return prepared
|
||||
@@ -764,44 +761,25 @@ class HistoryRoutesMixin:
|
||||
_check_current(scope, self, fence, epoch, generation)
|
||||
raise
|
||||
|
||||
async def execute():
|
||||
pending = asyncio.create_task(self._run_in_executor(project))
|
||||
try:
|
||||
return await asyncio.shield(pending)
|
||||
except asyncio.CancelledError:
|
||||
pending.add_done_callback(_close_prepared_result)
|
||||
raise
|
||||
except (
|
||||
HistoryReadError,
|
||||
HistoryEpochChanged,
|
||||
sqlite3.Error,
|
||||
OSError,
|
||||
ValueError,
|
||||
):
|
||||
raise
|
||||
except Exception as error:
|
||||
self.logger.warning(
|
||||
"Current history media unavailable (%s)", type(error).__name__
|
||||
)
|
||||
raise _current_unavailable("temporary_unavailable") from None
|
||||
|
||||
# A missing preview (#249) is noted on the first read, made through
|
||||
# the serve route's admission and single flight, then read again. A
|
||||
# video whose preview cannot be made is the placeholder frame.
|
||||
needed = []
|
||||
media_options = {}
|
||||
if kind == "thumbnail":
|
||||
media_options = {
|
||||
"ensure_preview": lambda *preview: needed.append(preview),
|
||||
"placeholder": self._video_placeholder_bytes(),
|
||||
}
|
||||
prepared = await execute()
|
||||
if needed:
|
||||
if not isinstance(prepared, bytes):
|
||||
prepared.close()
|
||||
await self._admit_preview(*needed[0])
|
||||
media_options["ensure_preview"] = None
|
||||
prepared = await execute()
|
||||
pending = asyncio.create_task(self._run_in_executor(project))
|
||||
try:
|
||||
prepared = await asyncio.shield(pending)
|
||||
except asyncio.CancelledError:
|
||||
pending.add_done_callback(_close_prepared_result)
|
||||
raise
|
||||
except (
|
||||
HistoryReadError,
|
||||
HistoryEpochChanged,
|
||||
sqlite3.Error,
|
||||
OSError,
|
||||
ValueError,
|
||||
):
|
||||
raise
|
||||
except Exception as error:
|
||||
self.logger.warning(
|
||||
"Current history media unavailable (%s)", type(error).__name__
|
||||
)
|
||||
raise _current_unavailable("temporary_unavailable") from None
|
||||
if isinstance(prepared, bytes):
|
||||
return _success_response(prepared)
|
||||
return await _current_media_response(request, prepared)
|
||||
|
||||
@@ -5,7 +5,6 @@ from dataclasses import dataclass
|
||||
import errno
|
||||
import hashlib
|
||||
import hmac
|
||||
import io
|
||||
import json
|
||||
import mimetypes
|
||||
import os
|
||||
@@ -310,35 +309,17 @@ class _OpenedCurrentMedia:
|
||||
self.handle.close()
|
||||
|
||||
|
||||
def _preview_writable(root: Path, preview: Path) -> bool:
|
||||
"""A preview may be written: ``thumbnails`` and every link segment of the
|
||||
preview's folder resolve inside *root* (the serve route's rule, #249)."""
|
||||
folder = root / "thumbnails"
|
||||
return _canonical(folder).is_relative_to(_canonical(root)) and _canonical(
|
||||
preview
|
||||
).is_relative_to(_canonical(folder))
|
||||
|
||||
|
||||
def _prepare_current_media(
|
||||
resolved: ResolvedCurrentCandidate,
|
||||
scope: CurrentOutputScope,
|
||||
*,
|
||||
thumbnail: bool,
|
||||
ensure_preview=None,
|
||||
placeholder: bytes | None = None,
|
||||
resolved: ResolvedCurrentCandidate, scope: CurrentOutputScope, *, thumbnail: bool
|
||||
) -> _OpenedCurrentMedia:
|
||||
"""Open the current file, or its preview for ``thumbnail``.
|
||||
|
||||
A missing preview that may be written is handed to
|
||||
``ensure_preview(src, dst, is_video)`` (the route makes it and reads
|
||||
again, #249). Without one an image's original is opened and a video is
|
||||
the ``placeholder`` JPEG, never the video inside an ``<img>``.
|
||||
"""
|
||||
target = resolved.canonical_target
|
||||
if thumbnail:
|
||||
relative = target.relative_to(resolved.gallery_root)
|
||||
is_video = target.suffix.lower() in VIDEO_EXTENSIONS
|
||||
extension = ".jpg" if is_video else target.suffix.lower()
|
||||
extension = (
|
||||
".jpg"
|
||||
if target.suffix.lower() in VIDEO_EXTENSIONS
|
||||
else target.suffix.lower()
|
||||
)
|
||||
preview = (
|
||||
resolved.gallery_root
|
||||
/ "thumbnails"
|
||||
@@ -347,14 +328,7 @@ def _prepare_current_media(
|
||||
try:
|
||||
preview.stat()
|
||||
except FileNotFoundError:
|
||||
if ensure_preview is not None and _preview_writable(
|
||||
resolved.gallery_root, preview
|
||||
):
|
||||
ensure_preview(target, preview, is_video)
|
||||
if is_video and placeholder is not None:
|
||||
return _OpenedCurrentMedia(
|
||||
io.BytesIO(placeholder), len(placeholder), "image/jpeg"
|
||||
)
|
||||
pass
|
||||
else:
|
||||
target = _canonical(preview)
|
||||
handle = _open_current_target(target, resolved.namespace, scope)
|
||||
|
||||
+137
-259
@@ -110,8 +110,6 @@ def _card_header(path: Path) -> tuple:
|
||||
|
||||
# The legacy gallery page's thumbnail generation, as Activity lists it (D80).
|
||||
THUMBNAIL_GENERATION_TITLE = "Thumbnail generation (gallery page)"
|
||||
# Its clear: automatic previews are not written while it runs.
|
||||
THUMBNAIL_CLEAR_TITLE = "Thumbnail clear (gallery page)"
|
||||
|
||||
|
||||
class _CardFacts(NamedTuple):
|
||||
@@ -263,32 +261,13 @@ def _temp_sibling(dst):
|
||||
return dst.with_name(f".{dst.stem}.part-{token}{dst.suffix}")
|
||||
|
||||
|
||||
class PreviewWriteHeld(Exception):
|
||||
"""A preview write stopped before its rename: ``guard()`` said no (#249)."""
|
||||
|
||||
|
||||
# The result error of a write held before its rename (never remembered).
|
||||
PREVIEW_HELD = "Preview writes are held while another job runs"
|
||||
|
||||
|
||||
def _replace_unless_held(tmp, dst, guard=None):
|
||||
"""Rename *tmp* over *dst*, unless ``guard()`` is false (PreviewWriteHeld)."""
|
||||
if guard is not None and not guard():
|
||||
raise PreviewWriteHeld(PREVIEW_HELD)
|
||||
os.replace(tmp, dst)
|
||||
|
||||
|
||||
def _atomic_save(img, dst, fmt=None, *, guard=None, **save_kwargs):
|
||||
def _atomic_save(img, dst, fmt=None, **save_kwargs):
|
||||
"""Save a PIL image to *dst* without a partially written file ever being
|
||||
visible at that path: write a temp sibling, then rename over it.
|
||||
|
||||
``guard()`` (automatic previews) is asked again just before the rename;
|
||||
false leaves *dst* untouched and raises PreviewWriteHeld.
|
||||
"""
|
||||
visible at that path: write a temp sibling, then rename over it."""
|
||||
tmp = _temp_sibling(dst)
|
||||
try:
|
||||
img.save(tmp, fmt, **save_kwargs)
|
||||
_replace_unless_held(tmp, dst, guard)
|
||||
os.replace(tmp, dst)
|
||||
finally:
|
||||
if tmp.exists():
|
||||
tmp.unlink()
|
||||
@@ -326,11 +305,11 @@ def _ffmpeg_qscale(encoding):
|
||||
return str(_FFMPEG_QSCALE.get(encoding["quality"], 4))
|
||||
|
||||
|
||||
def _write_image_thumbnail(src, dst, thumbnail_size, encoding=None, guard=None):
|
||||
def _write_image_thumbnail(src, dst, thumbnail_size, encoding=None):
|
||||
"""Resize *src* into *dst* (blocking PIL work), atomically.
|
||||
|
||||
``encoding`` is one of THUMBNAIL_ENCODINGS (medium when None); a PNG
|
||||
preview keeps only its ``optimize`` flag. ``guard`` see _atomic_save.
|
||||
preview keeps only its ``optimize`` flag.
|
||||
"""
|
||||
encoding = encoding or DEFAULT_THUMBNAIL_ENCODING
|
||||
with Image.open(src) as img:
|
||||
@@ -340,50 +319,7 @@ def _write_image_thumbnail(src, dst, thumbnail_size, encoding=None, guard=None):
|
||||
save_kwargs = dict(encoding)
|
||||
if dst.suffix.lower() == ".png":
|
||||
save_kwargs = {"optimize": encoding["optimize"]}
|
||||
_atomic_save(img, dst, guard=guard, **save_kwargs)
|
||||
|
||||
|
||||
def _video_placeholder_image(thumbnail_size):
|
||||
"""The placeholder frame for a video no decoder could read: a play mark."""
|
||||
from PIL import ImageDraw, ImageFont
|
||||
|
||||
img = Image.new("RGB", thumbnail_size, color=(50, 50, 50))
|
||||
draw = ImageDraw.Draw(img)
|
||||
center_x, center_y = thumbnail_size[0] // 2, thumbnail_size[1] // 2
|
||||
triangle_size = min(thumbnail_size) // 4
|
||||
points = [
|
||||
(center_x - triangle_size // 2, center_y - triangle_size // 2),
|
||||
(center_x - triangle_size // 2, center_y + triangle_size // 2),
|
||||
(center_x + triangle_size // 2, center_y),
|
||||
]
|
||||
draw.polygon(points, fill=(255, 255, 255))
|
||||
try:
|
||||
font = ImageFont.load_default()
|
||||
text = "VIDEO"
|
||||
bbox = draw.textbbox((0, 0), text, font=font)
|
||||
text_width = bbox[2] - bbox[0]
|
||||
draw.text(
|
||||
(center_x - text_width // 2, center_y + triangle_size // 2 + 10),
|
||||
text,
|
||||
fill=(255, 255, 255),
|
||||
font=font,
|
||||
)
|
||||
except (OSError, AttributeError):
|
||||
pass
|
||||
return img
|
||||
|
||||
|
||||
def _is_decode_error(exc):
|
||||
"""True when *exc* says the file itself cannot be decoded (#249 M4).
|
||||
|
||||
Pillow's decoder errors are OSErrors without an errno; a real I/O error
|
||||
(disk full, a share gone) carries one and is worth another try.
|
||||
"""
|
||||
if isinstance(exc, (UnidentifiedImageError, Image.DecompressionBombError)):
|
||||
return True
|
||||
if isinstance(exc, SyntaxError):
|
||||
return True
|
||||
return isinstance(exc, OSError) and exc.errno is None
|
||||
_atomic_save(img, dst, **save_kwargs)
|
||||
|
||||
|
||||
def _new_thumbnail_stats():
|
||||
@@ -774,51 +710,23 @@ def _iter_media_files(output_path, max_files):
|
||||
yield root_path / name
|
||||
|
||||
|
||||
def media_thumbnail_rel(rel_path):
|
||||
"""``thumbnails/<rel>_thumb<ext>`` (POSIX) for the media file *rel_path*.
|
||||
|
||||
The extension is lower case, and a video's preview is a ``.jpg``: the
|
||||
name every thumbnail writer (the plan, the job, automatic previews) uses.
|
||||
"""
|
||||
def media_thumbnail_url(output_path, rel_path, root_index, *, root_scope=None):
|
||||
"""The Images listing's thumbnail URL for *rel_path* under *output_path*,
|
||||
or None when no thumbnail exists (one blocking stat). The Needs review
|
||||
queue cards use the same URL (#249 F, M10)."""
|
||||
rel_path = Path(rel_path)
|
||||
extension = rel_path.suffix.lower()
|
||||
return (
|
||||
thumb_rel = (
|
||||
f"thumbnails/{rel_path.with_suffix('').as_posix()}_thumb"
|
||||
f"{'.jpg' if extension in VIDEO_EXTENSIONS else extension}"
|
||||
)
|
||||
|
||||
|
||||
def media_thumbnail(output_path, rel_path, root_index, *, root_scope=None):
|
||||
"""``(thumbnail_url, ready)`` for *rel_path* under *output_path* (one
|
||||
blocking stat). The Images listing, the Needs review queue cards and the
|
||||
Library use the same URL (#249 F, M10).
|
||||
|
||||
An existing preview's URL carries its mtime as ``v=`` and ``ready`` is
|
||||
True. A missing one is still listed (without ``v=``): the serve route
|
||||
makes it on first view and falls back when it cannot
|
||||
(py/api/thumbnail_auto.py), so no card loads the full file because its
|
||||
preview was not made yet. Media inside a nested ``thumbnails`` folder is
|
||||
never previewed: ``(None, False)``, and its card shows the original.
|
||||
"""
|
||||
rel_path = Path(rel_path)
|
||||
if "thumbnails" in rel_path.parts[:-1]:
|
||||
return None, False
|
||||
thumb_rel = media_thumbnail_rel(rel_path)
|
||||
thumbnail_url = thumbnail_url_for(thumb_rel, Path(output_path) / thumb_rel)
|
||||
ready = thumbnail_url is not None
|
||||
if thumbnail_url is None:
|
||||
quoted = urllib.parse.quote(thumb_rel, safe="/")
|
||||
thumbnail_url = f"/prompt_manager/images/serve/{quoted}?root={root_index}"
|
||||
else:
|
||||
thumbnail_url += f"&root={root_index}"
|
||||
return None
|
||||
thumbnail_url += f"&root={root_index}"
|
||||
if root_scope is not None:
|
||||
thumbnail_url += f"&root_scope={root_scope}"
|
||||
return thumbnail_url, ready
|
||||
|
||||
|
||||
def media_thumbnail_url(output_path, rel_path, root_index, *, root_scope=None):
|
||||
"""The thumbnail URL of ``media_thumbnail`` (None for nested previews)."""
|
||||
return media_thumbnail(output_path, rel_path, root_index, root_scope=root_scope)[0]
|
||||
return thumbnail_url
|
||||
|
||||
|
||||
def _output_image_entry(
|
||||
@@ -836,7 +744,7 @@ def _output_image_entry(
|
||||
extension = media_path.suffix.lower()
|
||||
is_video = extension in VIDEO_EXTENSIONS
|
||||
|
||||
thumbnail_url, thumbnail_ready = media_thumbnail(
|
||||
thumbnail_url = media_thumbnail_url(
|
||||
output_path, rel_path, root_index, root_scope=root_scope
|
||||
)
|
||||
encoded_path = urllib.parse.quote(rel_path.as_posix(), safe="/")
|
||||
@@ -855,8 +763,6 @@ def _output_image_entry(
|
||||
"root_index": root_index,
|
||||
"url": url,
|
||||
"thumbnail_url": thumbnail_url,
|
||||
# Whether the preview file exists yet (the legacy gallery's dot, #249).
|
||||
"thumbnail_ready": thumbnail_ready,
|
||||
"size": stat.st_size,
|
||||
"modified_time": stat.st_mtime,
|
||||
"extension": extension,
|
||||
@@ -1776,11 +1682,6 @@ class ImageRoutesMixin:
|
||||
if image_path.is_file():
|
||||
return self._file_response(image_path)
|
||||
|
||||
# A listed preview that was not made yet is made now (#249).
|
||||
preview = await self._serve_missing_thumbnail(output_dirs, filepath)
|
||||
if preview is not None:
|
||||
return preview
|
||||
|
||||
# ComfyUI's cache middleware would otherwise let browsers cache this
|
||||
# miss for an hour, hiding a thumbnail generated a moment later.
|
||||
return web.json_response(
|
||||
@@ -2008,8 +1909,6 @@ class ImageRoutesMixin:
|
||||
"total_images": pending + plan["skipped"],
|
||||
"pending": pending,
|
||||
"ready": plan["skipped"],
|
||||
# Files whose automatic preview failed this session (#249).
|
||||
"failed": self._auto_thumbnail_failures(),
|
||||
"collection_id": hashlib.sha256(collection.encode()).hexdigest()[
|
||||
:16
|
||||
],
|
||||
@@ -2104,8 +2003,6 @@ class ImageRoutesMixin:
|
||||
encoding=None,
|
||||
rebuild=False,
|
||||
stop=None,
|
||||
guard=None,
|
||||
track=None,
|
||||
):
|
||||
"""Create one thumbnail (blocking). Never raises; returns a result dict.
|
||||
|
||||
@@ -2113,11 +2010,7 @@ class ImageRoutesMixin:
|
||||
and touches nothing (a job's file boundary). ``rebuild`` writes a
|
||||
fresh preview again and never replaces an existing video preview with
|
||||
the placeholder. Every write is atomic, so a failure keeps the old
|
||||
preview. ``guard()`` false just before the rename gives ``action:
|
||||
"error"`` with ``held: True`` and writes nothing; ``track(process)``
|
||||
is handed the ffmpeg process (both for automatic previews, #249). An
|
||||
error result carries ``decode: True`` when the file itself cannot be
|
||||
decoded (an I/O error does not).
|
||||
preview.
|
||||
"""
|
||||
result = {
|
||||
"file": src.name,
|
||||
@@ -2134,43 +2027,22 @@ class ImageRoutesMixin:
|
||||
dst.parent.mkdir(parents=True, exist_ok=True)
|
||||
if is_video:
|
||||
made = self._generate_video_thumbnail(
|
||||
src,
|
||||
dst,
|
||||
thumbnail_size,
|
||||
encoding,
|
||||
keep_existing=rebuild,
|
||||
guard=guard,
|
||||
track=track,
|
||||
src, dst, thumbnail_size, encoding, keep_existing=rebuild
|
||||
)
|
||||
if not made:
|
||||
return {
|
||||
**result,
|
||||
"action": "error",
|
||||
"error": "Video decode failed",
|
||||
"decode": True,
|
||||
}
|
||||
return {**result, "action": "error", "error": "Video decode failed"}
|
||||
else:
|
||||
_write_image_thumbnail(src, dst, thumbnail_size, encoding, guard)
|
||||
_write_image_thumbnail(src, dst, thumbnail_size, encoding)
|
||||
return {**result, "action": "generated"}
|
||||
except PreviewWriteHeld:
|
||||
return {**result, "action": "error", "error": PREVIEW_HELD, "held": True}
|
||||
except UnidentifiedImageError:
|
||||
return {
|
||||
**result,
|
||||
"action": "error",
|
||||
"error": "Not a valid image file",
|
||||
"decode": True,
|
||||
}
|
||||
return {**result, "action": "error", "error": "Not a valid image file"}
|
||||
except PermissionError as e:
|
||||
message = _preview_write_error(e)
|
||||
self.logger.warning(f"Thumbnail failed for {src.name}: {message}")
|
||||
return {**result, "action": "error", "error": message}
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Thumbnail failed for {src.name}", exc_info=True)
|
||||
failed = {**result, "action": "error", "error": _safe_error(e)}
|
||||
if _is_decode_error(e):
|
||||
failed["decode"] = True
|
||||
return failed
|
||||
return {**result, "action": "error", "error": _safe_error(e)}
|
||||
|
||||
async def _emit_sse(self, response, event, data) -> bool:
|
||||
"""Write one SSE frame. Returns False once the client has gone away."""
|
||||
@@ -2189,59 +2061,147 @@ class ImageRoutesMixin:
|
||||
thumbnail_size,
|
||||
encoding=None,
|
||||
keep_existing=False,
|
||||
*,
|
||||
guard=None,
|
||||
track=None,
|
||||
):
|
||||
"""Generate thumbnail from video file. Returns True if successful.
|
||||
|
||||
Tries OpenCV, then ffmpeg, then writes the placeholder frame: a
|
||||
decoder that cannot open or read the clip (an AV1 file OpenCV was
|
||||
built without) falls through to the next one (#249 I1). ``encoding``
|
||||
is one of THUMBNAIL_ENCODINGS (medium when None); it encodes the
|
||||
OpenCV frame, ffmpeg's frame (as ``-q:v``) and the placeholder alike.
|
||||
Both decoders fit the frame inside ``thumbnail_size`` keeping its
|
||||
aspect ratio. With ``keep_existing`` an existing preview is kept, and
|
||||
the run fails, when no frame could be decoded, rather than replaced
|
||||
by the placeholder. A PermissionError while replacing the preview
|
||||
propagates (``_generate_one`` names it), as does PreviewWriteHeld
|
||||
when ``guard()`` refuses the rename. ``track(process)`` is handed the
|
||||
ffmpeg process so a shutdown can kill it.
|
||||
``encoding`` is one of THUMBNAIL_ENCODINGS (medium when None); it
|
||||
encodes the OpenCV frame, ffmpeg's frame (as ``-q:v``) and the
|
||||
placeholder alike. Both decoders fit the frame inside
|
||||
``thumbnail_size`` keeping its aspect ratio. With ``keep_existing`` an
|
||||
existing preview is kept, and the run fails, when no frame could be
|
||||
decoded, rather than replaced by the placeholder. A PermissionError
|
||||
while replacing the preview propagates (``_generate_one`` names it).
|
||||
"""
|
||||
encoding = encoding or DEFAULT_THUMBNAIL_ENCODING
|
||||
try:
|
||||
frame = self._opencv_frame(video_path, thumbnail_size)
|
||||
if frame is not None:
|
||||
_atomic_save(frame, thumbnail_path, "JPEG", guard=guard, **encoding)
|
||||
# Try using OpenCV first (most reliable)
|
||||
try:
|
||||
import cv2
|
||||
|
||||
cap = cv2.VideoCapture(str(video_path))
|
||||
if not cap.isOpened():
|
||||
return False
|
||||
|
||||
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
target_frame = max(1, int(frame_count * 0.1))
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, target_frame)
|
||||
|
||||
ret, frame = cap.read()
|
||||
cap.release()
|
||||
|
||||
if not ret or frame is None:
|
||||
return False
|
||||
|
||||
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
|
||||
img = Image.fromarray(frame_rgb)
|
||||
img.thumbnail(thumbnail_size, Image.Resampling.LANCZOS)
|
||||
_atomic_save(img, thumbnail_path, "JPEG", **encoding)
|
||||
self.logger.debug(
|
||||
f"Generated video thumbnail using OpenCV: {thumbnail_path}"
|
||||
)
|
||||
return True
|
||||
if self._ffmpeg_frame(
|
||||
video_path, thumbnail_path, thumbnail_size, encoding, guard, track
|
||||
):
|
||||
return True
|
||||
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# Fallback to ffmpeg
|
||||
try:
|
||||
import subprocess
|
||||
|
||||
tmp_path = _temp_sibling(thumbnail_path)
|
||||
cmd = [
|
||||
"ffmpeg",
|
||||
"-i",
|
||||
str(video_path),
|
||||
"-ss",
|
||||
"00:00:01",
|
||||
"-vframes",
|
||||
"1",
|
||||
"-vf",
|
||||
f"scale={thumbnail_size[0]}:{thumbnail_size[1]}"
|
||||
":force_original_aspect_ratio=decrease",
|
||||
"-q:v",
|
||||
_ffmpeg_qscale(encoding),
|
||||
"-y",
|
||||
str(tmp_path),
|
||||
]
|
||||
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
|
||||
if result.returncode == 0 and tmp_path.exists():
|
||||
os.replace(tmp_path, thumbnail_path)
|
||||
self.logger.debug(
|
||||
f"Generated video thumbnail using ffmpeg: {thumbnail_path}"
|
||||
)
|
||||
return True
|
||||
else:
|
||||
self.logger.warning(
|
||||
f"ffmpeg failed for {video_path}: {result.stderr}"
|
||||
)
|
||||
|
||||
except (ImportError, subprocess.TimeoutExpired, FileNotFoundError):
|
||||
pass
|
||||
finally:
|
||||
try:
|
||||
tmp_path.unlink()
|
||||
except (NameError, OSError):
|
||||
pass
|
||||
|
||||
if keep_existing and thumbnail_path.exists():
|
||||
self.logger.warning(
|
||||
f"No frame decoded for {video_path.name}; keeping its preview"
|
||||
)
|
||||
return False
|
||||
# Last resort: the placeholder frame.
|
||||
|
||||
# Last resort: create a placeholder thumbnail
|
||||
try:
|
||||
img = _video_placeholder_image(thumbnail_size)
|
||||
_atomic_save(img, thumbnail_path, "JPEG", guard=guard, **encoding)
|
||||
from PIL import ImageDraw, ImageFont
|
||||
|
||||
img = Image.new("RGB", thumbnail_size, color=(50, 50, 50))
|
||||
draw = ImageDraw.Draw(img)
|
||||
|
||||
center_x, center_y = thumbnail_size[0] // 2, thumbnail_size[1] // 2
|
||||
triangle_size = min(thumbnail_size) // 4
|
||||
|
||||
points = [
|
||||
(center_x - triangle_size // 2, center_y - triangle_size // 2),
|
||||
(center_x - triangle_size // 2, center_y + triangle_size // 2),
|
||||
(center_x + triangle_size // 2, center_y),
|
||||
]
|
||||
draw.polygon(points, fill=(255, 255, 255))
|
||||
|
||||
try:
|
||||
font = ImageFont.load_default()
|
||||
text = "VIDEO"
|
||||
bbox = draw.textbbox((0, 0), text, font=font)
|
||||
text_width = bbox[2] - bbox[0]
|
||||
draw.text(
|
||||
(
|
||||
center_x - text_width // 2,
|
||||
center_y + triangle_size // 2 + 10,
|
||||
),
|
||||
text,
|
||||
fill=(255, 255, 255),
|
||||
font=font,
|
||||
)
|
||||
except (OSError, AttributeError):
|
||||
pass
|
||||
|
||||
_atomic_save(img, thumbnail_path, "JPEG", **encoding)
|
||||
self.logger.debug(
|
||||
f"Generated placeholder video thumbnail: {thumbnail_path}"
|
||||
)
|
||||
return True
|
||||
except (PermissionError, PreviewWriteHeld):
|
||||
|
||||
except PermissionError:
|
||||
raise
|
||||
except Exception as e:
|
||||
self.logger.warning(
|
||||
f"Failed to create placeholder thumbnail for {video_path}: {e}"
|
||||
)
|
||||
return False
|
||||
except (PermissionError, PreviewWriteHeld):
|
||||
|
||||
except PermissionError:
|
||||
raise
|
||||
except Exception as e:
|
||||
self.logger.error(
|
||||
@@ -2249,88 +2209,6 @@ class ImageRoutesMixin:
|
||||
)
|
||||
return False
|
||||
|
||||
def _opencv_frame(self, video_path, thumbnail_size):
|
||||
"""A frame about 10% in, fitted to *thumbnail_size*, or None (no
|
||||
OpenCV, or it could not open or read the clip)."""
|
||||
try:
|
||||
import cv2
|
||||
except ImportError:
|
||||
return None
|
||||
try:
|
||||
cap = cv2.VideoCapture(str(video_path))
|
||||
try:
|
||||
if not cap.isOpened():
|
||||
return None
|
||||
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, max(1, int(frame_count * 0.1)))
|
||||
ret, frame = cap.read()
|
||||
finally:
|
||||
cap.release()
|
||||
if not ret or frame is None:
|
||||
return None
|
||||
img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||||
img.thumbnail(thumbnail_size, Image.Resampling.LANCZOS)
|
||||
return img
|
||||
except Exception as e:
|
||||
self.logger.debug(f"OpenCV could not read {video_path}: {e}")
|
||||
return None
|
||||
|
||||
def _ffmpeg_frame(
|
||||
self, video_path, thumbnail_path, thumbnail_size, encoding, guard, track
|
||||
):
|
||||
"""Write ffmpeg's frame at 1 s into *thumbnail_path*; True on success."""
|
||||
import subprocess # nosec B404 - fixed argv, no shell
|
||||
|
||||
tmp_path = _temp_sibling(thumbnail_path)
|
||||
cmd = [
|
||||
"ffmpeg",
|
||||
"-i",
|
||||
str(video_path),
|
||||
"-ss",
|
||||
"00:00:01",
|
||||
"-vframes",
|
||||
"1",
|
||||
"-vf",
|
||||
f"scale={thumbnail_size[0]}:{thumbnail_size[1]}"
|
||||
":force_original_aspect_ratio=decrease",
|
||||
"-q:v",
|
||||
_ffmpeg_qscale(encoding),
|
||||
"-y",
|
||||
str(tmp_path),
|
||||
]
|
||||
try:
|
||||
process = subprocess.Popen( # nosec B603 B607 - fixed argv
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
)
|
||||
if track is not None:
|
||||
track(process)
|
||||
try:
|
||||
_out, stderr = process.communicate(timeout=30)
|
||||
except subprocess.TimeoutExpired:
|
||||
process.kill()
|
||||
process.communicate()
|
||||
return False
|
||||
if process.returncode == 0 and tmp_path.exists():
|
||||
_replace_unless_held(tmp_path, thumbnail_path, guard)
|
||||
self.logger.debug(
|
||||
f"Generated video thumbnail using ffmpeg: {thumbnail_path}"
|
||||
)
|
||||
return True
|
||||
self.logger.warning(f"ffmpeg failed for {video_path}: {stderr}")
|
||||
return False
|
||||
except (FileNotFoundError, OSError) as e:
|
||||
if isinstance(e, PermissionError) and tmp_path.exists():
|
||||
raise
|
||||
return False
|
||||
finally:
|
||||
try:
|
||||
tmp_path.unlink()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
async def clear_thumbnails(self, request: web.Request) -> web.Response:
|
||||
"""Clear generated thumbnails across gallery roots, preserving originals.
|
||||
|
||||
@@ -2340,7 +2218,7 @@ class ImageRoutesMixin:
|
||||
"""
|
||||
return await self._run_held(
|
||||
THUMBNAILS_JOB_KIND,
|
||||
THUMBNAIL_CLEAR_TITLE,
|
||||
"Thumbnail clear (gallery page)",
|
||||
"guarded",
|
||||
self._clear_thumbnails,
|
||||
)
|
||||
|
||||
+5
-22
@@ -19,11 +19,7 @@ ones. One job of a kind runs at a time, and each job has an admission class
|
||||
- ``guarded`` must not overlap record changes but may overlap other guarded
|
||||
work (thumbnail generation and clearing, later backup). It is refused while
|
||||
a ``records`` job runs.
|
||||
- ``read`` only reads; it is refused only by a running job of its own kind
|
||||
or an ``exclusive`` one.
|
||||
- ``exclusive`` runs alone (Optimize's VACUUM, which rewrites the whole
|
||||
database under the write lock): it is refused while any job runs, and every
|
||||
job is refused while it runs.
|
||||
- ``read`` only reads; it is refused only by a running job of its own kind.
|
||||
|
||||
A short legacy-page operation takes a ``hold``: a running, non-cancellable
|
||||
job that lasts as long as the request, so the Activity page lists it and a
|
||||
@@ -79,14 +75,12 @@ IMAGE_TAG_GROUP_NAMES = ("tagged", "unchanged", "skipped", "errors")
|
||||
# Thumbnail generation (#249 System S5): previews written, already fresh, failed.
|
||||
THUMBNAIL_GROUP_NAMES = ("generated", "skipped", "errors")
|
||||
|
||||
ADMISSIONS = ("read", "guarded", "records", "exclusive")
|
||||
ADMISSIONS = ("read", "guarded", "records")
|
||||
# What each admission class is refused beside (besides its own kind).
|
||||
# ``exclusive`` (Optimize's VACUUM) runs alone: beside no job of any class.
|
||||
_REFUSED_BESIDE = {
|
||||
"read": frozenset({"exclusive"}),
|
||||
"guarded": frozenset({"records", "exclusive"}),
|
||||
"records": frozenset({"records", "guarded", "exclusive"}),
|
||||
"exclusive": frozenset(ADMISSIONS),
|
||||
"read": frozenset(),
|
||||
"guarded": frozenset({"records"}),
|
||||
"records": frozenset({"records", "guarded"}),
|
||||
}
|
||||
|
||||
# Kinds of the legacy-page operations that hold the registry (D80; titles in
|
||||
@@ -97,11 +91,6 @@ DUPLICATE_CLEANUP_JOB_KIND = "duplicate_cleanup"
|
||||
THUMBNAILS_JOB_KIND = "thumbnails"
|
||||
LORA_IMPORT_JOB_KIND = "lora_import"
|
||||
AUTOTAG_BATCH_JOB_KIND = "autotag"
|
||||
# Duplicate image files (#249 System S6, py/api/duplicate_files.py): the
|
||||
# find is a cancellable ``read`` job; the reviewed removal is ``records``.
|
||||
DUPLICATES_JOB_KIND = "duplicates"
|
||||
DUPLICATE_GROUP_NAMES = ("groups_found",)
|
||||
DUPLICATE_CLEANUP_GROUP_NAMES = ("removed", "kept", "errors")
|
||||
# The System health check (py/api/system.py): a read-only ``read`` job with
|
||||
# no file outcomes, so it may run beside a scan.
|
||||
HEALTH_JOB_KIND = "health"
|
||||
@@ -515,12 +504,6 @@ class JobRegistry:
|
||||
def get(self, job_id):
|
||||
return self._jobs.get(job_id)
|
||||
|
||||
def running(self):
|
||||
"""The running jobs, oldest first (a snapshot, safe off the event loop)."""
|
||||
jobs = [job for job in tuple(self._jobs.values()) if job.running]
|
||||
jobs.sort(key=lambda job: job._started_seq)
|
||||
return jobs
|
||||
|
||||
def active(self, kind):
|
||||
"""The running job of ``kind``, or None (safe off the event loop)."""
|
||||
for job in tuple(self._jobs.values()):
|
||||
|
||||
@@ -1,420 +0,0 @@
|
||||
"""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)
|
||||
@@ -1,815 +0,0 @@
|
||||
"""Maintenance routes: reviewed, read-only plans (#249 System S6).
|
||||
|
||||
``POST /prompt_manager/maintenance/preview {operation, options:{}}`` reviews
|
||||
one operation (``duplicate_prompts``, ``optimize``, ``unused_components`` or
|
||||
``missing_links``; ``database/maintenance.py``) off the event loop and
|
||||
answers a ``Plan``: ``{success, token, operation, expires_at, scope:{count,
|
||||
label}, protected:{count, reasons}, backup:{policy, destination_label},
|
||||
media:{deletes_files, label}, details:[≤50 {label, detail}]}``. The preview
|
||||
writes nothing. The plan's fingerprint, its targets and the sources it judged
|
||||
stay in the server's ``PlanStore`` under the token, bound to the database the
|
||||
preview read (a restore that replaces it invalidates the plan).
|
||||
|
||||
The operation enum is strict, ``options`` must be an empty object and the
|
||||
body is capped at ``PREVIEW_BODY_MAX`` bytes. Errors are ``{success:false,
|
||||
code, error}`` with ``invalid-request`` (400), ``too-large`` (413: the body,
|
||||
or a plan over the store's budget), ``storage-changed`` (409: the database
|
||||
was replaced while the preview read it) or a null code (500); none carries a
|
||||
server path or exception text.
|
||||
|
||||
``POST /prompt_manager/maintenance/apply {token}`` applies a reviewed plan as
|
||||
an Activity job (202 ``{job}``): kind ``maintenance``, titled after the
|
||||
operation, admission ``records`` (refused while a scan, restore, LoRA import
|
||||
or other record-changing job runs, and refusing them meanwhile), not
|
||||
cancellable, with outcome groups ``changed``, ``protected`` and ``errors``.
|
||||
The token is consumed when the job starts (on the same loop turn, before its
|
||||
runner runs); a refused start keeps it. The runner:
|
||||
|
||||
1. takes a safety backup when the plan was reviewed with the policy on or the
|
||||
policy (``maintenance.safety_backup``, default on) is on now: SQLite's
|
||||
backup API into a temp file beside the database, ``verify_database_file``,
|
||||
fsync, ``0o600`` on POSIX, then ``os.replace`` onto a unique
|
||||
``<db>.backup_<YYYYmmdd_HHMMSS_ffffff>[_n]`` name. Any failure removes
|
||||
what it wrote and ends the job ``backup-failed`` with nothing changed;
|
||||
2. applies the plan, which first rechecks the reviewed fingerprint inside
|
||||
``BEGIN IMMEDIATE`` (``plan-changed`` writes nothing);
|
||||
3. counts the changed and protected records, and prunes old safety backups.
|
||||
|
||||
The job log names the backup by its file name only, and the job says "No
|
||||
safety backup" when none was taken. Apply errors: ``invalid-request`` (400),
|
||||
``plan-expired`` (410: unknown, expired or forgotten after a restart),
|
||||
``plan-used`` (409, with the job that applied it), ``storage-changed`` (409:
|
||||
a restore replaced the database since the review) and ``busy`` (409).
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sqlite3
|
||||
import tempfile
|
||||
from contextlib import closing
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
from .admin import _prune_safety_backups, _scan_location, _scan_roots
|
||||
from .job_routes import busy_response
|
||||
from .jobs import MAINTENANCE_JOB_KIND, JobBusy
|
||||
from .maintenance_plans import PlanExpired, PlanUsed, ScopeTooLarge
|
||||
|
||||
try:
|
||||
from ...database.workflow_storage import PRE_MIGRATION_LABEL, pre_migration_backups
|
||||
from ...database.health import readonly_uri
|
||||
from ...database.maintenance import MAINTENANCE_OPERATIONS, PlanChanged
|
||||
except ImportError:
|
||||
from database.workflow_storage import PRE_MIGRATION_LABEL, pre_migration_backups
|
||||
from database.health import readonly_uri
|
||||
from database.maintenance import MAINTENANCE_OPERATIONS, PlanChanged
|
||||
|
||||
PREVIEW_BODY_MAX = 4096
|
||||
# Duplicate-file previews carry the reviewed keepers (py/api/duplicate_files.py).
|
||||
DUPLICATE_PREVIEW_BODY_MAX = 1024 * 1024
|
||||
DUPLICATE_FILES = "duplicate_files"
|
||||
WORKFLOW_STORAGE = "workflow_storage" # applied by py/api/workflow_storage.py
|
||||
PREVIEW_OPERATIONS = MAINTENANCE_OPERATIONS + (DUPLICATE_FILES,)
|
||||
# Room left over beyond what a step writes, so a disk is never filled to 0.
|
||||
DISK_MARGIN_BYTES = 16 * 2**20
|
||||
# The concept's Maintenance row titles, which name each operation's job.
|
||||
MAINTENANCE_TITLES = {
|
||||
"duplicate_prompts": "Duplicate prompt text",
|
||||
"optimize": "Optimize database",
|
||||
"unused_components": "Unused prompt components",
|
||||
"missing_links": "Missing-file links",
|
||||
"workflow_storage": "Store each workflow once",
|
||||
}
|
||||
_OPERATION_ERROR = "operation must be one of: " + ", ".join(PREVIEW_OPERATIONS)
|
||||
MAINTENANCE_GROUPS = ("changed", "protected", "errors")
|
||||
MAX_TOKEN_LENGTH = 128
|
||||
BACKUP_FAILED = (
|
||||
"The safety backup could not be taken and verified; nothing was changed."
|
||||
)
|
||||
PLAN_CHANGED = (
|
||||
"The library changed since the review; nothing was changed. "
|
||||
"Review the plan again."
|
||||
)
|
||||
APPLY_FAILED = "Maintenance failed; nothing was changed."
|
||||
NOTHING_TO_DO = "Every image record already stores its workflow once; nothing to do."
|
||||
NO_BACKUP = "No safety backup"
|
||||
_NAME_ATTEMPTS = 1000
|
||||
|
||||
|
||||
class SafetyBackupFailed(Exception):
|
||||
"""The safety backup could not be written, verified or put in place.
|
||||
|
||||
``reason`` is the verify reason when the copy did not verify.
|
||||
"""
|
||||
|
||||
def __init__(self, message, reason=None):
|
||||
super().__init__(message)
|
||||
self.reason = reason
|
||||
|
||||
|
||||
def backup_failed_message(error):
|
||||
"""The job line for a failed backup, naming a damaged workflow document."""
|
||||
reason = getattr(error, "reason", None) or ""
|
||||
if reason.startswith("Workflow document "):
|
||||
return f"{BACKUP_FAILED} {reason}"
|
||||
return BACKUP_FAILED
|
||||
|
||||
|
||||
def _now():
|
||||
return datetime.datetime.now()
|
||||
|
||||
|
||||
def _free_bytes(path):
|
||||
return shutil.disk_usage(path).free
|
||||
|
||||
|
||||
def _device(path):
|
||||
try:
|
||||
return os.stat(path).st_dev
|
||||
except OSError:
|
||||
return path
|
||||
|
||||
|
||||
def sqlite_temp_dir(environ=None, *, windows=None, usable=None):
|
||||
"""The folder SQLite builds its temporary files in (VACUUM's copy).
|
||||
|
||||
Mirrors SQLite's own search, which differs from ``tempfile``'s. On Unix
|
||||
(``unixTempFileDir``): ``SQLITE_TMPDIR``, ``TMPDIR``, then ``/var/tmp``,
|
||||
``/usr/tmp`` and ``/tmp``, the first that is a writable directory, else
|
||||
the current directory. On Windows (``GetTempPathW``): ``TMP``, ``TEMP``,
|
||||
then ``USERPROFILE``, else the current directory.
|
||||
"""
|
||||
environ = os.environ if environ is None else environ
|
||||
if windows is None:
|
||||
windows = os.name == "nt"
|
||||
if windows:
|
||||
for name in ("TMP", "TEMP", "USERPROFILE"):
|
||||
if environ.get(name):
|
||||
return environ[name]
|
||||
return os.path.abspath(".")
|
||||
if usable is None:
|
||||
|
||||
def usable(path):
|
||||
return os.path.isdir(path) and os.access(path, os.W_OK | os.X_OK)
|
||||
|
||||
candidates = [environ.get("SQLITE_TMPDIR"), environ.get("TMPDIR")]
|
||||
candidates += ["/var/tmp", "/usr/tmp", "/tmp"]
|
||||
for candidate in candidates:
|
||||
if candidate and usable(candidate):
|
||||
return candidate
|
||||
return os.path.abspath(".")
|
||||
|
||||
|
||||
def database_bytes(db_path):
|
||||
"""The size a safety backup of the database will be (blocking).
|
||||
|
||||
SQLite's backup API copies every page, committed WAL content included, so
|
||||
the copy is ``page_count * page_size`` read on one read-only snapshot.
|
||||
"""
|
||||
conn = sqlite3.connect(readonly_uri(db_path), uri=True)
|
||||
try:
|
||||
page_size = conn.execute("PRAGMA page_size").fetchone()[0]
|
||||
return page_size * conn.execute("PRAGMA page_count").fetchone()[0]
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _megabytes(size):
|
||||
return f"{size / 2**20:,.1f} MB"
|
||||
|
||||
|
||||
def disk_plan(db_path, *, backup, vacuum, vacuum_bytes=None):
|
||||
"""Whether the disk has room for an apply's writes (blocking).
|
||||
|
||||
``backup``: the safety backup, one database-sized file beside the
|
||||
database. ``vacuum``: VACUUM writes the rebuilt database through the WAL
|
||||
beside it and builds it in a temporary copy in SQLite's temporary folder
|
||||
(``sqlite_temp_dir``; SQLite's documentation: up to twice the database
|
||||
in free space). Needs
|
||||
on one device add up. ``vacuum_bytes``, when given, is what VACUUM
|
||||
copies (the live pages, for Reclaim space) instead of the whole file.
|
||||
Returns ``(public, short)``: ``public`` is
|
||||
``{free_bytes, required_bytes, sufficient}`` for the database folder,
|
||||
``short`` None or ``(place, what, required, free)`` of the first place
|
||||
that is short, and ``size`` the database's size (0 when nothing is
|
||||
written).
|
||||
"""
|
||||
db_path = os.path.abspath(db_path)
|
||||
folder = os.path.dirname(db_path)
|
||||
size = database_bytes(db_path) if backup or vacuum else 0
|
||||
needs = [] # (place, folder, bytes, what)
|
||||
if backup:
|
||||
needs.append(("database folder", folder, size, "the safety backup"))
|
||||
if vacuum:
|
||||
copy = size if vacuum_bytes is None else vacuum_bytes
|
||||
what = "Optimize" if vacuum_bytes is None else "Reclaim space"
|
||||
needs.append(("database folder", folder, copy, what))
|
||||
needs.append(("temporary folder", sqlite_temp_dir(), copy, what))
|
||||
devices = {}
|
||||
for place, path, needed, what in needs:
|
||||
entry = devices.setdefault(_device(path), {"place": place, "path": path})
|
||||
entry["bytes"] = entry.get("bytes", 0) + needed
|
||||
entry.setdefault("what", [])
|
||||
if what not in entry["what"]:
|
||||
entry["what"].append(what)
|
||||
short = None
|
||||
public = {"free_bytes": _free_bytes(folder), "required_bytes": 0}
|
||||
for device, entry in devices.items():
|
||||
required = entry["bytes"] + DISK_MARGIN_BYTES
|
||||
free = (
|
||||
public["free_bytes"]
|
||||
if entry["path"] == folder
|
||||
else _free_bytes(entry["path"])
|
||||
)
|
||||
if device == _device(folder):
|
||||
public["required_bytes"] = required
|
||||
if free < required and short is None:
|
||||
short = (entry["place"], " and ".join(entry["what"]), required, free)
|
||||
public["sufficient"] = short is None
|
||||
return public, short, size
|
||||
|
||||
|
||||
# A restore upload's hard ceiling; below it, free disk space decides.
|
||||
RESTORE_MAX_BYTES = 64 * 2**30
|
||||
|
||||
|
||||
def size_label(size):
|
||||
"""``1.6 GB`` or ``120.0 MB``: a backup's size for a message."""
|
||||
if size >= 2**30:
|
||||
return f"{size / 2**30:,.1f} GB"
|
||||
return _megabytes(size)
|
||||
|
||||
|
||||
def restore_room(db_path):
|
||||
"""The largest restore upload the disks have room for (blocking).
|
||||
|
||||
The upload is written to a temp file and staged once more in the
|
||||
temporary folder (2x), then copied over the live database beside it,
|
||||
whose current contents are first kept as a safety backup (1x the upload
|
||||
for the WAL, plus the current database). Needs on one device add up,
|
||||
each with ``DISK_MARGIN_BYTES``.
|
||||
"""
|
||||
db_path = os.path.abspath(db_path)
|
||||
folder = os.path.dirname(db_path)
|
||||
current = database_bytes(db_path) if os.path.exists(db_path) else 0
|
||||
devices = {}
|
||||
for path, per_upload, fixed in (
|
||||
(tempfile.gettempdir(), 2, 0),
|
||||
(folder, 1, current),
|
||||
):
|
||||
entry = devices.setdefault(_device(path), {"path": path, "per": 0, "fixed": 0})
|
||||
entry["per"] += per_upload
|
||||
entry["fixed"] += fixed
|
||||
room = min(
|
||||
(_free_bytes(e["path"]) - e["fixed"] - DISK_MARGIN_BYTES) // e["per"]
|
||||
for e in devices.values()
|
||||
)
|
||||
return max(0, room)
|
||||
|
||||
|
||||
def restore_space_message(size, room, *, known):
|
||||
"""The truthful refusal of a restore the disks have no room for."""
|
||||
if known:
|
||||
return (
|
||||
f"Not enough free disk space to restore a {size_label(size)} backup; "
|
||||
f"there is room for about {size_label(room)}. Nothing was changed."
|
||||
)
|
||||
return (
|
||||
"Not enough free disk space to restore this backup: it is larger than "
|
||||
f"the {size_label(room)} there is room for. Nothing was changed."
|
||||
)
|
||||
|
||||
|
||||
def short_space_message(short):
|
||||
"""The truthful refusal for ``disk_plan``'s ``short``."""
|
||||
place, what, required, free = short
|
||||
verb = "need" if " and " in what else "needs"
|
||||
return (
|
||||
f"Not enough free disk space: {what} {verb} about {_megabytes(required)} "
|
||||
f"in the {place} and {_megabytes(free)} is free. Nothing was changed."
|
||||
)
|
||||
|
||||
|
||||
def _changed_line(operation, count):
|
||||
if operation == "optimize":
|
||||
return "Optimized the database"
|
||||
if operation == "duplicate_prompts":
|
||||
noun = "duplicate prompt" if count == 1 else "duplicate prompts"
|
||||
return f"Merged {count:,} {noun}"
|
||||
if operation == "unused_components":
|
||||
noun = "unused prompt component" if count == 1 else "unused prompt components"
|
||||
return f"Removed {count:,} {noun}"
|
||||
noun = "link to a missing file" if count == 1 else "links to missing files"
|
||||
return f"Removed {count:,} {noun}"
|
||||
|
||||
|
||||
def _remove_quietly(*paths):
|
||||
for path in paths:
|
||||
try:
|
||||
if path and os.path.exists(path):
|
||||
os.remove(path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _reserve_backup_name(db_path, label="backup"):
|
||||
"""Create and return an unused ``<db>.<label>_<stamp>[_n]`` (O_EXCL)."""
|
||||
for candidate in _backup_names(db_path, label):
|
||||
try:
|
||||
fd = os.open(candidate, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
|
||||
except FileExistsError:
|
||||
continue
|
||||
os.close(fd)
|
||||
return candidate
|
||||
raise SafetyBackupFailed("No free backup name")
|
||||
|
||||
|
||||
def _backup_names(db_path, label="backup"):
|
||||
base = f"{db_path}.{label}_{_now().strftime('%Y%m%d_%H%M%S_%f')}"
|
||||
for attempt in range(_NAME_ATTEMPTS):
|
||||
yield base if attempt == 0 else f"{base}_{attempt}"
|
||||
|
||||
|
||||
def _place_backup(temp, db_path, label="backup"):
|
||||
"""Give the verified ``temp`` its unique backup name; return that path.
|
||||
|
||||
``os.link`` is atomic and refuses an existing name, so no empty or half
|
||||
file ever carries a backup name, even across a crash. A filesystem
|
||||
without hard links falls back to reserving the name (``O_EXCL``) and
|
||||
``os.replace``. The temp file is removed by the caller.
|
||||
"""
|
||||
for candidate in _backup_names(db_path, label):
|
||||
try:
|
||||
os.link(temp, candidate)
|
||||
except FileExistsError:
|
||||
continue
|
||||
except (OSError, NotImplementedError):
|
||||
break # no hard links here
|
||||
else:
|
||||
return candidate
|
||||
else:
|
||||
raise SafetyBackupFailed("No free backup name")
|
||||
final = _reserve_backup_name(db_path, label)
|
||||
try:
|
||||
os.replace(temp, final)
|
||||
except OSError:
|
||||
_remove_quietly(final)
|
||||
raise
|
||||
return final
|
||||
|
||||
|
||||
SAFETY_TEMP_PREFIX = ".pm-safety-"
|
||||
|
||||
|
||||
def sweep_safety_temps(db_path):
|
||||
"""Remove the temp copies an interrupted safety backup left (blocking).
|
||||
|
||||
Called once at startup, before any job can take a backup. Returns the
|
||||
names removed.
|
||||
"""
|
||||
folder = os.path.dirname(os.path.abspath(db_path))
|
||||
removed = []
|
||||
try:
|
||||
names = os.listdir(folder)
|
||||
except OSError:
|
||||
return removed
|
||||
for name in names:
|
||||
if not (name.startswith(SAFETY_TEMP_PREFIX) and ".tmp" in name):
|
||||
continue
|
||||
try:
|
||||
os.remove(os.path.join(folder, name))
|
||||
removed.append(name)
|
||||
except OSError:
|
||||
continue
|
||||
return removed
|
||||
|
||||
|
||||
def _standalone(path):
|
||||
"""Make the copy one self-contained file (rollback journal, not WAL).
|
||||
|
||||
The backup API copies the live database's WAL flag; a WAL file opened
|
||||
read-only later (a verify, a restore preview) would leave ``-wal`` and
|
||||
``-shm`` sidecars beside it.
|
||||
"""
|
||||
try:
|
||||
with closing(sqlite3.connect(path)) as conn:
|
||||
conn.execute("PRAGMA journal_mode=DELETE").fetchone()
|
||||
except sqlite3.Error as exc:
|
||||
raise SafetyBackupFailed("The backup could not be finished") from exc
|
||||
|
||||
|
||||
def _fsync(path):
|
||||
with open(path, "rb+") as handle:
|
||||
os.fsync(handle.fileno())
|
||||
|
||||
|
||||
def take_safety_backup(model, label="backup"):
|
||||
"""Write and verify a safety backup; return its file name (blocking).
|
||||
|
||||
The copy is made in a temp file beside the database and only renamed
|
||||
onto its final, unique name ``<db>.<label>_<stamp>`` once it verified,
|
||||
so a backup file on disk is always a verified one. ``label`` "backup"
|
||||
joins the pruned pool; another label (the workflow storage migration's
|
||||
``pre-workflow-storage``) is kept until the user removes it. Raises
|
||||
SafetyBackupFailed, leaving no file.
|
||||
"""
|
||||
db_path = os.path.abspath(model.db_path)
|
||||
fd, temp = tempfile.mkstemp(
|
||||
dir=os.path.dirname(db_path), prefix=SAFETY_TEMP_PREFIX, suffix=".tmp"
|
||||
)
|
||||
os.close(fd)
|
||||
final = None
|
||||
try:
|
||||
if not model.backup_database(temp):
|
||||
raise SafetyBackupFailed("The backup could not be written")
|
||||
_standalone(temp)
|
||||
ok, reason = model.verify_database_file(temp)
|
||||
if not ok:
|
||||
raise SafetyBackupFailed("The backup did not verify", reason)
|
||||
_fsync(temp)
|
||||
if os.name != "nt":
|
||||
os.chmod(temp, 0o600)
|
||||
final = _place_backup(temp, db_path, label)
|
||||
return os.path.basename(final)
|
||||
except OSError as exc:
|
||||
_remove_quietly(final)
|
||||
raise SafetyBackupFailed("The backup could not be put in place") from exc
|
||||
except SafetyBackupFailed:
|
||||
_remove_quietly(final)
|
||||
raise
|
||||
finally:
|
||||
_remove_quietly(temp, temp + "-journal", temp + "-wal", temp + "-shm")
|
||||
|
||||
|
||||
def _failure(message, status, code=None, **extra):
|
||||
return web.json_response(
|
||||
{"success": False, "code": code, "error": message, **extra}, status=status
|
||||
)
|
||||
|
||||
|
||||
def _invalid(message):
|
||||
return _failure(message, 400, "invalid-request")
|
||||
|
||||
|
||||
def _iso(timestamp):
|
||||
return datetime.datetime.fromtimestamp(timestamp, datetime.timezone.utc).isoformat()
|
||||
|
||||
|
||||
class MaintenanceRoutesMixin:
|
||||
"""Mixin providing the /prompt_manager/maintenance plan endpoints."""
|
||||
|
||||
def _register_maintenance_routes(self, routes):
|
||||
@routes.post("/prompt_manager/maintenance/preview")
|
||||
async def maintenance_preview_route(request):
|
||||
return await self.preview_maintenance(request)
|
||||
|
||||
@routes.post("/prompt_manager/maintenance/apply")
|
||||
async def maintenance_apply_route(request):
|
||||
return await self.apply_maintenance_plan(request)
|
||||
|
||||
def _database_identity(self, db):
|
||||
"""What binds a plan to one database: the swap generation and its file."""
|
||||
path = os.path.normcase(os.path.realpath(db.model.db_path))
|
||||
return (self._db_generation, path)
|
||||
|
||||
@staticmethod
|
||||
def _safety_backup_policy():
|
||||
"""``on`` or ``off``: whether an apply takes a safety backup first."""
|
||||
from ..config import PromptManagerConfig
|
||||
|
||||
return "on" if PromptManagerConfig.SAFETY_BACKUP else "off"
|
||||
|
||||
async def _read_maintenance_body(self, request, allowed, cap=None):
|
||||
"""``(data, None)`` for a capped JSON object with only ``allowed`` keys.
|
||||
|
||||
``cap(data)``, when given, is the body's byte limit for what it asks
|
||||
(read up to ``DUPLICATE_PREVIEW_BODY_MAX``); else PREVIEW_BODY_MAX.
|
||||
"""
|
||||
limit = PREVIEW_BODY_MAX if cap is None else DUPLICATE_PREVIEW_BODY_MAX
|
||||
data, error = await self._read_json_body(request, max_bytes=limit)
|
||||
if error is not None:
|
||||
if error.status == 413:
|
||||
return None, _failure("Request body too large", 413, "too-large")
|
||||
return None, _invalid("Request body must be a JSON object")
|
||||
if cap is not None:
|
||||
size = len(json.dumps(data, separators=(",", ":")).encode("utf-8"))
|
||||
if size > cap(data):
|
||||
return None, _failure("Request body too large", 413, "too-large")
|
||||
unknown = sorted(set(data) - set(allowed))
|
||||
if unknown:
|
||||
return None, _invalid("Unknown field: " + ", ".join(unknown)[:200])
|
||||
return data, None
|
||||
|
||||
async def preview_maintenance(self, request):
|
||||
data, error = await self._read_maintenance_body(
|
||||
request,
|
||||
("operation", "options"),
|
||||
cap=lambda body: (
|
||||
DUPLICATE_PREVIEW_BODY_MAX
|
||||
if body.get("operation") == DUPLICATE_FILES
|
||||
else PREVIEW_BODY_MAX
|
||||
),
|
||||
)
|
||||
if error is not None:
|
||||
return error
|
||||
operation = data.get("operation")
|
||||
options = data.get("options", {})
|
||||
if not isinstance(operation, str) or operation not in PREVIEW_OPERATIONS:
|
||||
return _invalid(_OPERATION_ERROR)
|
||||
if operation == DUPLICATE_FILES:
|
||||
return await self.preview_duplicate_files(options)
|
||||
if not isinstance(options, dict) or options:
|
||||
return _invalid(f"{operation} takes no options")
|
||||
db = self.db
|
||||
identity = self._database_identity(db)
|
||||
try:
|
||||
roots = (
|
||||
await self._run_in_executor(self._health_roots)
|
||||
if operation == "missing_links"
|
||||
else []
|
||||
)
|
||||
draft = await self._run_in_executor(
|
||||
db.preview_maintenance, operation, roots
|
||||
)
|
||||
except Exception:
|
||||
self.logger.exception(f"Maintenance preview of {operation} failed")
|
||||
return _failure("Failed to review maintenance", 500)
|
||||
return await self._plan_response(
|
||||
db,
|
||||
identity,
|
||||
draft.operation,
|
||||
draft.fingerprint,
|
||||
draft.public(),
|
||||
{
|
||||
"targets": draft.targets,
|
||||
"roots": roots,
|
||||
"protected": draft.protected["count"],
|
||||
"rows": draft.scope["count"],
|
||||
},
|
||||
)
|
||||
|
||||
async def _plan_response(
|
||||
self, db, identity, operation, fingerprint, public, payload
|
||||
):
|
||||
"""Issue the reviewed plan's token and answer the ``Plan``.
|
||||
|
||||
Refused with ``storage-changed`` when the database was replaced while
|
||||
it was reviewed. The plan reports the safety backup's expected size
|
||||
and the disk space the apply needs (``disk_plan``).
|
||||
"""
|
||||
if identity != self._database_identity(self.db):
|
||||
return _failure(
|
||||
"The database was replaced while it was reviewed. Review again.",
|
||||
409,
|
||||
"storage-changed",
|
||||
)
|
||||
policy = self._safety_backup_policy()
|
||||
destination = "backup"
|
||||
if operation == WORKFLOW_STORAGE:
|
||||
# Always backed up first, whatever the policy (the way back);
|
||||
# one kept pre-migration backup ever, pooled ones after it.
|
||||
policy = "required"
|
||||
kept = await self._run_in_executor(pre_migration_backups, db.model.db_path)
|
||||
destination = "backup" if kept else PRE_MIGRATION_LABEL
|
||||
try:
|
||||
disk, _short, size = await self._run_in_executor(
|
||||
disk_plan,
|
||||
db.model.db_path,
|
||||
backup=policy != "off",
|
||||
vacuum=operation == "optimize",
|
||||
)
|
||||
except (OSError, sqlite3.Error):
|
||||
self.logger.exception("Could not measure the disk space for a plan")
|
||||
return _failure("Failed to review maintenance", 500)
|
||||
store = self._maintenance_plans
|
||||
try:
|
||||
token = store.issue(
|
||||
operation, fingerprint, {**payload, "backup": policy}, identity
|
||||
)
|
||||
except ScopeTooLarge:
|
||||
return _failure(
|
||||
"This plan is too large to review at once.", 413, "too-large"
|
||||
)
|
||||
plan = store.peek(token, operation)
|
||||
name = os.path.basename(db.model.db_path)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"token": token,
|
||||
"expires_at": _iso(plan.expires_at),
|
||||
**public,
|
||||
"backup": {
|
||||
"policy": policy,
|
||||
"destination_label": f"{name}.{destination}_<date and time>",
|
||||
"expected_bytes": size if policy != "off" else None,
|
||||
},
|
||||
"disk": disk,
|
||||
},
|
||||
headers={"Cache-Control": "no-store"},
|
||||
)
|
||||
|
||||
# ── apply ─────────────────────────────────────────────────────────
|
||||
|
||||
async def apply_maintenance_plan(self, request):
|
||||
data, error = await self._read_maintenance_body(request, ("token",))
|
||||
if error is not None:
|
||||
return error
|
||||
token = data.get("token")
|
||||
if not isinstance(token, str) or not 0 < len(token) <= MAX_TOKEN_LENGTH:
|
||||
return _invalid("token must be the plan's token")
|
||||
plan, refused = self._peek_plan(token)
|
||||
if refused is not None:
|
||||
return refused
|
||||
store = self._maintenance_plans
|
||||
db = self.db
|
||||
if plan.db_identity != self._database_identity(db):
|
||||
store.discard(token)
|
||||
return _failure(
|
||||
"The database was replaced since the review. Review again.",
|
||||
409,
|
||||
"storage-changed",
|
||||
)
|
||||
if plan.operation == WORKFLOW_STORAGE and not plan.payload.get("rows"):
|
||||
store.discard(token) # no job, no backup: nothing to rewrite
|
||||
return _failure(NOTHING_TO_DO, 409, "nothing-to-do")
|
||||
backup = self._backup_promised(plan)
|
||||
try:
|
||||
_disk, short, _size = await self._run_in_executor(
|
||||
disk_plan,
|
||||
db.model.db_path,
|
||||
backup=backup,
|
||||
vacuum=plan.operation == "optimize",
|
||||
)
|
||||
except (OSError, sqlite3.Error):
|
||||
self.logger.exception("Could not measure the disk space for an apply")
|
||||
return _failure(APPLY_FAILED, 500)
|
||||
if short is not None:
|
||||
# Refused before any job starts: nothing fails halfway for space.
|
||||
return _failure(short_space_message(short), 507, "insufficient-space")
|
||||
# The token may have been used, expired or dropped while the disk was
|
||||
# measured: look again with no await before the job starts.
|
||||
plan, refused = self._peek_plan(token)
|
||||
if refused is not None:
|
||||
return refused
|
||||
try:
|
||||
if plan.operation == DUPLICATE_FILES:
|
||||
job = self._start_duplicate_cleanup(db, plan, backup)
|
||||
elif plan.operation == WORKFLOW_STORAGE:
|
||||
job = self._start_workflow_storage(db, backup)
|
||||
else:
|
||||
job = self._jobs.start(
|
||||
MAINTENANCE_JOB_KIND,
|
||||
MAINTENANCE_TITLES[plan.operation],
|
||||
# VACUUM rewrites the whole file under the write lock: it
|
||||
# runs alone, beside no other job of any class.
|
||||
admission=(
|
||||
"exclusive" if plan.operation == "optimize" else "records"
|
||||
),
|
||||
cancellable=False,
|
||||
scope={"operation": plan.operation},
|
||||
groups=MAINTENANCE_GROUPS,
|
||||
runner=lambda job: self._run_maintenance_plan(
|
||||
job, db, plan, backup
|
||||
),
|
||||
)
|
||||
except JobBusy as busy:
|
||||
return busy_response(busy) # the token stays usable
|
||||
# Same loop turn as the peek and the start: the runner has not run.
|
||||
store.consume(token, plan.operation, job_id=job.id)
|
||||
return web.json_response({"success": True, "job": job.summary()}, status=202)
|
||||
|
||||
def _peek_plan(self, token):
|
||||
"""``(plan, None)`` for a usable token, else ``(None, refusal)``."""
|
||||
try:
|
||||
return self._maintenance_plans.peek(token), None
|
||||
except PlanUsed as used:
|
||||
job = self._jobs.get(used.job_id)
|
||||
return None, _failure(
|
||||
"This plan was already applied.",
|
||||
409,
|
||||
"plan-used",
|
||||
job={"id": used.job_id, "title": job.title if job else None},
|
||||
)
|
||||
except PlanExpired:
|
||||
return None, _failure(
|
||||
"This review expired. Review the plan again.", 410, "plan-expired"
|
||||
)
|
||||
|
||||
def _backup_promised(self, plan):
|
||||
"""A plan reviewed with the policy on keeps that promise; so does on now.
|
||||
|
||||
Decided once, when the apply starts (the disk check measured it), and
|
||||
handed to the job.
|
||||
"""
|
||||
if plan.operation == WORKFLOW_STORAGE:
|
||||
return True # never without its pre-migration backup
|
||||
return plan.payload["backup"] == "on" or self._safety_backup_policy() == "on"
|
||||
|
||||
async def _safety_backup_step(self, job, db, backup_on):
|
||||
"""``(ok, name)``: take the promised safety backup, or report none.
|
||||
|
||||
``name`` is the backup's file name, or None when no backup was
|
||||
promised. A failed backup ends the job ``backup-failed`` (``ok``
|
||||
False).
|
||||
"""
|
||||
if not backup_on:
|
||||
job.log_line("No safety backup (the policy is off)")
|
||||
return True, None
|
||||
job.update(status_line="Taking a safety backup")
|
||||
job.log_line("Taking a safety backup of the database")
|
||||
try:
|
||||
backup = await self._run_in_executor(take_safety_backup, db.model)
|
||||
except SafetyBackupFailed as error:
|
||||
self.logger.error("Safety backup failed", exc_info=True)
|
||||
_fail_job(job, "backup-failed", backup_failed_message(error))
|
||||
return False, None
|
||||
job.log_line(f"Safety backup {backup} verified")
|
||||
return True, backup
|
||||
|
||||
async def _prune_backup_pool(self, job, db):
|
||||
"""Keep the one ``<db>.backup_*`` pool (restore's too) at its cap."""
|
||||
removed = await self._run_in_executor(_prune_safety_backups, db.model.db_path)
|
||||
if removed:
|
||||
job.log_line(f"Removed {len(removed)} older safety backups")
|
||||
|
||||
async def _run_maintenance_plan(self, job, db, plan, backup_on):
|
||||
ok, backup = await self._safety_backup_step(job, db, backup_on)
|
||||
if not ok:
|
||||
return
|
||||
try:
|
||||
await self._apply_maintenance_plan(job, db, plan, backup)
|
||||
finally:
|
||||
if backup:
|
||||
await self._prune_backup_pool(job, db)
|
||||
|
||||
async def _apply_maintenance_plan(self, job, db, plan, backup):
|
||||
payload = plan.payload
|
||||
job.update(status_line=f"Applying: {job.title}")
|
||||
try:
|
||||
result = await self._run_in_executor(
|
||||
db.apply_maintenance,
|
||||
plan.operation,
|
||||
plan.fingerprint,
|
||||
payload["targets"],
|
||||
payload["roots"],
|
||||
)
|
||||
except PlanChanged:
|
||||
return _fail_job(job, "plan-changed", PLAN_CHANGED)
|
||||
except Exception:
|
||||
self.logger.exception(f"Maintenance {plan.operation} failed")
|
||||
return _fail_job(job, None, APPLY_FAILED)
|
||||
rows = await self._run_in_executor(
|
||||
self._maintenance_rows, plan.operation, result["rows"]
|
||||
)
|
||||
job.record_counted("changed", result["changed"], rows)
|
||||
job.record_counted("protected", payload["protected"])
|
||||
line = _changed_line(plan.operation, result["changed"])
|
||||
job.log_line(line)
|
||||
job.update(progress=100, status_line=f"{line} · {backup_note(backup)}")
|
||||
|
||||
def _maintenance_rows(self, operation, described):
|
||||
"""Outcome rows for the changed records (root-relative paths only)."""
|
||||
if operation == "missing_links":
|
||||
roots = _scan_roots(self._get_all_output_dirs())
|
||||
rows = []
|
||||
for item in described:
|
||||
path, label, index = _scan_location(item["image_path"], roots)
|
||||
rows.append(
|
||||
{
|
||||
"path": path,
|
||||
"root": label,
|
||||
"root_index": index,
|
||||
"detail": "Link removed; the file was already missing",
|
||||
}
|
||||
)
|
||||
return rows
|
||||
if operation == "duplicate_prompts":
|
||||
return [
|
||||
{
|
||||
"path": "",
|
||||
"detail": (
|
||||
f"Prompt {item['prompt_id']} merged into "
|
||||
f"prompt {item['kept_id']}"
|
||||
),
|
||||
}
|
||||
for item in described
|
||||
]
|
||||
return [
|
||||
{"path": "", "detail": f"Removed: {item.get('text') or item['prompt_id']}"}
|
||||
for item in described
|
||||
]
|
||||
|
||||
|
||||
def backup_note(backup):
|
||||
"""The status line's backup part: the file name, or that there is none."""
|
||||
return f"Safety backup {backup}" if backup else NO_BACKUP
|
||||
|
||||
|
||||
def _fail_job(job, code, message):
|
||||
"""End ``job`` failed with ``message`` (one error row) and its code."""
|
||||
job.record("errors", {"path": "", "detail": message}, f"! {message}")
|
||||
job.update(status_line=message)
|
||||
job.publish({"type": "error", "code": code, "message": message})
|
||||
job.finish("failed")
|
||||
@@ -1,138 +0,0 @@
|
||||
"""The reviewed-plan store behind maintenance previews (#249 System S6).
|
||||
|
||||
A preview issues a plan: its operation, the fingerprint of exactly what was
|
||||
reviewed, the payload an apply needs (target ids, the sources it judged) and
|
||||
the identity of the database it describes. Everything is held here, in
|
||||
server memory, under a single-use token (``secrets.token_urlsafe(32)``); the
|
||||
client only ever sees the token and the public plan. Plans expire after
|
||||
``ttl`` seconds and are gone after a restart.
|
||||
|
||||
The store keeps at most ``max_plans`` plans and ``max_bytes`` of payload
|
||||
(measured as compact JSON). A plan larger than ``max_bytes`` is refused with
|
||||
``ScopeTooLarge``: it is never cut down to fit, since an apply must do
|
||||
exactly what was reviewed. A new plan that fits makes room by dropping the
|
||||
oldest ones, which then read as expired.
|
||||
|
||||
A consumed token is remembered (the last ``MAX_USED``) with the job it
|
||||
started, so a second apply of it is told which job ran it.
|
||||
"""
|
||||
|
||||
import json
|
||||
import secrets
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
class PlanExpired(Exception):
|
||||
"""No live plan has the token: unknown, expired, dropped or another operation's."""
|
||||
|
||||
|
||||
class PlanUsed(Exception):
|
||||
"""The token was already applied; ``job_id`` is the job that applied it."""
|
||||
|
||||
def __init__(self, job_id):
|
||||
super().__init__("This plan was already applied")
|
||||
self.job_id = job_id
|
||||
|
||||
|
||||
class ScopeTooLarge(Exception):
|
||||
"""The plan's payload exceeds the store's byte budget."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class StoredPlan:
|
||||
"""One reviewed plan, as only the server sees it."""
|
||||
|
||||
operation: str
|
||||
fingerprint: str
|
||||
payload: dict
|
||||
db_identity: tuple
|
||||
expires_at: float # wall-clock seconds since the epoch, for the response
|
||||
deadline: float = field(repr=False) # on the store's monotonic clock
|
||||
size: int = field(repr=False)
|
||||
|
||||
|
||||
def _payload_size(payload):
|
||||
return len(json.dumps(payload, separators=(",", ":"), default=str))
|
||||
|
||||
|
||||
class PlanStore:
|
||||
"""Reviewed maintenance plans under single-use tokens."""
|
||||
|
||||
MAX_USED = 64 # consumed tokens remembered for plan-used answers
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
ttl=900,
|
||||
max_plans=8,
|
||||
max_bytes=16 * 2**20,
|
||||
*,
|
||||
clock=time.monotonic,
|
||||
wall_clock=time.time,
|
||||
):
|
||||
self.ttl = ttl
|
||||
self.max_plans = max_plans
|
||||
self.max_bytes = max_bytes
|
||||
self._clock = clock
|
||||
self._wall_clock = wall_clock
|
||||
self._plans = OrderedDict() # token -> StoredPlan, oldest first
|
||||
self._used = OrderedDict() # token -> job id, oldest first
|
||||
|
||||
def issue(self, operation, fingerprint, payload, db_identity) -> str:
|
||||
"""Hold a plan and return its token; ScopeTooLarge when it cannot fit."""
|
||||
size = _payload_size(payload)
|
||||
if size > self.max_bytes:
|
||||
raise ScopeTooLarge(f"The plan needs {size} bytes")
|
||||
self._drop_expired()
|
||||
while self._plans and (
|
||||
len(self._plans) >= self.max_plans
|
||||
or self._held_bytes() + size > self.max_bytes
|
||||
):
|
||||
self._plans.popitem(last=False)
|
||||
token = secrets.token_urlsafe(32)
|
||||
self._plans[token] = StoredPlan(
|
||||
operation=operation,
|
||||
fingerprint=fingerprint,
|
||||
payload=payload,
|
||||
db_identity=tuple(db_identity),
|
||||
expires_at=self._wall_clock() + self.ttl,
|
||||
deadline=self._clock() + self.ttl,
|
||||
size=size,
|
||||
)
|
||||
return token
|
||||
|
||||
def peek(self, token, operation=None) -> StoredPlan:
|
||||
"""The live plan of ``token``, left in place; raises PlanUsed or PlanExpired."""
|
||||
if not isinstance(token, str):
|
||||
raise PlanExpired("No such plan")
|
||||
if token in self._used:
|
||||
raise PlanUsed(self._used[token])
|
||||
plan = self._plans.get(token)
|
||||
if plan is None or (operation is not None and plan.operation != operation):
|
||||
raise PlanExpired("No such plan")
|
||||
if self._clock() > plan.deadline:
|
||||
del self._plans[token]
|
||||
raise PlanExpired("The plan expired")
|
||||
return plan
|
||||
|
||||
def consume(self, token, operation=None, *, job_id=None) -> StoredPlan:
|
||||
"""Take the plan of ``token`` once; later calls raise PlanUsed(job_id)."""
|
||||
plan = self.peek(token, operation)
|
||||
del self._plans[token]
|
||||
self._used[token] = job_id
|
||||
while len(self._used) > self.MAX_USED:
|
||||
self._used.popitem(last=False)
|
||||
return plan
|
||||
|
||||
def discard(self, token):
|
||||
"""Forget a plan that can no longer apply (its database was replaced)."""
|
||||
self._plans.pop(token, None)
|
||||
|
||||
def _held_bytes(self):
|
||||
return sum(plan.size for plan in self._plans.values())
|
||||
|
||||
def _drop_expired(self):
|
||||
now = self._clock()
|
||||
for token in [t for t, plan in self._plans.items() if now > plan.deadline]:
|
||||
del self._plans[token]
|
||||
+4
-11
@@ -38,17 +38,13 @@ try:
|
||||
from ...database.library_commands import LibraryCommandError
|
||||
from ...database.review_decisions import ReviewCommandError
|
||||
from ...utils.media_capture import CaptureError
|
||||
from ...utils.review_capture import (
|
||||
captured_review,
|
||||
item_profile,
|
||||
review_evidence,
|
||||
)
|
||||
from ...utils.review_capture import captured_review, review_evidence
|
||||
except ImportError:
|
||||
from database.library import display_title
|
||||
from database.library_commands import LibraryCommandError
|
||||
from database.review_decisions import ReviewCommandError
|
||||
from utils.media_capture import CaptureError
|
||||
from utils.review_capture import captured_review, item_profile, review_evidence
|
||||
from utils.review_capture import captured_review, review_evidence
|
||||
|
||||
ITEMS_URL = "/prompt_manager/review/items"
|
||||
DEFAULT_LIMIT = 40
|
||||
@@ -645,8 +641,7 @@ def _verify_and_read(store, capture, item, root, need_metadata: bool):
|
||||
"""Worker thread: ``(relative, facts, captured)``, or None off the root.
|
||||
|
||||
Formats the item's target under the capture's root (resolving the root's
|
||||
real path), re-hashes the current file, then reads the capture and
|
||||
resolves it with the profile the item's evidence was recorded with.
|
||||
real path), re-hashes the current file, then reads the capture.
|
||||
"""
|
||||
found = format_review_target(
|
||||
item["spelled_path"],
|
||||
@@ -657,9 +652,7 @@ def _verify_and_read(store, capture, item, root, need_metadata: bool):
|
||||
return None
|
||||
relative = found[1]
|
||||
facts = store.verify(root, relative)
|
||||
captured = None
|
||||
if need_metadata:
|
||||
captured = captured_review(store, capture, item_profile(item["evidence"]))
|
||||
captured = captured_review(store, capture) if need_metadata else None
|
||||
return relative, facts, captured
|
||||
|
||||
|
||||
|
||||
+31
-179
@@ -8,25 +8,21 @@ already. A file the fallback links, an exact pair the ingester stored and an
|
||||
already linked file are never queued. Neither is an ambiguous file with
|
||||
nothing to choose: fewer than two of its candidates carry a prompt
|
||||
(``utils.review_capture.reviewable``); the scan reports it unresolved ("could
|
||||
not be traced") unless the inspection profile finds a choice in it.
|
||||
|
||||
Only scans add items (maintainer ruling, 2026-10-23). The monitor's pass
|
||||
records nothing: it only removes the pending items of an output it linked.
|
||||
not be traced").
|
||||
|
||||
The queue heals itself: a pending item whose file is now linked (by the
|
||||
fallback, a later resolution or the monitor) is removed by the next scan
|
||||
that lists the file or by the monitor's next pass over it. The next scan
|
||||
also removes a pending item without a choice whose file is still ambiguous
|
||||
without one (``settle_untraced_review_items``), and an older version's
|
||||
item of bytes now untraced (``settle_stale_review_items``). Only an undecided item goes
|
||||
that lists the file or by the monitor's next pass over it, and so is a
|
||||
pending item without a choice whose file is still ambiguous without one
|
||||
(``settle_untraced_review_items``). Only an undecided item goes
|
||||
(``pending``, revision 0, no decision row); a decision is never deleted or
|
||||
changed (``settle_linked_review_items``).
|
||||
|
||||
The scan records through ``_record_scanned_review`` (best effort: a failure
|
||||
is logged and counted), which goes through
|
||||
``utils.review_capture.record_ambiguous_file`` on the API's capture store;
|
||||
the monitor calls ``record_review_file`` from its worker thread after its
|
||||
own link. Nothing here links a file.
|
||||
is logged and counted); the monitor calls ``record_review_file`` from its
|
||||
worker thread after its own link. Both go through
|
||||
``utils.review_capture.record_ambiguous_file`` on the API's capture store.
|
||||
Nothing here links a file.
|
||||
|
||||
Recording changes records, so it follows the registry's ``records``
|
||||
admission (#289). The scan records under its own ``records`` job. The
|
||||
@@ -39,7 +35,6 @@ defer the recording (#249 S5 M7): a long rebuild no longer holds Needs review
|
||||
back for its whole run.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
@@ -55,29 +50,10 @@ QUEUED_OUTCOMES = ("recorded", "limited", "duplicate")
|
||||
NEW_OUTCOMES = ("recorded", "limited")
|
||||
|
||||
try:
|
||||
from ...utils import review_capture
|
||||
from ...utils.review_capture import record_ambiguous_file
|
||||
except ImportError:
|
||||
from utils import review_capture
|
||||
from utils.review_capture import record_ambiguous_file
|
||||
|
||||
# Untraced files remembered per process so a rescan of an unchanged file does
|
||||
# not run the inspection profile again (oldest forgotten first). A plain dict
|
||||
# of 16-byte digests (the locator, and its size, mtime_ns and item version):
|
||||
# measured at 8.6 MiB with 65,536 entries (tracemalloc, CPython 3.14).
|
||||
MAX_UNTRACED_MEMORY = 65_536
|
||||
|
||||
|
||||
def _digest(text: str) -> bytes:
|
||||
data = text.encode("utf-8", "surrogatepass")
|
||||
return hashlib.blake2b(data, digest_size=16).digest()
|
||||
|
||||
|
||||
def _memory_key(locator: str) -> bytes:
|
||||
"""A 16-byte digest of ``locator`` (the untraced memory's key)."""
|
||||
return _digest(locator)
|
||||
|
||||
|
||||
# The scan log's reason for items that leave the queue with nothing to choose.
|
||||
NOTHING_TO_CHOOSE = "nothing to choose"
|
||||
|
||||
@@ -85,15 +61,10 @@ NOTHING_TO_CHOOSE = "nothing to choose"
|
||||
class ReviewRecordingMixin:
|
||||
"""Record the scan's and the monitor's ambiguous files for Needs review."""
|
||||
|
||||
def _record_scanned_review(self, media_file, roots, review=None, scanned=None):
|
||||
"""Record one ambiguous, unlinked scanned file (logged and counted).
|
||||
|
||||
``scanned`` is the scan's ``(metadata, inspection_resolution)``.
|
||||
"""
|
||||
def _record_scanned_review(self, media_file, roots, review=None):
|
||||
"""Record one ambiguous, unlinked scanned file (logged and counted)."""
|
||||
try:
|
||||
outcome = self._record_review_sync(
|
||||
media_file, roots, review, inspection=True, scanned=scanned
|
||||
)
|
||||
outcome = self._record_review_sync(media_file, roots, review)
|
||||
except Exception:
|
||||
self.logger.warning(
|
||||
f"{media_file.name} could not be added to Needs review",
|
||||
@@ -109,13 +80,9 @@ class ReviewRecordingMixin:
|
||||
if outcome in QUEUED_OUTCOMES:
|
||||
review["queued"] += 1
|
||||
|
||||
def _record_review_sync(
|
||||
self, media_file, roots, review=None, inspection=False, scanned=None
|
||||
):
|
||||
def _record_review_sync(self, media_file, roots, review=None):
|
||||
"""``record_ambiguous_file`` for a file under one of ``roots`` (blocking).
|
||||
|
||||
``inspection`` (the scan only) lets a file without a choice under the
|
||||
default profile be recorded from the inspection profile's choices.
|
||||
An ``untraced`` file (its captured bytes offer nothing to choose) also
|
||||
loses its pending items that offer nothing either.
|
||||
"""
|
||||
@@ -124,134 +91,12 @@ class ReviewRecordingMixin:
|
||||
return "skipped"
|
||||
root, index, relative, _spelled = located
|
||||
outcome = record_ambiguous_file(
|
||||
self.db,
|
||||
self._capture_store(),
|
||||
Path(root),
|
||||
index,
|
||||
relative,
|
||||
inspection=inspection,
|
||||
scanned=scanned,
|
||||
self.db, self._capture_store(), Path(root), index, relative
|
||||
)
|
||||
if outcome == "untraced":
|
||||
self._settle_untraced_reviews([media_file], roots, review)
|
||||
return outcome
|
||||
|
||||
def _untraced_scan_outcome(self, media_file, roots, canonical, review=None):
|
||||
"""The scan status of an unlinked file the default profile left
|
||||
ambiguous without a choice (blocking).
|
||||
|
||||
``ambiguous`` when the inspection profile offers a choice: the file is
|
||||
recorded for review from that profile (never linked), or ``unchanged``
|
||||
when a saved row links it already. Otherwise ``unresolved``, and its
|
||||
pending items without a choice leave the queue. The profile runs at
|
||||
most once per file state: not for bytes already queued (or decided)
|
||||
at the file's current size and ``mtime_ns`` under the version that
|
||||
its records depend on, and not again for a state it already found
|
||||
untraced in this process.
|
||||
"""
|
||||
try:
|
||||
choice = self._inspection_choice(media_file, roots, canonical)
|
||||
except Exception:
|
||||
# One bad file never stops the scan: it stays "could not be traced".
|
||||
self.logger.warning(
|
||||
f"{media_file.name} could not be inspected for Needs review",
|
||||
exc_info=True,
|
||||
)
|
||||
if review is not None:
|
||||
review["errors"] += 1
|
||||
return "unresolved"
|
||||
if choice is None:
|
||||
self._settle_untraced_reviews([media_file], roots, review)
|
||||
self._settle_stale_reviews(media_file, roots, canonical, review)
|
||||
return "unresolved"
|
||||
if self._review_file_linked(media_file, roots):
|
||||
return "unchanged"
|
||||
self._record_scanned_review(media_file, roots, review, scanned=choice)
|
||||
return "ambiguous"
|
||||
|
||||
def _inspection_choice(self, media_file, roots, canonical):
|
||||
"""``(canonical, resolution_or_None)`` when the file offers a choice.
|
||||
|
||||
The resolution is None when the file's current state is already in
|
||||
the queue (so the profile did not run). None when it offers none.
|
||||
"""
|
||||
if canonical is None:
|
||||
return None
|
||||
version = review_capture.inspection_review_version(canonical)
|
||||
located = self._review_location(media_file, roots)
|
||||
if version is None or located is None:
|
||||
return None
|
||||
locator = os.path.normcase(located[3])
|
||||
try:
|
||||
facts = os.lstat(located[3])
|
||||
except OSError:
|
||||
facts = None
|
||||
state = None
|
||||
if facts is not None:
|
||||
state = _digest(f"{facts.st_size}:{facts.st_mtime_ns}:{version}")
|
||||
if self.db.review_decided_at(
|
||||
locator, facts.st_size, facts.st_mtime_ns
|
||||
) or self.db.review_fingerprint_queued(
|
||||
locator, facts.st_size, facts.st_mtime_ns, version
|
||||
):
|
||||
return canonical, None
|
||||
memory = self._untraced_memory()
|
||||
key = _memory_key(locator)
|
||||
if state is not None and memory.get(key) == state:
|
||||
return None
|
||||
resolution = review_capture.inspection_choice(canonical)
|
||||
if resolution is not None:
|
||||
memory.pop(key, None)
|
||||
return canonical, resolution
|
||||
if state is not None:
|
||||
memory.pop(key, None) # re-inserted last: the newest state
|
||||
memory[key] = state
|
||||
while len(memory) > MAX_UNTRACED_MEMORY:
|
||||
del memory[next(iter(memory))]
|
||||
return None
|
||||
|
||||
def _settle_stale_reviews(self, media_file, roots, canonical, review=None):
|
||||
"""Remove the pending, undecided items of an untraced file's current
|
||||
bytes recorded under another version (best effort, counted in
|
||||
``review["untraced"]``).
|
||||
|
||||
The current resolver and table offer nothing to choose in these
|
||||
bytes, so an item of an older version would only refuse Apply.
|
||||
"""
|
||||
located = self._review_location(media_file, roots)
|
||||
if located is None or canonical is None:
|
||||
return 0
|
||||
# An unusable semantics table says nothing about the current
|
||||
# versions; never let it retire items a working table recorded.
|
||||
if review_capture.inspection_review_version(canonical) is None:
|
||||
return 0
|
||||
try:
|
||||
facts = os.lstat(located[3])
|
||||
removed = self.db.settle_stale_review_items(
|
||||
os.path.normcase(located[3]),
|
||||
facts.st_size,
|
||||
facts.st_mtime_ns,
|
||||
review_capture.item_versions(canonical),
|
||||
)
|
||||
except Exception:
|
||||
self.logger.warning(
|
||||
f"Stale Needs review items of {media_file.name} could not be "
|
||||
"removed",
|
||||
exc_info=True,
|
||||
)
|
||||
if review is not None:
|
||||
review["errors"] += 1
|
||||
return 0
|
||||
if review is not None:
|
||||
review["untraced"] += removed
|
||||
return removed
|
||||
|
||||
def _untraced_memory(self):
|
||||
memory = getattr(self, "_untraced_inspections", None)
|
||||
if memory is None:
|
||||
memory = self._untraced_inspections = {}
|
||||
return memory
|
||||
|
||||
@staticmethod
|
||||
def _review_location(media_file, roots, real_roots=None):
|
||||
"""``(root, root_index, relative, spelled)`` of a listed file, or None.
|
||||
@@ -333,19 +178,24 @@ class ReviewRecordingMixin:
|
||||
review["untraced"] += removed
|
||||
return removed
|
||||
|
||||
def _fallback_links(self, metadata):
|
||||
"""True when the scan's legacy fallback would link these records."""
|
||||
return bool(metadata.get("prompt")) and bool(self._graph_prompt_text(metadata))
|
||||
|
||||
def record_review_file(self, image_path):
|
||||
"""Apply the Needs review rule to one monitor output (its worker thread).
|
||||
|
||||
Only scans add items (maintainer ruling, 2026-10-23): the monitor
|
||||
only clears the items of outputs it links. Maps ``image_path`` to the
|
||||
gallery root that lists it; a file under no current root is
|
||||
``skipped``. A linked file's pending items are removed and the answer
|
||||
is ``linked``. Any other file is ``skipped``: nothing is captured or
|
||||
recorded, and the next scan reads it. While a job runs that a
|
||||
``records`` job may not run beside, other than thumbnail work
|
||||
(``PREVIEW_ONLY_KINDS``), nothing is written and the answer is
|
||||
``deferred``. Errors propagate: the monitor logs them and keeps its
|
||||
link.
|
||||
Maps ``image_path`` to the gallery root that lists it; a file under no
|
||||
current root is ``skipped``. A linked file is never queued: its
|
||||
pending items are removed and the answer is ``linked``. An unlinked
|
||||
file the scan's fallback would link is ``skipped`` (the next scan
|
||||
links it); any other file is recorded when it is ambiguous with a
|
||||
choice, and is ``untraced`` (its pending items without a choice
|
||||
removed) when it is ambiguous without one. While a
|
||||
job runs that a ``records`` job may not run beside, other than
|
||||
thumbnail work (``PREVIEW_ONLY_KINDS``), nothing is written and the
|
||||
answer is ``deferred``. Errors propagate: the monitor logs
|
||||
them and keeps its link.
|
||||
"""
|
||||
busy = self._jobs.blocker(
|
||||
REVIEW_RECORDING_KIND, admission="records", compatible=PREVIEW_ONLY_KINDS
|
||||
@@ -365,7 +215,9 @@ class ReviewRecordingMixin:
|
||||
if self._review_file_linked(located, roots):
|
||||
self._settle_scanned_reviews([located], roots)
|
||||
return "linked"
|
||||
return "skipped"
|
||||
if self._fallback_links(self._extract_comfyui_metadata(str(located))):
|
||||
return "skipped"
|
||||
return self._record_review_sync(located, roots)
|
||||
|
||||
@staticmethod
|
||||
def _log_review_items(job, review):
|
||||
|
||||
@@ -40,7 +40,6 @@ _SETTINGS_FILE_KEYS = {
|
||||
"INFINITE_SCROLL": ("web_ui", "infinite_scroll"),
|
||||
"WORKER_THREADS": ("performance", "worker_threads"),
|
||||
"MONITORING_ENABLED": ("gallery", "monitoring", "enabled"),
|
||||
"SAFETY_BACKUP": ("maintenance", "safety_backup"),
|
||||
}
|
||||
# Attributes that live on GalleryConfig; the others on PromptManagerConfig.
|
||||
_GALLERY_SETTINGS = frozenset({"MONITORING_ENABLED"})
|
||||
@@ -188,7 +187,6 @@ class SettingsRoutesMixin:
|
||||
"webui_display_mode": PromptManagerConfig.WEBUI_DISPLAY_MODE,
|
||||
"infinite_scroll": PromptManagerConfig.INFINITE_SCROLL,
|
||||
"worker_threads": PromptManagerConfig.WORKER_THREADS,
|
||||
"safety_backup": PromptManagerConfig.SAFETY_BACKUP,
|
||||
"cpu_count": PromptManagerConfig.max_worker_threads(),
|
||||
"gallery_root_paths": root_paths,
|
||||
"gallery_root_path": root_paths[0] if root_paths else "",
|
||||
@@ -483,10 +481,6 @@ class SettingsRoutesMixin:
|
||||
if not isinstance(data["monitoring_enabled"], bool):
|
||||
raise ValueError("monitoring_enabled must be true or false")
|
||||
updates["MONITORING_ENABLED"] = data["monitoring_enabled"]
|
||||
if "safety_backup" in data:
|
||||
if not isinstance(data["safety_backup"], bool):
|
||||
raise ValueError("safety_backup must be true or false")
|
||||
updates["SAFETY_BACKUP"] = data["safety_backup"]
|
||||
if "worker_threads" in data:
|
||||
updates["WORKER_THREADS"] = validate_worker_threads(
|
||||
data["worker_threads"], PromptManagerConfig.max_worker_threads()
|
||||
|
||||
+1
-25
@@ -13,9 +13,7 @@
|
||||
- ``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; ``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 public path (``_public_path``), never an absolute one.
|
||||
|
||||
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
|
||||
@@ -211,27 +209,6 @@ 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).
|
||||
|
||||
@@ -256,7 +233,6 @@ 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)
|
||||
]
|
||||
|
||||
+1
-4
@@ -168,10 +168,9 @@ class SystemRoutesMixin:
|
||||
from ..config import GalleryConfig
|
||||
|
||||
try:
|
||||
(output_dirs, auto_detected), counts, location = await _gather_settled(
|
||||
(output_dirs, auto_detected), counts = 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")
|
||||
@@ -195,8 +194,6 @@ 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"},
|
||||
)
|
||||
|
||||
@@ -1,843 +0,0 @@
|
||||
"""Automatic thumbnails (#249): previews are a cache the server keeps itself.
|
||||
|
||||
The approved mock has no repeated global thumbnail prompt
|
||||
(``docs/concepts/workspace/README.md``), so the workspace never asks. Every
|
||||
listing gives an image or video a thumbnail URL (``media_thumbnail``), and the
|
||||
previews behind those URLs are made here, in three ways:
|
||||
|
||||
- **On first view.** ``GET /prompt_manager/images/serve/thumbnails/<rel>_thumb
|
||||
<ext>`` for a preview that does not exist yet is admitted on the event loop
|
||||
(``_PreviewGate``): at most ``ON_DEMAND_LIMIT`` previews are written at once,
|
||||
on this API's own small pool (never the default executor), one flight runs
|
||||
per file, and the requests waiting are capped (``MAX_WAITING``); past the
|
||||
cap a request gets its fallback at once. Inside the admission the preview is
|
||||
mapped back to its media file under the same root (the request text is
|
||||
checked first: no traversal, drive, UNC or empty segment; then a regular
|
||||
file inside the root, never through a link), made and served. When it
|
||||
cannot be made an image's original is served and a video's placeholder
|
||||
frame, both ``Cache-Control: no-store``: a video is never sent to an
|
||||
``<img>``. A file that cannot be decoded is remembered for this session
|
||||
(keyed by its path, size and mtime) and counted for the System Overview
|
||||
(``failed`` in ``/images/thumbnails/status``); I/O errors and a preview
|
||||
another program holds open (Windows) are not remembered. The recorded
|
||||
history's thumbnail route goes through the same admission and flight key.
|
||||
- **On capture.** The image monitor hands every new output it settles to
|
||||
``queue_capture_thumbnail``, which queues its preview for one background
|
||||
worker thread (``AutoThumbnailer``). Each wake takes at most Worker threads
|
||||
(``WORKER_THREADS``) files and writes them on that many threads; enqueuing
|
||||
is O(1), so the monitor and ComfyUI's queue never wait. A file already
|
||||
queued is not queued again, the queue is bounded (``MAX_QUEUED``) and a
|
||||
preview left out is made on its first view instead.
|
||||
- **After a scan.** A scan that added files checks the missing previews; from
|
||||
``SCAN_THUMBNAIL_JOB_THRESHOLD`` up it starts the thumbnail Activity job
|
||||
(missing only, visible and cancellable) instead of writing them silently,
|
||||
and fewer go through the worker.
|
||||
|
||||
Every preview is written by ``_generate_one`` with the thumbnail job's default
|
||||
size and quality (``DEFAULT_THUMBNAIL_EDGE``, medium), into a temp sibling
|
||||
renamed over the old one, so a reader never sees a partial file.
|
||||
|
||||
Holds: automatic previews only write files under ``thumbnails/`` and never
|
||||
touch the database, like the thumbnail job, so they run beside a scan, image
|
||||
tagging and the Needs review recording (S5 M7). They are not written while a
|
||||
job runs that replaces the database or removes media and previews: a restore,
|
||||
a duplicate removal, a sources change, the gallery's thumbnail clear or an
|
||||
``exclusive`` job (``_auto_thumbnail_blocker``). The holds are checked when a
|
||||
write starts (after its slot) and again just before its rename.
|
||||
|
||||
Shutdown: ``close_auto_thumbnails`` stops the worker at a file boundary,
|
||||
kills a running ffmpeg, waits a bounded time (and logs when the wait ran
|
||||
out) and shuts the view pool down; after the app's cleanup
|
||||
(``final=True``) nothing starts a worker again. Importing this module again
|
||||
(a package reload) stops the workers the previous import left.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import collections
|
||||
import functools
|
||||
import io
|
||||
import os
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
from .images import (
|
||||
GALLERY_IMAGE_EXTENSIONS,
|
||||
MEDIA_EXTENSIONS,
|
||||
THUMBNAIL_CLEAR_TITLE,
|
||||
THUMBNAIL_ENCODINGS,
|
||||
VIDEO_EXTENSIONS,
|
||||
_has_traversal,
|
||||
_is_within,
|
||||
_thumbnail_is_fresh,
|
||||
_video_placeholder_image,
|
||||
_worker_threads,
|
||||
media_thumbnail_rel,
|
||||
)
|
||||
from .jobs import (
|
||||
DUPLICATE_CLEANUP_JOB_KIND,
|
||||
RESTORE_JOB_KIND,
|
||||
THUMBNAILS_JOB_KIND,
|
||||
JobBusy,
|
||||
)
|
||||
from .settings_routes import SOURCE_CHANGE_JOB_KIND
|
||||
from .thumbnail_jobs import DEFAULT_THUMBNAIL_EDGE, ThumbnailOptions
|
||||
|
||||
try:
|
||||
from ...utils.parallel import map_parallel
|
||||
except ImportError: # tests import the package top-level
|
||||
from utils.parallel import map_parallel
|
||||
|
||||
AUTO_THUMBNAIL_EDGE = DEFAULT_THUMBNAIL_EDGE
|
||||
AUTO_THUMBNAIL_SIZE = (AUTO_THUMBNAIL_EDGE, AUTO_THUMBNAIL_EDGE)
|
||||
AUTO_THUMBNAIL_ENCODING = THUMBNAIL_ENCODINGS["medium"]
|
||||
# Previews views may write at once (the size of the view pool).
|
||||
ON_DEMAND_LIMIT = 3
|
||||
# View requests that may wait on the gate; one more gets its fallback at once.
|
||||
MAX_WAITING = 64
|
||||
# How long a pool thread waits for the capture worker's flight of the same file.
|
||||
FLIGHT_WAIT_SECONDS = 20
|
||||
# How long close() waits for the worker's batch in hand (ffmpeg's own limit).
|
||||
CLOSE_WAIT_SECONDS = 30
|
||||
# Files the capture worker holds; more are made on their first view.
|
||||
MAX_QUEUED = 10_000
|
||||
# Undecodable files remembered per session (oldest forgotten first).
|
||||
MAX_FAILURES_REMEMBERED = 4096
|
||||
# Missing previews after a scan from which the visible thumbnail job starts:
|
||||
# a fresh install or a newly added source. Fewer go through the worker.
|
||||
SCAN_THUMBNAIL_JOB_THRESHOLD = 200
|
||||
# Jobs beside which no automatic preview is written (see the module docstring).
|
||||
AUTO_REFUSED_KINDS = frozenset(
|
||||
{RESTORE_JOB_KIND, DUPLICATE_CLEANUP_JOB_KIND, SOURCE_CHANGE_JOB_KIND}
|
||||
)
|
||||
WORKER_NAME = "PromptManagerThumbnailWorker"
|
||||
POOL_PREFIX = "pm-preview"
|
||||
# Identifies this import's workers: a reload stops the ones it did not start.
|
||||
_MODULE_TOKEN = object()
|
||||
|
||||
READY, FAILED, BUSY, DEFERRED = "ready", "failed", "busy", "deferred"
|
||||
# What the gate answers a request past MAX_WAITING.
|
||||
OVER = object()
|
||||
|
||||
|
||||
def _fingerprint(src):
|
||||
"""``(path, size, mtime_ns)`` of *src*, or None when it cannot be read."""
|
||||
try:
|
||||
stat = os.stat(src)
|
||||
except OSError:
|
||||
return None
|
||||
return (os.path.normcase(str(src)), stat.st_size, stat.st_mtime_ns)
|
||||
|
||||
|
||||
def flight_key(root, rel_parts):
|
||||
"""The single-flight key of ``root/rel_parts``: its folder resolved, case
|
||||
folded. The serve route, the history route and the worker share it."""
|
||||
path = Path(root).joinpath(*rel_parts)
|
||||
return os.path.normcase(os.path.join(os.path.realpath(path.parent), path.name))
|
||||
|
||||
|
||||
def _dst_key(dst):
|
||||
return flight_key(Path(dst).parent, (Path(dst).name,))
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def video_placeholder_jpeg():
|
||||
"""The placeholder frame as JPEG bytes (a video whose preview is missing)."""
|
||||
buffer = io.BytesIO()
|
||||
_video_placeholder_image(AUTO_THUMBNAIL_SIZE).save(
|
||||
buffer, "JPEG", **AUTO_THUMBNAIL_ENCODING
|
||||
)
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
def _is_preview_request(filepath):
|
||||
"""Text check only: ``thumbnails/<...>/<name>_thumb<lower-case ext>``.
|
||||
|
||||
Refuses, as text and before any path is joined, what a joined path could
|
||||
take as an anchor: a backslash, a drive or stream colon, and empty, ``.``
|
||||
or ``..`` segments (so a leading ``/`` and ``//server/share`` too).
|
||||
"""
|
||||
if "\\" in filepath or ":" in filepath:
|
||||
return False
|
||||
parts = filepath.split("/")
|
||||
if len(parts) < 2 or parts[0] != "thumbnails":
|
||||
return False
|
||||
if any(part in ("", ".", "..") for part in parts) or "thumbnails" in parts[1:-1]:
|
||||
return False
|
||||
stem, extension = os.path.splitext(parts[-1])
|
||||
return (
|
||||
stem.endswith("_thumb")
|
||||
and len(stem) > len("_thumb")
|
||||
and extension == extension.lower()
|
||||
and extension in GALLERY_IMAGE_EXTENSIONS
|
||||
)
|
||||
|
||||
|
||||
def _regular_file(path):
|
||||
"""True for a regular file that is not a link (one lstat)."""
|
||||
try:
|
||||
return not path.is_symlink() and path.is_file()
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
|
||||
def _candidates(parts):
|
||||
"""``(folder parts, base, extensions)`` of a checked preview request."""
|
||||
stem, extension = os.path.splitext(parts[-1])
|
||||
wanted = [extension] + (sorted(VIDEO_EXTENSIONS) if extension == ".jpg" else [])
|
||||
return parts[1:-1], stem[: -len("_thumb")], wanted
|
||||
|
||||
|
||||
def _quick_source(root, parts):
|
||||
"""``(src, is_video)`` by exact name only (no folder listing), or None."""
|
||||
root = Path(root)
|
||||
folder_parts, base, wanted = _candidates(parts)
|
||||
folder = root.joinpath(*folder_parts)
|
||||
for extension in wanted:
|
||||
for name in (base + extension, base + extension.upper()):
|
||||
path = folder / name
|
||||
if _regular_file(path) and _is_within(path, root):
|
||||
return path, extension in VIDEO_EXTENSIONS
|
||||
return None
|
||||
|
||||
|
||||
def _thumbnail_source(root, parts):
|
||||
"""``(src, dst, is_video)`` for the preview *parts* under *root*, or None.
|
||||
|
||||
*parts* passed ``_is_preview_request``. The source has the preview's name
|
||||
without ``_thumb`` and an extension that previews as the requested one
|
||||
(any case: exact names first, then one listing of the folder); a ``.jpg``
|
||||
may also be a video's. It must be a regular file inside *root*. ``dst``
|
||||
is None when the preview's folder resolves outside ``root/thumbnails`` (a
|
||||
link): nothing is written there.
|
||||
"""
|
||||
root = Path(root)
|
||||
folder_parts, base, wanted = _candidates(parts)
|
||||
quick = _quick_source(root, parts)
|
||||
if quick is not None:
|
||||
found = quick[0]
|
||||
else:
|
||||
found = _scan_folder(root.joinpath(*folder_parts), base, set(wanted))
|
||||
if found is None or not _is_within(found, root):
|
||||
return None
|
||||
rel = Path(*folder_parts, found.name)
|
||||
dst = root / media_thumbnail_rel(rel)
|
||||
if dst != root.joinpath(*parts):
|
||||
return None
|
||||
if not _is_within(root / "thumbnails", root) or not _is_within(
|
||||
dst, root / "thumbnails"
|
||||
):
|
||||
dst = None
|
||||
return found, dst, found.suffix.lower() in VIDEO_EXTENSIONS
|
||||
|
||||
|
||||
def _scan_folder(folder, base, extensions):
|
||||
"""A mixed-case extension (``.Png``): one listing of *folder*, or None."""
|
||||
try:
|
||||
with os.scandir(folder) as entries:
|
||||
for entry in entries:
|
||||
stem, extension = os.path.splitext(entry.name)
|
||||
if stem == base and extension.lower() in extensions:
|
||||
path = Path(entry.path)
|
||||
if _regular_file(path):
|
||||
return path
|
||||
except OSError:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _stop_stale_workers():
|
||||
"""Stop capture workers a previous import of this module started."""
|
||||
for thread in threading.enumerate():
|
||||
token = getattr(thread, "pm_preview_token", None)
|
||||
close = getattr(thread, "pm_preview_close", None)
|
||||
if token is not None and token is not _MODULE_TOKEN and close is not None:
|
||||
close()
|
||||
|
||||
|
||||
class AutoThumbnailer:
|
||||
"""Single-flight preview writes and the capture worker (thread-safe).
|
||||
|
||||
``generate(src, dst, is_video, *, stop, guard, track)`` writes one
|
||||
preview and returns ``_generate_one``'s result dict; ``may_write()``
|
||||
says whether automatic previews may be written now. Views call
|
||||
``ensure`` on the view pool, the monitor calls ``enqueue`` from its
|
||||
thread.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
generate,
|
||||
may_write,
|
||||
logger,
|
||||
*,
|
||||
workers=_worker_threads,
|
||||
max_queued=MAX_QUEUED,
|
||||
start_worker=True,
|
||||
):
|
||||
self._generate = generate
|
||||
self._may_write = may_write
|
||||
self._logger = logger
|
||||
self._workers = workers
|
||||
self._max_queued = max_queued
|
||||
self._start_worker = start_worker
|
||||
self._changed = threading.Condition()
|
||||
self._inflight = {}
|
||||
self._failures = collections.OrderedDict()
|
||||
self._queue = collections.deque()
|
||||
self._queued = set()
|
||||
self._processes = []
|
||||
self._worker = None
|
||||
# True from a worker's start until it leaves (under the lock), so an
|
||||
# idle worker can exit and the next enqueue starts a fresh one.
|
||||
self._worker_running = False
|
||||
self._closed = False
|
||||
self._busy = False
|
||||
self._batch = 0
|
||||
|
||||
@property
|
||||
def failure_count(self):
|
||||
with self._changed:
|
||||
return len(self._failures)
|
||||
|
||||
@property
|
||||
def closed(self):
|
||||
return self._closed
|
||||
|
||||
def _allowed(self):
|
||||
return not self._closed and self._may_write()
|
||||
|
||||
def _track(self, process):
|
||||
"""Keep an ffmpeg process so close() can kill it."""
|
||||
with self._changed:
|
||||
self._processes = [p for p in self._processes if p.poll() is None]
|
||||
self._processes.append(process)
|
||||
closed = self._closed
|
||||
if closed:
|
||||
process.kill()
|
||||
|
||||
# ── one preview ───────────────────────────────────────────────────
|
||||
|
||||
def ensure(self, src, dst, is_video):
|
||||
"""Make *dst* for *src* unless it is fresh; returns the outcome.
|
||||
|
||||
``ready`` (the preview is there), ``failed`` (the file cannot be
|
||||
decoded, remembered), ``busy`` (an I/O error, the file is gone, or
|
||||
another flight did not finish) or ``deferred`` (a hold refuses
|
||||
automatic writes, checked when the write starts and before its
|
||||
rename). Blocking: call it off the event loop.
|
||||
"""
|
||||
if _thumbnail_is_fresh(src, dst):
|
||||
return READY
|
||||
fingerprint = _fingerprint(src)
|
||||
if fingerprint is None:
|
||||
return BUSY
|
||||
key = _dst_key(dst)
|
||||
with self._changed:
|
||||
if fingerprint in self._failures:
|
||||
return FAILED
|
||||
flight = self._inflight.get(key)
|
||||
if flight is None:
|
||||
self._inflight[key] = done = threading.Event()
|
||||
if flight is not None:
|
||||
flight.wait(FLIGHT_WAIT_SECONDS)
|
||||
if _thumbnail_is_fresh(src, dst):
|
||||
return READY
|
||||
with self._changed:
|
||||
return FAILED if fingerprint in self._failures else BUSY
|
||||
try:
|
||||
return self._lead(src, dst, is_video, fingerprint)
|
||||
finally:
|
||||
with self._changed:
|
||||
self._inflight.pop(key, None)
|
||||
self._changed.notify_all()
|
||||
done.set()
|
||||
|
||||
def _lead(self, src, dst, is_video, fingerprint):
|
||||
if not self._allowed():
|
||||
return DEFERRED
|
||||
if _thumbnail_is_fresh(src, dst):
|
||||
return READY
|
||||
result = self._generate(
|
||||
src,
|
||||
dst,
|
||||
is_video,
|
||||
stop=lambda: self._closed,
|
||||
guard=self._allowed,
|
||||
track=self._track,
|
||||
)
|
||||
if result.get("action") in ("generated", "skipped"):
|
||||
return READY
|
||||
if result.get("held"):
|
||||
return DEFERRED
|
||||
if not result.get("decode"):
|
||||
return BUSY # an I/O error says nothing about the file
|
||||
error = result.get("error") or "failed"
|
||||
with self._changed:
|
||||
self._failures[fingerprint] = error
|
||||
while len(self._failures) > MAX_FAILURES_REMEMBERED:
|
||||
self._failures.popitem(last=False)
|
||||
self._logger.info(
|
||||
f"No preview for {os.path.basename(str(src))} ({error}); "
|
||||
"the original is shown"
|
||||
)
|
||||
return FAILED
|
||||
|
||||
# ── the capture worker ────────────────────────────────────────────
|
||||
|
||||
def enqueue(self, src, dst, is_video):
|
||||
"""Queue *dst*'s preview for the worker; False when already queued,
|
||||
full or closed. Never blocks on a write."""
|
||||
key = _dst_key(dst)
|
||||
with self._changed:
|
||||
if self._closed or key in self._queued:
|
||||
return False
|
||||
if len(self._queue) >= self._max_queued:
|
||||
return False
|
||||
self._queued.add(key)
|
||||
self._queue.append((key, Path(src), Path(dst), bool(is_video)))
|
||||
self._busy = True
|
||||
if self._start_worker and not self._worker_running:
|
||||
self._start_worker_locked()
|
||||
self._changed.notify_all()
|
||||
return True
|
||||
|
||||
def _start_worker_locked(self):
|
||||
"""Start the worker (lock held). A start that fails leaves the file
|
||||
queued and the next enqueue tries again."""
|
||||
worker = threading.Thread(target=self._run, name=WORKER_NAME, daemon=True)
|
||||
worker.pm_preview_token = _MODULE_TOKEN
|
||||
worker.pm_preview_close = self.close
|
||||
self._worker_running = True
|
||||
try:
|
||||
worker.start()
|
||||
except RuntimeError:
|
||||
self._worker_running = False
|
||||
self._busy = False
|
||||
self._logger.warning(
|
||||
"The automatic preview worker could not start; the next output "
|
||||
"tries again",
|
||||
exc_info=True,
|
||||
)
|
||||
return
|
||||
self._worker = worker
|
||||
|
||||
def _run(self):
|
||||
try:
|
||||
self._drain()
|
||||
except Exception:
|
||||
self._logger.warning("The automatic preview worker stopped", exc_info=True)
|
||||
finally:
|
||||
# Whatever ended the loop, the next enqueue starts a fresh worker.
|
||||
with self._changed:
|
||||
self._worker_running = False
|
||||
self._busy = False
|
||||
self._batch = 0
|
||||
self._changed.notify_all()
|
||||
|
||||
def _drain(self):
|
||||
while True:
|
||||
with self._changed:
|
||||
if self._closed or not self._queue:
|
||||
# Idle or closed: leave, so no thread outlives its work.
|
||||
self._busy = False
|
||||
self._worker_running = False
|
||||
self._changed.notify_all()
|
||||
return
|
||||
take = min(len(self._queue), max(1, int(self._workers() or 1)))
|
||||
batch = [self._queue.popleft() for _ in range(take)]
|
||||
self._batch = len(batch)
|
||||
self._changed.notify_all()
|
||||
try:
|
||||
map_parallel(self._write_queued, batch, len(batch))
|
||||
finally:
|
||||
with self._changed:
|
||||
self._batch = 0
|
||||
for key, *_rest in batch:
|
||||
self._queued.discard(key)
|
||||
|
||||
def _write_queued(self, item):
|
||||
_key, src, dst, is_video = item
|
||||
try:
|
||||
self.ensure(src, dst, is_video)
|
||||
except Exception: # one file never stops the worker
|
||||
self._logger.warning(
|
||||
f"Automatic preview failed for {src.name}", exc_info=True
|
||||
)
|
||||
|
||||
def close(self):
|
||||
"""Stop the worker at its next file boundary (queued files are
|
||||
dropped), kill a running ffmpeg and wait up to CLOSE_WAIT_SECONDS;
|
||||
a worker still running after that is logged."""
|
||||
with self._changed:
|
||||
self._closed = True
|
||||
self._queue.clear()
|
||||
self._queued.clear()
|
||||
processes, self._processes = self._processes, []
|
||||
self._changed.notify_all()
|
||||
worker = self._worker
|
||||
for process in processes:
|
||||
try:
|
||||
process.kill()
|
||||
except OSError:
|
||||
pass
|
||||
if worker is None or worker is threading.current_thread():
|
||||
return
|
||||
worker.join(CLOSE_WAIT_SECONDS)
|
||||
if worker.is_alive():
|
||||
self._logger.warning(
|
||||
"The automatic preview worker is still running after "
|
||||
f"{CLOSE_WAIT_SECONDS} s; it stops at its next file"
|
||||
)
|
||||
|
||||
# ── waits (tests and shutdown; never timers) ──────────────────────
|
||||
|
||||
def _wait_for(self, predicate, timeout):
|
||||
with self._changed:
|
||||
return self._changed.wait_for(predicate, timeout)
|
||||
|
||||
def wait_idle(self, timeout=None):
|
||||
"""True once the worker has nothing queued or in hand."""
|
||||
return self._wait_for(lambda: not self._busy, timeout)
|
||||
|
||||
def wait_batch(self, size, timeout=None):
|
||||
"""True once the worker holds a batch of *size* files."""
|
||||
return self._wait_for(lambda: self._batch == size, timeout)
|
||||
|
||||
|
||||
class _PreviewGate:
|
||||
"""Admission of view requests on the event loop (#249 I2).
|
||||
|
||||
Waiters are coroutines: at most ``limit`` previews are written at once
|
||||
on the view pool, one flight runs per key, and past ``MAX_WAITING``
|
||||
waiting requests ``run`` answers ``OVER`` at once.
|
||||
"""
|
||||
|
||||
def __init__(self, limit, logger):
|
||||
self.loop = asyncio.get_running_loop()
|
||||
self._slots = asyncio.Semaphore(limit)
|
||||
self._flights = {}
|
||||
self._tasks = set()
|
||||
self._logger = logger
|
||||
self._tick = asyncio.Event()
|
||||
self.waiting = 0
|
||||
|
||||
def _changed(self):
|
||||
tick, self._tick = self._tick, asyncio.Event()
|
||||
tick.set()
|
||||
|
||||
async def settled(self, predicate):
|
||||
"""Return once ``predicate()`` holds (re-checked at each change)."""
|
||||
while not predicate():
|
||||
await self._tick.wait()
|
||||
|
||||
async def run(self, key, work, pool):
|
||||
"""``work()``'s answer for *key* (shared by every request of the
|
||||
flight), None when it raised, or OVER past the cap."""
|
||||
if self.waiting >= MAX_WAITING:
|
||||
return OVER
|
||||
flight = self._flights.get(key)
|
||||
if flight is None:
|
||||
flight = self.loop.create_future()
|
||||
self._flights[key] = flight
|
||||
task = self.loop.create_task(self._lead(key, flight, work, pool))
|
||||
self._tasks.add(task)
|
||||
task.add_done_callback(self._tasks.discard)
|
||||
self.waiting += 1
|
||||
self._changed()
|
||||
try:
|
||||
return await asyncio.shield(flight)
|
||||
finally:
|
||||
self.waiting -= 1
|
||||
self._changed()
|
||||
|
||||
async def _lead(self, key, flight, work, pool):
|
||||
result = None
|
||||
try:
|
||||
async with self._slots:
|
||||
result = await self.loop.run_in_executor(pool, work)
|
||||
except Exception:
|
||||
self._logger.warning("Automatic preview failed", exc_info=True)
|
||||
finally:
|
||||
self._flights.pop(key, None)
|
||||
if not flight.done():
|
||||
flight.set_result(result)
|
||||
|
||||
|
||||
class AutoThumbnailMixin:
|
||||
"""Automatic previews on the API: views, the monitor's queue, scans."""
|
||||
|
||||
_auto_thumbnailer = None
|
||||
_auto_thumbnails_shut = False
|
||||
_preview_gate_obj = None
|
||||
_preview_pool_obj = None
|
||||
_auto_thumbnailer_lock = threading.Lock()
|
||||
|
||||
def _auto_thumbnails(self):
|
||||
"""This API's AutoThumbnailer, created on first use (no worker after
|
||||
the app's cleanup)."""
|
||||
with AutoThumbnailMixin._auto_thumbnailer_lock:
|
||||
if self._auto_thumbnailer is None:
|
||||
self._auto_thumbnailer = AutoThumbnailer(
|
||||
self._write_auto_thumbnail,
|
||||
lambda: self._auto_thumbnail_blocker() is None,
|
||||
self.logger,
|
||||
start_worker=not self._auto_thumbnails_shut,
|
||||
)
|
||||
return self._auto_thumbnailer
|
||||
|
||||
def close_auto_thumbnails(self, final=False):
|
||||
"""Stop the capture worker and the view pool, and forget them.
|
||||
|
||||
Bounded (see AutoThumbnailer.close). ``final`` (the app's cleanup)
|
||||
also keeps a late monitor event from starting a worker again; a
|
||||
test reset passes False and the next use starts fresh ones.
|
||||
"""
|
||||
with AutoThumbnailMixin._auto_thumbnailer_lock:
|
||||
if final:
|
||||
self._auto_thumbnails_shut = True
|
||||
thumbnailer, self._auto_thumbnailer = self._auto_thumbnailer, None
|
||||
pool, self._preview_pool_obj = self._preview_pool_obj, None
|
||||
self._preview_gate_obj = None
|
||||
if thumbnailer is not None:
|
||||
thumbnailer.close()
|
||||
if pool is not None:
|
||||
pool.shutdown(wait=False, cancel_futures=True)
|
||||
|
||||
def _auto_thumbnail_failures(self):
|
||||
thumbnailer = self._auto_thumbnailer
|
||||
return 0 if thumbnailer is None else thumbnailer.failure_count
|
||||
|
||||
def _write_auto_thumbnail(self, src, dst, is_video, *, stop, guard, track):
|
||||
return self._generate_one(
|
||||
src,
|
||||
dst,
|
||||
AUTO_THUMBNAIL_SIZE,
|
||||
is_video,
|
||||
encoding=AUTO_THUMBNAIL_ENCODING,
|
||||
stop=stop,
|
||||
guard=guard,
|
||||
track=track,
|
||||
)
|
||||
|
||||
def _auto_thumbnail_blocker(self):
|
||||
"""The running job automatic previews wait for, or None (any thread)."""
|
||||
for job in self._jobs.running():
|
||||
if job.admission == "exclusive" or job.kind in AUTO_REFUSED_KINDS:
|
||||
return job
|
||||
if job.kind == THUMBNAILS_JOB_KIND and job.title == THUMBNAIL_CLEAR_TITLE:
|
||||
return job
|
||||
return None
|
||||
|
||||
# ── the view admission ────────────────────────────────────────────
|
||||
|
||||
def _preview_gate(self):
|
||||
"""The gate of the running loop (a new loop gets a new one)."""
|
||||
gate = self._preview_gate_obj
|
||||
if gate is None or gate.loop is not asyncio.get_running_loop():
|
||||
gate = self._preview_gate_obj = _PreviewGate(ON_DEMAND_LIMIT, self.logger)
|
||||
return gate
|
||||
|
||||
def _preview_pool(self):
|
||||
with AutoThumbnailMixin._auto_thumbnailer_lock:
|
||||
if self._preview_pool_obj is None:
|
||||
self._preview_pool_obj = ThreadPoolExecutor(
|
||||
ON_DEMAND_LIMIT, thread_name_prefix=POOL_PREFIX
|
||||
)
|
||||
return self._preview_pool_obj
|
||||
|
||||
def _preview_pool_threads(self):
|
||||
pool = self._preview_pool_obj
|
||||
return [] if pool is None else list(getattr(pool, "_threads", ()))
|
||||
|
||||
def _video_placeholder_bytes(self):
|
||||
return video_placeholder_jpeg()
|
||||
|
||||
async def _admit_preview(self, src, dst, is_video):
|
||||
"""Make one preview through the view admission (the history route).
|
||||
|
||||
Shares the serve route's flight key; every flight answers the
|
||||
``ensure`` outcome (or None, or OVER), whichever route started it.
|
||||
"""
|
||||
return await self._preview_gate().run(
|
||||
flight_key(Path(dst).parent, (Path(dst).name,)),
|
||||
functools.partial(self._ensure_view, src, dst, is_video),
|
||||
self._preview_pool(),
|
||||
)
|
||||
|
||||
def _ensure_view(self, src, dst, is_video):
|
||||
"""``ensure``, reporting a preview removed under us as busy."""
|
||||
outcome = self._auto_thumbnails().ensure(src, dst, is_video)
|
||||
if outcome == READY and not Path(dst).is_file():
|
||||
return BUSY # removed under us: the fallback this once
|
||||
return outcome
|
||||
|
||||
# ── on first view ─────────────────────────────────────────────────
|
||||
|
||||
def _preview_on_demand(self, root, parts):
|
||||
"""The ``ensure`` outcome for the preview *parts* under *root*
|
||||
(``deferred`` when its folder is linked out of the root), or None
|
||||
when no media file previews as *parts*."""
|
||||
located = _thumbnail_source(root, parts)
|
||||
if located is None:
|
||||
return None
|
||||
src, dst, is_video = located
|
||||
if dst is None:
|
||||
return DEFERRED
|
||||
return self._ensure_view(src, dst, is_video)
|
||||
|
||||
def _preview_response(self, outcome, src, dst, is_video):
|
||||
if outcome == READY and dst is not None:
|
||||
return self._file_response(dst)
|
||||
if is_video:
|
||||
return web.Response(
|
||||
body=video_placeholder_jpeg(),
|
||||
content_type="image/jpeg",
|
||||
headers={"Cache-Control": "no-store"},
|
||||
)
|
||||
response = web.FileResponse(src)
|
||||
response.headers["Cache-Control"] = "no-store"
|
||||
return response
|
||||
|
||||
async def _serve_missing_thumbnail(self, output_dirs, filepath):
|
||||
"""The made preview, or its fallback; None when *filepath* names no
|
||||
preview of a media file in *output_dirs*.
|
||||
|
||||
The flight answers only the outcome (a history flight may be the one
|
||||
joined); the file it previews is derived here: the preview path is
|
||||
the request's own, and a fallback looks its source up again on the
|
||||
view pool (bounded by it, never the default executor).
|
||||
"""
|
||||
if _has_traversal(filepath) or not _is_preview_request(filepath):
|
||||
return None
|
||||
parts = tuple(filepath.split("/"))
|
||||
gate = self._preview_gate()
|
||||
pool = self._preview_pool()
|
||||
for root in output_dirs:
|
||||
root = Path(root)
|
||||
outcome = await gate.run(
|
||||
flight_key(root, parts),
|
||||
functools.partial(self._preview_on_demand, root, parts),
|
||||
pool,
|
||||
)
|
||||
if outcome is None:
|
||||
continue
|
||||
if outcome is OVER:
|
||||
# Past the cap: exact names only, no listing, no write.
|
||||
quick = _quick_source(root, parts)
|
||||
if quick is None:
|
||||
continue
|
||||
return self._preview_response(BUSY, quick[0], None, quick[1])
|
||||
dst = root.joinpath(*parts)
|
||||
if outcome == READY:
|
||||
return self._file_response(dst)
|
||||
located = await gate.loop.run_in_executor(
|
||||
pool, _thumbnail_source, root, parts
|
||||
)
|
||||
if located is None:
|
||||
continue
|
||||
src, _dst, is_video = located
|
||||
return self._preview_response(outcome, src, None, is_video)
|
||||
return None
|
||||
|
||||
# ── on capture ────────────────────────────────────────────────────
|
||||
|
||||
def queue_capture_thumbnail(self, image_path):
|
||||
"""Queue the preview of a new output (the monitor's thread).
|
||||
|
||||
Returns ``queued``, ``fresh`` (its preview is current) or ``skipped``
|
||||
(not a media file under a current root, a link, a preview itself,
|
||||
already queued, or after the app's cleanup). Never writes on the
|
||||
caller's thread.
|
||||
"""
|
||||
if self._auto_thumbnails_shut:
|
||||
return "skipped"
|
||||
located = self._capture_target(image_path)
|
||||
if located is None:
|
||||
return "skipped"
|
||||
src, dst, is_video = located
|
||||
if _thumbnail_is_fresh(src, dst):
|
||||
return "fresh"
|
||||
if self._auto_thumbnails().enqueue(src, dst, is_video):
|
||||
return "queued"
|
||||
return "skipped"
|
||||
|
||||
def _capture_target(self, image_path):
|
||||
path = Path(os.path.abspath(str(image_path)))
|
||||
if path.suffix.lower() not in MEDIA_EXTENSIONS or not _regular_file(path):
|
||||
return None
|
||||
folded = os.path.normcase(str(path))
|
||||
for root in self._get_all_output_dirs():
|
||||
root = Path(os.path.abspath(str(root)))
|
||||
base = os.path.normcase(str(root)).rstrip("\\/") + os.sep
|
||||
if not folded.startswith(base) or not _is_within(path, root):
|
||||
continue
|
||||
rel = Path(os.path.relpath(path, root))
|
||||
if "thumbnails" in rel.parts[:-1]:
|
||||
return None
|
||||
dst = root / media_thumbnail_rel(rel)
|
||||
if not _is_within(root / "thumbnails", root) or not _is_within(
|
||||
dst, root / "thumbnails"
|
||||
):
|
||||
return None
|
||||
return path, dst, path.suffix.lower() in VIDEO_EXTENSIONS
|
||||
return None
|
||||
|
||||
# ── after a scan ──────────────────────────────────────────────────
|
||||
|
||||
async def _thumbnails_after_scan(self, job, output_dirs=None):
|
||||
"""Hand a finished scan's missing previews on; returns what it did.
|
||||
|
||||
None (the scan added nothing), ``job`` (the thumbnail job started for
|
||||
``SCAN_THUMBNAIL_JOB_THRESHOLD`` or more missing previews), ``busy``
|
||||
(a job refused it; the previews are made on view), ``queued`` (fewer,
|
||||
through the worker) or ``none`` (nothing missing). It counts the
|
||||
sources the scan walked (``output_dirs``, narrowed to the job's
|
||||
``roots``; every source when None). The server log says which.
|
||||
"""
|
||||
if job.counters.get("added", 0) <= 0:
|
||||
return None
|
||||
if output_dirs is None:
|
||||
output_dirs = await self._run_in_executor(self._get_all_output_dirs)
|
||||
roots = job.scope.get("roots")
|
||||
if roots:
|
||||
output_dirs = [output_dirs[i] for i in roots if 0 <= i < len(output_dirs)]
|
||||
plan = await self._run_in_executor(self._plan_gallery_thumbnails, output_dirs)
|
||||
missing = len(plan["targets"]) + plan["remaining"]
|
||||
if missing >= SCAN_THUMBNAIL_JOB_THRESHOLD:
|
||||
options = ThumbnailOptions("missing", AUTO_THUMBNAIL_EDGE, "medium")
|
||||
try:
|
||||
self._start_thumbnail_job(options)
|
||||
except JobBusy as busy:
|
||||
self.logger.info(
|
||||
f"{missing:,} previews missing after the scan; the thumbnail "
|
||||
f"job did not start while {busy.running.title} runs"
|
||||
)
|
||||
return "busy"
|
||||
self.logger.info(
|
||||
f"{missing:,} previews missing after the scan; thumbnail job started"
|
||||
)
|
||||
return "job"
|
||||
if not missing:
|
||||
return "none"
|
||||
thumbnailer = self._auto_thumbnails()
|
||||
for src, dst, is_video in plan["targets"]:
|
||||
thumbnailer.enqueue(src, dst, is_video)
|
||||
self.logger.info(f"{missing:,} previews missing after the scan; queued")
|
||||
return "queued"
|
||||
|
||||
async def _hand_off_scan_thumbnails(self, job, output_dirs=None):
|
||||
"""``_thumbnails_after_scan`` for the scan runner: never raises."""
|
||||
try:
|
||||
return await self._thumbnails_after_scan(job, output_dirs)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
self.logger.warning("Previews after the scan failed", exc_info=True)
|
||||
return None
|
||||
|
||||
|
||||
# A package reload re-runs this module: stop the workers the old one left.
|
||||
_stop_stale_workers()
|
||||
@@ -1,277 +0,0 @@
|
||||
"""Store each workflow once: the migration and Reclaim space jobs (#249 S6).
|
||||
|
||||
Reviewed like every maintenance operation (``POST /prompt_manager/
|
||||
maintenance/preview {operation:"workflow_storage"}``, a ``Plan`` with the
|
||||
records to rewrite, the backup and the disk check) and applied with its token
|
||||
(``POST /prompt_manager/maintenance/apply``). The apply starts:
|
||||
|
||||
1. **Store each workflow once**: kind ``maintenance``, admission ``records``,
|
||||
cancellable, groups ``changed``/``protected``/``errors``. It always takes
|
||||
a verified pre-migration backup, whatever the safety backup policy, named
|
||||
``<db>.pre-workflow-storage_<stamp>`` outside the pruned ``<db>.backup_*``
|
||||
pool and kept until the user removes it (``backup-failed`` ends the job
|
||||
with nothing rewritten). It then rewrites old-form rows batch by batch
|
||||
(``database/workflow_storage.py``): one transaction per batch of at most
|
||||
``BATCH_SECONDS``, then ``BATCH_PAUSE`` with the write lock free, progress
|
||||
after each, Cancel honoured between batches. A row that cannot be
|
||||
rewritten stays inline and is listed in ``errors``; any other failure
|
||||
rolls its batch back whole and fails the job. Every record still reads as
|
||||
before, and the next apply continues from the first old-form row.
|
||||
2. **Reclaim space**, once every record was looked at: kind ``maintenance``,
|
||||
admission ``exclusive`` (it runs alone; VACUUM holds the write lock), not
|
||||
cancellable. It checks the disk for VACUUM's copies of the live pages
|
||||
(``disk_plan``), then sweeps, runs VACUUM and truncates the WAL. When
|
||||
another job holds the registry, it does not start and the finished
|
||||
migration says so and names Optimize database.
|
||||
|
||||
Only the database changes: no image or video file is opened.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import sqlite3
|
||||
|
||||
from .jobs import MAINTENANCE_JOB_KIND, JobBusy
|
||||
from .maintenance import (
|
||||
MAINTENANCE_GROUPS,
|
||||
PRE_MIGRATION_LABEL,
|
||||
SafetyBackupFailed,
|
||||
backup_failed_message,
|
||||
_fail_job,
|
||||
_megabytes,
|
||||
disk_plan,
|
||||
short_space_message,
|
||||
take_safety_backup,
|
||||
)
|
||||
|
||||
try:
|
||||
from ...database import workflow_storage
|
||||
from ...database.health import readonly_uri
|
||||
except ImportError:
|
||||
from database import workflow_storage
|
||||
from database.health import readonly_uri
|
||||
|
||||
OPERATION = workflow_storage.OPERATION
|
||||
TITLE = "Store each workflow once"
|
||||
RECLAIM_TITLE = "Reclaim space"
|
||||
REWRITE_FAILED = (
|
||||
"Storing workflows once stopped on an error. The batch in progress was "
|
||||
"rolled back; every record reads as before. Run it again to continue."
|
||||
)
|
||||
RECLAIM_FAILED = "Reclaim space failed; the database is unchanged."
|
||||
ROW_KEPT = "could not be rewritten; it keeps its inline text and reads as before"
|
||||
|
||||
|
||||
NOTHING_LEFT = "Every image record already stores its workflow once"
|
||||
|
||||
|
||||
def _plural(count, one, many=None):
|
||||
return f"{count:,} {one if count == 1 else (many or one + 's')}"
|
||||
|
||||
|
||||
def live_bytes(db_path):
|
||||
"""The pages VACUUM copies: page_count minus the free list (blocking)."""
|
||||
conn = sqlite3.connect(readonly_uri(db_path), uri=True)
|
||||
try:
|
||||
page_size = conn.execute("PRAGMA page_size").fetchone()[0]
|
||||
pages = conn.execute("PRAGMA page_count").fetchone()[0]
|
||||
free = conn.execute("PRAGMA freelist_count").fetchone()[0]
|
||||
return page_size * (pages - free)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
class WorkflowStorageRoutesMixin:
|
||||
"""Mixin for the maintenance routes: the workflow storage jobs."""
|
||||
|
||||
def _start_workflow_storage(self, db, backup):
|
||||
"""Start the migration job (raises JobBusy); the token is the caller's."""
|
||||
return self._jobs.start(
|
||||
MAINTENANCE_JOB_KIND,
|
||||
TITLE,
|
||||
admission="records",
|
||||
cancellable=True,
|
||||
scope={"operation": OPERATION},
|
||||
groups=MAINTENANCE_GROUPS,
|
||||
runner=lambda job: self._run_workflow_storage(job, db, backup),
|
||||
)
|
||||
|
||||
async def _run_workflow_storage(self, job, db, _backup_on):
|
||||
# Always backed up first, whatever the policy: the way back. The first
|
||||
# apply keeps a pre-migration backup; later ones (after a Cancel, a
|
||||
# rerun) take a pooled safety backup, so only one copy is never pruned.
|
||||
counts = await self._run_in_executor(db.workflow_storage_status)
|
||||
if not counts["rows"]:
|
||||
job.log_line(NOTHING_LEFT)
|
||||
job.update(progress=100, status_line=NOTHING_LEFT)
|
||||
return
|
||||
kept = await self._run_in_executor(
|
||||
workflow_storage.pre_migration_backups, db.model.db_path
|
||||
)
|
||||
label = "backup" if kept else PRE_MIGRATION_LABEL
|
||||
backup = await self._backup_before_rewrite(job, db, label, kept)
|
||||
if backup is None:
|
||||
return
|
||||
try:
|
||||
done = await self._rewrite_workflows(job, db, backup)
|
||||
finally:
|
||||
if kept:
|
||||
await self._prune_backup_pool(job, db)
|
||||
if done:
|
||||
job.finish("finished")
|
||||
self._start_reclaim_space(job, db)
|
||||
|
||||
async def _backup_before_rewrite(self, job, db, label, kept):
|
||||
"""The verified backup's status note, or None (the job failed).
|
||||
|
||||
``label`` ``pre-workflow-storage``: kept until the user removes it.
|
||||
``backup``: a pooled safety backup, while the earliest kept
|
||||
pre-migration backup (``kept[0]``) stays. A failure ends the job
|
||||
``backup-failed``, nothing changed.
|
||||
"""
|
||||
noun = "safety backup" if kept else "pre-migration backup"
|
||||
job.update(status_line=f"Taking the {noun}")
|
||||
job.log_line(f"Taking the {noun} of the database")
|
||||
try:
|
||||
name = await self._run_in_executor(take_safety_backup, db.model, label)
|
||||
except SafetyBackupFailed as error:
|
||||
self.logger.error("Backup before the rewrite failed", exc_info=True)
|
||||
_fail_job(job, "backup-failed", backup_failed_message(error))
|
||||
return None
|
||||
if kept:
|
||||
job.log_line(
|
||||
f"Safety backup {name} verified; the pre-migration backup "
|
||||
f"{kept[0]} is kept"
|
||||
)
|
||||
return f"Safety backup {name}"
|
||||
job.log_line(
|
||||
f"Pre-migration backup {name} verified; it is kept until you remove it"
|
||||
)
|
||||
return f"Pre-migration backup {name}"
|
||||
|
||||
async def _rewrite_workflows(self, job, db, backup):
|
||||
"""Rewrite every old-form row; True when every row was looked at."""
|
||||
counts = await self._run_in_executor(db.workflow_storage_status)
|
||||
job.update(total=counts["rows"], processed=0)
|
||||
job.update(status_line="Rewriting image records")
|
||||
job.log_line(f"{_plural(counts['rows'], 'image record')} to rewrite")
|
||||
rows = protected = errors = 0
|
||||
last_id = 0
|
||||
while True:
|
||||
if job.cancelled:
|
||||
return self._stopped(job, rows, backup)
|
||||
try:
|
||||
batch = await self._run_in_executor(
|
||||
db.migrate_workflow_batch,
|
||||
last_id,
|
||||
workflow_storage.BATCH_ROWS,
|
||||
workflow_storage.BATCH_BYTES,
|
||||
)
|
||||
except Exception:
|
||||
self.logger.exception("Storing workflows once failed")
|
||||
_fail_job(job, None, REWRITE_FAILED)
|
||||
return False
|
||||
if batch["done"]:
|
||||
break
|
||||
last_id = batch["last_id"]
|
||||
rows += batch["rows"]
|
||||
protected += batch["protected"]
|
||||
errors += len(batch["errors"])
|
||||
job.record_counted("changed", batch["rows"])
|
||||
job.record_counted("protected", batch["protected"])
|
||||
for image_id in batch["errors"]:
|
||||
job.record(
|
||||
"errors",
|
||||
{"path": "", "detail": f"Image record {image_id}: {ROW_KEPT}"},
|
||||
f"! Image record {image_id}: {ROW_KEPT}",
|
||||
)
|
||||
total = max(job.total, rows)
|
||||
job.update(
|
||||
processed=rows,
|
||||
total=total,
|
||||
progress=100 * rows // total if total else 100,
|
||||
status_line=f"Rewrote {_plural(rows, 'image record')}",
|
||||
)
|
||||
job.publish({"type": "progress", "processed": rows, "total": total})
|
||||
# Leave the write lock free between batches (monitor, run counts).
|
||||
await asyncio.sleep(workflow_storage.BATCH_PAUSE)
|
||||
removed = await self._run_in_executor(db.sweep_workflow_documents)
|
||||
line = f"Rewrote {_plural(rows, 'image record')}"
|
||||
if protected:
|
||||
line += f"; {_plural(protected, 'stored value')} not text, kept as is"
|
||||
if errors:
|
||||
kept = "keeps its" if errors == 1 else "keep their"
|
||||
line += (
|
||||
f"; {_plural(errors, 'image record')} could not be rewritten and "
|
||||
f"{kept} inline text"
|
||||
)
|
||||
job.log_line(line)
|
||||
if removed:
|
||||
job.log_line(f"Removed {_plural(removed, 'unused workflow document')}")
|
||||
job.update(progress=100, status_line=f"{line} · {backup}")
|
||||
return True
|
||||
|
||||
def _stopped(self, job, rows, backup):
|
||||
line = (
|
||||
f"Stopped after {_plural(rows, 'image record')}; "
|
||||
"run it again to continue"
|
||||
)
|
||||
job.log_line(line)
|
||||
job.update(status_line=f"{line} · {backup}")
|
||||
job.finish("cancelled")
|
||||
return False
|
||||
|
||||
def _start_reclaim_space(self, finished, db):
|
||||
"""Start Reclaim space after the migration ``finished``, or say why not."""
|
||||
try:
|
||||
self._jobs.start(
|
||||
MAINTENANCE_JOB_KIND,
|
||||
RECLAIM_TITLE,
|
||||
admission="exclusive",
|
||||
cancellable=False,
|
||||
scope={"operation": OPERATION},
|
||||
groups=MAINTENANCE_GROUPS,
|
||||
runner=lambda job: self._run_reclaim_space(job, db),
|
||||
)
|
||||
except JobBusy as busy:
|
||||
line = (
|
||||
f"Reclaim space did not start: {busy.running.title} is running. "
|
||||
"Run Optimize database to give the space back."
|
||||
)
|
||||
finished.log_line(line)
|
||||
finished.update(status_line=f"{finished.status_line} · {line}")
|
||||
return
|
||||
finished.log_line("Reclaim space started")
|
||||
|
||||
async def _run_reclaim_space(self, job, db):
|
||||
job.update(status_line="Checking free disk space")
|
||||
try:
|
||||
live = await self._run_in_executor(live_bytes, db.model.db_path)
|
||||
_disk, short, _size = await self._run_in_executor(
|
||||
disk_plan,
|
||||
db.model.db_path,
|
||||
backup=False,
|
||||
vacuum=True,
|
||||
vacuum_bytes=live,
|
||||
)
|
||||
except (OSError, sqlite3.Error):
|
||||
self.logger.exception("Could not measure the disk space for Reclaim space")
|
||||
return _fail_job(job, None, RECLAIM_FAILED)
|
||||
if short is not None:
|
||||
return _fail_job(job, "insufficient-space", short_space_message(short))
|
||||
job.update(status_line="Reclaiming space (the database is locked)")
|
||||
job.log_line("Rebuilding the database file (VACUUM)")
|
||||
try:
|
||||
result = await self._run_in_executor(db.reclaim_space)
|
||||
except Exception:
|
||||
self.logger.exception("Reclaim space failed")
|
||||
return _fail_job(job, None, RECLAIM_FAILED)
|
||||
saved = max(0, result["before_bytes"] - result["after_bytes"])
|
||||
line = (
|
||||
f"Reclaimed {_megabytes(saved)} · the database is now "
|
||||
f"{_megabytes(result['after_bytes'])}"
|
||||
)
|
||||
if result["checkpoint_busy"]:
|
||||
line += " (the rest after the next restart)"
|
||||
job.record_counted("changed", 1)
|
||||
job.log_line(line)
|
||||
job.update(progress=100, status_line=line)
|
||||
+7
-29
@@ -36,14 +36,12 @@ 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")
|
||||
@@ -661,14 +659,6 @@ class PromptManagerConfig:
|
||||
DEFAULT_WORKER_THREADS = max(1, (os.cpu_count() or 1) // 2)
|
||||
WORKER_THREADS = DEFAULT_WORKER_THREADS
|
||||
|
||||
# Maintenance: a reviewed apply takes a verified safety backup of the
|
||||
# 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."""
|
||||
@@ -723,8 +713,6 @@ class PromptManagerConfig:
|
||||
"auto_backup_interval": cls.AUTO_BACKUP_INTERVAL,
|
||||
"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(),
|
||||
@@ -734,12 +722,14 @@ class PromptManagerConfig:
|
||||
def get_config_path(cls) -> str:
|
||||
"""Path of the persisted config.json.
|
||||
|
||||
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``.
|
||||
Honours ``PROMPT_MANAGER_CONFIG_PATH`` when set; otherwise the file
|
||||
lives at the repository root next to ``pyproject.toml``.
|
||||
"""
|
||||
return library_location.config_location()[0]
|
||||
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")
|
||||
|
||||
@classmethod
|
||||
def load_from_file(cls, config_path: Optional[str] = None):
|
||||
@@ -858,18 +848,6 @@ class PromptManagerConfig:
|
||||
performance["worker_threads"], cls.WORKER_THREADS
|
||||
)
|
||||
|
||||
maintenance = new_config.get("maintenance", {})
|
||||
if isinstance(maintenance, dict) and isinstance(
|
||||
maintenance.get("safety_backup"), bool
|
||||
):
|
||||
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"])
|
||||
|
||||
+3
-14
@@ -567,21 +567,10 @@ def download_civitai_images(
|
||||
|
||||
|
||||
def get_lora_image_cache_dir() -> Path:
|
||||
"""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
|
||||
|
||||
"""Get the directory used to cache downloaded LoRA example images."""
|
||||
# Store in the extension's own directory
|
||||
ext_root = Path(__file__).resolve().parent.parent
|
||||
cache = Path(
|
||||
library_location.cache_dir("lora_images", extension_root=str(ext_root))
|
||||
)
|
||||
cache = ext_root / "data" / "lora_images"
|
||||
cache.mkdir(parents=True, exist_ok=True)
|
||||
return cache
|
||||
|
||||
|
||||
@@ -3,7 +3,6 @@
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import socket
|
||||
@@ -387,8 +386,6 @@ def create_app(root: Path) -> web.Application:
|
||||
async def reset(request: web.Request) -> web.Response:
|
||||
nonlocal history_fixture
|
||||
request_holds.release()
|
||||
# No preview queued by the last test may land in the new fixture.
|
||||
await api._run_in_executor(api.close_auto_thumbnails)
|
||||
library_editor_commands.reset()
|
||||
data = await request.json() if request.can_read_body else {}
|
||||
if not isinstance(data, dict):
|
||||
@@ -657,18 +654,6 @@ 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(
|
||||
@@ -693,7 +678,6 @@ 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,17 +181,6 @@ 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.
|
||||
@@ -199,7 +188,8 @@ 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 },
|
||||
columns: { ...reference.columns, surfaces: reference.columns.surfaces.map(surface => ({ ...surface,
|
||||
rows: surface.rows.map(row => (rows[row.title] ?? (value => value))(row)) })) },
|
||||
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.' },
|
||||
@@ -242,21 +232,13 @@ 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; 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}`);
|
||||
// Save; D122 the nav label. (The fixture's first three metrics are the concept's own counts.)
|
||||
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 });
|
||||
@@ -921,85 +903,3 @@ 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,13 +19,6 @@ 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();
|
||||
|
||||
@@ -398,20 +398,18 @@ test('preview uses actual Viewer.js and cleans up when closing or leaving Images
|
||||
test('real video media has a usable card and native playback cleans up on close and navigation', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset', { data: { include_video: true } })).toBeOK();
|
||||
const media = await (await request.get('/prompt_manager/images/output?limit=40')).json();
|
||||
// #249: a video with no preview yet still lists its preview URL; the server makes the frame on first view.
|
||||
expect(media.images.find(image => image.filename === 'clip.webm').thumbnail_url).toMatch(/\/thumbnails\/clip_thumb\.jpg\?root=0/);
|
||||
expect(media.images.find(image => image.filename === 'clip.webm').thumbnail_url).toBeNull();
|
||||
await page.goto(workspace);
|
||||
await refine(page, { Media: 'video' });
|
||||
await expect(cards(page)).toHaveCount(1);
|
||||
const card = cards(page).first();
|
||||
await expect(card).toHaveAttribute('data-image-key', '0:clip.webm');
|
||||
await expect(card.locator('img')).toHaveCount(0);
|
||||
await expect(card.getByRole('button', { name: 'Open clip.webm', exact: true })).toBeVisible();
|
||||
// The card shows the preview frame, never the video itself; D11 names the video's own extension.
|
||||
await expect(card.locator('video')).toHaveCount(0);
|
||||
await expect(card.locator('.media-art > .image-open > img')).toHaveCount(1);
|
||||
// Legacy video art inside the concept card; D11 names the video's own extension.
|
||||
await expect(card.locator('.media-art > .image-open > video')).toHaveCount(1);
|
||||
await expect(card.locator('.media-info p')).toHaveText(/(?:^|· )WEBM$/);
|
||||
await expect.poll(() => card.locator('img').evaluate(node => node.complete && node.naturalHeight > 0)).toBe(true);
|
||||
expect(await card.locator('img').evaluate(node => node.getBoundingClientRect().height)).toBeGreaterThan(20);
|
||||
expect(await card.locator('video').evaluate(node => node.getBoundingClientRect().height)).toBeGreaterThan(20);
|
||||
await card.getByRole('button', { name: 'Open clip.webm', exact: true }).click();
|
||||
const video = page.locator('dialog[open] video');
|
||||
await expect(video).toBeVisible();
|
||||
@@ -715,45 +713,16 @@ test('selected layout stays blue with readable text on hover in either theme', a
|
||||
}
|
||||
});
|
||||
|
||||
test('refresh releases the previous inspector cover source', async ({ page, request }) => {
|
||||
test('refresh releases the previous inspector video source', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset', { data: { include_video: true } })).toBeOK();
|
||||
await page.goto(workspace);
|
||||
const card = cards(page).filter({ hasText: 'clip.webm' });
|
||||
await card.getByRole('button', { name: 'Inspect clip.webm', exact: true }).click();
|
||||
// #249: a video's cover is its preview frame, not the video.
|
||||
const cover = details(page).locator('img.detail-cover');
|
||||
await expect(cover).toHaveAttribute('src', /\/thumbnails\/clip_thumb\.jpg\?/);
|
||||
await expect(details(page).locator('video')).toHaveCount(0);
|
||||
await cover.evaluate(element => { window.__oldInspectorCover = element; });
|
||||
await rescan(page);
|
||||
await expect.poll(() => page.evaluate(() => !window.__oldInspectorCover.isConnected)).toBe(true);
|
||||
expect(await page.evaluate(() => window.__oldInspectorCover.getAttribute('src'))).toBeNull();
|
||||
});
|
||||
|
||||
/** A listing that names no preview for its videos (the fallback a card and the cover keep for one). */
|
||||
async function videosWithoutPreview(page) {
|
||||
await page.route(url => url.pathname === '/prompt_manager/images/output', async route => {
|
||||
const response = await route.fetch();
|
||||
const body = await response.json();
|
||||
if (Array.isArray(body.images)) body.images = body.images.map(image => image.is_video ? { ...image, thumbnail_url: null } : image);
|
||||
await route.fulfill({ response, json: body });
|
||||
});
|
||||
}
|
||||
|
||||
test('a video listed without a preview keeps its own element, released on refresh', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset', { data: { include_video: true } })).toBeOK();
|
||||
await videosWithoutPreview(page);
|
||||
await page.goto(workspace);
|
||||
const card = cards(page).filter({ hasText: 'clip.webm' });
|
||||
await expect(card.locator('.media-art > .image-open > video')).toHaveCount(1);
|
||||
await expect(card.locator('img')).toHaveCount(0);
|
||||
await card.getByRole('button', { name: 'Inspect clip.webm', exact: true }).click();
|
||||
await expect(details(page).locator('video')).toBeVisible();
|
||||
await details(page).locator('video').evaluate(element => { window.__oldInspectorVideo = element; });
|
||||
await card.locator('video').evaluate(element => { window.__oldCardVideo = element; });
|
||||
await rescan(page);
|
||||
await expect.poll(() => page.evaluate(() => !window.__oldInspectorVideo.isConnected && !window.__oldCardVideo.isConnected)).toBe(true);
|
||||
expect(await page.evaluate(() => [window.__oldInspectorVideo.getAttribute('src'), window.__oldCardVideo.getAttribute('src')])).toEqual([null, null]);
|
||||
await expect.poll(() => page.evaluate(() => !window.__oldInspectorVideo.isConnected)).toBe(true);
|
||||
expect(await page.evaluate(() => window.__oldInspectorVideo.getAttribute('src'))).toBeNull();
|
||||
});
|
||||
|
||||
test('Import Rescan rescans cached files and preserves the collection when rescan fails', async ({ page, request }) => {
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
// Needs review (#249 PR F, plan Tasks F1 and F2, Mock inventory M4): the SAVED VIEWS row and its count, the Metadata
|
||||
// review page over the real review routes, and Activity's Review metadata. Items are recorded by the real recorder
|
||||
// (tests/browser/workspace_review_fixture.py) through the scan's recorder (only scans add items); the conflict item's file was linked
|
||||
// (tests/browser/workspace_review_fixture.py) through the monitor's entry point; the conflict item's file was linked
|
||||
// to another prompt after it was queued. An open page re-reads its queue from the review change stream (F2).
|
||||
import { startCoverage, saveCoverage } from '../../tools/vendor/browser-coverage.mjs';
|
||||
import { test, expect } from './fixtures.mjs';
|
||||
@@ -28,17 +28,6 @@ async function pending(request) {
|
||||
}
|
||||
const links = async request => (await (await request.get('/__test__/review/links')).json()).links;
|
||||
const streams = async request => (await (await request.get('/__test__/review/streams')).json());
|
||||
/**
|
||||
* Wait until the server has dropped every stream a page closed. It notices a closed stream only when a write to it
|
||||
* fails, and the close can reach it after any one write (WebKit under load): each poll sends one more change notice,
|
||||
* which changes nothing in the queue, until the count is 0.
|
||||
*/
|
||||
async function closedStreamsDropped(request) {
|
||||
await expect.poll(async () => {
|
||||
await expect(await request.post('/__test__/review/notify')).toBeOK();
|
||||
return (await streams(request)).streams;
|
||||
}).toBe(0);
|
||||
}
|
||||
const apply = page => panel(page).getByRole('button', { name: 'Apply association', exact: true });
|
||||
const warning = page => panel(page).locator('.warning-note');
|
||||
/** The page with its first item's sampler paths loaded. */
|
||||
@@ -329,13 +318,13 @@ test('keyboard focus labels and touch targets', async ({ page, request }) => {
|
||||
await expect(page).toHaveURL(/#\/images$/);
|
||||
});
|
||||
|
||||
test('an open page shows items a scan adds and keeps the open choice', async ({ page, request }) => {
|
||||
test('an open page shows items the monitor adds and keeps the open choice', async ({ page, request }) => {
|
||||
await seed(request, 1);
|
||||
await open(page);
|
||||
await expect.poll(async () => (await streams(request)).streams).toBe(1);
|
||||
await context(page).selectOption({ label: OPTIONS[1] });
|
||||
await expect(reviewRow(page)).toHaveText('Needs review1');
|
||||
await expect(await request.post('/__test__/review/add', { data: { kind: 'scan' } })).toBeOK();
|
||||
await expect(await request.post('/__test__/review/add', { data: { kind: 'monitor' } })).toBeOK();
|
||||
await expect(cards(page)).toHaveCount(2);
|
||||
await expect(main(page).locator('.notice h2')).toHaveText('2 images need a decision');
|
||||
await expect(reviewRow(page)).toHaveText('Needs review2');
|
||||
@@ -392,7 +381,7 @@ test('a visit reads its queue once, and a change reads it once more', async ({ p
|
||||
const reads = reviewReads(page);
|
||||
await open(page);
|
||||
await expect.poll(async () => (await streams(request)).streams).toBe(1);
|
||||
await expect(await request.post('/__test__/review/add', { data: { kind: 'scan' } })).toBeOK();
|
||||
await expect(await request.post('/__test__/review/add', { data: { kind: 'monitor' } })).toBeOK();
|
||||
await expect(cards(page)).toHaveCount(2);
|
||||
await expect(reviewRow(page)).toHaveText('Needs review2');
|
||||
// The page's first read and the change's re-read; the shell's count read at start, and the row follows the page.
|
||||
@@ -434,7 +423,7 @@ test('a stream the browser gives up on is opened again, and the page says it is
|
||||
// The page opens the stream again (after 3 s) and is live once it sends ready.
|
||||
await expect.poll(async () => (await streams(request)).streams, { timeout: 15000 }).toBe(1);
|
||||
await expect(live(page)).toBeHidden();
|
||||
await expect(await request.post('/__test__/review/add', { data: { kind: 'scan' } })).toBeOK();
|
||||
await expect(await request.post('/__test__/review/add', { data: { kind: 'monitor' } })).toBeOK();
|
||||
await expect(cards(page)).toHaveCount(2);
|
||||
expect(attempts).toBe(2);
|
||||
});
|
||||
@@ -455,8 +444,9 @@ test('a hidden tab closes its stream and a visible one reads only what changed m
|
||||
await open(page);
|
||||
await expect.poll(async () => (await streams(request)).streams).toBe(1);
|
||||
await setVisible(page, false);
|
||||
await expect(await request.post('/__test__/review/add', { data: { kind: 'scan' } })).toBeOK();
|
||||
await closedStreamsDropped(request);
|
||||
// The server drops a closed stream at its next write: this monitor addition is one.
|
||||
await expect(await request.post('/__test__/review/add', { data: { kind: 'monitor' } })).toBeOK();
|
||||
await expect.poll(async () => (await streams(request)).streams).toBe(0);
|
||||
await expect(cards(page)).toHaveCount(1);
|
||||
expect(reads.filter(search => search === '?limit=40')).toHaveLength(1);
|
||||
// Visible again: the reopened stream's ready names a newer version, so the page reads once.
|
||||
@@ -469,30 +459,6 @@ test('a hidden tab closes its stream and a visible one reads only what changed m
|
||||
expect(reads.filter(search => search === '?limit=40')).toHaveLength(2);
|
||||
});
|
||||
|
||||
test('a stream whose close reaches the server after a change is dropped at a later write', async ({ page, request }) => {
|
||||
// The CI ordering made certain: the hidden page's close arrives only after the scan addition's change was
|
||||
// written to its stream, so that write succeeds and a later one finds the stream closed.
|
||||
await seed(request, 1);
|
||||
await open(page);
|
||||
await expect.poll(async () => (await streams(request)).streams).toBe(1);
|
||||
await page.evaluate(() => {
|
||||
const close = EventSource.prototype.close;
|
||||
window.lateChanges = 0;
|
||||
window.heldCloses = [];
|
||||
EventSource.prototype.close = function () {
|
||||
this.addEventListener('change', () => { window.lateChanges += 1; });
|
||||
window.heldCloses.push(() => close.call(this));
|
||||
};
|
||||
});
|
||||
await setVisible(page, false);
|
||||
await expect(await request.post('/__test__/review/add', { data: { kind: 'scan' } })).toBeOK();
|
||||
await expect.poll(() => page.evaluate(() => window.lateChanges)).toBe(1);
|
||||
await page.evaluate(() => window.heldCloses.splice(0).forEach(close => close()));
|
||||
expect((await streams(request)).streams).toBe(1);
|
||||
await closedStreamsDropped(request);
|
||||
await expect(cards(page)).toHaveCount(1);
|
||||
});
|
||||
|
||||
test('a re-read refreshes a card whose media or target changed', async ({ page, request }) => {
|
||||
// M6 and M10: a root came or went between reads; the card follows, and a known thumbnail is its media.
|
||||
await seed(request, 2);
|
||||
@@ -510,7 +476,7 @@ test('a re-read refreshes a card whose media or target changed', async ({ page,
|
||||
const second = cards(page).nth(1);
|
||||
await expect(second.getByRole('button', { name: 'Open review_0001.png', exact: true })).not.toHaveAttribute('aria-disabled', 'true');
|
||||
rewrite = true;
|
||||
await expect(await request.post('/__test__/review/add', { data: { kind: 'scan' } })).toBeOK();
|
||||
await expect(await request.post('/__test__/review/add', { data: { kind: 'monitor' } })).toBeOK();
|
||||
await expect(cards(page)).toHaveCount(3);
|
||||
await expect(cards(page).first().locator('img')).toHaveAttribute('src', /&thumb=1$/);
|
||||
await expect(second.getByRole('button', { name: 'Open review_0001.png', exact: true })).toHaveAttribute('aria-disabled', 'true');
|
||||
@@ -556,7 +522,7 @@ test('a live re-read keeps the loaded depth and the scroll position', async ({ p
|
||||
const before = await scrollTop();
|
||||
expect(before).toBeGreaterThan(0);
|
||||
const reads = reviewReads(page);
|
||||
await expect(await request.post('/__test__/review/add', { data: { kind: 'scan' } })).toBeOK();
|
||||
await expect(await request.post('/__test__/review/add', { data: { kind: 'monitor' } })).toBeOK();
|
||||
await expect(main(page).locator('.notice h2')).toHaveText('46 images need a decision');
|
||||
await expect(cards(page)).toHaveCount(46);
|
||||
expect(await scrollTop()).toBe(before);
|
||||
@@ -583,6 +549,8 @@ test('teardown releases captures listeners and requests', async ({ page, request
|
||||
await expect(main(page).getByRole('heading', { level: 1 })).toHaveText('Tags & categories');
|
||||
await expect.poll(async () => (await streams(request)).captures).toBe(0);
|
||||
await route.abort().catch(() => {});
|
||||
await closedStreamsDropped(request);
|
||||
// The server drops a closed stream at its next write: a monitor addition is one.
|
||||
await expect(await request.post('/__test__/review/add', { data: { kind: 'monitor' } })).toBeOK();
|
||||
await expect.poll(async () => (await streams(request)).streams).toBe(0);
|
||||
await expect.poll(() => inFlight.size).toBe(0);
|
||||
});
|
||||
|
||||
@@ -206,9 +206,10 @@ 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, its row saying it is the default.
|
||||
// Saved: ComfyUI's output folder is detected and listed without Remove.
|
||||
await expect.poll(rowsOf.bind(null, page)).toEqual([
|
||||
['ComfyUI output', 'ComfyUI/output · manual scan · default', ['Scan folder']],
|
||||
['ComfyUI output', 'output · manual scan', ['Scan folder']],
|
||||
['p', 'small', 'No source is configured; ComfyUI’s output folder is used.'],
|
||||
]);
|
||||
const { sources } = await serverJson(request, SETTINGS_URL);
|
||||
expect([sources.configured, sources.auto_detected]).toEqual([[], true]);
|
||||
|
||||
@@ -68,8 +68,6 @@ 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([
|
||||
@@ -80,95 +78,6 @@ 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);
|
||||
|
||||
@@ -602,10 +602,6 @@ test('a save under a Tags filter re-reads the listing: the selection stays while
|
||||
|
||||
test('a save under a search keeps the scroll position and the loaded range', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/tagging/reset', { data: { extra_images: 140 } })).toBeOK();
|
||||
// Listing pages still in flight: one that lands after the scroll below would grow the loaded range.
|
||||
const inFlight = new Set();
|
||||
page.on('request', item => { if (item.url().includes('/images/output?')) inFlight.add(item); });
|
||||
for (const done of ['requestfinished', 'requestfailed']) page.on(done, item => inFlight.delete(item));
|
||||
await page.goto(images);
|
||||
const searched = listingOf(page, query => query.get('q') === 'bulk');
|
||||
await page.getByRole('searchbox', { name: 'Search workspace', exact: true }).fill('bulk');
|
||||
@@ -619,9 +615,6 @@ test('a save under a search keeps the scroll position and the loaded range', asy
|
||||
await expect.poll(async () => { await continuation.scrollIntoViewIfNeeded(); return loadedCount(); }).toBeGreaterThanOrEqual(80);
|
||||
await page.evaluate(() => window.scrollTo(0, Math.round(document.documentElement.scrollHeight / 3)));
|
||||
await expect(continuation).not.toBeInViewport();
|
||||
// A load the continuation scheduled for the next frame has started by now; wait for every page to land.
|
||||
await page.evaluate(() => new Promise(done => requestAnimationFrame(() => requestAnimationFrame(done))));
|
||||
await expect.poll(() => inFlight.size).toBe(0);
|
||||
const loaded = await continuation.textContent();
|
||||
await expect.poll(() => continuation.textContent()).toBe(loaded);
|
||||
const y = await page.evaluate(() => window.scrollY);
|
||||
|
||||
@@ -1,135 +0,0 @@
|
||||
import { startCoverage, saveCoverage } from '../../tools/vendor/browser-coverage.mjs';
|
||||
import { test, expect } from './fixtures.mjs';
|
||||
|
||||
/**
|
||||
* #249: cards and strips load previews, never the full file. The approved mock has no repeated thumbnail prompt
|
||||
* (docs/concepts/workspace/README.md), so a preview that was not made yet is made by the server on first view:
|
||||
* the listing still names the `_thumb` URL. Only the media viewer's main image (and downloads) load the original.
|
||||
*/
|
||||
const images = '/prompt_manager/workspace#/images';
|
||||
const library = '/prompt_manager/workspace#/library';
|
||||
const cards = page => page.locator('article[data-image-key]');
|
||||
const details = page => page.getByRole('complementary', { name: 'Selected image details', exact: true });
|
||||
const SERVE = '/prompt_manager/images/serve/';
|
||||
|
||||
/** Every media file the page asks the server for, as `{path, preview}` (path relative to its root). */
|
||||
function recordMedia(page) {
|
||||
const seen = [];
|
||||
page.on('request', request => {
|
||||
const url = new URL(request.url());
|
||||
if (!url.pathname.startsWith(SERVE)) return;
|
||||
const path = decodeURIComponent(url.pathname.slice(SERVE.length));
|
||||
seen.push({ path, preview: /^thumbnails\/.+_thumb\.[a-z0-9]+$/.test(path) });
|
||||
});
|
||||
return seen;
|
||||
}
|
||||
/** Every card image has loaded (or failed): what it asked for is all in `seen`. */
|
||||
const settled = page => page.evaluate(() => Promise.all([...document.querySelectorAll('article img')].map(image => image.decode().catch(() => {}))));
|
||||
const originals = seen => seen.filter(item => !item.preview).map(item => item.path);
|
||||
|
||||
test.beforeEach(async ({ page, browserName }) => {
|
||||
await startCoverage(page, browserName);
|
||||
page.thumbErrors = [];
|
||||
page.on('pageerror', error => page.thumbErrors.push(error.message));
|
||||
});
|
||||
test.afterEach(async ({ page, browserName }) => {
|
||||
await saveCoverage(page, browserName);
|
||||
expect(page.thumbErrors).toEqual([]);
|
||||
});
|
||||
|
||||
test('a page of the Images grid requests previews, never originals, in every layout and size', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset')).toBeOK();
|
||||
// Newest outputs with no preview yet: the live finding (844 such files listed with a null thumbnail).
|
||||
await expect(await request.post('/__test__/system/unthumbnailed', { data: { count: 6 } })).toBeOK();
|
||||
const listing = await (await request.get('/prompt_manager/images/output?limit=40')).json();
|
||||
const fresh = listing.images.filter(image => image.path.startsWith('unthumbnailed/'));
|
||||
expect(fresh).toHaveLength(6);
|
||||
for (const image of fresh) expect(image.thumbnail_url).toMatch(/\/thumbnails\/unthumbnailed\/plain_\d{3}_thumb\.png\?root=0/);
|
||||
const seen = recordMedia(page);
|
||||
await page.goto(images);
|
||||
await expect(cards(page)).toHaveCount(8);
|
||||
await settled(page);
|
||||
expect(originals(seen)).toEqual([]);
|
||||
expect(seen.filter(item => item.path.startsWith('thumbnails/unthumbnailed/'))).toHaveLength(6);
|
||||
// The server made them on that first view.
|
||||
await expect.poll(async () => (await (await request.get('/prompt_manager/images/thumbnails/status')).json()).pending).toBe(0);
|
||||
for (const layout of ['List', 'Masonry', 'Grid']) {
|
||||
await page.getByRole('button', { name: layout, exact: true }).click();
|
||||
await expect(page.getByRole('button', { name: layout, exact: true })).toHaveAttribute('aria-pressed', 'true');
|
||||
for (const density of ['large', 'compact', 'comfortable']) {
|
||||
await page.getByRole('combobox', { name: 'Thumbnail size', exact: true }).selectOption(density);
|
||||
await expect(cards(page)).toHaveCount(8);
|
||||
await settled(page);
|
||||
}
|
||||
}
|
||||
expect(originals(seen)).toEqual([]);
|
||||
});
|
||||
|
||||
test('the inspector cover is the preview and only the media viewer loads the original', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset')).toBeOK();
|
||||
await expect(await request.post('/__test__/system/unthumbnailed', { data: { count: 1 } })).toBeOK();
|
||||
const seen = recordMedia(page);
|
||||
await page.goto(images);
|
||||
const card = cards(page).filter({ hasText: 'plain_001.png' });
|
||||
await card.getByRole('button', { name: 'Inspect plain_001.png', exact: true }).click();
|
||||
const cover = details(page).locator('img.detail-cover');
|
||||
await expect(cover).toHaveAttribute('src', /\/thumbnails\/unthumbnailed\/plain_001_thumb\.png\?/);
|
||||
await settled(page);
|
||||
await cover.evaluate(image => image.decode().catch(() => {}));
|
||||
expect(originals(seen)).toEqual([]);
|
||||
await card.getByRole('button', { name: 'Open plain_001.png', exact: true }).click();
|
||||
await expect(page.locator('dialog#viewer-dialog')).toBeVisible();
|
||||
await expect.poll(() => originals(seen)).toEqual(['unthumbnailed/plain_001.png']);
|
||||
await page.getByRole('button', { name: 'Close image viewer', exact: true }).click();
|
||||
});
|
||||
|
||||
test('a video card shows its preview frame instead of loading the video', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset', { data: { include_video: true } })).toBeOK();
|
||||
const seen = recordMedia(page);
|
||||
await page.goto(images);
|
||||
const card = cards(page).filter({ hasText: 'clip.webm' });
|
||||
await expect(card.locator('.media-art > .image-open > img')).toHaveAttribute('src', /\/thumbnails\/clip_thumb\.jpg\?/);
|
||||
await expect(card.locator('video')).toHaveCount(0);
|
||||
await settled(page);
|
||||
expect(originals(seen)).toEqual([]);
|
||||
});
|
||||
|
||||
test('Library cards and the shared gallery request previews only', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset', { data: { native_library: true } })).toBeOK();
|
||||
const seen = recordMedia(page);
|
||||
await page.goto(library);
|
||||
const cat = page.locator('article[data-prompt-id]').filter({ has: page.getByText('Cat portrait', { exact: true }) });
|
||||
await expect(cat.locator('.card-art img')).not.toHaveCount(0);
|
||||
await settled(page);
|
||||
expect(seen.length).toBeGreaterThan(0);
|
||||
expect(originals(seen)).toEqual([]);
|
||||
});
|
||||
|
||||
test('a video whose preview cannot be loaded swaps back to its own element', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset', { data: { include_video: true } })).toBeOK();
|
||||
await page.route(url => url.pathname.endsWith('/thumbnails/clip_thumb.jpg'), route => route.fulfill({ status: 404, body: '' }));
|
||||
await page.goto(images);
|
||||
const card = cards(page).filter({ hasText: 'clip.webm' });
|
||||
await expect(card.locator('.media-art > .image-open > video')).toHaveCount(1);
|
||||
await expect(card.locator('.media-art > .image-open > img')).toHaveCount(0);
|
||||
await card.getByRole('button', { name: 'Inspect clip.webm', exact: true }).click();
|
||||
await expect(details(page).locator('video.detail-cover')).toHaveCount(1);
|
||||
await expect(details(page).locator('img.detail-cover')).toHaveCount(0);
|
||||
});
|
||||
|
||||
test('an image listed without a preview (a nested thumbnails folder) shows its original', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset')).toBeOK();
|
||||
// M10: such files are listed with a null thumbnail_url; the card falls back to the original as before.
|
||||
await page.route(url => url.pathname === '/prompt_manager/images/output', async route => {
|
||||
const response = await route.fetch();
|
||||
const body = await response.json();
|
||||
if (Array.isArray(body.images)) body.images = body.images.map(image => ({ ...image, thumbnail_url: null }));
|
||||
await route.fulfill({ response, json: body });
|
||||
});
|
||||
const seen = recordMedia(page);
|
||||
await page.goto(images);
|
||||
await expect(cards(page)).toHaveCount(2);
|
||||
await settled(page);
|
||||
expect(seen.length).toBeGreaterThan(0);
|
||||
expect(seen.every(item => !item.preview)).toBe(true);
|
||||
});
|
||||
@@ -5,9 +5,8 @@ encoders and one shared negative, so the resolver finds it ``ambiguous``. Its
|
||||
loader node carries ComfyUI's ``Infinity`` stamp, which the legacy pages'
|
||||
parser (the scan's fallback) cannot read, so nothing links the file and it is
|
||||
a Needs review case (#284: only ambiguous, unlinked files are queued). It is
|
||||
recorded through the scan's recorder (``_record_review_sync`` with
|
||||
``inspection=True``; only scans add items). No path or SQL comes from a
|
||||
request.
|
||||
recorded through the producers' entry point (``record_review_file``). No path
|
||||
or SQL comes from a request.
|
||||
|
||||
``POST /__test__/review/reset`` also takes the special items F2 needs, queued
|
||||
after the plain ones in this order:
|
||||
@@ -29,14 +28,12 @@ its first negative, the one the concept's panel shows), so the panel's text
|
||||
wraps as the concept's does.
|
||||
|
||||
``POST /__test__/review/add {kind}`` adds one file while a page is open:
|
||||
``scan`` records a new ambiguous, unlinked output from a worker thread
|
||||
through the scan's recorder (as a running scan's batch does; the monitor
|
||||
never adds items), ``file`` only writes one (a metadata scan records it). ``GET
|
||||
``monitor`` records a new ambiguous, unlinked output from a worker thread
|
||||
through ``record_review_file`` (the monitor's own entry point; the monitor
|
||||
never queues an output it linked itself), ``file`` only writes one (a
|
||||
metadata scan records it). ``GET
|
||||
/__test__/review/links`` answers each review file's prompt and pair links, and
|
||||
``GET /__test__/review/streams`` the open change streams and captures.
|
||||
``POST /__test__/review/notify`` sends one queue change notice without
|
||||
changing the queue: each open stream writes a ``change``, and a stream whose
|
||||
page closed is dropped at such a write.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -142,16 +139,8 @@ def _png(
|
||||
image.save(path, pnginfo=info)
|
||||
|
||||
|
||||
def _scan_record(api: PromptManagerAPI, path: Path) -> str:
|
||||
"""Record ``path`` as the scan's recorder does."""
|
||||
from py.api.admin import _scan_roots
|
||||
|
||||
roots = _scan_roots(api._get_all_output_dirs())
|
||||
return api._record_review_sync(Path(str(path)), roots, inspection=True)
|
||||
|
||||
|
||||
def _record(api: PromptManagerAPI, path: Path, expected: str) -> None:
|
||||
outcome = _scan_record(api, path)
|
||||
outcome = api.record_review_file(str(path))
|
||||
if outcome != expected:
|
||||
raise RuntimeError(f"{path.name} was {outcome}, not {expected}")
|
||||
|
||||
@@ -161,19 +150,19 @@ def _stamp(path: Path, age: float) -> None:
|
||||
os.utime(path, (stamp, stamp))
|
||||
|
||||
|
||||
class _ScanRecorder:
|
||||
"""The scan's recorder over this server's database (one per server)."""
|
||||
class _Monitor:
|
||||
"""The monitor's recorder over this server's database (one per server)."""
|
||||
|
||||
def __init__(self, api: PromptManagerAPI) -> None:
|
||||
self.api = api
|
||||
self.lock = threading.Lock()
|
||||
|
||||
def new_output(self, path: Path) -> None:
|
||||
"""The scan's Needs review recording of a new output, from a worker thread."""
|
||||
"""The monitor's Needs review pass for a new output, from a worker thread."""
|
||||
outcome = []
|
||||
with self.lock:
|
||||
worker = threading.Thread(
|
||||
target=lambda: outcome.append(_scan_record(self.api, path))
|
||||
target=lambda: outcome.append(self.api.record_review_file(str(path)))
|
||||
)
|
||||
worker.start()
|
||||
worker.join()
|
||||
@@ -181,9 +170,9 @@ class _ScanRecorder:
|
||||
raise RuntimeError(f"{path.name} was {outcome}, not recorded")
|
||||
|
||||
|
||||
def _scanned_output(api: PromptManagerAPI, recorder: _ScanRecorder, path: Path) -> None:
|
||||
def _monitored_output(api: PromptManagerAPI, monitor: _Monitor, path: Path) -> None:
|
||||
_png(path, (150, 60, 40))
|
||||
recorder.new_output(path)
|
||||
monitor.new_output(path)
|
||||
|
||||
|
||||
def _linked_after_queueing(api: PromptManagerAPI, path: Path) -> None:
|
||||
@@ -216,7 +205,7 @@ def seed_review(
|
||||
root: Path,
|
||||
items: int,
|
||||
special: tuple = (),
|
||||
recorder: _ScanRecorder | None = None,
|
||||
monitor: _Monitor | None = None,
|
||||
concept: bool = False,
|
||||
) -> dict:
|
||||
"""Replace this server's disposable rows and files with ``items`` queued files."""
|
||||
@@ -269,9 +258,9 @@ def seed_review(
|
||||
|
||||
|
||||
def add_review_file(
|
||||
api: PromptManagerAPI, root: Path, kind: str, recorder: _ScanRecorder
|
||||
api: PromptManagerAPI, root: Path, kind: str, monitor: _Monitor
|
||||
) -> dict:
|
||||
"""One more ambiguous output: recorded as a scan does, or only written."""
|
||||
"""One more ambiguous output: through the monitor, or only written."""
|
||||
output = root / "output"
|
||||
number = (
|
||||
len(list((output / "added").glob("*.png")))
|
||||
@@ -279,8 +268,8 @@ def add_review_file(
|
||||
else 0
|
||||
)
|
||||
relative = f"added/added_{number:04}.png"
|
||||
if kind == "scan":
|
||||
_scanned_output(api, recorder, output / relative)
|
||||
if kind == "monitor":
|
||||
_monitored_output(api, monitor, output / relative)
|
||||
else:
|
||||
_png(output / relative, (200, 200, 40))
|
||||
api.invalidate_gallery_cache()
|
||||
@@ -309,12 +298,12 @@ def register_review_commands(
|
||||
routes: web.RouteTableDef, api: PromptManagerAPI, root: Path
|
||||
) -> None:
|
||||
"""Register the review fixture's commands; no path or SQL comes from a request."""
|
||||
recorders = []
|
||||
monitors = []
|
||||
|
||||
def recorder() -> _ScanRecorder:
|
||||
if not recorders:
|
||||
recorders.append(_ScanRecorder(api))
|
||||
return recorders[0]
|
||||
def monitor() -> _Monitor:
|
||||
if not monitors:
|
||||
monitors.append(_Monitor(api))
|
||||
return monitors[0]
|
||||
|
||||
@routes.post("/__test__/review/reset")
|
||||
async def reset(request: web.Request) -> web.Response:
|
||||
@@ -342,7 +331,7 @@ def register_review_commands(
|
||||
root,
|
||||
items,
|
||||
special,
|
||||
recorder(),
|
||||
monitor(),
|
||||
data.get("concept", False),
|
||||
)
|
||||
return web.json_response({"success": True, **manifest})
|
||||
@@ -352,22 +341,13 @@ def register_review_commands(
|
||||
data = await request.json() if request.can_read_body else None
|
||||
if request.query or not isinstance(data, dict) or set(data) != {"kind"}:
|
||||
raise web.HTTPBadRequest()
|
||||
if data["kind"] not in ("scan", "file"):
|
||||
if data["kind"] not in ("monitor", "file"):
|
||||
raise web.HTTPBadRequest()
|
||||
added = await api._run_in_executor(
|
||||
add_review_file, api, root, data["kind"], recorder()
|
||||
add_review_file, api, root, data["kind"], monitor()
|
||||
)
|
||||
return web.json_response({"success": True, **added})
|
||||
|
||||
@routes.post("/__test__/review/notify")
|
||||
async def notify(request: web.Request) -> web.Response:
|
||||
if request.query or request.can_read_body:
|
||||
raise web.HTTPBadRequest()
|
||||
from database.review_changes import notify_review_change
|
||||
|
||||
notify_review_change(api.db.model.db_path)
|
||||
return web.json_response({"success": True})
|
||||
|
||||
@routes.get("/__test__/review/links")
|
||||
async def links(request: web.Request) -> web.Response:
|
||||
if request.query:
|
||||
|
||||
@@ -23,12 +23,6 @@ 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``).
|
||||
"""
|
||||
@@ -60,12 +54,7 @@ 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,
|
||||
@@ -103,9 +92,6 @@ 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:
|
||||
@@ -201,43 +187,6 @@ 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
|
||||
|
||||
|
||||
|
||||
@@ -1,73 +0,0 @@
|
||||
// #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(); }
|
||||
}
|
||||
});
|
||||
@@ -39,4 +39,3 @@ test('stale inspection echoes use the canonical recovery vocabulary and identity
|
||||
});
|
||||
}
|
||||
});
|
||||
test('the recorded output inspector cover is its preview, not the original (#249)',()=>{const a=adapter(),target=media('2',1);const record=a.toRecord(envelope(target),target,extractGenerationFacts);assert.deepEqual(record.preview,{url:target.thumbnail_url});assert.notEqual(record.preview.url,record.download.url);});
|
||||
|
||||
@@ -895,30 +895,3 @@ test('cache off never rereads a selection whose source changed, and refresh neve
|
||||
assert.equal(deferred.state.deferred, true);
|
||||
deferred.dispose();
|
||||
});
|
||||
|
||||
test('record carries the listing thumbnail as its preview, never another file (#249)', () => {
|
||||
const toRecord = contract('toInspectionRecord'), entry = image();
|
||||
const thumb = `/prompt_manager/images/serve/thumbnails/folder/a%20%23%25_thumb.png?root=0&root_scope=${scopeA}`;
|
||||
assert.deepEqual(toRecord(envelope(entry), { ...entry, thumbnail_url: thumb }, extractGenerationFacts).preview, { url: thumb });
|
||||
const versioned = `/prompt_manager/images/serve/thumbnails/folder/a%20%23%25_thumb.png?v=12&root=0&root_scope=${scopeA}`;
|
||||
assert.deepEqual(toRecord(envelope(entry), { ...entry, thumbnail_url: versioned }, extractGenerationFacts).preview, { url: versioned });
|
||||
// No thumbnail, or one naming another file, root or origin: no preview (the cover keeps the original).
|
||||
for (const thumbnail_url of [undefined, null, '', thumb.replace('a%20%23%25_thumb', 'b_thumb'), thumb.replace('root=0', 'root=1'),
|
||||
thumb.replace(scopeA, scopeB), `https://other.test${thumb}`, thumb.replace('thumbnails/', ''), `${thumb}#x`]) {
|
||||
assert.equal(toRecord(envelope(entry), { ...entry, thumbnail_url }, extractGenerationFacts).preview, null, String(thumbnail_url));
|
||||
}
|
||||
const video = { root_index: 0, path: 'clip.MP4', filename: 'clip.MP4', url: `/prompt_manager/images/serve/clip.MP4?root=0&root_scope=${scopeA}`, root_scope: scopeA,
|
||||
thumbnail_url: `/prompt_manager/images/serve/thumbnails/clip_thumb.jpg?root=0&root_scope=${scopeA}` };
|
||||
const data = envelope(video); data.file = { ...data.file, media_type: 'video', extension: '.mp4', format: 'MP4' };
|
||||
assert.deepEqual(toRecord(data, video, extractGenerationFacts).preview, { url: video.thumbnail_url });
|
||||
});
|
||||
|
||||
test("names the server quotes (parentheses, '!*') still match their original and preview (#249 M11)", () => {
|
||||
const toRecord = contract('toInspectionRecord');
|
||||
const entry = { root_index: 0, path: "folder/a (1)!'*.png", filename: "a (1)!'*.png", root_scope: scopeA,
|
||||
url: `/prompt_manager/images/serve/folder/a%20%281%29%21%27%2A.png?root=0&root_scope=${scopeA}`,
|
||||
thumbnail_url: `/prompt_manager/images/serve/thumbnails/folder/a%20%281%29%21%27%2A_thumb.png?root=0&root_scope=${scopeA}` };
|
||||
const record = toRecord(envelope(entry), entry, extractGenerationFacts);
|
||||
assert.equal(record.download.url, entry.url);
|
||||
assert.deepEqual(record.preview, { url: entry.thumbnail_url });
|
||||
});
|
||||
|
||||
@@ -1,159 +0,0 @@
|
||||
// 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);
|
||||
});
|
||||
@@ -253,11 +253,3 @@ test('waiting stops when its signal aborts', async () => {
|
||||
gone.abort(new Error('already'));
|
||||
await assert.rejects(waitForJob(fakeActivity(), JOB, { signal: gone.signal }), { message: 'already' });
|
||||
});
|
||||
|
||||
test('the thumbnail row counts previews that could not be made (#249)', async () => {
|
||||
const { thumbnailTitle } = await overviewModule;
|
||||
assert.equal(thumbnailTitle({ pending: 12, failed: 2 }), '12 missing thumbnails · 2 failed');
|
||||
assert.equal(thumbnailTitle({ pending: 1, failed: 1 }), '1 missing thumbnail · 1 failed');
|
||||
assert.equal(thumbnailTitle({ pending: 0, failed: 0 }), 'No missing thumbnails');
|
||||
assert.equal(thumbnailTitle({ pending: 3, failed: 'x' }), '3 missing thumbnails');
|
||||
});
|
||||
|
||||
@@ -72,35 +72,21 @@ 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 and read as the default: "ComfyUI/<path> · <watch> · default"', async () => {
|
||||
test('auto-detected sources are listed without Remove, with the line that says how they were chosen', async () => {
|
||||
const { sourceRows, NOTES } = await sourcesModule;
|
||||
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.
|
||||
const auto = { configured: [], effective: [{ label: 'ComfyUI output', path: 'output', root_index: 0, watch: 'manual' }], auto_detected: true };
|
||||
assert.deepEqual(sourceRows({ sources: auto, saved: [], current: [] }), { rows: [
|
||||
{ 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.
|
||||
{ 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.
|
||||
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.description, row.remove]),
|
||||
[['gone', 'gone · Gallery root does not exist', 'gone'], ['ComfyUI output', 'ComfyUI/output · watching · default', null]]);
|
||||
assert.deepEqual(result.rows.map(row => [row.title, row.remove]), [['gone', 'gone'], ['ComfyUI output', 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);
|
||||
|
||||
@@ -79,8 +79,6 @@ class AdminAPITestCase(AioHTTPTestCase):
|
||||
|
||||
async def tearDownAsync(self):
|
||||
await super().tearDownAsync() # closes the aiohttp test client
|
||||
# The automatic preview worker writes into the temp tree: stop it first.
|
||||
self.api.close_auto_thumbnails()
|
||||
GalleryConfig.MONITORING_DIRECTORIES = self._orig_gallery_dirs
|
||||
PromptManagerConfig.update_config(self._orig_pm_config)
|
||||
for key, value in self._orig_env.items():
|
||||
@@ -2245,12 +2243,7 @@ class TestActivityJobs(AdminAPITestCase):
|
||||
(self.output_dir / name).write_bytes(b"not an image")
|
||||
|
||||
def fake_ingest(
|
||||
media_file,
|
||||
metadata,
|
||||
detailed=False,
|
||||
ambiguous=None,
|
||||
untraced=None,
|
||||
records=None,
|
||||
media_file, metadata, detailed=False, ambiguous=None, untraced=None
|
||||
):
|
||||
key = Path(media_file).relative_to(self.output_dir).as_posix()
|
||||
outcome = names[key]
|
||||
|
||||
@@ -131,8 +131,6 @@ class AutotagAPITestCase(AioHTTPTestCase):
|
||||
|
||||
async def tearDownAsync(self):
|
||||
await super().tearDownAsync() # closes the aiohttp test client
|
||||
# The automatic preview worker writes into the temp tree: stop it first.
|
||||
self.api.close_auto_thumbnails()
|
||||
self._service_patch.stop()
|
||||
if model_gate is not None:
|
||||
leaked = model_gate.holder()
|
||||
|
||||
@@ -47,13 +47,11 @@ def _texts(path):
|
||||
|
||||
|
||||
class BackupRestoreTestCase(AioHTTPTestCase):
|
||||
app_options = {}
|
||||
|
||||
async def get_application(self):
|
||||
self.tmpdir = tempfile.mkdtemp()
|
||||
self.db_path = os.path.join(self.tmpdir, "prompts.db")
|
||||
|
||||
app = web.Application(**self.app_options)
|
||||
app = web.Application()
|
||||
routes = web.RouteTableDef()
|
||||
self.api = PromptManagerAPI(db=PromptDatabase(self.db_path))
|
||||
self.api.add_routes(routes)
|
||||
@@ -62,8 +60,6 @@ class BackupRestoreTestCase(AioHTTPTestCase):
|
||||
|
||||
async def tearDownAsync(self):
|
||||
await super().tearDownAsync() # closes the aiohttp test client
|
||||
# The automatic preview worker writes into the temp tree: stop it first.
|
||||
self.api.close_auto_thumbnails()
|
||||
PromptModel.close_all_instances() # restore may have swapped api.db
|
||||
assert_closed(self, self.tmpdir)
|
||||
for name in os.listdir(self.tmpdir):
|
||||
@@ -331,15 +327,6 @@ class TestRestoreRoute(BackupRestoreTestCase):
|
||||
self.assertIn("too large", data["error"].lower())
|
||||
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
|
||||
|
||||
async def test_restore_cap_keeps_the_legacy_shape(self):
|
||||
self.api.restore_max_bytes = 1024
|
||||
upload = self._write("big.db", b"SQLite format 3\x00" + b"\x00" * 8192)
|
||||
resp = await self._restore(upload)
|
||||
data = await resp.json()
|
||||
self.assertEqual(resp.status, 400)
|
||||
self.assertEqual(data["code"], "too-large")
|
||||
self.assertNotIn(self.tmpdir, data["error"])
|
||||
|
||||
async def test_non_multipart_request_is_rejected(self):
|
||||
resp = await self.client.post("/prompt_manager/restore", data=b"raw bytes")
|
||||
self.assertIn(resp.status, (400, 500))
|
||||
@@ -349,132 +336,3 @@ class TestRestoreRoute(BackupRestoreTestCase):
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
|
||||
MIB = 2**20
|
||||
|
||||
|
||||
class TestLargeRestores(BackupRestoreTestCase):
|
||||
"""A real library's backup restores: room on disk decides, not 100 MB."""
|
||||
|
||||
# As ComfyUI builds its app: client_max_size is --max-upload-size (100
|
||||
# MB by default). The restore streams with read_chunk, which it does not
|
||||
# cap.
|
||||
app_options = {"client_max_size": 100 * 2**20}
|
||||
|
||||
def _large_db_file(self, texts, pad_bytes):
|
||||
path = self._make_db_file(texts, name="large.db")
|
||||
conn = sqlite3.connect(path)
|
||||
try:
|
||||
conn.execute("CREATE TABLE padding(b BLOB)")
|
||||
conn.execute("INSERT INTO padding VALUES (zeroblob(?))", (pad_bytes,))
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
return path
|
||||
|
||||
async def _stream_restore(self, path, chunked=False):
|
||||
with open(path, "rb") as fh:
|
||||
data = fh.read()
|
||||
form = aiohttp.FormData()
|
||||
form.add_field(
|
||||
"database_file",
|
||||
data,
|
||||
filename=os.path.basename(path),
|
||||
content_type="application/octet-stream",
|
||||
)
|
||||
# chunked=False must not be passed: aiohttp 3.14 then frames the body
|
||||
# as chunked while also sending a Content-Length.
|
||||
options = {"chunked": True} if chunked else {}
|
||||
return await self.client.post("/prompt_manager/restore", data=form, **options)
|
||||
|
||||
def _spy_uploads(self):
|
||||
"""Patch the upload's mkstemp; the list collects each temp path made."""
|
||||
created = []
|
||||
real_mkstemp = tempfile.mkstemp
|
||||
|
||||
def spy(*args, **kwargs):
|
||||
result = real_mkstemp(*args, **kwargs)
|
||||
created.append(result[1])
|
||||
return result
|
||||
|
||||
return patch("py.api.admin.tempfile.mkstemp", spy), created
|
||||
|
||||
def _assert_removed(self, created):
|
||||
for path in created:
|
||||
for suffix in ("", "-wal", "-shm", "-journal"):
|
||||
self.assertFalse(os.path.exists(path + suffix), path + suffix)
|
||||
|
||||
async def test_a_150_mb_backup_restores(self):
|
||||
self._save("before")
|
||||
upload = self._large_db_file(["restored"], 150 * MIB)
|
||||
self.assertGreater(os.path.getsize(upload), 150 * MIB)
|
||||
spy, created = self._spy_uploads()
|
||||
with spy:
|
||||
resp = await self._stream_restore(upload)
|
||||
data = await resp.json()
|
||||
self.assertEqual(resp.status, 200, data)
|
||||
self.assertEqual(
|
||||
[p["text"] for p in self.api.db.search_prompts()], ["restored"]
|
||||
)
|
||||
self.assertEqual(len(created), 1)
|
||||
self._assert_removed(created)
|
||||
|
||||
async def test_a_short_disk_is_refused_up_front(self):
|
||||
from py.api import maintenance
|
||||
|
||||
self._save("keep me")
|
||||
upload = self._make_db_file(["restored"])
|
||||
spy, created = self._spy_uploads()
|
||||
with spy, patch.object(maintenance, "_free_bytes", return_value=MIB):
|
||||
resp = await self._restore(upload)
|
||||
data = await resp.json()
|
||||
self.assertEqual(resp.status, 507)
|
||||
self.assertEqual(data["code"], "insufficient-space")
|
||||
self.assertTrue(
|
||||
data["error"].startswith("Not enough free disk space to restore a "),
|
||||
data["error"],
|
||||
)
|
||||
self.assertTrue(data["error"].endswith("Nothing was changed."))
|
||||
self.assertNotIn(self.tmpdir, data["error"])
|
||||
self.assertNotIn(tempfile.gettempdir(), data["error"])
|
||||
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
|
||||
self.assertEqual(created, []) # refused before anything was written
|
||||
|
||||
async def test_a_short_disk_is_refused_while_streaming_without_a_length(self):
|
||||
from py.api import maintenance
|
||||
|
||||
self._save("keep me")
|
||||
upload = self._large_db_file(["restored"], 4 * MIB)
|
||||
spy, created = self._spy_uploads()
|
||||
# Room for 1 MiB only, and no Content-Length to refuse up front.
|
||||
with spy, patch.object(maintenance, "restore_room", return_value=MIB):
|
||||
resp = await self._stream_restore(upload, chunked=True)
|
||||
data = await resp.json()
|
||||
self.assertEqual(resp.status, 507)
|
||||
self.assertEqual(data["code"], "insufficient-space")
|
||||
self.assertTrue(data["error"].endswith("Nothing was changed."))
|
||||
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
|
||||
self.assertEqual(len(created), 1)
|
||||
self._assert_removed(created)
|
||||
|
||||
|
||||
class TestRestoreRoom(unittest.TestCase):
|
||||
def test_room_counts_the_upload_its_staged_copy_and_the_swap(self):
|
||||
from py.api import maintenance
|
||||
|
||||
tmp = tempfile.mkdtemp()
|
||||
try:
|
||||
db_path = os.path.join(tmp, "prompts.db")
|
||||
free = 10 * 2**30
|
||||
with (
|
||||
patch.object(maintenance, "_free_bytes", return_value=free),
|
||||
patch.object(maintenance, "_device", return_value="one disk"),
|
||||
):
|
||||
room = maintenance.restore_room(db_path)
|
||||
# Upload + staged copy + WAL growth on one disk, no database yet.
|
||||
self.assertEqual(room, (free - maintenance.DISK_MARGIN_BYTES) // 3)
|
||||
self.assertEqual(maintenance.size_label(int(1.6 * 2**30)), "1.6 GB")
|
||||
self.assertEqual(maintenance.size_label(120 * 2**20), "120.0 MB")
|
||||
finally:
|
||||
os.rmdir(tmp)
|
||||
|
||||
@@ -384,8 +384,6 @@ class CoreRoutesTestCase(AioHTTPTestCase):
|
||||
|
||||
async def tearDownAsync(self):
|
||||
await super().tearDownAsync() # closes the aiohttp test client
|
||||
# The automatic preview worker writes into the temp tree: stop it first.
|
||||
self.api.close_auto_thumbnails()
|
||||
self.api.db.close_all()
|
||||
assert_closed(self, self.tmpdir)
|
||||
shutil.rmtree(self.tmpdir, ignore_errors=True)
|
||||
|
||||
@@ -1,995 +0,0 @@
|
||||
"""Duplicate image files on the job registry (#249 System S6, Task 17).
|
||||
|
||||
``POST /prompt_manager/jobs/duplicates`` finds content-identical media files
|
||||
as a ``read`` Activity job; ``GET /prompt_manager/maintenance/duplicates/{id}``
|
||||
pages the groups it holds; ``duplicate_files`` previews a removal bound to the
|
||||
reviewed keepers, and its apply job removes only the reviewed, revalidated
|
||||
non-keepers, moving their database links to the kept copy. The legacy gallery
|
||||
routes keep their shapes on the same registry kinds.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
try:
|
||||
from tests.open_handles import open_handles_under
|
||||
from tests.test_api_maintenance import ApplyTestCase, _digest
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from open_handles import open_handles_under
|
||||
from test_api_maintenance import ApplyTestCase, _digest
|
||||
|
||||
from py.api import duplicate_files # noqa: E402
|
||||
from py.config import GalleryConfig, PromptManagerConfig # noqa: E402
|
||||
|
||||
WAIT = 10
|
||||
FIND = "/prompt_manager/jobs/duplicates"
|
||||
GROUPS = "/prompt_manager/maintenance/duplicates/"
|
||||
PREVIEW = "/prompt_manager/maintenance/preview"
|
||||
STREAM = "/prompt_manager/scan_duplicates/stream"
|
||||
STATUS = "/prompt_manager/scan_duplicates/status"
|
||||
LEGACY_DELETE = "/prompt_manager/delete_duplicate_images"
|
||||
|
||||
|
||||
def _sse_payloads(body):
|
||||
return [
|
||||
json.loads(line[len("data: ") :])
|
||||
for line in body.splitlines()
|
||||
if line.startswith("data: ")
|
||||
]
|
||||
|
||||
|
||||
def _sha(path):
|
||||
with open(path, "rb") as handle:
|
||||
return hashlib.sha256(handle.read()).hexdigest()
|
||||
|
||||
|
||||
class DuplicateFilesTestCase(ApplyTestCase):
|
||||
async def get_application(self):
|
||||
app = await super().get_application()
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]
|
||||
return app
|
||||
|
||||
def write(self, name, data, *, folder=None, mtime=None):
|
||||
path = Path(folder or self.output_dir) / name
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_bytes(data)
|
||||
if mtime is not None:
|
||||
os.utime(path, (mtime, mtime))
|
||||
return path
|
||||
|
||||
def second_source(self):
|
||||
"""A second valid source outside ComfyUI (an extra gallery root)."""
|
||||
second = self.outside_dir / "renders"
|
||||
second.mkdir(parents=True, exist_ok=True)
|
||||
os.environ["PROMPT_MANAGER_EXTRA_GALLERY_ROOTS"] = str(second)
|
||||
return second
|
||||
|
||||
def trio(self):
|
||||
"""a.png (oldest), b.png and c.png share bytes; d.png is unique."""
|
||||
self.write("a.png", b"same bytes", mtime=1_000_000)
|
||||
self.write("b.png", b"same bytes", mtime=2_000_000)
|
||||
self.write("sub/c.png", b"same bytes", mtime=3_000_000)
|
||||
self.write("d.png", b"unique", mtime=4_000_000)
|
||||
|
||||
async def find(self):
|
||||
response = await self.client.request("POST", FIND)
|
||||
text = await response.text()
|
||||
self.assertEqual(response.status, 202, text)
|
||||
self.assertNotIn(self.tmpdir, text)
|
||||
job = self.api._jobs.get(json.loads(text)["job"]["id"])
|
||||
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
|
||||
return job
|
||||
|
||||
async def groups(self, job_id, status=200, **query):
|
||||
params = "&".join(f"{key}={value}" for key, value in query.items())
|
||||
response = await self.client.request("GET", f"{GROUPS}{job_id}?{params}")
|
||||
text = await response.text()
|
||||
self.assertEqual(response.status, status, text)
|
||||
self.assertNotIn(self.tmpdir, text)
|
||||
return json.loads(text)
|
||||
|
||||
async def review(self, job_id, keep, status=200):
|
||||
return await self.post(
|
||||
PREVIEW,
|
||||
{
|
||||
"operation": "duplicate_files",
|
||||
"options": {"job_id": job_id, "keep": keep},
|
||||
},
|
||||
status,
|
||||
)
|
||||
|
||||
async def found_group(self):
|
||||
job = await self.find()
|
||||
[group] = (await self.groups(job.id))["groups"]
|
||||
return job, group
|
||||
|
||||
@staticmethod
|
||||
def keep_of(group, path, root_index=None):
|
||||
[entry] = [
|
||||
f
|
||||
for f in group["files"]
|
||||
if f["path"] == path and root_index in (None, f["root_index"])
|
||||
]
|
||||
return {
|
||||
"group_id": group["group_id"],
|
||||
"root_index": entry["root_index"],
|
||||
"path": entry["path"],
|
||||
}
|
||||
|
||||
async def remove_all_but(self, path):
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, path)])
|
||||
body = await self.apply(plan["token"])
|
||||
cleanup = self.api._jobs.get(body["job"]["id"])
|
||||
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
|
||||
return plan, cleanup
|
||||
|
||||
async def until(self, predicate):
|
||||
"""Wait for a server-side state (a frame at a time), never a fixed timer."""
|
||||
deadline = asyncio.get_running_loop().time() + WAIT
|
||||
while not predicate():
|
||||
if asyncio.get_running_loop().time() > deadline:
|
||||
self.fail("the expected server state never appeared")
|
||||
await asyncio.sleep(0)
|
||||
|
||||
def gate_hashing(self):
|
||||
gate = threading.Event()
|
||||
calls = []
|
||||
original = self.api._calculate_file_hash
|
||||
|
||||
def gated(path):
|
||||
calls.append(path)
|
||||
gate.wait(WAIT)
|
||||
return original(path)
|
||||
|
||||
self.api._calculate_file_hash = gated
|
||||
self.addCleanup(gate.set)
|
||||
return gate, calls
|
||||
|
||||
|
||||
class TestFindJob(DuplicateFilesTestCase):
|
||||
async def test_find_is_listed_in_activity_and_legacy_stream_attaches(self):
|
||||
self.trio()
|
||||
gate, calls = self.gate_hashing()
|
||||
response = await self.client.request("POST", FIND)
|
||||
self.assertEqual(response.status, 202)
|
||||
summary = (await response.json())["job"]
|
||||
listed = (
|
||||
await (await self.client.request("GET", "/prompt_manager/jobs")).json()
|
||||
)["jobs"]
|
||||
self.assertEqual([j["id"] for j in listed], [summary["id"]])
|
||||
self.assertEqual(summary["kind"], "duplicates")
|
||||
self.assertEqual(summary["title"], "Find duplicate image files")
|
||||
self.assertEqual(summary["admission"], "read")
|
||||
self.assertTrue(summary["cancellable"])
|
||||
self.assertEqual(summary["counters"], {"groups_found": 0})
|
||||
# The legacy gallery attaches to the same job, and its status says so.
|
||||
legacy = asyncio.ensure_future(self.client.request("POST", STREAM))
|
||||
job = self.api._jobs.get(summary["id"])
|
||||
await self.until(lambda: job.subscriber_count > 0)
|
||||
status = await (await self.client.request("GET", STATUS)).json()
|
||||
self.assertTrue(status["running"])
|
||||
gate.set()
|
||||
body = await (await legacy).text()
|
||||
job = self.api._jobs.get(summary["id"])
|
||||
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
|
||||
events = _sse_payloads(body)
|
||||
done = events[-1]
|
||||
self.assertEqual(done["type"], "complete")
|
||||
self.assertEqual(done["found"], 1)
|
||||
self.assertEqual(
|
||||
sorted(i["filename"] for i in done["duplicates"][0]["images"]),
|
||||
["a.png", "b.png", "c.png"],
|
||||
)
|
||||
self.assertEqual(len(calls), 4) # one pass: each file hashed once
|
||||
self.assertEqual(job.state, "finished")
|
||||
self.assertEqual(job.counters, {"groups_found": 1})
|
||||
[row] = job.outcomes["groups_found"]
|
||||
self.assertEqual(row["path"], "a.png")
|
||||
self.assertEqual(row["root_index"], 0)
|
||||
self.assertEqual(row["detail"], "3 identical files · 10 B each")
|
||||
self.assertEqual(job.status_line, "Found 1 group of duplicate files")
|
||||
finished = await (await self.client.request("GET", STATUS)).json()
|
||||
self.assertFalse(finished["running"])
|
||||
self.assertEqual(finished["last_event"]["found"], 1)
|
||||
|
||||
async def test_legacy_stream_starts_a_registry_job(self):
|
||||
self.trio()
|
||||
body = await (await self.client.request("POST", STREAM)).text()
|
||||
self.assertEqual(_sse_payloads(body)[-1]["type"], "complete")
|
||||
[job] = self.api._jobs.list()
|
||||
self.assertEqual(job.kind, "duplicates")
|
||||
self.assertEqual(job.title, "Find duplicate image files (gallery page)")
|
||||
self.assertNotIn(self.tmpdir, body)
|
||||
# Its groups are reviewable in the workspace too.
|
||||
page = await self.groups(job.id)
|
||||
self.assertEqual(page["total"], 1)
|
||||
|
||||
async def test_find_covers_every_source_with_its_index(self):
|
||||
second = self.second_source()
|
||||
self.write("x.png", b"twin", folder=second)
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [
|
||||
self._root(self.output_dir),
|
||||
self._root(second),
|
||||
]
|
||||
self.write("x.png", b"twin")
|
||||
job = await self.find()
|
||||
[group] = (await self.groups(job.id))["groups"]
|
||||
self.assertEqual(
|
||||
sorted((f["root_index"], f["path"]) for f in group["files"]),
|
||||
[(0, "x.png"), (1, "x.png")],
|
||||
)
|
||||
self.assertEqual({f["root"] for f in group["files"]}, {"output", "renders"})
|
||||
|
||||
async def test_find_can_be_cancelled(self):
|
||||
self.trio()
|
||||
gate, _calls = self.gate_hashing()
|
||||
response = await self.client.request("POST", FIND)
|
||||
job_id = (await response.json())["job"]["id"]
|
||||
stream = asyncio.ensure_future(self.client.request("POST", STREAM))
|
||||
found = self.api._jobs.get(job_id)
|
||||
await self.until(lambda: found.subscriber_count > 0)
|
||||
cancel = await self.client.request(
|
||||
"POST", f"/prompt_manager/jobs/{job_id}/cancel"
|
||||
)
|
||||
self.assertTrue((await cancel.json())["cancellation_requested"])
|
||||
gate.set()
|
||||
job = self.api._jobs.get(job_id)
|
||||
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
|
||||
self.assertEqual(job.state, "cancelled")
|
||||
last = _sse_payloads(await (await stream).text())[-1]
|
||||
self.assertEqual(last, {"type": "error", "message": "Duplicate scan cancelled"})
|
||||
await self.groups(job_id, 409)
|
||||
|
||||
async def test_second_find_is_busy_and_names_the_running_job(self):
|
||||
self.trio()
|
||||
gate, _calls = self.gate_hashing()
|
||||
first = (await (await self.client.request("POST", FIND)).json())["job"]
|
||||
second = await self.client.request("POST", FIND)
|
||||
self.assertEqual(second.status, 409)
|
||||
body = await second.json()
|
||||
self.assertEqual(body["code"], "busy")
|
||||
self.assertEqual(body["running"]["id"], first["id"])
|
||||
gate.set()
|
||||
|
||||
async def test_failed_find_is_coded_and_releases_the_kind(self):
|
||||
self.write("a.png", b"x")
|
||||
with patch.object(
|
||||
self.api, "_hash_media_groups", side_effect=OSError(f"boom {self.tmpdir}")
|
||||
):
|
||||
with self.assertLogs("prompt_manager", "ERROR"):
|
||||
job = await self.find()
|
||||
self.assertEqual(job.state, "failed")
|
||||
self.assertEqual(job.status_line, "The duplicate search failed.")
|
||||
self.assertNotIn("boom", json.dumps(job.report()))
|
||||
await self.groups(job.id, 409)
|
||||
self.assertEqual((await self.find()).state, "finished")
|
||||
|
||||
|
||||
class TestGroupPages(DuplicateFilesTestCase):
|
||||
async def test_groups_are_paged_ten_at_a_time(self):
|
||||
for index in range(12):
|
||||
for copy in range(2):
|
||||
self.write(f"g{index:02d}_{copy}.png", f"group {index}".encode())
|
||||
job = await self.find()
|
||||
first = await self.groups(job.id)
|
||||
self.assertEqual((first["total"], first["offset"], first["limit"]), (12, 0, 10))
|
||||
self.assertEqual(len(first["groups"]), 10)
|
||||
rest = await self.groups(job.id, offset=10, limit=10)
|
||||
self.assertEqual(len(rest["groups"]), 2)
|
||||
ids = [g["group_id"] for g in first["groups"] + rest["groups"]]
|
||||
self.assertEqual(len(set(ids)), 12)
|
||||
entry = first["groups"][0]["files"][0]
|
||||
for key in ("path", "root", "root_index", "url", "size", "media_type"):
|
||||
self.assertIn(key, entry)
|
||||
self.assertTrue(entry["url"].endswith("?root=0"))
|
||||
for query in ({"limit": 11}, {"limit": 0}, {"offset": -1}, {"offset": "x"}):
|
||||
with self.subTest(query=query):
|
||||
reply = await self.groups(job.id, 400, **query)
|
||||
self.assertEqual(reply["code"], "invalid-request")
|
||||
|
||||
async def test_unknown_or_running_job_has_no_groups(self):
|
||||
reply = await self.groups("f" * 32, 404)
|
||||
self.assertEqual(reply["code"], "not-found")
|
||||
self.trio()
|
||||
gate, _calls = self.gate_hashing()
|
||||
job_id = (await (await self.client.request("POST", FIND)).json())["job"]["id"]
|
||||
reply = await self.groups(job_id, 409)
|
||||
self.assertEqual(reply["code"], "job-running")
|
||||
gate.set()
|
||||
|
||||
|
||||
class TestReview(DuplicateFilesTestCase):
|
||||
async def test_preview_lists_only_non_keepers_and_writes_nothing(self):
|
||||
self.trio()
|
||||
prompt = self.save("Linked")
|
||||
self.api.db.link_image_to_prompt(prompt, str(self.output_dir / "b.png"))
|
||||
job, group = await self.found_group()
|
||||
db_path = self.api.db.model.db_path
|
||||
files = [db_path, db_path + "-wal"]
|
||||
before = [_digest(path) for path in files]
|
||||
media = {p: _digest(p) for p in self.output_dir.rglob("*") if p.is_file()}
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
self.assertEqual([_digest(path) for path in files], before)
|
||||
self.assertEqual(
|
||||
{p: _digest(p) for p in self.output_dir.rglob("*") if p.is_file()},
|
||||
media,
|
||||
)
|
||||
self.assertEqual(plan["operation"], "duplicate_files")
|
||||
self.assertEqual(
|
||||
plan["scope"], {"count": 2, "label": "2 duplicate files in 1 group"}
|
||||
)
|
||||
self.assertTrue(plan["media"]["deletes_files"])
|
||||
self.assertEqual(
|
||||
plan["media"]["label"],
|
||||
"Deletes 2 files (20 B) from disk; the kept copies stay. "
|
||||
"The safety backup covers the database, not these files.",
|
||||
)
|
||||
self.assertEqual(
|
||||
[row["label"] for row in plan["details"]], ["b.png", "sub/c.png"]
|
||||
)
|
||||
self.assertEqual(
|
||||
plan["details"][0]["detail"],
|
||||
"output · a copy of a.png · 1 database link moves to the kept copy",
|
||||
)
|
||||
self.assertEqual(plan["protected"], {"count": 0, "reasons": []})
|
||||
|
||||
async def test_cannot_delete_keeper_or_outside_group(self):
|
||||
self.trio()
|
||||
self.write("e.png", b"other pair")
|
||||
self.write("f.png", b"other pair")
|
||||
job = await self.find()
|
||||
groups = (await self.groups(job.id))["groups"]
|
||||
trio = next(g for g in groups if len(g["files"]) == 3)
|
||||
pair = next(g for g in groups if len(g["files"]) == 2)
|
||||
bad = [
|
||||
# a file of another group named as this group's keeper
|
||||
[{**self.keep_of(trio, "a.png"), "path": "e.png"}],
|
||||
# a file that is in no group
|
||||
[{**self.keep_of(trio, "a.png"), "path": "d.png"}],
|
||||
# the right path under another source index
|
||||
[{**self.keep_of(trio, "a.png"), "root_index": 1}],
|
||||
[{**self.keep_of(trio, "a.png"), "group_id": "0" * 64}],
|
||||
[{**self.keep_of(trio, "a.png"), "extra": 1}],
|
||||
[],
|
||||
"a.png",
|
||||
]
|
||||
for keep in bad:
|
||||
with self.subTest(keep=keep):
|
||||
reply = await self.review(job.id, keep, 400)
|
||||
self.assertEqual(reply["code"], "invalid-request")
|
||||
# Two keepers of one group: only the third copy goes.
|
||||
plan = await self.review(
|
||||
job.id,
|
||||
[self.keep_of(trio, "a.png"), self.keep_of(trio, "sub/c.png")],
|
||||
)
|
||||
self.assertEqual([row["label"] for row in plan["details"]], ["b.png"])
|
||||
body = await self.apply(plan["token"])
|
||||
cleanup = self.api._jobs.get(body["job"]["id"])
|
||||
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
|
||||
self.assertEqual(cleanup.counters["removed"], 1)
|
||||
self.assertTrue((self.output_dir / "a.png").is_file())
|
||||
self.assertTrue((self.output_dir / "sub" / "c.png").is_file())
|
||||
self.assertFalse((self.output_dir / "b.png").exists())
|
||||
# The untouched group is untouched.
|
||||
self.assertTrue((self.output_dir / "e.png").is_file())
|
||||
self.assertTrue((self.output_dir / "f.png").is_file())
|
||||
self.assertEqual(len(pair["files"]), 2)
|
||||
|
||||
async def test_a_source_removed_since_the_find_is_refused(self):
|
||||
self.trio()
|
||||
job, group = await self.found_group()
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.second_source())]
|
||||
reply = await self.review(job.id, [self.keep_of(group, "a.png")], 409)
|
||||
self.assertEqual(reply["code"], "sources-changed")
|
||||
|
||||
def test_containment_is_decided_as_text_before_the_filesystem(self):
|
||||
root = str(self.output_dir)
|
||||
outside = os.path.join(str(self.outside_dir), "x.png")
|
||||
touched = []
|
||||
real_lstat = os.lstat
|
||||
|
||||
def lstat(path, *args, **kwargs):
|
||||
touched.append(str(path))
|
||||
return real_lstat(path, *args, **kwargs)
|
||||
|
||||
with patch.object(duplicate_files.os, "lstat", side_effect=lstat):
|
||||
for candidate in (
|
||||
outside,
|
||||
os.path.join(root, "..", "outside", "x.png"),
|
||||
os.path.join(root, "thumbnails", "a_thumb.png"),
|
||||
os.path.join(root, "notes.txt"),
|
||||
):
|
||||
with self.subTest(candidate=candidate):
|
||||
reason = duplicate_files.file_problem(candidate, root, "0" * 64, 1)
|
||||
self.assertEqual(reason, duplicate_files.OUTSIDE_SOURCE)
|
||||
self.assertEqual(touched, [])
|
||||
|
||||
@staticmethod
|
||||
def _can_symlink(directory):
|
||||
probe = Path(directory) / "probe-link"
|
||||
try:
|
||||
os.symlink(str(Path(directory) / "x"), str(probe))
|
||||
except (OSError, NotImplementedError):
|
||||
return False
|
||||
probe.unlink()
|
||||
return True
|
||||
|
||||
async def test_symlink_rejected(self):
|
||||
if not self._can_symlink(self.tmpdir): # pragma: no cover - Windows w/o rights
|
||||
self.skipTest("symlinks are not available")
|
||||
self.trio()
|
||||
job, group = await self.found_group()
|
||||
# b.png becomes a link to the keeper after the find: never removed.
|
||||
(self.output_dir / "b.png").unlink()
|
||||
os.symlink(str(self.output_dir / "a.png"), str(self.output_dir / "b.png"))
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
self.assertEqual([row["label"] for row in plan["details"]], ["sub/c.png"])
|
||||
self.assertEqual(
|
||||
plan["protected"],
|
||||
{"count": 1, "reasons": ["1 file is a link, not a regular file"]},
|
||||
)
|
||||
# A keeper that became a link protects its whole group.
|
||||
(self.output_dir / "a.png").rename(self.output_dir / "real.png")
|
||||
os.symlink(str(self.output_dir / "real.png"), str(self.output_dir / "a.png"))
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
self.assertEqual(plan["scope"]["count"], 0)
|
||||
self.assertEqual(
|
||||
plan["protected"]["reasons"],
|
||||
["2 files kept: their kept copy changed or is gone"],
|
||||
)
|
||||
|
||||
|
||||
class TestCleanupJob(DuplicateFilesTestCase):
|
||||
async def test_cleanup_is_a_records_job_with_a_safety_backup(self):
|
||||
self.trio()
|
||||
plan, cleanup = await self.remove_all_but("a.png")
|
||||
summary = cleanup.summary()
|
||||
self.assertEqual(summary["kind"], "duplicate_cleanup")
|
||||
self.assertEqual(summary["title"], "Remove duplicate image files")
|
||||
self.assertEqual(summary["admission"], "records")
|
||||
self.assertFalse(summary["cancellable"])
|
||||
self.assertEqual(summary["scope"], {"operation": "duplicate_files"})
|
||||
self.assertEqual(cleanup.state, "finished", cleanup.log)
|
||||
self.assertEqual(cleanup.counters, {"removed": 2, "kept": 1, "errors": 0})
|
||||
self.assertEqual(
|
||||
sorted(row["path"] for row in cleanup.outcomes["removed"]),
|
||||
["b.png", "sub/c.png"],
|
||||
)
|
||||
self.assertEqual(cleanup.outcomes["kept"][0]["path"], "a.png")
|
||||
[backup] = self.backups()
|
||||
self.assertEqual(
|
||||
cleanup.status_line,
|
||||
f"Removed 2 duplicate files · Safety backup {backup}",
|
||||
)
|
||||
self.assertEqual(
|
||||
_sha(self.output_dir / "a.png"), hashlib.sha256(b"same bytes").hexdigest()
|
||||
)
|
||||
self.assertFalse((self.output_dir / "b.png").exists())
|
||||
self.assertFalse((self.output_dir / "sub" / "c.png").exists())
|
||||
self.assertNotIn(self.tmpdir, json.dumps(cleanup.report()))
|
||||
# Single use, like every plan.
|
||||
again = await self.apply(plan["token"], 409)
|
||||
self.assertEqual(again["code"], "plan-used")
|
||||
|
||||
async def test_thumbnails_of_removed_files_are_removed(self):
|
||||
self.trio()
|
||||
self.write("clip.mp4", b"video")
|
||||
self.write("copy.mp4", b"video")
|
||||
kept = self.write("thumbnails/a_thumb.png", b"t")
|
||||
gone = self.write("thumbnails/b_thumb.png", b"t")
|
||||
video_thumb = self.write("thumbnails/copy_thumb.jpg", b"t")
|
||||
job = await self.find()
|
||||
groups = (await self.groups(job.id))["groups"]
|
||||
keep = []
|
||||
for group in groups:
|
||||
paths = sorted(f["path"] for f in group["files"])
|
||||
keep.append(
|
||||
self.keep_of(group, "a.png" if "a.png" in paths else "clip.mp4")
|
||||
)
|
||||
plan = await self.review(job.id, keep)
|
||||
body = await self.apply(plan["token"])
|
||||
cleanup = self.api._jobs.get(body["job"]["id"])
|
||||
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
|
||||
self.assertTrue(kept.is_file())
|
||||
self.assertFalse(gone.exists())
|
||||
self.assertFalse(video_thumb.exists())
|
||||
|
||||
async def test_changed_content_skips_file(self):
|
||||
self.trio()
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
self.write("b.png", b"new bytes") # changed since the review
|
||||
body = await self.apply(plan["token"])
|
||||
cleanup = self.api._jobs.get(body["job"]["id"])
|
||||
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
|
||||
self.assertEqual(cleanup.counters, {"removed": 1, "kept": 2, "errors": 0})
|
||||
self.assertEqual((self.output_dir / "b.png").read_bytes(), b"new bytes")
|
||||
self.assertFalse((self.output_dir / "sub" / "c.png").exists())
|
||||
rows = {row["path"]: row["detail"] for row in cleanup.outcomes["kept"]}
|
||||
self.assertEqual(rows["b.png"], duplicate_files.CHANGED)
|
||||
|
||||
async def test_a_changed_keeper_keeps_its_whole_group(self):
|
||||
self.trio()
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
self.write("a.png", b"edited keeper")
|
||||
body = await self.apply(plan["token"])
|
||||
cleanup = self.api._jobs.get(body["job"]["id"])
|
||||
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
|
||||
self.assertEqual(cleanup.counters["removed"], 0)
|
||||
self.assertTrue((self.output_dir / "b.png").is_file())
|
||||
self.assertTrue((self.output_dir / "sub" / "c.png").is_file())
|
||||
self.assertTrue(
|
||||
cleanup.status_line.startswith(
|
||||
"Removed 0 duplicate files, 2 not removed · "
|
||||
)
|
||||
)
|
||||
rows = {row["path"]: row["detail"] for row in cleanup.outcomes["kept"]}
|
||||
self.assertEqual(rows["b.png"], duplicate_files.KEEPER_CHANGED)
|
||||
|
||||
async def test_partial_failure_reports_each_path(self):
|
||||
self.trio()
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
real = duplicate_files._unlink
|
||||
|
||||
def unlink(path):
|
||||
if os.path.basename(path) == "c.png":
|
||||
raise PermissionError(13, f"Permission denied: {path}")
|
||||
real(path)
|
||||
|
||||
self.write("z.png", b"same bytes") # never reviewed: not touched
|
||||
with patch.object(duplicate_files, "_unlink", side_effect=unlink):
|
||||
body = await self.apply(plan["token"])
|
||||
cleanup = self.api._jobs.get(body["job"]["id"])
|
||||
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
|
||||
self.assertEqual(cleanup.state, "failed")
|
||||
self.assertEqual(cleanup.counters, {"removed": 1, "kept": 1, "errors": 1})
|
||||
self.assertEqual(
|
||||
[row["path"] for row in cleanup.outcomes["removed"]], ["b.png"]
|
||||
)
|
||||
[error] = cleanup.outcomes["errors"]
|
||||
self.assertEqual(
|
||||
(error["path"], error["detail"]),
|
||||
("sub/c.png", duplicate_files.LOCKED_FILE),
|
||||
)
|
||||
self.assertEqual(
|
||||
cleanup.status_line,
|
||||
"Stopped after an error · Removed 1 duplicate file, 1 not removed · "
|
||||
"Safety backup " + self.backups()[0],
|
||||
)
|
||||
report = json.dumps(cleanup.report())
|
||||
self.assertNotIn("Permission denied", report)
|
||||
self.assertNotIn("Errno", report)
|
||||
self.assertTrue((self.output_dir / "z.png").is_file())
|
||||
|
||||
async def test_stop_on_failure_leaves_later_files_untouched(self):
|
||||
self.write("a.png", b"same", mtime=1)
|
||||
for name in ("b.png", "c.png", "e.png"):
|
||||
self.write(name, b"same", mtime=2)
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
with (
|
||||
patch.object(
|
||||
duplicate_files, "_unlink", side_effect=OSError(5, "I/O error")
|
||||
),
|
||||
self.assertLogs("prompt_manager", "WARNING"),
|
||||
):
|
||||
body = await self.apply(plan["token"])
|
||||
cleanup = self.api._jobs.get(body["job"]["id"])
|
||||
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
|
||||
self.assertEqual(cleanup.counters, {"removed": 0, "kept": 3, "errors": 1})
|
||||
self.assertEqual(
|
||||
cleanup.outcomes["errors"][0]["detail"], duplicate_files.DELETE_FAILED
|
||||
)
|
||||
not_attempted = [
|
||||
row
|
||||
for row in cleanup.outcomes["kept"]
|
||||
if row["detail"] == duplicate_files.NOT_ATTEMPTED
|
||||
]
|
||||
self.assertEqual(len(not_attempted), 2)
|
||||
for name in ("b.png", "c.png", "e.png"):
|
||||
self.assertTrue((self.output_dir / name).is_file())
|
||||
|
||||
async def test_locked_file_gives_a_coded_message(self):
|
||||
self.trio()
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
target = self.output_dir / "b.png"
|
||||
handle = None
|
||||
if os.name == "nt": # pragma: no cover - the native sharing violation
|
||||
handle = open(target, "rb")
|
||||
context = patch.object(duplicate_files, "_unlink", duplicate_files._unlink)
|
||||
else:
|
||||
context = patch.object(
|
||||
duplicate_files,
|
||||
"_unlink",
|
||||
side_effect=PermissionError(
|
||||
32, "The process cannot access the file", str(target)
|
||||
),
|
||||
)
|
||||
try:
|
||||
with context:
|
||||
body = await self.apply(plan["token"])
|
||||
cleanup = self.api._jobs.get(body["job"]["id"])
|
||||
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
|
||||
finally:
|
||||
if handle is not None: # pragma: no cover - Windows
|
||||
handle.close()
|
||||
[error] = cleanup.outcomes["errors"]
|
||||
self.assertEqual(error["detail"], duplicate_files.LOCKED_FILE)
|
||||
self.assertNotIn("WinError", json.dumps(cleanup.report()))
|
||||
self.assertTrue(target.is_file())
|
||||
|
||||
async def test_handles_closed_before_unlink(self):
|
||||
self.trio()
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
real = duplicate_files._unlink
|
||||
open_at_unlink = []
|
||||
|
||||
def unlink(path):
|
||||
open_at_unlink.extend(open_handles_under(path))
|
||||
real(path)
|
||||
|
||||
with patch.object(duplicate_files, "_unlink", side_effect=unlink):
|
||||
body = await self.apply(plan["token"])
|
||||
cleanup = self.api._jobs.get(body["job"]["id"])
|
||||
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
|
||||
self.assertEqual(cleanup.counters["removed"], 2)
|
||||
self.assertEqual(open_at_unlink, [])
|
||||
|
||||
async def test_db_links_reconciled_only_for_removed_files(self):
|
||||
self.trio()
|
||||
db = self.api.db
|
||||
moved = self.save("Moves to the keeper")
|
||||
merged = self.save("Already on the keeper")
|
||||
db.link_image_to_prompt(moved, str(self.output_dir / "b.png"))
|
||||
db.link_image_to_prompt(moved, str(self.output_dir / "d.png"))
|
||||
db.link_image_to_prompt(merged, str(self.output_dir / "a.png"))
|
||||
db.link_image_to_prompt(
|
||||
merged,
|
||||
str(self.output_dir / "sub" / "c.png"),
|
||||
{"workflow": {"only": "c"}, "prompt": {"only": "c"}},
|
||||
)
|
||||
db.set_image_favorite(
|
||||
duplicate_files.favorite_root_key(self.output_dir), "b.png"
|
||||
)
|
||||
plan, cleanup = await self.remove_all_but("a.png")
|
||||
# The merged link's own workflow document went with it (#249 S6 dedupe).
|
||||
documents = db.model.get_connection().execute(
|
||||
"SELECT COUNT(*) FROM workflow_documents"
|
||||
)
|
||||
self.assertEqual(documents.fetchone()[0], 1) # the shared "{}"
|
||||
self.assertEqual(
|
||||
plan["details"][0]["detail"],
|
||||
"output · a copy of a.png · 1 database link moves to the kept copy",
|
||||
)
|
||||
|
||||
def files_of(prompt_id):
|
||||
return sorted(
|
||||
os.path.relpath(image["image_path"], self.output_dir)
|
||||
for image in db.get_prompt_images(prompt_id)
|
||||
)
|
||||
|
||||
self.assertEqual(files_of(moved), ["a.png", "d.png"])
|
||||
self.assertEqual(files_of(merged), ["a.png"])
|
||||
root_key = duplicate_files.favorite_root_key(self.output_dir)
|
||||
self.assertEqual(db.get_root_favorites(root_key), {"a.png"})
|
||||
self.assertIn(
|
||||
"Moved 2 database links to the kept copies", "\n".join(cleanup.log)
|
||||
)
|
||||
|
||||
async def test_a_failed_removal_leaves_its_links_on_the_kept_copy(self):
|
||||
self.trio()
|
||||
prompt = self.save("Stays on c")
|
||||
self.api.db.link_image_to_prompt(prompt, str(self.output_dir / "sub" / "c.png"))
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
real = duplicate_files._unlink
|
||||
|
||||
def unlink(path):
|
||||
if os.path.basename(path) == "c.png":
|
||||
raise PermissionError(13, "denied")
|
||||
real(path)
|
||||
|
||||
with patch.object(duplicate_files, "_unlink", side_effect=unlink):
|
||||
body = await self.apply(plan["token"])
|
||||
cleanup = self.api._jobs.get(body["job"]["id"])
|
||||
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
|
||||
# Its records moved to the identical kept copy before the unlink failed.
|
||||
[image] = self.api.db.get_prompt_images(prompt)
|
||||
self.assertTrue(image["image_path"].endswith("a.png"))
|
||||
self.assertTrue((self.output_dir / "sub" / "c.png").is_file())
|
||||
|
||||
async def test_cleanup_and_find_exclude_each_other(self):
|
||||
self.trio()
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
gate, _calls = self.gate_hashing()
|
||||
find = (await (await self.client.request("POST", FIND)).json())["job"]
|
||||
busy = await self.apply(plan["token"], 409)
|
||||
self.assertEqual((busy["code"], busy["running"]["id"]), ("busy", find["id"]))
|
||||
gate.set()
|
||||
await asyncio.wait_for(
|
||||
asyncio.shield(self.api._jobs.get(find["id"]).task), WAIT
|
||||
)
|
||||
# The token survived the refusal; the cleanup holds off a new find.
|
||||
release = asyncio.Event()
|
||||
real = self.api._run_duplicate_cleanup
|
||||
|
||||
async def held(*args):
|
||||
await release.wait()
|
||||
return await real(*args)
|
||||
|
||||
with patch.object(self.api, "_run_duplicate_cleanup", side_effect=held):
|
||||
body = await self.apply(plan["token"])
|
||||
refused = await self.client.request("POST", FIND)
|
||||
self.assertEqual(refused.status, 409)
|
||||
legacy = await self.client.request("POST", STREAM)
|
||||
self.assertEqual(legacy.status, 409)
|
||||
release.set()
|
||||
cleanup = self.api._jobs.get(body["job"]["id"])
|
||||
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
|
||||
|
||||
async def test_cleanup_refused_up_front_when_the_backup_does_not_fit(self):
|
||||
from py.api import maintenance
|
||||
|
||||
self.trio()
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
with patch.object(maintenance, "_free_bytes", return_value=0):
|
||||
reply = await self.apply(plan["token"], 507)
|
||||
self.assertEqual(reply["code"], "insufficient-space")
|
||||
self.assertTrue((self.output_dir / "b.png").is_file())
|
||||
|
||||
async def test_policy_off_removes_without_a_backup(self):
|
||||
PromptManagerConfig.SAFETY_BACKUP = False
|
||||
self.trio()
|
||||
_plan, cleanup = await self.remove_all_but("a.png")
|
||||
self.assertEqual(self.backups(), [])
|
||||
self.assertEqual(
|
||||
cleanup.status_line, "Removed 2 duplicate files · No safety backup"
|
||||
)
|
||||
|
||||
|
||||
class TestLegacyDuplicateRoutes(DuplicateFilesTestCase):
|
||||
async def test_legacy_delete_duplicates_unchanged(self):
|
||||
self.trio()
|
||||
response = await self.client.request(
|
||||
"POST", LEGACY_DELETE, json={"image_paths": ["b.png", "nope.png"]}
|
||||
)
|
||||
body = await response.json()
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(
|
||||
set(body),
|
||||
{"success", "deleted_count", "failed_count", "message", "failed_files"},
|
||||
)
|
||||
self.assertEqual((body["deleted_count"], body["failed_count"]), (1, 1))
|
||||
self.assertEqual(body["failed_files"], ["nope.png (file not found)"])
|
||||
[hold] = self.api._jobs.list()
|
||||
self.assertEqual(
|
||||
(hold.kind, hold.title, hold.admission),
|
||||
("duplicate_cleanup", "Duplicate image delete (gallery page)", "records"),
|
||||
)
|
||||
|
||||
async def test_legacy_delete_refused_while_a_find_runs(self):
|
||||
self.trio()
|
||||
gate, _calls = self.gate_hashing()
|
||||
await self.client.request("POST", FIND)
|
||||
response = await self.client.request(
|
||||
"POST", LEGACY_DELETE, json={"image_paths": ["b.png"]}
|
||||
)
|
||||
self.assertEqual(response.status, 409)
|
||||
self.assertEqual(
|
||||
await response.json(),
|
||||
{"success": False, "error": "Find duplicate image files already running"},
|
||||
)
|
||||
self.assertTrue((self.output_dir / "b.png").is_file())
|
||||
gate.set()
|
||||
|
||||
async def test_legacy_scan_holds_the_duplicates_kind(self):
|
||||
self.trio()
|
||||
gate, _calls = self.gate_hashing()
|
||||
legacy = asyncio.ensure_future(
|
||||
self.client.request("GET", "/prompt_manager/scan_duplicates")
|
||||
)
|
||||
await self.until(lambda: self.api._jobs.active("duplicates") is not None)
|
||||
[hold] = self.api._jobs.list()
|
||||
self.assertEqual(
|
||||
(hold.kind, hold.title, hold.admission, hold.cancellable),
|
||||
("duplicates", "Find duplicate image files (gallery page)", "read", False),
|
||||
)
|
||||
refused = await self.client.request("POST", FIND)
|
||||
self.assertEqual(refused.status, 409)
|
||||
# A held legacy scan publishes no events: the stream does not attach.
|
||||
stream = await self.client.request("POST", STREAM)
|
||||
self.assertEqual(stream.status, 409)
|
||||
gate.set()
|
||||
body = await (await legacy).json()
|
||||
self.assertEqual(set(body), {"success", "duplicates", "message"})
|
||||
self.assertEqual(body["duplicates"][0]["count"], 3)
|
||||
|
||||
async def test_legacy_stream_does_not_attach_to_a_find_missing_its_folder(self):
|
||||
other = self.second_source()
|
||||
self.write("x.png", b"x", folder=other)
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [self._root(other)]
|
||||
self.api._find_comfyui_output_dir = lambda: str(self.output_dir)
|
||||
gate, _calls = self.gate_hashing()
|
||||
await self.client.request("POST", FIND)
|
||||
stream = await self.client.request("POST", STREAM)
|
||||
self.assertEqual(stream.status, 409)
|
||||
gate.set()
|
||||
|
||||
|
||||
class TestLinkQueries(DuplicateFilesTestCase):
|
||||
def test_link_queries_use_the_path_indexes(self):
|
||||
conn = sqlite3.connect(self.api.db.model.db_path)
|
||||
try:
|
||||
for sql in duplicate_files.query_plan_sql():
|
||||
plan = " ".join(
|
||||
row[3]
|
||||
for row in conn.execute(
|
||||
"EXPLAIN QUERY PLAN " + sql, [None] * sql.count("?")
|
||||
)
|
||||
)
|
||||
with self.subTest(sql=sql):
|
||||
self.assertNotIn("SCAN generated_images", plan)
|
||||
self.assertNotIn("SCAN image_favorites", plan)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
class TestReviewHelpers(DuplicateFilesTestCase):
|
||||
"""The pure checks of py/api/duplicate_review.py."""
|
||||
|
||||
def test_file_problem_on_disk(self):
|
||||
from py.api import duplicate_review as review
|
||||
|
||||
root = str(self.output_dir)
|
||||
good = self.write("ok.png", b"bytes")
|
||||
digest = hashlib.sha256(b"bytes").hexdigest()
|
||||
self.assertIsNone(review.file_problem(str(good), root, digest, 5))
|
||||
self.assertEqual(
|
||||
review.file_problem(str(good), root, digest, 6), review.CHANGED_CONTENT
|
||||
)
|
||||
self.assertEqual(
|
||||
review.file_problem(str(good), root, "0" * 64, 5), review.CHANGED_CONTENT
|
||||
)
|
||||
self.assertEqual(
|
||||
review.file_problem(os.path.join(root, "gone.png"), root, digest, 5),
|
||||
review.GONE,
|
||||
)
|
||||
folder = self.output_dir / "folder.png"
|
||||
folder.mkdir()
|
||||
self.assertEqual(review.file_problem(str(folder), root, digest, 5), review.LINK)
|
||||
with patch.object(review.os, "lstat", side_effect=OSError("io")):
|
||||
self.assertEqual(
|
||||
review.file_problem(str(good), root, digest, 5), review.UNREADABLE
|
||||
)
|
||||
|
||||
def unreadable(_path):
|
||||
raise OSError("io")
|
||||
|
||||
self.assertEqual(
|
||||
review.file_problem(str(good), root, digest, 5, unreadable),
|
||||
review.UNREADABLE,
|
||||
)
|
||||
|
||||
def vanished(_path):
|
||||
raise FileNotFoundError
|
||||
|
||||
self.assertEqual(
|
||||
review.file_problem(str(good), root, digest, 5, vanished), review.GONE
|
||||
)
|
||||
|
||||
def test_a_linked_folder_is_outside_its_source(self):
|
||||
from py.api import duplicate_review as review
|
||||
|
||||
if not TestReview._can_symlink(self.tmpdir): # pragma: no cover
|
||||
self.skipTest("symlinks are not available")
|
||||
target = self.write("x.png", b"bytes", folder=self.outside_dir)
|
||||
os.symlink(str(self.outside_dir), str(self.output_dir / "linked"))
|
||||
digest = hashlib.sha256(b"bytes").hexdigest()
|
||||
path = os.path.join(str(self.output_dir), "linked", "x.png")
|
||||
self.assertTrue(os.path.isfile(path))
|
||||
self.assertEqual(
|
||||
review.file_problem(path, str(self.output_dir), digest, 5),
|
||||
review.OUTSIDE_SOURCE,
|
||||
)
|
||||
self.assertTrue(target.is_file())
|
||||
|
||||
def test_options_are_validated(self):
|
||||
from py.api import duplicate_review as review
|
||||
|
||||
keep = {"group_id": "g", "root_index": 0, "path": "a.png"}
|
||||
cases = [
|
||||
None,
|
||||
{"job_id": "x"},
|
||||
{"job_id": "", "keep": [keep]},
|
||||
{"job_id": "x" * 65, "keep": [keep]},
|
||||
{"job_id": "x", "keep": {}},
|
||||
{"job_id": "x", "keep": [keep] * (review.MAX_KEEP + 1)},
|
||||
{"job_id": "x", "keep": [{**keep, "group_id": 1}]},
|
||||
{"job_id": "x", "keep": [{**keep, "root_index": True}]},
|
||||
{"job_id": "x", "keep": [{**keep, "root_index": -1}]},
|
||||
{"job_id": "x", "keep": [{**keep, "path": ""}]},
|
||||
{"job_id": "x", "keep": [["g", 0, "a.png"]]},
|
||||
]
|
||||
for options in cases:
|
||||
with self.subTest(options=options):
|
||||
self.assertIsNotNone(review._parse_options(options)[2])
|
||||
job_id, parsed, message = review._parse_options(
|
||||
{"job_id": "x", "keep": [{**keep, "root_index": None}]}
|
||||
)
|
||||
self.assertEqual((job_id, parsed, message), ("x", [("g", None, "a.png")], None))
|
||||
|
||||
def test_sizes_read_naturally(self):
|
||||
from py.api.duplicate_review import human_bytes
|
||||
|
||||
self.assertEqual(human_bytes(10), "10 B")
|
||||
self.assertEqual(human_bytes(1536), "1.5 KB")
|
||||
self.assertEqual(human_bytes(3 * 2**20), "3.0 MB")
|
||||
self.assertEqual(human_bytes(2**40 * 5000), "5,000.0 TB")
|
||||
|
||||
async def test_a_kept_copy_listed_twice_counts_once(self):
|
||||
self.trio()
|
||||
job, group = await self.found_group()
|
||||
keep = self.keep_of(group, "a.png")
|
||||
plan = await self.review(job.id, [keep, keep])
|
||||
self.assertEqual(plan["scope"]["count"], 2)
|
||||
|
||||
async def test_preview_failure_is_coded(self):
|
||||
self.trio()
|
||||
job, group = await self.found_group()
|
||||
with (
|
||||
patch.object(
|
||||
self.api, "_review_duplicate_files", side_effect=OSError("boom")
|
||||
),
|
||||
self.assertLogs("prompt_manager", "ERROR"),
|
||||
):
|
||||
reply = await self.review(job.id, [self.keep_of(group, "a.png")], 500)
|
||||
self.assertEqual(reply["error"], "Failed to review maintenance")
|
||||
|
||||
async def test_a_source_removed_before_the_apply_deletes_nothing(self):
|
||||
self.trio()
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.second_source())]
|
||||
body = await self.apply(plan["token"])
|
||||
cleanup = self.api._jobs.get(body["job"]["id"])
|
||||
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
|
||||
self.assertEqual(cleanup.state, "failed")
|
||||
self.assertEqual(cleanup.last_event["code"], "sources-changed")
|
||||
self.assertEqual(self.backups(), [])
|
||||
self.assertTrue((self.output_dir / "b.png").is_file())
|
||||
|
||||
async def test_a_records_move_failure_deletes_nothing(self):
|
||||
self.trio()
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
with (
|
||||
patch.object(
|
||||
self.api.db,
|
||||
"move_copy_records",
|
||||
side_effect=sqlite3.OperationalError("x"),
|
||||
),
|
||||
self.assertLogs("prompt_manager", "ERROR"),
|
||||
):
|
||||
body = await self.apply(plan["token"])
|
||||
cleanup = self.api._jobs.get(body["job"]["id"])
|
||||
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
|
||||
self.assertEqual(cleanup.state, "failed")
|
||||
self.assertEqual(
|
||||
cleanup.outcomes["errors"][-1]["detail"], duplicate_files.RECORDS_FAILED
|
||||
)
|
||||
self.assertEqual(cleanup.counters["removed"], 0)
|
||||
self.assertTrue((self.output_dir / "b.png").is_file())
|
||||
self.assertTrue((self.output_dir / "sub" / "c.png").is_file())
|
||||
@@ -1,558 +0,0 @@
|
||||
"""Duplicate-file removal safety (#249 System S6 final review: C1, I1, I2, I5, M9).
|
||||
|
||||
- C1: one physical file reached through two listed sources, a symlinked
|
||||
source or a hardlink is never offered or removed as a copy of itself.
|
||||
- I1: each file's records move to its verified kept copy before that file
|
||||
is unlinked, so no database row ever names a deleted file.
|
||||
- I2: links stored under another spelling of the source move too.
|
||||
- I5: image tags move to the kept copy (their union) and the removed copy's
|
||||
pending Needs review items are settled.
|
||||
- M9: the preview binds files by size, mtime and inode without hashing.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import os
|
||||
import sqlite3
|
||||
import types
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
try:
|
||||
from tests.open_handles import open_handles_under
|
||||
from tests import test_api_duplicate_files as dup_tests
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
import test_api_duplicate_files as dup_tests
|
||||
from open_handles import open_handles_under
|
||||
|
||||
from py.api import duplicate_files # noqa: E402
|
||||
from py.config import GalleryConfig # noqa: E402
|
||||
|
||||
WAIT = 10
|
||||
|
||||
|
||||
def _locator(path):
|
||||
return os.path.normcase(os.path.realpath(str(path)))
|
||||
|
||||
|
||||
class SafetyTestCase(dup_tests.DuplicateFilesTestCase):
|
||||
def need_symlinks(self):
|
||||
if not dup_tests.TestReview._can_symlink(self.tmpdir): # pragma: no cover
|
||||
self.skipTest("symlinks are not available")
|
||||
|
||||
def linked_files(self):
|
||||
"""Every image record's path that names a file that does not exist."""
|
||||
with self.api.db.model.get_connection() as conn:
|
||||
rows = conn.execute("SELECT image_path FROM generated_images").fetchall()
|
||||
return [row[0] for row in rows if not os.path.exists(row[0])]
|
||||
|
||||
async def apply_plan(self, plan):
|
||||
body = await self.apply(plan["token"])
|
||||
cleanup = self.api._jobs.get(body["job"]["id"])
|
||||
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
|
||||
return cleanup
|
||||
|
||||
|
||||
class TestOnePhysicalFile(SafetyTestCase):
|
||||
"""C1: the kept copy survives every alias of itself."""
|
||||
|
||||
async def test_a_folder_listed_twice_never_loses_its_only_copy(self):
|
||||
self.need_symlinks()
|
||||
only = self.write("a.png", b"the only copy")
|
||||
alias = Path(self.tmpdir) / "Output-alias"
|
||||
os.symlink(str(self.output_dir), str(alias))
|
||||
with patch.object(
|
||||
self.api, "_get_all_output_dirs", return_value=[self.output_dir, alias]
|
||||
):
|
||||
job = await self.find()
|
||||
page = await self.groups(job.id)
|
||||
self.assertEqual(page["total"], 0)
|
||||
self.assertEqual(only.read_bytes(), b"the only copy")
|
||||
|
||||
async def test_a_symlinked_source_inside_another_is_one_file(self):
|
||||
self.need_symlinks()
|
||||
only = self.write("sub/c.png", b"one file")
|
||||
linked = self.output_dir / "linked" # walked by neither as a link
|
||||
os.symlink(str(self.output_dir / "sub"), str(linked))
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [
|
||||
self._root(self.output_dir),
|
||||
str(linked),
|
||||
]
|
||||
self.assertEqual(len(self.api._get_all_output_dirs()), 2)
|
||||
job = await self.find()
|
||||
self.assertEqual((await self.groups(job.id))["total"], 0)
|
||||
self.assertTrue(only.is_file())
|
||||
|
||||
async def test_a_hardlink_is_never_a_copy_to_remove(self):
|
||||
first = self.write("a.png", b"same bytes", mtime=1)
|
||||
os.link(str(first), str(self.output_dir / "b.png"))
|
||||
self.write("c.png", b"same bytes", mtime=3)
|
||||
job, group = await self.found_group()
|
||||
self.assertEqual(len(group["files"]), 2) # a.png and b.png are one file
|
||||
# c.png becomes a hardlink of the keeper after the find: protected.
|
||||
(self.output_dir / "c.png").unlink()
|
||||
os.link(str(first), str(self.output_dir / "c.png"))
|
||||
kept = group["files"][0]["path"]
|
||||
plan = await self.review(job.id, [self.keep_of(group, kept)])
|
||||
self.assertEqual(plan["scope"]["count"], 0)
|
||||
self.assertEqual(
|
||||
plan["protected"]["reasons"], ["1 file is the same file as its kept copy"]
|
||||
)
|
||||
|
||||
async def test_the_job_refuses_a_file_that_became_its_kept_copy(self):
|
||||
first = self.write("a.png", b"same bytes", mtime=1)
|
||||
self.write("b.png", b"same bytes", mtime=2)
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
(self.output_dir / "b.png").unlink()
|
||||
os.link(str(first), str(self.output_dir / "b.png"))
|
||||
cleanup = await self.apply_plan(plan)
|
||||
self.assertEqual(cleanup.counters["removed"], 0)
|
||||
rows = {row["path"]: row["detail"] for row in cleanup.outcomes["kept"]}
|
||||
self.assertEqual(rows["b.png"], duplicate_files.SAME_FILE)
|
||||
self.assertTrue((self.output_dir / "b.png").is_file())
|
||||
|
||||
|
||||
class TestRecordsMoveBeforeTheUnlink(SafetyTestCase):
|
||||
"""I1: a deleted file is never named by a database row."""
|
||||
|
||||
async def test_a_failure_between_two_files_leaves_no_row_on_a_deleted_file(self):
|
||||
self.write("a.png", b"same", mtime=1)
|
||||
self.write("b.png", b"same", mtime=2)
|
||||
self.write("c.png", b"same", mtime=3)
|
||||
prompt = self.save("Linked to every copy")
|
||||
for name in ("b.png", "c.png"):
|
||||
self.api.db.link_image_to_prompt(prompt, str(self.output_dir / name))
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
real = self.api.db.move_copy_records
|
||||
calls = []
|
||||
|
||||
def move(*args, **kwargs):
|
||||
calls.append(args)
|
||||
if len(calls) == 2:
|
||||
raise sqlite3.OperationalError("database is locked")
|
||||
return real(*args, **kwargs)
|
||||
|
||||
with patch.object(self.api.db, "move_copy_records", side_effect=move):
|
||||
with self.assertLogs("prompt_manager", "ERROR"):
|
||||
cleanup = await self.apply_plan(plan)
|
||||
self.assertEqual(self.linked_files(), [])
|
||||
self.assertEqual(cleanup.counters["removed"], 1)
|
||||
self.assertEqual(cleanup.state, "failed")
|
||||
[error] = cleanup.outcomes["errors"]
|
||||
self.assertEqual(error["detail"], duplicate_files.RECORDS_FAILED)
|
||||
self.assertNotIn("Missing-file", cleanup.status_line)
|
||||
self.assertEqual(
|
||||
cleanup.status_line.split(" · ")[:2],
|
||||
["Stopped after an error", "Removed 1 duplicate file, 1 not removed"],
|
||||
)
|
||||
|
||||
async def test_a_failed_unlink_leaves_its_links_on_the_kept_copy(self):
|
||||
self.trio()
|
||||
prompt = self.save("On c")
|
||||
self.api.db.link_image_to_prompt(prompt, str(self.output_dir / "sub" / "c.png"))
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
real = duplicate_files._unlink
|
||||
|
||||
def unlink(path):
|
||||
if os.path.basename(path) == "c.png":
|
||||
raise PermissionError(13, "denied")
|
||||
real(path)
|
||||
|
||||
with patch.object(duplicate_files, "_unlink", side_effect=unlink):
|
||||
with self.assertLogs("prompt_manager", "WARNING"):
|
||||
cleanup = await self.apply_plan(plan)
|
||||
[image] = self.api.db.get_prompt_images(prompt)
|
||||
self.assertEqual(os.path.basename(image["image_path"]), "a.png")
|
||||
self.assertTrue((self.output_dir / "sub" / "c.png").is_file())
|
||||
self.assertEqual(self.linked_files(), [])
|
||||
self.assertEqual(
|
||||
cleanup.outcomes["errors"][0]["detail"], duplicate_files.LOCKED_FILE
|
||||
)
|
||||
|
||||
|
||||
class TestEverySpelling(SafetyTestCase):
|
||||
"""I2: links under the monitor's spelling of a source move too."""
|
||||
|
||||
async def test_links_under_a_symlinked_source_spelling_move(self):
|
||||
self.need_symlinks()
|
||||
self.trio()
|
||||
spelled = Path(self.comfy_dir) / "out-link"
|
||||
os.symlink(str(self.output_dir), str(spelled))
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [str(spelled)]
|
||||
prompt = self.save("Linked through the symlink")
|
||||
self.api.db.link_image_to_prompt(prompt, str(spelled / "b.png"))
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
self.assertIn("1 database link moves", plan["details"][0]["detail"])
|
||||
await self.apply_plan(plan)
|
||||
[image] = self.api.db.get_prompt_images(prompt)
|
||||
self.assertEqual(
|
||||
_locator(image["image_path"]), _locator(self.output_dir / "a.png")
|
||||
)
|
||||
self.assertEqual(self.linked_files(), [])
|
||||
|
||||
|
||||
class TestTagsAndReview(SafetyTestCase):
|
||||
"""I5: tags join the kept copy; pending review items of the copy settle."""
|
||||
|
||||
def tag(self, path, tags):
|
||||
data = Path(path).read_bytes()
|
||||
info = os.stat(path)
|
||||
return self.api.db.replace_image_tags(
|
||||
locator=_locator(path),
|
||||
sha256=hashlib.sha256(data).hexdigest(),
|
||||
size=info.st_size,
|
||||
mtime_ns=info.st_mtime_ns,
|
||||
expected_epoch=self.api.db.get_execution_epoch(),
|
||||
expected_tag_revision=0,
|
||||
tags=tags,
|
||||
)
|
||||
|
||||
def tags_of(self, path):
|
||||
data = Path(path).read_bytes()
|
||||
return self.api.db.read_image_tags(
|
||||
_locator(path), hashlib.sha256(data).hexdigest()
|
||||
)["tags"]
|
||||
|
||||
def queue(self, path):
|
||||
data = Path(path).read_bytes()
|
||||
info = os.stat(path)
|
||||
return self.api.db.record_review_item(
|
||||
locator=_locator(path),
|
||||
spelled_path=os.path.realpath(str(path)),
|
||||
sha256=hashlib.sha256(data).hexdigest(),
|
||||
size=info.st_size,
|
||||
mtime_ns=info.st_mtime_ns,
|
||||
epoch=self.api.db.get_execution_epoch(),
|
||||
resolver_version=1,
|
||||
evidence={"candidates": []},
|
||||
candidate_count=0,
|
||||
limited=True,
|
||||
)
|
||||
|
||||
def pending(self):
|
||||
return self.api.db.review_queue_page(limit=50)["items"]
|
||||
|
||||
async def test_tags_move_as_a_union_and_pending_items_settle(self):
|
||||
self.write("a.png", b"same", mtime=1)
|
||||
self.write("b.png", b"same", mtime=2)
|
||||
keeper, removed = self.output_dir / "a.png", self.output_dir / "b.png"
|
||||
self.tag(keeper, ["kept"])
|
||||
self.tag(removed, ["moved", "kept"])
|
||||
self.assertTrue(self.queue(removed))
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
self.assertIn("image tags join the kept copy", plan["details"][0]["detail"])
|
||||
self.assertIn("1 Needs review item is settled", plan["details"][0]["detail"])
|
||||
await self.apply_plan(plan)
|
||||
self.assertEqual(self.tags_of(keeper), ["kept", "moved"])
|
||||
self.assertEqual(self.pending(), [])
|
||||
|
||||
|
||||
class TestPreviewDoesNotHash(SafetyTestCase):
|
||||
"""M9: the preview binds by size, mtime and inode; the job hashes."""
|
||||
|
||||
async def test_preview_reads_no_file_and_spots_a_touched_one(self):
|
||||
self.trio()
|
||||
job, group = await self.found_group()
|
||||
calls = []
|
||||
real = self.api._calculate_file_hash
|
||||
|
||||
def spy(path):
|
||||
calls.append(path)
|
||||
return real(path)
|
||||
|
||||
self.api._calculate_file_hash = spy
|
||||
os.utime(self.output_dir / "b.png", (5_000_000, 5_000_000))
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
self.assertEqual(calls, [])
|
||||
self.assertEqual(plan["scope"]["count"], 1)
|
||||
self.assertEqual(
|
||||
plan["protected"]["reasons"], ["1 file changed since the search"]
|
||||
)
|
||||
cleanup = await self.apply_plan(plan)
|
||||
self.assertEqual(cleanup.counters["removed"], 1)
|
||||
self.assertGreaterEqual(len(calls), 2) # the job hashes keeper and file
|
||||
|
||||
|
||||
def _fake_stat(real, **changes):
|
||||
"""A stat result with some fields changed (st_dev, st_ino, st_mtime_ns...)."""
|
||||
fields = {name: getattr(real, name) for name in dir(real) if name.startswith("st_")}
|
||||
fields.update(changes)
|
||||
return types.SimpleNamespace(**fields)
|
||||
|
||||
|
||||
class TestAnotherMount(SafetyTestCase):
|
||||
"""Round 2, R1: one file reached through two mounts is never removed.
|
||||
|
||||
The second source stands for the same folder seen through another
|
||||
filesystem: each test reports its file on another ``st_dev``.
|
||||
"""
|
||||
|
||||
MAYBE = "1 file may be the same file reached through another mount or share"
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.real_lstat, self.real_stat = os.lstat, os.stat
|
||||
|
||||
def mount(self, rel, keeper_rel, *, ino="same", mtime_shift=0):
|
||||
"""Copies at ``output/keeper_rel`` and ``renders/rel``; the patch."""
|
||||
second = self.second_source()
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [
|
||||
self._root(self.output_dir),
|
||||
self._root(second),
|
||||
]
|
||||
stamp = 1_700_000_000_123_456_789
|
||||
keeper = self.write(keeper_rel, b"one file")
|
||||
other = self.write(rel, b"one file", folder=second)
|
||||
os.utime(keeper, ns=(stamp, stamp))
|
||||
os.utime(other, ns=(stamp, stamp))
|
||||
target = os.path.normcase(str(other))
|
||||
keeper_ino = self.real_stat(keeper).st_ino
|
||||
|
||||
def changed(real):
|
||||
return _fake_stat(
|
||||
real,
|
||||
st_dev=real.st_dev + 7,
|
||||
st_ino=keeper_ino if ino == "same" else real.st_ino + 1000,
|
||||
st_mtime_ns=real.st_mtime_ns + mtime_shift,
|
||||
)
|
||||
|
||||
def lstat(path, *args, **kwargs):
|
||||
real = self.real_lstat(path, *args, **kwargs)
|
||||
return changed(real) if os.path.normcase(str(path)) == target else real
|
||||
|
||||
def stat(path, *args, **kwargs):
|
||||
real = self.real_stat(path, *args, **kwargs)
|
||||
return changed(real) if os.path.normcase(str(path)) == target else real
|
||||
|
||||
patches = (
|
||||
patch("os.lstat", side_effect=lstat),
|
||||
patch("os.stat", side_effect=stat),
|
||||
)
|
||||
return keeper, other, patches
|
||||
|
||||
async def review_with(self, patches, keeper_rel):
|
||||
# The mount reports its device and file id to the find too.
|
||||
with patches[0], patches[1]:
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, keeper_rel, 0)])
|
||||
return job, group, plan
|
||||
|
||||
async def test_nfs_and_smb_spellings_of_one_share_are_protected(self):
|
||||
# NFS and SMB report their own device and file id; SMB rounds mtime.
|
||||
keeper, other, patches = self.mount(
|
||||
"x.png", "x.png", ino="other", mtime_shift=-89
|
||||
)
|
||||
_job, _group, plan = await self.review_with(patches, "x.png")
|
||||
self.assertEqual(plan["scope"]["count"], 0)
|
||||
self.assertEqual(plan["protected"]["reasons"], [self.MAYBE])
|
||||
self.assertTrue(other.is_file() and keeper.is_file())
|
||||
|
||||
async def test_a_loopback_mount_with_the_same_inode_is_protected(self):
|
||||
keeper, other, patches = self.mount("y.png", "sub/y.png", ino="same")
|
||||
_job, _group, plan = await self.review_with(patches, "sub/y.png")
|
||||
self.assertEqual(plan["protected"]["reasons"], [self.MAYBE])
|
||||
|
||||
async def test_a_fuse_pool_path_is_protected(self):
|
||||
# mergerfs, rclone or sshfs synthesize both the device and the inode.
|
||||
keeper, other, patches = self.mount("x.png", "x.png", ino="other")
|
||||
_job, _group, plan = await self.review_with(patches, "x.png")
|
||||
self.assertEqual(plan["protected"]["reasons"], [self.MAYBE])
|
||||
|
||||
async def test_a_true_copy_on_another_disk_is_still_offered(self):
|
||||
keeper, other, patches = self.mount("elsewhere.png", "x.png", ino="other")
|
||||
_job, _group, plan = await self.review_with(patches, "x.png")
|
||||
self.assertEqual(plan["scope"]["count"], 1)
|
||||
|
||||
async def test_the_job_checks_again_before_the_unlink(self):
|
||||
keeper, other, patches = self.mount("x.png", "x.png", ino="other")
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "x.png", 0)])
|
||||
self.assertEqual(plan["scope"]["count"], 1) # looked separate then
|
||||
with patches[0], patches[1]:
|
||||
cleanup = await self.apply_plan(plan)
|
||||
self.assertEqual(cleanup.counters["removed"], 0)
|
||||
rows = {row["path"]: row["detail"] for row in cleanup.outcomes["kept"]}
|
||||
self.assertEqual(rows["x.png"], duplicate_files.MAYBE_SAME_FILE)
|
||||
self.assertTrue(other.is_file() and keeper.is_file())
|
||||
|
||||
async def test_a_volume_without_file_ids_is_named_truthfully(self):
|
||||
keeper, other, patches = self.mount("x.png", "x.png")
|
||||
|
||||
def zero(path, *args, **kwargs):
|
||||
return _fake_stat(self.real_stat(path), st_dev=0, st_ino=0)
|
||||
|
||||
with patch("os.lstat", side_effect=zero), patch("os.stat", side_effect=zero):
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "x.png", 0)])
|
||||
self.assertEqual(
|
||||
plan["protected"]["reasons"],
|
||||
["1 file: can't confirm this is a separate file on this drive"],
|
||||
)
|
||||
|
||||
async def test_the_kept_copy_is_held_open_during_the_unlink(self):
|
||||
self.trio()
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
real = duplicate_files._unlink
|
||||
held = []
|
||||
|
||||
def unlink(path):
|
||||
held.append(bool(open_handles_under(self.output_dir / "a.png")))
|
||||
real(path)
|
||||
|
||||
if not os.path.isdir("/proc/self/fd"): # pragma: no cover - no /proc
|
||||
self.skipTest("open handles are read from /proc")
|
||||
with patch.object(duplicate_files, "_unlink", side_effect=unlink):
|
||||
await self.apply_plan(plan)
|
||||
self.assertEqual(held, [True, True])
|
||||
|
||||
|
||||
class TestNestedSpellings(SafetyTestCase):
|
||||
"""Round 2, R2: spellings and favorites by containment, not equality."""
|
||||
|
||||
async def test_a_parent_spelling_of_a_nested_source(self):
|
||||
# The reviewer's probe: sources [output/sub, out-link -> output].
|
||||
self.need_symlinks()
|
||||
self.write("sub/a.png", b"same", mtime=1)
|
||||
self.write("sub/b.png", b"same", mtime=2)
|
||||
out_link = Path(self.comfy_dir) / "out-link"
|
||||
os.symlink(str(self.output_dir), str(out_link))
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [
|
||||
self._root(self.output_dir / "sub"),
|
||||
str(out_link),
|
||||
]
|
||||
prompt = self.save("Linked before sub was a source")
|
||||
self.api.db.link_image_to_prompt(prompt, str(out_link / "sub" / "b.png"))
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
self.assertIn("1 database link moves", plan["details"][0]["detail"])
|
||||
await self.apply_plan(plan)
|
||||
[image] = self.api.db.get_prompt_images(prompt)
|
||||
self.assertEqual(
|
||||
_locator(image["image_path"]), _locator(self.output_dir / "sub" / "a.png")
|
||||
)
|
||||
self.assertEqual(self.linked_files(), [])
|
||||
|
||||
async def test_nested_sources_subfolder_first_with_a_parent_favorite(self):
|
||||
self.write("sub/a.png", b"same", mtime=1)
|
||||
self.write("sub/b.png", b"same", mtime=2)
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [
|
||||
self._root(self.output_dir / "sub"),
|
||||
self._root(self.output_dir),
|
||||
]
|
||||
parent_key = duplicate_files.favorite_root_key(self.output_dir)
|
||||
self.api.db.set_image_favorite(parent_key, "sub/b.png")
|
||||
prompt = self.save("Linked under the parent")
|
||||
self.api.db.link_image_to_prompt(prompt, str(self.output_dir / "sub" / "b.png"))
|
||||
job, group = await self.found_group()
|
||||
self.assertEqual({f["root_index"] for f in group["files"]}, {0})
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
await self.apply_plan(plan)
|
||||
self.assertEqual(self.api.db.get_root_favorites(parent_key), {"sub/a.png"})
|
||||
self.assertEqual(self.linked_files(), [])
|
||||
|
||||
|
||||
class TestLockedRowSaysWhereLinksAre(SafetyTestCase):
|
||||
"""Round 2, m1."""
|
||||
|
||||
async def test_the_locked_row_names_the_kept_copy(self):
|
||||
self.trio()
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
with (
|
||||
patch.object(
|
||||
duplicate_files, "_unlink", side_effect=PermissionError(13, "denied")
|
||||
),
|
||||
self.assertLogs("prompt_manager", "WARNING"),
|
||||
):
|
||||
cleanup = await self.apply_plan(plan)
|
||||
self.assertIn(
|
||||
"its links now name the kept copy", cleanup.outcomes["errors"][0]["detail"]
|
||||
)
|
||||
|
||||
|
||||
class TestRound3(SafetyTestCase):
|
||||
"""Re-review 2: n1 (sources at different depths) and n2 (the kept copy
|
||||
must still be the held file right before the unlink)."""
|
||||
|
||||
MAYBE = TestAnotherMount.MAYBE
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.real_lstat, self.real_stat = os.lstat, os.stat
|
||||
|
||||
mount = TestAnotherMount.mount
|
||||
review_with = TestAnotherMount.review_with
|
||||
|
||||
async def test_a_deeper_source_on_another_mount_is_protected(self):
|
||||
# The reviewer's probe: ComfyUI/output -> /mnt/pool/output and
|
||||
# /mnt/disk1/output/2024 as an extra root, FUSE-style inodes.
|
||||
_keeper, other, patches = self.mount("x.png", "2024/x.png", ino="other")
|
||||
_job, _group, plan = await self.review_with(patches, "2024/x.png")
|
||||
self.assertEqual(plan["protected"]["reasons"], [self.MAYBE])
|
||||
self.assertTrue(other.is_file())
|
||||
|
||||
async def test_a_shallower_source_on_another_mount_is_protected(self):
|
||||
_keeper, _other, patches = self.mount("2024/x.png", "x.png", ino="other")
|
||||
_job, _group, plan = await self.review_with(patches, "x.png")
|
||||
self.assertEqual(plan["protected"]["reasons"], [self.MAYBE])
|
||||
|
||||
async def test_only_whole_folder_names_match(self):
|
||||
_keeper, _other, patches = self.mount("photo-x.png", "x.png", ino="other")
|
||||
_job, _group, plan = await self.review_with(patches, "x.png")
|
||||
self.assertEqual(plan["scope"]["count"], 1)
|
||||
|
||||
async def removal_while_the_keeper_goes(self, change):
|
||||
self.write("a.png", b"same", mtime=1)
|
||||
self.write("b.png", b"same", mtime=2)
|
||||
prompt = self.save("On b")
|
||||
self.api.db.link_image_to_prompt(prompt, str(self.output_dir / "b.png"))
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
real = self.api.db.move_copy_records
|
||||
refused = []
|
||||
|
||||
def move(*args, **kwargs):
|
||||
counts = real(*args, **kwargs)
|
||||
if args[0]["kept_path"].endswith("a.png"):
|
||||
try:
|
||||
change(self.output_dir / "a.png")
|
||||
except PermissionError:
|
||||
# Windows: the job holds the kept copy open (no
|
||||
# FILE_SHARE_DELETE), so it cannot be deleted or replaced.
|
||||
if os.name != "nt":
|
||||
raise
|
||||
refused.append(True)
|
||||
return counts
|
||||
|
||||
with patch.object(self.api.db, "move_copy_records", side_effect=move):
|
||||
cleanup = await self.apply_plan(plan)
|
||||
if refused: # pragma: no cover - native Windows
|
||||
self.assertEqual((self.output_dir / "a.png").read_bytes(), b"same")
|
||||
self.assertEqual(cleanup.counters["removed"], 1)
|
||||
self.assertEqual(self.linked_files(), [])
|
||||
return
|
||||
self.assertTrue((self.output_dir / "b.png").is_file())
|
||||
self.assertEqual(cleanup.counters["removed"], 0)
|
||||
rows = {row["path"]: row["detail"] for row in cleanup.outcomes["kept"]}
|
||||
self.assertEqual(rows["b.png"], duplicate_files.KEEPER_WENT)
|
||||
self.assertEqual(self.linked_files(), [])
|
||||
[image] = self.api.db.get_prompt_images(prompt)
|
||||
self.assertEqual(os.path.basename(image["image_path"]), "b.png")
|
||||
self.assertIn("1 not removed", cleanup.status_line)
|
||||
|
||||
async def test_a_kept_copy_deleted_during_the_move_keeps_the_file(self):
|
||||
await self.removal_while_the_keeper_goes(lambda path: path.unlink())
|
||||
|
||||
async def test_a_kept_copy_replaced_during_the_move_keeps_the_file(self):
|
||||
def replace(path):
|
||||
spare = path.with_name("spare.tmp")
|
||||
spare.write_bytes(b"other bytes")
|
||||
os.replace(spare, path)
|
||||
|
||||
await self.removal_while_the_keeper_goes(replace)
|
||||
@@ -401,7 +401,7 @@ class HistoryOutputAPITests(ImageAPITestCase):
|
||||
self.assertEqual(response.status, 400)
|
||||
await response.read()
|
||||
|
||||
async def test_media_ranges_head_and_a_preview_made_on_first_view(self):
|
||||
async def test_media_ranges_head_and_existing_thumbnail_without_generation(self):
|
||||
path = self.output_dir / "clip.mp4"
|
||||
payload = bytes(range(256)) * 600
|
||||
path.write_bytes(payload)
|
||||
@@ -430,14 +430,9 @@ class HistoryOutputAPITests(ImageAPITestCase):
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(int(response.headers["Content-Length"]), len(payload))
|
||||
self.assertEqual(await response.read(), b"")
|
||||
# A missing preview is made on first view (#249): an undecodable video
|
||||
# gets the placeholder frame, never the whole file.
|
||||
preview = self.output_dir / "thumbnails" / "clip_thumb.jpg"
|
||||
response = await self.current(descriptor["thumbnail_url"])
|
||||
self.assertEqual(response.headers["Content-Type"], "image/jpeg")
|
||||
self.assertEqual(await response.read(), preview.read_bytes())
|
||||
self.assertLess(preview.stat().st_size, len(payload))
|
||||
make_png(preview)
|
||||
self.assertEqual(await response.read(), payload)
|
||||
preview = make_png(self.output_dir / "thumbnails" / "clip_thumb.jpg")
|
||||
response = await self.current(descriptor["thumbnail_url"])
|
||||
self.assertEqual(await response.read(), preview.read_bytes())
|
||||
|
||||
@@ -1116,73 +1111,6 @@ class HistoryOutputAPITests(ImageAPITestCase):
|
||||
response = await self.current(descriptor["thumbnail_url"])
|
||||
self.assert_error(response, await response.json(), 403, "not-authorized")
|
||||
|
||||
async def test_missing_preview_is_never_made_through_a_linked_out_folder(self):
|
||||
original = make_png(self.output_dir / "image.png", size=(64, 64))
|
||||
elsewhere = self.temp_dir / "elsewhere"
|
||||
elsewhere.mkdir()
|
||||
(self.output_dir / "thumbnails").symlink_to(elsewhere, target_is_directory=True)
|
||||
start = self.capture([self.ref()])
|
||||
descriptor = await self.media(start)
|
||||
response = await self.current(descriptor["thumbnail_url"])
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(await response.read(), original.read_bytes())
|
||||
self.assertEqual(os.listdir(elsewhere), [])
|
||||
|
||||
async def test_missing_preview_is_never_made_through_a_linked_sub_folder(self):
|
||||
# M6: every link segment of the preview's folder is checked.
|
||||
original = make_png(self.output_dir / "sub" / "image.png", size=(64, 64))
|
||||
elsewhere = self.temp_dir / "elsewhere"
|
||||
elsewhere.mkdir()
|
||||
(self.output_dir / "thumbnails").mkdir()
|
||||
(self.output_dir / "thumbnails" / "sub").symlink_to(
|
||||
elsewhere, target_is_directory=True
|
||||
)
|
||||
start = self.capture([self.ref(subfolder="sub")])
|
||||
descriptor = await self.media(start)
|
||||
response = await self.current(descriptor["thumbnail_url"])
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(await response.read(), original.read_bytes())
|
||||
self.assertEqual(os.listdir(elsewhere), [])
|
||||
|
||||
async def test_a_video_preview_that_cannot_be_made_is_the_placeholder(self):
|
||||
# I1: never the video bytes inside an <img>.
|
||||
from py.api.jobs import RESTORE_JOB_KIND
|
||||
|
||||
(self.output_dir / "clip.webm").write_bytes(b"AV1 stream no decoder here")
|
||||
start = self.capture([self.ref(filename="clip.webm")])
|
||||
descriptor = await self.media(start)
|
||||
gate = asyncio.Event()
|
||||
|
||||
async def runner(job):
|
||||
await gate.wait()
|
||||
|
||||
job = self.api._jobs.start(
|
||||
RESTORE_JOB_KIND,
|
||||
"restore",
|
||||
admission="records",
|
||||
cancellable=False,
|
||||
runner=runner,
|
||||
)
|
||||
try:
|
||||
response = await self.current(descriptor["thumbnail_url"])
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.headers["Content-Type"], "image/jpeg")
|
||||
self.assertTrue((await response.read()).startswith(b"\xff\xd8"))
|
||||
finally:
|
||||
gate.set()
|
||||
await asyncio.wait_for(job.task, 5)
|
||||
self.assertFalse((self.output_dir / "thumbnails").exists())
|
||||
|
||||
async def test_a_preview_that_cannot_be_made_serves_the_original(self):
|
||||
broken = self.output_dir / "image.png"
|
||||
broken.write_bytes(b"\x89PNG\r\n\x1a\n not an image")
|
||||
start = self.capture([self.ref()])
|
||||
descriptor = await self.media(start)
|
||||
response = await self.current(descriptor["thumbnail_url"])
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(await response.read(), broken.read_bytes())
|
||||
self.assertFalse((self.output_dir / "thumbnails" / "image_thumb.png").exists())
|
||||
|
||||
async def test_disappearance_before_open_is_missing_and_closes_handle(
|
||||
self,
|
||||
):
|
||||
|
||||
+16
-69
@@ -80,8 +80,6 @@ class ImageAPITestCase(AioHTTPTestCase):
|
||||
|
||||
async def tearDownAsync(self):
|
||||
await super().tearDownAsync() # closes the aiohttp test client
|
||||
# The automatic preview worker writes into the temp tree: stop it first.
|
||||
self.api.close_auto_thumbnails()
|
||||
self._anchor_patch.stop()
|
||||
# Windows refuses to unlink a database that still has open handles.
|
||||
self.api.db.close_all()
|
||||
@@ -1567,65 +1565,39 @@ class TestVideoThumbnailFallbacks(ImageRouteCoverageCase):
|
||||
with Image.open(self.thumb) as img:
|
||||
self.assertEqual(img.format, "JPEG")
|
||||
|
||||
def test_cv2_cannot_open_or_read_falls_through_to_ffmpeg_then_placeholder(self):
|
||||
# #249 I1: an OpenCV that cannot decode the clip (AV1) no longer ends
|
||||
# the run; ffmpeg and then the placeholder are tried.
|
||||
for fake in (self._fake_cv2(opened=False), self._fake_cv2(read_ok=False)):
|
||||
self.thumb.unlink(missing_ok=True)
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": fake}),
|
||||
patch("subprocess.Popen", side_effect=FileNotFoundError("ffmpeg")),
|
||||
):
|
||||
ok = self.api._generate_video_thumbnail(
|
||||
self.video, self.thumb, self.SIZE
|
||||
)
|
||||
self.assertTrue(ok)
|
||||
with Image.open(self.thumb) as img:
|
||||
self.assertEqual((img.format, img.size), ("JPEG", self.SIZE))
|
||||
def test_cv2_cannot_open_or_read(self):
|
||||
with patch.dict(sys.modules, {"cv2": self._fake_cv2(opened=False)}):
|
||||
closed = self.api._generate_video_thumbnail(
|
||||
self.video, self.thumb, self.SIZE
|
||||
)
|
||||
with patch.dict(sys.modules, {"cv2": self._fake_cv2(read_ok=False)}):
|
||||
unread = self.api._generate_video_thumbnail(
|
||||
self.video, self.thumb, self.SIZE
|
||||
)
|
||||
self.assertEqual((closed, unread), (False, False))
|
||||
|
||||
def test_cv2_crash_is_contained(self):
|
||||
cv2 = MagicMock()
|
||||
cv2.VideoCapture.side_effect = RuntimeError("driver crash")
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": cv2}),
|
||||
patch("subprocess.Popen", side_effect=FileNotFoundError("ffmpeg")),
|
||||
):
|
||||
with patch.dict(sys.modules, {"cv2": cv2}):
|
||||
ok = self.api._generate_video_thumbnail(self.video, self.thumb, self.SIZE)
|
||||
self.assertTrue(ok, "the placeholder is written")
|
||||
|
||||
def test_cv2_failure_keeps_an_existing_preview_on_rebuild(self):
|
||||
make_png(self.thumb)
|
||||
before = self.thumb.read_bytes()
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": self._fake_cv2(opened=False)}),
|
||||
patch("subprocess.Popen", side_effect=FileNotFoundError("ffmpeg")),
|
||||
):
|
||||
ok = self.api._generate_video_thumbnail(
|
||||
self.video, self.thumb, self.SIZE, keep_existing=True
|
||||
)
|
||||
self.assertFalse(ok)
|
||||
self.assertEqual(self.thumb.read_bytes(), before)
|
||||
|
||||
def test_ffmpeg_absent_yields_placeholder(self):
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": None}),
|
||||
patch("subprocess.Popen", side_effect=FileNotFoundError("ffmpeg")),
|
||||
patch("subprocess.run", side_effect=FileNotFoundError("ffmpeg")),
|
||||
):
|
||||
ok = self.api._generate_video_thumbnail(self.video, self.thumb, self.SIZE)
|
||||
self.assertTrue(ok)
|
||||
with Image.open(self.thumb) as img:
|
||||
self.assertEqual(img.size, self.SIZE)
|
||||
|
||||
@staticmethod
|
||||
def _process(returncode, stderr=""):
|
||||
process = MagicMock(returncode=returncode)
|
||||
process.communicate.return_value = ("", stderr)
|
||||
return process
|
||||
|
||||
def test_ffmpeg_failure_yields_placeholder(self):
|
||||
failed = MagicMock(returncode=1, stderr="no stream")
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": None}),
|
||||
patch("subprocess.Popen", return_value=self._process(1, "no stream")),
|
||||
patch("subprocess.run", return_value=failed),
|
||||
):
|
||||
ok = self.api._generate_video_thumbnail(self.video, self.thumb, self.SIZE)
|
||||
self.assertTrue(ok)
|
||||
@@ -1634,41 +1606,16 @@ class TestVideoThumbnailFallbacks(ImageRouteCoverageCase):
|
||||
def test_ffmpeg_success(self):
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": None}),
|
||||
patch("subprocess.Popen", return_value=self._process(0)) as run,
|
||||
patch("subprocess.run", return_value=MagicMock(returncode=0)) as run,
|
||||
):
|
||||
ok = self.api._generate_video_thumbnail(self.video, self.thumb, self.SIZE)
|
||||
self.assertTrue(ok)
|
||||
self.assertEqual(run.call_args.args[0][0], "ffmpeg")
|
||||
|
||||
def test_ffmpeg_is_tracked_and_a_hold_before_the_rename_writes_nothing(self):
|
||||
tracked = []
|
||||
|
||||
def writing(cmd, **kwargs):
|
||||
Path(cmd[-1]).write_bytes(b"\xff\xd8 frame")
|
||||
return self._process(0)
|
||||
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": None}),
|
||||
patch("subprocess.Popen", side_effect=writing),
|
||||
):
|
||||
with self.assertRaises(images_module.PreviewWriteHeld):
|
||||
self.api._generate_video_thumbnail(
|
||||
self.video,
|
||||
self.thumb,
|
||||
self.SIZE,
|
||||
guard=lambda: False,
|
||||
track=tracked.append,
|
||||
)
|
||||
self.assertEqual(len(tracked), 1)
|
||||
self.assertFalse(self.thumb.exists())
|
||||
self.assertEqual(
|
||||
[p.name for p in self.thumb.parent.iterdir() if "part" in p.name], []
|
||||
)
|
||||
|
||||
def test_placeholder_failure_returns_false(self):
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": None}),
|
||||
patch("subprocess.Popen", side_effect=FileNotFoundError),
|
||||
patch("subprocess.run", side_effect=FileNotFoundError),
|
||||
patch.object(
|
||||
images_module.Image, "new", side_effect=RuntimeError("no PIL")
|
||||
),
|
||||
|
||||
@@ -1,427 +0,0 @@
|
||||
"""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)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -397,20 +397,15 @@ class TestQueueReads(ReviewRoutesTestCase):
|
||||
self._no_paths()
|
||||
|
||||
async def test_queue_cards_use_the_images_listing_thumbnail(self):
|
||||
# M10: the card's media is the Images listing's thumbnail. Without a
|
||||
# preview yet it is still the thumbnail URL (made on first view, #249).
|
||||
# M10: a known thumbnail is the card's media, as on the Images page;
|
||||
# without one the card falls back to the original.
|
||||
self._require_routes()
|
||||
from py.api.images import _output_image_entry
|
||||
|
||||
path = self._ambiguous("sub/f0.png")
|
||||
self._record(path)
|
||||
item = (await self._list())["items"][0]
|
||||
scope = item["target"]["root_scope"]
|
||||
self.assertEqual(
|
||||
item["thumbnail_url"],
|
||||
"/prompt_manager/images/serve/thumbnails/sub/f0_thumb.png"
|
||||
f"?root=0&root_scope={scope}",
|
||||
)
|
||||
self.assertIsNone(item["thumbnail_url"])
|
||||
thumb = self.output_dir / "thumbnails" / "sub" / "f0_thumb.png"
|
||||
thumb.parent.mkdir(parents=True)
|
||||
write_png(thumb, prompt=ambiguous_graph("cat", "dog"), color=(1, 2, 3))
|
||||
|
||||
@@ -12,7 +12,6 @@ import asyncio
|
||||
import json
|
||||
import sqlite3
|
||||
import threading
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
from tests.test_api_review import ReviewRoutesTestCase
|
||||
@@ -87,16 +86,11 @@ class ReviewEventsTestCase(ReviewRoutesTestCase):
|
||||
return path
|
||||
|
||||
async def _record_in_thread(self, path):
|
||||
"""Record as the scan does: from a worker thread, not the loop."""
|
||||
from py.api.admin import _scan_roots
|
||||
|
||||
"""Record as the monitor does: from a worker thread, not the loop."""
|
||||
self._roots()
|
||||
roots = _scan_roots(self.api._get_all_output_dirs())
|
||||
outcomes = []
|
||||
thread = threading.Thread(
|
||||
target=lambda: outcomes.append(
|
||||
self.api._record_review_sync(Path(str(path)), roots, inspection=True)
|
||||
)
|
||||
target=lambda: outcomes.append(self.api.record_review_file(str(path)))
|
||||
)
|
||||
thread.start()
|
||||
await asyncio.get_running_loop().run_in_executor(None, thread.join)
|
||||
@@ -104,7 +98,7 @@ class ReviewEventsTestCase(ReviewRoutesTestCase):
|
||||
|
||||
|
||||
class TestReviewEvents(ReviewEventsTestCase):
|
||||
async def test_scan_recording_sends_one_change(self):
|
||||
async def test_monitor_recording_sends_one_change(self):
|
||||
self._require_events()
|
||||
stream = await self._stream()
|
||||
self.assertEqual(await self._record_in_thread(self._queueable()), "recorded")
|
||||
|
||||
@@ -1,424 +0,0 @@
|
||||
"""Store each workflow once, as a maintenance job (py/api/workflow_storage.py).
|
||||
|
||||
``POST /prompt_manager/maintenance/preview {operation:"workflow_storage"}``
|
||||
reviews the rewrite; applying it starts a ``records`` job that takes the
|
||||
verified safety backup first, rewrites old-form rows batch by batch (Cancel
|
||||
stops after the current batch) and then starts "Reclaim space", an
|
||||
``exclusive`` job with its own disk check. No image file is ever opened.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import builtins
|
||||
import contextlib
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
from unittest.mock import patch
|
||||
|
||||
from PIL import Image, PngImagePlugin
|
||||
|
||||
from database import workflow_storage as ws
|
||||
from database.workflow_documents import DocumentReader, resolve_image_row, store_text
|
||||
from py.api import maintenance
|
||||
from py.api.jobs import MAINTENANCE_JOB_KIND
|
||||
|
||||
try:
|
||||
from tests.test_api_maintenance import ApplyTestCase
|
||||
from tests.test_workflow_documents import sample_prompt_graph, sample_workflow
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from test_api_maintenance import ApplyTestCase
|
||||
from test_workflow_documents import sample_prompt_graph, sample_workflow
|
||||
|
||||
WAIT = 20
|
||||
OPERATION = "workflow_storage"
|
||||
|
||||
|
||||
def _digest(path):
|
||||
with open(path, "rb") as handle:
|
||||
return hashlib.sha256(handle.read()).hexdigest()
|
||||
|
||||
|
||||
class WorkflowStorageTestCase(ApplyTestCase):
|
||||
def old_library(self, count=12, media=False):
|
||||
"""Old-form rows (text inline), four images per run."""
|
||||
prompt_id = self.save("a cat")
|
||||
conn = self.api.db.model.get_connection()
|
||||
for index in range(count):
|
||||
run = index // 4
|
||||
path = self.output_dir / f"{index}.png"
|
||||
workflow = json.dumps(sample_workflow(run, nodes=120))
|
||||
if media:
|
||||
info = PngImagePlugin.PngInfo()
|
||||
info.add_text("workflow", workflow)
|
||||
with Image.new("RGB", (4, 4)) as image:
|
||||
image.save(path, pnginfo=info)
|
||||
with conn:
|
||||
conn.execute(
|
||||
"INSERT INTO generated_images (prompt_id, image_path, filename, "
|
||||
"file_path, workflow_data, prompt_metadata, parameters) "
|
||||
"VALUES (?, ?, ?, ?, ?, ?, '{}')",
|
||||
(
|
||||
prompt_id,
|
||||
str(path),
|
||||
path.name,
|
||||
str(path),
|
||||
workflow,
|
||||
json.dumps(sample_prompt_graph(run)),
|
||||
),
|
||||
)
|
||||
return prompt_id
|
||||
|
||||
def texts(self):
|
||||
conn = self.api.db.model.get_connection()
|
||||
reader = DocumentReader(conn)
|
||||
return {
|
||||
row["id"]: resolve_image_row(reader, row)
|
||||
for row in conn.execute(
|
||||
"SELECT id, workflow_data, prompt_metadata, workflow_document_id, "
|
||||
"prompt_document_id FROM generated_images ORDER BY id"
|
||||
)
|
||||
}
|
||||
|
||||
def inline_count(self):
|
||||
conn = self.api.db.model.get_connection()
|
||||
return conn.execute(
|
||||
"SELECT COUNT(*) FROM generated_images WHERE workflow_data IS NOT NULL"
|
||||
).fetchone()[0]
|
||||
|
||||
def reclaim_job(self):
|
||||
jobs = [job for job in self.api._jobs.list() if job.title == "Reclaim space"]
|
||||
return jobs[0] if jobs else None
|
||||
|
||||
async def finish(self, job):
|
||||
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
|
||||
reclaim = self.reclaim_job()
|
||||
if reclaim is not None and reclaim.task is not None:
|
||||
await asyncio.wait_for(asyncio.shield(reclaim.task), WAIT)
|
||||
return reclaim
|
||||
|
||||
async def start(self):
|
||||
plan = await self.preview(OPERATION)
|
||||
body = await self.apply(plan["token"])
|
||||
return plan, body["job"], self.api._jobs.get(body["job"]["id"])
|
||||
|
||||
|
||||
class TestWorkflowStorageJob(WorkflowStorageTestCase):
|
||||
async def test_preview_reviews_the_rewrite_and_writes_nothing(self):
|
||||
self.old_library()
|
||||
db_path = self.api.db.model.db_path
|
||||
before = _digest(db_path), _digest(db_path + "-wal")
|
||||
plan = await self.preview(OPERATION)
|
||||
self.assertEqual((_digest(db_path), _digest(db_path + "-wal")), before)
|
||||
self.assertEqual(plan["operation"], OPERATION)
|
||||
self.assertEqual(plan["scope"]["count"], 12)
|
||||
self.assertEqual(plan["scope"]["label"], "12 image records to rewrite")
|
||||
self.assertFalse(plan["media"]["deletes_files"])
|
||||
self.assertEqual(
|
||||
plan["media"]["label"],
|
||||
"No image or video file is opened or changed; only the database.",
|
||||
)
|
||||
self.assertEqual(plan["backup"]["policy"], "required")
|
||||
self.assertEqual(
|
||||
plan["backup"]["destination_label"],
|
||||
"prompts.db.pre-workflow-storage_<date and time>",
|
||||
)
|
||||
self.assertGreater(plan["backup"]["expected_bytes"], 0)
|
||||
self.assertTrue(plan["disk"]["sufficient"])
|
||||
|
||||
async def test_apply_backs_up_rewrites_then_reclaims_space(self):
|
||||
prompt_id = self.old_library(count=40)
|
||||
before = self.texts()
|
||||
listed = await self.client.request(
|
||||
"GET", f"/prompt_manager/prompts/{prompt_id}/images"
|
||||
)
|
||||
listed_before = await listed.json()
|
||||
size_before = self.database_bytes()
|
||||
plan, summary, job = await self.start()
|
||||
self.assertEqual(summary["kind"], MAINTENANCE_JOB_KIND)
|
||||
self.assertEqual(summary["title"], "Store each workflow once")
|
||||
self.assertEqual(summary["admission"], "records")
|
||||
self.assertTrue(summary["cancellable"])
|
||||
self.assertEqual(summary["scope"], {"operation": OPERATION})
|
||||
reclaim = await self.finish(job)
|
||||
self.assertEqual(job.state, "finished", job.log)
|
||||
self.assertEqual(job.counters, {"changed": 40, "protected": 0, "errors": 0})
|
||||
self.assertEqual(job.total, 40)
|
||||
self.assertEqual(job.processed, 40)
|
||||
self.assertEqual(self.inline_count(), 0)
|
||||
self.assertEqual(self.texts(), before)
|
||||
self.assertEqual(self.backups(), []) # never in the pruned pool
|
||||
[backup] = [
|
||||
name
|
||||
for name in os.listdir(self.tmpdir)
|
||||
if name.startswith("prompts.db.pre-workflow-storage_")
|
||||
]
|
||||
self.assertIn(f"Pre-migration backup {backup} verified", "\n".join(job.log))
|
||||
self.assertEqual(
|
||||
job.status_line,
|
||||
f"Rewrote 40 image records · Pre-migration backup {backup}",
|
||||
)
|
||||
# The safety backup holds the old form, verified.
|
||||
conn = sqlite3.connect(os.path.join(self.tmpdir, backup))
|
||||
try:
|
||||
inline = conn.execute(
|
||||
"SELECT COUNT(*) FROM generated_images WHERE workflow_data IS NOT NULL"
|
||||
).fetchone()[0]
|
||||
finally:
|
||||
conn.close()
|
||||
self.assertEqual(inline, 40)
|
||||
# Reclaim space ran alone afterwards and gave the space back.
|
||||
self.assertIsNotNone(reclaim)
|
||||
self.assertEqual(reclaim.admission, "exclusive")
|
||||
self.assertFalse(reclaim.cancellable)
|
||||
self.assertEqual(reclaim.state, "finished", reclaim.log)
|
||||
self.assertIn("Reclaim space started", "\n".join(job.log))
|
||||
self.assertLess(self.database_bytes(), size_before / 3)
|
||||
wal = self.api.db.model.db_path + "-wal"
|
||||
self.assertEqual(os.path.getsize(wal) if os.path.exists(wal) else 0, 0)
|
||||
self.assertRegex(reclaim.status_line, r"^Reclaimed [\d.,]+ MB · ")
|
||||
listed = await self.client.request(
|
||||
"GET", f"/prompt_manager/prompts/{prompt_id}/images"
|
||||
)
|
||||
self.assertEqual(await listed.json(), listed_before)
|
||||
self.assertEqual(self.strays(), [backup]) # only the kept backup
|
||||
self.assertNotIn(self.tmpdir, json.dumps(job.report()))
|
||||
|
||||
async def test_a_backup_failure_rewrites_nothing(self):
|
||||
self.old_library()
|
||||
from py.api import workflow_storage as api_storage
|
||||
|
||||
with patch.object(
|
||||
api_storage,
|
||||
"take_safety_backup",
|
||||
side_effect=maintenance.SafetyBackupFailed("no"),
|
||||
):
|
||||
_plan, _summary, job = await self.start()
|
||||
await self.finish(job)
|
||||
self.assertEqual(job.state, "failed")
|
||||
self.assertEqual(job.last_event["code"], "backup-failed")
|
||||
self.assertEqual(self.inline_count(), 12)
|
||||
self.assertIsNone(self.reclaim_job())
|
||||
|
||||
async def test_cancel_stops_after_the_current_batch_and_a_rerun_continues(self):
|
||||
self.old_library()
|
||||
before = self.texts()
|
||||
real = self.api.db.migrate_workflow_batch
|
||||
holder = {}
|
||||
|
||||
def first_then_cancel(*args, **kwargs):
|
||||
result = real(*args, **kwargs)
|
||||
holder["job"].cancelled = True # Cancel pressed during batch one
|
||||
return result
|
||||
|
||||
with (
|
||||
patch.object(ws, "BATCH_ROWS", 5),
|
||||
patch.object(
|
||||
self.api.db, "migrate_workflow_batch", side_effect=first_then_cancel
|
||||
),
|
||||
):
|
||||
plan = await self.preview(OPERATION)
|
||||
body = await self.apply(plan["token"])
|
||||
holder["job"] = job = self.api._jobs.get(body["job"]["id"])
|
||||
await self.finish(job)
|
||||
self.assertEqual(job.state, "cancelled")
|
||||
self.assertEqual(job.counters["changed"], 5)
|
||||
self.assertEqual(self.inline_count(), 7)
|
||||
self.assertIn("run it again to continue", job.status_line)
|
||||
self.assertIsNone(self.reclaim_job())
|
||||
self.assertEqual(self.texts(), before)
|
||||
_plan, _summary, again = await self.start()
|
||||
await self.finish(again)
|
||||
self.assertEqual(again.state, "finished", again.log)
|
||||
self.assertEqual(again.counters["changed"], 7)
|
||||
self.assertEqual(self.inline_count(), 0)
|
||||
self.assertEqual(self.texts(), before)
|
||||
|
||||
async def test_a_failed_batch_fails_the_job_and_every_record_still_reads(self):
|
||||
self.old_library()
|
||||
before = self.texts()
|
||||
real = self.api.db.migrate_workflow_batch
|
||||
calls = []
|
||||
|
||||
def failing(*args, **kwargs):
|
||||
calls.append(1)
|
||||
if len(calls) == 2:
|
||||
raise sqlite3.OperationalError("disk I/O error")
|
||||
return real(*args, **kwargs)
|
||||
|
||||
with (
|
||||
patch.object(ws, "BATCH_ROWS", 5),
|
||||
patch.object(self.api.db, "migrate_workflow_batch", side_effect=failing),
|
||||
):
|
||||
_plan, _summary, job = await self.start()
|
||||
await self.finish(job)
|
||||
self.assertEqual(job.state, "failed")
|
||||
self.assertIn(
|
||||
"The batch in progress was rolled back; every record reads as before.",
|
||||
job.status_line,
|
||||
)
|
||||
self.assertNotIn("disk I/O", job.status_line)
|
||||
self.assertEqual(job.counters["changed"], 5)
|
||||
self.assertEqual(job.counters["errors"], 1)
|
||||
self.assertEqual(self.texts(), before)
|
||||
self.assertIsNone(self.reclaim_job())
|
||||
|
||||
async def test_reclaim_space_that_cannot_start_says_why(self):
|
||||
self.old_library()
|
||||
release = asyncio.Event()
|
||||
|
||||
async def held(job):
|
||||
await release.wait()
|
||||
|
||||
reader = self.api._jobs.start(
|
||||
"health", "Health check", admission="read", runner=held
|
||||
)
|
||||
try:
|
||||
_plan, _summary, job = await self.start()
|
||||
await self.finish(job)
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.wait_for(reader.task, WAIT)
|
||||
self.assertEqual(job.state, "finished", job.log)
|
||||
self.assertIsNone(self.reclaim_job())
|
||||
self.assertIn(
|
||||
"Reclaim space did not start: Health check is running. "
|
||||
"Run Optimize database to give the space back.",
|
||||
"\n".join(job.log),
|
||||
)
|
||||
self.assertIn("Reclaim space did not start", job.status_line)
|
||||
|
||||
async def test_reclaim_space_checks_the_disk_first(self):
|
||||
self.old_library()
|
||||
before = self.texts()
|
||||
_plan, _summary, job = await self.start()
|
||||
free = {"bytes": 10**12}
|
||||
|
||||
def free_bytes(_path):
|
||||
return free["bytes"]
|
||||
|
||||
with patch.object(maintenance, "_free_bytes", side_effect=free_bytes):
|
||||
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
|
||||
reclaim = self.reclaim_job()
|
||||
free["bytes"] = 1024
|
||||
# The reclaim job checks when it runs, after the rewrite.
|
||||
await asyncio.wait_for(asyncio.shield(reclaim.task), WAIT)
|
||||
self.assertEqual(reclaim.state, "failed")
|
||||
self.assertEqual(reclaim.last_event["code"], "insufficient-space")
|
||||
self.assertIn("Not enough free disk space", reclaim.status_line)
|
||||
self.assertEqual(self.texts(), before)
|
||||
|
||||
async def test_apply_refuses_up_front_when_the_backup_does_not_fit(self):
|
||||
self.old_library()
|
||||
plan = await self.preview(OPERATION)
|
||||
with patch.object(maintenance, "_free_bytes", return_value=1024):
|
||||
body = await self.apply(plan["token"], status=507)
|
||||
self.assertEqual(body["code"], "insufficient-space")
|
||||
self.assertEqual(self.inline_count(), 12)
|
||||
|
||||
async def test_the_jobs_never_open_a_media_file(self):
|
||||
self.old_library(media=True)
|
||||
media = str(self.output_dir)
|
||||
before = {
|
||||
name: (
|
||||
_digest(os.path.join(media, name)),
|
||||
os.stat(os.path.join(media, name)),
|
||||
)
|
||||
for name in os.listdir(media)
|
||||
}
|
||||
opened = []
|
||||
real_open, real_os_open = builtins.open, os.open
|
||||
|
||||
def tracking_open(file, *args, **kwargs):
|
||||
opened.append(str(file))
|
||||
return real_open(file, *args, **kwargs)
|
||||
|
||||
def tracking_os_open(path, *args, **kwargs):
|
||||
opened.append(str(path))
|
||||
return real_os_open(path, *args, **kwargs)
|
||||
|
||||
with patch("builtins.open", tracking_open), patch("os.open", tracking_os_open):
|
||||
with patch.object(Image, "open", side_effect=AssertionError("opened")):
|
||||
_plan, _summary, job = await self.start()
|
||||
reclaim = await self.finish(job)
|
||||
self.assertEqual(job.state, "finished", job.log)
|
||||
self.assertEqual(reclaim.state, "finished", reclaim.log)
|
||||
self.assertEqual([path for path in opened if path.startswith(media)], [])
|
||||
for name, (digest, stat) in before.items():
|
||||
path = os.path.join(media, name)
|
||||
self.assertEqual(_digest(path), digest)
|
||||
self.assertEqual(os.stat(path).st_mtime_ns, stat.st_mtime_ns)
|
||||
|
||||
async def test_a_second_apply_of_the_token_is_plan_used(self):
|
||||
self.old_library(count=4)
|
||||
plan = await self.preview(OPERATION)
|
||||
body = await self.apply(plan["token"])
|
||||
await self.finish(self.api._jobs.get(body["job"]["id"]))
|
||||
again = await self.apply(plan["token"], status=409)
|
||||
self.assertEqual(again["code"], "plan-used")
|
||||
|
||||
|
||||
class TestWorkflowStorageOutcomes(WorkflowStorageTestCase):
|
||||
async def test_kept_values_and_swept_documents_are_logged(self):
|
||||
prompt_id = self.old_library(count=4)
|
||||
conn = self.api.db.model.get_connection()
|
||||
with conn:
|
||||
conn.execute(
|
||||
"INSERT INTO generated_images (prompt_id, image_path, filename, "
|
||||
"file_path, workflow_data) VALUES (?, '/x/b', 'b', '/x/b', x'00')",
|
||||
(prompt_id,),
|
||||
)
|
||||
store_text(conn, "a document no record references")
|
||||
_plan, _summary, job = await self.start()
|
||||
await self.finish(job)
|
||||
log = "\n".join(job.log)
|
||||
self.assertEqual(job.counters["protected"], 1)
|
||||
self.assertIn(
|
||||
"Rewrote 4 image records; 1 stored value not text, kept as is", log
|
||||
)
|
||||
self.assertIn("Removed 1 unused workflow document", log)
|
||||
|
||||
async def reclaim_with(self, **patches):
|
||||
self.old_library(count=4)
|
||||
before = self.texts()
|
||||
with contextlib.ExitStack() as stack:
|
||||
for name, value in patches.items():
|
||||
stack.enter_context(patch.object(self.api.db, name, value))
|
||||
_plan, _summary, job = await self.start()
|
||||
reclaim = await self.finish(job)
|
||||
self.assertEqual(self.texts(), before)
|
||||
return reclaim
|
||||
|
||||
async def test_a_failed_reclaim_says_the_database_is_unchanged(self):
|
||||
reclaim = await self.reclaim_with(
|
||||
reclaim_space=lambda: (_ for _ in ()).throw(
|
||||
sqlite3.OperationalError("database is locked")
|
||||
)
|
||||
)
|
||||
self.assertEqual(reclaim.state, "failed")
|
||||
self.assertEqual(
|
||||
reclaim.status_line, "Reclaim space failed; the database is unchanged."
|
||||
)
|
||||
|
||||
async def test_a_reclaim_whose_disk_cannot_be_measured_fails_cleanly(self):
|
||||
from py.api import workflow_storage as api_storage
|
||||
|
||||
with patch.object(api_storage, "live_bytes", side_effect=OSError("gone")):
|
||||
reclaim = await self.reclaim_with()
|
||||
self.assertEqual(reclaim.state, "failed")
|
||||
self.assertNotIn("gone", reclaim.status_line)
|
||||
|
||||
async def test_a_busy_checkpoint_says_the_rest_follows_a_restart(self):
|
||||
reclaim = await self.reclaim_with(
|
||||
reclaim_space=lambda: {
|
||||
"before_bytes": 3 * 2**20,
|
||||
"after_bytes": 2**20,
|
||||
"documents_removed": 0,
|
||||
"checkpoint_busy": True,
|
||||
}
|
||||
)
|
||||
self.assertEqual(
|
||||
reclaim.status_line,
|
||||
"Reclaimed 2.0 MB · the database is now 1.0 MB "
|
||||
"(the rest after the next restart)",
|
||||
)
|
||||
@@ -1374,17 +1374,13 @@ class TestModelHousekeeping(DatabaseTestCase):
|
||||
):
|
||||
self.assertEqual(self.db.model.get_database_info(), {})
|
||||
|
||||
def test_vacuum_logs_and_raises_its_failure(self):
|
||||
# S6: a failed VACUUM is reported (legacy /maintenance success:false),
|
||||
# never logged as done.
|
||||
def test_vacuum_logs_errors_instead_of_raising(self):
|
||||
self._save("x")
|
||||
self.db.model.vacuum_database()
|
||||
with patch(
|
||||
"database.models.sqlite3.connect", side_effect=sqlite3.OperationalError
|
||||
):
|
||||
with self.assertLogs("prompt_manager", "ERROR"):
|
||||
with self.assertRaises(sqlite3.OperationalError):
|
||||
self.db.model.vacuum_database()
|
||||
self.db.model.vacuum_database() # logged, not raised
|
||||
self.assertEqual(len(self.db.search_prompts()), 1)
|
||||
|
||||
def test_close_tolerates_a_connection_that_fails_to_close(self):
|
||||
|
||||
@@ -293,8 +293,6 @@ class TestRestoreFromFile(BackupTestCase):
|
||||
super().backup(target, *args, **kwargs)
|
||||
|
||||
def connect_with_failure(*args, **kwargs) -> sqlite3.Connection:
|
||||
if "factory" in kwargs: # the model's own managed connections
|
||||
return connect(*args, **kwargs)
|
||||
return connect(*args, **kwargs, factory=FailingLiveBackup)
|
||||
|
||||
with patch("database.models.sqlite3.connect", side_effect=connect_with_failure):
|
||||
|
||||
@@ -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_or_raise.assert_called_once()
|
||||
handler.db_manager.link_image_to_prompt.assert_called_once()
|
||||
|
||||
def test_admin_scan_ingestion_uses_the_default_profile(self):
|
||||
from py.api import admin
|
||||
|
||||
@@ -63,7 +63,7 @@ class TestAdmissionClasses(unittest.IsolatedAsyncioTestCase):
|
||||
)
|
||||
|
||||
async def test_admission_names_and_explicit_read(self):
|
||||
self.assertEqual(ADMISSIONS, ("read", "guarded", "records", "exclusive"))
|
||||
self.assertEqual(ADMISSIONS, ("read", "guarded", "records"))
|
||||
registry = JobRegistry()
|
||||
job = registry.start(
|
||||
"report", "Report", admission="read", scope={}, runner=_noop
|
||||
@@ -72,7 +72,7 @@ class TestAdmissionClasses(unittest.IsolatedAsyncioTestCase):
|
||||
self.assertEqual(job.admission, "read")
|
||||
self.assertFalse(job.record_changing)
|
||||
with self.assertRaises(ValueError):
|
||||
registry.start("x", "X", admission="bogus", scope={}, runner=_noop)
|
||||
registry.start("x", "X", admission="exclusive", scope={}, runner=_noop)
|
||||
with self.assertRaises(ValueError):
|
||||
registry.start(
|
||||
"x",
|
||||
|
||||
@@ -1,369 +0,0 @@
|
||||
"""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()
|
||||
@@ -1,169 +0,0 @@
|
||||
"""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()
|
||||
@@ -1,497 +0,0 @@
|
||||
"""Move library (#249 S6b): copy, verify, switch, then retire the old files.
|
||||
|
||||
Each test builds a node-folder library (a real PromptDatabase with prompts,
|
||||
images and workflow documents, a config.json holding a key, a pooled safety
|
||||
backup and the pre-migration backup) and a ComfyUI user folder in a temp
|
||||
directory, then moves it with ``database.library_move``. Failures and crash
|
||||
windows are injected through the move's ``checkpoint`` hook: an ``Exception``
|
||||
is a failure the move cleans up after, ``SimulatedCrash`` (a BaseException)
|
||||
stops it dead like a killed process.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import types
|
||||
import unittest
|
||||
from unittest import mock
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from database import library_move as lm # noqa: E402
|
||||
from database.execution_history import ( # noqa: E402
|
||||
notify_history_restore,
|
||||
subscribe_history_restore,
|
||||
)
|
||||
from database.models import PromptModel # noqa: E402
|
||||
from database.operations import PromptDatabase # noqa: E402
|
||||
from database.review_changes import ( # noqa: E402
|
||||
notify_review_change,
|
||||
subscribe_review_changes,
|
||||
)
|
||||
from tests.open_handles import assert_closed, open_handles_under # noqa: E402
|
||||
from utils import library_location as ll # noqa: E402
|
||||
|
||||
KEY = "civitai-secret-key"
|
||||
POOL = "prompts.db.backup_20261001_101010_000001"
|
||||
PRE = "prompts.db.pre-workflow-storage_20261002_101010_000002"
|
||||
|
||||
|
||||
class SimulatedCrash(BaseException):
|
||||
"""The process died here: nothing after this point runs."""
|
||||
|
||||
|
||||
def _digest(path):
|
||||
with open(path, "rb") as handle:
|
||||
return hashlib.sha256(handle.read()).hexdigest()
|
||||
|
||||
|
||||
def _rows(path):
|
||||
"""Every table's rows of ``path`` (read-only), to compare libraries."""
|
||||
uri = "file:" + path.replace(os.sep, "/") + "?mode=ro"
|
||||
conn = sqlite3.connect(uri, uri=True)
|
||||
try:
|
||||
tables = [
|
||||
row[0]
|
||||
for row in conn.execute(
|
||||
"SELECT name FROM sqlite_master WHERE type='table' "
|
||||
"AND name NOT LIKE 'sqlite_%' ORDER BY name"
|
||||
)
|
||||
]
|
||||
return {
|
||||
table: sorted(
|
||||
repr(tuple(row)) for row in conn.execute(f'SELECT * FROM "{table}"')
|
||||
)
|
||||
for table in tables
|
||||
}
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _workflow(seed):
|
||||
return {
|
||||
"nodes": [
|
||||
{"id": n, "type": f"Node{seed}", "widgets": [seed, n]} for n in range(30)
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
class _MoveCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self._tmp = tempfile.TemporaryDirectory()
|
||||
self.base = os.path.realpath(self._tmp.name)
|
||||
self.comfy = os.path.join(self.base, "ComfyUI")
|
||||
self.ext = os.path.join(self.comfy, "custom_nodes", "ComfyUI_PromptManager")
|
||||
os.makedirs(self.ext)
|
||||
self.fp = types.SimpleNamespace()
|
||||
for name in ("output", "input", "temp", "user"):
|
||||
folder = os.path.join(self.comfy, name)
|
||||
os.makedirs(folder)
|
||||
setattr(self.fp, f"get_{name}_directory", lambda f=folder: f)
|
||||
self.target = os.path.join(self.comfy, "user", "__promptmanager")
|
||||
self.old_db = os.path.join(self.ext, "prompts.db")
|
||||
self.new_db = os.path.join(self.target, "prompts.db")
|
||||
ll.clear_override()
|
||||
self.db = PromptDatabase(self.old_db)
|
||||
for index in range(3):
|
||||
prompt_id = self.db.save_prompt(f"a cat number {index}", category="cats")
|
||||
self.db.link_image_to_prompt(
|
||||
str(prompt_id),
|
||||
os.path.join(self.comfy, "output", f"cat{index}.png"),
|
||||
{"workflow": _workflow(index % 2), "prompt": {"1": {"seed": index}}},
|
||||
)
|
||||
self.db.model.close_all()
|
||||
self.config = os.path.join(self.ext, "config.json")
|
||||
with open(self.config, "w", encoding="utf-8") as handle:
|
||||
json.dump(
|
||||
{"integrations": {"lora_manager": {"civitai_api_key": KEY}}}, handle
|
||||
)
|
||||
for name, stamp in ((POOL, 1_700_000_000), (PRE, 1_700_000_100)):
|
||||
path = os.path.join(self.ext, name)
|
||||
source = sqlite3.connect(self.old_db)
|
||||
copy = sqlite3.connect(path)
|
||||
try:
|
||||
source.backup(copy)
|
||||
finally:
|
||||
copy.close()
|
||||
source.close()
|
||||
os.utime(path, (stamp, stamp))
|
||||
self.before = {
|
||||
"rows": _rows(self.old_db),
|
||||
"config": _digest(self.config),
|
||||
"pool": _digest(os.path.join(self.ext, POOL)),
|
||||
"pre": _digest(os.path.join(self.ext, PRE)),
|
||||
}
|
||||
self.config_lock = threading.RLock()
|
||||
|
||||
def tearDown(self):
|
||||
import gc
|
||||
|
||||
gc.collect()
|
||||
PromptModel.close_all_instances()
|
||||
ll.clear_override()
|
||||
try:
|
||||
assert_closed(self, self.ext, self.target)
|
||||
finally:
|
||||
self._tmp.cleanup()
|
||||
|
||||
def location(self, **kwargs):
|
||||
return ll.resolve_library(
|
||||
None, environ={}, extension_root=self.ext, folder_paths=self.fp, **kwargs
|
||||
)
|
||||
|
||||
def move(self, **kwargs):
|
||||
plan = lm.plan_move(self.location())
|
||||
kwargs.setdefault("config_lock", self.config_lock)
|
||||
return lm.move_library(plan, **kwargs)
|
||||
|
||||
def assert_old_library_untouched(self):
|
||||
self.assertTrue(os.path.isfile(self.old_db))
|
||||
self.assertEqual(_rows(self.old_db), self.before["rows"])
|
||||
self.assertEqual(_digest(self.config), self.before["config"])
|
||||
self.assertEqual(_digest(os.path.join(self.ext, POOL)), self.before["pool"])
|
||||
self.assertEqual(_digest(os.path.join(self.ext, PRE)), self.before["pre"])
|
||||
self.assertEqual(
|
||||
os.path.normcase(self.db.model.db_path), os.path.normcase(self.old_db)
|
||||
)
|
||||
self.assertEqual(self.location().source, ll.SOURCE_NODE_FOLDER)
|
||||
|
||||
def assert_target_clean(self):
|
||||
leftovers = os.listdir(self.target) if os.path.isdir(self.target) else []
|
||||
self.assertEqual(leftovers, [])
|
||||
|
||||
|
||||
class TestMove(_MoveCase):
|
||||
def test_move_copies_verifies_switches_and_retires_the_old_files(self):
|
||||
result = self.move()
|
||||
|
||||
self.assertEqual(_rows(self.new_db), self.before["rows"])
|
||||
self.assertEqual(result.counts["prompts"], 3)
|
||||
self.assertEqual(result.counts["generated_images"], 3)
|
||||
self.assertEqual(
|
||||
result.counts["workflow_documents"], 5
|
||||
) # 2 workflows, 3 graphs
|
||||
self.assertEqual(
|
||||
sorted(result.files), sorted(["config.json", POOL, PRE, "prompts.db"])
|
||||
)
|
||||
self.assertIsNone(result.retire_problem)
|
||||
# The live database object now reads and writes the new file.
|
||||
self.assertEqual(self.db.model.db_path, self.new_db)
|
||||
self.db.save_prompt("written after the move")
|
||||
self.assertEqual(len(_rows(self.new_db)["prompts"]), 4)
|
||||
image = self.db.get_prompt_images(
|
||||
str(self.db.search_prompts(text="number 1")[0]["id"])
|
||||
)[0]
|
||||
self.assertEqual(image["workflow_data"], _workflow(1))
|
||||
# config.json, the pooled backup and the pre-migration backup: same bytes.
|
||||
self.assertEqual(
|
||||
_digest(os.path.join(self.target, "config.json")), self.before["config"]
|
||||
)
|
||||
self.assertEqual(_digest(os.path.join(self.target, POOL)), self.before["pool"])
|
||||
self.assertEqual(_digest(os.path.join(self.target, PRE)), self.before["pre"])
|
||||
self.assertEqual(
|
||||
int(os.path.getmtime(os.path.join(self.target, POOL))), 1_700_000_000
|
||||
)
|
||||
# The old files are renamed, never deleted; the backups stay.
|
||||
renamed = sorted(os.listdir(self.ext))
|
||||
self.assertIn("prompts.db.moved-to-__promptmanager", renamed)
|
||||
self.assertIn("config.json.moved-to-__promptmanager", renamed)
|
||||
self.assertNotIn("prompts.db", renamed)
|
||||
self.assertNotIn("config.json", renamed)
|
||||
self.assertFalse(any(name.endswith(("-wal", "-shm")) for name in renamed))
|
||||
self.assertEqual(
|
||||
_rows(os.path.join(self.ext, "prompts.db.moved-to-__promptmanager")),
|
||||
self.before["rows"],
|
||||
)
|
||||
self.assertEqual(ll.read_marker(self.target)["state"], ll.MARKER_MOVED)
|
||||
|
||||
def test_new_files_are_private(self):
|
||||
if os.name != "posix":
|
||||
self.skipTest("POSIX permissions")
|
||||
self.move()
|
||||
self.assertEqual(os.stat(self.target).st_mode & 0o777, 0o700)
|
||||
for name in ("prompts.db", "config.json", POOL, PRE, ll.MARKER_NAME):
|
||||
mode = os.stat(os.path.join(self.target, name)).st_mode & 0o777
|
||||
self.assertEqual(mode, 0o600, name)
|
||||
|
||||
def test_restart_uses_the_new_library_and_never_the_renamed_file(self):
|
||||
self.move()
|
||||
self.assertEqual(self.location().db_path, self.new_db)
|
||||
ll.clear_override()
|
||||
restarted = self.location()
|
||||
self.assertEqual(restarted.source, ll.SOURCE_DEFAULT)
|
||||
self.assertEqual(restarted.db_path, self.new_db)
|
||||
self.assertEqual(
|
||||
restarted.config_path, os.path.join(self.target, "config.json")
|
||||
)
|
||||
|
||||
def test_every_open_database_object_is_repointed(self):
|
||||
node = PromptDatabase(self.old_db)
|
||||
monitor = PromptDatabase(self.old_db)
|
||||
monitor.get_prompt_by_id(1)
|
||||
other = PromptDatabase(os.path.join(self.base, "other.db"))
|
||||
self.move()
|
||||
self.assertEqual(node.model.db_path, self.new_db)
|
||||
self.assertEqual(monitor.model.db_path, self.new_db)
|
||||
self.assertEqual(monitor.get_prompt_by_id(1)["text"], "a cat number 0")
|
||||
self.assertEqual(other.model.db_path, os.path.join(self.base, "other.db"))
|
||||
assert_closed(
|
||||
self, os.path.join(self.ext, "prompts.db.moved-to-__promptmanager")
|
||||
)
|
||||
|
||||
def test_change_observers_follow_the_library(self):
|
||||
reviews, restores = [], []
|
||||
unsubscribe_review = subscribe_review_changes(
|
||||
self.old_db, lambda: reviews.append(1)
|
||||
)
|
||||
unsubscribe_restore = subscribe_history_restore(
|
||||
self.old_db, lambda: restores.append(1)
|
||||
)
|
||||
self.addCleanup(unsubscribe_review)
|
||||
self.addCleanup(unsubscribe_restore)
|
||||
self.move()
|
||||
self.assertEqual(restores, [1]) # the history readers start a new epoch
|
||||
notify_review_change(self.new_db)
|
||||
notify_history_restore(self.new_db)
|
||||
self.assertEqual(reviews, [1])
|
||||
self.assertEqual(restores, [1, 1])
|
||||
|
||||
def test_files_are_only_renamed_inside_one_folder(self):
|
||||
real_replace, real_rename = os.replace, os.rename
|
||||
moves = []
|
||||
|
||||
def replace(src, dst, *args, **kwargs):
|
||||
moves.append((src, dst))
|
||||
return real_replace(src, dst, *args, **kwargs)
|
||||
|
||||
def rename(src, dst, *args, **kwargs):
|
||||
moves.append((src, dst))
|
||||
return real_rename(src, dst, *args, **kwargs)
|
||||
|
||||
with (
|
||||
mock.patch.object(lm.os, "replace", replace),
|
||||
mock.patch.object(lm.os, "rename", rename),
|
||||
):
|
||||
self.move()
|
||||
self.assertTrue(moves)
|
||||
for src, dst in moves:
|
||||
self.assertEqual(os.path.dirname(src), os.path.dirname(dst), (src, dst))
|
||||
|
||||
def test_a_write_during_the_copy_is_copied_again(self):
|
||||
writes = []
|
||||
|
||||
def checkpoint(step):
|
||||
if step == "files" and not writes:
|
||||
writes.append(PromptDatabase(self.old_db).save_prompt("late write"))
|
||||
|
||||
progress = []
|
||||
self.move(checkpoint=checkpoint, progress=progress.append)
|
||||
self.assertIn("late write", repr(_rows(self.new_db)["prompts"]))
|
||||
self.assertTrue(any("changed" in line for line in progress), progress)
|
||||
|
||||
def test_a_library_that_keeps_changing_is_not_moved(self):
|
||||
def checkpoint(step):
|
||||
if step == "files":
|
||||
PromptDatabase(self.old_db).save_prompt(f"write {os.urandom(4).hex()}")
|
||||
|
||||
with self.assertRaises(lm.MoveFailed) as caught:
|
||||
self.move(checkpoint=checkpoint, attempts=2)
|
||||
self.assertEqual(caught.exception.code, "library-busy")
|
||||
self.assertEqual(self.location().source, ll.SOURCE_NODE_FOLDER)
|
||||
self.assert_target_clean()
|
||||
|
||||
def test_old_file_still_open_elsewhere_is_left_in_place(self):
|
||||
reader = sqlite3.connect(self.old_db)
|
||||
self.addCleanup(reader.close)
|
||||
reader.execute("BEGIN")
|
||||
reader.execute("SELECT COUNT(*) FROM prompts").fetchone()
|
||||
PromptDatabase(self.old_db).save_prompt("in the wal")
|
||||
result = self.move()
|
||||
self.assertEqual(result.retire_problem, "old-in-use")
|
||||
self.assertTrue(os.path.isfile(self.old_db))
|
||||
self.assertEqual(self.db.model.db_path, self.new_db)
|
||||
reader.rollback()
|
||||
reader.close()
|
||||
ll.clear_override()
|
||||
self.assertEqual(self.location().source, ll.SOURCE_DEFAULT)
|
||||
|
||||
|
||||
class TestFailures(_MoveCase):
|
||||
def _fail_at(self, step):
|
||||
def checkpoint(name):
|
||||
if name == step:
|
||||
raise RuntimeError(f"failure at {step}")
|
||||
|
||||
return checkpoint
|
||||
|
||||
def test_a_failure_at_any_step_leaves_the_old_library_in_use_and_untouched(self):
|
||||
for step in ("copied", "verified", "files", "locked", "pending", "published"):
|
||||
with self.subTest(step=step):
|
||||
with self.assertRaises(RuntimeError):
|
||||
self.move(checkpoint=self._fail_at(step))
|
||||
self.assert_old_library_untouched()
|
||||
self.assert_target_clean()
|
||||
assert_closed(self, self.target)
|
||||
|
||||
def test_a_copy_that_does_not_verify_is_refused(self):
|
||||
def checkpoint(step):
|
||||
if step == "copied":
|
||||
(temp,) = [
|
||||
name for name in os.listdir(self.target) if name.endswith(".db")
|
||||
]
|
||||
with open(os.path.join(self.target, temp), "r+b") as handle:
|
||||
handle.truncate(4096)
|
||||
|
||||
with self.assertRaises(lm.MoveFailed) as caught:
|
||||
self.move(checkpoint=checkpoint)
|
||||
self.assertEqual(caught.exception.code, "verify-failed")
|
||||
self.assertNotIn(self.base, caught.exception.message)
|
||||
self.assert_old_library_untouched()
|
||||
self.assert_target_clean()
|
||||
|
||||
def test_a_copied_file_that_differs_is_refused(self):
|
||||
def checkpoint(step):
|
||||
if step == "files":
|
||||
for name in os.listdir(self.target):
|
||||
if name.endswith(".bak"):
|
||||
with open(os.path.join(self.target, name), "ab") as handle:
|
||||
handle.write(b"x")
|
||||
|
||||
with self.assertRaises(lm.MoveFailed) as caught:
|
||||
self.move(checkpoint=checkpoint)
|
||||
self.assertEqual(caught.exception.code, "copy-failed")
|
||||
self.assert_old_library_untouched()
|
||||
self.assert_target_clean()
|
||||
|
||||
def test_not_movable_locations_are_refused(self):
|
||||
os.makedirs(self.target)
|
||||
with open(self.new_db, "wb") as handle:
|
||||
handle.write(b"someone else's library")
|
||||
with self.assertRaises(lm.MoveFailed) as caught:
|
||||
lm.plan_move(self.location())
|
||||
self.assertEqual(caught.exception.code, "target-has-library")
|
||||
os.remove(self.new_db)
|
||||
default = ll.LibraryLocation(
|
||||
ll.SOURCE_DEFAULT, self.new_db, "c", self.target, self.target
|
||||
)
|
||||
with self.assertRaises(lm.MoveFailed) as caught:
|
||||
lm.plan_move(default)
|
||||
self.assertEqual(caught.exception.code, "not-movable")
|
||||
|
||||
def test_target_that_is_the_node_folder_is_refused(self):
|
||||
link = os.path.join(self.base, "link-to-node-folder")
|
||||
try:
|
||||
os.symlink(self.ext, link)
|
||||
except (OSError, NotImplementedError):
|
||||
self.skipTest("symlinks are not available")
|
||||
location = ll.LibraryLocation(
|
||||
ll.SOURCE_NODE_FOLDER, self.old_db, self.config, self.ext, link
|
||||
)
|
||||
with self.assertRaises(lm.MoveFailed) as caught:
|
||||
lm.plan_move(location)
|
||||
self.assertEqual(caught.exception.code, "same-folder")
|
||||
|
||||
def test_disk_check_measures_the_target_volume(self):
|
||||
measured = []
|
||||
|
||||
def free(path):
|
||||
measured.append(path)
|
||||
return 1024
|
||||
|
||||
with mock.patch.object(lm, "_free_bytes", free):
|
||||
plan = lm.plan_move(self.location())
|
||||
with self.assertRaises(lm.MoveFailed) as caught:
|
||||
lm.move_library(plan, config_lock=self.config_lock)
|
||||
self.assertEqual(caught.exception.code, "insufficient-space")
|
||||
self.assertTrue(measured)
|
||||
for path in measured:
|
||||
self.assertTrue(
|
||||
os.path.normcase(self.comfy) in os.path.normcase(path)
|
||||
and "custom_nodes" not in path,
|
||||
path,
|
||||
)
|
||||
self.assertGreater(plan.required_bytes, os.path.getsize(self.old_db))
|
||||
self.assert_old_library_untouched()
|
||||
|
||||
|
||||
class TestCrashWindows(_MoveCase):
|
||||
"""A crash between steps never loses the library or picks a stale one."""
|
||||
|
||||
def _crash_at(self, step):
|
||||
def checkpoint(name):
|
||||
if name == step:
|
||||
raise SimulatedCrash(step)
|
||||
|
||||
return checkpoint
|
||||
|
||||
def _restart(self):
|
||||
PromptModel.close_all_instances()
|
||||
ll.clear_override()
|
||||
return self.location()
|
||||
|
||||
def test_crash_before_the_switch_keeps_the_old_library_and_a_retry_succeeds(self):
|
||||
for step in ("copied", "files", "pending", "published"):
|
||||
with self.subTest(step=step):
|
||||
with self.assertRaises(SimulatedCrash):
|
||||
self.move(checkpoint=self._crash_at(step))
|
||||
restarted = self._restart()
|
||||
self.assertEqual(restarted.source, ll.SOURCE_NODE_FOLDER)
|
||||
self.assertEqual(restarted.db_path, self.old_db)
|
||||
self.assertEqual(_rows(self.old_db), self.before["rows"])
|
||||
result = self.move()
|
||||
self.assertEqual(_rows(self.new_db), self.before["rows"])
|
||||
self.assertIsNone(result.retire_problem)
|
||||
self.assertFalse(
|
||||
[name for name in os.listdir(self.target) if name.startswith(".pm-move-")]
|
||||
)
|
||||
|
||||
def test_crash_after_the_marker_uses_the_new_library(self):
|
||||
for step in ("marked", "switched"):
|
||||
with self.subTest(step=step):
|
||||
with self.assertRaises(SimulatedCrash):
|
||||
self.move(checkpoint=self._crash_at(step))
|
||||
restarted = self._restart()
|
||||
self.assertEqual(restarted.source, ll.SOURCE_DEFAULT)
|
||||
self.assertEqual(restarted.db_path, self.new_db)
|
||||
self.assertEqual(_rows(self.new_db), self.before["rows"])
|
||||
# The old file was not renamed yet and is never picked.
|
||||
self.assertTrue(os.path.isfile(self.old_db))
|
||||
# Reset for the next window: back to the node folder only.
|
||||
for name in os.listdir(self.target):
|
||||
os.remove(os.path.join(self.target, name))
|
||||
|
||||
def test_crash_after_the_rename(self):
|
||||
with self.assertRaises(SimulatedCrash):
|
||||
self.move(checkpoint=self._crash_at("retired"))
|
||||
restarted = self._restart()
|
||||
self.assertEqual(restarted.db_path, self.new_db)
|
||||
self.assertFalse(os.path.exists(self.old_db))
|
||||
|
||||
|
||||
class TestEnvironmentTarget(_MoveCase):
|
||||
def test_move_into_the_environment_folder(self):
|
||||
data = os.path.join(self.base, "pm-data")
|
||||
environ = {ll.DATA_DIR_ENV: data}
|
||||
location = ll.resolve_library(
|
||||
None, environ=environ, extension_root=self.ext, folder_paths=self.fp
|
||||
)
|
||||
self.assertEqual(location.problem, "data-dir-empty")
|
||||
result = lm.move_library(lm.plan_move(location), config_lock=self.config_lock)
|
||||
self.assertEqual(result.location.source, ll.SOURCE_ENVIRONMENT)
|
||||
self.assertEqual(_rows(os.path.join(data, "prompts.db")), self.before["rows"])
|
||||
ll.clear_override()
|
||||
self.assertEqual(
|
||||
ll.resolve_library(
|
||||
None, environ=environ, extension_root=self.ext, folder_paths=self.fp
|
||||
).source,
|
||||
ll.SOURCE_ENVIRONMENT,
|
||||
)
|
||||
self.assertEqual(open_handles_under(self.old_db), [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,458 +0,0 @@
|
||||
"""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.
|
||||
- The switching thread never closes a connection another thread is inside
|
||||
a statement on (a use-after-free in SQLite: "Windows fatal exception:
|
||||
access violation"); it interrupts it and the owner closes it.
|
||||
- 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 _CrossThreadCloseProbe:
|
||||
"""Record every ``close()`` this thread calls on a connection in use elsewhere.
|
||||
|
||||
The writer marks its connection with ``using`` when its statement starts.
|
||||
Closing that connection from the switching thread while the writer is
|
||||
inside the statement frees memory SQLite is still using; on Windows the
|
||||
process dies with an access violation, on Linux the writer sees garbage.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.in_use = set()
|
||||
self.violations = []
|
||||
|
||||
def using(self, conn):
|
||||
self.in_use.add(conn)
|
||||
|
||||
def _profile(self, _frame, event, arg):
|
||||
if event != "c_call" or getattr(arg, "__name__", "") != "close":
|
||||
return
|
||||
target = getattr(arg, "__self__", None)
|
||||
if isinstance(target, sqlite3.Connection) and target in self.in_use:
|
||||
self.violations.append(threading.current_thread().name)
|
||||
|
||||
@contextlib.contextmanager
|
||||
def watching(self):
|
||||
previous = sys.getprofile()
|
||||
sys.setprofile(self._profile)
|
||||
try:
|
||||
yield self
|
||||
finally:
|
||||
sys.setprofile(previous)
|
||||
|
||||
|
||||
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")
|
||||
writer_db = PromptDatabase(self.old_db)
|
||||
handler = ImageGenerationHandler(writer_db, mock.MagicMock())
|
||||
started = threading.Event()
|
||||
probe = _CrossThreadCloseProbe()
|
||||
outcome = {}
|
||||
|
||||
def write():
|
||||
conn = writer_db.model.get_connection()
|
||||
|
||||
def trace(sql):
|
||||
if sql.lstrip().upper().startswith(("BEGIN", "INSERT")):
|
||||
probe.using(conn)
|
||||
started.set()
|
||||
|
||||
conn.set_trace_callback(trace)
|
||||
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))
|
||||
|
||||
with probe.watching():
|
||||
result = self.move(checkpoint=checkpoint)
|
||||
thread.join(WAIT)
|
||||
self.assertFalse(thread.is_alive())
|
||||
self.assertEqual(probe.violations, [])
|
||||
self.assertEqual(outcome, {"ok": True})
|
||||
self.assertIsNone(result.retire_problem)
|
||||
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()
|
||||
writer_db = PromptDatabase(self.old_db)
|
||||
thread, started, probe, outcome = self._racing_run_count(writer_db, text)
|
||||
|
||||
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))
|
||||
|
||||
with probe.watching():
|
||||
result = self.move(checkpoint=checkpoint)
|
||||
# The writer's old connection was closed before the old file was
|
||||
# renamed: Windows refuses to rename an open file.
|
||||
self.assertIsNone(result.retire_problem)
|
||||
thread.join(WAIT)
|
||||
self.assertFalse(thread.is_alive())
|
||||
self.assertEqual(probe.violations, [])
|
||||
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)
|
||||
assert_closed(self, renamed)
|
||||
|
||||
def test_a_write_racing_a_failed_switch_lands_in_the_old_library(self):
|
||||
text = "a prompt racing a failed switch"
|
||||
self.db.save_prompt(text, prompt_hash=generate_prompt_hash(text))
|
||||
self.db.model.close_all()
|
||||
writer_db = PromptDatabase(self.old_db)
|
||||
thread, started, probe, outcome = self._racing_run_count(writer_db, text)
|
||||
|
||||
def checkpoint(step):
|
||||
if step == "locked":
|
||||
thread.start()
|
||||
self.assertTrue(started.wait(WAIT))
|
||||
|
||||
def switch(old_db, location, detached):
|
||||
lm.switch_process(old_db, location, detached)
|
||||
raise RuntimeError("switch failed at its very end")
|
||||
|
||||
with probe.watching():
|
||||
with self.assertRaises(lm.MoveFailed):
|
||||
self.move(checkpoint=checkpoint, switch=switch)
|
||||
# The copy was removed: no connection was left open on it.
|
||||
self.assert_target_clean()
|
||||
assert_closed(self, self.target)
|
||||
thread.join(WAIT)
|
||||
self.assertFalse(thread.is_alive())
|
||||
self.assertEqual(probe.violations, [])
|
||||
self.assertEqual(outcome, {"ok": True})
|
||||
self.assertEqual(_run_count(self.old_db, text), 1)
|
||||
|
||||
def test_a_retire_failure_after_a_complete_switch_does_not_fail_the_move(self):
|
||||
node = PromptDatabase(self.old_db)
|
||||
node.get_prompt_by_id(1)
|
||||
real = models.PromptModel.retire_detached
|
||||
calls = []
|
||||
|
||||
def failing(connections, logger=None):
|
||||
calls.append(len(connections))
|
||||
real(connections, logger) # closed, but the call still fails
|
||||
raise RuntimeError("retiring failed")
|
||||
|
||||
with mock.patch.object(
|
||||
models.PromptModel, "retire_detached", staticmethod(failing)
|
||||
):
|
||||
with self.assertLogs("prompt_manager.database.library_move", "ERROR"):
|
||||
result = self.move()
|
||||
self.assertEqual(len(calls), 1)
|
||||
self.assertGreater(calls[0], 0)
|
||||
self.assertEqual(node.model.db_path, self.new_db)
|
||||
self.assertEqual(result.location.db_path, self.new_db)
|
||||
self.assertTrue(
|
||||
os.path.isfile(
|
||||
os.path.join(self.ext, "prompts.db.moved-to-__promptmanager")
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _racing_run_count(writer_db, text):
|
||||
"""A run count whose UPDATE marks its connection in use for the probe."""
|
||||
started = threading.Event()
|
||||
probe = _CrossThreadCloseProbe()
|
||||
outcome = {}
|
||||
|
||||
def write():
|
||||
conn = writer_db.model.get_connection()
|
||||
|
||||
def trace(sql):
|
||||
if sql.startswith("UPDATE"):
|
||||
probe.using(conn)
|
||||
started.set()
|
||||
|
||||
conn.set_trace_callback(trace)
|
||||
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()
|
||||
|
||||
return threading.Thread(target=write), started, probe, outcome
|
||||
|
||||
|
||||
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()
|
||||
@@ -1,546 +0,0 @@
|
||||
"""Maintenance plans (#249 System S6): the plan store and the read-only previews.
|
||||
|
||||
``PlanStore`` holds each reviewed plan server-side under a single-use token
|
||||
for ``ttl`` seconds, within a count and byte budget it never truncates to fit.
|
||||
``PromptDatabase.preview_maintenance`` builds a plan on its own ``mode=ro``
|
||||
connection: it writes nothing (the TESTING.md read-only standard) and reuses
|
||||
the legacy maintenance queries, so a preview counts what the legacy cleanup
|
||||
would change, except that a missing-file link is only counted under a source
|
||||
that can be read now.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from database import maintenance # noqa: E402
|
||||
from database.operations import PromptDatabase # noqa: E402
|
||||
from py.api.maintenance_plans import ( # noqa: E402
|
||||
PlanExpired,
|
||||
PlanStore,
|
||||
PlanUsed,
|
||||
ScopeTooLarge,
|
||||
)
|
||||
from utils.hashing import generate_prompt_hash # noqa: E402
|
||||
|
||||
try:
|
||||
from tests import test_execution_history_database as history_tests # noqa: E402
|
||||
from tests.open_handles import assert_closed # noqa: E402
|
||||
from tests.test_execution_context_database import bundle # noqa: E402
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
import test_execution_history_database as history_tests # noqa: E402
|
||||
from open_handles import assert_closed # noqa: E402
|
||||
from test_execution_context_database import bundle # noqa: E402
|
||||
|
||||
OPERATIONS = ("duplicate_prompts", "optimize", "unused_components", "missing_links")
|
||||
|
||||
|
||||
def _digest(path):
|
||||
if not os.path.exists(path):
|
||||
return None
|
||||
with open(path, "rb") as handle:
|
||||
return hashlib.sha256(handle.read()).hexdigest()
|
||||
|
||||
|
||||
class FakeClock:
|
||||
def __init__(self):
|
||||
self.now = 1000.0
|
||||
|
||||
def __call__(self):
|
||||
return self.now
|
||||
|
||||
|
||||
class TestPlanStore(unittest.TestCase):
|
||||
def test_plan_expiry_budget_and_single_use(self):
|
||||
clock = FakeClock()
|
||||
store = PlanStore(ttl=900, max_plans=2, max_bytes=200, clock=clock)
|
||||
first = store.issue("optimize", "sha256:a", {"targets": []}, ("db", 0))
|
||||
self.assertGreaterEqual(len(first), 40) # token_urlsafe(32)
|
||||
plan = store.peek(first, "optimize")
|
||||
self.assertEqual(
|
||||
(plan.operation, plan.fingerprint, plan.db_identity),
|
||||
("optimize", "sha256:a", ("db", 0)),
|
||||
)
|
||||
# Single use: the second consume names the job the first one started.
|
||||
self.assertIs(store.consume(first, "optimize", job_id="job-1"), plan)
|
||||
with self.assertRaises(PlanUsed) as used:
|
||||
store.consume(first, "optimize", job_id="job-2")
|
||||
self.assertEqual(used.exception.job_id, "job-1")
|
||||
with self.assertRaises(PlanUsed):
|
||||
store.peek(first, "optimize")
|
||||
# Another operation's token, an unknown token and a non-string are expired.
|
||||
second = store.issue("missing_links", "sha256:b", {"targets": [1]}, ("db", 0))
|
||||
for token, operation in (
|
||||
(second, "optimize"),
|
||||
("nope", "missing_links"),
|
||||
(None, "missing_links"),
|
||||
):
|
||||
with self.subTest(token=token), self.assertRaises(PlanExpired):
|
||||
store.peek(token, operation)
|
||||
# Expiry: valid up to ttl, gone after it.
|
||||
clock.now += 900
|
||||
self.assertEqual(store.peek(second, "missing_links").fingerprint, "sha256:b")
|
||||
clock.now += 0.001
|
||||
with self.assertRaises(PlanExpired):
|
||||
store.consume(second, "missing_links", job_id="job-3")
|
||||
# Budget: a plan over max_bytes is refused whole, never truncated.
|
||||
with self.assertRaises(ScopeTooLarge):
|
||||
store.issue("unused_components", "x", {"targets": list(range(100))}, ())
|
||||
# Count budget: the oldest plan makes room for a new one and is expired.
|
||||
a = store.issue("optimize", "a", {}, ())
|
||||
b = store.issue("optimize", "b", {}, ())
|
||||
c = store.issue("optimize", "c", {}, ())
|
||||
with self.assertRaises(PlanExpired):
|
||||
store.peek(a, "optimize")
|
||||
self.assertEqual(
|
||||
[store.peek(t, "optimize").fingerprint for t in (b, c)], ["b", "c"]
|
||||
)
|
||||
# Byte budget across plans: older plans are dropped until the new one fits.
|
||||
sized = PlanStore(max_plans=8, max_bytes=100, clock=clock)
|
||||
pad = {"pad": "x" * 40} # 50 bytes as compact JSON
|
||||
p1, p2, p3 = (sized.issue("optimize", str(n), pad, ()) for n in range(3))
|
||||
with self.assertRaises(PlanExpired):
|
||||
sized.peek(p1, "optimize")
|
||||
self.assertEqual([sized.peek(t).fingerprint for t in (p2, p3)], ["1", "2"])
|
||||
|
||||
def test_expires_at_is_wall_clock(self):
|
||||
store = PlanStore(ttl=60, clock=FakeClock(), wall_clock=lambda: 5000.0)
|
||||
token = store.issue("optimize", "f", {}, ())
|
||||
self.assertEqual(store.peek(token, "optimize").expires_at, 5060.0)
|
||||
|
||||
def test_discard_forgets_a_plan(self):
|
||||
store = PlanStore()
|
||||
token = store.issue("optimize", "f", {}, ())
|
||||
store.discard(token)
|
||||
store.discard("unknown")
|
||||
with self.assertRaises(PlanExpired):
|
||||
store.peek(token, "optimize")
|
||||
|
||||
def test_used_tokens_are_remembered_within_a_bound(self):
|
||||
store = PlanStore(max_plans=2)
|
||||
tokens = []
|
||||
for index in range(store.MAX_USED + 1):
|
||||
token = store.issue("optimize", str(index), {}, ())
|
||||
store.consume(token, "optimize", job_id=f"job-{index}")
|
||||
tokens.append(token)
|
||||
with self.assertRaises(PlanExpired): # the oldest record was dropped
|
||||
store.peek(tokens[0], "optimize")
|
||||
with self.assertRaises(PlanUsed):
|
||||
store.peek(tokens[-1], "optimize")
|
||||
|
||||
|
||||
class MaintenanceDatabaseCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temp = tempfile.TemporaryDirectory()
|
||||
self.root = Path(os.path.realpath(self.temp.name))
|
||||
self.output = self.root / "output"
|
||||
self.output.mkdir()
|
||||
self.db_path = str(self.root / "prompts.db")
|
||||
self.db = PromptDatabase(self.db_path)
|
||||
|
||||
def tearDown(self):
|
||||
self.db.close_all()
|
||||
assert_closed(self, self.root)
|
||||
self.temp.cleanup()
|
||||
|
||||
def save(self, text, **metadata):
|
||||
return self.db.save_prompt(
|
||||
text, prompt_hash=generate_prompt_hash(text), **metadata
|
||||
)
|
||||
|
||||
def image(self, prompt_id, name, *, exists=True, folder=None):
|
||||
path = Path(folder or self.output) / name
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
if exists:
|
||||
path.write_bytes(b"png")
|
||||
return self.db.link_image_to_prompt(prompt_id, str(path))
|
||||
|
||||
def duplicate(self, text):
|
||||
"""A second prompt row whose text differs only in case (legacy duplicate)."""
|
||||
with self.db.model.get_connection() as conn:
|
||||
cursor = conn.execute(
|
||||
"INSERT INTO prompts(text, created_at) VALUES(?, '2999-01-01')",
|
||||
(text,),
|
||||
)
|
||||
return cursor.lastrowid
|
||||
|
||||
def preview(self, operation, roots=None):
|
||||
roots = [str(self.output)] if roots is None else roots
|
||||
return self.db.preview_maintenance(operation, roots)
|
||||
|
||||
def ids(self, sql, *args):
|
||||
with self.db.model.get_connection() as conn:
|
||||
return [row[0] for row in conn.execute(sql, args)]
|
||||
|
||||
|
||||
class TestPreviews(MaintenanceDatabaseCase):
|
||||
def populate(self):
|
||||
kept = self.save("Kept prompt")
|
||||
self.image(kept, "kept.png")
|
||||
self.image(kept, "gone.png", exists=False)
|
||||
self.save("Unused prompt")
|
||||
self.duplicate("KEPT PROMPT")
|
||||
return kept
|
||||
|
||||
def test_preview_writes_nothing(self):
|
||||
self.populate()
|
||||
# The app's own connection stays open, as in the running server.
|
||||
self.db.get_statistics()
|
||||
files = [self.db_path, self.db_path + "-wal"]
|
||||
real_connect = sqlite3.connect
|
||||
|
||||
class Recording(sqlite3.Connection):
|
||||
changes_at_close = []
|
||||
|
||||
def close(self):
|
||||
Recording.changes_at_close.append(self.total_changes)
|
||||
super().close()
|
||||
|
||||
opened = []
|
||||
|
||||
def connect(*args, **kwargs):
|
||||
conn = real_connect(*args, factory=Recording, **kwargs)
|
||||
opened.append((args, kwargs, conn))
|
||||
return conn
|
||||
|
||||
for operation in OPERATIONS:
|
||||
with self.subTest(operation=operation):
|
||||
before = [_digest(path) for path in files]
|
||||
shm_size = os.path.getsize(self.db_path + "-shm")
|
||||
media = {
|
||||
name: _digest(self.output / name)
|
||||
for name in os.listdir(self.output)
|
||||
}
|
||||
prompts = self.ids("SELECT id FROM prompts ORDER BY id")
|
||||
images = self.ids("SELECT id FROM generated_images ORDER BY id")
|
||||
opened.clear()
|
||||
Recording.changes_at_close.clear()
|
||||
with patch.object(maintenance.sqlite3, "connect", side_effect=connect):
|
||||
draft = self.preview(operation)
|
||||
self.assertTrue(draft.fingerprint.startswith("sha256:"))
|
||||
self.assertEqual(len(opened), 1)
|
||||
args, kwargs, conn = opened[0]
|
||||
self.assertTrue(str(args[0]).endswith("?mode=ro"), args)
|
||||
self.assertTrue(kwargs.get("uri"))
|
||||
self.assertEqual(Recording.changes_at_close, [0])
|
||||
with self.assertRaises(sqlite3.ProgrammingError): # closed
|
||||
conn.execute("SELECT 1")
|
||||
self.assertEqual([_digest(path) for path in files], before)
|
||||
self.assertEqual(os.path.getsize(self.db_path + "-shm"), shm_size)
|
||||
self.assertEqual(
|
||||
{
|
||||
name: _digest(self.output / name)
|
||||
for name in os.listdir(self.output)
|
||||
},
|
||||
media,
|
||||
)
|
||||
self.assertEqual(
|
||||
self.ids("SELECT id FROM prompts ORDER BY id"), prompts
|
||||
)
|
||||
self.assertEqual(
|
||||
self.ids("SELECT id FROM generated_images ORDER BY id"), images
|
||||
)
|
||||
|
||||
def test_unknown_operation_raises(self):
|
||||
with self.assertRaises(ValueError):
|
||||
self.db.preview_maintenance("reboot", [])
|
||||
|
||||
def test_same_library_gives_the_same_fingerprint(self):
|
||||
self.populate()
|
||||
for operation in OPERATIONS:
|
||||
with self.subTest(operation=operation):
|
||||
self.assertEqual(
|
||||
self.preview(operation).fingerprint,
|
||||
self.preview(operation).fingerprint,
|
||||
)
|
||||
|
||||
def test_preview_details_are_capped_and_carry_no_server_path(self):
|
||||
prompt = self.save("Many")
|
||||
for index in range(60):
|
||||
self.image(prompt, f"sub/gone-{index}.png", exists=False)
|
||||
draft = self.preview("missing_links")
|
||||
self.assertEqual(draft.scope["count"], 60)
|
||||
self.assertEqual(len(draft.details), maintenance.MAX_DETAILS)
|
||||
self.assertEqual(draft.details[0]["label"], "sub/gone-0.png")
|
||||
self.assertEqual(draft.details[0]["detail"], "output · the file is missing")
|
||||
self.assertNotIn(str(self.root), str(draft.public()))
|
||||
self.assertEqual(len(draft.targets), 60)
|
||||
|
||||
def test_duplicate_prompts_plan_counts_the_legacy_groups(self):
|
||||
kept = self.save("A cat")
|
||||
self.duplicate("A CAT")
|
||||
self.duplicate("a cat")
|
||||
paired = self.db.get_prompt_pair(self.db.save_prompt_pair("a dog", "blur"))
|
||||
self.duplicate("A DOG")
|
||||
draft = self.preview("duplicate_prompts")
|
||||
self.assertEqual(draft.scope["count"], 2)
|
||||
self.assertEqual(draft.scope["label"], "2 duplicate prompts in 1 group")
|
||||
self.assertEqual(draft.protected["count"], 2)
|
||||
self.assertEqual(
|
||||
draft.protected["reasons"], ["1 group belongs to a saved pair"]
|
||||
)
|
||||
self.assertEqual(draft.targets[0][0], kept)
|
||||
self.assertEqual(len(draft.targets), 1)
|
||||
self.assertEqual(
|
||||
[row["detail"] for row in draft.details],
|
||||
["Keeps the oldest of 3 copies", "Kept: belongs to a saved pair"],
|
||||
)
|
||||
self.assertEqual(draft.details[0]["label"], "A cat")
|
||||
self.assertIsNotNone(paired)
|
||||
|
||||
def test_unused_components_plan_protects_pairs_history_and_titled_parents(self):
|
||||
orphan = self.save("Orphan")
|
||||
pair = self.db.get_prompt_pair(self.db.save_prompt_pair("paired", "neg"))
|
||||
titled = self.save("Titled parent")
|
||||
untitled = self.save("Explicit parent")
|
||||
protected = self.save("Protected", tags=["__protected__"])
|
||||
with self.db.model.get_connection() as conn:
|
||||
conn.execute(
|
||||
"INSERT INTO prompt_library_state(prompt_id,title) "
|
||||
"VALUES(?,?),(?,NULL)",
|
||||
(titled, "My recipe", untitled),
|
||||
)
|
||||
hinted = self.save("Recorded run")
|
||||
draft_before_hint = self.preview("unused_components")
|
||||
self.assertEqual(sorted(draft_before_hint.targets), sorted([orphan, hinted]))
|
||||
self.record_run_hint(hinted)
|
||||
draft = self.preview("unused_components")
|
||||
self.assertEqual(draft.targets, [orphan])
|
||||
self.assertEqual(draft.scope["count"], 1)
|
||||
self.assertEqual(draft.scope["label"], "1 unused prompt component")
|
||||
self.assertEqual(draft.protected["count"], 6)
|
||||
self.assertEqual(
|
||||
draft.protected["reasons"],
|
||||
[
|
||||
"2 in saved pairs",
|
||||
"1 tagged __protected__",
|
||||
"2 with Library state (a title or an explicit parent)",
|
||||
"1 used by a recorded run",
|
||||
],
|
||||
)
|
||||
self.assertEqual(
|
||||
draft.details,
|
||||
[
|
||||
{
|
||||
"label": "Orphan",
|
||||
"detail": maintenance.UNUSED_DETAIL,
|
||||
}
|
||||
],
|
||||
)
|
||||
# The plan counts exactly what the legacy prune removes.
|
||||
self.assertIsNotNone(pair)
|
||||
self.assertEqual(protected > 0, True)
|
||||
self.assertEqual(self.db.prune_orphaned_prompts(), 1)
|
||||
self.assertEqual(self.preview("unused_components").targets, [])
|
||||
|
||||
def record_run_hint(self, prompt_id, run=0):
|
||||
"""Make ``prompt_id`` the live positive hint of a recorded run."""
|
||||
text = self.db.get_prompt_by_id(prompt_id)["text"]
|
||||
history = history_tests.HistoryDatabaseTests
|
||||
start = history.start(self, self.db, run)
|
||||
self.db.write_execution_start(start)
|
||||
self.db.write_execution_finish(history.finish(self, start))
|
||||
self.assertTrue(
|
||||
self.db.write_execution_contexts(
|
||||
start.run_id, start.storage_epoch, bundle(positive=text, negative="")
|
||||
)
|
||||
)
|
||||
|
||||
def test_missing_links_plan_skips_unavailable_roots(self):
|
||||
prompt = self.save("Links")
|
||||
self.image(prompt, "here.png")
|
||||
gone = self.image(prompt, "gone.png", exists=False)
|
||||
unmounted = self.root / "unmounted"
|
||||
self.image(prompt, "away.png", folder=unmounted)
|
||||
os.remove(unmounted / "away.png")
|
||||
os.rmdir(unmounted) # the drive went away: its links look missing
|
||||
self.image(
|
||||
prompt, "elsewhere.png", folder=self.root / "elsewhere", exists=False
|
||||
)
|
||||
draft = self.preview("missing_links", [str(self.output), str(unmounted)])
|
||||
self.assertEqual(draft.targets, [gone])
|
||||
self.assertEqual(draft.scope["count"], 1)
|
||||
self.assertEqual(draft.scope["label"], "1 link to a missing file")
|
||||
self.assertEqual(draft.protected["count"], 2)
|
||||
self.assertEqual(
|
||||
draft.protected["reasons"],
|
||||
[
|
||||
"1 link under an unavailable source",
|
||||
"1 link outside the image sources",
|
||||
],
|
||||
)
|
||||
self.assertFalse(draft.media["deletes_files"])
|
||||
# No readable source: nothing is counted, every link is protected.
|
||||
none = self.preview("missing_links", [str(unmounted)])
|
||||
self.assertEqual((none.targets, none.scope["count"]), ([], 0))
|
||||
self.assertEqual(none.protected["count"], 4)
|
||||
|
||||
def test_missing_links_plan_matches_the_health_check(self):
|
||||
prompt = self.save("Links")
|
||||
self.image(prompt, "a.png", exists=False)
|
||||
self.image(prompt, "b.png", exists=False)
|
||||
self.image(prompt, "c.png")
|
||||
report = self.db.run_health_checks([str(self.output)])
|
||||
missing = next(c for c in report["checks"] if c["id"] == "missing_files")
|
||||
self.assertEqual(self.preview("missing_links").scope["count"], missing["count"])
|
||||
|
||||
def test_optimize_plan_reports_sizes_without_paths(self):
|
||||
self.populate()
|
||||
draft = self.preview("optimize")
|
||||
self.assertEqual(draft.targets, [])
|
||||
self.assertEqual(draft.scope, {"count": 1, "label": "The whole database"})
|
||||
self.assertEqual(
|
||||
[row["label"] for row in draft.details],
|
||||
["Database file", "Free pages", "Lock", "Free space needed"],
|
||||
)
|
||||
self.assertNotIn(str(self.root), str(draft.public()))
|
||||
|
||||
def test_legacy_cleanup_missing_images_skips_unavailable_roots(self):
|
||||
prompt = self.save("Links")
|
||||
gone = self.image(prompt, "gone.png", exists=False)
|
||||
unmounted = self.root / "unmounted"
|
||||
away = self.image(prompt, "away.png", folder=unmounted, exists=False)
|
||||
outside = self.image(prompt, "x.png", folder=self.root / "other", exists=False)
|
||||
unmounted.rmdir()
|
||||
self.assertEqual(
|
||||
self.db.cleanup_missing_images(roots=[str(self.output), str(unmounted)]), 1
|
||||
)
|
||||
remaining = self.ids("SELECT id FROM generated_images ORDER BY id")
|
||||
self.assertEqual(remaining, sorted([away, outside]))
|
||||
self.assertNotIn(gone, remaining)
|
||||
# Without roots it keeps the legacy rule: every link whose file is gone.
|
||||
self.assertEqual(self.db.cleanup_missing_images(), 2)
|
||||
|
||||
def test_reader_closes_its_connection_when_a_read_raises(self):
|
||||
self.populate()
|
||||
with patch.object(
|
||||
maintenance, "_orphaned_ids", side_effect=sqlite3.OperationalError("boom")
|
||||
):
|
||||
with self.assertRaises(sqlite3.OperationalError):
|
||||
self.preview("unused_components")
|
||||
# The teardown's assert_closed proves the read-only handle was closed.
|
||||
|
||||
|
||||
class TestApply(MaintenanceDatabaseCase):
|
||||
"""``apply_maintenance`` rechecks the fingerprint inside BEGIN IMMEDIATE."""
|
||||
|
||||
def apply(self, draft, roots=None):
|
||||
roots = [str(self.output)] if roots is None else roots
|
||||
return self.db.apply_maintenance(
|
||||
draft.operation, draft.fingerprint, draft.targets, roots
|
||||
)
|
||||
|
||||
def test_unused_components_apply_removes_exactly_the_reviewed_prompts(self):
|
||||
first, second = self.save("One"), self.save("Two")
|
||||
draft = self.preview("unused_components")
|
||||
with patch.object(maintenance, "ID_BATCH", 1): # one id per statement
|
||||
result = self.apply(draft)
|
||||
self.assertEqual(result["changed"], 2)
|
||||
self.assertEqual(
|
||||
[row["prompt_id"] for row in result["rows"]], sorted([first, second])
|
||||
)
|
||||
self.assertIsNone(self.db.get_prompt_by_id(first))
|
||||
self.assertIsNone(self.db.get_prompt_by_id(second))
|
||||
|
||||
def test_changed_targets_raise_before_anything_is_written(self):
|
||||
orphan, protected = self.save("Orphan"), self.save("Soon protected")
|
||||
prompt = self.save("Links")
|
||||
gone = self.image(prompt, "gone.png", exists=False)
|
||||
kept = self.save("Kept")
|
||||
self.duplicate("KEPT")
|
||||
cases = {
|
||||
"unused_components": lambda: self.db.model.get_connection().execute(
|
||||
"INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)",
|
||||
(protected, "Now a recipe"),
|
||||
),
|
||||
"missing_links": lambda: (self.output / "gone.png").write_bytes(b"back"),
|
||||
"duplicate_prompts": lambda: self.db.update_prompt_metadata(kept, rating=4),
|
||||
}
|
||||
for operation, change in cases.items():
|
||||
with self.subTest(operation=operation):
|
||||
draft = self.preview(operation)
|
||||
self.assertTrue(draft.targets)
|
||||
change()
|
||||
self.db.model.get_connection().commit()
|
||||
with self.assertRaises(maintenance.PlanChanged):
|
||||
self.apply(draft)
|
||||
self.assertIsNotNone(self.db.get_prompt_by_id(orphan))
|
||||
self.assertIsNotNone(self.db.get_image_by_id(gone))
|
||||
self.assertEqual(len(self.db.search_prompts(text="kept")), 2)
|
||||
|
||||
def test_a_link_whose_source_went_away_is_never_removed(self):
|
||||
folder = self.root / "drive"
|
||||
prompt = self.save("Links")
|
||||
link = self.image(prompt, "gone.png", folder=folder, exists=False)
|
||||
roots = [str(folder)]
|
||||
draft = self.preview("missing_links", roots)
|
||||
self.assertEqual(draft.targets, [link])
|
||||
folder.rmdir() # unmounted between the review and the apply
|
||||
with self.assertRaises(maintenance.PlanChanged):
|
||||
self.apply(draft, roots)
|
||||
self.assertIsNotNone(self.db.get_image_by_id(link))
|
||||
|
||||
def test_missing_links_apply_removes_rows_only(self):
|
||||
prompt = self.save("Links")
|
||||
kept = self.image(prompt, "kept.png")
|
||||
gone = [self.image(prompt, f"gone-{n}.png", exists=False) for n in range(3)]
|
||||
media = {p.name: _digest(p) for p in self.output.iterdir()}
|
||||
draft = self.preview("missing_links")
|
||||
with patch.object(maintenance, "ID_BATCH", 2):
|
||||
result = self.apply(draft)
|
||||
self.assertEqual(result["changed"], 3)
|
||||
self.assertEqual(
|
||||
[os.path.basename(row["image_path"]) for row in result["rows"]],
|
||||
[f"gone-{n}.png" for n in range(3)],
|
||||
)
|
||||
self.assertEqual({p.name: _digest(p) for p in self.output.iterdir()}, media)
|
||||
self.assertIsNotNone(self.db.get_image_by_id(kept))
|
||||
for image_id in gone:
|
||||
self.assertIsNone(self.db.get_image_by_id(image_id))
|
||||
|
||||
def test_duplicate_prompts_apply_merges_into_the_oldest(self):
|
||||
kept = self.save("A cat")
|
||||
copy = self.duplicate("A CAT")
|
||||
image = self.image(copy, "cat.png")
|
||||
draft = self.preview("duplicate_prompts")
|
||||
result = self.apply(draft)
|
||||
self.assertEqual(result["changed"], 1)
|
||||
self.assertEqual(result["rows"], [{"prompt_id": copy, "kept_id": kept}])
|
||||
self.assertIsNone(self.db.get_prompt_by_id(copy))
|
||||
self.assertEqual(self.db.get_image_by_id(image)["prompt_id"], kept)
|
||||
|
||||
def test_optimize_apply_vacuums(self):
|
||||
self.save("Some")
|
||||
draft = self.preview("optimize")
|
||||
with patch.object(self.db.model, "vacuum_database") as vacuum:
|
||||
self.assertEqual(self.apply(draft), {"changed": 1, "rows": []})
|
||||
vacuum.assert_called_once_with()
|
||||
self.assertEqual(self.apply(self.preview("optimize"))["changed"], 1)
|
||||
with self.assertRaises(maintenance.PlanChanged):
|
||||
self.db.apply_maintenance("optimize", "sha256:other", [], [])
|
||||
|
||||
def test_apply_of_an_unknown_operation_raises(self):
|
||||
with self.assertRaises(ValueError):
|
||||
self.db.apply_maintenance("reboot", "x", [], [])
|
||||
|
||||
|
||||
class TestVacuum(MaintenanceDatabaseCase):
|
||||
def test_vacuum_failure_raises(self):
|
||||
with patch(
|
||||
"database.models.sqlite3.connect",
|
||||
side_effect=sqlite3.OperationalError("disk I/O error"),
|
||||
):
|
||||
with self.assertLogs("prompt_manager", "ERROR"):
|
||||
with self.assertRaises(sqlite3.OperationalError):
|
||||
self.db.model.vacuum_database()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,332 +0,0 @@
|
||||
"""ManagedConnection: never closed under the thread using it.
|
||||
|
||||
Closing a sqlite3 connection from one thread while another is inside a
|
||||
statement on it is a use-after-free in SQLite ("Windows fatal exception:
|
||||
access violation" in CI). A retired connection is closed at once when idle,
|
||||
otherwise interrupted and closed by its owner as the call ends; it never
|
||||
commits again.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
import weakref
|
||||
from unittest import mock
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from database.managed_connection import ( # noqa: E402
|
||||
ManagedConnection,
|
||||
ManagedCursor,
|
||||
await_closed,
|
||||
retire_connections,
|
||||
)
|
||||
from database.models import PromptModel # noqa: E402
|
||||
from tests.open_handles import assert_closed # noqa: E402
|
||||
|
||||
WAIT = 30
|
||||
|
||||
|
||||
class _Case(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self._tmp = tempfile.TemporaryDirectory()
|
||||
self.path = os.path.join(self._tmp.name, "managed.db")
|
||||
with contextlib.closing(sqlite3.connect(self.path)) as conn:
|
||||
conn.execute("PRAGMA journal_mode=WAL")
|
||||
conn.execute("CREATE TABLE t(x INTEGER)")
|
||||
conn.execute("INSERT INTO t VALUES (0)")
|
||||
conn.commit()
|
||||
self.opened = []
|
||||
|
||||
def tearDown(self):
|
||||
for conn in self.opened:
|
||||
conn.close()
|
||||
try:
|
||||
assert_closed(self, self._tmp.name)
|
||||
finally:
|
||||
self._tmp.cleanup()
|
||||
|
||||
def connect(self, **kwargs):
|
||||
conn = sqlite3.connect(
|
||||
self.path,
|
||||
check_same_thread=False,
|
||||
timeout=WAIT,
|
||||
factory=ManagedConnection,
|
||||
**kwargs,
|
||||
)
|
||||
self.opened.append(conn)
|
||||
return conn
|
||||
|
||||
def value(self):
|
||||
with contextlib.closing(sqlite3.connect(self.path)) as conn:
|
||||
return conn.execute("SELECT x FROM t").fetchone()[0]
|
||||
|
||||
|
||||
class TestRetire(_Case):
|
||||
def test_an_idle_connection_retired_from_another_thread_closes_at_once(self):
|
||||
conn = self.connect()
|
||||
conn.execute("SELECT 1").fetchone()
|
||||
closed = []
|
||||
thread = threading.Thread(target=lambda: closed.append(conn.retire()))
|
||||
thread.start()
|
||||
thread.join(WAIT)
|
||||
self.assertEqual(closed, [True])
|
||||
self.assertTrue(conn.wait_closed(0))
|
||||
with self.assertRaises(sqlite3.ProgrammingError):
|
||||
conn.execute("SELECT 1")
|
||||
self.assertTrue(conn.retire()) # retiring again is harmless
|
||||
|
||||
def test_a_connection_in_use_is_closed_by_its_owner_when_the_call_ends(self):
|
||||
conn = self.connect()
|
||||
holder = sqlite3.connect(self.path, isolation_level=None)
|
||||
self.opened.append(holder)
|
||||
holder.execute("BEGIN IMMEDIATE")
|
||||
started = threading.Event()
|
||||
outcome = {}
|
||||
conn.set_trace_callback(
|
||||
lambda sql: started.set() if sql.startswith("UPDATE") else None
|
||||
)
|
||||
|
||||
def write():
|
||||
try:
|
||||
with conn:
|
||||
conn.execute("UPDATE t SET x = 1")
|
||||
outcome["ok"] = True
|
||||
except sqlite3.Error as error:
|
||||
outcome["error"] = error
|
||||
|
||||
thread = threading.Thread(target=write)
|
||||
thread.start()
|
||||
self.assertTrue(started.wait(WAIT))
|
||||
# The owner is waiting on our lock inside its UPDATE: not closed here.
|
||||
self.assertFalse(conn.retire())
|
||||
self.assertFalse(conn.retire())
|
||||
self.assertFalse(conn.wait_closed(0))
|
||||
self.assertEqual(retire_connections([conn]), [conn])
|
||||
holder.execute("ROLLBACK")
|
||||
thread.join(WAIT)
|
||||
self.assertFalse(thread.is_alive())
|
||||
self.assertTrue(conn.wait_closed(WAIT))
|
||||
self.assertIsInstance(outcome.get("error"), sqlite3.Error)
|
||||
self.assertEqual(self.value(), 0)
|
||||
|
||||
def test_a_write_retired_before_its_commit_is_rolled_back(self):
|
||||
conn = self.connect()
|
||||
conn.execute("UPDATE t SET x = 2") # an implicit transaction is open
|
||||
thread = threading.Thread(target=conn.retire)
|
||||
thread.start()
|
||||
thread.join(WAIT)
|
||||
with self.assertRaises(sqlite3.ProgrammingError):
|
||||
conn.commit()
|
||||
with self.assertRaises(sqlite3.ProgrammingError):
|
||||
with conn:
|
||||
pass
|
||||
with self.assertRaises(sqlite3.ProgrammingError):
|
||||
conn.rollback()
|
||||
self.assertEqual(self.value(), 0)
|
||||
|
||||
def test_close_is_retire(self):
|
||||
conn = self.connect()
|
||||
conn.close()
|
||||
self.assertTrue(conn.wait_closed(0))
|
||||
with self.assertRaises(sqlite3.ProgrammingError):
|
||||
conn.commit()
|
||||
|
||||
def test_plain_connections_are_closed_directly(self):
|
||||
plain = sqlite3.connect(self.path)
|
||||
self.assertEqual(retire_connections([plain]), [])
|
||||
with self.assertRaises(sqlite3.ProgrammingError):
|
||||
plain.execute("SELECT 1")
|
||||
|
||||
def test_awaiting_a_connection_still_in_use_times_out_with_a_warning(self):
|
||||
conn = self.connect()
|
||||
conn._enter() # its owner is inside a call
|
||||
self.assertEqual(retire_connections([conn]), [conn])
|
||||
logger = mock.Mock()
|
||||
self.assertFalse(await_closed([conn], 0, logger))
|
||||
logger.warning.assert_called_once()
|
||||
conn._leave()
|
||||
self.assertTrue(await_closed([conn], WAIT, logger))
|
||||
|
||||
|
||||
class TestCalls(_Case):
|
||||
def test_every_cursor_call_works_and_is_refused_once_retired(self):
|
||||
conn = self.connect()
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
self.assertIsInstance(cursor, ManagedCursor)
|
||||
conn.executemany("INSERT INTO t VALUES (?)", [(1,), (2,)])
|
||||
conn.executescript("INSERT INTO t VALUES (3);")
|
||||
self.assertEqual(
|
||||
[row["x"] for row in conn.execute("SELECT x FROM t ORDER BY x")],
|
||||
[0, 1, 2, 3],
|
||||
)
|
||||
cursor.execute("SELECT x FROM t ORDER BY x")
|
||||
self.assertEqual(cursor.fetchone()["x"], 0)
|
||||
self.assertEqual([r["x"] for r in cursor.fetchmany(2)], [1, 2])
|
||||
self.assertEqual([r["x"] for r in cursor.fetchall()], [3])
|
||||
cursor.executemany("UPDATE t SET x = x WHERE x = ?", [(1,)])
|
||||
cursor.executescript("SELECT 1;")
|
||||
cursor.close()
|
||||
conn.commit()
|
||||
conn.rollback()
|
||||
with conn:
|
||||
conn.execute("DELETE FROM t WHERE x = 3")
|
||||
with self.assertRaises(ZeroDivisionError):
|
||||
with conn:
|
||||
conn.execute("DELETE FROM t")
|
||||
raise ZeroDivisionError
|
||||
self.assertEqual(self.value(), 0)
|
||||
copy = sqlite3.connect(":memory:")
|
||||
self.opened.append(copy)
|
||||
conn.backup(copy)
|
||||
self.assertEqual(copy.execute("SELECT COUNT(*) FROM t").fetchone()[0], 3)
|
||||
|
||||
reader = conn.execute("SELECT x FROM t")
|
||||
conn.retire()
|
||||
for call in (
|
||||
lambda: reader.fetchone(),
|
||||
lambda: reader.fetchmany(1),
|
||||
lambda: reader.fetchall(),
|
||||
lambda: next(reader),
|
||||
lambda: reader.execute("SELECT 1"),
|
||||
lambda: reader.executemany("SELECT ?", [(1,)]),
|
||||
lambda: reader.executescript("SELECT 1;"),
|
||||
lambda: conn.backup(copy),
|
||||
):
|
||||
with self.assertRaises(sqlite3.ProgrammingError):
|
||||
call()
|
||||
# As with sqlite3's own cursors (the result varies by Python version).
|
||||
with contextlib.suppress(sqlite3.ProgrammingError):
|
||||
reader.close()
|
||||
|
||||
|
||||
class _CursorOpenedDuringClose(weakref.WeakSet):
|
||||
"""The cursor set of a connection being closed; a cursor opens mid-iteration.
|
||||
|
||||
``hook`` runs after the first cursor is yielded, where a thread switch
|
||||
could let the owner call ``conn.cursor()`` (Python 3.10-3.12 copy a
|
||||
WeakSet one item at a time). It runs in the closing thread itself, so
|
||||
the result does not depend on scheduling.
|
||||
"""
|
||||
|
||||
def __init__(self, items, hook):
|
||||
super().__init__(items)
|
||||
self.hook = hook
|
||||
|
||||
def __iter__(self):
|
||||
hooked = False
|
||||
for item in super().__iter__():
|
||||
yield item
|
||||
if not hooked:
|
||||
hooked = True
|
||||
self.hook()
|
||||
|
||||
|
||||
class TestReviewRound(_Case):
|
||||
"""PR #321 review: I1 (cursor set race), M1 and M2 (unguarded paths)."""
|
||||
|
||||
def test_a_cursor_opened_while_the_connection_closes_cannot_break_retire(self):
|
||||
conn = self.connect()
|
||||
kept = [conn.cursor() for _ in range(3)]
|
||||
refused = []
|
||||
|
||||
def open_cursor():
|
||||
try:
|
||||
kept.append(conn.cursor())
|
||||
except sqlite3.ProgrammingError as error:
|
||||
refused.append(error)
|
||||
|
||||
conn._cursors = _CursorOpenedDuringClose(conn._cursors, open_cursor)
|
||||
self.assertTrue(conn.retire()) # no "Set changed size during iteration"
|
||||
self.assertTrue(conn.wait_closed(0))
|
||||
self.assertEqual(len(refused), 1)
|
||||
self.assertEqual(len(kept), 3)
|
||||
|
||||
def test_a_retired_connection_refuses_new_cursors_and_scripts(self):
|
||||
conn = self.connect()
|
||||
conn._enter() # its owner is inside a call: retired, not yet closed
|
||||
try:
|
||||
self.assertFalse(conn.retire())
|
||||
for call in (
|
||||
conn.cursor,
|
||||
lambda: conn.executescript("CREATE TABLE late(x);"),
|
||||
lambda: conn.in_transaction,
|
||||
lambda: conn.create_function("f", 0, lambda: 1),
|
||||
):
|
||||
with self.subTest(call=call):
|
||||
with self.assertRaises(sqlite3.ProgrammingError):
|
||||
call()
|
||||
finally:
|
||||
conn._leave()
|
||||
self.assertTrue(conn.wait_closed(0))
|
||||
with contextlib.closing(sqlite3.connect(self.path)) as check:
|
||||
tables = check.execute(
|
||||
"SELECT name FROM sqlite_master WHERE name='late'"
|
||||
).fetchall()
|
||||
self.assertEqual(tables, [])
|
||||
|
||||
def test_one_connection_failing_to_retire_does_not_skip_the_rest(self):
|
||||
class Broken(ManagedConnection):
|
||||
def retire(self):
|
||||
raise RuntimeError("retire failed")
|
||||
|
||||
broken = sqlite3.connect(self.path, check_same_thread=False, factory=Broken)
|
||||
try:
|
||||
after = self.connect()
|
||||
logger = mock.Mock()
|
||||
pending = retire_connections([broken, after], logger)
|
||||
self.assertEqual(pending, [])
|
||||
self.assertTrue(after.wait_closed(0))
|
||||
logger.exception.assert_called_once()
|
||||
finally:
|
||||
ManagedConnection.retire(broken)
|
||||
|
||||
def test_a_restore_marks_its_backup_target_in_use(self):
|
||||
model = PromptModel(os.path.join(self._tmp.name, "live.db"))
|
||||
try:
|
||||
self._restore_marks_its_target(model)
|
||||
finally:
|
||||
model.close_all()
|
||||
|
||||
def _restore_marks_its_target(self, model):
|
||||
source_path = os.path.join(self._tmp.name, "source.db")
|
||||
PromptModel(source_path).close_all()
|
||||
target = model.get_connection()
|
||||
depths = []
|
||||
connect = sqlite3.connect
|
||||
|
||||
class Source(sqlite3.Connection):
|
||||
def backup(self, into, *args, **kwargs):
|
||||
depths.append(into._depth)
|
||||
return super().backup(into, *args, **kwargs)
|
||||
|
||||
def connect_source(*args, **kwargs):
|
||||
if "factory" in kwargs:
|
||||
return connect(*args, **kwargs)
|
||||
return connect(*args, **kwargs, factory=Source)
|
||||
|
||||
with mock.patch("database.models.sqlite3.connect", connect_source):
|
||||
model._replace_contents_from(source_path)
|
||||
self.assertEqual(depths, [1])
|
||||
self.assertEqual(target._depth, 0)
|
||||
|
||||
|
||||
class TestModelConnections(_Case):
|
||||
def test_the_model_hands_out_managed_connections_and_closes_them(self):
|
||||
model = PromptModel(os.path.join(self._tmp.name, "model.db"))
|
||||
conn = model.get_connection()
|
||||
self.assertIsInstance(conn, ManagedConnection)
|
||||
model.close()
|
||||
self.assertTrue(conn.wait_closed(0))
|
||||
conn = model.get_connection()
|
||||
model.close_all()
|
||||
self.assertTrue(conn.wait_closed(0))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -99,16 +99,11 @@ class StrictRows:
|
||||
"""Assertions over the rows a write left behind."""
|
||||
|
||||
def image_rows(self, db):
|
||||
"""The stored texts, read through the one workflow document accessor."""
|
||||
from database.workflow_documents import DocumentReader, resolve_image_row
|
||||
|
||||
with db.model.get_connection() as conn:
|
||||
reader = DocumentReader(conn)
|
||||
return [
|
||||
resolve_image_row(reader, row)
|
||||
dict(row)
|
||||
for row in conn.execute(
|
||||
"SELECT prompt_id, workflow_data, prompt_metadata, parameters, "
|
||||
"workflow_document_id, prompt_document_id "
|
||||
"SELECT prompt_id, workflow_data, prompt_metadata, parameters "
|
||||
"FROM generated_images ORDER BY id"
|
||||
)
|
||||
]
|
||||
|
||||
+18
-29
@@ -171,13 +171,6 @@ class ReviewTestCase(AdminAPITestCase):
|
||||
def _require_module(self):
|
||||
self.assertIsNotNone(review_capture, "utils.review_capture is not implemented")
|
||||
|
||||
def _scan_record(self, path):
|
||||
"""Record ``path`` as the scan's recorder does (only scans add items)."""
|
||||
from py.api.admin import _scan_roots
|
||||
|
||||
roots = _scan_roots(self.api._get_all_output_dirs())
|
||||
return self.api._record_review_sync(Path(str(path)), roots, inspection=True)
|
||||
|
||||
|
||||
class TestScanRecording(ReviewTestCase):
|
||||
async def test_scan_persists_only_ambiguous_files_with_capture_digest(self):
|
||||
@@ -509,13 +502,13 @@ class TestNonFiniteGraphs(ReviewTestCase):
|
||||
["cat NaN", "dog Infinity"],
|
||||
)
|
||||
|
||||
async def test_scan_recorder_records_a_nan_stamped_graph(self):
|
||||
async def test_record_review_file_records_a_nan_stamped_graph(self):
|
||||
self._require_module()
|
||||
self._roots()
|
||||
path = self.output_dir / "nan.png"
|
||||
write_png(path, prompt=self._nan_stamped("cat", "dog"))
|
||||
|
||||
self.assertEqual(self._scan_record(path), "recorded")
|
||||
self.assertEqual(self.api.record_review_file(str(path)), "recorded")
|
||||
(item,) = self._items()
|
||||
self.assertEqual(item["candidate_count"], 2)
|
||||
|
||||
@@ -863,13 +856,12 @@ class TestMonitorRecording(ReviewTestCase):
|
||||
async def test_monitor_recording_is_refused_beside_record_changing_jobs(self):
|
||||
"""S1 admission: a ``records`` recording waits for no job; it is refused.
|
||||
|
||||
While a ``records`` or ``guarded`` job runs, the monitor's pass
|
||||
writes nothing and answers ``deferred`` (its fallback link stands and
|
||||
While a ``records`` or ``guarded`` job runs, the monitor's recording
|
||||
records nothing and answers ``deferred`` (its fallback link stands and
|
||||
a later scan reads the ambiguous file again); a ``read`` job does not
|
||||
refuse it. Thumbnail work (the job and the gallery's holds) is
|
||||
``guarded`` but writes only previews, never records, so it does not
|
||||
defer the pass either (S5 M7). Not deferred, the pass still records
|
||||
nothing: only scans add items, so an unlinked output is ``skipped``.
|
||||
refuse it, and it records once the blocking job has ended. Thumbnail
|
||||
work (the job and the gallery's holds) is ``guarded`` but writes only
|
||||
previews, never records, so it does not defer the recording (S5 M7).
|
||||
"""
|
||||
self._require_module()
|
||||
from py.api.jobs import RESTORE_JOB_KIND, THUMBNAILS_JOB_KIND
|
||||
@@ -906,11 +898,11 @@ class TestMonitorRecording(ReviewTestCase):
|
||||
outcome = await asyncio.to_thread(
|
||||
self.api.record_review_file, str(path)
|
||||
)
|
||||
self.assertEqual(outcome, "skipped")
|
||||
self.assertEqual(self._items(), [])
|
||||
self.assertEqual(outcome, "recorded")
|
||||
self.assertEqual(len(self._items()), 1)
|
||||
|
||||
async def test_a_running_thumbnail_job_does_not_defer_monitor_recording(self):
|
||||
"""S5 M7: a long rebuild never defers the monitor's pass (no records)."""
|
||||
"""S5 M7: a long rebuild never holds Needs review back (no records)."""
|
||||
self._require_module()
|
||||
from py.api.jobs import THUMBNAILS_JOB_KIND
|
||||
|
||||
@@ -935,9 +927,8 @@ class TestMonitorRecording(ReviewTestCase):
|
||||
finally:
|
||||
gate.set()
|
||||
await asyncio.wait_for(job.task, 5)
|
||||
# Not deferred (thumbnail work writes no records); only scans add items.
|
||||
self.assertEqual(outcome, "skipped")
|
||||
self.assertEqual(self._items(), [])
|
||||
self.assertEqual(outcome, "recorded")
|
||||
self.assertEqual(len(self._items()), 1)
|
||||
|
||||
async def test_monitor_wiring_reaches_new_and_existing_handlers(self):
|
||||
self._require_module()
|
||||
@@ -1005,12 +996,12 @@ class TestRecordedSpelling(ReviewTestCase):
|
||||
upper = self.output_dir / "Two.PNG"
|
||||
write_png(upper, prompt=unlinkable_graph("cat", "dog"))
|
||||
with _case_folding():
|
||||
self.assertEqual(self._scan_record(upper), "recorded")
|
||||
self.assertEqual(self.api.record_review_file(str(upper)), "recorded")
|
||||
(before,) = self._items()
|
||||
revision = self._review_revision()
|
||||
lower = self.output_dir / "two.png"
|
||||
os.rename(upper, lower) # keeps size and mtime_ns
|
||||
self.assertEqual(self._scan_record(lower), "duplicate")
|
||||
self.assertEqual(self.api.record_review_file(str(lower)), "duplicate")
|
||||
(after,) = self._items()
|
||||
self.assertEqual(self._spelled(), [str(lower)])
|
||||
self.assertEqual(self._review_revision(), revision + 1)
|
||||
@@ -1038,7 +1029,7 @@ class TestRecordedSpelling(ReviewTestCase):
|
||||
self.assertEqual(capture.locator, after["locator"])
|
||||
self.api._captures.release(capture.token)
|
||||
# Seeing the same spelling again writes nothing.
|
||||
self.assertEqual(self._scan_record(lower), "duplicate")
|
||||
self.assertEqual(self.api.record_review_file(str(lower)), "duplicate")
|
||||
self.assertEqual(self._review_revision(), revision + 1)
|
||||
|
||||
async def test_duplicate_insert_refreshes_the_spelling(self):
|
||||
@@ -1079,8 +1070,7 @@ class TestRecordedSpelling(ReviewTestCase):
|
||||
|
||||
class TestCaptureStoreShutdown(ReviewTestCase):
|
||||
async def test_recording_after_close_creates_no_store(self):
|
||||
"""A scan recording still in flight at shutdown leaks nothing (M2),
|
||||
and the monitor's pass after it captures nothing at all."""
|
||||
"""A monitor recording still in flight at shutdown leaks nothing (M2)."""
|
||||
self._require_module()
|
||||
self._roots()
|
||||
path = self.output_dir / "two.png"
|
||||
@@ -1089,12 +1079,11 @@ class TestCaptureStoreShutdown(ReviewTestCase):
|
||||
self.assertIsNone(self.api._captures)
|
||||
with patch("utils.media_capture.tempfile.mkdtemp") as mkdtemp:
|
||||
with self.assertRaises(CaptureError):
|
||||
self._scan_record(path)
|
||||
# The monitor's pass records nothing, so it needs no store.
|
||||
self.api.record_review_file(str(path))
|
||||
# The monitor logs the failure and keeps its link.
|
||||
handler = ImageGenerationHandler(self.api.db, PromptTracker(self.api.db))
|
||||
handler.review_recorder = self.api.record_review_file
|
||||
handler._record_review(str(path))
|
||||
self.assertEqual(self.api.record_review_file(str(path)), "skipped")
|
||||
mkdtemp.assert_not_called()
|
||||
self.assertIsNone(self.api._captures)
|
||||
self.assertIsNone(self.api._captures_unsubscribe)
|
||||
|
||||
@@ -202,7 +202,7 @@ class TestNothingToChooseIsNotQueued(ChoosableTestCase):
|
||||
self._roots()
|
||||
path = self._write("none.png", opaque_graph("cat", "dog"))
|
||||
|
||||
self.assertEqual(self.api.record_review_file(str(path)), "skipped")
|
||||
self.assertEqual(self.api.record_review_file(str(path)), "untraced")
|
||||
self.assertEqual(self._items(), [])
|
||||
|
||||
|
||||
@@ -320,18 +320,14 @@ class TestUntracedItemsLeaveTheQueue(ChoosableTestCase):
|
||||
await self._scan("all")
|
||||
self.assertEqual(self._pending(), 0)
|
||||
|
||||
async def test_the_scan_not_the_monitor_removes_an_untraced_outputs_item(self):
|
||||
"""The monitor only clears items of outputs it links; the next scan
|
||||
removes the item of an output that offers nothing to choose."""
|
||||
def test_monitor_removes_the_pending_item_of_an_untraced_output(self):
|
||||
self._roots()
|
||||
path = self._write("none.png", opaque_graph("cat", "dog"))
|
||||
self._record_pre_fix("none.png")
|
||||
# A pre-fix item of other bytes keeps the stat shortcut out of the way.
|
||||
self._write("none.png", opaque_graph("a cow", "dog"))
|
||||
|
||||
self.assertEqual(self.api.record_review_file(str(path)), "skipped")
|
||||
self.assertEqual(len(self._items()), 1)
|
||||
await self._scan("all")
|
||||
self.assertEqual(self.api.record_review_file(str(path)), "untraced")
|
||||
self.assertEqual(self._items(), [])
|
||||
|
||||
|
||||
|
||||
@@ -1,971 +0,0 @@
|
||||
"""Needs review choices from the inspection profile (follow-up to #284/#288).
|
||||
|
||||
A Qwen image-edit batch runs several KSamplers, each with its own
|
||||
``TextEncodeQwenImageEdit(prompt=...)``, and every SaveImage shares one
|
||||
prefix: the file cannot say which prompt made it. The scan's default
|
||||
(ingestion) profile does not know the encoder or the string nodes, so every
|
||||
candidate is unresolved and the file used to leave the queue as "could not
|
||||
be traced". The inspection profile resolves every text, so these are real
|
||||
choices. Ruling (option a): when the default profile is ambiguous with fewer
|
||||
than two candidates that carry text, the scan asks the inspection profile;
|
||||
with at least two choosable candidates it records the item from that
|
||||
profile, and Apply re-resolves with the same profile. Nothing links
|
||||
automatically from the inspection profile, and the live monitor never queues.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
try:
|
||||
from tests.test_api_review import ReviewRoutesTestCase
|
||||
from tests.test_review import unlinkable_graph, write_png
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from test_api_review import ReviewRoutesTestCase
|
||||
from test_review import unlinkable_graph, write_png
|
||||
|
||||
from utils import inspection_semantics, review_capture
|
||||
from utils.hashing import generate_prompt_hash
|
||||
from utils.image_monitor import ImageGenerationHandler
|
||||
from utils.prompt_tracker import PromptTracker
|
||||
|
||||
NAME = "kittybear_00001_.png"
|
||||
PROMPTS = ("make it night", "add a red hat", "turn the cat into a bear")
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _counted(module, name):
|
||||
"""Count the calls of ``module.name`` while calling through."""
|
||||
real = getattr(module, name)
|
||||
calls = []
|
||||
|
||||
def spy(*args, **kwargs):
|
||||
calls.append(args)
|
||||
return real(*args, **kwargs)
|
||||
|
||||
with patch.object(module, name, side_effect=spy):
|
||||
yield calls
|
||||
|
||||
|
||||
def qwen_edit_graph(*prompts, joined=()):
|
||||
"""One Qwen image-edit sampler per prompt, all saved under one prefix.
|
||||
|
||||
Each sampler has its own ``TextEncodeQwenImageEdit``, used for both
|
||||
positive and negative as the real batches do. ``joined`` prompts are
|
||||
``(first, second)`` pairs joined by ``JoinStringMulti`` over two
|
||||
``StringConstantMultiline`` nodes. Every SaveImage takes its prefix from
|
||||
one linked string node.
|
||||
"""
|
||||
graph = {
|
||||
"1": {"class_type": "UNETLoader", "inputs": {"unet_name": "q.safetensors"}},
|
||||
"2": {"class_type": "CLIPLoader", "inputs": {"clip_name": "c.safetensors"}},
|
||||
"3": {"class_type": "VAELoader", "inputs": {"vae_name": "v.safetensors"}},
|
||||
"4": {"class_type": "LoadImage", "inputs": {"image": "source.png"}},
|
||||
"5": {"class_type": "PrimitiveString", "inputs": {"value": "kittybear"}},
|
||||
"6": {
|
||||
"class_type": "EmptyLatentImage",
|
||||
"inputs": {"width": 512, "height": 512, "batch_size": 1},
|
||||
},
|
||||
}
|
||||
sides = [("literal", text) for text in prompts]
|
||||
sides += [("joined", parts) for parts in joined]
|
||||
for index, (kind, text) in enumerate(sides):
|
||||
base = 100 + 10 * index
|
||||
prompt = text
|
||||
if kind == "joined":
|
||||
for offset, part in enumerate(text, start=1):
|
||||
graph[str(base + offset)] = {
|
||||
"class_type": "StringConstantMultiline",
|
||||
"inputs": {"string": part, "strip_newlines": True},
|
||||
}
|
||||
graph[str(base + 3)] = {
|
||||
"class_type": "JoinStringMulti",
|
||||
"inputs": {
|
||||
"inputcount": 2,
|
||||
"delimiter": ", ",
|
||||
"return_list": False,
|
||||
"string_1": [str(base + 1), 0],
|
||||
"string_2": [str(base + 2), 0],
|
||||
},
|
||||
}
|
||||
prompt = [str(base + 3), 0]
|
||||
encoder, sampler = str(base), str(base + 5)
|
||||
graph[encoder] = {
|
||||
"class_type": "TextEncodeQwenImageEdit",
|
||||
"inputs": {
|
||||
"clip": ["2", 0],
|
||||
"prompt": prompt,
|
||||
"vae": ["3", 0],
|
||||
"image": ["4", 0],
|
||||
},
|
||||
}
|
||||
graph[sampler] = {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {
|
||||
"model": ["1", 0],
|
||||
"positive": [encoder, 0],
|
||||
"negative": [encoder, 0],
|
||||
"latent_image": ["6", 0],
|
||||
"seed": index,
|
||||
},
|
||||
}
|
||||
graph[str(base + 6)] = {
|
||||
"class_type": "VAEDecode",
|
||||
"inputs": {"samples": [sampler, 0], "vae": ["3", 0]},
|
||||
}
|
||||
graph[str(base + 7)] = {
|
||||
"class_type": "SaveImage",
|
||||
"inputs": {"images": [str(base + 6), 0], "filename_prefix": ["5", 0]},
|
||||
}
|
||||
return graph
|
||||
|
||||
|
||||
def _prefixed(graph, prefix):
|
||||
"""``graph`` with every node id (and link to one) prefixed."""
|
||||
renamed = {}
|
||||
for node_id, node in graph.items():
|
||||
inputs = {
|
||||
key: (
|
||||
[prefix + value[0], value[1]]
|
||||
if isinstance(value, list) and len(value) == 2
|
||||
else value
|
||||
)
|
||||
for key, value in node["inputs"].items()
|
||||
}
|
||||
renamed[prefix + node_id] = {**node, "inputs": inputs}
|
||||
return renamed
|
||||
|
||||
|
||||
def _version(graph=None):
|
||||
"""The item version of a Qwen-edit graph under the current table."""
|
||||
graph = graph if graph is not None else qwen_edit_graph(*PROMPTS)
|
||||
return review_capture.inspection_review_version({"prompt": graph})
|
||||
|
||||
|
||||
def _retarget_unet_loader(doc):
|
||||
"""A table edit for a class the Qwen-edit graph uses (no prompt effect)."""
|
||||
doc["facts"]["model_outputs"]["UNETLoader"] = "unet"
|
||||
|
||||
|
||||
def _unrelated_encoder(doc):
|
||||
"""A table edit for a class the Qwen-edit graph does not use."""
|
||||
doc["encoders"].append(
|
||||
{
|
||||
"class": "UnrelatedTextEncode",
|
||||
"slot": 0,
|
||||
"strategy": "single",
|
||||
"inputs": ["text"],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _texts(evidence):
|
||||
return sorted(
|
||||
c["positive"]["text"]
|
||||
for c in evidence["candidates"]
|
||||
if isinstance(c.get("positive"), dict)
|
||||
)
|
||||
|
||||
|
||||
class InspectionTestCase(ReviewRoutesTestCase):
|
||||
def _write(self, name=NAME, graph=None, color=(255, 0, 0)):
|
||||
path = self.output_dir / name
|
||||
write_png(
|
||||
path,
|
||||
prompt=graph if graph is not None else qwen_edit_graph(*PROMPTS),
|
||||
color=color,
|
||||
)
|
||||
return path
|
||||
|
||||
def _pending(self):
|
||||
return self.api.db.library_counts()["review_pending"]
|
||||
|
||||
def _evidence(self, item):
|
||||
return json.loads(item["evidence"])
|
||||
|
||||
def _table(self, edit):
|
||||
"""Point the inspection profile at an edited copy of the table."""
|
||||
doc = json.loads(inspection_semantics.SEMANTICS_PATH.read_text("utf-8"))
|
||||
edit(doc)
|
||||
folder = tempfile.mkdtemp(dir=self.tmpdir)
|
||||
table = Path(folder) / "inspection-semantics.json"
|
||||
table.write_text(json.dumps(doc), encoding="utf-8")
|
||||
patcher = patch.object(inspection_semantics, "SEMANTICS_PATH", table)
|
||||
patcher.start()
|
||||
inspection_semantics.get_semantics.cache_clear()
|
||||
self.addCleanup(inspection_semantics.get_semantics.cache_clear)
|
||||
self.addCleanup(patcher.stop)
|
||||
|
||||
async def _apply(self, positive):
|
||||
item, detail, observed = await self._open()
|
||||
candidate = self._candidate(detail, positive)
|
||||
return item, await self._decide(
|
||||
item["id"], self._body(detail, observed, candidate["hash"])
|
||||
)
|
||||
|
||||
def _linked_text(self, path):
|
||||
row = (
|
||||
self.api.db.model.get_connection()
|
||||
.execute(
|
||||
"SELECT p.text FROM generated_images gi "
|
||||
"JOIN prompts p ON p.id=gi.prompt_id WHERE gi.image_path=?",
|
||||
(str(path),),
|
||||
)
|
||||
.fetchone()
|
||||
)
|
||||
return row[0] if row else None
|
||||
|
||||
|
||||
class TestQwenEditBatchIsQueued(InspectionTestCase):
|
||||
async def test_qwen_edit_batch_is_queued_with_its_texts(self):
|
||||
self._roots()
|
||||
graph = qwen_edit_graph(*PROMPTS, joined=[("a red scarf", "snow")])
|
||||
self._write(graph=graph)
|
||||
|
||||
job = await self._scan()
|
||||
|
||||
(item,) = self._items()
|
||||
evidence = self._evidence(item)
|
||||
self.assertEqual(_texts(evidence), sorted(PROMPTS + ("a red scarf, snow",)))
|
||||
self.assertEqual(item["candidate_count"], 4)
|
||||
self.assertEqual(evidence["profile"], "inspection")
|
||||
self.assertEqual(
|
||||
evidence["semantics"], review_capture.semantics_version({"prompt": graph})
|
||||
)
|
||||
self.assertEqual(
|
||||
item["resolver_version"],
|
||||
review_capture.inspection_review_version({"prompt": graph}),
|
||||
)
|
||||
self.assertNotEqual(item["resolver_version"], review_capture.RESOLVER_VERSION)
|
||||
self.assertEqual(job["scope"]["review_items"], 1)
|
||||
(row,) = job["outcomes"]["review"]
|
||||
self.assertEqual(row["path"], NAME)
|
||||
self.assertIn("more than one prompt could match", row["detail"])
|
||||
# Review choices only: nothing linked the file.
|
||||
self.assertIsNone(self._linked_text(self.output_dir / NAME))
|
||||
self.assertEqual(self._capture_files(), [])
|
||||
|
||||
async def test_rescan_keeps_one_item_and_skips_the_copy(self):
|
||||
self._roots()
|
||||
self._write()
|
||||
await self._scan()
|
||||
(item,) = self._items()
|
||||
|
||||
with _counted(self.api._captures, "capture") as copies:
|
||||
job = await self._scan("all")
|
||||
|
||||
self.assertEqual(copies, [])
|
||||
self.assertEqual([i["id"] for i in self._items()], [item["id"]])
|
||||
self.assertEqual(job["scope"]["review_items"], 1)
|
||||
self.assertEqual(job["scope"]["review_untraced"], 0)
|
||||
|
||||
async def test_more_than_sixteen_qwen_samplers_make_a_limited_item(self):
|
||||
self._roots()
|
||||
self._write(graph=qwen_edit_graph(*(f"edit {n}" for n in range(17))))
|
||||
|
||||
await self._scan()
|
||||
|
||||
(item,) = self._items()
|
||||
evidence = self._evidence(item)
|
||||
self.assertTrue(item["limited"])
|
||||
self.assertEqual(item["candidate_count"], 0)
|
||||
self.assertEqual(evidence["candidate_total"], 17)
|
||||
self.assertEqual(evidence["profile"], "inspection")
|
||||
|
||||
async def test_evidence_over_the_cap_makes_a_limited_item(self):
|
||||
self._roots()
|
||||
self._write()
|
||||
real = review_capture.canonical_evidence
|
||||
|
||||
def capped(value):
|
||||
if value.get("candidates"):
|
||||
raise ValueError("evidence exceeds the bound")
|
||||
return real(value)
|
||||
|
||||
with patch.object(review_capture, "canonical_evidence", side_effect=capped):
|
||||
await self._scan()
|
||||
|
||||
(item,) = self._items()
|
||||
self.assertTrue(item["limited"])
|
||||
self.assertEqual(self._evidence(item)["candidate_total"], 3)
|
||||
self.assertEqual(self._evidence(item)["profile"], "inspection")
|
||||
|
||||
|
||||
class TestDefaultProfileStillDecides(InspectionTestCase):
|
||||
async def test_ambiguous_file_with_texts_keeps_the_default_profile(self):
|
||||
"""A default-ambiguous file with a choice never runs the inspection
|
||||
profile, even when that profile would offer more candidates."""
|
||||
self._roots()
|
||||
graph = {
|
||||
**_prefixed(qwen_edit_graph("make it night"), "q"),
|
||||
**unlinkable_graph("cat", "dog"),
|
||||
}
|
||||
self._write(graph=graph)
|
||||
|
||||
with patch.object(
|
||||
review_capture,
|
||||
"resolve_inspection",
|
||||
side_effect=AssertionError("inspection profile ran"),
|
||||
):
|
||||
await self._scan()
|
||||
|
||||
(item,) = self._items()
|
||||
evidence = self._evidence(item)
|
||||
self.assertNotIn("profile", evidence)
|
||||
self.assertNotIn("semantics", evidence)
|
||||
self.assertEqual(item["resolver_version"], review_capture.RESOLVER_VERSION)
|
||||
self.assertEqual(_texts(evidence), ["cat", "dog"])
|
||||
|
||||
async def test_apply_of_a_default_item_never_runs_the_inspection_profile(self):
|
||||
self._roots()
|
||||
write_png(self.output_dir / "two.png", prompt=unlinkable_graph("cat", "dog"))
|
||||
await self._scan()
|
||||
with patch.object(
|
||||
review_capture,
|
||||
"resolve_inspection",
|
||||
side_effect=AssertionError("inspection profile ran"),
|
||||
):
|
||||
_item, (status, data) = await self._apply("cat")
|
||||
self.assertEqual(status, 200, data)
|
||||
self.assertEqual(self._linked_text(self.output_dir / "two.png"), "cat")
|
||||
|
||||
|
||||
class TestApplyInspectionCandidate(InspectionTestCase):
|
||||
async def test_apply_links_the_chosen_qwen_edit_text(self):
|
||||
self._require_routes()
|
||||
self._roots()
|
||||
path = self._write()
|
||||
await self._scan()
|
||||
|
||||
item, (status, data) = await self._apply("add a red hat")
|
||||
|
||||
self.assertEqual(status, 200, data)
|
||||
self.assertEqual(data["item"]["state"], "selected")
|
||||
self.assertEqual(data["total"], 0)
|
||||
self.assertEqual(self._linked_text(path), "add a red hat")
|
||||
links = self.api.db.get_image_prompt_pairs(data["pair"]["image_id"])
|
||||
(link,) = links
|
||||
provenance = link["provenance"]
|
||||
self.assertEqual(provenance["kind"], "user-selection")
|
||||
self.assertEqual(provenance["resolver_version"], _version())
|
||||
self.assertEqual(provenance["positive"]["text"], "add a red hat")
|
||||
(decision,) = self._decisions()
|
||||
self.assertEqual(decision["item_id"], item["id"])
|
||||
self._no_paths()
|
||||
|
||||
async def test_apply_rechecks_with_the_recorded_profile(self):
|
||||
"""The default profile would not offer the candidate: Apply must use
|
||||
the profile the item recorded, and refuses when that profile no
|
||||
longer offers it."""
|
||||
self._require_routes()
|
||||
self._roots()
|
||||
self._write()
|
||||
await self._scan()
|
||||
calls = []
|
||||
real = review_capture.resolve_inspection
|
||||
|
||||
def spy(metadata):
|
||||
calls.append(metadata)
|
||||
return real(metadata)
|
||||
|
||||
with patch.object(review_capture, "resolve_inspection", side_effect=spy):
|
||||
_item, (status, data) = await self._apply("make it night")
|
||||
|
||||
self.assertEqual(status, 200, data)
|
||||
self.assertEqual(len(calls), 1)
|
||||
|
||||
async def test_table_change_keeps_the_item_applicable(self):
|
||||
"""A table change for a class the graph uses moves the item's
|
||||
version; the pending item stays applicable while its candidate is
|
||||
still offered."""
|
||||
self._require_routes()
|
||||
self._roots()
|
||||
path = self._write()
|
||||
await self._scan()
|
||||
(old,) = self._items()
|
||||
version = _version()
|
||||
item, detail, observed = await self._open()
|
||||
|
||||
self._table(_retarget_unet_loader)
|
||||
self.assertNotEqual(_version(), version)
|
||||
|
||||
candidate = self._candidate(detail, "add a red hat")
|
||||
status, data = await self._decide(
|
||||
item["id"], self._body(detail, observed, candidate["hash"])
|
||||
)
|
||||
self.assertEqual(status, 200, data)
|
||||
self.assertEqual(item["id"], old["id"])
|
||||
self.assertEqual(self._linked_text(path), "add a red hat")
|
||||
|
||||
async def test_table_change_that_drops_the_encoder_is_a_review_conflict(self):
|
||||
self._require_routes()
|
||||
self._roots()
|
||||
self._write()
|
||||
await self._scan()
|
||||
item, detail, observed = await self._open()
|
||||
candidate = self._candidate(detail, "add a red hat")
|
||||
|
||||
def drop_encoder(doc):
|
||||
doc["encoders"] = [
|
||||
entry
|
||||
for entry in doc["encoders"]
|
||||
if entry["class"] != "TextEncodeQwenImageEdit"
|
||||
]
|
||||
|
||||
version = _version()
|
||||
self._table(drop_encoder)
|
||||
self.assertNotEqual(_version(), version)
|
||||
before = self._snapshot()
|
||||
status, data = await self._decide(
|
||||
item["id"], self._body(detail, observed, candidate["hash"])
|
||||
)
|
||||
|
||||
self.assertEqual(status, 409, data)
|
||||
self.assertEqual(data["code"], "review-conflict")
|
||||
self.assertEqual(self._snapshot(), before)
|
||||
# The rescan finds nothing to choose under the new version: the
|
||||
# stale, undecided item leaves the queue instead of staying stuck.
|
||||
job = await self._scan("all")
|
||||
self.assertEqual(self._items(), [])
|
||||
self.assertEqual(job["scope"]["review_untraced"], 1)
|
||||
|
||||
|
||||
class TestMonitorNeverQueues(InspectionTestCase):
|
||||
def _handler(self):
|
||||
handler = ImageGenerationHandler(self.api.db, PromptTracker(self.api.db))
|
||||
handler.review_recorder = self.api.record_review_file
|
||||
return handler
|
||||
|
||||
async def test_monitor_pass_never_queues_but_the_next_scan_does(self):
|
||||
self._roots()
|
||||
path = self._write()
|
||||
|
||||
self.assertEqual(self.api.record_review_file(str(path)), "skipped")
|
||||
self.assertEqual(self._items(), [])
|
||||
self.assertEqual(self._capture_files(), [])
|
||||
|
||||
await self._scan()
|
||||
(item,) = self._items()
|
||||
self.assertEqual(self._evidence(item)["profile"], "inspection")
|
||||
|
||||
async def test_monitor_keeps_the_inspection_item_of_an_unlinked_output(self):
|
||||
"""The monitor's untraced settle judges each item by its own
|
||||
evidence: an inspection item offers a choice and stays."""
|
||||
self._roots()
|
||||
path = self._write()
|
||||
await self._scan()
|
||||
(item,) = self._items()
|
||||
# Other bytes, so the stat shortcut does not answer for the file.
|
||||
self._write(color=(0, 0, 255))
|
||||
|
||||
self.assertEqual(self.api.record_review_file(str(path)), "skipped")
|
||||
self.assertEqual([i["id"] for i in self._items()], [item["id"]])
|
||||
|
||||
async def test_monitor_link_removes_the_inspection_item(self):
|
||||
self._roots()
|
||||
path = self._write()
|
||||
await self._scan()
|
||||
self.assertEqual(len(self._items()), 1)
|
||||
prompt_id = self.api.db.save_prompt(
|
||||
text="latest", prompt_hash=generate_prompt_hash("latest")
|
||||
)
|
||||
|
||||
self._handler().link_image_to_prompt(
|
||||
str(path),
|
||||
{"id": prompt_id, "text": "latest", "fallback": True},
|
||||
{"prompt": qwen_edit_graph(*PROMPTS)},
|
||||
)
|
||||
|
||||
self.assertEqual(self._items(), [])
|
||||
self.assertEqual(os.path.basename(str(path)), NAME)
|
||||
|
||||
|
||||
class TestTableVersion(InspectionTestCase):
|
||||
def test_table_version_follows_content_not_layout(self):
|
||||
doc = {"schema": 1, "encoders": [{"class": "A", "slot": 0}]}
|
||||
same = json.loads(json.dumps(doc, indent=4))
|
||||
other = {"schema": 1, "encoders": [{"class": "B", "slot": 0}]}
|
||||
|
||||
version = inspection_semantics.table_version(doc)
|
||||
|
||||
self.assertRegex(version, "^[0-9a-f]{64}$")
|
||||
self.assertEqual(inspection_semantics.table_version(same), version)
|
||||
self.assertNotEqual(inspection_semantics.table_version(other), version)
|
||||
self.assertEqual(
|
||||
inspection_semantics.load_semantics().version,
|
||||
inspection_semantics.table_version(
|
||||
json.loads(inspection_semantics.SEMANTICS_PATH.read_text("utf-8"))
|
||||
),
|
||||
)
|
||||
|
||||
async def test_an_unusable_table_queues_nothing(self):
|
||||
self._roots()
|
||||
self._write()
|
||||
unusable = inspection_semantics.SemanticsError("bad table")
|
||||
with (
|
||||
patch.object(inspection_semantics, "get_semantics", side_effect=unusable),
|
||||
patch.object(review_capture, "get_semantics", side_effect=unusable),
|
||||
):
|
||||
self.assertIsNone(review_capture.inspection_review_version())
|
||||
self.assertIsNone(_version())
|
||||
job = await self._scan()
|
||||
|
||||
self.assertEqual(self._items(), [])
|
||||
(row,) = job["outcomes"]["review"]
|
||||
self.assertIn("could not be traced", row["detail"])
|
||||
|
||||
async def test_an_unusable_table_keeps_existing_items_on_rescan(self):
|
||||
self._require_routes()
|
||||
self._roots()
|
||||
self._write()
|
||||
await self._scan()
|
||||
(item,) = self._items()
|
||||
|
||||
unusable = inspection_semantics.SemanticsError("bad table")
|
||||
with (
|
||||
patch.object(inspection_semantics, "get_semantics", side_effect=unusable),
|
||||
patch.object(review_capture, "get_semantics", side_effect=unusable),
|
||||
):
|
||||
await self._scan("all")
|
||||
|
||||
self.assertEqual([i["id"] for i in self._items()], [item["id"]])
|
||||
|
||||
|
||||
class TestVersionChanges(InspectionTestCase):
|
||||
"""A version change never leaves two pending items for one file."""
|
||||
|
||||
async def test_table_change_for_a_graph_class_leaves_one_item(self):
|
||||
self._require_routes()
|
||||
self._roots()
|
||||
path = self._write()
|
||||
await self._scan()
|
||||
(old,) = self._items()
|
||||
|
||||
self._table(_retarget_unet_loader)
|
||||
job = await self._scan("all")
|
||||
|
||||
(item,) = self._items()
|
||||
self.assertNotEqual(item["id"], old["id"])
|
||||
self.assertEqual(item["resolver_version"], _version())
|
||||
self.assertNotEqual(item["resolver_version"], old["resolver_version"])
|
||||
self.assertEqual(job["scope"]["review_items"], 1)
|
||||
self.assertEqual(self._pending(), 1)
|
||||
_item, (status, data) = await self._apply("make it night")
|
||||
self.assertEqual(status, 200, data)
|
||||
self.assertEqual(self._linked_text(path), "make it night")
|
||||
|
||||
async def test_unrelated_table_edit_keeps_the_item_and_does_no_work(self):
|
||||
self._roots()
|
||||
self._write()
|
||||
await self._scan()
|
||||
(old,) = self._items()
|
||||
version = _version()
|
||||
|
||||
self._table(_unrelated_encoder)
|
||||
self.assertEqual(_version(), version)
|
||||
self.assertNotEqual(
|
||||
inspection_semantics.get_semantics().version,
|
||||
json.loads(old["evidence"])["semantics"],
|
||||
)
|
||||
with (
|
||||
_counted(self.api._captures, "capture") as copies,
|
||||
_counted(review_capture, "resolve_inspection") as calls,
|
||||
):
|
||||
await self._scan("all")
|
||||
|
||||
self.assertEqual(copies, [])
|
||||
self.assertEqual([i["id"] for i in self._items()], [old["id"]])
|
||||
self.assertEqual(calls, [])
|
||||
|
||||
async def test_resolver_change_leaves_one_default_item(self):
|
||||
self._roots()
|
||||
write_png(self.output_dir / "two.png", prompt=unlinkable_graph("cat", "dog"))
|
||||
await self._scan()
|
||||
(old,) = self._items()
|
||||
self.assertEqual(old["resolver_version"], 1)
|
||||
|
||||
with patch.object(review_capture, "RESOLVER_VERSION", 2):
|
||||
await self._scan("all")
|
||||
|
||||
(item,) = self._items()
|
||||
self.assertNotEqual(item["id"], old["id"])
|
||||
self.assertEqual(item["resolver_version"], 2)
|
||||
self.assertEqual(self._pending(), 1)
|
||||
|
||||
async def test_decided_old_item_is_kept_and_no_duplicate_is_made(self):
|
||||
self._require_routes()
|
||||
self._roots()
|
||||
self._write()
|
||||
await self._scan()
|
||||
item, detail, observed = await self._open()
|
||||
status, data = await self._decide(
|
||||
item["id"], self._body(detail, observed, "unresolved")
|
||||
)
|
||||
self.assertEqual(status, 200, data)
|
||||
before = (self._items(), self._decisions())
|
||||
|
||||
self._table(_retarget_unet_loader)
|
||||
await self._scan("all")
|
||||
self.assertEqual((self._items(), self._decisions()), before)
|
||||
# The decided bytes are not copied again on the next rescan.
|
||||
with _counted(self.api._captures, "capture") as copies:
|
||||
await self._scan("all")
|
||||
|
||||
self.assertEqual(copies, [])
|
||||
self.assertEqual((self._items(), self._decisions()), before)
|
||||
self.assertEqual(self._pending(), 0)
|
||||
|
||||
async def test_advanced_pending_item_is_kept_and_no_duplicate_is_made(self):
|
||||
self._roots()
|
||||
self._write()
|
||||
await self._scan()
|
||||
(old,) = self._items()
|
||||
conn = self.api.db.model.get_connection()
|
||||
with conn:
|
||||
conn.execute(
|
||||
"UPDATE media_review_items SET revision=1 WHERE id=?", (old["id"],)
|
||||
)
|
||||
before = self._items()
|
||||
|
||||
self._table(_retarget_unet_loader)
|
||||
await self._scan("all")
|
||||
|
||||
self.assertEqual(self._items(), before)
|
||||
|
||||
|
||||
class TestInspectionRuns(InspectionTestCase):
|
||||
"""The heavier profile runs once per new file state, never again."""
|
||||
|
||||
async def test_new_file_is_inspected_once_and_not_on_rescan(self):
|
||||
self._roots()
|
||||
self._write()
|
||||
with _counted(review_capture, "resolve_inspection") as first:
|
||||
await self._scan()
|
||||
self.assertEqual(len(self._items()), 1)
|
||||
self.assertEqual(len(first), 1)
|
||||
|
||||
with _counted(review_capture, "resolve_inspection") as again:
|
||||
await self._scan("all")
|
||||
|
||||
self.assertEqual(again, [])
|
||||
self.assertEqual(len(self._items()), 1)
|
||||
|
||||
async def test_untraced_file_is_not_inspected_again_until_it_changes(self):
|
||||
self._roots()
|
||||
graph = qwen_edit_graph(*PROMPTS)
|
||||
for node in graph.values():
|
||||
if node["class_type"] == "TextEncodeQwenImageEdit":
|
||||
node["class_type"] = "OpaqueTextEncode"
|
||||
self._write(graph=graph)
|
||||
with _counted(review_capture, "resolve_inspection") as first:
|
||||
job = await self._scan()
|
||||
self.assertEqual(len(first), 1)
|
||||
(row,) = job["outcomes"]["review"]
|
||||
self.assertIn("could not be traced", row["detail"])
|
||||
|
||||
with _counted(review_capture, "resolve_inspection") as again:
|
||||
await self._scan("all")
|
||||
self.assertEqual(again, [])
|
||||
|
||||
self._write(graph=graph, color=(0, 255, 0))
|
||||
with _counted(review_capture, "resolve_inspection") as changed:
|
||||
await self._scan("all")
|
||||
self.assertEqual(len(changed), 1)
|
||||
self.assertEqual(self._items(), [])
|
||||
|
||||
async def test_capture_with_other_records_is_inspected_again(self):
|
||||
"""The scan's result is reused only for identical records."""
|
||||
self._roots()
|
||||
self._write()
|
||||
real = review_capture.captured_review
|
||||
|
||||
def edited(store, capture, profile=review_capture.DEFAULT):
|
||||
metadata, resolution = real(store, capture, profile)
|
||||
graph = json.loads(json.dumps(metadata["prompt"]))
|
||||
graph["100"]["inputs"]["prompt"] = "a captured edit"
|
||||
return {**metadata, "prompt": graph}, resolution
|
||||
|
||||
with (
|
||||
patch.object(review_capture, "captured_review", side_effect=edited),
|
||||
_counted(review_capture, "resolve_inspection") as calls,
|
||||
):
|
||||
await self._scan()
|
||||
|
||||
self.assertEqual(len(calls), 2)
|
||||
(item,) = self._items()
|
||||
self.assertIn("a captured edit", _texts(self._evidence(item)))
|
||||
|
||||
|
||||
class TestClassProjection(unittest.TestCase):
|
||||
def test_projection_over_every_named_class_is_the_whole_table(self):
|
||||
doc = json.loads(inspection_semantics.SEMANTICS_PATH.read_text("utf-8"))
|
||||
names = set(inspection_semantics.load_semantics().all_classes)
|
||||
names |= set(doc["facts"]["configuration_outputs"])
|
||||
names |= set(doc["facts"]["model_outputs"])
|
||||
names |= set(doc["facts"]["model_passthroughs"])
|
||||
names |= set(doc["facts"].get("scalar_outputs", {}))
|
||||
names |= {entry["class"] for entry in doc["lineage"]}
|
||||
|
||||
self.assertEqual(inspection_semantics.class_projection(doc, names), doc)
|
||||
|
||||
def test_class_version_moves_only_for_the_classes_it_covers(self):
|
||||
doc = json.loads(inspection_semantics.SEMANTICS_PATH.read_text("utf-8"))
|
||||
classes = frozenset({"KSampler", "UNETLoader"})
|
||||
version = inspection_semantics.class_version(doc, classes)
|
||||
_unrelated_encoder(doc)
|
||||
self.assertEqual(inspection_semantics.class_version(doc, classes), version)
|
||||
_retarget_unet_loader(doc)
|
||||
self.assertNotEqual(inspection_semantics.class_version(doc, classes), version)
|
||||
|
||||
|
||||
class TestRecordingEdges(InspectionTestCase):
|
||||
async def test_store_refuses_a_new_version_beside_a_decided_item(self):
|
||||
self._roots()
|
||||
self._write()
|
||||
await self._scan()
|
||||
(old,) = self._items()
|
||||
conn = self.api.db.model.get_connection()
|
||||
with conn:
|
||||
conn.execute(
|
||||
"UPDATE media_review_items SET state='unresolved', revision=1 "
|
||||
"WHERE id=?",
|
||||
(old["id"],),
|
||||
)
|
||||
evidence = {**self._evidence(old), "resolver_version": 7}
|
||||
|
||||
inserted = self.api.db.record_review_item(
|
||||
locator=old["locator"],
|
||||
spelled_path=old["spelled_path"],
|
||||
sha256=old["sha256"],
|
||||
size=old["size"],
|
||||
mtime_ns=old["mtime_ns"],
|
||||
epoch=self.api.db.get_execution_epoch(),
|
||||
resolver_version=7,
|
||||
evidence=evidence,
|
||||
candidate_count=old["candidate_count"],
|
||||
limited=bool(old["limited"]),
|
||||
)
|
||||
|
||||
self.assertFalse(inserted)
|
||||
self.assertEqual([i["id"] for i in self._items()], [old["id"]])
|
||||
|
||||
async def test_already_linked_file_with_an_inspection_choice_is_unchanged(self):
|
||||
self._roots()
|
||||
path = self._write()
|
||||
prompt_id = self.api.db.save_prompt(
|
||||
text="latest", prompt_hash=generate_prompt_hash("latest")
|
||||
)
|
||||
handler = ImageGenerationHandler(self.api.db, PromptTracker(self.api.db))
|
||||
handler.link_image_to_prompt(
|
||||
str(path),
|
||||
{"id": prompt_id, "text": "latest", "fallback": True},
|
||||
{"prompt": qwen_edit_graph(*PROMPTS)},
|
||||
)
|
||||
|
||||
job = await self._scan("all")
|
||||
|
||||
self.assertEqual(self._items(), [])
|
||||
self.assertEqual(job["outcomes"]["review"], [])
|
||||
(row,) = job["outcomes"]["known"]
|
||||
self.assertIn("Already linked", row["detail"])
|
||||
self.assertEqual(self._linked_text(path), "latest")
|
||||
|
||||
async def test_untraced_memory_forgets_the_oldest_state(self):
|
||||
from py.api import review_recording
|
||||
|
||||
self._roots()
|
||||
graph = qwen_edit_graph(*PROMPTS)
|
||||
for node in graph.values():
|
||||
if node["class_type"] == "TextEncodeQwenImageEdit":
|
||||
node["class_type"] = "OpaqueTextEncode"
|
||||
self._write("a.png", graph=graph)
|
||||
self._write("b.png", graph=graph)
|
||||
with patch.object(review_recording, "MAX_UNTRACED_MEMORY", 1):
|
||||
await self._scan()
|
||||
with _counted(review_capture, "resolve_inspection") as calls:
|
||||
await self._scan("all")
|
||||
|
||||
# One remembered state: each file evicts the other, so both run.
|
||||
self.assertEqual(len(calls), 2)
|
||||
self.assertEqual(len(self.api._untraced_inspections), 1)
|
||||
|
||||
def test_graph_classes(self):
|
||||
self.assertEqual(review_capture.graph_classes({"parameters": "x"}), set())
|
||||
self.assertEqual(
|
||||
review_capture.graph_classes({"prompt": {"1": {"class_type": "A"}}}),
|
||||
{"A"},
|
||||
)
|
||||
many = {
|
||||
str(n): {"class_type": f"C{n}"}
|
||||
for n in range(review_capture.MAX_FINGERPRINT_CLASSES + 1)
|
||||
}
|
||||
self.assertIsNone(review_capture.graph_classes({"prompt": many}))
|
||||
self.assertEqual(
|
||||
review_capture.semantics_version({"prompt": many}),
|
||||
inspection_semantics.get_semantics().version,
|
||||
)
|
||||
|
||||
|
||||
def _move_qwen_encoder_slot(doc):
|
||||
"""An edit of the Qwen-edit encoder's own table entry."""
|
||||
for entry in doc["encoders"]:
|
||||
if entry["class"] == "TextEncodeQwenImageEdit":
|
||||
entry["slot"] = 1
|
||||
|
||||
|
||||
def _unused_sampler(doc):
|
||||
"""A new dict-keyed sampler class the Qwen-edit graph does not use."""
|
||||
doc["samplers"]["UnusedSampler"] = dict(doc["samplers"]["KSampler"])
|
||||
|
||||
|
||||
class TestFingerprintCoverage(InspectionTestCase):
|
||||
def test_editing_the_encoders_entry_moves_the_version(self):
|
||||
version = _version()
|
||||
self._table(_move_qwen_encoder_slot)
|
||||
self.assertIsNotNone(_version())
|
||||
self.assertNotEqual(_version(), version)
|
||||
|
||||
def test_an_unused_dict_keyed_class_leaves_the_version(self):
|
||||
version = _version()
|
||||
table = inspection_semantics.get_semantics().version
|
||||
self._table(_unused_sampler)
|
||||
self.assertNotEqual(inspection_semantics.get_semantics().version, table)
|
||||
self.assertEqual(_version(), version)
|
||||
|
||||
|
||||
class TestInspectionResolverVersion(InspectionTestCase):
|
||||
async def test_inspection_resolver_change_supersedes_the_item(self):
|
||||
"""An inspection-only resolver change moves the item version, so the
|
||||
rescan re-inspects and the new item is applicable (never a stuck
|
||||
409)."""
|
||||
self._require_routes()
|
||||
self._roots()
|
||||
path = self._write()
|
||||
await self._scan()
|
||||
(old,) = self._items()
|
||||
version = _version()
|
||||
real = review_capture.resolve_inspection
|
||||
|
||||
def newer(metadata):
|
||||
resolution = real(metadata)
|
||||
return {
|
||||
**resolution,
|
||||
"candidates": [
|
||||
{**c, "reason": "a newer inspection resolver"}
|
||||
for c in resolution["candidates"]
|
||||
],
|
||||
}
|
||||
|
||||
with (
|
||||
patch.object(review_capture, "INSPECTION_RESOLVER_VERSION", 2),
|
||||
patch.object(review_capture, "resolve_inspection", side_effect=newer),
|
||||
):
|
||||
self.assertNotEqual(_version(), version)
|
||||
await self._scan("all")
|
||||
(item,) = self._items()
|
||||
self.assertNotEqual(item["id"], old["id"])
|
||||
_item, (status, data) = await self._apply("add a red hat")
|
||||
|
||||
self.assertEqual(status, 200, data)
|
||||
self.assertEqual(self._linked_text(path), "add a red hat")
|
||||
|
||||
def test_profile_budgets_are_folded_into_the_version(self):
|
||||
from dataclasses import replace
|
||||
|
||||
version = _version()
|
||||
wider = replace(
|
||||
review_capture.INSPECTION_PROFILE,
|
||||
total_visits=review_capture.INSPECTION_PROFILE.total_visits + 1,
|
||||
)
|
||||
with patch.object(review_capture, "INSPECTION_PROFILE", wider):
|
||||
self.assertNotEqual(_version(), version)
|
||||
|
||||
|
||||
class TestScanContainment(InspectionTestCase):
|
||||
async def test_an_inspection_failure_is_counted_and_the_scan_continues(self):
|
||||
self._roots()
|
||||
self._write("a.png")
|
||||
self._write("b.png", graph=qwen_edit_graph("snow", "rain"))
|
||||
real = review_capture.resolve_inspection
|
||||
|
||||
def fails_for_a(metadata):
|
||||
if "make it night" in json.dumps(metadata):
|
||||
raise RuntimeError("inspection blew up")
|
||||
return real(metadata)
|
||||
|
||||
with (
|
||||
patch.object(review_capture, "resolve_inspection", side_effect=fails_for_a),
|
||||
self.assertLogs("prompt_manager.api", "WARNING") as logs,
|
||||
):
|
||||
job = await self._scan()
|
||||
|
||||
self.assertEqual(job["state"], "finished")
|
||||
(item,) = self._items()
|
||||
self.assertEqual(os.path.basename(item["locator"]), "b.png")
|
||||
self.assertEqual(job["scope"]["review_errors"], 1)
|
||||
rows = {row["path"]: row for row in job["outcomes"]["review"]}
|
||||
self.assertIn("could not be traced", rows["a.png"]["detail"])
|
||||
self.assertIn("more than one prompt could match", rows["b.png"]["detail"])
|
||||
self.assertTrue(any("a.png" in line for line in logs.output))
|
||||
|
||||
|
||||
class TestUntracedMemoryKeys(InspectionTestCase):
|
||||
async def test_untraced_memory_is_keyed_by_a_16_byte_digest(self):
|
||||
self._roots()
|
||||
graph = qwen_edit_graph(*PROMPTS)
|
||||
for node in graph.values():
|
||||
if node["class_type"] == "TextEncodeQwenImageEdit":
|
||||
node["class_type"] = "OpaqueTextEncode"
|
||||
self._write(graph=graph)
|
||||
|
||||
await self._scan()
|
||||
|
||||
(key,) = self.api._untraced_inspections
|
||||
self.assertIsInstance(key, bytes)
|
||||
self.assertEqual(len(key), 16)
|
||||
|
||||
|
||||
class TestMonitorNeverRecords(InspectionTestCase):
|
||||
async def test_monitor_never_records_an_unlinked_ambiguous_output(self):
|
||||
"""Only scans add items (maintainer ruling); the monitor only clears
|
||||
the items of outputs it links."""
|
||||
self._roots()
|
||||
path = self.output_dir / "two.png"
|
||||
write_png(path, prompt=unlinkable_graph("cat", "dog"))
|
||||
|
||||
with _counted(self.api._captures, "capture") as copies:
|
||||
self.assertEqual(self.api.record_review_file(str(path)), "skipped")
|
||||
|
||||
self.assertEqual(copies, [])
|
||||
self.assertEqual(self._items(), [])
|
||||
await self._scan()
|
||||
self.assertEqual(len(self._items()), 1)
|
||||
|
||||
|
||||
class TestStaleSettleContainment(InspectionTestCase):
|
||||
async def test_a_failed_stale_settle_is_logged_and_counted(self):
|
||||
import sqlite3
|
||||
|
||||
self._roots()
|
||||
graph = qwen_edit_graph(*PROMPTS)
|
||||
for node in graph.values():
|
||||
if node["class_type"] == "TextEncodeQwenImageEdit":
|
||||
node["class_type"] = "OpaqueTextEncode"
|
||||
self._write(graph=graph)
|
||||
with (
|
||||
patch.object(
|
||||
self.api.db,
|
||||
"settle_stale_review_items",
|
||||
side_effect=sqlite3.OperationalError("database is locked"),
|
||||
),
|
||||
self.assertLogs("prompt_manager.api", "WARNING") as logs,
|
||||
):
|
||||
job = await self._scan()
|
||||
|
||||
self.assertEqual(job["state"], "finished")
|
||||
self.assertEqual(job["scope"]["review_errors"], 1)
|
||||
self.assertTrue(any("Stale Needs review" in line for line in logs.output))
|
||||
@@ -355,9 +355,7 @@ class TestMonitorRule(OutcomeTestCase):
|
||||
self.assertEqual((self._items(), self._decisions()), before)
|
||||
self.assertEqual(self._pending(), 0)
|
||||
|
||||
async def test_record_review_file_only_clears_what_it_links(self):
|
||||
"""Only scans add items: the monitor's pass skips every unlinked file
|
||||
and settles the items of a linked one."""
|
||||
async def test_record_review_file_applies_the_scan_rule(self):
|
||||
self._roots()
|
||||
write_png(self.output_dir / "fallback.png", prompt=ambiguous_graph("a", "b"))
|
||||
write_png(self.output_dir / "two.png", prompt=unlinkable_graph("cat", "dog"))
|
||||
@@ -368,8 +366,8 @@ class TestMonitorRule(OutcomeTestCase):
|
||||
record = self.api.record_review_file
|
||||
self.assertEqual(record(str(self.output_dir / "fallback.png")), "skipped")
|
||||
self.assertEqual(record(str(linked)), "linked")
|
||||
self.assertEqual(record(str(self.output_dir / "two.png")), "skipped")
|
||||
self.assertEqual(self._items(), [])
|
||||
self.assertEqual(record(str(self.output_dir / "two.png")), "recorded")
|
||||
self.assertEqual(len(self._items()), 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -150,15 +150,7 @@ class TestSourceStatus(SourcesTestCase):
|
||||
)
|
||||
self.assertEqual(
|
||||
data["sources"]["effective"],
|
||||
[
|
||||
{
|
||||
"label": "a",
|
||||
"path": "output/a",
|
||||
"root_index": 0,
|
||||
"watch": "manual",
|
||||
"in_comfyui": True,
|
||||
}
|
||||
],
|
||||
[{"label": "a", "path": "output/a", "root_index": 0, "watch": "manual"}],
|
||||
)
|
||||
self.assertFalse(data["sources"]["auto_detected"])
|
||||
# The legacy keys are unchanged.
|
||||
@@ -191,21 +183,10 @@ 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,7 +277,6 @@ class TestOverviewRoute(SystemAPITestCase):
|
||||
"backup",
|
||||
"needs_review",
|
||||
"sources",
|
||||
"location",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -324,7 +324,7 @@ class TestThumbnailJob(ThumbnailJobCase):
|
||||
old = preview.read_bytes()
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": None}),
|
||||
patch("subprocess.Popen", side_effect=FileNotFoundError("no ffmpeg")),
|
||||
patch("subprocess.run", side_effect=FileNotFoundError("no ffmpeg")),
|
||||
):
|
||||
_, body = await self._start(mode="rebuild")
|
||||
job = await self._finished(body["job"]["id"])
|
||||
@@ -566,7 +566,7 @@ class TestThumbnailJob(ThumbnailJobCase):
|
||||
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": None}),
|
||||
patch("subprocess.Popen", side_effect=FileNotFoundError("no ffmpeg")),
|
||||
patch("subprocess.run", side_effect=FileNotFoundError("no ffmpeg")),
|
||||
patch.object(images_module, "_atomic_save", refuse),
|
||||
):
|
||||
_, body = await self._start()
|
||||
@@ -594,14 +594,12 @@ class TestThumbnailJob(ThumbnailJobCase):
|
||||
def fake_ffmpeg(cmd, **_kwargs):
|
||||
commands.append(cmd)
|
||||
Path(cmd[-1]).write_bytes(b"jpeg")
|
||||
done = type("Done", (), {"returncode": 0})()
|
||||
done.communicate = lambda timeout=None: ("", "")
|
||||
return done
|
||||
return type("Done", (), {"returncode": 0, "stderr": ""})()
|
||||
|
||||
for quality, qscale in (("high", "2"), ("medium", "4"), ("low", "8")):
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": None}),
|
||||
patch("subprocess.Popen", side_effect=fake_ffmpeg),
|
||||
patch("subprocess.run", side_effect=fake_ffmpeg),
|
||||
):
|
||||
_, body = await self._start(mode="rebuild", size=640, quality=quality)
|
||||
job = await self._finished(body["job"]["id"])
|
||||
@@ -620,13 +618,12 @@ class TestThumbnailJob(ThumbnailJobCase):
|
||||
real_save = images_module._atomic_save
|
||||
|
||||
def recording(img, dst, fmt=None, **kwargs):
|
||||
# The encoder settings (guard is the automatic previews' hold check).
|
||||
saved.append({k: v for k, v in kwargs.items() if k != "guard"})
|
||||
saved.append(kwargs)
|
||||
return real_save(img, dst, fmt, **kwargs)
|
||||
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": None}),
|
||||
patch("subprocess.Popen", side_effect=FileNotFoundError("no ffmpeg")),
|
||||
patch("subprocess.run", side_effect=FileNotFoundError("no ffmpeg")),
|
||||
patch.object(images_module, "_atomic_save", recording),
|
||||
):
|
||||
_, body = await self._start(quality="low")
|
||||
|
||||
@@ -1,817 +0,0 @@
|
||||
"""Each workflow and prompt graph stored once (database/workflow_documents.py).
|
||||
|
||||
``generated_images.workflow_data`` and ``prompt_metadata`` held the same large
|
||||
JSON text again for every image of a run. The text now lives once in
|
||||
``workflow_documents`` (keyed by the sha256 of its exact UTF-8 bytes, zlib
|
||||
compressed) and image rows reference it. Old rows keep their inline text until
|
||||
the migration rewrites them; every reader goes through one accessor, which
|
||||
reads both forms, and gives back the exact text that was stored.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import shutil
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
import zlib
|
||||
from unittest.mock import patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from database import workflow_documents as wd # noqa: E402
|
||||
from database.models import PromptModel # noqa: E402
|
||||
from database.operations import PromptDatabase # noqa: E402
|
||||
from utils.hashing import generate_prompt_hash # noqa: E402
|
||||
from utils.json_strict import dumps_strict # noqa: E402
|
||||
|
||||
try:
|
||||
from tests.open_handles import assert_closed
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from open_handles import assert_closed
|
||||
|
||||
NAN = float("nan")
|
||||
INF = float("inf")
|
||||
|
||||
|
||||
def sample_workflow(seed=0, nodes=40):
|
||||
"""A ComfyUI-shaped UI workflow: nodes, links, widgets and groups."""
|
||||
rng = random.Random(seed)
|
||||
return {
|
||||
"last_node_id": nodes,
|
||||
"last_link_id": nodes * 2,
|
||||
"nodes": [
|
||||
{
|
||||
"id": index,
|
||||
"type": rng.choice(
|
||||
["KSampler", "CLIPTextEncode", "VAEDecode", "LoraLoader"]
|
||||
),
|
||||
"pos": [rng.uniform(-2000, 2000), rng.uniform(-2000, 2000)],
|
||||
"size": {"0": 315, "1": 262},
|
||||
"flags": {},
|
||||
"order": index,
|
||||
"mode": 0,
|
||||
"inputs": [{"name": "model", "type": "MODEL", "link": index}],
|
||||
"outputs": [{"name": "LATENT", "type": "LATENT", "links": [index]}],
|
||||
"properties": {"Node name for S&R": "KSampler"},
|
||||
"widgets_values": [seed * 1000 + index, "fixed", 20, 7.5, "euler"],
|
||||
}
|
||||
for index in range(nodes)
|
||||
],
|
||||
"links": [[i, i, 0, i + 1, 0, "MODEL"] for i in range(nodes)],
|
||||
"groups": [{"title": "Sampler ✨ 采样", "bounding": [0, 0, 100, 100]}],
|
||||
"config": {},
|
||||
"extra": {"ds": {"scale": 0.5, "offset": [0, 0]}},
|
||||
"version": 0.4,
|
||||
}
|
||||
|
||||
|
||||
def sample_prompt_graph(seed=0):
|
||||
return {
|
||||
"3": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {"seed": seed, "steps": 20, "cfg": 7.0, "model": ["4", 0]},
|
||||
},
|
||||
"6": {"class_type": "CLIPTextEncode", "inputs": {"text": f"a cat {seed}"}},
|
||||
}
|
||||
|
||||
|
||||
def edge_texts():
|
||||
"""Texts whose bytes a parse-and-serialize round trip would change."""
|
||||
nested = "[" * 500 + "]" * 500
|
||||
return [
|
||||
"",
|
||||
"{}",
|
||||
"null",
|
||||
" ",
|
||||
"\n",
|
||||
'{"a": 1}',
|
||||
'{"a":1}',
|
||||
'{"a": 1, "a": 2}', # duplicate key: a parse would keep one
|
||||
'{"b": 1, "a": 2}', # key order
|
||||
'{"v": 1.0}',
|
||||
'{"v": 1e3}', # number spelling
|
||||
'{"v": 1.10}',
|
||||
'{"v": NaN, "w": Infinity, "x": -Infinity}', # ComfyUI's tokens
|
||||
'{"text": "NaN"}',
|
||||
'{"u": "\\u00e9\\ud83d\\ude00"}', # escaped
|
||||
'{"u": "é 😀 中文 עברית é"}', # raw unicode, astral, RTL, combining
|
||||
"trailing space ",
|
||||
'{"nul": "a\u0000b"}',
|
||||
"a\u0000b",
|
||||
"{}", # byte-order mark
|
||||
"
",
|
||||
"not json at all",
|
||||
nested,
|
||||
json.dumps(sample_workflow(1)),
|
||||
json.dumps(sample_workflow(1), indent=2),
|
||||
json.dumps(sample_workflow(1), separators=(",", ":")),
|
||||
json.dumps({"v": NAN, "w": [INF, -INF]}), # Python's lenient output
|
||||
dumps_strict({"v": NAN, "w": [INF, -INF]}),
|
||||
]
|
||||
|
||||
|
||||
def random_text(rng):
|
||||
"""Arbitrary text of valid code points (no surrogates), ASCII to astral."""
|
||||
length = rng.choice([0, 1, 2, 7, 64, 1000, 20000])
|
||||
ranges = [(0x00, 0x7F), (0x80, 0x7FF), (0x800, 0xD7FF), (0xE000, 0xFFFF)]
|
||||
ranges.append((0x10000, 0x10FFFF))
|
||||
chars = []
|
||||
for _ in range(length):
|
||||
low, high = rng.choice(ranges)
|
||||
chars.append(chr(rng.randint(low, high)))
|
||||
return "".join(chars)
|
||||
|
||||
|
||||
class DocumentsTestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmpdir = os.path.realpath(tempfile.mkdtemp())
|
||||
self.db_path = os.path.join(self.tmpdir, "prompts.db")
|
||||
self.db = PromptDatabase(self.db_path)
|
||||
self.addCleanup(self._cleanup)
|
||||
|
||||
def _cleanup(self):
|
||||
PromptModel.close_all_instances()
|
||||
PromptModel.reset_schema_cache()
|
||||
assert_closed(self, self.tmpdir)
|
||||
shutil.rmtree(self.tmpdir)
|
||||
|
||||
def conn(self):
|
||||
return self.db.model.get_connection()
|
||||
|
||||
def prompt(self, text="a cat"):
|
||||
return self.db.save_prompt(text, prompt_hash=generate_prompt_hash(text))
|
||||
|
||||
def old_form_image(self, prompt_id, name, workflow, prompt_graph):
|
||||
"""A row as every version before the dedupe wrote it: text inline."""
|
||||
conn = self.conn()
|
||||
with conn:
|
||||
path = os.path.join(self.tmpdir, name)
|
||||
cursor = conn.execute(
|
||||
"INSERT INTO generated_images (prompt_id, image_path, filename, "
|
||||
"file_path, workflow_data, prompt_metadata, parameters) "
|
||||
"VALUES (?, ?, ?, ?, ?, ?, ?)",
|
||||
(prompt_id, path, name, path, workflow, prompt_graph, "{}"),
|
||||
)
|
||||
return cursor.lastrowid
|
||||
|
||||
def raw(self, image_id):
|
||||
return dict(
|
||||
self.conn()
|
||||
.execute(
|
||||
"SELECT workflow_data, prompt_metadata, workflow_document_id, "
|
||||
"prompt_document_id FROM generated_images WHERE id=?",
|
||||
(image_id,),
|
||||
)
|
||||
.fetchone()
|
||||
)
|
||||
|
||||
def document_count(self):
|
||||
return (
|
||||
self.conn().execute("SELECT COUNT(*) FROM workflow_documents").fetchone()[0]
|
||||
)
|
||||
|
||||
def linked(self, prompt_id, name, seed):
|
||||
return self.db.link_image_to_prompt(
|
||||
prompt_id,
|
||||
os.path.join(self.tmpdir, name),
|
||||
{"workflow": sample_workflow(seed), "prompt": sample_prompt_graph(seed)},
|
||||
)
|
||||
|
||||
def store(self, text):
|
||||
conn = self.conn()
|
||||
with conn:
|
||||
return wd.store_text(conn, text)
|
||||
|
||||
|
||||
class TestRoundTrip(DocumentsTestCase):
|
||||
def assert_round_trip(self, text):
|
||||
document_id = self.store(text)
|
||||
self.assertIsInstance(document_id, int)
|
||||
back = wd.read_stored_text(self.conn(), None, document_id)
|
||||
self.assertIsInstance(back, str)
|
||||
self.assertEqual(back, text)
|
||||
self.assertEqual(back.encode("utf-8"), text.encode("utf-8"))
|
||||
return document_id
|
||||
|
||||
def test_edge_texts_come_back_byte_identical(self):
|
||||
for text in edge_texts():
|
||||
with self.subTest(text=text[:40]):
|
||||
self.assert_round_trip(text)
|
||||
|
||||
def test_texts_that_parse_alike_stay_distinct(self):
|
||||
spaced = self.assert_round_trip('{"a": 1}')
|
||||
tight = self.assert_round_trip('{"a":1}')
|
||||
self.assertNotEqual(spaced, tight)
|
||||
|
||||
def test_random_texts_come_back_byte_identical(self):
|
||||
rng = random.Random(20261024)
|
||||
for case in range(400):
|
||||
text = random_text(rng)
|
||||
with self.subTest(case=case):
|
||||
self.assert_round_trip(text)
|
||||
|
||||
def test_random_workflows_come_back_byte_identical(self):
|
||||
rng = random.Random(7)
|
||||
for case in range(150):
|
||||
value = sample_workflow(rng.randint(0, 10**9), nodes=rng.randint(0, 60))
|
||||
if rng.random() < 0.3:
|
||||
value["nodes"].append({"widgets_values": [NAN, INF, -INF]})
|
||||
for text in (json.dumps(value), dumps_strict(value)):
|
||||
with self.subTest(case=case):
|
||||
self.assert_round_trip(text)
|
||||
|
||||
def test_huge_workflow_round_trips(self):
|
||||
text = json.dumps(sample_workflow(3, nodes=4000))
|
||||
self.assertGreater(len(text), 1_200_000)
|
||||
self.assert_round_trip(text)
|
||||
|
||||
def test_identical_text_is_stored_once(self):
|
||||
text = json.dumps(sample_workflow(5))
|
||||
first = self.store(text)
|
||||
second = self.store(text)
|
||||
self.assertEqual(first, second)
|
||||
self.assertEqual(self.document_count(), 1)
|
||||
|
||||
def test_null_stays_null(self):
|
||||
self.assertIsNone(self.store(None))
|
||||
self.assertIsNone(wd.read_stored_text(self.conn(), None, None))
|
||||
self.assertEqual(self.document_count(), 0)
|
||||
|
||||
def test_document_is_compressed_utf8_keyed_by_sha256(self):
|
||||
text = "é" + json.dumps(sample_workflow(9))
|
||||
document_id = self.store(text)
|
||||
row = (
|
||||
self.conn()
|
||||
.execute(
|
||||
"SELECT sha256, byte_length, data FROM workflow_documents WHERE id=?",
|
||||
(document_id,),
|
||||
)
|
||||
.fetchone()
|
||||
)
|
||||
raw = text.encode("utf-8")
|
||||
import hashlib
|
||||
|
||||
self.assertEqual(bytes(row["sha256"]), hashlib.sha256(raw).digest())
|
||||
self.assertEqual(row["byte_length"], len(raw))
|
||||
self.assertEqual(zlib.decompress(row["data"]), raw)
|
||||
self.assertLess(len(row["data"]), len(raw) / 3)
|
||||
|
||||
def test_inline_text_wins_over_a_document(self):
|
||||
document_id = self.store("from the document")
|
||||
self.assertEqual(
|
||||
wd.read_stored_text(self.conn(), "inline", document_id), "inline"
|
||||
)
|
||||
self.assertEqual(wd.read_stored_text(self.conn(), "", document_id), "")
|
||||
|
||||
def test_a_tampered_document_reads_as_unreadable_never_as_other_text(self):
|
||||
document_id = self.store('{"a": 1}')
|
||||
conn = self.conn()
|
||||
for data in (zlib.compress(b'{"a": 2}'), b"not zlib"):
|
||||
with conn:
|
||||
conn.execute(
|
||||
"UPDATE workflow_documents SET data=? WHERE id=?",
|
||||
(data, document_id),
|
||||
)
|
||||
with self.subTest(data=data), self.assertLogs(wd.logger, level="WARNING"):
|
||||
self.assertIsNone(wd.read_stored_text(conn, None, document_id))
|
||||
|
||||
def test_a_missing_document_reads_as_unreadable(self):
|
||||
with self.assertLogs(wd.logger, level="WARNING"):
|
||||
self.assertIsNone(wd.read_stored_text(self.conn(), None, 424242))
|
||||
|
||||
def test_reader_memo_returns_the_same_text(self):
|
||||
text = json.dumps(sample_workflow(2))
|
||||
document_id = self.store(text)
|
||||
reader = wd.DocumentReader(self.conn())
|
||||
self.assertEqual(reader.text(None, document_id), text)
|
||||
self.assertEqual(reader.text(None, document_id), text)
|
||||
self.assertEqual(reader.text("inline", document_id), "inline")
|
||||
|
||||
|
||||
class TestSchema(DocumentsTestCase):
|
||||
def test_fresh_database_has_the_table_columns_and_indexes(self):
|
||||
conn = self.conn()
|
||||
self.assertIsNone(wd.workflow_documents_schema_error(conn, allow_missing=False))
|
||||
columns = {
|
||||
row[1] for row in conn.execute("PRAGMA table_info(generated_images)")
|
||||
}
|
||||
self.assertTrue(
|
||||
{"workflow_document_id", "prompt_document_id"}.issubset(columns)
|
||||
)
|
||||
indexes = {
|
||||
row[1] for row in conn.execute("PRAGMA index_list(generated_images)")
|
||||
}
|
||||
self.assertIn("idx_image_workflow_document", indexes)
|
||||
self.assertIn("idx_image_prompt_document", indexes)
|
||||
|
||||
def test_no_trigger_is_installed(self):
|
||||
# Every restore validator refuses triggers; GC is done in code.
|
||||
row = (
|
||||
self.conn()
|
||||
.execute("SELECT name FROM sqlite_master WHERE type='trigger'")
|
||||
.fetchone()
|
||||
)
|
||||
self.assertIsNone(row)
|
||||
|
||||
def test_current_database_verifies_and_restores(self):
|
||||
image = self.db.link_image_to_prompt(
|
||||
self.prompt(), os.path.join(self.tmpdir, "a.png"), {"workflow": {"n": 1}}
|
||||
)
|
||||
self.assertTrue(image)
|
||||
ok, reason = self.db.model.verify_database_file(
|
||||
self.db_path, require_library_state=True
|
||||
)
|
||||
self.assertTrue(ok, reason)
|
||||
|
||||
def test_an_old_database_opens_migrates_and_keeps_its_inline_rows(self):
|
||||
PromptModel.close_all_instances()
|
||||
PromptModel.reset_schema_cache()
|
||||
old_path = os.path.join(self.tmpdir, "old.db")
|
||||
conn = sqlite3.connect(old_path)
|
||||
try:
|
||||
conn.executescript(
|
||||
"CREATE TABLE prompts (id INTEGER PRIMARY KEY AUTOINCREMENT, "
|
||||
"text TEXT NOT NULL, created_at TIMESTAMP, updated_at TIMESTAMP, "
|
||||
"category TEXT, tags TEXT, rating INTEGER, notes TEXT, "
|
||||
"hash TEXT UNIQUE);"
|
||||
"CREATE TABLE generated_images (id INTEGER PRIMARY KEY "
|
||||
"AUTOINCREMENT, prompt_id INTEGER NOT NULL, image_path TEXT NOT "
|
||||
"NULL, filename TEXT NOT NULL, generation_time TIMESTAMP, "
|
||||
"file_size INTEGER, width INTEGER, height INTEGER, format TEXT, "
|
||||
"workflow_data TEXT, prompt_metadata TEXT, parameters TEXT, "
|
||||
"file_path TEXT, FOREIGN KEY (prompt_id) REFERENCES prompts(id) "
|
||||
"ON DELETE CASCADE, UNIQUE(prompt_id, file_path));"
|
||||
"INSERT INTO prompts (id, text, hash) VALUES (1, 'a cat', 'h');"
|
||||
"INSERT INTO generated_images (prompt_id, image_path, filename, "
|
||||
"file_path, workflow_data, prompt_metadata) VALUES (1, '/x/a.png', "
|
||||
"'a.png', '/x/a.png', '{\"n\": NaN}', '{\"3\": {}}');"
|
||||
)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
ok, reason = PromptModel.verify_database_file(old_path)
|
||||
self.assertTrue(ok, reason)
|
||||
old = PromptDatabase(old_path)
|
||||
[image] = old.get_prompt_images(1)
|
||||
self.assertEqual(image["workflow_data"], {"n": None})
|
||||
self.assertEqual(image["prompt_metadata"], {"3": {}})
|
||||
row = (
|
||||
old.model.get_connection()
|
||||
.execute("SELECT workflow_data, workflow_document_id FROM generated_images")
|
||||
.fetchone()
|
||||
)
|
||||
self.assertEqual(tuple(row), ('{"n": NaN}', None))
|
||||
old.close_all()
|
||||
|
||||
def test_verify_refuses_a_malformed_documents_table(self):
|
||||
bad = os.path.join(self.tmpdir, "bad.db")
|
||||
PromptModel.close_all_instances()
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
with sqlite3.connect(bad) as target:
|
||||
conn.backup(target)
|
||||
target.close()
|
||||
finally:
|
||||
conn.close()
|
||||
cases = {
|
||||
"table": "DROP TABLE workflow_documents; CREATE TABLE "
|
||||
"workflow_documents (id INTEGER PRIMARY KEY, sha256 TEXT, data TEXT)",
|
||||
"values": "INSERT INTO workflow_documents (sha256, byte_length, data) "
|
||||
"VALUES (zeroblob(32), -1, x'00')",
|
||||
"reference": "PRAGMA foreign_keys=OFF; INSERT INTO prompts (text, hash) "
|
||||
"VALUES ('p', 'h2'); INSERT INTO generated_images (prompt_id, image_path,"
|
||||
" filename, workflow_document_id) VALUES (last_insert_rowid(), '/a', 'a',"
|
||||
" 99)",
|
||||
"index": "DROP INDEX idx_image_workflow_document",
|
||||
}
|
||||
for name, script in cases.items():
|
||||
copy = os.path.join(self.tmpdir, f"bad-{name}.db")
|
||||
with open(bad, "rb") as source, open(copy, "wb") as target:
|
||||
target.write(source.read())
|
||||
conn = sqlite3.connect(copy)
|
||||
try:
|
||||
conn.executescript(script)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
with self.subTest(name=name):
|
||||
ok, reason = PromptModel.verify_database_file(copy)
|
||||
self.assertFalse(ok)
|
||||
self.assertIn("orkflow", reason)
|
||||
|
||||
|
||||
class TestSchemaErrors(unittest.TestCase):
|
||||
"""Each way a restored file's document schema can be wrong is named."""
|
||||
|
||||
IMAGES = (
|
||||
"CREATE TABLE generated_images (id INTEGER PRIMARY KEY, "
|
||||
"workflow_data TEXT, prompt_metadata TEXT{extra})"
|
||||
)
|
||||
TABLE = (
|
||||
"CREATE TABLE workflow_documents (id INTEGER PRIMARY KEY, sha256 BLOB NOT "
|
||||
"NULL UNIQUE, byte_length INTEGER NOT NULL, data BLOB NOT NULL{extra})"
|
||||
)
|
||||
REFS = (
|
||||
", workflow_document_id {kind}, prompt_document_id "
|
||||
"INTEGER REFERENCES workflow_documents(id)"
|
||||
)
|
||||
GOOD_REF = "INTEGER REFERENCES workflow_documents(id)"
|
||||
|
||||
def schema(self, images_extra="", table_extra="", indexes=True, table=True):
|
||||
conn = sqlite3.connect(":memory:")
|
||||
self.addCleanup(conn.close)
|
||||
conn.execute(self.IMAGES.format(extra=images_extra))
|
||||
if table:
|
||||
conn.execute(self.TABLE.format(extra=table_extra))
|
||||
if indexes and images_extra:
|
||||
for column, index in wd.INDEXES.items():
|
||||
conn.execute(f"CREATE INDEX {index} ON generated_images({column})")
|
||||
return conn
|
||||
|
||||
def test_each_error_is_named(self):
|
||||
refs = self.REFS.format(kind=self.GOOD_REF)
|
||||
cases = {
|
||||
"Missing workflow document schema": self.schema(table=False),
|
||||
"Incomplete workflow document schema": self.schema(),
|
||||
"Workflow documents have an invalid table definition": self.schema(
|
||||
refs, table_extra=", extra TEXT"
|
||||
),
|
||||
"Workflow document references have an invalid column": self.schema(
|
||||
self.REFS.format(kind="TEXT REFERENCES workflow_documents(id)")
|
||||
),
|
||||
"Workflow document references have invalid foreign keys": self.schema(
|
||||
self.REFS.format(kind="INTEGER")
|
||||
),
|
||||
"Workflow document references are missing their index": self.schema(
|
||||
refs, indexes=False
|
||||
),
|
||||
}
|
||||
for message, conn in cases.items():
|
||||
with self.subTest(message=message):
|
||||
self.assertEqual(
|
||||
wd.workflow_documents_schema_error(conn, allow_missing=False),
|
||||
message,
|
||||
)
|
||||
self.assertIsNone(wd.workflow_documents_schema_error(self.schema(table=False)))
|
||||
self.assertIsNone(
|
||||
wd.workflow_documents_schema_error(self.schema(refs), allow_missing=False)
|
||||
)
|
||||
|
||||
def test_an_index_on_another_table_is_not_the_index(self):
|
||||
conn = self.schema(self.REFS.format(kind=self.GOOD_REF), indexes=False)
|
||||
conn.execute("CREATE TABLE other (workflow_document_id INTEGER)")
|
||||
conn.execute(
|
||||
"CREATE INDEX idx_image_workflow_document ON other(workflow_document_id)"
|
||||
)
|
||||
self.assertEqual(
|
||||
wd.workflow_documents_schema_error(conn),
|
||||
"Workflow document references are missing their index",
|
||||
)
|
||||
|
||||
def test_migrate_refuses_a_bad_schema_and_leaves_it_alone(self):
|
||||
conn = self.schema(table_extra="") # the table without the columns
|
||||
with self.assertRaisesRegex(sqlite3.DatabaseError, "Incomplete"):
|
||||
wd.migrate_workflow_documents(conn)
|
||||
columns = {
|
||||
row[1] for row in conn.execute("PRAGMA table_info(generated_images)")
|
||||
}
|
||||
self.assertNotIn("workflow_document_id", columns)
|
||||
|
||||
def test_migrate_validates_what_it_installed(self):
|
||||
conn = self.schema(table=False)
|
||||
with patch.object(
|
||||
wd,
|
||||
"workflow_documents_schema_error",
|
||||
side_effect=[None, "Workflow documents broke"],
|
||||
):
|
||||
with self.assertRaisesRegex(sqlite3.DatabaseError, "broke"):
|
||||
wd.migrate_workflow_documents(conn)
|
||||
tables = {row[0] for row in conn.execute("SELECT name FROM sqlite_master")}
|
||||
self.assertNotIn("workflow_documents", tables)
|
||||
|
||||
|
||||
class TestEdgeErrors(DocumentsTestCase):
|
||||
def test_store_refuses_values_that_cannot_be_documents(self):
|
||||
conn = self.conn()
|
||||
with self.assertRaises(TypeError):
|
||||
wd.store_text(conn, b"bytes")
|
||||
with patch.object(wd, "MAX_DOCUMENT_BYTES", 8):
|
||||
with self.assertRaises(ValueError):
|
||||
wd.store_text(conn, "123456789")
|
||||
self.assertFalse(wd.storable("12345678é"))
|
||||
self.assertTrue(wd.storable("1234567"))
|
||||
self.assertEqual(
|
||||
wd.stored_values(conn, {"workflow_data": "123456789"}),
|
||||
{"workflow_data": "123456789", "workflow_document_id": None},
|
||||
)
|
||||
conn.rollback()
|
||||
|
||||
def test_a_hash_match_with_another_length_is_a_mismatch(self):
|
||||
document_id = self.store("abc")
|
||||
conn = self.conn()
|
||||
with conn:
|
||||
conn.execute(
|
||||
"UPDATE workflow_documents SET byte_length=5 WHERE id=?",
|
||||
(document_id,),
|
||||
)
|
||||
with conn, self.assertLogs(wd.logger, "WARNING"): # repaired from the text
|
||||
self.assertEqual(wd.store_text(conn, "abc"), document_id)
|
||||
self.assertEqual(wd.read_stored_text(conn, None, document_id), "abc")
|
||||
for length in (-1, 2, 4):
|
||||
with conn:
|
||||
conn.execute(
|
||||
"UPDATE workflow_documents SET byte_length=? WHERE id=?",
|
||||
(length, document_id),
|
||||
)
|
||||
with self.subTest(length=length), self.assertLogs(wd.logger, "WARNING"):
|
||||
self.assertIsNone(wd.read_stored_text(conn, None, document_id))
|
||||
|
||||
def test_the_reader_memo_is_bounded(self):
|
||||
ids = [self.store(f"text {index}") for index in range(40)]
|
||||
reader = wd.DocumentReader(self.conn())
|
||||
for index, document_id in enumerate(ids):
|
||||
self.assertEqual(reader.text(None, document_id), f"text {index}")
|
||||
self.assertEqual(len(reader._memo), wd.DocumentReader.MEMO)
|
||||
|
||||
def test_a_row_without_text_columns_is_left_alone(self):
|
||||
reader = wd.DocumentReader(self.conn())
|
||||
self.assertEqual(wd.resolve_image_row(reader, {"id": 1}), {"id": 1})
|
||||
self.assertEqual(
|
||||
wd.resolve_image_row(reader, {"workflow_document_id": None}),
|
||||
{"workflow_data": None},
|
||||
)
|
||||
|
||||
def test_documents_of_rows_takes_only_known_columns(self):
|
||||
with self.assertRaises(ValueError):
|
||||
wd.documents_of_rows(self.conn(), "image_path", ["x"])
|
||||
|
||||
|
||||
class TestWritersStoreNewForm(DocumentsTestCase):
|
||||
def test_link_image_stores_documents_not_inline_text(self):
|
||||
workflow = sample_workflow(1)
|
||||
graph = sample_prompt_graph(1)
|
||||
prompt_id = self.prompt()
|
||||
first = self.db.link_image_to_prompt(
|
||||
prompt_id,
|
||||
os.path.join(self.tmpdir, "a.png"),
|
||||
{"workflow": workflow, "prompt": graph},
|
||||
)
|
||||
second = self.db.link_image_to_prompt(
|
||||
prompt_id,
|
||||
os.path.join(self.tmpdir, "b.png"),
|
||||
{"workflow": workflow, "prompt": graph},
|
||||
)
|
||||
a, b = self.raw(first), self.raw(second)
|
||||
self.assertIsNone(a["workflow_data"])
|
||||
self.assertIsNone(a["prompt_metadata"])
|
||||
self.assertIsNotNone(a["workflow_document_id"])
|
||||
self.assertEqual(a["workflow_document_id"], b["workflow_document_id"])
|
||||
self.assertEqual(a["prompt_document_id"], b["prompt_document_id"])
|
||||
self.assertEqual(self.document_count(), 2)
|
||||
conn = self.conn()
|
||||
self.assertEqual(
|
||||
wd.read_stored_text(conn, None, a["workflow_document_id"]),
|
||||
dumps_strict(workflow),
|
||||
)
|
||||
|
||||
def test_readers_answer_the_same_for_both_forms(self):
|
||||
workflow, graph = sample_workflow(4), sample_prompt_graph(4)
|
||||
workflow["nodes"][0]["widgets_values"] = [NAN, INF]
|
||||
prompt_id = self.prompt()
|
||||
new_id = self.db.link_image_to_prompt(
|
||||
prompt_id,
|
||||
os.path.join(self.tmpdir, "new.png"),
|
||||
{"workflow": workflow, "prompt": graph},
|
||||
)
|
||||
old_id = self.old_form_image(
|
||||
prompt_id, "old.png", dumps_strict(workflow), dumps_strict(graph)
|
||||
)
|
||||
new, old = self.db.get_image_by_id(new_id), self.db.get_image_by_id(old_id)
|
||||
for key in ("workflow_data", "prompt_metadata", "parameters"):
|
||||
self.assertEqual(new[key], old[key])
|
||||
self.assertEqual(new["workflow_data"]["nodes"][0]["widgets_values"], [None] * 2)
|
||||
self.assertNotIn("workflow_document_id", new)
|
||||
self.assertNotIn("prompt_document_id", old)
|
||||
listed = {
|
||||
image["id"]: image
|
||||
for reader in (
|
||||
lambda: self.db.get_prompt_images(prompt_id),
|
||||
lambda: self.db.get_recent_images(10),
|
||||
lambda: self.db.get_all_images(),
|
||||
lambda: self.db.search_images_by_prompt("cat"),
|
||||
)
|
||||
for image in reader()
|
||||
}
|
||||
self.assertEqual(listed[new_id]["workflow_data"], new["workflow_data"])
|
||||
self.assertEqual(listed[old_id]["prompt_metadata"], graph)
|
||||
info_new = self.db.get_image_prompt_info(os.path.join(self.tmpdir, "new.png"))
|
||||
info_old = self.db.get_image_prompt_info(os.path.join(self.tmpdir, "old.png"))
|
||||
self.assertEqual(info_new["workflow_data"], info_old["workflow_data"])
|
||||
self.assertEqual(info_new["prompt_metadata"], graph)
|
||||
|
||||
def test_diagnostics_recent_images_resolve_documents(self):
|
||||
prompt_id = self.prompt()
|
||||
self.db.link_image_to_prompt(
|
||||
prompt_id,
|
||||
os.path.join(self.tmpdir, "d.png"),
|
||||
{"workflow": {"n": 1}, "prompt": {"3": {}}},
|
||||
)
|
||||
from utils import diagnostics
|
||||
|
||||
tool = diagnostics.GalleryDiagnostics(self.db_path)
|
||||
result = tool.check_images_table()
|
||||
tool_images = result["recent_images"]
|
||||
self.assertEqual(tool_images[0]["workflow_data"], '{"n": 1}')
|
||||
self.assertNotIn("workflow_document_id", tool_images[0])
|
||||
|
||||
|
||||
class TestGarbage(DocumentsTestCase):
|
||||
"""A document no row references is removed with its last reference."""
|
||||
|
||||
def test_delete_image_releases_only_unshared_documents(self):
|
||||
prompt_id = self.prompt()
|
||||
shared_a = self.linked(prompt_id, "a.png", 1)
|
||||
self.linked(prompt_id, "b.png", 1)
|
||||
alone = self.linked(prompt_id, "c.png", 2)
|
||||
self.assertEqual(self.document_count(), 4)
|
||||
self.db.delete_image(shared_a)
|
||||
self.assertEqual(self.document_count(), 4)
|
||||
self.db.delete_image(alone)
|
||||
self.assertEqual(self.document_count(), 2)
|
||||
|
||||
def test_deleting_prompts_releases_their_images_documents(self):
|
||||
first, second = self.prompt("one"), self.prompt("two")
|
||||
self.linked(first, "a.png", 1)
|
||||
self.linked(second, "b.png", 2)
|
||||
self.db.delete_prompt(first)
|
||||
self.assertEqual(self.document_count(), 2)
|
||||
self.db.update_prompt_metadata(second, category="gone")
|
||||
self.assertEqual(self.db.delete_prompts_by_category("gone"), 1)
|
||||
self.assertEqual(self.document_count(), 0)
|
||||
|
||||
def test_cleanup_missing_images_releases_documents(self):
|
||||
prompt_id = self.prompt()
|
||||
self.linked(prompt_id, "gone.png", 1)
|
||||
self.db.cleanup_missing_images()
|
||||
self.assertEqual(self.document_count(), 0)
|
||||
self.linked(prompt_id, "gone-too.png", 2)
|
||||
self.assertEqual(self.db.cleanup_missing_images(roots=[self.tmpdir]), 1)
|
||||
self.assertEqual(self.document_count(), 0)
|
||||
|
||||
def test_sweep_removes_what_a_path_without_release_left(self):
|
||||
prompt_id = self.prompt()
|
||||
image = self.linked(prompt_id, "a.png", 1)
|
||||
conn = self.conn()
|
||||
with conn:
|
||||
conn.execute("DELETE FROM generated_images WHERE id=?", (image,))
|
||||
self.assertEqual(self.document_count(), 2)
|
||||
with conn:
|
||||
self.assertEqual(wd.sweep_documents(conn), 2)
|
||||
self.assertEqual(self.document_count(), 0)
|
||||
|
||||
def test_startup_sweeps_unreferenced_documents(self):
|
||||
prompt_id = self.prompt()
|
||||
image = self.linked(prompt_id, "a.png", 1)
|
||||
conn = self.conn()
|
||||
with conn:
|
||||
conn.execute("DELETE FROM generated_images WHERE id=?", (image,))
|
||||
PromptModel.close_all_instances()
|
||||
PromptModel.reset_schema_cache()
|
||||
self.db = PromptDatabase(self.db_path)
|
||||
self.assertEqual(self.document_count(), 0)
|
||||
|
||||
def test_release_keeps_a_document_the_other_column_references(self):
|
||||
prompt_id = self.prompt()
|
||||
text = '{"same": true}'
|
||||
first = self.db.link_image_to_prompt(
|
||||
prompt_id, os.path.join(self.tmpdir, "a.png"), {"workflow": {"same": True}}
|
||||
)
|
||||
second = self.db.link_image_to_prompt(
|
||||
prompt_id, os.path.join(self.tmpdir, "b.png"), {"prompt": {"same": True}}
|
||||
)
|
||||
self.assertEqual(
|
||||
self.raw(first)["workflow_document_id"],
|
||||
self.raw(second)["prompt_document_id"],
|
||||
)
|
||||
self.db.delete_image(first)
|
||||
self.assertEqual(
|
||||
wd.read_stored_text(
|
||||
self.conn(), None, self.raw(second)["prompt_document_id"]
|
||||
),
|
||||
text,
|
||||
)
|
||||
|
||||
|
||||
class TestMoreGarbagePaths(DocumentsTestCase):
|
||||
def test_bulk_delete_prompts_releases_documents(self):
|
||||
first, second = self.prompt("one"), self.prompt("two")
|
||||
self.linked(first, "a.png", 1)
|
||||
self.linked(second, "b.png", 2)
|
||||
self.assertEqual(self.db.bulk_delete_prompts([first, second]), 2)
|
||||
self.assertEqual(self.document_count(), 0)
|
||||
|
||||
def test_duplicate_prompt_merge_releases_images_it_could_not_move(self):
|
||||
primary = self.prompt("A cat")
|
||||
copy = self.db.save_prompt("a cat ", prompt_hash="other-hash")
|
||||
self.linked(primary, "same.png", 1)
|
||||
self.linked(copy, "same.png", 2) # collides with the primary's link
|
||||
self.assertEqual(self.document_count(), 4)
|
||||
self.db.cleanup_duplicates()
|
||||
self.assertEqual(self.document_count(), 2)
|
||||
|
||||
def test_maintenance_delete_ids_releases_documents(self):
|
||||
from database.maintenance import _delete_ids
|
||||
|
||||
prompt_id = self.prompt()
|
||||
image = self.linked(prompt_id, "a.png", 1)
|
||||
other = self.prompt("other")
|
||||
self.linked(other, "b.png", 2)
|
||||
conn = self.conn()
|
||||
with conn:
|
||||
_delete_ids(conn, "generated_images", [image])
|
||||
self.assertEqual(self.document_count(), 2)
|
||||
with conn:
|
||||
_delete_ids(conn, "prompts", [other])
|
||||
self.assertEqual(self.document_count(), 0)
|
||||
|
||||
def test_failed_insert_of_a_known_link_leaves_no_new_document(self):
|
||||
prompt_id = self.prompt()
|
||||
path = os.path.join(self.tmpdir, "a.png")
|
||||
self.db.link_image_to_prompt(prompt_id, path, {"workflow": {"v": 1}})
|
||||
before = self.document_count()
|
||||
self.assertEqual(
|
||||
self.db.link_image_to_prompt(prompt_id, path, {"workflow": {"v": 2}}), 0
|
||||
)
|
||||
self.assertEqual(self.document_count(), before)
|
||||
|
||||
|
||||
class TestPairIngestionDocuments(DocumentsTestCase):
|
||||
def resolution(self):
|
||||
return {
|
||||
"state": "resolved",
|
||||
"positive": {"text": "cat", "provenance": {"source": "api"}},
|
||||
"negative": {"text": "blur", "provenance": {"source": "api"}},
|
||||
}
|
||||
|
||||
def test_ingestion_stores_new_form_and_replaces_documents_it_updates(self):
|
||||
from database.pair_ingestion import ingest_resolved_pair
|
||||
|
||||
path = os.path.join(self.tmpdir, "scan.png")
|
||||
first = {"workflow": sample_workflow(1), "prompt": sample_prompt_graph(1)}
|
||||
added = ingest_resolved_pair(self.db, path, first, self.resolution())
|
||||
self.assertEqual(added["status"], "added")
|
||||
row = self.raw(added["image_id"])
|
||||
self.assertIsNone(row["workflow_data"])
|
||||
self.assertIsNotNone(row["workflow_document_id"])
|
||||
again = ingest_resolved_pair(self.db, path, first, self.resolution())
|
||||
self.assertEqual(again["status"], "unchanged")
|
||||
second = {"workflow": sample_workflow(2), "prompt": sample_prompt_graph(1)}
|
||||
updated = ingest_resolved_pair(self.db, path, second, self.resolution())
|
||||
self.assertEqual(updated["status"], "enriched")
|
||||
self.assertEqual(self.document_count(), 2) # the first workflow went
|
||||
image = self.db.get_image_by_id(added["image_id"])
|
||||
self.assertEqual(image["workflow_data"], sample_workflow(2))
|
||||
|
||||
def test_an_old_form_row_with_the_same_text_is_unchanged(self):
|
||||
from database.pair_ingestion import ingest_resolved_pair
|
||||
|
||||
path = os.path.join(self.tmpdir, "old.png")
|
||||
metadata = {"workflow": sample_workflow(3), "prompt": sample_prompt_graph(3)}
|
||||
added = ingest_resolved_pair(self.db, path, metadata, self.resolution())
|
||||
conn = self.conn()
|
||||
with conn: # rewrite it as an older version stored it
|
||||
conn.execute(
|
||||
"UPDATE generated_images SET workflow_data=?, prompt_metadata=?, "
|
||||
"workflow_document_id=NULL, prompt_document_id=NULL WHERE id=?",
|
||||
(
|
||||
dumps_strict(metadata["workflow"], sort_keys=True),
|
||||
dumps_strict(metadata["prompt"], sort_keys=True),
|
||||
added["image_id"],
|
||||
),
|
||||
)
|
||||
wd.sweep_documents(conn)
|
||||
again = ingest_resolved_pair(self.db, path, metadata, self.resolution())
|
||||
self.assertEqual(again["status"], "unchanged")
|
||||
self.assertEqual(self.document_count(), 0)
|
||||
|
||||
|
||||
class TestQueryPlans(DocumentsTestCase):
|
||||
"""Every statement the store and the migration run is an index search."""
|
||||
|
||||
def plan(self, sql, params=()):
|
||||
rows = self.conn().execute("EXPLAIN QUERY PLAN " + sql, params).fetchall()
|
||||
return " | ".join(row[3] for row in rows)
|
||||
|
||||
def test_statements_use_indexes(self):
|
||||
for name, (sql, params) in wd.QUERY_PLAN_STATEMENTS.items():
|
||||
with self.subTest(name=name):
|
||||
plan = self.plan(sql, params)
|
||||
self.assertNotRegex(plan, r"\bSCAN (generated_images|gi)\b(?! USING)")
|
||||
self.assertTrue("USING" in plan or "SEARCH" in plan, f"{name}: {plan}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,469 +0,0 @@
|
||||
"""The workflow storage migration (database/workflow_storage.py).
|
||||
|
||||
Rewrites old-form rows (text inline) to the new form (a document id) in
|
||||
batches. Each batch is one transaction, so a crash leaves every row in one
|
||||
form or the other and readable either way; a rerun continues where the last
|
||||
committed batch stopped. Rows are read one at a time, the stored text comes
|
||||
back byte for byte, and no image or video file is ever opened.
|
||||
"""
|
||||
|
||||
import builtins
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import sqlite3
|
||||
import tracemalloc
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
from PIL import Image, PngImagePlugin
|
||||
|
||||
from database import workflow_documents as wd
|
||||
from database import workflow_storage as ws
|
||||
from database.models import PromptModel
|
||||
|
||||
try:
|
||||
from tests.test_workflow_documents import (
|
||||
DocumentsTestCase,
|
||||
edge_texts,
|
||||
random_text,
|
||||
sample_prompt_graph,
|
||||
sample_workflow,
|
||||
)
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from test_workflow_documents import (
|
||||
DocumentsTestCase,
|
||||
edge_texts,
|
||||
random_text,
|
||||
sample_prompt_graph,
|
||||
sample_workflow,
|
||||
)
|
||||
|
||||
|
||||
def _digest(path):
|
||||
if not os.path.exists(path):
|
||||
return None
|
||||
with open(path, "rb") as handle:
|
||||
return hashlib.sha256(handle.read()).hexdigest()
|
||||
|
||||
|
||||
class MigrationTestCase(DocumentsTestCase):
|
||||
def library(self, count=40, seed=11):
|
||||
"""Old-form rows: duplicated runs, edge texts and random text."""
|
||||
rng = random.Random(seed)
|
||||
prompt_id = self.prompt()
|
||||
texts = edge_texts()
|
||||
expected = {}
|
||||
for index in range(count):
|
||||
run = index // 4 # four images per run share their texts
|
||||
if index % 10 == 9:
|
||||
workflow = texts[index % len(texts)]
|
||||
graph = random_text(rng)
|
||||
else:
|
||||
workflow = json.dumps(sample_workflow(run, nodes=20))
|
||||
graph = json.dumps(sample_prompt_graph(run))
|
||||
if index % 13 == 5:
|
||||
graph = None
|
||||
image_id = self.old_form_image(prompt_id, f"{index}.png", workflow, graph)
|
||||
expected[image_id] = (workflow, graph)
|
||||
return prompt_id, expected
|
||||
|
||||
def texts(self):
|
||||
conn = self.conn()
|
||||
reader = wd.DocumentReader(conn)
|
||||
rows = conn.execute(
|
||||
"SELECT id, workflow_data, prompt_metadata, workflow_document_id, "
|
||||
"prompt_document_id FROM generated_images ORDER BY id"
|
||||
)
|
||||
result = {}
|
||||
for row in rows:
|
||||
data = wd.resolve_image_row(reader, row)
|
||||
result[row["id"]] = (data["workflow_data"], data["prompt_metadata"])
|
||||
return result
|
||||
|
||||
def inline_ids(self):
|
||||
return {
|
||||
row[0]
|
||||
for row in self.conn().execute(
|
||||
"SELECT id FROM generated_images WHERE workflow_data IS NOT NULL "
|
||||
"OR prompt_metadata IS NOT NULL"
|
||||
)
|
||||
}
|
||||
|
||||
def payloads(self, prompt_id):
|
||||
images = sorted(self.db.get_prompt_images(prompt_id), key=lambda i: i["id"])
|
||||
for image in images:
|
||||
image.pop("generation_time", None)
|
||||
return images
|
||||
|
||||
def migrate(self, **kwargs):
|
||||
kwargs.setdefault("max_rows", 7)
|
||||
return self.db.migrate_workflow_storage(**kwargs)
|
||||
|
||||
|
||||
class TestMigration(MigrationTestCase):
|
||||
def test_every_row_is_rewritten_and_reads_back_identical(self):
|
||||
prompt_id, expected = self.library()
|
||||
before = self.payloads(prompt_id)
|
||||
result = self.migrate()
|
||||
self.assertEqual(self.texts(), expected)
|
||||
self.assertEqual(self.inline_ids(), set())
|
||||
self.assertEqual(self.payloads(prompt_id), before)
|
||||
distinct = {t for pair in expected.values() for t in pair if t is not None}
|
||||
self.assertEqual(self.document_count(), len(distinct))
|
||||
self.assertEqual(result["rows"], len(expected))
|
||||
self.assertGreater(result["batches"], 1)
|
||||
self.assertFalse(result["stopped"])
|
||||
|
||||
def test_a_second_run_changes_nothing(self):
|
||||
_prompt_id, expected = self.library(count=12)
|
||||
self.migrate()
|
||||
conn = self.conn()
|
||||
documents = list(conn.execute("SELECT * FROM workflow_documents ORDER BY id"))
|
||||
result = self.migrate()
|
||||
self.assertEqual(result["rows"], 0)
|
||||
self.assertEqual(
|
||||
[tuple(r) for r in documents],
|
||||
[tuple(r) for r in conn.execute("SELECT * FROM workflow_documents")],
|
||||
)
|
||||
self.assertEqual(self.texts(), expected)
|
||||
|
||||
def test_a_crash_mid_migration_leaves_every_row_readable_and_resumes(self):
|
||||
prompt_id, expected = self.library()
|
||||
before = self.payloads(prompt_id)
|
||||
real = wd.store_text
|
||||
calls = []
|
||||
|
||||
def crashing(conn, text, *args):
|
||||
calls.append(text)
|
||||
if len(calls) == 23:
|
||||
raise RuntimeError("power cut")
|
||||
return real(conn, text, *args)
|
||||
|
||||
with patch.object(ws, "store_text", side_effect=crashing):
|
||||
with self.assertRaises(RuntimeError):
|
||||
self.migrate()
|
||||
self.assertFalse(self.conn().in_transaction)
|
||||
self.assertEqual(self.texts(), expected)
|
||||
self.assertEqual(self.payloads(prompt_id), before)
|
||||
remaining = self.inline_ids()
|
||||
self.assertTrue(0 < len(remaining) < len(expected))
|
||||
# The failed batch rolled back whole: what is left is a suffix by id.
|
||||
self.assertEqual(remaining, {i for i in expected if i >= min(remaining)})
|
||||
self.migrate()
|
||||
self.assertEqual(self.inline_ids(), set())
|
||||
self.assertEqual(self.texts(), expected)
|
||||
|
||||
def test_a_crash_inside_a_batch_writes_nothing_of_that_batch(self):
|
||||
_prompt_id, expected = self.library(count=10)
|
||||
with patch.object(ws, "store_text", side_effect=RuntimeError("full")):
|
||||
with self.assertRaises(RuntimeError):
|
||||
self.migrate(max_rows=50)
|
||||
self.assertEqual(self.inline_ids(), set(expected))
|
||||
self.assertEqual(self.document_count(), 0)
|
||||
|
||||
def test_a_document_that_reads_back_differently_keeps_its_row(self):
|
||||
_prompt_id, expected = self.library(count=6)
|
||||
with patch.object(ws, "read_stored_text", return_value="something else"):
|
||||
result = self.migrate(max_rows=50)
|
||||
self.assertEqual(result["errors"], sorted(expected))
|
||||
self.assertEqual(self.inline_ids(), set(expected))
|
||||
self.assertEqual(self.texts(), expected)
|
||||
self.assertEqual(self.document_count(), 0)
|
||||
|
||||
def test_stop_between_batches_then_continue(self):
|
||||
_prompt_id, expected = self.library(count=30)
|
||||
seen = []
|
||||
result = self.migrate(on_batch=seen.append, should_stop=lambda: len(seen) >= 2)
|
||||
self.assertTrue(result["stopped"])
|
||||
self.assertEqual(result["rows"], 14)
|
||||
self.assertEqual(len(self.inline_ids()), len(expected) - 14)
|
||||
self.assertEqual(seen[-1]["rows"], 14)
|
||||
self.migrate()
|
||||
self.assertEqual(self.inline_ids(), set())
|
||||
self.assertEqual(self.texts(), expected)
|
||||
|
||||
def test_values_that_are_not_text_stay_inline(self):
|
||||
# A TEXT column turns numbers into text; only a blob stays other.
|
||||
prompt_id = self.prompt()
|
||||
blob = self.old_form_image(prompt_id, "blob.png", b"\x00raw", "{}")
|
||||
other = self.old_form_image(prompt_id, "other.png", "42", b"\x01")
|
||||
result = self.migrate()
|
||||
self.assertEqual(result["protected"], 2)
|
||||
conn = self.conn()
|
||||
rows = {
|
||||
row[0]: tuple(row[1:])
|
||||
for row in conn.execute(
|
||||
"SELECT id, workflow_data, prompt_metadata, workflow_document_id, "
|
||||
"prompt_document_id FROM generated_images"
|
||||
)
|
||||
}
|
||||
self.assertEqual(rows[blob][:2], (b"\x00raw", None))
|
||||
self.assertIsNone(rows[blob][2])
|
||||
self.assertIsNotNone(rows[blob][3])
|
||||
self.assertEqual(rows[other][:2], (None, b"\x01"))
|
||||
self.assertEqual(self.texts()[other], ("42", b"\x01"))
|
||||
self.assertEqual(self.migrate()["protected"], 2) # still there, still kept
|
||||
|
||||
def test_a_row_an_older_version_rewrote_inline_takes_the_inline_text(self):
|
||||
prompt_id = self.prompt()
|
||||
image = self.db.link_image_to_prompt(
|
||||
prompt_id, os.path.join(self.tmpdir, "a.png"), {"workflow": {"v": 1}}
|
||||
)
|
||||
conn = self.conn()
|
||||
with conn: # an older version updates the inline column only
|
||||
conn.execute(
|
||||
"UPDATE generated_images SET workflow_data=? WHERE id=?",
|
||||
('{"v": 2}', image),
|
||||
)
|
||||
self.assertEqual(self.texts()[image][0], '{"v": 2}')
|
||||
self.migrate()
|
||||
self.assertEqual(self.texts()[image][0], '{"v": 2}')
|
||||
self.assertEqual(self.inline_ids(), set())
|
||||
stored = {row[0] for row in conn.execute("SELECT data FROM workflow_documents")}
|
||||
self.assertNotIn(b'{"v": 1}', {bytes(wd_data) for wd_data in stored})
|
||||
texts = [
|
||||
wd.read_stored_text(conn, None, row[0])
|
||||
for row in conn.execute("SELECT id FROM workflow_documents")
|
||||
]
|
||||
self.assertNotIn('{"v": 1}', texts)
|
||||
|
||||
def test_new_writes_during_the_migration_are_kept(self):
|
||||
prompt_id, expected = self.library(count=20)
|
||||
added = []
|
||||
|
||||
def between(_progress):
|
||||
if not added:
|
||||
added.append(
|
||||
self.db.link_image_to_prompt(
|
||||
prompt_id,
|
||||
os.path.join(self.tmpdir, "new.png"),
|
||||
{"workflow": {"new": True}},
|
||||
)
|
||||
)
|
||||
added.append(
|
||||
self.old_form_image(prompt_id, "late.png", '{"late": 1}', None)
|
||||
)
|
||||
|
||||
self.migrate(on_batch=between)
|
||||
texts = self.texts()
|
||||
self.assertEqual(texts[added[0]][0], '{"new": true}')
|
||||
self.assertEqual(texts[added[1]], ('{"late": 1}', None))
|
||||
self.assertEqual(self.inline_ids(), set())
|
||||
for image_id, pair in expected.items():
|
||||
self.assertEqual(texts[image_id], pair)
|
||||
|
||||
def test_memory_stays_bounded_by_one_row(self):
|
||||
prompt_id = self.prompt()
|
||||
big = json.dumps(sample_workflow(1, nodes=3000)) # about 0.9 MB
|
||||
for index in range(40):
|
||||
self.old_form_image(
|
||||
prompt_id, f"{index}.png", big[:-1] + f', "n": {index}}}', None
|
||||
)
|
||||
tracemalloc.start()
|
||||
try:
|
||||
self.migrate(max_rows=200, max_bytes=4 * 2**20)
|
||||
_current, peak = tracemalloc.get_traced_memory()
|
||||
finally:
|
||||
tracemalloc.stop()
|
||||
self.assertLess(peak, 12 * 2**20, f"peak {peak / 2**20:.1f} MB")
|
||||
self.assertEqual(self.inline_ids(), set())
|
||||
|
||||
def test_batches_are_bounded_by_bytes(self):
|
||||
prompt_id = self.prompt()
|
||||
text = "x" * 100_000
|
||||
for index in range(10):
|
||||
self.old_form_image(prompt_id, f"{index}.png", text + str(index), None)
|
||||
seen = []
|
||||
self.migrate(max_rows=200, max_bytes=250_000, on_batch=seen.append)
|
||||
self.assertEqual([batch["rows"] for batch in seen], [3, 6, 9, 10])
|
||||
|
||||
def test_unreferenced_documents_are_swept_at_the_end(self):
|
||||
prompt_id = self.prompt()
|
||||
image = self.db.link_image_to_prompt(
|
||||
prompt_id, os.path.join(self.tmpdir, "a.png"), {"workflow": {"v": 1}}
|
||||
)
|
||||
conn = self.conn()
|
||||
with conn:
|
||||
conn.execute("DELETE FROM generated_images WHERE id=?", (image,))
|
||||
result = self.migrate()
|
||||
self.assertEqual(result["documents_removed"], 2)
|
||||
self.assertEqual(self.document_count(), 0)
|
||||
|
||||
def test_no_media_file_is_opened_or_changed(self):
|
||||
media = os.path.join(self.tmpdir, "output")
|
||||
os.makedirs(media)
|
||||
prompt_id = self.prompt()
|
||||
for index in range(5):
|
||||
path = os.path.join(media, f"{index}.png")
|
||||
info = PngImagePlugin.PngInfo()
|
||||
info.add_text("workflow", json.dumps(sample_workflow(index)))
|
||||
with Image.new("RGB", (4, 4)) as image:
|
||||
image.save(path, pnginfo=info)
|
||||
conn = self.conn()
|
||||
with conn:
|
||||
conn.execute(
|
||||
"INSERT INTO generated_images (prompt_id, image_path, filename, "
|
||||
"file_path, workflow_data) VALUES (?, ?, ?, ?, ?)",
|
||||
(prompt_id, path, f"{index}.png", path, '{"n": 1}'),
|
||||
)
|
||||
before = {
|
||||
name: (
|
||||
_digest(os.path.join(media, name)),
|
||||
os.stat(os.path.join(media, name)),
|
||||
)
|
||||
for name in os.listdir(media)
|
||||
}
|
||||
opened = []
|
||||
real_open, real_os_open = builtins.open, os.open
|
||||
|
||||
def tracking_open(file, *args, **kwargs):
|
||||
opened.append(str(file))
|
||||
return real_open(file, *args, **kwargs)
|
||||
|
||||
def tracking_os_open(path, *args, **kwargs):
|
||||
opened.append(str(path))
|
||||
return real_os_open(path, *args, **kwargs)
|
||||
|
||||
with patch("builtins.open", tracking_open), patch("os.open", tracking_os_open):
|
||||
with patch.object(Image, "open", side_effect=AssertionError("opened")):
|
||||
self.migrate()
|
||||
self.assertEqual([p for p in opened if p.startswith(media)], [])
|
||||
for name, (digest, stat) in before.items():
|
||||
now = os.stat(os.path.join(media, name))
|
||||
self.assertEqual(_digest(os.path.join(media, name)), digest)
|
||||
self.assertEqual(now.st_mtime_ns, stat.st_mtime_ns)
|
||||
self.assertEqual(self.inline_ids(), set())
|
||||
|
||||
|
||||
class TestPreview(MigrationTestCase):
|
||||
def test_preview_counts_and_writes_nothing(self):
|
||||
prompt_id, expected = self.library(count=20)
|
||||
self.old_form_image(prompt_id, "blob.png", b"raw", None)
|
||||
files = [self.db_path + suffix for suffix in ("", "-wal")]
|
||||
before = [_digest(path) for path in files]
|
||||
draft = self.db.preview_maintenance("workflow_storage")
|
||||
self.assertEqual([_digest(path) for path in files], before)
|
||||
self.assertEqual(draft.operation, "workflow_storage")
|
||||
self.assertEqual(draft.scope["count"], len(expected))
|
||||
self.assertEqual(draft.protected["count"], 1)
|
||||
self.assertFalse(draft.media["deletes_files"])
|
||||
self.assertIn("No image or video file", draft.media["label"])
|
||||
labels = [row["label"] for row in draft.details]
|
||||
self.assertIn("Database file", labels)
|
||||
self.assertIn("Older versions", labels)
|
||||
self.assertIn("Reclaim space", labels)
|
||||
|
||||
def test_preview_without_octet_length_omits_the_inline_bytes(self):
|
||||
# SQLite before 3.43 (older Windows Python builds) has no octet_length.
|
||||
self.library(count=4)
|
||||
with patch.object(ws.sqlite3, "sqlite_version_info", (3, 37, 2)):
|
||||
draft = self.db.preview_maintenance("workflow_storage")
|
||||
labels = [row["label"] for row in draft.details]
|
||||
self.assertNotIn("Inline text", labels)
|
||||
self.assertEqual(draft.scope["count"], 4)
|
||||
|
||||
def test_preview_of_a_migrated_library_has_nothing_to_do(self):
|
||||
self.library(count=5)
|
||||
self.migrate()
|
||||
draft = self.db.preview_maintenance("workflow_storage")
|
||||
self.assertEqual(draft.scope["count"], 0)
|
||||
|
||||
def test_the_batch_reads_are_index_searches(self):
|
||||
conn = self.conn()
|
||||
for name, (sql, params) in ws.QUERY_PLAN_STATEMENTS.items():
|
||||
plan = " | ".join(
|
||||
row[3] for row in conn.execute("EXPLAIN QUERY PLAN " + sql, params)
|
||||
)
|
||||
with self.subTest(name=name):
|
||||
self.assertIn("USING INTEGER PRIMARY KEY", plan, plan)
|
||||
|
||||
|
||||
class TestReclaimSpace(MigrationTestCase):
|
||||
def test_reclaim_space_gives_the_freed_pages_back(self):
|
||||
prompt_id = self.prompt()
|
||||
big = json.dumps(sample_workflow(1, nodes=600))
|
||||
for index in range(60):
|
||||
self.old_form_image(prompt_id, f"{index}.png", big, big)
|
||||
expected = self.texts()
|
||||
self.migrate()
|
||||
conn = self.conn()
|
||||
with conn: # a document a path forgot to release
|
||||
conn.execute(
|
||||
"INSERT INTO workflow_documents (sha256, byte_length, data) "
|
||||
"VALUES (zeroblob(32), 0, x'789c030000000001')"
|
||||
)
|
||||
result = self.db.reclaim_space()
|
||||
self.assertEqual(result["documents_removed"], 1)
|
||||
self.assertFalse(result["checkpoint_busy"])
|
||||
self.assertLess(result["after_bytes"], result["before_bytes"] / 4)
|
||||
wal = self.db_path + "-wal"
|
||||
self.assertEqual(os.path.getsize(wal) if os.path.exists(wal) else 0, 0)
|
||||
self.assertEqual(self.texts(), expected)
|
||||
|
||||
|
||||
class TestRestore(MigrationTestCase):
|
||||
def test_restoring_a_pre_dedupe_backup_reads_and_migrates(self):
|
||||
prompt_id, expected = self.library(count=8)
|
||||
backup = os.path.join(self.tmpdir, "old-backup.db")
|
||||
conn = sqlite3.connect(backup)
|
||||
try: # as a version before the dedupe wrote it: no table, no columns
|
||||
conn.executescript(
|
||||
"CREATE TABLE prompts (id INTEGER PRIMARY KEY AUTOINCREMENT, "
|
||||
"text TEXT NOT NULL, created_at TIMESTAMP, updated_at TIMESTAMP, "
|
||||
"category TEXT, tags TEXT, rating INTEGER, notes TEXT, "
|
||||
"hash TEXT UNIQUE);"
|
||||
"CREATE TABLE generated_images (id INTEGER PRIMARY KEY "
|
||||
"AUTOINCREMENT, prompt_id INTEGER NOT NULL, image_path TEXT NOT "
|
||||
"NULL, filename TEXT NOT NULL, generation_time TIMESTAMP, "
|
||||
"file_size INTEGER, width INTEGER, height INTEGER, format TEXT, "
|
||||
"workflow_data TEXT, prompt_metadata TEXT, parameters TEXT, "
|
||||
"file_path TEXT, FOREIGN KEY (prompt_id) REFERENCES prompts(id) "
|
||||
"ON DELETE CASCADE, UNIQUE(prompt_id, file_path));"
|
||||
)
|
||||
conn.execute(
|
||||
"INSERT INTO prompts (id, text, hash) VALUES (?, 'a cat', 'h')",
|
||||
(prompt_id,),
|
||||
)
|
||||
for image_id, (workflow, graph) in expected.items():
|
||||
conn.execute(
|
||||
"INSERT INTO generated_images (id, prompt_id, image_path, "
|
||||
"filename, file_path, workflow_data, prompt_metadata) "
|
||||
"VALUES (?, ?, ?, ?, ?, ?, ?)",
|
||||
(
|
||||
image_id,
|
||||
prompt_id,
|
||||
f"/x/{image_id}",
|
||||
"x",
|
||||
f"/x/{image_id}",
|
||||
workflow,
|
||||
graph,
|
||||
),
|
||||
)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
self.migrate() # the live library is new form now
|
||||
self.db.model.restore_from_file(backup)
|
||||
self.assertEqual(self.texts(), expected)
|
||||
self.assertEqual(self.inline_ids(), set(expected))
|
||||
self.migrate()
|
||||
self.assertEqual(self.texts(), expected)
|
||||
self.assertEqual(self.inline_ids(), set())
|
||||
|
||||
def test_restoring_a_deduped_backup_reads_the_same(self):
|
||||
prompt_id, expected = self.library(count=8)
|
||||
self.migrate()
|
||||
backup = os.path.join(self.tmpdir, "new-backup.db")
|
||||
self.assertTrue(self.db.model.backup_database(backup))
|
||||
ok, reason = PromptModel.verify_database_file(backup)
|
||||
self.assertTrue(ok, reason)
|
||||
self.db.delete_prompt(prompt_id)
|
||||
self.assertEqual(self.document_count(), 0)
|
||||
self.db.model.restore_from_file(backup)
|
||||
self.assertEqual(self.texts(), expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,267 +0,0 @@
|
||||
"""PR #315 scoped re-review fixes: R1, R3, R4, R5 and R6 (R2 is wording).
|
||||
|
||||
- R1: only a row's own cause (undecodable text, a text that would not come
|
||||
back) keeps that row inline; disk-full, I/O and corruption fail the batch.
|
||||
- R3: nothing to rewrite is refused up front; one kept pre-migration backup
|
||||
ever, later applies take a pooled safety backup instead.
|
||||
- R4: the live job loop lets a writer on another connection in.
|
||||
- R5: the digest tool loads no package, adds no log handler and leaves no
|
||||
``-wal``/``-shm`` beside a stopped database.
|
||||
- R6: a damaged document is named, with what to do, in the verify reason and
|
||||
in the job's backup failure.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sqlite3
|
||||
import subprocess # nosec B404
|
||||
import sys
|
||||
import threading
|
||||
from unittest.mock import patch
|
||||
|
||||
from database import workflow_documents as wd
|
||||
from database import workflow_storage as ws
|
||||
from database.models import PromptModel
|
||||
|
||||
try:
|
||||
from tests.test_api_workflow_storage import OPERATION, WorkflowStorageTestCase
|
||||
from tests.test_workflow_migration import MigrationTestCase
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from test_api_workflow_storage import OPERATION, WorkflowStorageTestCase
|
||||
from test_workflow_migration import MigrationTestCase
|
||||
|
||||
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
TOOL = os.path.join(ROOT, "tools", "workflow_text_digests.py")
|
||||
KEPT = "prompts.db.pre-workflow-storage_"
|
||||
WAIT = 20
|
||||
|
||||
|
||||
def disk_full(*_args, **_kwargs):
|
||||
raise sqlite3.OperationalError("database or disk is full")
|
||||
|
||||
|
||||
class TestSystemErrorsFailTheBatch(MigrationTestCase):
|
||||
def test_a_disk_error_fails_the_batch_instead_of_listing_rows(self):
|
||||
_prompt_id, expected = self.library(count=6)
|
||||
for error in (
|
||||
sqlite3.OperationalError("database or disk is full"),
|
||||
sqlite3.OperationalError("disk I/O error"),
|
||||
sqlite3.DatabaseError("database disk image is malformed"),
|
||||
TypeError("a bug"),
|
||||
):
|
||||
with self.subTest(error=str(error)):
|
||||
with patch.object(ws, "store_text", side_effect=error):
|
||||
with self.assertRaises(type(error)):
|
||||
self.migrate(max_rows=50)
|
||||
self.assertFalse(self.conn().in_transaction)
|
||||
self.assertEqual(self.inline_ids(), set(expected))
|
||||
self.assertEqual(self.texts(), expected)
|
||||
|
||||
|
||||
class TestSystemErrorsInTheJob(WorkflowStorageTestCase):
|
||||
async def test_a_disk_error_ends_the_job_failed_with_a_truthful_line(self):
|
||||
self.old_library(count=4)
|
||||
before = self.texts()
|
||||
with patch.object(ws, "store_text", side_effect=disk_full):
|
||||
_plan, _summary, job = await self.start()
|
||||
await self.finish(job)
|
||||
self.assertEqual(job.state, "failed")
|
||||
self.assertIn("The batch in progress was rolled back", job.status_line)
|
||||
self.assertEqual(job.counters["errors"], 1) # the job's one error row
|
||||
self.assertNotIn("could not be rewritten", job.status_line)
|
||||
self.assertIsNone(self.reclaim_job())
|
||||
self.assertEqual(self.texts(), before)
|
||||
|
||||
|
||||
class TestOneKeptBackup(WorkflowStorageTestCase):
|
||||
def kept(self):
|
||||
return sorted(n for n in os.listdir(self.tmpdir) if n.startswith(KEPT))
|
||||
|
||||
async def test_nothing_to_rewrite_is_said_and_refused_without_a_backup(self):
|
||||
self.save("no images")
|
||||
plan = await self.preview(OPERATION)
|
||||
self.assertEqual(plan["scope"]["count"], 0)
|
||||
self.assertEqual(plan["scope"]["label"], "Nothing to rewrite")
|
||||
reply = await self.apply(plan["token"], status=409)
|
||||
self.assertEqual(reply["code"], "nothing-to-do")
|
||||
self.assertEqual(
|
||||
reply["error"],
|
||||
"Every image record already stores its workflow once; nothing to do.",
|
||||
)
|
||||
self.assertEqual(self.kept(), [])
|
||||
self.assertEqual(self.backups(), [])
|
||||
self.assertEqual(self.api._jobs.list(), [])
|
||||
|
||||
async def test_a_cancel_and_rerun_keep_one_kept_backup_and_pool_the_next(self):
|
||||
self.old_library(count=12)
|
||||
real = self.api.db.migrate_workflow_batch
|
||||
holder = {}
|
||||
|
||||
def first_then_cancel(*args, **kwargs):
|
||||
result = real(*args, **kwargs)
|
||||
holder["job"].cancelled = True
|
||||
return result
|
||||
|
||||
with (
|
||||
patch.object(ws, "BATCH_ROWS", 5),
|
||||
patch.object(
|
||||
self.api.db, "migrate_workflow_batch", side_effect=first_then_cancel
|
||||
),
|
||||
):
|
||||
plan = await self.preview(OPERATION)
|
||||
body = await self.apply(plan["token"])
|
||||
holder["job"] = job = self.api._jobs.get(body["job"]["id"])
|
||||
await self.finish(job)
|
||||
self.assertEqual(job.state, "cancelled")
|
||||
[kept] = self.kept()
|
||||
plan = await self.preview(OPERATION)
|
||||
self.assertEqual(plan["backup"]["policy"], "required")
|
||||
self.assertEqual(
|
||||
plan["backup"]["destination_label"], "prompts.db.backup_<date and time>"
|
||||
)
|
||||
self.assertIn(
|
||||
f"the pre-migration backup {kept} is already kept",
|
||||
" ".join(row["detail"] for row in plan["details"]),
|
||||
)
|
||||
body = await self.apply(plan["token"])
|
||||
again = self.api._jobs.get(body["job"]["id"])
|
||||
await self.finish(again)
|
||||
self.assertEqual(again.state, "finished", again.log)
|
||||
self.assertEqual(self.kept(), [kept]) # still one never-pruned copy
|
||||
[pooled] = self.backups()
|
||||
self.assertIn(f"Safety backup {pooled} verified", "\n".join(again.log))
|
||||
self.assertIn(f"pre-migration backup {kept} is kept", "\n".join(again.log))
|
||||
|
||||
|
||||
class TestLiveLoopLetsWritersIn(WorkflowStorageTestCase):
|
||||
async def test_a_writer_on_another_connection_gets_in_during_the_job(self):
|
||||
prompt_id = self.save("writer")
|
||||
conn = self.api.db.model.get_connection()
|
||||
with conn:
|
||||
for index in range(200):
|
||||
conn.execute(
|
||||
"INSERT INTO generated_images (prompt_id, image_path, filename, "
|
||||
"file_path, workflow_data) VALUES (?, ?, 'x', ?, ?)",
|
||||
(prompt_id, f"/x/{index}", f"/x/{index}", f'{{"n": {index}}}'),
|
||||
)
|
||||
started, done = threading.Event(), threading.Event()
|
||||
errors, written = [], []
|
||||
real = ws.store_text
|
||||
|
||||
def slow(conn, text, *args, **kwargs):
|
||||
started.set()
|
||||
done.wait(0.01) # a slow disk: each row takes a while
|
||||
return real(conn, text, *args, **kwargs)
|
||||
|
||||
def writer():
|
||||
started.wait(WAIT)
|
||||
own = self.api.db.model.get_connection() # this thread's connection
|
||||
own.execute("PRAGMA busy_timeout = 2000")
|
||||
try:
|
||||
for index in range(5):
|
||||
written.append(
|
||||
self.api.db.link_image_to_prompt(
|
||||
prompt_id,
|
||||
os.path.join(self.tmpdir, f"w{index}.png"),
|
||||
{"workflow": {"w": index}},
|
||||
)
|
||||
)
|
||||
except sqlite3.OperationalError as error: # pragma: no cover - RED
|
||||
errors.append(str(error))
|
||||
finally:
|
||||
self.api.db.model.close()
|
||||
|
||||
thread = threading.Thread(target=writer)
|
||||
thread.start()
|
||||
try:
|
||||
with patch.object(ws, "store_text", side_effect=slow):
|
||||
_plan, _summary, job = await self.start()
|
||||
await asyncio.wait_for(asyncio.shield(job.task), 120)
|
||||
await self.finish(job) # and Reclaim space after it
|
||||
finally:
|
||||
done.set()
|
||||
await asyncio.get_running_loop().run_in_executor(None, thread.join, WAIT)
|
||||
self.assertEqual(job.state, "finished", job.log)
|
||||
self.assertEqual(errors, [])
|
||||
self.assertEqual(len([image for image in written if image]), 5)
|
||||
|
||||
|
||||
class TestDigestToolSideEffects(MigrationTestCase):
|
||||
def run_tool(self, path, cwd):
|
||||
out = os.path.join(cwd, "out.sha")
|
||||
code = (
|
||||
"import logging, runpy, sys\n"
|
||||
"sys.argv = ['tool', sys.argv[1], sys.argv[2]]\n"
|
||||
"try:\n"
|
||||
f" runpy.run_path({TOOL!r}, run_name='__main__')\n"
|
||||
"except SystemExit:\n"
|
||||
" pass\n"
|
||||
"loaded = sorted(m for m in sys.modules\n"
|
||||
" if m.split('.')[0] in ('database', 'utils', 'py'))\n"
|
||||
"handlers = [h for n in ('', 'prompt_manager')\n"
|
||||
" for h in logging.getLogger(n).handlers]\n"
|
||||
"print('LOADED', loaded, 'HANDLERS', len(handlers))\n"
|
||||
)
|
||||
result = subprocess.run( # nosec B603
|
||||
[sys.executable, "-c", code, path, out],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
cwd=cwd,
|
||||
check=False,
|
||||
)
|
||||
self.assertEqual(result.returncode, 0, result.stderr)
|
||||
return result.stdout
|
||||
|
||||
def test_a_stopped_database_gets_no_sidecars_and_no_log(self):
|
||||
self.library(count=3)
|
||||
self.migrate()
|
||||
PromptModel.close_all_instances()
|
||||
for suffix in ("-wal", "-shm"):
|
||||
self.assertFalse(os.path.exists(self.db_path + suffix))
|
||||
cwd = os.path.join(self.tmpdir, "run")
|
||||
os.mkdir(cwd)
|
||||
before = sorted(os.listdir(self.tmpdir))
|
||||
printed = self.run_tool(self.db_path, cwd)
|
||||
self.assertEqual(sorted(os.listdir(self.tmpdir)), before)
|
||||
self.assertEqual(sorted(os.listdir(cwd)), ["out.sha"]) # no logs/ folder
|
||||
self.assertIn("LOADED [] HANDLERS 0", printed)
|
||||
with open(os.path.join(cwd, "out.sha"), encoding="utf-8") as handle:
|
||||
self.assertEqual(len(handle.read().splitlines()), 3)
|
||||
|
||||
def test_a_live_database_is_read_with_its_wal(self):
|
||||
prompt_id = self.prompt()
|
||||
self.old_form_image(prompt_id, "a.png", '{"in": "wal"}', None)
|
||||
# The connection stays open: the row may still be only in the WAL.
|
||||
cwd = os.path.join(self.tmpdir, "run")
|
||||
os.mkdir(cwd)
|
||||
self.run_tool(self.db_path, cwd)
|
||||
with open(os.path.join(cwd, "out.sha"), encoding="utf-8") as handle:
|
||||
self.assertEqual(len(handle.read().splitlines()), 1)
|
||||
|
||||
|
||||
class TestDamagedDocumentIsNamed(WorkflowStorageTestCase):
|
||||
def damage(self):
|
||||
conn = self.api.db.model.get_connection()
|
||||
with conn:
|
||||
document_id = wd.store_text(conn, '{"damaged": true}')
|
||||
conn.execute(
|
||||
"UPDATE workflow_documents SET data=x'00' WHERE id=?", (document_id,)
|
||||
)
|
||||
return document_id
|
||||
|
||||
async def test_verify_names_the_document_and_what_to_do(self):
|
||||
document_id = self.damage()
|
||||
ok, reason = PromptModel.verify_database_file(self.api.db.model.db_path)
|
||||
self.assertFalse(ok)
|
||||
self.assertEqual(reason, wd.damaged_message(document_id))
|
||||
self.assertIn(f"Workflow document {document_id} is damaged", reason)
|
||||
self.assertIn("restore a backup", reason)
|
||||
|
||||
async def test_the_backup_failure_says_which_document(self):
|
||||
self.old_library(count=2)
|
||||
document_id = self.damage()
|
||||
_plan, _summary, job = await self.start()
|
||||
await self.finish(job)
|
||||
self.assertEqual(job.state, "failed")
|
||||
self.assertEqual(job.last_event["code"], "backup-failed")
|
||||
self.assertIn(f"Workflow document {document_id} is damaged", job.status_line)
|
||||
@@ -1,374 +0,0 @@
|
||||
"""PR #315 final review fixes: I1, I2, I3, M1, M2, M3 and M5.
|
||||
|
||||
- I1: the migration always takes a verified backup, whatever the policy.
|
||||
- I2: that backup has its own name outside the pruned pool.
|
||||
- I3: ``tools/workflow_text_digests.py`` proves byte identity across the rewrite.
|
||||
- M1: a writer on another connection is never starved by the batches.
|
||||
- M2: a row that cannot be rewritten stays inline, is counted, and the rest go on.
|
||||
- M3: a damaged document is refused on restore and repaired before reuse.
|
||||
- M5: the legacy table rebuilds refuse a file with document id columns.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import os
|
||||
import re
|
||||
import sqlite3
|
||||
import subprocess # nosec B404
|
||||
import sys
|
||||
import threading
|
||||
import zlib
|
||||
from unittest.mock import patch
|
||||
|
||||
from database import workflow_documents as wd
|
||||
from database import workflow_storage as ws
|
||||
from database.models import PromptModel
|
||||
|
||||
try:
|
||||
from tests.test_api_workflow_storage import OPERATION, WorkflowStorageTestCase
|
||||
from tests.test_workflow_migration import MigrationTestCase
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from test_api_workflow_storage import OPERATION, WorkflowStorageTestCase
|
||||
from test_workflow_migration import MigrationTestCase
|
||||
|
||||
# After the API test modules, which stand in ComfyUI's ``server`` module.
|
||||
from py.api.admin import _prune_safety_backups # noqa: E402
|
||||
from py.config import PromptManagerConfig # noqa: E402
|
||||
|
||||
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
PRE_MIGRATION = re.compile(
|
||||
r"^prompts\.db\.pre-workflow-storage_\d{8}_\d{6}_\d{6}(_\d+)?$"
|
||||
)
|
||||
WAIT = 20
|
||||
|
||||
|
||||
class TestMigrationBackup(WorkflowStorageTestCase):
|
||||
def pre_migration(self):
|
||||
return sorted(
|
||||
name for name in os.listdir(self.tmpdir) if PRE_MIGRATION.match(name)
|
||||
)
|
||||
|
||||
async def test_the_backup_is_taken_even_with_the_policy_off(self):
|
||||
PromptManagerConfig.SAFETY_BACKUP = False
|
||||
self.old_library(count=4)
|
||||
plan = await self.preview(OPERATION)
|
||||
self.assertEqual(plan["backup"]["policy"], "required")
|
||||
self.assertGreater(plan["backup"]["expected_bytes"], 0)
|
||||
self.assertGreater(plan["disk"]["required_bytes"], 0)
|
||||
body = await self.apply(plan["token"])
|
||||
job = self.api._jobs.get(body["job"]["id"])
|
||||
await self.finish(job)
|
||||
self.assertEqual(job.state, "finished", job.log)
|
||||
[backup] = self.pre_migration()
|
||||
self.assertNotIn("No safety backup", "\n".join(job.log))
|
||||
self.assertIn(f"Pre-migration backup {backup} verified", "\n".join(job.log))
|
||||
|
||||
async def test_the_backup_is_outside_the_pool_never_pruned_and_restores(self):
|
||||
self.old_library(count=4)
|
||||
before = self.texts()
|
||||
plan = await self.preview(OPERATION)
|
||||
self.assertIn(
|
||||
"prompts.db.pre-workflow-storage_<date and time>",
|
||||
" ".join(row["detail"] for row in plan["details"]),
|
||||
)
|
||||
_plan, _summary, job = await self.start()
|
||||
await self.finish(job)
|
||||
[backup] = self.pre_migration()
|
||||
self.assertEqual(self.backups(), []) # not in the pruned pool
|
||||
self.assertIn(f"Pre-migration backup {backup}", job.status_line)
|
||||
path = os.path.join(self.tmpdir, backup)
|
||||
ok, reason = PromptModel.verify_database_file(path)
|
||||
self.assertTrue(ok, reason)
|
||||
_prune_safety_backups(self.api.db.model.db_path, keep=0)
|
||||
self.assertEqual(self.pre_migration(), [backup])
|
||||
# The way back: the legacy restore accepts it, old form included.
|
||||
self.api.db.model.restore_from_file(path)
|
||||
self.assertEqual(self.texts(), before)
|
||||
self.assertEqual(self.inline_count(), 4)
|
||||
|
||||
|
||||
class TestTextDigestsTool(MigrationTestCase):
|
||||
def digests(self, path, out):
|
||||
result = subprocess.run( # nosec B603
|
||||
[sys.executable, os.path.join(ROOT, "tools", "workflow_text_digests.py")]
|
||||
+ [path, out],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
cwd=self.tmpdir,
|
||||
check=False,
|
||||
)
|
||||
self.assertEqual(result.returncode, 0, result.stderr)
|
||||
with open(out, encoding="utf-8") as handle:
|
||||
return handle.read()
|
||||
|
||||
def test_digests_match_across_the_rewrite_and_read_only(self):
|
||||
_prompt_id, expected = self.library(count=12)
|
||||
PromptModel.close_all_instances()
|
||||
before = self.digests(self.db_path, os.path.join(self.tmpdir, "a.sha"))
|
||||
self.assertEqual(len(before.splitlines()), len(expected))
|
||||
self.migrate()
|
||||
PromptModel.close_all_instances()
|
||||
with open(self.db_path, "rb") as handle:
|
||||
frozen = handle.read()
|
||||
after = self.digests(self.db_path, os.path.join(self.tmpdir, "b.sha"))
|
||||
with open(self.db_path, "rb") as handle:
|
||||
self.assertEqual(handle.read(), frozen)
|
||||
self.assertEqual(after, before)
|
||||
|
||||
def test_a_changed_text_shows_in_the_diff(self):
|
||||
prompt_id = self.prompt()
|
||||
self.old_form_image(prompt_id, "a.png", '{"a": 1}', None)
|
||||
PromptModel.close_all_instances()
|
||||
before = self.digests(self.db_path, os.path.join(self.tmpdir, "a.sha"))
|
||||
conn = self.conn()
|
||||
with conn:
|
||||
conn.execute("UPDATE generated_images SET workflow_data='{\"a\":1}'")
|
||||
PromptModel.close_all_instances()
|
||||
after = self.digests(self.db_path, os.path.join(self.tmpdir, "b.sha"))
|
||||
self.assertNotEqual(after, before)
|
||||
|
||||
def test_main_in_process_hashes_undecodable_text_too(self):
|
||||
import contextlib
|
||||
import importlib.util
|
||||
import io
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"workflow_text_digests",
|
||||
os.path.join(ROOT, "tools", "workflow_text_digests.py"),
|
||||
)
|
||||
tool = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(tool)
|
||||
prompt_id = self.prompt()
|
||||
conn = self.conn()
|
||||
with conn:
|
||||
conn.execute(
|
||||
"INSERT INTO generated_images (prompt_id, image_path, filename, "
|
||||
"file_path, workflow_data) VALUES (?, '/x', 'x', '/x', "
|
||||
"CAST(x'ff' AS TEXT))",
|
||||
(prompt_id,),
|
||||
)
|
||||
out = os.path.join(self.tmpdir, "c.sha")
|
||||
printed = io.StringIO()
|
||||
with contextlib.redirect_stdout(printed):
|
||||
self.assertEqual(tool.main([self.db_path, out]), 0)
|
||||
self.assertEqual(printed.getvalue(), "1 image records\n")
|
||||
with open(out, encoding="utf-8") as handle:
|
||||
line = handle.read().split()
|
||||
self.assertEqual(line[1], hashlib.sha256(b"\xff").hexdigest())
|
||||
self.assertEqual(line[2], "-")
|
||||
|
||||
|
||||
class TestWriterIsNotStarved(MigrationTestCase):
|
||||
def test_a_writer_on_another_connection_gets_in_between_batches(self):
|
||||
prompt_id = self.prompt()
|
||||
for index in range(200):
|
||||
self.old_form_image(prompt_id, f"{index}.png", f'{{"n": {index}}}', None)
|
||||
started, done = threading.Event(), threading.Event()
|
||||
errors, written = [], []
|
||||
real = ws.store_text
|
||||
|
||||
def slow(conn, text, *args, **kwargs):
|
||||
started.set()
|
||||
done.wait(0.01) # a slow disk: each row takes a while
|
||||
return real(conn, text, *args, **kwargs)
|
||||
|
||||
def writer():
|
||||
started.wait(WAIT)
|
||||
conn = self.db.model.get_connection() # this thread's own connection
|
||||
conn.execute("PRAGMA busy_timeout = 2000")
|
||||
try:
|
||||
for index in range(5):
|
||||
written.append(
|
||||
self.db.link_image_to_prompt(
|
||||
prompt_id,
|
||||
os.path.join(self.tmpdir, f"w{index}.png"),
|
||||
{"workflow": {"w": index}},
|
||||
)
|
||||
)
|
||||
except sqlite3.OperationalError as error: # pragma: no cover - RED
|
||||
errors.append(str(error))
|
||||
finally:
|
||||
self.db.model.close()
|
||||
|
||||
thread = threading.Thread(target=writer)
|
||||
thread.start()
|
||||
with patch.object(ws, "store_text", side_effect=slow):
|
||||
self.db.migrate_workflow_storage()
|
||||
done.set()
|
||||
thread.join(WAIT)
|
||||
self.assertEqual(errors, [])
|
||||
self.assertEqual(len([image for image in written if image]), 5)
|
||||
|
||||
def test_a_batch_has_a_time_budget(self):
|
||||
prompt_id = self.prompt()
|
||||
for index in range(30):
|
||||
self.old_form_image(prompt_id, f"{index}.png", f'{{"n": {index}}}', None)
|
||||
ticks = iter(range(10**6))
|
||||
with patch.object(ws, "_clock", side_effect=lambda: next(ticks)):
|
||||
batch = ws.migrate_batch(self.conn(), 0, max_seconds=5)
|
||||
self.assertLess(batch["rows"], 30)
|
||||
self.assertGreater(batch["rows"], 0)
|
||||
|
||||
|
||||
class TestRowErrors(MigrationTestCase):
|
||||
def test_an_unreadable_row_stays_inline_and_the_rest_are_rewritten(self):
|
||||
prompt_id, expected = self.library(count=10)
|
||||
conn = self.conn()
|
||||
with conn:
|
||||
bad = conn.execute(
|
||||
"INSERT INTO generated_images (prompt_id, image_path, filename, "
|
||||
"file_path, workflow_data) VALUES (?, '/x/bad', 'bad', '/x/bad', "
|
||||
"CAST(x'7b22ff227d' AS TEXT))",
|
||||
(prompt_id,),
|
||||
).lastrowid
|
||||
result = self.migrate()
|
||||
self.assertEqual(result["errors"], [bad])
|
||||
self.assertEqual(result["rows"], len(expected))
|
||||
self.assertEqual(self.inline_ids_without(bad), set())
|
||||
raw = conn.execute(
|
||||
"SELECT CAST(workflow_data AS BLOB), workflow_document_id "
|
||||
"FROM generated_images WHERE id=?",
|
||||
(bad,),
|
||||
).fetchone()
|
||||
self.assertEqual((bytes(raw[0]), raw[1]), (b'{"\xff"}', None))
|
||||
self.assertEqual(result["documents_removed"], 0)
|
||||
|
||||
def inline_ids_without(self, image_id):
|
||||
return {
|
||||
row[0]
|
||||
for row in self.conn().execute(
|
||||
"SELECT id FROM generated_images WHERE id != ? AND "
|
||||
"(workflow_data IS NOT NULL OR prompt_metadata IS NOT NULL)",
|
||||
(image_id,),
|
||||
)
|
||||
}
|
||||
|
||||
def test_a_row_whose_text_will_not_come_back_is_kept_and_counted(self):
|
||||
prompt_id = self.prompt()
|
||||
good = self.old_form_image(prompt_id, "good.png", '{"good": 1}', None)
|
||||
bad = self.old_form_image(prompt_id, "bad.png", '{"bad": 1}', None)
|
||||
real = ws.read_stored_text
|
||||
|
||||
def lying(conn, inline, document_id):
|
||||
text = real(conn, inline, document_id)
|
||||
return "other" if text == '{"bad": 1}' else text
|
||||
|
||||
with patch.object(ws, "read_stored_text", side_effect=lying):
|
||||
result = self.migrate()
|
||||
self.assertEqual(result["errors"], [bad])
|
||||
self.assertEqual(self.inline_ids(), {bad})
|
||||
self.assertEqual(self.texts()[good][0], '{"good": 1}')
|
||||
self.assertEqual(self.texts()[bad][0], '{"bad": 1}')
|
||||
self.assertEqual(self.document_count(), 1) # the bad row's went
|
||||
|
||||
|
||||
class TestRowErrorsInTheJob(WorkflowStorageTestCase):
|
||||
async def test_the_job_reports_kept_rows_and_reclaim_space_still_runs(self):
|
||||
prompt_id = self.old_library(count=4)
|
||||
conn = self.api.db.model.get_connection()
|
||||
with conn:
|
||||
conn.execute(
|
||||
"INSERT INTO generated_images (prompt_id, image_path, filename, "
|
||||
"file_path, workflow_data) VALUES (?, '/x/bad', 'bad', '/x/bad', "
|
||||
"CAST(x'7b22ff227d' AS TEXT))",
|
||||
(prompt_id,),
|
||||
)
|
||||
_plan, _summary, job = await self.start()
|
||||
reclaim = await self.finish(job)
|
||||
self.assertEqual(job.state, "finished", job.log)
|
||||
self.assertEqual(job.counters["changed"], 4)
|
||||
self.assertEqual(job.counters["errors"], 1)
|
||||
self.assertIn(
|
||||
"1 image record could not be rewritten and keeps its inline text",
|
||||
job.status_line,
|
||||
)
|
||||
self.assertIsNotNone(reclaim)
|
||||
self.assertEqual(reclaim.state, "finished", reclaim.log)
|
||||
|
||||
|
||||
class TestDamagedDocuments(MigrationTestCase):
|
||||
def damage(self, document_id, data=b"not zlib"):
|
||||
conn = self.conn()
|
||||
with conn:
|
||||
conn.execute(
|
||||
"UPDATE workflow_documents SET data=? WHERE id=?", (data, document_id)
|
||||
)
|
||||
|
||||
def test_verify_refuses_a_damaged_document(self):
|
||||
document_id = self.store('{"a": 1}')
|
||||
ok, _reason = PromptModel.verify_database_file(self.db_path)
|
||||
self.assertTrue(ok)
|
||||
for data in (b"not zlib", zlib.compress(b'{"a": 2}')):
|
||||
self.damage(document_id, data)
|
||||
with self.subTest(data=data):
|
||||
ok, reason = PromptModel.verify_database_file(self.db_path)
|
||||
self.assertFalse(ok)
|
||||
self.assertEqual(reason, wd.damaged_message(document_id))
|
||||
|
||||
def test_a_damaged_document_is_repaired_before_it_is_reused(self):
|
||||
prompt_id = self.prompt()
|
||||
first = self.db.link_image_to_prompt(
|
||||
prompt_id, os.path.join(self.tmpdir, "a.png"), {"workflow": {"v": 1}}
|
||||
)
|
||||
document_id = self.raw(first)["workflow_document_id"]
|
||||
self.damage(document_id)
|
||||
with self.assertLogs(wd.logger, "WARNING"):
|
||||
second = self.db.link_image_to_prompt(
|
||||
prompt_id, os.path.join(self.tmpdir, "b.png"), {"workflow": {"v": 1}}
|
||||
)
|
||||
self.assertEqual(self.raw(second)["workflow_document_id"], document_id)
|
||||
for image_id in (first, second):
|
||||
self.assertEqual(self.texts()[image_id][0], '{"v": 1}')
|
||||
|
||||
def test_store_verifies_a_reused_document_once_per_given_set(self):
|
||||
document_id = self.store("abc")
|
||||
verified = set()
|
||||
conn = self.conn()
|
||||
with conn:
|
||||
self.assertEqual(wd.store_text(conn, "abc", verified), document_id)
|
||||
self.assertEqual(verified, {document_id})
|
||||
self.damage(document_id)
|
||||
with conn: # already verified in this set: not read again
|
||||
self.assertEqual(wd.store_text(conn, "abc", verified), document_id)
|
||||
|
||||
|
||||
class TestRebuildGuards(MigrationTestCase):
|
||||
def test_the_legacy_rebuilds_refuse_document_columns(self):
|
||||
for column, kind in (("prompt_id", "TEXT"), ("file_path", None)):
|
||||
path = os.path.join(self.tmpdir, f"guard-{column}.db")
|
||||
conn = sqlite3.connect(path)
|
||||
try:
|
||||
columns = [
|
||||
"id INTEGER PRIMARY KEY AUTOINCREMENT",
|
||||
f"prompt_id {kind or 'INTEGER'} NOT NULL",
|
||||
"image_path TEXT NOT NULL",
|
||||
"filename TEXT NOT NULL",
|
||||
"generation_time TIMESTAMP",
|
||||
"workflow_data TEXT",
|
||||
"prompt_metadata TEXT",
|
||||
"parameters TEXT",
|
||||
"workflow_document_id INTEGER",
|
||||
"prompt_document_id INTEGER",
|
||||
]
|
||||
if kind is not None:
|
||||
columns.append("file_path TEXT")
|
||||
conn.executescript(
|
||||
"CREATE TABLE prompts (id INTEGER PRIMARY KEY AUTOINCREMENT, "
|
||||
"text TEXT NOT NULL, created_at TIMESTAMP, updated_at TIMESTAMP, "
|
||||
"category TEXT, tags TEXT, rating INTEGER, notes TEXT, "
|
||||
"hash TEXT UNIQUE);"
|
||||
f"CREATE TABLE generated_images ({', '.join(columns)});"
|
||||
)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
with self.subTest(column=column):
|
||||
with self.assertRaisesRegex(sqlite3.DatabaseError, "workflow document"):
|
||||
PromptModel(path)
|
||||
PromptModel.close_all_instances()
|
||||
PromptModel.reset_schema_cache()
|
||||
|
||||
|
||||
if __name__ == "__main__": # pragma: no cover
|
||||
import unittest
|
||||
|
||||
unittest.main()
|
||||
Vendored
+1
-3
@@ -25,16 +25,14 @@ 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: 48 tests in 3 files/); // + the fresh-host library location case
|
||||
assert.match(executions, /Total: 46 tests in 2 files/);
|
||||
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', 'library-location.spec.mjs'],
|
||||
testIgnore: ['execution-history.spec.mjs', 'execution-outputs.spec.mjs'],
|
||||
outputDir: './test-results/comfyui',
|
||||
timeout: 60000,
|
||||
workers: 1,
|
||||
|
||||
@@ -335,12 +335,6 @@
|
||||
"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', 'library-location.spec.mjs'],
|
||||
testMatch: ['execution-history.spec.mjs', 'execution-outputs.spec.mjs'],
|
||||
outputDir: './test-results/executions/cases',
|
||||
timeout: 60000,
|
||||
workers: 1,
|
||||
|
||||
@@ -1,94 +0,0 @@
|
||||
"""Digests of every image record's stored workflow and prompt graph text.
|
||||
|
||||
Proves the workflow storage rewrite (#249 S6) kept every text byte for byte.
|
||||
Run it from this repository on the database before the rewrite (or on its
|
||||
``prompts.db.pre-workflow-storage_*`` backup) and again after, then compare::
|
||||
|
||||
python3 tools/workflow_text_digests.py prompts.db.pre-workflow-storage_<stamp> \\
|
||||
before.sha
|
||||
python3 tools/workflow_text_digests.py prompts.db after.sha
|
||||
diff before.sha after.sha # prints nothing when every text is identical
|
||||
|
||||
Each line is ``<image id> <sha256 of workflow_data> <sha256 of prompt_metadata>``
|
||||
(``-`` for none), hashing the exact stored bytes: the inline column before the
|
||||
rewrite, the document read through the one accessor after it.
|
||||
|
||||
Read-only, with no side effects: the database is opened with ``mode=ro``,
|
||||
and with ``immutable=1`` when no ``-wal``/``-shm`` file exists beside it (a
|
||||
stopped database or a backup), so SQLite creates none. A live database
|
||||
already has them and is read through its WAL. Only
|
||||
``database/workflow_documents.py`` is loaded (by path; it needs only the
|
||||
standard library), so no package start-up runs and no log file or handler is
|
||||
added. Run it as the user that runs ComfyUI.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import importlib.util
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def _load_documents():
|
||||
"""``database/workflow_documents.py`` alone, without the ``database`` package."""
|
||||
path = _ROOT / "database" / "workflow_documents.py"
|
||||
spec = importlib.util.spec_from_file_location("_pm_workflow_documents", path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def database_uri(db_path):
|
||||
"""``mode=ro``, plus ``immutable=1`` when nothing has the database open."""
|
||||
path = os.path.abspath(db_path)
|
||||
uri = Path(path).as_uri() + "?mode=ro"
|
||||
if not any(os.path.exists(path + suffix) for suffix in ("-wal", "-shm")):
|
||||
uri += "&immutable=1"
|
||||
return uri
|
||||
|
||||
|
||||
def _digest(value):
|
||||
if value is None:
|
||||
return "-"
|
||||
raw = value if isinstance(value, bytes) else value.encode("utf-8")
|
||||
return hashlib.sha256(raw).hexdigest()
|
||||
|
||||
|
||||
def write_digests(db_path, out):
|
||||
"""Write one digest line per image record of ``db_path`` to ``out``."""
|
||||
documents = _load_documents()
|
||||
conn = sqlite3.connect(database_uri(db_path), uri=True)
|
||||
try:
|
||||
conn.row_factory = sqlite3.Row
|
||||
conn.text_factory = bytes # inline text exactly as stored, even undecodable
|
||||
reader = documents.DocumentReader(conn)
|
||||
count = 0
|
||||
with open(out, "w", encoding="utf-8") as handle:
|
||||
for row in conn.execute("SELECT * FROM generated_images ORDER BY id"):
|
||||
texts = documents.resolve_image_row(reader, row)
|
||||
handle.write(
|
||||
f"{row['id']} {_digest(texts.get('workflow_data'))} "
|
||||
f"{_digest(texts.get('prompt_metadata'))}\n"
|
||||
)
|
||||
count += 1
|
||||
return count
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def main(argv=None):
|
||||
parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
|
||||
parser.add_argument("database", help="prompts.db or one of its backups")
|
||||
parser.add_argument("out", help="file to write the digests to")
|
||||
args = parser.parse_args(argv)
|
||||
count = write_digests(args.database, args.out)
|
||||
print(f"{count} image records")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
+14
-14
@@ -38,11 +38,6 @@ from typing import Dict, Any
|
||||
|
||||
from .logging_config import get_logger
|
||||
|
||||
try:
|
||||
from ..database.workflow_documents import DocumentReader, resolve_image_row
|
||||
except ImportError: # loaded as a top-level ``utils`` package
|
||||
from database.workflow_documents import DocumentReader, resolve_image_row
|
||||
|
||||
|
||||
class GalleryDiagnostics:
|
||||
"""Diagnostics for the gallery system.
|
||||
@@ -69,8 +64,8 @@ class GalleryDiagnostics:
|
||||
|
||||
Args:
|
||||
db_path: Path to the SQLite database file to diagnose.
|
||||
If None, the library's database (see
|
||||
``utils.library_location``).
|
||||
If None, resolves from config.json or defaults to
|
||||
prompts.db in the extension root.
|
||||
"""
|
||||
if db_path is None:
|
||||
db_path = self._resolve_db_path()
|
||||
@@ -79,12 +74,18 @@ class GalleryDiagnostics:
|
||||
|
||||
@staticmethod
|
||||
def _resolve_db_path() -> str:
|
||||
"""The library's database: the same resolver as PromptDatabase."""
|
||||
"""Resolve database path from config, matching PromptDatabase logic."""
|
||||
extension_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
try:
|
||||
from ..database.operations import resolve_library_location
|
||||
except ImportError:
|
||||
from database.operations import resolve_library_location
|
||||
return resolve_library_location().db_path
|
||||
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
|
||||
|
||||
def run_full_diagnostic(self) -> Dict[str, Any]:
|
||||
"""Run a complete diagnostic check.
|
||||
@@ -220,8 +221,7 @@ class GalleryDiagnostics:
|
||||
ORDER BY gi.generation_time DESC
|
||||
LIMIT 5
|
||||
""")
|
||||
reader = DocumentReader(conn)
|
||||
recent_images = [resolve_image_row(reader, row) for row in cursor]
|
||||
recent_images = [dict(row) for row in cursor.fetchall()]
|
||||
|
||||
self.logger.info(f" [TIME] Recent images: {len(recent_images)}")
|
||||
for img in recent_images:
|
||||
|
||||
+17
-67
@@ -36,7 +36,6 @@ 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
|
||||
@@ -80,10 +79,6 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
# ``record_review_file``, wired at start-up): a linked output is
|
||||
# never queued and its pending items leave the queue. None: no pass.
|
||||
self.review_recorder: Optional[Callable[[str], str]] = None
|
||||
# Queues a settled output's preview (the API's
|
||||
# ``queue_capture_thumbnail``, wired at start-up; #249). It only
|
||||
# queues: the preview is written on its own worker. None: no queue.
|
||||
self.thumbnail_queue: Optional[Callable[[str], str]] = None
|
||||
|
||||
# Read from GalleryConfig if available, otherwise use defaults
|
||||
try:
|
||||
@@ -225,30 +220,10 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
)
|
||||
return
|
||||
self.process_new_image(image_path, prompt_snapshot=prompt_snapshot)
|
||||
self._queue_thumbnail(image_path)
|
||||
finally:
|
||||
with self._pending_lock:
|
||||
self._pending.pop(image_path, None)
|
||||
|
||||
def _queue_thumbnail(self, image_path: str) -> None:
|
||||
"""Queue a settled output's preview, best effort (linked or not).
|
||||
|
||||
A failure is logged, never raised: the preview is then made on its
|
||||
first view.
|
||||
"""
|
||||
queue_preview = getattr(self, "thumbnail_queue", None)
|
||||
if queue_preview is None:
|
||||
return
|
||||
try:
|
||||
outcome = queue_preview(image_path)
|
||||
except Exception:
|
||||
self.logger.warning(
|
||||
f"Preview of {os.path.basename(image_path)} could not be queued",
|
||||
exc_info=True,
|
||||
)
|
||||
return
|
||||
self.logger.debug(f"Preview {outcome}: {os.path.basename(image_path)}")
|
||||
|
||||
def wait_until_settled(self, image_path: str, stop=None) -> bool:
|
||||
"""True once the file's size is non-zero and unchanged across two reads.
|
||||
|
||||
@@ -588,38 +563,24 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
"""
|
||||
try:
|
||||
resolution = resolve_prompt_pair(metadata)
|
||||
|
||||
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"])
|
||||
)
|
||||
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
|
||||
)
|
||||
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}"
|
||||
)
|
||||
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}"
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to link image to prompt: {e}")
|
||||
return
|
||||
@@ -675,16 +636,6 @@ class ImageMonitor:
|
||||
self.monitored_directories = []
|
||||
self.logger = get_logger("prompt_manager.image_monitor")
|
||||
self.review_recorder: Optional[Callable[[str], str]] = None
|
||||
self.thumbnail_queue: Optional[Callable[[str], str]] = None
|
||||
|
||||
def set_thumbnail_queue(
|
||||
self, queue_preview: Optional[Callable[[str], str]]
|
||||
) -> None:
|
||||
"""Queue every settled output's preview through ``queue_preview``
|
||||
(None stops it), on the running handler and every later one."""
|
||||
self.thumbnail_queue = queue_preview
|
||||
if self.handler is not None:
|
||||
self.handler.thumbnail_queue = queue_preview
|
||||
|
||||
def set_review_recorder(self, recorder: Optional[Callable[[str], str]]) -> None:
|
||||
"""Record ambiguous outputs through ``recorder`` (None stops recording).
|
||||
@@ -752,7 +703,6 @@ class ImageMonitor:
|
||||
# leaves the monitor stuck "already running" with nothing to stop.
|
||||
handler = ImageGenerationHandler(self.db_manager, self.prompt_tracker)
|
||||
handler.review_recorder = self.review_recorder
|
||||
handler.thumbnail_queue = self.thumbnail_queue
|
||||
observer = Observer()
|
||||
scheduled = []
|
||||
for output_dir in valid_directories:
|
||||
|
||||
@@ -16,21 +16,10 @@ nesting before it is validated. Every class appears in exactly one prompt
|
||||
category; classes that are not listed stay ``unsupported-node``. ``lineage``
|
||||
(image and latent pass-through outputs for multi-stage chains) describes no
|
||||
prompt, so its classes are not claimed by a category.
|
||||
|
||||
``Semantics.version`` is the SHA-256 of the table's canonical JSON (sorted
|
||||
keys, no whitespace): any change of content, never of layout, moves it.
|
||||
``class_version(doc, classes)`` is the narrower fingerprint Needs review
|
||||
folds into the version of the items it records from the inspection profile
|
||||
(``utils.review_capture``): the SHA-256 of ``class_projection``, the table
|
||||
restricted to what it says about ``classes`` (the node classes of one graph)
|
||||
plus the class-independent ``schema`` and ``stage_inputs``. The resolver looks
|
||||
the table up only by the class of a graph node, so an edit about other
|
||||
classes cannot change how that graph resolves and leaves the version alone.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
from dataclasses import dataclass, replace
|
||||
from dataclasses import dataclass
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
@@ -129,10 +118,6 @@ class Semantics:
|
||||
stage_inputs: Mapping[str, Tuple[str, ...]]
|
||||
facts: Mapping[str, Any]
|
||||
all_classes: FrozenSet[str]
|
||||
# SHA-256 of the canonical table JSON (set by ``load_semantics``).
|
||||
version: str = ""
|
||||
# The canonical table JSON itself, for ``class_version``.
|
||||
source: str = ""
|
||||
|
||||
|
||||
def _freeze(value: Any) -> Frozen:
|
||||
@@ -536,71 +521,7 @@ def load_semantics(path: Optional[Union[str, Path]] = None) -> Semantics:
|
||||
except (OSError, UnicodeDecodeError, ValueError) as error:
|
||||
raise SemanticsError(f"cannot read semantics table {source.name}: {error}")
|
||||
_check_nesting(doc)
|
||||
return replace(
|
||||
_validate(doc), version=table_version(doc), source=_canonical_json(doc)
|
||||
)
|
||||
|
||||
|
||||
# Table parts that hold no class names: they apply to every graph.
|
||||
_CLASS_INDEPENDENT = frozenset({"schema", "stage_inputs"})
|
||||
|
||||
|
||||
def _canonical_json(doc: Any) -> str:
|
||||
return json.dumps(doc, sort_keys=True, separators=(",", ":"), ensure_ascii=True)
|
||||
|
||||
|
||||
def table_version(doc: Any) -> str:
|
||||
"""SHA-256 of ``doc`` as canonical JSON: content, not layout."""
|
||||
return hashlib.sha256(_canonical_json(doc).encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def _about(value: Any, classes: FrozenSet[str]) -> Any:
|
||||
"""The part of one category that names a class in ``classes``.
|
||||
|
||||
Categories are keyed by class (objects), list classes (strings) or list
|
||||
entries that carry a ``class``.
|
||||
"""
|
||||
if isinstance(value, dict):
|
||||
return {key: item for key, item in value.items() if key in classes}
|
||||
if isinstance(value, list):
|
||||
return [
|
||||
entry
|
||||
for entry in value
|
||||
if (isinstance(entry, str) and entry in classes)
|
||||
or (isinstance(entry, dict) and entry.get("class") in classes)
|
||||
]
|
||||
return value
|
||||
|
||||
|
||||
def class_projection(doc: Mapping[str, Any], classes) -> dict:
|
||||
"""What ``doc`` says about ``classes``, plus its class-independent parts."""
|
||||
classes = frozenset(classes)
|
||||
projected = {}
|
||||
for category, value in doc.items():
|
||||
if category in _CLASS_INDEPENDENT:
|
||||
projected[category] = value
|
||||
elif category == "facts" and isinstance(value, dict):
|
||||
projected[category] = {
|
||||
table: _about(entries, classes) for table, entries in value.items()
|
||||
}
|
||||
else:
|
||||
projected[category] = _about(value, classes)
|
||||
return projected
|
||||
|
||||
|
||||
def class_version(doc: Mapping[str, Any], classes) -> str:
|
||||
"""SHA-256 of ``class_projection(doc, classes)`` as canonical JSON."""
|
||||
return table_version(class_projection(doc, classes))
|
||||
|
||||
|
||||
@lru_cache(maxsize=256)
|
||||
def _cached_class_version(source: str, classes: FrozenSet[str]) -> str:
|
||||
return class_version(json.loads(source), classes)
|
||||
|
||||
|
||||
def semantics_class_version(semantics: Semantics, classes) -> str:
|
||||
"""``class_version`` of a loaded table (cached per table and class set)."""
|
||||
return _cached_class_version(semantics.source, frozenset(classes))
|
||||
return _validate(doc)
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
|
||||
@@ -1,480 +0,0 @@
|
||||
"""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)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user