Compare commits

..
Author SHA1 Message Date
Vito Sansevero c1e74bd0f1 fix(thumbnails): answer shared flights with the outcome only and restart a failed worker
The history and serve routes share a flight key, so a flight now answers
only the ensure outcome whichever route started it; the serve route
derives the preview path and, for a fallback, looks its source up again
on the view pool. A view joining a history flight no longer gets a 500
or an image answered with the video placeholder.

The capture worker clears its running state however its loop ends, and a
worker that cannot start leaves the file queued for the next enqueue.
2026-10-24 08:18:02 -07:00
Vito Sansevero e43157e684 fix(thumbnails): let an idle preview worker leave and fit case-insensitive filesystems
The capture worker exits once its queue is empty, so no thread outlives
its work; the next output starts a fresh one. The mixed-case source test
expects no folder listing where the exact name already resolves (macOS,
Windows).
2026-10-24 07:53:13 -07:00
Vito Sansevero 44c2fc4133 fix(thumbnails): cover the image fallback card and document the admission
An image listed without a preview (a nested thumbnails folder) still
shows its original; TESTING.md describes the view admission, the video
placeholder, the holds before the rename and the bounded shutdown.
2026-10-24 07:19:05 -07:00
Vito Sansevero 78e6ed7a2d Merge remote-tracking branch 'origin/main' into fix/249-thumbnails-automatic 2026-10-24 06:53:53 -07:00
Vito Sansevero 297b2060ee fix(thumbnails): admit previews on the loop and never send a video to an img
Final review fixes for automatic previews:
- A video preview that cannot be made is the placeholder frame (no-store),
  never the video; OpenCV failures fall through to ffmpeg and the
  placeholder. Cards and the inspector cover swap back to the video
  element when the preview image fails.
- View requests are admitted on the event loop (3 writers, one flight per
  file, at most 64 waiting; past that the fallback at once). Writes run on
  a small dedicated pool, never the default executor. The history route
  shares the admission, flight key and link checks.
- Drive, UNC, empty and dot segments are refused as text; the folder
  listing runs inside the admission.
- Holds are checked after the slot and again before the rename; only
  decode failures are remembered.
- Bounded, logged shutdown: ffmpeg is killed, no worker restarts after
  cleanup, and a reload stops the old worker.
- The listing says whether the preview file exists (thumbnail_ready) for
  the legacy gallery's dot; media in a nested thumbnails folder keeps its
  original; file names the server quotes match in the inspector.
2026-10-24 06:53:51 -07:00
Vito 06099579e6 feat(maintenance): store each workflow once (#315)
* feat(maintenance): reviewed, read-only maintenance plans

POST /prompt_manager/maintenance/preview reviews duplicate_prompts,
optimize, unused_components or missing_links on its own mode=ro snapshot
and answers a Plan (scope, protected reasons, backup policy, media, at most
50 detail rows). The fingerprint, targets and judged sources stay in a
server-side PlanStore under a single-use token (15 min, 8 plans, 16 MiB,
never truncated; a restart forgets them), bound to the database it read.

The previews reuse the legacy queries: duplicate groups and the merge are
shared with cleanup_duplicates, ORPHANED_PROMPTS_SQL is built from the
#239 protection predicates the plan reports, and missing links are judged
by the health check's link_states. Legacy /maintenance keeps its shape;
vacuum now reports success:false when VACUUM fails, and
cleanup_orphaned_images skips links under unavailable sources.

* feat(maintenance): apply reviewed plans as Activity jobs

POST /prompt_manager/maintenance/apply {token} starts a maintenance job
(records admission, not cancellable, groups changed/protected/errors).
The token is consumed on the loop turn the job starts: a second apply is
409 plan-used naming the job, a busy refusal keeps the token, an unknown
or expired one is 410 plan-expired, and a restore since the review is
409 storage-changed.

With the safety backup policy on (maintenance.safety_backup, default true,
saved through the settings route with the revision), the job first writes
a backup into a temp file beside the database, makes it a standalone
rollback-journal file, verifies it, fsyncs it and renames it onto a unique
<db>.backup_<YYYYmmdd_HHMMSS_ffffff>[_n] name; any failure leaves no file
and nothing changed. The apply then rechecks the reviewed fingerprint
inside BEGIN IMMEDIATE (plan-changed writes nothing), deletes in batches of
400 ids, counts changed and protected records, logs the backup by name
only and prunes old safety backups. With the policy off the log and the
status line say "No safety backup".

* feat(maintenance): find and remove duplicate files through Activity

Duplicate image files move onto the job registry:

- POST /prompt_manager/jobs/duplicates finds content-identical media
  files over every effective source as a cancellable read job
  ("Find duplicate image files", group groups_found); the server holds
  the groups per job and GET /prompt_manager/maintenance/duplicates/{id}
  pages them ten at a time.
- The duplicate_files preview binds a removal to the reviewed keepers.
  Every keeper and removal is checked again: containment as text before
  any filesystem access, then not a link, still inside its source, same
  size and hash. The plan lists the database links that move.
- Apply runs duplicate_cleanup (records, not cancellable): the promised
  safety backup, then per file a recheck, the unlink with every handle
  closed and its thumbnail. A changed file is kept, the first failed
  unlink stops the job with a coded message (never the OS text), and only
  removed files have their image links and favorites moved to the kept
  copy.
- The legacy scan, stream, status and delete routes keep their shapes on
  the same registry kinds; the hashing, grouping and thumbnail removal
  are shared with them.

Coordinator rulings on Tasks 15-16:

- Safety backups share the restore pool of 10 `<db>.backup_*` files; the
  pool is now also pruned when an apply fails after its backup.
- Plans report the safety backup's size and the free disk space, and
  apply refuses up front (507 insufficient-space) when the backup or
  VACUUM does not fit, so no job fails halfway for space.

* fix(maintenance): never remove the kept copy; move records before unlink

Final review of #313:

- C1: the duplicate find tells files apart by (st_dev, st_ino), so a
  source listed twice, a symlinked source inside another and a hardlink
  are one file. The preview and the job refuse a removal that
  os.path.samestat matches with its kept copy.
- I1: each file's records move to its verified kept copy in one short
  transaction before that file is unlinked; a failed move deletes
  nothing and a failed unlink leaves the records on the identical kept
  copy. The status line counts files removed and not removed.
- I2: links are looked up under every spelling of the source (the real
  path, the configured and monitored folders, ComfyUI's output).
- I5: image tags join the kept copy as their union and the removed
  copy's pending, undecided Needs review items are settled; the plan
  says so.
- I3: the VACUUM disk check measures SQLite's real temporary folder.
- M1: the token is looked at again with no await before the job starts.
- M2: the backup decision is made once and handed to the job.
- M3: backups are linked into place; startup sweeps interrupted temps.
- M4: Optimize runs alone (new "exclusive" admission) and its plan says
  the database is locked while VACUUM runs.
- M5: the reviewed duplicate merge is strict: collisions and Library
  titles are protected, and errors roll back. The legacy cleanup is
  unchanged.
- M7: 2,000 keep entries within a 1 MiB body. M8: event waits in tests.
- M9: the preview binds files by size, mtime and inode without hashing;
  the job hashes right before each unlink.

* fix(maintenance): protect one file reached through two mounts

Re-review of #313, round 2:

- R1: a removal on another device than its kept copy, with the same
  size and mtimes within 2 s, is protected when its inode is equal or its
  path within its source is the same ("May be the same file reached
  through another mount or share"), in the preview and again right before
  each unlink. The job holds the kept copy open during the unlink, so on
  Windows deleting the kept copy itself through another path fails.
- R2: a file's spellings are found by containment: every configured,
  monitored, ComfyUI-output or listed folder whose stat identity is a
  folder above the file. Parent spellings of nested sources and case
  variants match; favorites move under every listed root holding the
  file. Lookups stay index searches.
- m1: a failed unlink's row says the links now name the kept copy.
- m2: a volume with no file ids reads "Can't confirm this is a separate
  file on this drive".

* fix(maintenance): stat duplicate candidates with os.stat

Python 3.10's Path.stat uses the os.stat bound when pathlib was imported,
so the find and the per-file checks could see different stat calls. The
find now calls os.stat directly, like every other check; the mount tests
failed on py3.10 only.

* fix(maintenance): recheck the kept copy right before each unlink

Re-review 2 of #313:

- n1: the cross-mount guard also matches when one path within its source
  ends with the other on a whole folder name (sources listed at different
  depths, e.g. 2024/x.png and x.png on a FUSE or noserverino mount).
- n2: after a file's records move and right before its unlink, the held
  kept copy must still be the file at its path (fstat of the hold against
  os.stat) and the cross-mount guard must still hold; otherwise the file
  is kept, its records move back to it, and it counts as not removed.

* test(maintenance): the kept-copy race tests hold on Windows

On Windows the removal job holds the kept copy open without
FILE_SHARE_DELETE, so the tests' attempt to delete or replace it during
the record move is refused. The tests now assert that native outcome:
the kept copy stays and the removal goes ahead.

* feat(database): store each workflow and prompt graph once

generated_images.workflow_data and prompt_metadata repeated the same large
JSON text for every image of a run. Each distinct text now lives once in
workflow_documents, keyed by the sha256 of its exact UTF-8 bytes and
zlib-compressed; image rows reference it through workflow_document_id and
prompt_document_id.

- New writes (link_image_to_prompt, pair ingestion for scan, monitor and
  Needs review) store the new form. Old rows keep their inline text and
  read the same until the migration rewrites them.
- Every reader goes through one accessor (read_stored_text,
  resolve_image_row): gallery and legacy image lists, image by id, prompt
  info for the inspector and metadata page, diagnostics. Payloads keep their
  shape. Inline text wins when present; a document that fails its length or
  hash check reads as unreadable, never as other text.
- No trigger (restore validators refuse them): a document goes in the same
  transaction as its last reference (image and prompt deletes, missing-link
  cleanup, duplicate merges, duplicate-file link merges, ingestion updates),
  and a sweep at startup catches anything a path missed.
- verify_database_file validates the table, the id columns, their foreign
  keys and indexes, and the stored values.

* feat(database): migrate old workflow rows in batches, then reclaim space

The workflow storage migration rewrites old-form image rows (text inline)
to document ids, batch by batch:

- Each batch is one BEGIN IMMEDIATE transaction of at most 200 rows or
  16 MB of text, found by a rowid range search and read one row at a time.
- Each text is stored or reused, then read back through the one accessor
  before its row changes; a text that would not come back identical fails
  and rolls back the whole batch.
- A value that is not text stays inline and is counted as protected.
- A crash or a stop leaves committed batches new form and the rest old
  form; both read the same, and a rerun continues. No media file is opened.
- Unreferenced documents are swept at the end.

The maintenance preview gains a read-only workflow_storage plan (records to
rewrite, inline bytes where SQLite has octet_length, the Reclaim space and
older-version notes). reclaim_space sweeps, runs VACUUM and truncates the
WAL so the freed pages go back to the disk.

* feat(maintenance): store each workflow once as an Activity job

The reviewed maintenance plan gains workflow_storage. Applying it starts
"Store each workflow once": kind maintenance, records admission,
cancellable, groups changed/protected/errors.

- The promised verified safety backup comes first; a failed backup ends the
  job with nothing rewritten. The apply's up-front disk check covers it.
- Old-form rows are rewritten batch by batch on the event loop's executor,
  with progress after each committed batch. Cancel stops between batches;
  a failed batch rolls back whole and fails the job with a coded message.
  Either way every record reads as before and the next apply continues.
- When every row was looked at, "Reclaim space" starts: exclusive (runs
  alone), not cancellable. It checks the disk for VACUUM's copies of the
  live pages, then sweeps, runs VACUUM and truncates the WAL. If another job
  holds the registry it does not start, and the finished migration says so
  and names Optimize database.
- No image or video file is opened; tests track every open under the
  output folder and compare the media bytes and mtimes.

* feat(maintenance): document the workflow storage contract

TESTING.md's maintenance contract gains the workflow storage rules: the
documents table and its two forms, the one accessor, garbage in the same
transaction plus the sweeps, the batched migration and Reclaim space, the
media and memory guarantees, and that a deduped backup cannot be restored
into an older version.

A test covers the preview on SQLite before 3.43 (no octet_length, as in
older Windows Python builds): the inline-bytes detail is left out.

* feat(database): keep the migration safe row by row and beside other writers

Final review fixes M1, M2, M3 and M5 for the workflow storage dedupe:

- M1: a batch also ends after 0.5 s, and the loop leaves the write lock
  free for 0.15 s after each committed batch (longer than SQLite's longest
  busy poll), so the image monitor and run counting on another connection
  get in between. RED: a threaded writer with a 2 s busy timeout hit
  "database is locked" (4 of 5 links landed).
- M2: each row runs in its own savepoint. A row that cannot be rewritten
  (undecodable text, a text that would not come back) is rolled back alone,
  stays inline and is listed in errors; the rest of the batch goes on.
- M3: a reused document is decompressed and checked before its id is handed
  out, and repaired from the text (whose hash is its key) when damaged;
  verify_database_file reads every document and refuses a damaged one.
- M5: the two legacy generated_images rebuilds refuse a file that already
  has the document id columns instead of dropping them.
- The preview names the pre-migration backup and says it is always taken.

* feat(maintenance): always back up before storing workflows once

Final review fixes I1, I2 and I3 for the workflow storage dedupe:

- I1: the migration always takes a verified backup, whatever
  maintenance.safety_backup says; the plan reports backup.policy
  "required" and the disk check counts it.
- I2: that backup is prompts.db.pre-workflow-storage_<stamp>, outside the
  pruned prompts.db.backup_* pool and kept until the user removes it; the
  job's last line and the preview name it, and the legacy restore accepts
  it. README's changelog warns that an older version opens and restores a
  rewritten database but shows no workflows, and says how to go back.
- I3: tools/workflow_text_digests.py writes the sha256 of every record's
  stored texts (read-only; inline before the rewrite, through the accessor
  after), so a diff before and after proves byte identity.
- A row the migration could not rewrite is listed in the job's errors and
  named in its last line, and Reclaim space still runs.

* feat(database): fail the batch on system errors and name a damaged document

Re-review fixes R1 and R6 (database side), plus R3's helpers:

- R1: a row is kept inline only for its own cause: text SQLite cannot
  decode, or a text that would not come back. A full disk, an I/O error, a
  malformed database or a bug now fails the batch and rolls it back,
  instead of listing every remaining row as "could not be rewritten".
- R6: verify_database_file names the damaged document and says what to do
  (restore a backup from before the damage, or store an image with the same
  workflow, which repairs it).
- R3: pre_migration_backups lists the kept backups, earliest first; the plan
  says "Nothing to rewrite" and names an existing kept backup, and the older
  version note says new images use the new form from the first start and
  names the earliest backup as the way back.

* feat(maintenance): one kept pre-migration backup and nothing-to-do refusal

Re-review fixes R2, R3, R4, R5 and R6 (job side):

- R3: Apply refuses a plan with nothing to rewrite (409 nothing-to-do, no
  job, no backup). Only the first apply keeps a pre-workflow-storage
  backup; later applies (after a Cancel, a rerun) take a pooled safety
  backup, forced on whatever the policy, so one copy is never pruned.
- R4: a second-connection writer test runs against the live job loop; with
  its pause removed the writer hits "database is locked" (4 of 5 links).
- R5: the digest tool loads only database/workflow_documents.py by path (no
  package start-up, no log handler) and opens a stopped database with
  immutable=1, so it leaves no -wal/-shm; a live one is read through its WAL.
- R6: a backup that fails on a damaged document says which one in the job
  line.
- R2: README's changelog says images recorded by this version use the new
  form from the first start, and the way back is the earliest
  pre-workflow-storage backup (or a pre-update backup if the job never ran).
2026-10-24 06:50:53 -07:00
Vito Sansevero cfa47ccac8 fix(thumbnails): keep the video fallback tested and stop the preview worker before cleanup
A video listed without a preview still renders its own element and is
released on refresh. Test teardowns and the browser reset stop the
automatic preview worker before the temp tree goes, so no preview lands
after cleanup (Windows). TESTING.md documents automatic previews.
2026-10-24 06:16:49 -07:00
Vito Sansevero 64f5b91eee fix(thumbnails): generate thumbnails automatically and use them everywhere
New outputs never got a preview, so the workspace grid loaded full-size
originals (844 of 19,907 images on a live install, exactly the newest).

- Listings always give images and videos a thumbnail URL; a missing
  preview is made by the serve route on first view (bounded, single
  flight per file) and falls back to the original, remembering files
  that cannot be decoded.
- The image monitor queues each settled output's preview on one
  background worker that follows Worker threads.
- A scan that leaves 200 or more previews missing starts the visible
  thumbnail job; fewer go through the worker.
- Automatic previews wait for a restore, a duplicate removal, a sources
  change, the gallery's clear and exclusive jobs.
- The inspector cover uses the preview; the Overview row counts
  previews that failed.
2026-10-24 06:01:37 -07:00
Vito f68e92b3b6 feat(system): maintenance backend (#313)
* feat(maintenance): reviewed, read-only maintenance plans

POST /prompt_manager/maintenance/preview reviews duplicate_prompts,
optimize, unused_components or missing_links on its own mode=ro snapshot
and answers a Plan (scope, protected reasons, backup policy, media, at most
50 detail rows). The fingerprint, targets and judged sources stay in a
server-side PlanStore under a single-use token (15 min, 8 plans, 16 MiB,
never truncated; a restart forgets them), bound to the database it read.

The previews reuse the legacy queries: duplicate groups and the merge are
shared with cleanup_duplicates, ORPHANED_PROMPTS_SQL is built from the
#239 protection predicates the plan reports, and missing links are judged
by the health check's link_states. Legacy /maintenance keeps its shape;
vacuum now reports success:false when VACUUM fails, and
cleanup_orphaned_images skips links under unavailable sources.

* feat(maintenance): apply reviewed plans as Activity jobs

POST /prompt_manager/maintenance/apply {token} starts a maintenance job
(records admission, not cancellable, groups changed/protected/errors).
The token is consumed on the loop turn the job starts: a second apply is
409 plan-used naming the job, a busy refusal keeps the token, an unknown
or expired one is 410 plan-expired, and a restore since the review is
409 storage-changed.

With the safety backup policy on (maintenance.safety_backup, default true,
saved through the settings route with the revision), the job first writes
a backup into a temp file beside the database, makes it a standalone
rollback-journal file, verifies it, fsyncs it and renames it onto a unique
<db>.backup_<YYYYmmdd_HHMMSS_ffffff>[_n] name; any failure leaves no file
and nothing changed. The apply then rechecks the reviewed fingerprint
inside BEGIN IMMEDIATE (plan-changed writes nothing), deletes in batches of
400 ids, counts changed and protected records, logs the backup by name
only and prunes old safety backups. With the policy off the log and the
status line say "No safety backup".

* feat(maintenance): find and remove duplicate files through Activity

Duplicate image files move onto the job registry:

- POST /prompt_manager/jobs/duplicates finds content-identical media
  files over every effective source as a cancellable read job
  ("Find duplicate image files", group groups_found); the server holds
  the groups per job and GET /prompt_manager/maintenance/duplicates/{id}
  pages them ten at a time.
- The duplicate_files preview binds a removal to the reviewed keepers.
  Every keeper and removal is checked again: containment as text before
  any filesystem access, then not a link, still inside its source, same
  size and hash. The plan lists the database links that move.
- Apply runs duplicate_cleanup (records, not cancellable): the promised
  safety backup, then per file a recheck, the unlink with every handle
  closed and its thumbnail. A changed file is kept, the first failed
  unlink stops the job with a coded message (never the OS text), and only
  removed files have their image links and favorites moved to the kept
  copy.
- The legacy scan, stream, status and delete routes keep their shapes on
  the same registry kinds; the hashing, grouping and thumbnail removal
  are shared with them.

Coordinator rulings on Tasks 15-16:

- Safety backups share the restore pool of 10 `<db>.backup_*` files; the
  pool is now also pruned when an apply fails after its backup.
- Plans report the safety backup's size and the free disk space, and
  apply refuses up front (507 insufficient-space) when the backup or
  VACUUM does not fit, so no job fails halfway for space.

* fix(maintenance): never remove the kept copy; move records before unlink

Final review of #313:

- C1: the duplicate find tells files apart by (st_dev, st_ino), so a
  source listed twice, a symlinked source inside another and a hardlink
  are one file. The preview and the job refuse a removal that
  os.path.samestat matches with its kept copy.
- I1: each file's records move to its verified kept copy in one short
  transaction before that file is unlinked; a failed move deletes
  nothing and a failed unlink leaves the records on the identical kept
  copy. The status line counts files removed and not removed.
- I2: links are looked up under every spelling of the source (the real
  path, the configured and monitored folders, ComfyUI's output).
- I5: image tags join the kept copy as their union and the removed
  copy's pending, undecided Needs review items are settled; the plan
  says so.
- I3: the VACUUM disk check measures SQLite's real temporary folder.
- M1: the token is looked at again with no await before the job starts.
- M2: the backup decision is made once and handed to the job.
- M3: backups are linked into place; startup sweeps interrupted temps.
- M4: Optimize runs alone (new "exclusive" admission) and its plan says
  the database is locked while VACUUM runs.
- M5: the reviewed duplicate merge is strict: collisions and Library
  titles are protected, and errors roll back. The legacy cleanup is
  unchanged.
- M7: 2,000 keep entries within a 1 MiB body. M8: event waits in tests.
- M9: the preview binds files by size, mtime and inode without hashing;
  the job hashes right before each unlink.

* fix(maintenance): protect one file reached through two mounts

Re-review of #313, round 2:

- R1: a removal on another device than its kept copy, with the same
  size and mtimes within 2 s, is protected when its inode is equal or its
  path within its source is the same ("May be the same file reached
  through another mount or share"), in the preview and again right before
  each unlink. The job holds the kept copy open during the unlink, so on
  Windows deleting the kept copy itself through another path fails.
- R2: a file's spellings are found by containment: every configured,
  monitored, ComfyUI-output or listed folder whose stat identity is a
  folder above the file. Parent spellings of nested sources and case
  variants match; favorites move under every listed root holding the
  file. Lookups stay index searches.
- m1: a failed unlink's row says the links now name the kept copy.
- m2: a volume with no file ids reads "Can't confirm this is a separate
  file on this drive".

* fix(maintenance): stat duplicate candidates with os.stat

Python 3.10's Path.stat uses the os.stat bound when pathlib was imported,
so the find and the per-file checks could see different stat calls. The
find now calls os.stat directly, like every other check; the mount tests
failed on py3.10 only.

* fix(maintenance): recheck the kept copy right before each unlink

Re-review 2 of #313:

- n1: the cross-mount guard also matches when one path within its source
  ends with the other on a whole folder name (sources listed at different
  depths, e.g. 2024/x.png and x.png on a FUSE or noserverino mount).
- n2: after a file's records move and right before its unlink, the held
  kept copy must still be the file at its path (fstat of the hold against
  os.stat) and the cross-mount guard must still hold; otherwise the file
  is kept, its records move back to it, and it counts as not removed.

* test(maintenance): the kept-copy race tests hold on Windows

On Windows the removal job holds the kept copy open without
FILE_SHARE_DELETE, so the tests' attempt to delete or replace it during
the record move is refused. The tests now assert that native outcome:
the kept copy stays and the removal goes ahead.

* fix(maintenance): restore a backup of any size the disk has room for

The legacy restore refused the maintainer's own 1.6 GB safety backup with
"File too large. Maximum size is 100MB", a fixed cap since #153. Free
disk space now decides: the upload, its staged copy and the swap must fit,
under a 64 GiB ceiling. A request with a Content-Length that does not fit
is refused before anything is written, one without as soon as its running
size does not, with 507 insufficient-space and a message without paths.
The upload still streams and its temp file is removed on every path.
2026-10-24 05:28:44 -07:00
50 changed files with 6857 additions and 287 deletions
+5
View File
@@ -970,6 +970,11 @@ MIT License - see LICENSE file for details.
## Changelog
### 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))
+124 -3
View File
@@ -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 when the server knows one
(`thumbnail_url`, the same URL the Images page uses), else the original; a re-read
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
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
@@ -2230,6 +2230,59 @@ 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
@@ -2277,7 +2330,8 @@ 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`) or "Thumbnail status unavailable"; Review thumbnails
`/images/thumbnails/status`, followed by " · M failed" when automatic previews
failed this session) 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
@@ -2708,6 +2762,73 @@ 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).
## Workspace data fields for the concept pages (#249)
These fields are additive. The legacy `/prompt_manager/images/output` listing (no
+8 -1
View File
@@ -132,11 +132,18 @@ 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):
async def _stop_review_recording(
app, monitor=_global_image_monitor, api=_review_api
):
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)
+3
View File
@@ -25,6 +25,7 @@ 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 = (
@@ -155,7 +156,9 @@ class DuplicateCopyOperations:
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(
+18 -2
View File
@@ -42,12 +42,15 @@ 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)
@@ -260,16 +263,24 @@ def _orphaned_ids(conn):
def _delete_ids(conn, table, ids):
"""Delete ``table`` rows by id in batches; returns the rows removed."""
"""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
@@ -379,10 +390,13 @@ class MaintenanceOperations:
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
@@ -408,6 +422,8 @@ class MaintenanceOperations:
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)
@@ -656,7 +672,7 @@ class MaintenanceOperations:
records changed and up to ``MAX_RESULT_ROWS`` descriptions of them
(``{prompt_id[, kept_id]}`` or ``{image_path}``).
"""
if operation not in MAINTENANCE_OPERATIONS:
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.
+29
View File
@@ -24,6 +24,13 @@ 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,
@@ -50,6 +57,19 @@ 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):
@@ -316,6 +336,8 @@ 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")
@@ -711,6 +733,7 @@ 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):
@@ -876,6 +899,7 @@ 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):
@@ -1117,6 +1141,11 @@ 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}"
+79 -22
View File
@@ -25,6 +25,14 @@ 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,
@@ -339,6 +347,7 @@ class PromptDatabase(
HealthOperations,
MaintenanceOperations,
DuplicateCopyOperations,
WorkflowStorageOperations,
):
"""Database operations class for managing prompts."""
@@ -776,10 +785,12 @@ 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()
@@ -803,10 +814,12 @@ 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,
@@ -1607,6 +1620,14 @@ class PromptDatabase(
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)
@@ -1614,8 +1635,9 @@ class PromptDatabase(
"""
INSERT OR IGNORE INTO generated_images
(prompt_id, image_path, filename, file_path, file_size, width,
height, format, workflow_data, prompt_metadata, parameters)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
height, format, workflow_data, prompt_metadata, parameters,
workflow_document_id, prompt_document_id)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
(
prompt_id_int,
@@ -1626,11 +1648,22 @@ class PromptDatabase(
width,
height,
image_format,
dumps_strict(metadata.get("workflow", {})),
dumps_strict(metadata.get("prompt", {})),
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:
@@ -1664,7 +1697,8 @@ class PromptDatabase(
""",
(prompt_id,),
)
return [self._image_row_to_dict(row) for row in cursor.fetchall()]
reader = DocumentReader(conn)
return [self._image_row_to_dict(row, reader) for row in cursor]
def get_recent_images(
self, limit: int = 50, offset: int = 0
@@ -1690,7 +1724,8 @@ class PromptDatabase(
""",
(int(limit), max(0, int(offset))),
)
return [self._image_row_to_dict(row) for row in cursor.fetchall()]
reader = DocumentReader(conn)
return [self._image_row_to_dict(row, reader) for row in cursor]
def get_all_images(
self,
@@ -1724,10 +1759,10 @@ class PromptDatabase(
with self.model.get_connection() as conn:
cursor = conn.execute(sql, params)
rows = cursor.fetchall()
reader = DocumentReader(conn)
result = []
for row in rows:
data = self._image_row_to_dict(row)
for row in cursor:
data = self._image_row_to_dict(row, reader)
if row["_prompt_tags_list"]:
data["prompt_tags"] = row["_prompt_tags_list"].split("|||")
else:
@@ -1765,7 +1800,8 @@ class PromptDatabase(
max(0, int(offset)),
),
)
return [self._image_row_to_dict(row) for row in cursor.fetchall()]
reader = DocumentReader(conn)
return [self._image_row_to_dict(row, reader) for row in cursor]
def get_image_by_id(self, image_id: int) -> Optional[Dict[str, Any]]:
"""
@@ -1782,7 +1818,7 @@ class PromptDatabase(
"SELECT * FROM generated_images WHERE id = ?", (image_id,)
)
row = cursor.fetchone()
return self._image_row_to_dict(row) if row else None
return self._image_row_to_dict(row, DocumentReader(conn)) if row else None
def delete_image(self, image_id: int) -> bool:
"""
@@ -1795,9 +1831,11 @@ 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
@@ -1831,12 +1869,16 @@ class PromptDatabase(
cursor = conn.execute("SELECT id, image_path FROM generated_images")
images = cursor.fetchall()
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
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)
conn.commit()
@@ -2127,11 +2169,13 @@ 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
@@ -2205,7 +2249,8 @@ 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.parameters, gi.id, gi.image_path, gi.file_path,
gi.workflow_document_id, gi.prompt_document_id
FROM generated_images gi
JOIN prompts p ON gi.prompt_id = p.id
WHERE gi.image_path = ? OR gi.file_path = ?
@@ -2285,12 +2330,21 @@ class PromptDatabase(
"text_source": source,
}
tags = self._get_prompt_tags(conn, row[0])
texts = resolve_image_row(DocumentReader(conn), row)
try:
workflow = json.loads(row[5]) if row[5] else None
workflow = (
json.loads(texts["workflow_data"])
if texts["workflow_data"]
else None
)
except (json.JSONDecodeError, TypeError):
workflow = None
try:
prompt_meta = json.loads(row[6]) if row[6] else None
prompt_meta = (
json.loads(texts["prompt_metadata"])
if texts["prompt_metadata"]
else None
)
except (json.JSONDecodeError, TypeError):
prompt_meta = None
return {
@@ -2540,17 +2594,20 @@ class PromptDatabase(
with self.model.get_connection() as conn:
return list(iter_consistency_issues(conn))
def _image_row_to_dict(self, row: sqlite3.Row) -> Dict[str, Any]:
def _image_row_to_dict(
self, row: sqlite3.Row, reader: DocumentReader
) -> 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 = dict(row)
data = resolve_image_row(reader, row)
# Parse JSON fields
for field in ["workflow_data", "prompt_metadata", "parameters"]:
+35 -3
View File
@@ -5,6 +5,13 @@ 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
@@ -116,25 +123,35 @@ 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 existing[key] != value
key: value for key, value in snapshots.items() if current[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:
columns = ["prompt_id", "image_path", "filename", "file_path", *snapshots]
stored = _document_form(conn, snapshots)
columns = ["prompt_id", "image_path", "filename", "file_path", *stored]
values = [
pair["positive_prompt_id"],
image_path,
os.path.basename(image_path),
normalized,
*snapshots.values(),
*stored.values(),
]
image_id = conn.execute(
"INSERT INTO generated_images("
@@ -162,3 +179,18 @@ 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)}
+477
View File
@@ -0,0 +1,477 @@
"""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
+455
View File
@@ -0,0 +1,455 @@
"""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."
)
+6
View File
@@ -9,11 +9,13 @@ 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
- MaintenanceRoutesMixin: reviewed maintenance plans (preview, apply)
- WorkflowStorageRoutesMixin: store each workflow once, then Reclaim space
- DuplicateFilesMixin: duplicate image files as find and removal jobs
"""
@@ -47,8 +49,10 @@ 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 .maintenance import MaintenanceRoutesMixin
from .workflow_storage import WorkflowStorageRoutesMixin
from .duplicate_files import DuplicateFilesMixin
from .maintenance_plans import PlanStore
from .jobs import JobRegistry
@@ -289,8 +293,10 @@ class PromptManagerAPI(
JobRoutesMixin,
ImageTagJobsMixin,
ThumbnailJobsMixin,
AutoThumbnailMixin,
SystemRoutesMixin,
MaintenanceRoutesMixin,
WorkflowStorageRoutesMixin,
DuplicateFilesMixin,
LoggingRoutesMixin,
AutotagRoutesMixin,
+3
View File
@@ -1504,6 +1504,9 @@ 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."""
+42 -20
View File
@@ -734,7 +734,10 @@ class HistoryRoutesMixin:
raise HistoryReadError(422, "response-limit") from None
else:
prepared = _prepare_current_media(
resolved, scope, thumbnail=kind == "thumbnail"
resolved,
scope,
thumbnail=kind == "thumbnail",
**media_options,
)
_check_current(scope, self, fence, epoch, generation)
return prepared
@@ -761,25 +764,44 @@ class HistoryRoutesMixin:
_check_current(scope, self, fence, epoch, generation)
raise
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
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()
if isinstance(prepared, bytes):
return _success_response(prepared)
return await _current_media_response(request, prepared)
+33 -7
View File
@@ -5,6 +5,7 @@ from dataclasses import dataclass
import errno
import hashlib
import hmac
import io
import json
import mimetypes
import os
@@ -309,17 +310,35 @@ 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
resolved: ResolvedCurrentCandidate,
scope: CurrentOutputScope,
*,
thumbnail: bool,
ensure_preview=None,
placeholder: bytes | None = None,
) -> _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)
extension = (
".jpg"
if target.suffix.lower() in VIDEO_EXTENSIONS
else target.suffix.lower()
)
is_video = target.suffix.lower() in VIDEO_EXTENSIONS
extension = ".jpg" if is_video else target.suffix.lower()
preview = (
resolved.gallery_root
/ "thumbnails"
@@ -328,7 +347,14 @@ def _prepare_current_media(
try:
preview.stat()
except FileNotFoundError:
pass
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"
)
else:
target = _canonical(preview)
handle = _open_current_target(target, resolved.namespace, scope)
+259 -137
View File
@@ -110,6 +110,8 @@ 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):
@@ -261,13 +263,32 @@ def _temp_sibling(dst):
return dst.with_name(f".{dst.stem}.part-{token}{dst.suffix}")
def _atomic_save(img, dst, fmt=None, **save_kwargs):
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):
"""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."""
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.
"""
tmp = _temp_sibling(dst)
try:
img.save(tmp, fmt, **save_kwargs)
os.replace(tmp, dst)
_replace_unless_held(tmp, dst, guard)
finally:
if tmp.exists():
tmp.unlink()
@@ -305,11 +326,11 @@ def _ffmpeg_qscale(encoding):
return str(_FFMPEG_QSCALE.get(encoding["quality"], 4))
def _write_image_thumbnail(src, dst, thumbnail_size, encoding=None):
def _write_image_thumbnail(src, dst, thumbnail_size, encoding=None, guard=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.
preview keeps only its ``optimize`` flag. ``guard`` see _atomic_save.
"""
encoding = encoding or DEFAULT_THUMBNAIL_ENCODING
with Image.open(src) as img:
@@ -319,7 +340,50 @@ def _write_image_thumbnail(src, dst, thumbnail_size, encoding=None):
save_kwargs = dict(encoding)
if dst.suffix.lower() == ".png":
save_kwargs = {"optimize": encoding["optimize"]}
_atomic_save(img, dst, **save_kwargs)
_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
def _new_thumbnail_stats():
@@ -710,23 +774,51 @@ def _iter_media_files(output_path, max_files):
yield root_path / name
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)."""
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.
"""
rel_path = Path(rel_path)
extension = rel_path.suffix.lower()
thumb_rel = (
return (
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:
return None
thumbnail_url += f"&root={root_index}"
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}"
if root_scope is not None:
thumbnail_url += f"&root_scope={root_scope}"
return thumbnail_url
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]
def _output_image_entry(
@@ -744,7 +836,7 @@ def _output_image_entry(
extension = media_path.suffix.lower()
is_video = extension in VIDEO_EXTENSIONS
thumbnail_url = media_thumbnail_url(
thumbnail_url, thumbnail_ready = media_thumbnail(
output_path, rel_path, root_index, root_scope=root_scope
)
encoded_path = urllib.parse.quote(rel_path.as_posix(), safe="/")
@@ -763,6 +855,8 @@ 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,
@@ -1682,6 +1776,11 @@ 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(
@@ -1909,6 +2008,8 @@ 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
],
@@ -2003,6 +2104,8 @@ class ImageRoutesMixin:
encoding=None,
rebuild=False,
stop=None,
guard=None,
track=None,
):
"""Create one thumbnail (blocking). Never raises; returns a result dict.
@@ -2010,7 +2113,11 @@ 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.
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).
"""
result = {
"file": src.name,
@@ -2027,22 +2134,43 @@ 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
src,
dst,
thumbnail_size,
encoding,
keep_existing=rebuild,
guard=guard,
track=track,
)
if not made:
return {**result, "action": "error", "error": "Video decode failed"}
return {
**result,
"action": "error",
"error": "Video decode failed",
"decode": True,
}
else:
_write_image_thumbnail(src, dst, thumbnail_size, encoding)
_write_image_thumbnail(src, dst, thumbnail_size, encoding, guard)
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"}
return {
**result,
"action": "error",
"error": "Not a valid image file",
"decode": True,
}
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)
return {**result, "action": "error", "error": _safe_error(e)}
failed = {**result, "action": "error", "error": _safe_error(e)}
if _is_decode_error(e):
failed["decode"] = True
return failed
async def _emit_sse(self, response, event, data) -> bool:
"""Write one SSE frame. Returns False once the client has gone away."""
@@ -2061,147 +2189,59 @@ class ImageRoutesMixin:
thumbnail_size,
encoding=None,
keep_existing=False,
*,
guard=None,
track=None,
):
"""Generate thumbnail from video file. Returns True if successful.
``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).
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 = encoding or DEFAULT_THUMBNAIL_ENCODING
try:
# 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)
frame = self._opencv_frame(video_path, thumbnail_size)
if frame is not None:
_atomic_save(frame, thumbnail_path, "JPEG", guard=guard, **encoding)
self.logger.debug(
f"Generated video thumbnail using OpenCV: {thumbnail_path}"
)
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 self._ffmpeg_frame(
video_path, thumbnail_path, thumbnail_size, encoding, guard, track
):
return True
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: create a placeholder thumbnail
# Last resort: the placeholder frame.
try:
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)
img = _video_placeholder_image(thumbnail_size)
_atomic_save(img, thumbnail_path, "JPEG", guard=guard, **encoding)
self.logger.debug(
f"Generated placeholder video thumbnail: {thumbnail_path}"
)
return True
except PermissionError:
except (PermissionError, PreviewWriteHeld):
raise
except Exception as e:
self.logger.warning(
f"Failed to create placeholder thumbnail for {video_path}: {e}"
)
return False
except PermissionError:
except (PermissionError, PreviewWriteHeld):
raise
except Exception as e:
self.logger.error(
@@ -2209,6 +2249,88 @@ 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.
@@ -2218,7 +2340,7 @@ class ImageRoutesMixin:
"""
return await self._run_held(
THUMBNAILS_JOB_KIND,
"Thumbnail clear (gallery page)",
THUMBNAIL_CLEAR_TITLE,
"guarded",
self._clear_thumbnails,
)
+6
View File
@@ -515,6 +515,12 @@ 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()):
+66 -24
View File
@@ -58,9 +58,11 @@ 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
@@ -68,6 +70,7 @@ 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
@@ -77,6 +80,7 @@ MAINTENANCE_TITLES = {
"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")
@@ -89,12 +93,28 @@ PLAN_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."""
"""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():
@@ -160,7 +180,7 @@ def _megabytes(size):
return f"{size / 2**20:,.1f} MB"
def disk_plan(db_path, *, backup, vacuum):
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
@@ -168,7 +188,9 @@ def disk_plan(db_path, *, backup, vacuum):
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. Returns ``(public, short)``: ``public`` is
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
@@ -181,8 +203,10 @@ def disk_plan(db_path, *, backup, vacuum):
if backup:
needs.append(("database folder", folder, size, "the safety backup"))
if vacuum:
needs.append(("database folder", folder, size, "Optimize"))
needs.append(("temporary folder", sqlite_temp_dir(), size, "Optimize"))
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})
@@ -290,9 +314,9 @@ def _remove_quietly(*paths):
pass
def _reserve_backup_name(db_path):
"""Create and return an unused ``<db>.backup_<stamp>[_n]`` (O_EXCL)."""
for candidate in _backup_names(db_path):
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:
@@ -302,13 +326,13 @@ def _reserve_backup_name(db_path):
raise SafetyBackupFailed("No free backup name")
def _backup_names(db_path):
base = f"{db_path}.backup_{_now().strftime('%Y%m%d_%H%M%S_%f')}"
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):
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
@@ -316,7 +340,7 @@ def _place_backup(temp, db_path):
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):
for candidate in _backup_names(db_path, label):
try:
os.link(temp, candidate)
except FileExistsError:
@@ -327,7 +351,7 @@ def _place_backup(temp, db_path):
return candidate
else:
raise SafetyBackupFailed("No free backup name")
final = _reserve_backup_name(db_path)
final = _reserve_backup_name(db_path, label)
try:
os.replace(temp, final)
except OSError:
@@ -381,12 +405,15 @@ def _fsync(path):
os.fsync(handle.fileno())
def take_safety_backup(model):
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 once it verified, so a backup file on disk
is always a verified one. Raises SafetyBackupFailed, leaving no file.
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(
@@ -398,13 +425,13 @@ def take_safety_backup(model):
if not model.backup_database(temp):
raise SafetyBackupFailed("The backup could not be written")
_standalone(temp)
ok, _reason = model.verify_database_file(temp)
ok, reason = model.verify_database_file(temp)
if not ok:
raise SafetyBackupFailed("The backup did not verify")
raise SafetyBackupFailed("The backup did not verify", reason)
_fsync(temp)
if os.name != "nt":
os.chmod(temp, 0o600)
final = _place_backup(temp, db_path)
final = _place_backup(temp, db_path, label)
return os.path.basename(final)
except OSError as exc:
_remove_quietly(final)
@@ -519,6 +546,7 @@ class MaintenanceRoutesMixin:
"targets": draft.targets,
"roots": roots,
"protected": draft.protected["count"],
"rows": draft.scope["count"],
},
)
@@ -538,11 +566,18 @@ class MaintenanceRoutesMixin:
"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 == "on",
backup=policy != "off",
vacuum=operation == "optimize",
)
except (OSError, sqlite3.Error):
@@ -567,8 +602,8 @@ class MaintenanceRoutesMixin:
**public,
"backup": {
"policy": policy,
"destination_label": f"{name}.backup_<date and time>",
"expected_bytes": size if policy == "on" else None,
"destination_label": f"{name}.{destination}_<date and time>",
"expected_bytes": size if policy != "off" else None,
},
"disk": disk,
},
@@ -596,6 +631,9 @@ class MaintenanceRoutesMixin:
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(
@@ -618,6 +656,8 @@ class MaintenanceRoutesMixin:
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,
@@ -663,6 +703,8 @@ class MaintenanceRoutesMixin:
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):
@@ -679,9 +721,9 @@ class MaintenanceRoutesMixin:
job.log_line("Taking a safety backup of the database")
try:
backup = await self._run_in_executor(take_safety_backup, db.model)
except SafetyBackupFailed:
except SafetyBackupFailed as error:
self.logger.error("Safety backup failed", exc_info=True)
_fail_job(job, "backup-failed", BACKUP_FAILED)
_fail_job(job, "backup-failed", backup_failed_message(error))
return False, None
job.log_line(f"Safety backup {backup} verified")
return True, backup
+843
View File
@@ -0,0 +1,843 @@
"""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()
+277
View File
@@ -0,0 +1,277 @@
"""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)
+2
View File
@@ -386,6 +386,8 @@ 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):
+39 -8
View File
@@ -398,18 +398,20 @@ 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();
expect(media.images.find(image => image.filename === 'clip.webm').thumbnail_url).toBeNull();
// #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/);
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();
// 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);
// 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);
await expect(card.locator('.media-info p')).toHaveText(/(?:^|· )WEBM$/);
expect(await card.locator('video').evaluate(node => node.getBoundingClientRect().height)).toBeGreaterThan(20);
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);
await card.getByRole('button', { name: 'Open clip.webm', exact: true }).click();
const video = page.locator('dialog[open] video');
await expect(video).toBeVisible();
@@ -713,16 +715,45 @@ test('selected layout stays blue with readable text on hover in either theme', a
}
});
test('refresh releases the previous inspector video source', async ({ page, request }) => {
test('refresh releases the previous inspector cover 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)).toBe(true);
expect(await page.evaluate(() => window.__oldInspectorVideo.getAttribute('src'))).toBeNull();
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]);
});
test('Import Rescan rescans cached files and preserves the collection when rescan fails', async ({ page, request }) => {
+135
View File
@@ -0,0 +1,135 @@
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);
});
+1
View File
@@ -39,3 +39,4 @@ 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,3 +895,30 @@ 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 });
});
@@ -253,3 +253,11 @@ 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');
});
+2
View File
@@ -79,6 +79,8 @@ 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():
+2
View File
@@ -131,6 +131,8 @@ 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()
+2
View File
@@ -62,6 +62,8 @@ 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):
+2
View File
@@ -384,6 +384,8 @@ 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)
+10 -1
View File
@@ -654,11 +654,20 @@ class TestCleanupJob(DuplicateFilesTestCase):
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"))
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",
+75 -3
View File
@@ -401,7 +401,7 @@ class HistoryOutputAPITests(ImageAPITestCase):
self.assertEqual(response.status, 400)
await response.read()
async def test_media_ranges_head_and_existing_thumbnail_without_generation(self):
async def test_media_ranges_head_and_a_preview_made_on_first_view(self):
path = self.output_dir / "clip.mp4"
payload = bytes(range(256)) * 600
path.write_bytes(payload)
@@ -430,9 +430,14 @@ 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(await response.read(), payload)
preview = make_png(self.output_dir / "thumbnails" / "clip_thumb.jpg")
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)
response = await self.current(descriptor["thumbnail_url"])
self.assertEqual(await response.read(), preview.read_bytes())
@@ -1111,6 +1116,73 @@ 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,
):
+69 -16
View File
@@ -80,6 +80,8 @@ 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()
@@ -1565,39 +1567,65 @@ class TestVideoThumbnailFallbacks(ImageRouteCoverageCase):
with Image.open(self.thumb) as img:
self.assertEqual(img.format, "JPEG")
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_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_crash_is_contained(self):
cv2 = MagicMock()
cv2.VideoCapture.side_effect = RuntimeError("driver crash")
with patch.dict(sys.modules, {"cv2": cv2}):
with (
patch.dict(sys.modules, {"cv2": cv2}),
patch("subprocess.Popen", side_effect=FileNotFoundError("ffmpeg")),
):
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.run", side_effect=FileNotFoundError("ffmpeg")),
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.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.run", return_value=failed),
patch("subprocess.Popen", return_value=self._process(1, "no stream")),
):
ok = self.api._generate_video_thumbnail(self.video, self.thumb, self.SIZE)
self.assertTrue(ok)
@@ -1606,16 +1634,41 @@ class TestVideoThumbnailFallbacks(ImageRouteCoverageCase):
def test_ffmpeg_success(self):
with (
patch.dict(sys.modules, {"cv2": None}),
patch("subprocess.run", return_value=MagicMock(returncode=0)) as run,
patch("subprocess.Popen", return_value=self._process(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.run", side_effect=FileNotFoundError),
patch("subprocess.Popen", side_effect=FileNotFoundError),
patch.object(
images_module.Image, "new", side_effect=RuntimeError("no PIL")
),
+8 -3
View File
@@ -397,15 +397,20 @@ class TestQueueReads(ReviewRoutesTestCase):
self._no_paths()
async def test_queue_cards_use_the_images_listing_thumbnail(self):
# M10: a known thumbnail is the card's media, as on the Images page;
# without one the card falls back to the original.
# 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).
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]
self.assertIsNone(item["thumbnail_url"])
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}",
)
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))
+424
View File
@@ -0,0 +1,424 @@
"""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)",
)
+7 -2
View File
@@ -99,11 +99,16 @@ 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 [
dict(row)
resolve_image_row(reader, row)
for row in conn.execute(
"SELECT prompt_id, workflow_data, prompt_metadata, parameters "
"SELECT prompt_id, workflow_data, prompt_metadata, parameters, "
"workflow_document_id, prompt_document_id "
"FROM generated_images ORDER BY id"
)
]
File diff suppressed because it is too large Load Diff
+9 -6
View File
@@ -324,7 +324,7 @@ class TestThumbnailJob(ThumbnailJobCase):
old = preview.read_bytes()
with (
patch.dict(sys.modules, {"cv2": None}),
patch("subprocess.run", side_effect=FileNotFoundError("no ffmpeg")),
patch("subprocess.Popen", 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.run", side_effect=FileNotFoundError("no ffmpeg")),
patch("subprocess.Popen", side_effect=FileNotFoundError("no ffmpeg")),
patch.object(images_module, "_atomic_save", refuse),
):
_, body = await self._start()
@@ -594,12 +594,14 @@ class TestThumbnailJob(ThumbnailJobCase):
def fake_ffmpeg(cmd, **_kwargs):
commands.append(cmd)
Path(cmd[-1]).write_bytes(b"jpeg")
return type("Done", (), {"returncode": 0, "stderr": ""})()
done = type("Done", (), {"returncode": 0})()
done.communicate = lambda timeout=None: ("", "")
return done
for quality, qscale in (("high", "2"), ("medium", "4"), ("low", "8")):
with (
patch.dict(sys.modules, {"cv2": None}),
patch("subprocess.run", side_effect=fake_ffmpeg),
patch("subprocess.Popen", side_effect=fake_ffmpeg),
):
_, body = await self._start(mode="rebuild", size=640, quality=quality)
job = await self._finished(body["job"]["id"])
@@ -618,12 +620,13 @@ class TestThumbnailJob(ThumbnailJobCase):
real_save = images_module._atomic_save
def recording(img, dst, fmt=None, **kwargs):
saved.append(kwargs)
# The encoder settings (guard is the automatic previews' hold check).
saved.append({k: v for k, v in kwargs.items() if k != "guard"})
return real_save(img, dst, fmt, **kwargs)
with (
patch.dict(sys.modules, {"cv2": None}),
patch("subprocess.run", side_effect=FileNotFoundError("no ffmpeg")),
patch("subprocess.Popen", side_effect=FileNotFoundError("no ffmpeg")),
patch.object(images_module, "_atomic_save", recording),
):
_, body = await self._start(quality="low")
+817
View File
@@ -0,0 +1,817 @@
"""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()
+469
View File
@@ -0,0 +1,469 @@
"""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()
+267
View File
@@ -0,0 +1,267 @@
"""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)
+374
View File
@@ -0,0 +1,374 @@
"""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()
+94
View File
@@ -0,0 +1,94 @@
"""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())
+7 -1
View File
@@ -38,6 +38,11 @@ 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.
@@ -221,7 +226,8 @@ class GalleryDiagnostics:
ORDER BY gi.generation_time DESC
LIMIT 5
""")
recent_images = [dict(row) for row in cursor.fetchall()]
reader = DocumentReader(conn)
recent_images = [resolve_image_row(reader, row) for row in cursor]
self.logger.info(f" [TIME] Recent images: {len(recent_images)}")
for img in recent_images:
+35
View File
@@ -79,6 +79,10 @@ 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:
@@ -220,10 +224,30 @@ 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.
@@ -636,6 +660,16 @@ 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).
@@ -703,6 +737,7 @@ 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:
+4 -2
View File
@@ -488,7 +488,8 @@
const index = from + offset;
// Use thumbnail for display if available, fallback to original
const displayUrl = image.thumbnail_url || image.url;
const hasThumb = !!image.thumbnail_url;
// The listing names a preview URL for every file (made on view, #249); the dot means the file exists.
const hasThumb = typeof image.thumbnail_ready === 'boolean' ? image.thumbnail_ready : !!image.thumbnail_url;
const isVideo = image.is_video || false;
const mediaType = image.media_type || 'image';
@@ -551,7 +552,8 @@
const index = from + offset;
// Use thumbnail for display if available, fallback to original
const displayUrl = image.thumbnail_url || image.url;
const hasThumb = !!image.thumbnail_url;
// The listing names a preview URL for every file (made on view, #249); the dot means the file exists.
const hasThumb = typeof image.thumbnail_ready === 'boolean' ? image.thumbnail_ready : !!image.thumbnail_url;
const isVideo = image.is_video || false;
const mediaType = image.media_type || 'image';
+18 -11
View File
@@ -318,18 +318,25 @@ export function mountFileInspector(container, {api, sender, parentScope, variant
});
return button;
}
/** The concept's plain cover; the card image opens the viewer (D19). Video keeps its own element. */
/**
* The concept's plain cover; the card image opens the viewer (D19). #249: it is the file's preview (a video's
* is a JPEG frame) whenever the listing gave one; without one the original, and a video keeps its own element.
*/
function cover(record) {
const {file} = record, video = file?.media_type === 'video';
const node = el(video ? 'video' : 'img', undefined, 'detail-cover'); node.src = record.download.url;
if (video) {
node.muted = true; node.playsInline = true; node.preload = 'metadata';
contentScope.cleanup(() => {node.pause(); node.removeAttribute('src'); node.load();});
} else {
node.alt = ''; node.decoding = 'async';
if (file?.width && file?.height) {node.width = file.width; node.height = file.height;}
contentScope.cleanup(() => node.removeAttribute('src'));
}
const {file} = record, preview = record.preview?.url, isVideo = file?.media_type === 'video';
if (isVideo && !preview) return videoCover(record);
const node = el('img', undefined, 'detail-cover'); node.src = preview || record.download.url;
node.alt = ''; node.decoding = 'async';
if (file?.width && file?.height) {node.width = file.width; node.height = file.height;}
contentScope.cleanup(() => node.removeAttribute('src'));
// #249: a video whose preview cannot be loaded swaps back to its own element.
if (isVideo) contentScope.listen(node, 'error', () => { if (node.isConnected) node.replaceWith(videoCover(record)); }, {once:true});
return node;
}
function videoCover(record) {
const node = el('video', undefined, 'detail-cover'); node.src = record.download.url;
node.muted = true; node.playsInline = true; node.preload = 'metadata';
contentScope.cleanup(() => {node.pause(); node.removeAttribute('src'); node.load();});
return node;
}
/** "File & provenance": the file facts, the D21 prompt evidence, the occurrence and the saved association (D66 for a local file). */
+4 -1
View File
@@ -64,7 +64,10 @@ export function toHistoryInspectionRecord(result, target, extractFacts) {
const path = [target.reference.subfolder, target.reference.filename].filter(Boolean).join('/');
const identity = {kind:'reported-location', key:target.location_key, path:target.reference.path_style === 'windows' ? path.replaceAll('\\', '/') : path,
occurrence:{...target.occurrence}, history_epoch:target.history_epoch, current_scope:target.current_scope, reference:{...target.reference}};
return projectInspectionPayload(result, identity, {url:urlFor(target, 'url', 'media'), filename:target.filename}, extractFacts);
const record = projectInspectionPayload(result, identity, {url:urlFor(target, 'url', 'media'), filename:target.filename}, extractFacts);
// #249: the inspector's cover is the recorded output's preview (made by the server on first view).
record.preview = {url:urlFor(target, 'thumbnail_url', 'thumbnail')};
return record;
}
export const reportedLocationAdapter = Object.freeze({
validate,
+20 -10
View File
@@ -17,19 +17,29 @@ function el(tag, text, className) {
return node;
}
/** The card's media: a video's first frame when it has no thumbnail, else the thumbnail or the image. */
/** A video's own element (its first frame), for a video without a usable preview. */
function videoMedia(image) {
const node = el('video');
node.src = window.ApiClient.resolveUrl(image.url);
node.muted = true; node.playsInline = true; node.preload = 'metadata';
return node;
}
/**
* The card's media: the thumbnail (a video's is its preview frame), else the image, else a video's own element.
* #249: a video whose preview cannot be loaded swaps back to its own element.
*/
export function cardMedia(image, lazy = false) {
const { resolveUrl } = window.ApiClient;
const video = image.is_video && !image.thumbnail_url;
const node = el(video ? 'video' : 'img');
if (image.is_video && !image.thumbnail_url) return videoMedia(image);
const node = el('img');
node.src = resolveUrl(image.thumbnail_url || image.url);
if (video) { node.muted = true; node.playsInline = true; node.preload = 'metadata'; }
else {
// The concept card image: alt is the card title, with the listing's intrinsic size when known.
node.alt = cardTitle(image); node.decoding = 'async'; if (lazy) node.loading = 'lazy';
if ([image.width, image.height].every(value => Number.isFinite(value) && value > 0)) {
node.width = image.width; node.height = image.height;
}
// replaceWith does nothing for a card no longer in the page.
if (image.is_video) node.addEventListener('error', () => node.replaceWith(videoMedia(image)), { once: true });
// The concept card image: alt is the card title, with the listing's intrinsic size when known.
node.alt = cardTitle(image); node.decoding = 'async'; if (lazy) node.loading = 'lazy';
if ([image.width, image.height].every(value => Number.isFinite(value) && value > 0)) {
node.width = image.width; node.height = image.height;
}
return node;
}
+41 -3
View File
@@ -44,6 +44,9 @@ export function inspectPath(image) {
return `/prompt_manager/images/inspect/${image.path.split('/').map(encodeURIComponent).join('/')}?${query}`;
}
/** One path segment as the server quotes it (Python `quote(..., safe="/")` also escapes `!'()*`, #249 M11). */
const quoteSegment = segment => encodeURIComponent(segment).replace(/[!'()*]/g, c => `%${c.charCodeAt(0).toString(16).toUpperCase()}`);
function downloadFor(image, required) {
try {
if (typeof image.url !== 'string' || !image.url) throw new TypeError('Missing original media URL');
@@ -55,7 +58,7 @@ function downloadFor(image, required) {
const sameOrigin = 'Origin' in globalThis ?
globalThis.Origin.from(url).isSameOrigin(globalThis.Origin.from(reference)) :
url.origin !== 'null' && url.origin === reference.origin;
const expected = new URL(resolveUrl(`/prompt_manager/images/serve/${image.path.split('/').map(encodeURIComponent).join('/')}`), reference);
const expected = new URL(resolveUrl(`/prompt_manager/images/serve/${image.path.split('/').map(quoteSegment).join('/')}`), reference);
const scope = descriptorScope(image);
if (!sameOrigin || !['http:', 'https:'].includes(url.protocol) || url.username || url.password || url.hash ||
url.pathname !== expected.pathname || url.searchParams.getAll('root').length !== 1 || url.searchParams.get('root') !== String(image.root_index) ||
@@ -68,6 +71,39 @@ function downloadFor(image, required) {
}
}
// The server's video extensions (py/api/inspection_service.py VIDEO_EXTENSIONS): their preview is a JPEG.
const VIDEO_EXTENSIONS = new Set(['.mp4', '.webm', '.avi', '.mov', '.mkv', '.m4v', '.wmv']);
/** `thumbnails/<path>_thumb<ext>` for a listed file, as the server names its preview (images.media_thumbnail_rel). */
function previewPath(path) {
const slash = path.lastIndexOf('/'), dot = path.lastIndexOf('.');
const cut = dot > slash ? dot : path.length, extension = path.slice(cut).toLowerCase();
return `thumbnails/${path.slice(0, cut)}_thumb${VIDEO_EXTENSIONS.has(extension) ? '.jpg' : extension}`;
}
/**
* #249: the listing's thumbnail of `image` as `{url}` for the inspector's cover, or null. Only the same-origin
* serve URL of this file's own preview under the same root and source scope is used; anything else keeps the
* original (a missing preview is made by the server on first view).
*/
function previewFor(image) {
try {
if (typeof image.thumbnail_url !== 'string' || !image.thumbnail_url) return null;
const location = globalThis.location || globalThis.window?.location;
const reference = new URL(location?.href || 'http://localhost/');
const resolveUrl = globalThis.window?.ApiClient?.resolveUrl || globalThis.ApiClient?.resolveUrl || (url => url);
const value = resolveUrl(image.thumbnail_url), url = new URL(value, reference);
const expected = new URL(resolveUrl(`/prompt_manager/images/serve/${previewPath(image.path).split('/').map(quoteSegment).join('/')}`), reference);
const scope = descriptorScope(image);
if (url.origin !== reference.origin || url.username || url.password || url.hash || url.pathname !== expected.pathname ||
url.searchParams.getAll('root').length !== 1 || url.searchParams.get('root') !== String(image.root_index) ||
url.searchParams.getAll('root_scope').length > 1 || (scope !== null && url.searchParams.get('root_scope') !== scope)) return null;
return { url: value };
} catch {
return null;
}
}
export function originalMediaDownload(image) {
validateIdentity(image);
return downloadFor(image, true);
@@ -284,16 +320,18 @@ export function toInspectionRecord(result, image, extractFacts = globalThis.Comf
if (!legacy && result.inspection_version !== 1) throw new TypeError('Unsupported inspection version');
if (!legacy && result.success !== true) throw new TypeError('Invalid inspection success status');
if ((!legacy || result.image !== undefined) && (!object(result.image) || result.image.root_index !== image.root_index || result.image.path !== image.path || result.image.filename !== image.path.split('/').at(-1))) throw new TypeError('Mismatched inspection identity');
return projectInspectionPayload(result,
const record = projectInspectionPayload(result,
{ kind: 'rooted-file', key: imageKey(image), root_index: image.root_index, path: image.path, root_scope: scope },
downloadFor(image, !legacy || scope !== null), extractFacts, legacy, parser);
record.preview = record.download ? previewFor(image) : null;
return record;
}
/** Shared bounded payload projection; adapters validate identity before entering here. */
export function projectInspectionPayload(result, identity, download, extractFacts = globalThis.ComfyMetadata?.extractGenerationFacts || globalThis.window?.ComfyMetadata?.extractGenerationFacts, legacy = false, parser = comfyMetadata()) {
const record = {
identity, file: null, prompts: promptsFor(result.prompts), library: null, metadata: null,
generation: unavailableFacts(), historyTarget: null, download, imageTags: null,
generation: unavailableFacts(), historyTarget: null, download, imageTags: null, preview: null,
};
if (legacy) return record;
record.imageTags = imageTagsFor(result.image_tags);
+3 -1
View File
@@ -88,7 +88,9 @@ export function thumbnailTitle(status, sources) {
const missing = status?.pending;
if (!Number.isSafeInteger(missing) || sources?.effective === 0) return 'Thumbnail status unavailable';
if (missing === 0) return 'No missing thumbnails';
return `${number(missing)} missing ${noun(missing, 'thumbnail')}`;
// #249: files whose automatic preview could not be made this session (they show the original).
const failed = Number.isSafeInteger(status.failed) && status.failed > 0 ? ` · ${number(status.failed)} failed` : '';
return `${number(missing)} missing ${noun(missing, 'thumbnail')}${failed}`;
}
const SOURCES_SHARED = 'Monitoring and enrichment share one source list.';