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.
This commit is contained in:
Vito
2026-10-24 05:28:44 -07:00
committed by GitHub
parent e34bd70e13
commit f68e92b3b6
23 changed files with 7110 additions and 264 deletions
+204
View File
@@ -2499,11 +2499,215 @@ node --test tests/js/workspace-system-state.test.js tests/js/workspace-system-co
tests/js/workspace-inspection-state.test.js tests/js/workspace-images-state.test.js \
tests/js/workspace-system-sources.test.js tests/js/workspace-thumbnail-review.test.js
python3 -m unittest tests.test_settings_sources tests.test_scan_scope -v
python3 -m unittest tests.test_maintenance_plans tests.test_api_maintenance -v
npm run --prefix tools/vendor test:workflows -- system-concept.spec.mjs workspace-system.spec.mjs \
workspace-system-sources.spec.mjs \
--project chromium --project touch --project firefox --project webkit
```
### Maintenance contract (#249 System S6)
`POST /prompt_manager/maintenance/preview {operation, options:{}}` reviews one
operation and answers a `Plan` (`py/api/maintenance.py`, `database/maintenance.py`):
`duplicate_prompts` (legacy `cleanup_duplicates`: equal text ignoring case and outer
spaces, the oldest prompt of each group kept; a group with a member in a saved pair is
kept whole), `optimize` (VACUUM, with the database size, its free pages and the free
space it needs), `unused_components` (legacy `prune_orphaned_prompts`; the #239
protections are the predicates `ORPHANED_PROMPTS_SQL` is built from, so the plan's
protected reasons are exactly the prune's rules) and `missing_links` (legacy
`cleanup_orphaned_images`, judged as the health check's `missing_files` through
`health.link_states`: only a link under a source that can be read now counts; one under
an unavailable source or outside every source is protected). The operation enum is
strict, `options` must be `{}` and the body is capped at 4 KiB (512 KiB for
`duplicate_files`: 1 MiB, below). A preview meets the
read-only standard above on its own `mode=ro` snapshot (`test_preview_writes_nothing`,
per operation, plus the media bytes) and stats files after the snapshot ended. The
plan lists at most 50 detail rows and no server path. Its fingerprint (the ids, text
hashes, metadata, tags, Library state and pair refs of every prompt involved, or each
link's row and source), its targets and the sources it judged stay in the server's
`PlanStore` under a single-use `secrets.token_urlsafe(32)` token: 15 minutes, at most 8
plans and 16 MiB of payload. A plan over that budget is refused with `too-large`, never
truncated; a new plan drops the oldest ones, which then read as expired; a restart
forgets them all. A plan is bound to the database it read: a restore during the preview
answers `storage-changed`. Errors are coded (`invalid-request`, `too-large`,
`storage-changed`, or a null code with "Failed to review maintenance") and carry no
path or exception text.
`POST /prompt_manager/maintenance/apply {token}` starts the reviewed plan as an
Activity job (202 `{job}`): kind `maintenance`, titled after the concept's row
("Duplicate prompt text", "Optimize database", "Unused prompt components", "Missing-file
links"), admission `records` (it never runs beside a scan, restore, LoRA import or other
record-changing job, legacy-page holds included: `test_apply_refused_while_restore_runs`),
not cancellable, with the groups `changed`, `protected` and `errors`. The token is consumed
on the loop turn the job starts, so a second apply answers 409 `plan-used` with that job;
a refused start (`busy`) keeps it. Other errors: `invalid-request`, `plan-expired` (410:
unknown, expired, or forgotten after a restart) and `storage-changed` (409: a restore
replaced the database since the review; the plan is dropped). The job:
1. takes a safety backup when the plan was reviewed with the policy on or the policy is
on now (`maintenance.safety_backup` in `config.json`, default `true`; GET/POST
`/prompt_manager/settings` field `safety_backup` with the revision). The copy is made
with SQLite's backup API into a temp file beside the database, switched to a
rollback journal (so a later read-only open leaves no `-wal`/`-shm` beside it),
verified with `verify_database_file`, fsynced, made `0o600` on POSIX and only then
renamed with `os.replace` onto `<db>.backup_<YYYYmmdd_HHMMSS_ffffff>`, reserved with
`O_EXCL` (`_1`, `_2`… when a coarse clock repeats a stamp:
`test_backup_names_unique_within_one_second`). A backup file on disk is therefore
always a verified one. Any failure removes what it wrote and ends the job
`backup-failed` before anything is changed (`test_backup_failure_aborts_before_write`);
2. applies the plan: inside `BEGIN IMMEDIATE` the fingerprint is computed again on the
write connection and a difference ends the job `plan-changed` with nothing written
(`test_apply_rejects_changed_target`; a missing link whose file is back or whose source
went away is never removed). The targets are deleted in batches of 400 ids in that one
transaction. Missing-file links remove database rows only; every media file's bytes
are unchanged (`test_missing_links_never_deletes_files`). VACUUM runs outside a
transaction and binds to its database only;
3. counts the changed and protected records (outcome rows hold root-relative paths or
prompt ids only), logs the backup by its file name, and prunes safety backups to the
newest 10. With no backup the log and the status line say "No safety backup".
The legacy restore (`POST /prompt_manager/restore`) takes a backup of any size the disks
have room for, so a real library's safety backup restores: the upload, its staged copy
and the swap (the upload twice in the temporary folder, once more plus the current
database beside the database, added up on one disk, each with the 16 MiB margin) must
fit, under a 64 GiB ceiling (`restore_max_bytes`). A request whose Content-Length does
not fit is refused before anything is written, one without is refused as soon as its
running size does not, both with 507 `insufficient-space` and "Not enough free disk
space to restore a 1.6 GB backup; there is room for about …. Nothing was changed."
(no path); past the ceiling it is the legacy 400 "File too large" with code `too-large`.
The upload streams with `read_chunk`, which aiohttp's `client_max_size` (ComfyUI's
`--max-upload-size`, 100 MB by default) does not cap, and its temp file is closed and
removed on every path (`TestLargeRestores`, a 150 MB backup under a 100 MB app limit).
Safety backups and restore backups are one pool: every `<db>.backup_*` file beside the
database counts toward the same `MAX_SAFETY_BACKUPS` (10), whichever wrote it, and the
pool is pruned after every backup a job took, also when the apply then failed
(`test_safety_backups_share_the_restore_pool`,
`test_the_pool_is_pruned_when_the_apply_fails_after_the_backup`). The legacy
`GET /backup` download writes no file beside the database, so it adds nothing to the
pool. Each plan reports `backup.expected_bytes` (the backup's size, `page_count *
page_size`, or null with the policy off) and `disk: {free_bytes, required_bytes,
sufficient}` for the database folder: the backup needs its size there, and Optimize
needs the database's size twice (the WAL beside it and SQLite's temporary copy, added
up when the temporary folder is on the same disk), each plus a 16 MiB margin. The
temporary folder is the one SQLite really uses (`sqlite_temp_dir`): on Unix
`SQLITE_TMPDIR`, `TMPDIR`, then the first writable of `/var/tmp`, `/usr/tmp` and
`/tmp` (not Python's `tempfile` order, which picks a RAM-backed `/tmp` on many Linux
systems); on Windows `TMP`, `TEMP`, then `USERPROFILE`. Whether a backup is taken is
decided once, when Apply measured the disk, and handed to the job. Apply looks at the
token again with no await before it starts the job, so a token used or expired during
the measurement answers its code, never a 500. A backup is put in place with
`os.link` (atomic, refuses an existing name; reserve-and-replace where a filesystem
has no hard links), and startup removes the `.pm-safety-*.tmp` copies an interrupted
backup left. Optimize runs alone: its job has the `exclusive` admission class (refused
beside any job, and every job is refused while it runs), and its plan says the database
is locked while VACUUM runs and that run counts or captures waiting longer than the
30-second busy timeout are not recorded. The duplicate prompt merge is strict inside
the apply: a group whose copies link one image twice, or whose removed copy has its own
Library title, is protected in the plan, and any merge error rolls the transaction back
(`MergeConflict`); the legacy cleanup keeps its lenient behaviour. Apply
measures again before it starts a job and refuses a short disk up front with 507
`insufficient-space` and a truthful message ("Not enough free disk space: the safety
backup needs about … in the database folder and … is free. Nothing was changed."); no
job starts and the token stays usable
(`test_apply_refuses_up_front_when_the_backup_does_not_fit`,
`test_apply_refuses_up_front_when_vacuum_does_not_fit`).
**Duplicate image files** (`py/api/duplicate_files.py`, `py/api/duplicate_review.py`;
`tests/test_api_duplicate_files.py`). `POST /prompt_manager/jobs/duplicates` starts the
find: kind `duplicates`, "Find duplicate image files", admission `read`, cancellable,
group `groups_found`, over every effective source (the legacy scan's
`_hash_media_groups`: no symlinks, no `thumbnails`, bounded, each file once). The
find lists a physical file once, by `(st_dev, st_ino)`: a source listed twice under two
spellings, a symlinked source inside another one and a hardlink are one file, never a
copy of themselves (`test_a_folder_listed_twice_never_loses_its_only_copy`,
`test_a_symlinked_source_inside_another_is_one_file`,
`test_a_hardlink_is_never_a_copy_to_remove`). One file reached through two mounts is
protected, never guessed at: a removal on another `st_dev` than its kept copy, with the
same size and mtimes within 2 s, is kept as "may be the same file reached through another
mount or share" when its `st_ino` is equal or its path within its source is the same, or
one ends with the other on a whole folder name (sources listed at different depths:
`test_a_deeper_source_on_another_mount_is_protected`)
(NFS and SMB, a loopback mount, a FUSE pool: `TestAnotherMount`). Some true copies on
another disk are therefore never offered. A volume that reports no file id protects
every pair ("Can't confirm this is a separate file on this drive"). The job checks both
before each file's records move, holds the kept copy open, and right before the unlink
checks again that the held kept copy is still the file at its path (`fstat` of the held
handle against `os.stat` of the path) and that the guard still holds; otherwise the file
is kept, its records move back to it, and it counts as not removed
(`test_a_kept_copy_deleted_during_the_move_keeps_the_file`). On Windows the hold also
makes deleting the kept copy itself through another path fail. The server holds the groups (content
hash, root-relative paths, source index, size, mtime and inode) for the last three
finds; `GET /prompt_manager/maintenance/duplicates/{job_id}?offset&limit` pages
them ten at a time (`not-found` 404, `job-running` or `no-groups` 409). The preview
`{operation:"duplicate_files", options:{job_id, keep:[{group_id, root_index, path}]}}`
(body cap 1 MiB, at most 2,000 keep entries; `test_the_keep_limit_fits_the_body_cap`) removes, in each group named, every copy
not kept; a keeper that is not in its group is `invalid-request`, and a source no longer
listed at its index is `sources-changed`. Every keeper and removal is checked with
`file_problem`: containment as text first, with no filesystem access (inside its
source, not under `thumbnails`, a media extension), then `lstat` (a link or any
non-regular file is refused), the resolved path still inside the resolved source, not
the kept copy itself (`os.path.samestat`), and the size, mtime and inode the find saw.
The preview reads no file; the removal job hashes each file and its kept copy again
right before the unlink (`test_preview_reads_no_file_and_spots_a_touched_one`). A removal that fails, or whose kept copies all fail, is
protected with its reason. The plan names each removal, its kept copy and how many
database links move, says the safety backup covers the database and not the files,
and writes nothing. Apply runs `duplicate_cleanup` (records, not cancellable, groups
`removed`, `kept`, `errors`): sources unchanged, the promised safety backup, then per
file the kept copy and the file are checked again (hashes included, and never the same
file), the file's records move to the kept copy in one short `BEGIN IMMEDIATE`
transaction, and only then is the file removed (every hash handle closed first:
`test_handles_closed_before_unlink`), with its thumbnail (a video's `.jpg` too). A crash
or a failed unlink therefore leaves the records on the identical kept copy, and a failed
move leaves the file and its records alone: no row ever names a deleted file
(`test_a_failure_between_two_files_leaves_no_row_on_a_deleted_file`). The status line
counts the files removed and not removed. A changed file is kept (`test_changed_content_skips_file`); the
first failed unlink stops the job, a locked or read-only file with "The file is open in
another program or read-only; it was not deleted, and its links now name the kept copy."
and never the OS text, and the
files after it are listed as not attempted (`test_partial_failure_reports_each_path`).
Records move for the reviewed files only (`database/duplicate_copies.py`): each image
record, looked up under every spelling of the file by containment
(`py/api/duplicate_spellings.py`: each configured, monitored, ComfyUI output or listed
folder whose `os.stat` identity is a folder above the file, plus the path below it, so a
parent spelling of a nested source and a case variant match:
`test_links_under_a_symlinked_source_spelling_move`,
`test_a_parent_spelling_of_a_nested_source`), moves to the kept copy under the same
spelling, or, when its prompt already links the kept copy, hands its saved-pair
links to that record and is deleted; a favorite under any listed root that holds the
file moves with it (`test_nested_sources_subfolder_first_with_a_parent_favorite`)
(`test_db_links_reconciled_only_for_removed_files`); the copy's image tags join the kept
copy's as their union; and its pending Needs review items (still `pending` at revision
0 with no decision row) are settled (`test_tags_move_as_a_union_and_pending_items_settle`).
The plan states each of these. The find refuses while a removal
runs and the removal while a find runs.
The legacy gallery routes keep their shapes on the same kinds: `/scan_duplicates`
holds `duplicates` (read, "Find duplicate image files (gallery page)"),
`/scan_duplicates/stream` starts the find over the gallery's output folder or attaches
to a running one that covers it (409 otherwise), `/scan_duplicates/status` reports the
latest find, and `/delete_duplicate_images` holds `duplicate_cleanup` (records) and is
refused while a find runs (`test_legacy_delete_duplicates_unchanged`). The link
statements are index searches (`test_link_queries_use_the_path_indexes`; the fixture
check is in the PR). Windows: the hash reads close before `os.remove`; a sharing
violation is a `PermissionError` answered with the coded message (native on the
Windows runners, `test_locked_file_gives_a_coded_message`).
A failure the job cannot name ends it "Maintenance failed; nothing was changed." with a
null code and no exception text. Windows: every connection the backup and the previews
open is closed before the temp copy is renamed or removed, and the suites' teardown
asserts no handle under the temp root is still open.
The legacy `POST /prompt_manager/maintenance` keeps its shape. Its `vacuum` result is
`success:false` with the error when VACUUM fails (`vacuum_database` now raises after
logging), and `cleanup_orphaned_images` removes only links under a source that can be
read now (the health check's sources), never every link of an unmounted drive.
`cleanup_missing_images()` without sources keeps the old rule for its other callers.
Query plans on a fixture of 5,100 prompts, 8,000 images, 300 pairs and 200 recorded
runs: every per-prompt protection and fingerprint read is an index search, the
duplicate groups and the link rows are one scan each, and a delete batch (400 ids)
searches by rowid with its cascades by index.
## Workspace data fields for the concept pages (#249)
These fields are additive. The legacy `/prompt_manager/images/output` listing (no
+212
View File
@@ -0,0 +1,212 @@
"""Records of duplicate image files: count them and move them (#249 S6, T17).
A removed copy's records move to its verified kept copy before the copy is
unlinked (``py/api/duplicate_files.py``), one file per short ``BEGIN
IMMEDIATE`` transaction, so no row ever names a deleted file:
- its image records (``generated_images``), found under every spelling the
file is linked by. A record moves to the kept copy, unless its prompt
already links the kept copy: then its saved-pair links join that record
and it is deleted;
- its favorite;
- its image tags (``media_image_tags`` of its byte record), joined onto the
kept copy's byte record as their union, in order;
- its pending Needs review items, settled: only an item still ``pending`` at
revision 0 with no decision row is removed (``review._SETTLE_ITEM``), so a
decided item is never touched.
Every statement is an index search (see ``FILE_LINK_SQL``).
"""
import os
from .annotation_schema import BYTES_TABLE, FILES_TABLE
from .annotations import _BUMP_INDEX
from .models import normalize_image_path, utc_now_iso
from .review import _BUMP_REVIEW_INDEX, _PENDING_OF_LOCATOR, _SETTLE_ITEM
from .review_changes import notify_review_change
LINKS_BY_FILE_SQL = "SELECT id, prompt_id FROM generated_images WHERE file_path = ?"
LEGACY_LINKS_SQL = (
"SELECT id, prompt_id FROM generated_images "
"WHERE image_path = ? AND file_path IS NULL"
)
KEEPER_LINK_SQL = (
"SELECT id FROM generated_images WHERE prompt_id = ? AND file_path = ?"
)
MOVE_PAIR_LINKS_SQL = (
"UPDATE OR IGNORE image_prompt_pairs SET image_id = ? WHERE image_id = ?"
)
MOVE_LINK_SQL = (
"UPDATE generated_images SET image_path = ?, filename = ?, file_path = ? "
"WHERE id = ?"
)
DELETE_LINK_SQL = "DELETE FROM generated_images WHERE id = ?"
FAVORITE_SQL = (
"SELECT created_at FROM image_favorites WHERE root_key = ? AND relative_path = ?"
)
INSERT_FAVORITE_SQL = "INSERT OR IGNORE INTO image_favorites VALUES (?, ?, ?)"
DELETE_FAVORITE_SQL = (
"DELETE FROM image_favorites WHERE root_key = ? AND relative_path = ?"
)
COPY_BYTE_SQL = (
f"SELECT b.id FROM {FILES_TABLE} f JOIN {BYTES_TABLE} b " # nosec B608
"ON b.file_id=f.id WHERE f.locator=? AND b.sha256=?"
)
COPY_TAG_COUNT_SQL = (
"SELECT COUNT(*) FROM media_image_tags WHERE byte_id=?" # nosec B608
)
PENDING_COUNT_SQL = (
f"SELECT COUNT(*) FROM {FILES_TABLE} f " # nosec B608
f"JOIN {BYTES_TABLE} b ON b.file_id=f.id "
"JOIN media_review_items i ON i.byte_id=b.id "
"WHERE f.locator=? AND +i.state='pending' AND i.revision=0 "
"AND NOT EXISTS (SELECT 1 FROM media_review_decisions d WHERE d.item_id=i.id)"
)
FILE_LINK_SQL = (
LINKS_BY_FILE_SQL,
LEGACY_LINKS_SQL,
KEEPER_LINK_SQL,
MOVE_PAIR_LINKS_SQL,
MOVE_LINK_SQL,
DELETE_LINK_SQL,
FAVORITE_SQL,
DELETE_FAVORITE_SQL,
COPY_BYTE_SQL,
COPY_TAG_COUNT_SQL,
PENDING_COUNT_SQL,
)
def _file_link_rows(conn, spelled):
"""``(id, prompt_id)`` of the image records linking one spelling."""
rows = list(conn.execute(LINKS_BY_FILE_SQL, (normalize_image_path(spelled),)))
rows += list(conn.execute(LEGACY_LINKS_SQL, (spelled,)))
return [tuple(row) for row in rows]
def _copy_byte(conn, locator, sha256):
row = conn.execute(COPY_BYTE_SQL, (locator, sha256)).fetchone()
return row[0] if row else None
class DuplicateCopyOperations:
"""PromptDatabase mixin: the records of duplicate copies."""
def count_copy_records(self, copies) -> dict:
"""``{path: {links, tags, review}}`` for each copy (one read snapshot).
``copies`` maps each path to ``{spellings, locator, sha256}``:
``links`` counts the image records under any spelling, ``tags`` the
image tags of its bytes and ``review`` its pending, undecided Needs
review items.
"""
from .maintenance import _read_snapshot
def read(conn):
found = {}
for path, copy in copies.items():
ids = {
row[0]
for spelled in copy["spellings"]
for row in _file_link_rows(conn, spelled)
}
byte_id = _copy_byte(conn, copy["locator"], copy["sha256"])
tags = (
conn.execute(COPY_TAG_COUNT_SQL, (byte_id,)).fetchone()[0]
if byte_id
else 0
)
review = conn.execute(PENDING_COUNT_SQL, (copy["locator"],))
found[path] = {
"links": len(ids),
"tags": tags,
"review": review.fetchone()[0],
}
return found
return _read_snapshot(self.model.db_path, read)
def move_copy_records(self, move) -> dict:
"""Move one removed copy's records to its kept copy (one transaction).
``move``: ``removed`` and ``kept`` (the spellings each file is linked
by), ``targets`` (``{removed spelling: kept spelling}``, the spelling
a moved record takes; else ``kept_path``), ``favorites`` (``(root key,
removed key, kept root key, kept key)`` per listed root holding the
file), ``removed_locator``,
``kept_locator``, ``sha256`` and the kept copy's ``kept_size`` and
``kept_mtime_ns``. Returns ``{moved, merged, favorites, tags,
review}``.
"""
counts = dict.fromkeys(("moved", "merged", "favorites", "tags", "review"), 0)
kept_keys = [normalize_image_path(path) for path in move["kept"]]
with self._pair_transaction() as conn:
for spelled in dict.fromkeys(move["removed"]):
target = move["targets"].get(spelled, move["kept_path"])
keys = list(dict.fromkeys(kept_keys + [normalize_image_path(target)]))
for image_id, prompt_id in _file_link_rows(conn, spelled):
existing = None
for key in keys:
existing = conn.execute(
KEEPER_LINK_SQL, (prompt_id, key)
).fetchone()
if existing is not None:
break
if existing is not None:
conn.execute(MOVE_PAIR_LINKS_SQL, (existing[0], image_id))
conn.execute(DELETE_LINK_SQL, (image_id,))
counts["merged"] += 1
continue
conn.execute(
MOVE_LINK_SQL,
(
target,
os.path.basename(target),
normalize_image_path(target),
image_id,
),
)
counts["moved"] += 1
for favorite in move.get("favorites", ()):
counts["favorites"] += _move_favorite(conn, *favorite)
counts["tags"] = self._move_copy_tags(conn, move)
for (item_id,) in conn.execute(
_PENDING_OF_LOCATOR, (move["removed_locator"],)
).fetchall():
counts["review"] += conn.execute(_SETTLE_ITEM, (item_id,)).rowcount
if counts["review"]:
conn.execute(_BUMP_REVIEW_INDEX)
if counts["review"]:
notify_review_change(self.model.db_path)
return counts
def _move_copy_tags(self, conn, move):
"""Join the removed copy's image tags onto the kept copy's; the count added."""
removed = _copy_byte(conn, move["removed_locator"], move["sha256"])
if removed is None:
return 0
tags = self._byte_tags(conn, removed)
if not tags:
return 0
now = utc_now_iso()
file_id = self._annotation_owner(conn, move["kept_locator"], now)
byte_id, _revision, created, moved = self._annotation_byte(
conn, file_id, move["sha256"], move["kept_size"], move["kept_mtime_ns"], now
)
current = self._byte_tags(conn, byte_id)
union = current + [tag for tag in tags if tag not in current]
if union != current:
self._write_image_tags(conn, byte_id, union)
elif created or moved:
conn.execute(_BUMP_INDEX)
return len(union) - len(current)
def _move_favorite(conn, root_key, removed_key, kept_root, kept_key):
row = conn.execute(FAVORITE_SQL, (root_key, removed_key)).fetchone()
if row is None:
return 0
conn.execute(INSERT_FAVORITE_SQL, (kept_root, kept_key, row[0]))
conn.execute(DELETE_FAVORITE_SQL, (root_key, removed_key))
return 1
+32 -6
View File
@@ -189,6 +189,33 @@ def _image_owner(image_path, file_path, prefixes):
return owner
def readable_roots(roots):
"""Whether each source can be read now (an unmounted drive cannot)."""
return [_readable(root) for root in roots]
def link_states(rows, roots, readable):
"""``(row, owner, state)`` for each image row ``(id, image_path, filename,
file_path)``: ``owner`` is the index of the source it lies under (None
outside every source) and ``state`` is ``missing``, ``present``,
``unavailable`` (its source cannot be read, so the file is not looked at)
or ``outside``. ``readable`` is ``readable_roots(roots)``. The health
check's ``missing_files`` and the maintenance missing-link plan both
judge links this way, so their counts agree.
"""
prefixes = [_root_prefix(root) for root in roots]
for row in rows:
owner = _image_owner(row[1], row[3], prefixes)
if owner is None:
yield row, None, "outside"
elif not readable[owner]:
yield row, owner, "unavailable"
elif os.path.exists(row[1]):
yield row, owner, "present"
else:
yield row, owner, "missing"
def _unavailable_line(unavailable, *, partial):
noun = "source is" if unavailable == 1 else "sources are"
effect = "its files were" if unavailable == 1 else "their files were"
@@ -224,7 +251,7 @@ def _missing_files(conn, roots):
def _check_files(rows, roots):
"""The ``missing_files`` result of the image *rows*, from the file system."""
readable = [_readable(root) for root in roots]
readable = readable_roots(roots)
sources = [
{
"name": _source_name(root, index),
@@ -239,13 +266,12 @@ def _check_files(rows, roots):
None, [_unavailable_line(unavailable, partial=False)], status="unavailable"
)
return {**result, "partial": False, "sources": sources}
prefixes = [_root_prefix(root) for root in roots]
outside, missing = 0, []
for image_id, image_path, filename, file_path in rows:
owner = _image_owner(image_path, file_path, prefixes)
if owner is None:
for row, owner, state in link_states(rows, roots, readable):
if state == "outside":
outside += 1
elif readable[owner] and not os.path.exists(image_path):
elif state == "missing":
image_id, image_path, filename, _file_path = row
sources[owner]["count"] += 1
if len(missing) < MAX_DETAIL:
name = filename or os.path.basename(image_path)
+724
View File
@@ -0,0 +1,724 @@
"""Reviewed maintenance plans: read-only previews and their exact apply (#249 S6).
Four operations, each the native owner of a legacy ``/maintenance`` one and
built on the same queries, so a preview counts what the legacy cleanup would
change:
- ``duplicate_prompts``: prompts whose text is equal ignoring case and outer
spaces (``DUPLICATE_PROMPTS_SQL``, legacy ``cleanup_duplicates``); each
group keeps its oldest prompt, which takes over the others' images and
metadata. A group with a member in a saved pair is kept whole.
- ``optimize``: VACUUM (legacy ``vacuum``).
- ``unused_components``: prompts nothing references (``ORPHANED_PROMPTS_SQL``,
legacy ``prune_orphaned_prompts``, with the #239 protections).
- ``missing_links``: image rows whose file is gone (legacy
``cleanup_orphaned_images``), counted only under a source that can be read
now, as the health check's ``missing_files`` judges them
(``health.link_states``). It deletes database rows only, never a file.
``preview_maintenance`` reads one snapshot on its own ``mode=ro`` connection
(``health.readonly_uri``) and closes it in ``finally``; files are looked at
after the snapshot ended. It returns a ``MaintenanceDraft``: the public plan
fields, the targets an apply needs and the fingerprint of exactly what was
reviewed (ids, text hashes, metadata, tags, Library state and pair refs of
every prompt involved, or each link's row and source). The fingerprint never
leaves the server.
``apply_maintenance`` recomputes that fingerprint on the write connection
inside ``BEGIN IMMEDIATE`` and raises ``PlanChanged`` before writing anything
when it differs; otherwise it applies the reviewed targets in batches of
``ID_BATCH`` in that one transaction. VACUUM cannot run in a transaction, so
``optimize`` binds to its database only (the caller checks that).
"""
import hashlib
import json
import os
import shutil
import sqlite3
from dataclasses import dataclass, field
from pathlib import Path
from .execution_contexts import LIVE_HINT_CTES
from .health import BATCH_ROWS, link_states, readable_roots, readonly_uri
from .pairs import PairDependencyError
MAINTENANCE_OPERATIONS = (
"duplicate_prompts",
"optimize",
"unused_components",
"missing_links",
)
MAX_DETAILS = 50 # detail rows a plan lists
MAX_RESULT_ROWS = 200 # changed records an apply describes (the job lists 200)
EXCERPT_CHARS = 80
# Ids bound per statement: the pair-reference read binds each id twice, which
# stays under SQLite's historical 999-variable limit.
ID_BATCH = 400
# Why a prompt no image links to is kept (#239): predicates on ``p``, in the
# order the plan names them. ORPHANED_PROMPTS_SQL is built from them, so the
# reasons a plan lists are exactly the rules the prune applies.
PROTECTIONS = (
(
"EXISTS (SELECT 1 FROM prompt_pairs pp "
"WHERE pp.positive_prompt_id = p.id OR pp.negative_prompt_id = p.id)",
"in saved pairs",
"in a saved pair",
),
(
"EXISTS (SELECT 1 FROM prompt_tags pt JOIN tags t ON pt.tag_id = t.id "
"WHERE pt.prompt_id = p.id AND t.name = '__protected__')",
"tagged __protected__",
"tagged __protected__",
),
(
"EXISTS (SELECT 1 FROM prompt_library_state ls WHERE ls.prompt_id = p.id)",
"with Library state (a title or an explicit parent)",
"with Library state (a title or an explicit parent)",
),
(
"p.id IN (SELECT prompt_id FROM live_hint_prompts)",
"used by recorded runs",
"used by a recorded run",
),
)
# Exempt references: images, pairs, the __protected__ tag, Library parent
# state, and live recorded-run hints (computed once, see LIVE_HINT_CTES).
ORPHANED_PROMPTS_SQL = (
"WITH "
+ LIVE_HINT_CTES
+ " SELECT p.id FROM prompts p"
+ " LEFT JOIN generated_images gi ON p.id = gi.prompt_id"
+ " WHERE gi.prompt_id IS NULL"
+ "".join(f" AND NOT ({predicate})" for predicate, _many, _one in PROTECTIONS)
)
# The prompts no image links to, with which protections keep each one.
_UNLINKED_PROTECTIONS_SQL = (
"WITH "
+ LIVE_HINT_CTES
+ " SELECT "
+ ", ".join(
f"COALESCE(SUM({predicate}), 0)" for predicate, _many, _one in PROTECTIONS
)
+ ", COALESCE(SUM("
+ " OR ".join(f"({predicate})" for predicate, _many, _one in PROTECTIONS)
+ "), 0) FROM prompts p WHERE NOT EXISTS "
"(SELECT 1 FROM generated_images gi WHERE gi.prompt_id = p.id)"
)
# Duplicate prompt groups, as the legacy cleanup finds them (case-insensitive
# text; GROUP_CONCAT order is sorted by created_at afterwards).
DUPLICATE_PROMPTS_SQL = """
SELECT LOWER(TRIM(text)) as normalized_text, COUNT(*) as count,
GROUP_CONCAT(id) as ids,
GROUP_CONCAT(created_at) as created_dates
FROM prompts
GROUP BY LOWER(TRIM(text))
HAVING COUNT(*) > 1
"""
_LINK_ROWS_SQL = (
"SELECT id, image_path, filename, file_path FROM generated_images ORDER BY id"
)
# What a prompt's fingerprint covers, per batch of ids.
_PROMPT_STATE_SQL = (
"SELECT id, hash, text, category, rating, notes, created_at, updated_at "
"FROM prompts WHERE id IN ({marks}) ORDER BY id",
"SELECT prompt_id, tag_id FROM prompt_tags WHERE prompt_id IN ({marks}) "
"ORDER BY prompt_id, tag_id",
"SELECT * FROM prompt_library_state WHERE prompt_id IN ({marks}) "
"ORDER BY prompt_id",
"SELECT id, positive_prompt_id, negative_prompt_id FROM prompt_pairs "
"WHERE positive_prompt_id IN ({marks}) OR negative_prompt_id IN ({marks}) "
"ORDER BY id",
)
NO_FILES_DELETED = "No image file is deleted."
LINKS_ONLY = "Database links only; no image file is deleted."
UNUSED_DETAIL = "No image, saved pair, Library state or recorded run uses it"
# Why a duplicate group is kept whole (``_duplicate_conflict``).
CONFLICT_REASONS = {
"collision": ("group links one image twice", "groups link one image twice"),
"title": (
"group has its own Library title",
"groups have their own Library titles",
),
}
CONFLICT_DETAILS = {
"collision": "Kept: two copies link the same image",
"title": "Kept: a copy has its own Library title",
}
OPTIMIZE_LOCK_NOTE = (
"Runs alone: the database is locked while it runs, and on a large library "
"this takes a while. Prompt runs counted and images captured meanwhile "
"wait up to 30 seconds each; anything that waits longer is not recorded."
)
class PlanChanged(Exception):
"""What a plan was reviewed on changed before it was applied."""
class MergeConflict(Exception):
"""A strict duplicate merge would drop an image link or a Library title."""
@dataclass
class MaintenanceDraft:
"""A preview: the public plan fields plus what only the server keeps."""
operation: str
fingerprint: str
targets: list
scope: dict
protected: dict
media: dict
details: list = field(default_factory=list)
def public(self):
"""The plan fields a response may carry (no fingerprint, no targets)."""
return {
"operation": self.operation,
"scope": dict(self.scope),
"protected": {
"count": self.protected["count"],
"reasons": list(self.protected["reasons"]),
},
"media": dict(self.media),
"details": [dict(row) for row in self.details],
}
def _plural(count, one, many=None):
return f"{count:,} {one if count == 1 else (many or one + 's')}"
def _excerpt(text):
text = " ".join(str(text or "").split())
if len(text) <= EXCERPT_CHARS:
return text
return text[: EXCERPT_CHARS - 1].rstrip() + "…"
def _megabytes(size):
return f"{size / 2**20:,.1f} MB"
class _Fingerprint:
"""A streaming sha256 over JSON-encoded parts."""
def __init__(self, operation):
self._hash = hashlib.sha256()
self.add(operation)
def add(self, part):
encoded = json.dumps(part, separators=(",", ":"), default=str)
self._hash.update(encoded.encode("utf-8") + b"\n")
def hexdigest(self):
return "sha256:" + self._hash.hexdigest()
def _batches(ids):
ids = list(ids)
for start in range(0, len(ids), ID_BATCH):
yield ids[start : start + ID_BATCH]
def _marks(batch):
return ",".join("?" for _ in batch)
def _add_prompt_state(fingerprint, conn, ids):
"""Fold every reviewed fact about the prompts ``ids`` into ``fingerprint``."""
for batch in _batches(sorted(ids)):
marks = _marks(batch)
for sql in _PROMPT_STATE_SQL:
params = batch * sql.count("{marks}")
for row in conn.execute(sql.format(marks=marks), params):
fingerprint.add(list(row))
def _texts(conn, ids, limit=MAX_DETAILS):
"""``{id: text}`` of the first ``limit`` prompts, for detail rows."""
ids = list(ids)[:limit]
if not ids:
return {}
rows = conn.execute(
f"SELECT id, text FROM prompts WHERE id IN ({_marks(ids)})", ids
)
return {row[0]: row[1] for row in rows}
def _orphaned_ids(conn):
return sorted(row[0] for row in conn.execute(ORPHANED_PROMPTS_SQL))
def _delete_ids(conn, table, ids):
"""Delete ``table`` rows by id in batches; returns the rows removed."""
if table not in ("prompts", "generated_images"):
raise ValueError(f"Unsupported table {table!r}")
removed = 0
for batch in _batches(ids):
cursor = conn.execute(
f"DELETE FROM {table} WHERE id IN ({_marks(batch)})", # nosec B608
batch,
)
removed += cursor.rowcount
return removed
def _read_snapshot(db_path, read):
"""``read(conn)`` on one ``mode=ro`` snapshot; the connection always closes."""
conn = sqlite3.connect(readonly_uri(db_path), uri=True)
try:
conn.row_factory = sqlite3.Row
conn.execute("BEGIN") # one read snapshot for the whole plan
return read(conn)
finally:
conn.close()
def _relative(image_path, root):
try:
return Path(os.path.relpath(image_path, root)).as_posix()
except ValueError: # another drive (Windows)
return os.path.basename(image_path)
def _source_name(root, index):
return os.path.basename(str(root).rstrip("\\/")) or f"Source {index + 1}"
class MaintenanceOperations:
"""Mixin for PromptDatabase: maintenance previews and their apply."""
# ── shared with the legacy cleanup ───────────────────────────────
def _duplicate_groups(self, conn):
"""``(groups, paired)``: duplicate id groups, oldest first.
``paired`` are the groups with a member in a saved pair, which the
cleanup keeps whole.
"""
groups, paired = [], []
for duplicate in conn.execute(DUPLICATE_PROMPTS_SQL).fetchall():
ids = duplicate["ids"].split(",")
created_dates = duplicate["created_dates"].split(",")
# Sort IDs by created_at date to keep the oldest
id_date_pairs = list(zip(ids, created_dates))
id_date_pairs.sort(key=lambda x: x[1])
sorted_ids = [int(pair[0]) for pair in id_date_pairs]
try:
self._ensure_prompts_unpaired(conn, sorted_ids)
except PairDependencyError:
paired.append(sorted_ids)
continue
groups.append(sorted_ids)
return groups, paired
@staticmethod
def _duplicate_conflict(conn, sorted_ids):
"""Why merging the group would lose something, or None.
``collision``: two members link the same file, so moving the images
to the oldest prompt would hit ``UNIQUE(prompt_id, file_path)``.
``title``: a member that would be deleted has its own Library title.
"""
marks = _marks(sorted_ids)
collision = conn.execute(
"SELECT 1 FROM generated_images WHERE prompt_id IN " # nosec B608
f"({marks}) AND file_path IS NOT NULL GROUP BY file_path "
"HAVING COUNT(DISTINCT prompt_id) > 1 LIMIT 1",
sorted_ids,
).fetchone()
if collision:
return "collision"
duplicates = sorted_ids[1:]
titled = conn.execute(
"SELECT 1 FROM prompt_library_state WHERE prompt_id IN " # nosec B608
f"({_marks(duplicates)}) AND trim(coalesce(title, '')) != '' LIMIT 1",
duplicates,
).fetchone()
return "title" if titled else None
def _merge_duplicate_group(self, conn, sorted_ids, strict=False):
"""Keep the oldest prompt of a group; returns (removed, images moved).
``strict`` (the reviewed apply): a merge that would drop an image
link or a Library title raises MergeConflict, and every helper error
propagates, so the transaction rolls back. The legacy cleanup keeps
its old, lenient behaviour.
"""
primary_id, duplicate_ids = sorted_ids[0], sorted_ids[1:]
if strict:
conflict = self._duplicate_conflict(conn, sorted_ids)
if conflict:
raise MergeConflict(conflict)
self.logger.debug(
f"Merging duplicates: keeping {primary_id}, removing {duplicate_ids}"
)
primary = conn.execute(
"SELECT * FROM prompts WHERE id = ?", (primary_id,)
).fetchone()
if not primary:
return 0, 0
merged_metadata = self._merge_duplicate_metadata(
conn, primary_id, duplicate_ids, strict=strict
)
moved = self._transfer_images_to_primary(
conn, primary_id, duplicate_ids, strict=strict
)
if merged_metadata:
self._update_primary_with_merged_metadata(
conn, primary_id, merged_metadata, strict=strict
)
removed = 0
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}")
return removed, moved
@staticmethod
def _missing_link_ids(rows, roots):
"""Ids of the ``rows`` whose file is gone under a readable source."""
readable = readable_roots(roots)
return [
row[0]
for row, _owner, state in link_states(rows, roots, readable)
if state == "missing"
]
# ── previews ──────────────────────────────────────────────────────
def preview_maintenance(self, operation, roots=()) -> MaintenanceDraft:
"""The reviewed plan of ``operation``; writes nothing (blocking).
``roots`` are the sources missing links are judged under (the health
check's: the effective sources plus configured ones that are gone).
"""
if operation not in MAINTENANCE_OPERATIONS:
raise ValueError(f"Unknown maintenance operation {operation!r}")
roots = [str(root) for root in roots]
if operation == "optimize":
return self._plan_optimize()
if operation == "missing_links":
rows = _read_snapshot(self.model.db_path, self._link_rows)
return self._plan_missing_links(rows, roots)
reader = (
self._plan_duplicate_prompts
if operation == "duplicate_prompts"
else self._plan_unused_components
)
return _read_snapshot(self.model.db_path, reader)
@staticmethod
def _link_rows(conn):
rows = []
cursor = conn.execute(_LINK_ROWS_SQL)
while batch := cursor.fetchmany(BATCH_ROWS):
rows.extend(tuple(row) for row in batch)
return rows
def _duplicates_fingerprint(self, conn):
"""``(groups, paired, conflicted, fingerprint)`` of the duplicate merge.
``conflicted`` holds ``(ids, reason)`` of the groups a merge would
lose an image link or a Library title in (``_duplicate_conflict``);
the plan protects them.
"""
candidates, paired = self._duplicate_groups(conn)
groups, conflicted = [], []
for ids in candidates:
reason = self._duplicate_conflict(conn, ids)
if reason:
conflicted.append((ids, reason))
else:
groups.append(ids)
fingerprint = _Fingerprint("duplicate_prompts")
fingerprint.add({"groups": groups, "paired": paired, "conflicted": conflicted})
_add_prompt_state(fingerprint, conn, [i for group in groups for i in group])
return groups, paired, conflicted, fingerprint.hexdigest()
def _plan_duplicate_prompts(self, conn):
groups, paired, conflicted, fingerprint = self._duplicates_fingerprint(conn)
removed = sum(len(group) - 1 for group in groups)
held = [ids for ids, _reason in conflicted]
texts = _texts(conn, [group[0] for group in (groups + paired + held)])
details = (
[
{
"label": _excerpt(texts.get(group[0])),
"detail": f"Keeps the oldest of {len(group)} copies",
}
for group in groups
]
+ [
{
"label": _excerpt(texts.get(group[0])),
"detail": "Kept: belongs to a saved pair",
}
for group in paired
]
+ [
{
"label": _excerpt(texts.get(ids[0])),
"detail": CONFLICT_DETAILS[reason],
}
for ids, reason in conflicted
]
)
if groups:
label = (
f"{_plural(removed, 'duplicate prompt')} in "
f"{_plural(len(groups), 'group')}"
)
else:
label = "No duplicate prompt text"
reasons = []
if len(paired) == 1:
reasons.append("1 group belongs to a saved pair")
elif paired:
reasons.append(f"{len(paired):,} groups belong to saved pairs")
for reason, (one, many) in CONFLICT_REASONS.items():
count = sum(1 for _ids, why in conflicted if why == reason)
if count:
reasons.append(f"{count:,} {one if count == 1 else many}")
return MaintenanceDraft(
operation="duplicate_prompts",
fingerprint=fingerprint,
targets=groups,
scope={"count": removed, "label": label},
protected={
"count": sum(len(group) for group in paired + held),
"reasons": reasons,
},
media={"deletes_files": False, "label": NO_FILES_DELETED},
details=details[:MAX_DETAILS],
)
def _unused_fingerprint(self, conn):
ids = _orphaned_ids(conn)
fingerprint = _Fingerprint("unused_components")
fingerprint.add(ids)
_add_prompt_state(fingerprint, conn, ids)
return ids, fingerprint.hexdigest()
def _plan_unused_components(self, conn):
ids, fingerprint = self._unused_fingerprint(conn)
*by_reason, kept = conn.execute(_UNLINKED_PROTECTIONS_SQL).fetchone()
reasons = [
f"{count:,} {many if count != 1 else one}"
for count, (_predicate, many, one) in zip(by_reason, PROTECTIONS)
if count
]
texts = _texts(conn, ids)
count = len(ids)
return MaintenanceDraft(
operation="unused_components",
fingerprint=fingerprint,
targets=ids,
scope={
"count": count,
"label": (
_plural(count, "unused prompt component")
if count
else "No unused prompt components"
),
},
protected={"count": kept, "reasons": reasons},
media={"deletes_files": False, "label": NO_FILES_DELETED},
details=[
{"label": _excerpt(texts.get(prompt_id)), "detail": UNUSED_DETAIL}
for prompt_id in ids[:MAX_DETAILS]
],
)
@staticmethod
def _links_fingerprint(candidates):
fingerprint = _Fingerprint("missing_links")
for row, owner in candidates:
fingerprint.add([row[0], row[1], row[3], owner])
return fingerprint.hexdigest()
def _plan_missing_links(self, rows, roots):
readable = readable_roots(roots)
candidates, unavailable, outside = [], 0, 0
for row, owner, state in link_states(rows, roots, readable):
if state == "missing":
candidates.append((row, owner))
elif state == "unavailable":
unavailable += 1
elif state == "outside":
outside += 1
reasons = []
if unavailable:
reasons.append(
_plural(
unavailable,
"link under an unavailable source",
"links under unavailable sources",
)
)
if outside:
reasons.append(
_plural(
outside,
"link outside the image sources",
"links outside the image sources",
)
)
count = len(candidates)
return MaintenanceDraft(
operation="missing_links",
fingerprint=self._links_fingerprint(candidates),
targets=[row[0] for row, _owner in candidates],
scope={
"count": count,
"label": (
_plural(count, "link to a missing file", "links to missing files")
if count
else "No links to missing files"
),
},
protected={"count": unavailable + outside, "reasons": reasons},
media={"deletes_files": False, "label": LINKS_ONLY},
details=[
{
"label": _relative(row[1], roots[owner]),
"detail": (
f"{_source_name(roots[owner], owner)} · the file is missing"
),
}
for row, owner in candidates[:MAX_DETAILS]
],
)
def _plan_optimize(self):
def read(conn):
return [
conn.execute(f"PRAGMA {name}").fetchone()[0]
for name in ("page_size", "page_count", "freelist_count")
]
page_size, page_count, free_pages = _read_snapshot(self.model.db_path, read)
db_path = self.model.db_path
wal = db_path + "-wal"
size = page_size * page_count
wal_size = os.path.getsize(wal) if os.path.exists(wal) else 0
free = shutil.disk_usage(os.path.dirname(os.path.abspath(db_path))).free
# VACUUM writes the rebuilt database through the WAL beside it.
return MaintenanceDraft(
operation="optimize",
fingerprint=_Fingerprint("optimize").hexdigest(),
targets=[],
scope={"count": 1, "label": "The whole database"},
protected={"count": 0, "reasons": []},
media={"deletes_files": False, "label": NO_FILES_DELETED},
details=[
{
"label": "Database file",
"detail": f"{_megabytes(size)}, plus {_megabytes(wal_size)} "
"not yet checkpointed",
},
{
"label": "Free pages",
"detail": f"About {_megabytes(free_pages * page_size)} "
"can be reclaimed",
},
{
"label": "Lock",
"detail": OPTIMIZE_LOCK_NOTE,
},
{
"label": "Free space needed",
"detail": f"About {_megabytes(size)} beside the database and "
f"{_megabytes(size)} in SQLite's temporary folder; "
f"{_megabytes(free)} available beside the database",
},
],
)
# ── apply ─────────────────────────────────────────────────────────
def apply_maintenance(self, operation, fingerprint, targets, roots=()) -> dict:
"""Apply a reviewed plan exactly, or raise PlanChanged (blocking).
Inside ``BEGIN IMMEDIATE`` the fingerprint is computed again on the
write connection; when it differs nothing is written. ``targets`` and
``roots`` are the plan's. Returns ``{changed, rows}``: how many
records changed and up to ``MAX_RESULT_ROWS`` descriptions of them
(``{prompt_id[, kept_id]}`` or ``{image_path}``).
"""
if operation not in MAINTENANCE_OPERATIONS:
raise ValueError(f"Unknown maintenance operation {operation!r}")
if operation == "optimize":
# VACUUM cannot run in a transaction; the plan binds the database.
if fingerprint != _Fingerprint("optimize").hexdigest():
raise PlanChanged("optimize")
self.model.vacuum_database()
return {"changed": 1, "rows": []}
appliers = {
"duplicate_prompts": self._apply_duplicate_prompts,
"unused_components": self._apply_unused_components,
"missing_links": self._apply_missing_links,
}
with self._pair_transaction() as conn:
return appliers[operation](conn, fingerprint, targets, roots)
def _apply_duplicate_prompts(self, conn, fingerprint, targets, _roots):
groups, _paired, _conflicted, current = self._duplicates_fingerprint(conn)
if current != fingerprint or groups != targets:
raise PlanChanged("duplicate_prompts")
changed, rows = 0, []
for sorted_ids in groups:
removed, _moved = self._merge_duplicate_group(conn, sorted_ids, strict=True)
changed += removed
rows.extend(
{"prompt_id": duplicate_id, "kept_id": sorted_ids[0]}
for duplicate_id in sorted_ids[1:]
)
return {"changed": changed, "rows": rows[:MAX_RESULT_ROWS]}
def _apply_unused_components(self, conn, fingerprint, targets, _roots):
ids, current = self._unused_fingerprint(conn)
if current != fingerprint or ids != targets:
raise PlanChanged("unused_components")
texts = _texts(conn, ids, limit=MAX_RESULT_ROWS)
changed = _delete_ids(conn, "prompts", ids)
rows = [
{"prompt_id": prompt_id, "text": _excerpt(texts.get(prompt_id))}
for prompt_id in ids[:MAX_RESULT_ROWS]
]
return {"changed": changed, "rows": rows}
def _apply_missing_links(self, conn, fingerprint, targets, roots):
rows = []
for batch in _batches(sorted(targets)):
rows.extend(
tuple(row)
for row in conn.execute(
"SELECT id, image_path, filename, file_path FROM generated_images "
f"WHERE id IN ({_marks(batch)}) ORDER BY id", # nosec B608
batch,
)
)
readable = readable_roots(roots)
candidates = []
for row, owner, state in link_states(rows, roots, readable):
if state != "missing": # back, or its source went away
raise PlanChanged("missing_links")
candidates.append((row, owner))
if self._links_fingerprint(candidates) != fingerprint:
raise PlanChanged("missing_links")
changed = _delete_ids(conn, "generated_images", [row[0] for row in rows])
return {
"changed": changed,
"rows": [{"image_path": row[1]} for row in rows[:MAX_RESULT_ROWS]],
}
+5
View File
@@ -941,6 +941,10 @@ class PromptModel:
Reclaims unused space and defragments the database file,
improving query performance and reducing file size.
Raises:
Exception: whatever stopped VACUUM (logged first); the database is
then unchanged.
"""
try:
with closing(sqlite3.connect(self.db_path)) as conn:
@@ -948,6 +952,7 @@ class PromptModel:
conn.commit()
except Exception as e:
self.logger.error(f"Error vacuuming database: {e}")
raise
def get_database_info(self) -> dict:
"""
+60 -107
View File
@@ -10,7 +10,7 @@ import os
from typing import Optional, List, Dict, Any, Set, Tuple
from .models import PromptModel, normalize_image_path, utc_now_iso
from .pairs import PairDependencyError, PairOperations
from .pairs import PairDependencyError, PairOperations # noqa: F401 (re-exported)
from .library import LibraryOperations, titled
from .image_favorites import ImageFavoriteOperations
from .annotations import AnnotationOperations
@@ -20,10 +20,18 @@ from .saved_views import SavedViewOperations
from .library_commands import LibraryCommandOperations
from .vocabulary import VocabularyOperations
from .execution_history import ExecutionHistoryOperations
from .execution_contexts import ExecutionContextOperations, LIVE_HINT_CTES
from .execution_contexts import ExecutionContextOperations
from .execution_outputs import ExecutionOutputOperations
from .history_reads import HistoryReadOperations
from .health import HealthOperations, iter_consistency_issues
from .duplicate_copies import DuplicateCopyOperations
from .maintenance import ( # noqa: F401 (ORPHANED_PROMPTS_SQL is re-exported)
ORPHANED_PROMPTS_SQL,
_LINK_ROWS_SQL,
MaintenanceOperations,
_delete_ids,
_orphaned_ids,
)
# Import logging system
try:
@@ -166,28 +174,6 @@ def _coerce_file_info(metadata: Any) -> tuple:
)
# Exempt references: images, pairs, the __protected__ tag, Library parent
# state, and live recorded-run hints (computed once, see LIVE_HINT_CTES).
ORPHANED_PROMPTS_SQL = "WITH " + LIVE_HINT_CTES + """
SELECT p.id FROM prompts p
LEFT JOIN generated_images gi ON p.id = gi.prompt_id
WHERE gi.prompt_id IS NULL
AND NOT EXISTS (
SELECT 1 FROM prompt_pairs pp
WHERE pp.positive_prompt_id = p.id OR pp.negative_prompt_id = p.id
)
AND NOT EXISTS (
SELECT 1 FROM prompt_tags pt
JOIN tags t ON pt.tag_id = t.id
WHERE pt.prompt_id = p.id AND t.name = '__protected__'
)
AND NOT EXISTS (
SELECT 1 FROM prompt_library_state ls WHERE ls.prompt_id = p.id
)
AND p.id NOT IN (SELECT prompt_id FROM live_hint_prompts)
"""
# The scan's known set (#249 Activity-2): every image row, with its stored file
# size and whether its pair link is resolved: exactly one distinct pair, the rule
# _LINK_SUMMARY_SQL and get_image_prompt_info use. A file is known and resolved
@@ -351,6 +337,8 @@ class PromptDatabase(
ReviewDecisionOperations,
SavedViewOperations,
HealthOperations,
MaintenanceOperations,
DuplicateCopyOperations,
):
"""Database operations class for managing prompts."""
@@ -1541,80 +1529,14 @@ class PromptDatabase(
self.logger.info("Starting duplicate cleanup process with image preservation")
try:
with self._pair_transaction() as conn:
# Find duplicates by text content (case-insensitive)
# Note: Removed ORDER BY from GROUP_CONCAT for SQLite compatibility
# We'll sort the IDs manually after fetching
cursor = conn.execute("""
SELECT LOWER(TRIM(text)) as normalized_text, COUNT(*) as count,
GROUP_CONCAT(id) as ids,
GROUP_CONCAT(created_at) as created_dates
FROM prompts
GROUP BY LOWER(TRIM(text))
HAVING COUNT(*) > 1
""")
duplicates = cursor.fetchall()
self.logger.debug(
f"Found {len(duplicates)} groups of duplicate prompts"
)
groups, _paired = self._duplicate_groups(conn)
self.logger.debug(f"Found {len(groups)} groups of duplicate prompts")
total_removed = 0
total_images_transferred = 0
for duplicate in duplicates:
ids = duplicate["ids"].split(",")
created_dates = duplicate["created_dates"].split(",")
# Sort IDs by created_at date to keep the oldest
id_date_pairs = list(zip(ids, created_dates))
id_date_pairs.sort(key=lambda x: x[1]) # Sort by date
sorted_ids = [int(pair[0]) for pair in id_date_pairs]
try:
self._ensure_prompts_unpaired(conn, sorted_ids)
except PairDependencyError:
continue
# Keep the oldest one (first), merge and delete the rest
primary_id = sorted_ids[0] # Keep the oldest
duplicate_ids = sorted_ids[1:]
self.logger.debug(
f"Merging duplicates: keeping {primary_id}, removing "
f"{duplicate_ids}"
)
# Get primary prompt details
cursor = conn.execute(
"SELECT * FROM prompts WHERE id = ?", (primary_id,)
)
primary_prompt = cursor.fetchone()
if not primary_prompt:
continue
# Collect and merge metadata from all duplicates
merged_metadata = self._merge_duplicate_metadata(
conn, primary_id, duplicate_ids
)
# Transfer all images from duplicates to primary prompt
images_transferred = self._transfer_images_to_primary(
conn, primary_id, duplicate_ids
)
total_images_transferred += images_transferred
# Update primary prompt with merged metadata
if merged_metadata:
self._update_primary_with_merged_metadata(
conn, primary_id, merged_metadata
)
# Delete duplicate prompts (images already transferred)
for duplicate_id in duplicate_ids:
conn.execute(
"DELETE FROM prompts WHERE id = ?", (duplicate_id,)
)
total_removed += 1
self.logger.debug(f"Removed duplicate prompt {duplicate_id}")
for sorted_ids in groups:
removed, moved = self._merge_duplicate_group(conn, sorted_ids)
total_removed += removed
total_images_transferred += moved
conn.commit()
@@ -1879,13 +1801,30 @@ class PromptDatabase(
conn.commit()
return cursor.rowcount > 0
def cleanup_missing_images(self) -> int:
def cleanup_missing_images(self, roots=None) -> int:
"""
Remove image records where the actual file no longer exists.
Args:
roots: The image sources links are judged under (the health
check's). When given, only a link under a source that can be
read now is removed: a link under an unavailable source (an
unmounted drive) or outside every source is kept. None keeps
the old rule: every link whose file does not exist.
Returns:
int: Number of orphaned records removed
"""
if roots is not None:
with self.model.get_connection() as conn:
rows = [tuple(row) for row in conn.execute(_LINK_ROWS_SQL)]
# Files are looked at before the write transaction starts.
ids = self._missing_link_ids(rows, [str(root) for root in roots])
if not ids:
return 0
with self._pair_transaction() as conn:
return _delete_ids(conn, "generated_images", ids)
removed_count = 0
with self.model.get_connection() as conn:
@@ -1923,7 +1862,11 @@ class PromptDatabase(
return obj
def _merge_duplicate_metadata(
self, conn: sqlite3.Connection, primary_id: int, duplicate_ids: List[int]
self,
conn: sqlite3.Connection,
primary_id: int,
duplicate_ids: List[int],
strict: bool = False,
) -> Dict[str, Any]:
"""
Merge metadata from duplicate prompts, prioritizing non-empty values.
@@ -1983,11 +1926,17 @@ class PromptDatabase(
return merged
except Exception as e:
if strict: # a reviewed apply rolls back instead
raise
self.logger.error(f"Error merging metadata: {e}", exc_info=True)
return {}
def _transfer_images_to_primary(
self, conn: sqlite3.Connection, primary_id: int, duplicate_ids: List[int]
self,
conn: sqlite3.Connection,
primary_id: int,
duplicate_ids: List[int],
strict: bool = False,
) -> int:
"""
Transfer all images from duplicate prompts to the primary prompt.
@@ -2019,11 +1968,17 @@ class PromptDatabase(
return transferred_count
except Exception as e:
if strict: # a reviewed apply rolls back instead
raise
self.logger.error(f"Error transferring images: {e}", exc_info=True)
return 0
def _update_primary_with_merged_metadata(
self, conn: sqlite3.Connection, primary_id: int, merged_metadata: Dict[str, Any]
self,
conn: sqlite3.Connection,
primary_id: int,
merged_metadata: Dict[str, Any],
strict: bool = False,
) -> None:
"""
Update the primary prompt with merged metadata.
@@ -2054,6 +2009,8 @@ class PromptDatabase(
)
except Exception as e:
if strict: # a reviewed apply rolls back instead
raise
self.logger.error(
f"Error updating primary prompt metadata: {e}", exc_info=True
)
@@ -2567,16 +2524,12 @@ class PromptDatabase(
Returns number of prompts removed.
"""
with self._pair_transaction() as conn:
cursor = conn.execute(ORPHANED_PROMPTS_SQL)
orphaned = [row["id"] for row in cursor.fetchall()]
orphaned = _orphaned_ids(conn)
if not orphaned:
return 0
placeholders = ",".join(["?"] * len(orphaned))
cursor = conn.execute(
f"DELETE FROM prompts WHERE id IN ({placeholders})", orphaned
)
removed = _delete_ids(conn, "prompts", orphaned)
conn.commit()
return cursor.rowcount
return removed
def check_consistency(self) -> List[str]:
"""Run consistency checks and return a list of issue descriptions.
+25 -2
View File
@@ -13,6 +13,8 @@ Split into domain-specific mixins for maintainability:
- ReviewRecordingMixin: Needs review recording for the scan and the monitor
- ReviewEventsMixin: the Needs review change stream an open page re-reads on
- SystemRoutesMixin: System health check, Overview and usage statistics
- MaintenanceRoutesMixin: reviewed maintenance plans (preview, apply)
- DuplicateFilesMixin: duplicate image files as find and removal jobs
"""
import asyncio
@@ -46,6 +48,9 @@ from .job_routes import JobRoutesMixin
from .image_tag_jobs import ImageTagJobsMixin
from .thumbnail_jobs import ThumbnailJobsMixin
from .system import SystemRoutesMixin
from .maintenance import MaintenanceRoutesMixin
from .duplicate_files import DuplicateFilesMixin
from .maintenance_plans import PlanStore
from .jobs import JobRegistry
from .logging_routes import LoggingRoutesMixin
from .autotag_routes import AutotagRoutesMixin
@@ -285,6 +290,8 @@ class PromptManagerAPI(
ImageTagJobsMixin,
ThumbnailJobsMixin,
SystemRoutesMixin,
MaintenanceRoutesMixin,
DuplicateFilesMixin,
LoggingRoutesMixin,
AutotagRoutesMixin,
LoraIntegrationMixin,
@@ -341,13 +348,27 @@ class PromptManagerAPI(
self._jobs = JobRegistry(self.logger)
# The latest metadata scan job (also in self._jobs), if one was started.
self._scan_job = None
# The current duplicate scan (a py/api/jobs.Job outside the registry).
self._duplicate_job = None
# Duplicate-file finds by job id, with the groups each one holds for
# review (py/api/duplicate_files.py); the last few only.
self._duplicate_finds = {}
# The last System health report this session and the database
# generation it describes, bumped when a restore replaces self.db
# (py/api/system.py).
self._health_report = None
self._db_generation = 0
# Reviewed maintenance plans under single-use tokens; a restart
# forgets them (py/api/maintenance.py).
self._maintenance_plans = PlanStore()
# A safety backup interrupted by a crash leaves its temp copy behind.
try:
from .maintenance import sweep_safety_temps
swept = sweep_safety_temps(self.db.model.db_path)
if swept:
self.logger.info(f"Removed {len(swept)} interrupted backup copies")
except Exception as e:
self.logger.warning(f"Could not sweep interrupted backup copies: {e}")
# Run cleanup on initialization to remove any existing duplicates
try:
@@ -508,6 +529,8 @@ class PromptManagerAPI(
self._register_settings_routes(routes)
self._register_job_routes(routes)
self._register_system_routes(routes)
self._register_maintenance_routes(routes)
self._register_duplicate_routes(routes)
self._register_logging_routes(routes)
self._register_autotag_routes(routes)
self._register_lora_routes(routes)
+235 -139
View File
@@ -24,9 +24,9 @@ from .jobs import (
SCAN_GROUPS,
SCAN_SCOPES,
DUPLICATE_CLEANUP_JOB_KIND,
DUPLICATES_JOB_KIND,
MAINTENANCE_JOB_KIND,
RESTORE_JOB_KIND,
Job,
JobBusy,
)
@@ -69,7 +69,9 @@ except ImportError:
IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp", ".tiff")
VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv", ".webm", ".m4v", ".wmv")
SCAN_BATCH_SIZE = 50 # Files per executor call during the output scan
MAX_SAFETY_BACKUPS = 10 # '<db>.backup_*' files kept next to the live database
# '<db>.backup_*' files kept next to the live database: one pool for the
# restore backups and the maintenance safety backups (py/api/maintenance.py).
MAX_SAFETY_BACKUPS = 10
MAX_SCAN_FILES = 50_000 # Upper bound on media files one scan will look at
SCAN_COUNT_KEYS = (
"processed",
@@ -90,6 +92,7 @@ MAX_SCAN_DEPTH = 12 # Directory levels below an output root that are scanned
# the common single-f TIFF spelling. Anything else inside the output
# directory (databases, scripts, notes) is never a "duplicate image".
DELETABLE_MEDIA_EXTENSIONS = frozenset(IMAGE_EXTENSIONS + VIDEO_EXTENSIONS + (".tif",))
DUPLICATE_FIND_LEGACY_TITLE = "Find duplicate image files (gallery page)"
def _collect_media_files(output_dirs, extensions):
@@ -132,11 +135,52 @@ def _collect_media_files(output_dirs, extensions):
return found
def _physical_file_key(path, stat):
"""What tells two physical files apart: ``(st_dev, st_ino)``.
A filesystem that reports no inode (``st_ino`` 0, some network or FAT
volumes on Windows) falls back to the resolved, case-folded path.
"""
if stat.st_ino:
return ("inode", stat.st_dev, stat.st_ino)
return ("path", os.path.normcase(os.path.realpath(str(path))))
def _thumbnail_rel_path(rel_path, thumbnail_ext):
"""Relative path of the thumbnail the gallery would generate for ``rel_path``."""
return f"thumbnails/{rel_path.with_suffix('').as_posix()}_thumb{thumbnail_ext}"
def _media_thumbnail_rel_path(rel_path):
"""The thumbnail of a media file: a video's preview is a ``.jpg``."""
extension = rel_path.suffix.lower()
return _thumbnail_rel_path(
rel_path, ".jpg" if extension in VIDEO_EXTENSIONS else rel_path.suffix
)
def _remove_thumbnail(output_path, rel_path, logger=None):
"""Remove the thumbnail of ``rel_path`` under ``output_path``, if any.
Shared by the legacy duplicate delete and the reviewed removal job. A
thumbnail that is a link, or one that cannot be removed, is left alone
(a stale preview never fails a removal). Returns whether one was removed.
"""
output_path = Path(output_path)
thumbnail_path = output_path / _media_thumbnail_rel_path(Path(rel_path))
try:
if thumbnail_path.is_symlink() or not thumbnail_path.is_file():
return False
os.remove(thumbnail_path)
except OSError as e:
if logger is not None:
logger.warning(f"Could not delete thumbnail for {rel_path}: {e}")
return False
if logger is not None:
logger.debug(f"Deleted associated thumbnail: {thumbnail_path}")
return True
async def _read_json(request):
"""Parse a JSON object body through the package's capped reader.
@@ -561,12 +605,21 @@ class AdminRoutesMixin:
LONG_JOB_DUPLICATES = "duplicates"
async def scan_duplicates_endpoint(self, request):
"""Scan for duplicate images without removing them (single flight)."""
job = self._acquire_long_job(self.LONG_JOB_DUPLICATES)
if job is None:
return self._long_job_busy_response("Duplicate scan")
async with job:
return await self._scan_duplicates(request)
"""Scan for duplicate images without removing them.
Holds the ``duplicates`` kind on the registry (a ``read`` hold, so it
is listed in Activity), so it is single flight with every duplicate
find, and it is refused while a duplicate removal runs.
"""
cleanup = self._jobs.active(DUPLICATE_CLEANUP_JOB_KIND)
if cleanup is not None:
return self._long_job_busy_response(cleanup.title)
return await self._run_held(
DUPLICATES_JOB_KIND,
DUPLICATE_FIND_LEGACY_TITLE,
"read",
functools.partial(self._scan_duplicates, request),
)
async def _scan_duplicates(self, request):
try:
@@ -588,24 +641,29 @@ class AdminRoutesMixin:
)
async def scan_duplicates_stream(self, request):
"""Run the duplicate scan as a background job and stream its progress.
"""Stream a duplicate find to the legacy gallery (its event shapes).
A request while a scan is running attaches to it. The job holds the
duplicates lock, so a delete is refused (409) until it finishes.
Starts the registry's ``duplicates`` job over the gallery's output
folder, or attaches to the running one when it covers that folder
(one started in the workspace, for instance). A held legacy scan, a
find that does not cover the folder or a running removal answers the
legacy 409. The job outlives the request.
"""
job = self._duplicate_job
if job is None or job.finished:
guard = self._acquire_long_job(self.LONG_JOB_DUPLICATES)
if guard is None:
return self._long_job_busy_response("Duplicate scan or delete")
job = Job("duplicates", "Duplicate scan")
self._duplicate_job = job
job.task = asyncio.ensure_future(self._run_duplicate_job(job, guard))
job = self._jobs.active(DUPLICATES_JOB_KIND)
if job is None:
try:
job = self._start_duplicate_find(origin="gallery")
except JobBusy as busy:
return self._long_job_busy_response(busy.running.title)
else:
legacy_dir = await self._run_in_executor(self._find_comfyui_output_dir)
if not self._find_covers(job, legacy_dir):
return self._long_job_busy_response(job.title)
return await self._stream_scan_job(request, job)
async def scan_duplicates_status(self, request):
"""Whether a duplicate scan is running, and its latest event."""
job = self._duplicate_job
"""Whether a duplicate find is running, and its latest legacy event."""
job = self._latest_duplicate_find()
return web.json_response(
{
"success": True,
@@ -614,62 +672,6 @@ class AdminRoutesMixin:
}
)
async def _run_duplicate_job(self, job, guard):
loop = asyncio.get_running_loop()
counts = {"processed": 0}
def report(done, total):
# Called from worker threads; hand the event to the loop.
counts["processed"] = done
pct = int(done / total * 100) if total else 100
loop.call_soon_threadsafe(
job.publish,
{
"type": "progress",
"progress": pct,
"status": f"Hashing file {done}/{total}...",
"processed": done,
"found": 0,
},
)
try:
async with guard:
job.publish(
{
"type": "progress",
"progress": 0,
"status": "Collecting media files...",
"processed": 0,
"found": 0,
}
)
output_dir = await self._run_in_executor(self._find_comfyui_output_dir)
groups = await self._run_in_executor(
self._find_duplicate_images_sync, output_dir, report
)
job.publish(
{
"type": "complete",
"processed": counts["processed"],
"found": len(groups),
"duplicates": groups,
"message": f"Found {len(groups)} groups of duplicate images",
}
)
except Exception:
self.logger.exception("Duplicate scan error")
job.publish(
{
"type": "error",
"message": (
"An internal error occurred. Check server logs for details."
),
}
)
finally:
job.finish()
async def cleanup_duplicates_endpoint(self, request):
"""Cleanup duplicate prompts endpoint (holds the registry: records)."""
return await self._run_held(
@@ -706,31 +708,109 @@ class AdminRoutesMixin:
def _find_duplicate_images_sync(self, output_dir, progress=None):
"""Group media files under ``output_dir`` by content hash (blocking).
Files are hashed on ``PromptManagerConfig.WORKER_THREADS`` threads.
``progress(done, total)`` is called from those threads about every
percent and once at the end.
The legacy response shape over ``_hash_media_groups``, which the
workspace's duplicate find shares.
"""
from ..config import PromptManagerConfig
if not output_dir or not os.path.isdir(output_dir):
self.logger.warning("ComfyUI output directory not found")
return []
groups = self._hash_media_groups([os.path.abspath(output_dir)], progress)
duplicates = self._legacy_duplicate_groups(groups, output_dir)
self.logger.info(f"Found {len(duplicates)} groups of duplicate images")
return duplicates
def _legacy_duplicate_groups(self, groups, output_dir):
"""``_hash_media_groups`` groups as the legacy gallery reads them.
Only the files under ``output_dir`` count (paths are relative to it);
a group needs two of them. A file that vanished since it was hashed
is left out.
"""
output_path = Path(os.path.abspath(output_dir))
base = os.path.normcase(str(output_path)).rstrip("\\/") + os.sep
duplicates = []
for group in groups:
images = []
for record in group["files"]:
if not os.path.normcase(record["abs"]).startswith(base):
continue
try:
images.append(
self._describe_media_file(
Path(record["abs"]), output_path, file_hash=group["hash"]
)
)
except (OSError, ValueError) as e:
self.logger.error(
f"Error describing {Path(record['abs']).name}: {e}"
)
if len(images) > 1:
duplicates.append(
{"hash": group["hash"], "images": images, "count": len(images)}
)
return duplicates
def _hash_media_groups(self, roots, progress=None, stop=None):
"""Content-identical media files under ``roots`` (blocking).
Each root is walked as ``_collect_media_files`` walks (no symlinks,
no ``thumbnails``, bounded). A physical file is listed once, under
the first root that holds it: files are told apart by ``(st_dev,
st_ino)``, so a source listed twice under two spellings (a symlink, a
bind mount, a case variant), a source inside another one or a
hardlink is one file, never a copy of itself. At most MAX_SCAN_FILES
files are read. Files
are hashed on ``PromptManagerConfig.WORKER_THREADS`` threads;
``progress(done, total)`` is called from those threads about every
percent and once at the end, and ``stop()`` true skips the files not
yet read. Returns ``[{hash, size, files:[{slot, abs, rel, size,
mtime, mtime_ns, inode}]}]`` for every hash shared by two or more
files (``inode`` is ``[st_dev, st_ino]``), ``slot``
being the root's place in ``roots`` and the files oldest first.
"""
from ..config import PromptManagerConfig
output_path = Path(output_dir)
extensions = IMAGE_EXTENSIONS + VIDEO_EXTENSIONS
media_files = _collect_media_files([output_path], extensions)
media_files = []
seen = set()
for slot, root in enumerate(roots):
for path in _collect_media_files([root], extensions):
try:
# os.stat, not Path.stat: one stat call for every check
# (Python 3.10's Path.stat bypasses a patched os.stat).
stat = os.stat(path)
except OSError:
continue
key = _physical_file_key(path, stat)
if key not in seen:
seen.add(key)
media_files.append((slot, root, path, stat))
if len(media_files) >= MAX_SCAN_FILES:
break
media_files = media_files[:MAX_SCAN_FILES]
total = len(media_files)
self.logger.info(f"Found {total} media files to analyze")
step = max(1, total // 100)
counter = {"done": 0}
lock = threading.Lock()
def describe(media_path):
try:
info = self._describe_media_file(media_path, output_path)
except Exception as e:
self.logger.error(f"Error processing file {media_path.name}: {e}")
info = None
def describe(item):
slot, root, media_path, stat = item
info = None
if stop is None or not stop():
try:
info = {
"slot": slot,
"abs": str(media_path),
"rel": Path(os.path.relpath(str(media_path), root)).as_posix(),
"size": stat.st_size,
"mtime": stat.st_mtime,
"mtime_ns": stat.st_mtime_ns,
"inode": [stat.st_dev, stat.st_ino],
"hash": self._calculate_file_hash(media_path),
}
except Exception as e:
self.logger.error(f"Error processing file {media_path.name}: {e}")
with lock:
counter["done"] += 1
done = counter["done"]
@@ -744,24 +824,22 @@ class AdminRoutesMixin:
infos = map_parallel(describe, media_files, PromptManagerConfig.WORKER_THREADS)
file_hashes = {}
by_hash = {}
for info in infos:
if info is not None:
file_hashes.setdefault(info["hash"], []).append(info)
duplicates = []
for file_hash, images in file_hashes.items():
if len(images) > 1:
images.sort(key=lambda x: x["modified_time"])
duplicates.append(
{"hash": file_hash, "images": images, "count": len(images)}
by_hash.setdefault(info.pop("hash"), []).append(info)
groups = []
for file_hash, files in by_hash.items():
if len(files) > 1:
files.sort(key=lambda record: (record["mtime"], record["abs"]))
groups.append(
{"hash": file_hash, "size": files[0]["size"], "files": files}
)
return groups
self.logger.info(f"Found {len(duplicates)} groups of duplicate images")
return duplicates
def _describe_media_file(self, media_path, output_path):
"""Hash one media file and describe it for the duplicates response."""
def _describe_media_file(self, media_path, output_path, file_hash=None):
"""Describe one media file for the duplicates response (hashing it
unless ``file_hash`` is given)."""
stat = media_path.stat()
rel_path = media_path.relative_to(output_path)
extension = media_path.suffix.lower()
@@ -782,7 +860,7 @@ class AdminRoutesMixin:
"extension": extension,
"media_type": "video" if is_video else "image",
"is_video": is_video,
"hash": self._calculate_file_hash(media_path),
"hash": file_hash or self._calculate_file_hash(media_path),
}
def _calculate_file_hash(self, file_path):
@@ -796,9 +874,9 @@ class AdminRoutesMixin:
async def delete_duplicate_images_endpoint(self, request):
"""Delete duplicate image files from disk.
Single flight with the duplicate scan, and a records hold on the
registry, so no scan, restore or tagging job runs meanwhile. A body
without paths is refused before the hold (it is no job).
Refused while a duplicate find runs, and a records hold on the
registry, so no scan, restore, removal or tagging job runs meanwhile.
A body without paths is refused before the hold (it is no job).
"""
data, error_response = await _read_json(request)
if error_response is not None:
@@ -809,20 +887,14 @@ class AdminRoutesMixin:
{"success": False, "error": "No image paths provided"},
status=400,
)
guard = self._acquire_long_job(self.LONG_JOB_DUPLICATES)
if guard is None:
return self._long_job_busy_response("Duplicate scan or delete")
async def delete():
# No await since the check above: the lock is still free.
async with guard:
return await self._delete_duplicate_images(image_paths)
find = self._jobs.active(DUPLICATES_JOB_KIND)
if find is not None:
return self._long_job_busy_response(find.title)
return await self._run_held(
DUPLICATE_CLEANUP_JOB_KIND,
"Duplicate image delete (gallery page)",
"records",
delete,
functools.partial(self._delete_duplicate_images, image_paths),
)
async def _delete_duplicate_images(self, image_paths):
@@ -902,14 +974,7 @@ class AdminRoutesMixin:
os.remove(file_path)
self.logger.info(f"Deleted duplicate image: {image_path}")
thumbnail_path = output_path / _thumbnail_rel_path(rel_path, file_path.suffix)
try:
if thumbnail_path.exists():
os.remove(thumbnail_path)
self.logger.debug(f"Deleted associated thumbnail: {thumbnail_path}")
except OSError as e:
self.logger.warning(f"Could not delete thumbnail for {image_path}: {e}")
_remove_thumbnail(output_path.resolve(), rel_path, self.logger)
return None
async def get_statistics(self, request):
@@ -1112,8 +1177,10 @@ class AdminRoutesMixin:
db.cleanup_duplicates(), "duplicate prompts"
),
"vacuum": vacuum,
# Only under sources that can be read now (never an unmounted drive).
"cleanup_orphaned_images": lambda: count_result(
db.cleanup_missing_images(), "orphaned image records"
db.cleanup_missing_images(self._health_roots()),
"orphaned image records",
),
"check_hash_duplicates": lambda: self._hash_duplicates_result(
db.check_hash_duplicates()
@@ -1207,7 +1274,11 @@ class AdminRoutesMixin:
async def restore_database(self, request):
"""Restore the prompts database from an uploaded SQLite file.
The upload is streamed to a temp file (capped at restore_max_bytes),
The upload is streamed to a temp file. Free disk space decides how
large it may be (``maintenance.restore_room``), under a 64 GiB ceiling
(``restore_max_bytes``): a request whose Content-Length does not fit
is refused before anything is written, and one without is refused as
soon as its running size does not. The file is
verified with PromptModel.verify_database_file and only then copied
over the live database, which is backed up first. The whole request
holds the registry (records): it is refused while a scan or another
@@ -1220,10 +1291,44 @@ class AdminRoutesMixin:
lambda: self._restore_database(request),
)
def _restore_refusal(self, size, ceiling, room, *, known):
"""``400 too-large`` past the ceiling, else ``507 insufficient-space``."""
from .maintenance import restore_space_message
if size > ceiling:
return web.json_response(
{
"success": False,
"code": "too-large",
"error": (
"File too large. Maximum size is "
f"{ceiling // (1024 * 1024)}MB"
),
},
status=400,
)
return web.json_response(
{
"success": False,
"code": "insufficient-space",
"error": restore_space_message(size, room, known=known),
},
status=507,
)
async def _restore_database(self, request):
max_bytes = getattr(self, "restore_max_bytes", 100 * 1024 * 1024)
from .maintenance import RESTORE_MAX_BYTES, restore_room
ceiling = getattr(self, "restore_max_bytes", RESTORE_MAX_BYTES)
temp_path = None
try:
# Free disk space decides how large a backup may be (#313): the
# upload, its staged copy and the swap must all fit.
room = await self._run_in_executor(restore_room, self.db.model.db_path)
max_bytes = min(ceiling, room)
length = request.content_length
if length is not None and length > max_bytes:
return self._restore_refusal(length, ceiling, room, known=True)
reader = await request.multipart()
field = await reader.next()
@@ -1249,16 +1354,7 @@ class AdminRoutesMixin:
break
size += len(chunk)
if size > max_bytes:
return web.json_response(
{
"success": False,
"error": (
"File too large. Maximum size is "
f"{max_bytes // (1024 * 1024)}MB"
),
},
status=400,
)
return self._restore_refusal(size, ceiling, room, known=False)
out.write(chunk)
if size == 0:
+779
View File
@@ -0,0 +1,779 @@
"""Duplicate image files on the job registry (#249 System S6, Task 17).
- ``POST /prompt_manager/jobs/duplicates`` starts the find (202 ``{job}``):
kind ``duplicates``, "Find duplicate image files", admission ``read``,
cancellable, outcome group ``groups_found``. It hashes the media files of
every effective source (``_hash_media_groups``, shared with the legacy
scan) and the server keeps the groups it found under the job id: content
hash, root-relative paths and each file's source index. A find is refused
while a duplicate removal runs, and one find runs at a time.
- ``GET /prompt_manager/maintenance/duplicates/{job_id}?offset&limit`` pages
those groups, at most ``MAX_GROUP_PAGE`` at a time: ``{success, job:{id,
state}, total, offset, limit, groups:[{group_id, size, count, files}]}``.
Each file names its root-relative ``path``, its source's public label and
index and the Images serve URLs; never an absolute path.
- ``POST /prompt_manager/maintenance/preview {operation:"duplicate_files",
options:{job_id, keep:[{group_id, root_index, path}]}}`` reviews a removal:
in each group named, the files kept are the ones listed and every other
copy is removed. Each keeper and each removal is checked against the held
group, with containment decided as text before any filesystem access
(``file_problem``): inside its source, a media file outside
``thumbnails``, then a regular file (not a link) whose resolved path is
still inside the source, with the same size and content hash. A removal
that fails, or whose kept copy fails, is protected and listed with the
reason. The plan carries the files' database links, which move to the
kept copy.
- ``POST /prompt_manager/maintenance/apply {token}`` runs the reviewed plan
as the ``duplicate_cleanup`` job: admission ``records``, not cancellable,
groups ``removed``, ``kept`` and ``errors``. It checks the sources are
unchanged, takes the safety backup the policy promises (the database only;
removed files are not in it), then for each removal rechecks its kept copy
and the file itself (``file_problem``) and removes it with every handle
closed (``_unlink``), together with its thumbnail. A file that changed is
kept; the first unlink that fails stops the job, a locked file with the
coded ``LOCKED_FILE`` message (never the OS text), and the files after it
are not attempted. The database links of the files actually removed then
move to their kept copy (``PromptDatabase.move_file_links``).
The legacy gallery routes (``/scan_duplicates``, its stream and status, and
``/delete_duplicate_images``) keep their shapes on the same registry kinds;
see ``py/api/admin.py``.
"""
import asyncio
import os
from pathlib import Path
from aiohttp import web
from .admin import DUPLICATE_FIND_LEGACY_TITLE, _remove_thumbnail
from .duplicate_review import ( # noqa: F401 (CHANGED, OUTSIDE_SOURCE: public)
CHANGED,
DELETE_FAILED,
FIND_CANCELLED,
FIND_ERROR_EVENT,
FIND_FAILED,
GONE,
KEEPER,
KEEPER_CHANGED,
KEEPER_GONE,
KEEPER_WENT,
KEPT_COPY,
LOCKED_FILE,
MAYBE_SAME_FILE,
NO_FILE_ID_ROW,
NOT_ATTEMPTED,
OUTSIDE_SOURCE,
RECORDS_FAILED,
ROW_DETAILS,
SAME_FILE,
SOURCES_CHANGED,
SOURCES_CHANGED_REVIEW,
HeldFind,
SourcesChanged,
_fingerprint,
_group_order,
_identity,
_parse_options,
_plural,
_public_plan,
_select,
file_problem,
human_bytes,
still_kept,
)
from .duplicate_spellings import Spellings, copy_move
from .images import _favorite_root_key, _output_image_entry
from .job_routes import busy_response
from .jobs import (
DUPLICATE_CLEANUP_GROUP_NAMES,
DUPLICATE_CLEANUP_JOB_KIND,
DUPLICATE_GROUP_NAMES,
DUPLICATES_JOB_KIND,
JobBusy,
)
from .maintenance import DUPLICATE_FILES, _fail_job, _failure, _invalid, backup_note
try:
from ...database.duplicate_copies import FILE_LINK_SQL
except ImportError:
from database.duplicate_copies import FILE_LINK_SQL
DUPLICATE_FIND_TITLE = "Find duplicate image files"
DUPLICATE_CLEANUP_TITLE = "Remove duplicate image files"
MAX_GROUP_PAGE = 10 # groups a page lists (D69: ten at a time)
MAX_HELD_FINDS = 3 # finds whose groups the server keeps for review
favorite_root_key = _favorite_root_key
def query_plan_sql():
"""The link reconciliation's statements (EXPLAIN QUERY PLAN checks)."""
return list(FILE_LINK_SQL)
def _unlink(path):
"""Remove one file (the caller closed every handle on it)."""
os.remove(path)
class DuplicateFilesMixin:
"""Mixin providing the duplicate-file find, its groups and its removal."""
def _register_duplicate_routes(self, routes):
@routes.post("/prompt_manager/jobs/duplicates")
async def start_duplicate_find_route(request):
return await self.start_duplicate_find(request)
@routes.get("/prompt_manager/maintenance/duplicates/{job_id}")
async def duplicate_groups_route(request):
return await self.duplicate_groups(request)
# ── find ──────────────────────────────────────────────────────────
async def start_duplicate_find(self, request):
"""Start the duplicate find (202), or 409 naming the job in the way."""
try:
job = self._start_duplicate_find(origin="workspace")
except JobBusy as busy:
return busy_response(busy)
return web.json_response({"success": True, "job": job.summary()}, status=202)
def _start_duplicate_find(self, origin):
"""Register a find job (on the loop); raises JobBusy.
A workspace find walks every effective source; a gallery one the
gallery's output folder, as the legacy scan did. Both resolve their
folders when the job begins, so a source change later cannot move it.
"""
cleanup = self._jobs.active(DUPLICATE_CLEANUP_JOB_KIND)
if cleanup is not None:
raise JobBusy(cleanup)
held = HeldFind(origin)
if origin == "workspace":
title = DUPLICATE_FIND_TITLE
else:
title = DUPLICATE_FIND_LEGACY_TITLE
job = self._jobs.start(
DUPLICATES_JOB_KIND,
title,
admission="read",
cancellable=True,
groups=DUPLICATE_GROUP_NAMES,
runner=lambda job: self._run_duplicate_find(job, held),
)
held.job_id = job.id
self._remember_find(held)
return job
def _remember_find(self, held):
finds = self._duplicate_finds
finds[held.job_id] = held
for job_id in list(finds):
if self._jobs.get(job_id) is None:
del finds[job_id]
while len(finds) > MAX_HELD_FINDS:
del finds[next(iter(finds))]
def _find_covers(self, job, legacy_dir):
held = self._duplicate_finds.get(job.id)
return held is not None and held.covers(legacy_dir)
def _latest_duplicate_find(self):
"""The running find, else the newest finished one this session."""
for job in self._jobs.list():
if job.kind == DUPLICATES_JOB_KIND and job.id in self._duplicate_finds:
return job
return None
def _resolve_find_roots(self, held):
"""Fill in where a find looks (blocking): sources and the gallery folder."""
listed = [os.path.abspath(str(d)) for d in self._get_all_output_dirs()]
held.legacy_dir = self._find_comfyui_output_dir()
if held.origin == "workspace":
held.roots = list(enumerate(listed))
return
if not held.legacy_dir or not os.path.isdir(held.legacy_dir):
self.logger.warning("ComfyUI output directory not found")
held.roots = []
return
legacy = os.path.abspath(held.legacy_dir)
index = next(
(i for i, d in enumerate(listed) if _identity(d) == _identity(legacy)),
None,
)
held.roots = [(index, legacy)]
def _find_progress(self, job, done, total):
if done < job.processed:
return # hashing threads report out of order; keep the count rising
pct = int(done / total * 100) if total else 100
job.update(
processed=done,
total=total,
status_line=f"Hashing file {done:,} of {total:,}",
)
job.publish(
{
"type": "progress",
"progress": pct,
"status": f"Hashing file {done}/{total}...",
"processed": done,
"found": 0,
}
)
async def _run_duplicate_find(self, job, held):
loop = asyncio.get_running_loop()
def report(done, total): # from the hashing threads
loop.call_soon_threadsafe(self._find_progress, job, done, total)
try:
job.update(status_line="Looking for media files")
job.publish(
{
"type": "progress",
"progress": 0,
"status": "Collecting media files...",
"processed": 0,
"found": 0,
}
)
await self._run_in_executor(self._resolve_find_roots, held)
job.scope["output_folders"] = len(held.roots)
job.log_line(f"Searching {_plural(len(held.roots), 'source')}")
groups = await self._run_in_executor(
self._hash_media_groups,
[root for _index, root in held.roots],
report,
lambda: job.cancelled,
)
if job.cancelled:
job.update(status_line="Cancelled; nothing was changed.")
job.publish({"type": "error", "message": FIND_CANCELLED})
job.finish("cancelled")
return
await self._finish_find(job, held, groups)
except Exception:
self.logger.exception("Duplicate search failed")
job.update(status_line=FIND_FAILED)
job.publish({"type": "error", "message": FIND_ERROR_EVENT})
job.finish("failed")
async def _finish_find(self, job, held, groups):
for group in groups:
group["group_id"] = group["hash"]
for record in group["files"]:
index, root = held.roots[record.pop("slot")]
record["root_index"] = index
record["root"] = root
groups.sort(key=_group_order)
held.groups = groups
labels = {root: self._public_path(root) for _index, root in held.roots}
rows = [
{
"path": group["files"][0]["rel"],
"root": labels[group["files"][0]["root"]],
"root_index": group["files"][0]["root_index"],
"detail": f"{len(group['files'])} identical files · "
f"{human_bytes(group['size'])} each",
}
for group in groups
]
job.record_counted("groups_found", len(groups), rows)
line = f"Found {_plural(len(groups), 'group')} of duplicate files"
job.log_line(line)
legacy = []
if held.covers(held.legacy_dir) and held.legacy_dir:
legacy = await self._run_in_executor(
self._legacy_duplicate_groups, groups, held.legacy_dir
)
# Every file was read by now (the last reports may still be queued).
job.update(progress=100, status_line=line, processed=job.total)
job.publish(
{
"type": "complete",
"processed": job.processed,
"found": len(legacy),
"duplicates": legacy,
"message": f"Found {len(legacy)} groups of duplicate images",
}
)
# ── groups ────────────────────────────────────────────────────────
def _held_result(self, job_id):
"""``(held, None)`` for a finished find's groups, else an error reply."""
held = self._duplicate_finds.get(job_id)
job = self._jobs.get(job_id)
if held is None or job is None or job.kind != DUPLICATES_JOB_KIND:
return None, _failure(
"No duplicate search with this id. Search again.", 404, "not-found"
)
if job.running:
return None, _failure(
"The duplicate search is still running.", 409, "job-running"
)
if held.groups is None:
return None, _failure(
"This search did not finish. Search again.", 409, "no-groups"
)
return held, None
async def duplicate_groups(self, request):
"""One page of a finished find's groups (at most MAX_GROUP_PAGE)."""
try:
offset = int(request.query.get("offset", "0"))
limit = int(request.query.get("limit", str(MAX_GROUP_PAGE)))
except ValueError:
return _invalid("offset and limit must be whole numbers")
if offset < 0 or not 0 < limit <= MAX_GROUP_PAGE:
return _invalid(f"limit must be 1 to {MAX_GROUP_PAGE}; offset 0 or more")
job_id = request.match_info["job_id"]
held, error = self._held_result(job_id)
if error is not None:
return error
page = held.groups[offset : offset + limit]
groups = await self._run_in_executor(self._describe_groups, page)
job = self._jobs.get(job_id)
return web.json_response(
{
"success": True,
"job": {"id": job.id, "state": job.state},
"total": len(held.groups),
"offset": offset,
"limit": limit,
"groups": groups,
},
headers={"Cache-Control": "no-store"},
)
def _describe_groups(self, groups):
"""Public group rows (blocking: one stat per file)."""
return [
{
"group_id": group["group_id"],
"size": group["size"],
"count": len(group["files"]),
"files": [self._describe_copy(record) for record in group["files"]],
}
for group in groups
]
def _describe_copy(self, record):
label = self._public_path(record["root"])
base = {
"filename": os.path.basename(record["rel"]),
"path": record["rel"],
"root": label,
"root_index": record["root_index"],
"url": None,
"thumbnail_url": None,
"size": record["size"],
"modified_time": record["mtime"],
"media_type": "image",
"missing": False,
}
if record["root_index"] is None:
return base
try:
entry = _output_image_entry(
Path(record["abs"]), Path(record["root"]), record["root_index"], label
)
except (OSError, ValueError):
return {**base, "missing": True}
return {
**base,
"url": entry["url"],
"thumbnail_url": entry["thumbnail_url"],
"size": entry["size"],
"modified_time": entry["modified_time"],
"media_type": entry["media_type"],
}
# ── review ────────────────────────────────────────────────────────
def _listed_sources(self):
"""``{identity: index}`` of the effective sources now (blocking)."""
return {
_identity(d): index for index, d in enumerate(self._get_all_output_dirs())
}
@staticmethod
def _source_listed(listed, index, root):
return index is not None and listed.get(_identity(root)) == index
async def preview_duplicate_files(self, options):
"""The ``duplicate_files`` plan for one finished find and its keepers."""
job_id, keep, message = _parse_options(options)
if message:
return _invalid(message)
held, error = self._held_result(job_id)
if error is not None:
return error
selection, message = _select(held, keep)
if message:
return _invalid(message)
db = self.db
identity = self._database_identity(db)
try:
review = await self._run_in_executor(
self._review_duplicate_files, db, selection
)
except SourcesChanged:
return _failure(SOURCES_CHANGED_REVIEW, 409, "sources-changed")
except Exception:
self.logger.exception("Duplicate file review failed")
return _failure("Failed to review maintenance", 500)
return await self._plan_response(
db,
identity,
DUPLICATE_FILES,
review["fingerprint"],
review["public"],
{
"targets": review["targets"],
"kept": review["kept"],
"protected": review["public"]["protected"]["count"],
},
)
def _review_duplicate_files(self, db, selection):
"""Check every keeper and removal now (blocking); the plan's parts.
No file is read: each copy is bound to what the search saw by its
size, mtime and inode, and never to its kept copy itself. The removal
job hashes each file again right before it unlinks it.
"""
listed = self._listed_sources()
targets, kept, protected = [], [], {}
for group, keepers in selection.values():
removals = [r for r in group["files"] if r not in keepers]
for record in removals:
if record["root_index"] is not None and not self._source_listed(
listed, record["root_index"], record["root"]
):
raise SourcesChanged(record["root"])
keeper = next(
(k for k in keepers if _bound_problem(k, group) is None), None
)
reviewed = []
for record in removals:
if keeper is None:
problem = KEEPER
elif record["root_index"] is None:
problem = OUTSIDE_SOURCE
else:
problem = _bound_problem(record, group, keeper)
if problem:
protected[problem] = protected.get(problem, 0) + 1
continue
reviewed.append(self._target(group, record, keeper))
if reviewed:
targets.extend(reviewed)
kept.append(self._copy_row(keeper, KEPT_COPY))
records = {}
if targets:
spellings = self._source_spellings()
records = db.count_copy_records(
{
t["abs"]: {
"spellings": list(spellings.of(t["abs"]).values()),
"locator": t["locator"],
"sha256": t["hash"],
}
for t in targets
}
)
return {
"targets": targets,
"kept": kept,
"fingerprint": _fingerprint(targets),
"public": _public_plan(targets, protected, records),
}
def _source_spellings(self):
"""The folder spellings records may name files by (blocking).
The search walks each source's real path, while a record names a file
as it was seen when written: under the folder the image monitor
watches, the configured spelling, ComfyUI's output folder, a symlink
to it, or a parent of a source added later. ``Spellings`` matches
them to each file by containment (``py/api/duplicate_spellings.py``).
"""
from ..config import GalleryConfig
listed = [os.path.abspath(str(d)) for d in self._get_all_output_dirs()]
candidates = list(self._monitored_directories())
candidates += list(GalleryConfig.MONITORING_DIRECTORIES)
try:
import folder_paths
candidates.append(folder_paths.get_output_directory())
except Exception: # nosec B110 - outside ComfyUI there is none
pass
candidates.append(self._find_comfyui_output_dir())
folders = [
os.path.abspath(str(c)) for c in candidates if c and os.path.isabs(str(c))
]
return Spellings(folders + listed, listed)
def _target(self, group, record, keeper):
return {
"group_id": group["group_id"],
"hash": group["hash"],
"size": record["size"],
"root_index": record["root_index"],
"root": record["root"],
"root_label": self._public_path(record["root"]),
"abs": record["abs"],
"rel": record["rel"],
"keeper_abs": keeper["abs"],
"keeper_rel": keeper["rel"],
"keeper_root": keeper["root"],
"keeper_size": keeper["size"],
"locator": os.path.normcase(record["abs"]),
"keeper_locator": os.path.normcase(keeper["abs"]),
}
def _copy_row(self, record, detail):
return {
"path": record["rel"],
"root": self._public_path(record["root"]),
"root_index": record["root_index"],
"detail": detail,
}
# ── removal ───────────────────────────────────────────────────────
def _start_duplicate_cleanup(self, db, plan, backup_on):
"""Register the removal job for a reviewed plan (on the loop)."""
find = self._jobs.active(DUPLICATES_JOB_KIND)
if find is not None:
raise JobBusy(find)
return self._jobs.start(
DUPLICATE_CLEANUP_JOB_KIND,
DUPLICATE_CLEANUP_TITLE,
admission="records",
cancellable=False,
scope={"operation": DUPLICATE_FILES},
groups=DUPLICATE_CLEANUP_GROUP_NAMES,
runner=lambda job: self._run_duplicate_cleanup(job, db, plan, backup_on),
)
def _sources_unchanged(self, targets):
listed = self._listed_sources()
return all(
self._source_listed(listed, t["root_index"], t["root"]) for t in targets
)
async def _run_duplicate_cleanup(self, job, db, plan, backup_on):
targets = plan.payload["targets"]
unchanged = await self._run_in_executor(self._sources_unchanged, targets)
if not unchanged:
return _fail_job(job, "sources-changed", SOURCES_CHANGED)
ok, backup = await self._safety_backup_step(job, db, backup_on)
if not ok:
return
try:
await self._remove_duplicates(job, db, plan, backup)
finally:
if backup:
await self._prune_backup_pool(job, db)
async def _remove_duplicates(self, job, db, plan, backup):
targets = plan.payload["targets"]
for row in plan.payload["kept"]:
job.record("kept", row, f"= {row['root']}/{row['path']} · {KEPT_COPY}")
job.update(total=len(targets), status_line="Removing duplicate files")
spellings = await self._run_in_executor(self._source_spellings)
removed, failure, keepers = 0, None, {}
totals = dict.fromkeys(("moved", "merged", "favorites", "tags", "review"), 0)
for done, target in enumerate(targets, start=1):
if failure is not None:
self._record_target(job, "kept", target, NOT_ATTEMPTED)
continue
outcome, detail, counts = await self._run_in_executor(
self._remove_one_duplicate, db, target, keepers, spellings
)
for key, value in (counts or {}).items():
totals[key] += value
if outcome == "removed":
removed += 1
text = f"A copy of {target['keeper_rel']}"
self._record_target(job, "removed", target, text, "-")
elif outcome == "kept":
self._record_target(job, "kept", target, ROW_DETAILS[detail])
else:
failure = detail
self._record_target(job, "errors", target, detail, "!")
job.update(processed=done, progress=int(done / len(targets) * 100))
if removed:
self.invalidate_gallery_cache()
self._log_moved_records(job, totals)
line = f"Removed {_plural(removed, 'duplicate file')}"
if removed < len(targets):
line += f", {len(targets) - removed:,} not removed"
job.log_line(line)
if failure is not None:
line = f"Stopped after an error · {line}"
job.update(progress=100, status_line=f"{line} · {backup_note(backup)}")
if failure is not None:
job.publish({"type": "error", "code": "delete-failed", "message": line})
job.finish("failed")
@staticmethod
def _log_moved_records(job, totals):
links = totals["moved"] + totals["merged"]
if links:
job.log_line(f"Moved {_plural(links, 'database link')} to the kept copies")
if totals["favorites"]:
job.log_line(
f"Moved {_plural(totals['favorites'], 'favorite')} to the kept copies"
)
if totals["tags"]:
job.log_line(
f"Added {_plural(totals['tags'], 'image tag')} to the kept copies"
)
if totals["review"]:
job.log_line(
f"Settled {_plural(totals['review'], 'Needs review item')} "
"of removed copies"
)
@staticmethod
def _record_target(job, group, target, detail, mark="="):
row = {
"path": target["rel"],
"root": target["root_label"],
"root_index": target["root_index"],
"detail": detail,
}
job.record(
group, row, f"{mark} {target['root_label']}/{target['rel']} · {detail}"
)
def _remove_one_duplicate(self, db, target, keepers, spellings):
"""``(outcome, detail, counts)`` for one reviewed removal (blocking).
The kept copy (once per group; cached in ``keepers``) and the file
are checked again, hashes included, and the file must not be its kept
copy. Then the file's records move to the kept copy in one short
transaction, and only then is the file unlinked: a crash or a failed
unlink leaves its records on the identical kept copy, never on a
deleted file. A failed move leaves the file and its records alone.
"""
hasher = self._calculate_file_hash
group_id = target["group_id"]
if group_id not in keepers:
keepers[group_id] = _kept_copy_state(target, hasher)
if keepers[group_id] is None:
return "kept", KEEPER, None
problem = file_problem(
target["abs"],
target["root"],
target["hash"],
target["size"],
hasher,
keeper=target["keeper_abs"],
keeper_root=target["keeper_root"],
)
if problem:
return "kept", problem, None
try:
# Held open until the unlink: on Windows, deleting the kept copy
# itself through another path then fails instead of removing it.
hold = open(target["keeper_abs"], "rb")
except OSError:
return "kept", KEEPER, None
with hold:
return self._move_then_unlink(
db, target, keepers[group_id], spellings, hold
)
def _move_then_unlink(self, db, target, kept_stat, spellings, hold):
try:
counts = db.move_copy_records(copy_move(target, kept_stat, spellings))
except Exception:
self.logger.exception(f"Could not move the records of {target['rel']}")
return "error", RECORDS_FAILED, None
if not still_kept(target, hold):
# The kept copy went away or changed while the records moved:
# keep this file and give its records back to it.
return "kept", KEEPER_GONE, self._move_back(db, target, spellings)
try:
_unlink(target["abs"])
except FileNotFoundError:
return "removed", None, counts
except PermissionError:
self.logger.warning(
f"Duplicate file is locked: {target['rel']}; "
"its records now name the kept copy"
)
return "error", LOCKED_FILE, counts
except OSError:
self.logger.warning(
f"Could not delete duplicate file {target['rel']}; "
"its records now name the kept copy",
exc_info=True,
)
return "error", DELETE_FAILED, counts
self.logger.info(f"Deleted duplicate file: {target['rel']}")
_remove_thumbnail(Path(target["root"]), Path(target["rel"]), self.logger)
return "removed", None, counts
def _move_back(self, db, target, spellings):
"""Give the kept copy's records to the file that stays (blocking)."""
swapped = {
**target,
"abs": target["keeper_abs"],
"root": target["keeper_root"],
"rel": target["keeper_rel"],
"locator": target["keeper_locator"],
"keeper_abs": target["abs"],
"keeper_root": target["root"],
"keeper_rel": target["rel"],
"keeper_locator": target["locator"],
}
try:
return db.move_copy_records(
copy_move(swapped, os.stat(target["abs"]), spellings)
)
except Exception:
self.logger.exception(f"Could not give the records back to {target['rel']}")
return None
def _bound_problem(record, group, keeper=None):
"""``file_problem`` without reading the file: size, mtime and inode.
``keeper`` is the kept copy's record, never offered as a copy of itself.
"""
return file_problem(
record["abs"],
record["root"],
group["hash"],
record["size"],
hash_check=False,
mtime_ns=record.get("mtime_ns"),
inode=record.get("inode"),
keeper=keeper["abs"] if keeper else None,
keeper_root=keeper["root"] if keeper else None,
)
def _kept_copy_state(target, hasher):
"""The kept copy's ``os.stat`` when it is still the reviewed bytes, else None."""
problem = file_problem(
target["keeper_abs"],
target["keeper_root"],
target["hash"],
target["keeper_size"],
hasher,
)
if problem:
return None
try:
return os.stat(target["keeper_abs"])
except OSError:
return None
+451
View File
@@ -0,0 +1,451 @@
"""Reviewing duplicate image files: the checks and the plan (#249 S6, T17).
Pure helpers of ``py/api/duplicate_files.py``: the keep-list parsing, the
per-file check ``file_problem`` (containment decided as text before any
filesystem access, then a regular file of the found size and hash) and the
public ``duplicate_files`` plan.
"""
import hashlib
import json
import os
import stat as stat_module
from dataclasses import dataclass, field
from pathlib import Path
from .admin import DELETABLE_MEDIA_EXTENSIONS
from .maintenance import DUPLICATE_FILES
# Keep entries one review takes; with paths of a few hundred characters the
# list stays inside the preview body cap (DUPLICATE_PREVIEW_BODY_MAX, 1 MiB).
MAX_KEEP = 2000
MAX_DETAILS = 50
MAX_JOB_ID = 64
HASH_CHUNK = 1024 * 1024
# What ``file_problem`` finds, and how a plan counts it as protected.
OUTSIDE_SOURCE = "outside"
LINK = "link"
CHANGED_CONTENT = "changed"
GONE = "gone"
UNREADABLE = "unreadable"
KEEPER = "keeper"
SAME = "same"
MAYBE_SAME = "maybe-same"
NO_FILE_ID = "no-file-id"
# Two mounts of one share may round mtimes differently (SMB keeps 100 ns,
# FAT 2 s): within this, equal sizes on another device may be one file.
MOUNT_MTIME_SLACK_NS = 2_000_000_000
PROTECTED_REASONS = {
MAYBE_SAME: (
"file may be the same file reached through another mount or share",
"files may be the same file reached through another mount or share",
),
NO_FILE_ID: (
"file: can't confirm this is a separate file on this drive",
"files: can't confirm these are separate files on this drive",
),
SAME: (
"file is the same file as its kept copy",
"files are the same file as their kept copy",
),
OUTSIDE_SOURCE: (
"file is not inside a configured source",
"files are not inside a configured source",
),
LINK: ("file is a link, not a regular file", "files are links, not regular files"),
CHANGED_CONTENT: (
"file changed since the search",
"files changed since the search",
),
GONE: ("file no longer exists", "files no longer exist"),
UNREADABLE: ("file could not be read", "files could not be read"),
KEEPER: (
"file kept: its kept copy changed or is gone",
"files kept: their kept copy changed or is gone",
),
}
# Job outcome rows (never an OS error text).
CHANGED = "The file changed since the review; it was kept."
KEEPER_CHANGED = "The kept copy changed or is gone; this file was kept."
SAME_FILE = "This is the same file as its kept copy; it was kept."
MAYBE_SAME_FILE = (
"May be the same file reached through another mount or share; it was kept."
)
NO_FILE_ID_ROW = "Can't confirm this is a separate file on this drive; it was kept."
KEEPER_GONE = "keeper-gone"
KEEPER_WENT = (
"The kept copy changed or went away during the removal; this file was kept "
"and its links stay with it."
)
ROW_DETAILS = {
KEEPER_GONE: None, # filled below
MAYBE_SAME: MAYBE_SAME_FILE,
NO_FILE_ID: NO_FILE_ID_ROW,
SAME: SAME_FILE,
OUTSIDE_SOURCE: "The file is no longer inside its source; it was kept.",
LINK: "The file is now a link, not a regular file; it was kept.",
CHANGED_CONTENT: CHANGED,
GONE: "The file was already gone.",
UNREADABLE: "The file could not be read; it was kept.",
KEEPER: KEEPER_CHANGED,
}
# An unlink fails only after the file's records moved to its kept copy.
LOCKED_FILE = (
"The file is open in another program or read-only; it was not deleted, "
"and its links now name the kept copy."
)
DELETE_FAILED = "The file could not be deleted; its links now name the kept copy."
NOT_ATTEMPTED = "Not attempted: the job stopped after an error."
KEPT_COPY = "Kept copy"
SOURCES_CHANGED = (
"The sources changed since the review; nothing was deleted. Review again."
)
RECORDS_FAILED = (
"The database could not be updated, so the file was not deleted. "
"Its links and tags are unchanged."
)
FIND_FAILED = "The duplicate search failed."
FIND_ERROR_EVENT = "An internal error occurred. Check server logs for details."
FIND_CANCELLED = "Duplicate scan cancelled"
class SourcesChanged(Exception):
"""A source a review or removal names is no longer listed at its index."""
ROW_DETAILS[KEEPER_GONE] = KEEPER_WENT
def still_kept(target, hold):
"""Right before an unlink: the held kept copy is still the file at its
path, and the file is still not the kept copy itself (``keeper_problem``).
"""
try:
if not os.path.samestat(os.fstat(hold.fileno()), os.stat(target["keeper_abs"])):
return False
info = os.lstat(target["abs"])
problem = keeper_problem(
target["abs"],
target["root"],
info,
target["keeper_abs"],
target["keeper_root"],
)
return problem is None
except OSError:
return False
def _plural(count, one, many=None):
return f"{count:,} {one if count == 1 else (many or one + 's')}"
def human_bytes(size):
"""``10 B``, ``1.5 KB``, ``2.0 MB``…"""
if size < 1024:
return f"{size:,} B"
for unit in ("KB", "MB", "GB", "TB"):
size /= 1024
if size < 1024 or unit == "TB":
return f"{size:,.1f} {unit}"
return f"{size:,.1f} TB" # pragma: no cover - unreachable
def _hash_file(path):
digest = hashlib.sha256()
with open(path, "rb") as handle:
for chunk in iter(lambda: handle.read(HASH_CHUNK), b""):
digest.update(chunk)
return digest.hexdigest()
def file_problem(
path,
root,
digest,
size,
hasher=None,
*,
hash_check=True,
mtime_ns=None,
inode=None,
keeper=None,
keeper_root=None,
):
"""Why ``path`` may not be treated as a copy in ``root``, or None (blocking).
Containment is decided as text first, with no filesystem access: the
normalized path must lie below the normalized root, outside
``thumbnails``, with a media extension. Only then is the file looked at:
a regular file, not a link, whose resolved path is still inside the
resolved root and that is not ``keeper`` itself (``keeper_problem``: a
hardlink, one file reached through two spellings or two mounts, or a
volume that cannot tell files apart), of ``size`` bytes.
``mtime_ns`` and ``inode`` (``[st_dev, st_ino]``), when given, bind it to
what the search saw; ``hash_check`` compares its content hash with
``digest`` (the removal job, right before an unlink). Every handle is
closed when this returns.
"""
# py.config imports ComfyUI's server module: only when it is in use.
from ..config import _canonical_path, _lexical_path, _lexically_inside
lexical = _lexical_path(str(path))
base = _lexical_path(str(root))
if lexical == base or not _lexically_inside(lexical, [base]):
return OUTSIDE_SOURCE
parts = Path(os.path.relpath(lexical, base)).parts
if "thumbnails" in parts or ".." in parts:
return OUTSIDE_SOURCE
if os.path.splitext(lexical)[1].lower() not in DELETABLE_MEDIA_EXTENSIONS:
return OUTSIDE_SOURCE
try:
info = os.lstat(path)
except FileNotFoundError:
return GONE
except OSError:
return UNREADABLE
if stat_module.S_ISLNK(info.st_mode):
return LINK
if not stat_module.S_ISREG(info.st_mode):
return LINK
try:
if not _lexically_inside(_canonical_path(path), [_canonical_path(root)]):
return OUTSIDE_SOURCE # a linked folder on the way
if keeper is not None:
same = keeper_problem(path, root, info, keeper, keeper_root)
if same:
return same
if info.st_size != size:
return CHANGED_CONTENT
if mtime_ns is not None and info.st_mtime_ns != mtime_ns:
return CHANGED_CONTENT
if inode is not None and [info.st_dev, info.st_ino] != list(inode):
return CHANGED_CONTENT
if hash_check and (hasher or _hash_file)(path) != digest:
return CHANGED_CONTENT
except FileNotFoundError:
return GONE
except OSError:
return UNREADABLE
return None
def keeper_problem(path, root, info, keeper, keeper_root=None):
"""Why ``path`` (its ``lstat`` is ``info``) may be ``keeper`` itself, or None.
- ``NO_FILE_ID``: a volume reports no file id (``st_ino`` 0), so the two
cannot be told apart;
- ``SAME``: ``os.path.samestat`` (a hardlink, or one file reached through
two spellings of one mount);
- ``MAYBE_SAME``: another device, the same size, mtimes within
``MOUNT_MTIME_SLACK_NS``, and the same ``st_ino`` (NFS and SMB report
the server's file id) or the same path within their sources, or one
ending with the other on a folder name (one share mounted twice, a
FUSE pool over a listed disk, sources listed at different depths).
Protect, never guess: a true copy on another disk at the same relative
path is kept too.
Raises OSError when the kept copy cannot be read (the caller's GONE or
UNREADABLE).
"""
kept = os.stat(keeper)
if not info.st_ino or not kept.st_ino:
return NO_FILE_ID
if os.path.samestat(info, kept):
return SAME
if info.st_dev == kept.st_dev or info.st_size != kept.st_size:
return None
if abs(info.st_mtime_ns - kept.st_mtime_ns) > MOUNT_MTIME_SLACK_NS:
return None
if info.st_ino == kept.st_ino:
return MAYBE_SAME
if keeper_root is not None and _same_tail(
_relative_key(path, root), _relative_key(keeper, keeper_root)
):
return MAYBE_SAME
return None
def _same_tail(one, other):
"""Whether one relative path is the other, or ends with it on a folder name.
Two mounts of one tree may be listed at different depths (``2024/x.png``
in one source, ``x.png`` in the other); ``ax.png`` never matches ``x.png``.
"""
if one is None or other is None:
return False
return one == other or one.endswith("/" + other) or other.endswith("/" + one)
def _relative_key(path, root):
try:
relative = os.path.relpath(str(path), str(root))
except ValueError: # another drive (Windows)
return None
return os.path.normcase(relative).replace("\\", "/")
@dataclass
class HeldFind:
"""One find: where it looked and, once it finished, what it found."""
origin: str # "workspace" or "gallery"
job_id: str = ""
roots: list = field(default_factory=list) # [(source index or None, root)]
legacy_dir: str = None
groups: list = None # [{group_id, hash, size, files:[record]}]
def covers(self, legacy_dir):
"""Whether its events answer the legacy gallery for ``legacy_dir``."""
if self.origin == "gallery":
return True
if not legacy_dir:
return False
wanted = _identity(legacy_dir)
return any(_identity(root) == wanted for _index, root in self.roots)
def _identity(path):
return os.path.normcase(os.path.realpath(str(path)))
def _group_order(group):
return (-(len(group["files"]) - 1) * group["size"], group["hash"])
def _keep_entry(entry):
"""``((group_id, root_index, path), None)`` or ``(None, message)``."""
if not isinstance(entry, dict) or set(entry) != {"group_id", "root_index", "path"}:
return None, "Each keep entry is {group_id, root_index, path}"
group_id, index, path = entry["group_id"], entry["root_index"], entry["path"]
if not isinstance(group_id, str) or not 0 < len(group_id) <= 128:
return None, "group_id must be a group's id"
if index is not None and (
isinstance(index, bool) or not isinstance(index, int) or index < 0
):
return None, "root_index must be a source index"
if not isinstance(path, str) or not 0 < len(path) <= 4096:
return None, "path must be a file's path in its group"
return (group_id, index, path), None
def _parse_options(options):
"""``(job_id, keep, None)`` for duplicate_files options, else a message."""
if not isinstance(options, dict) or set(options) != {"job_id", "keep"}:
return None, None, "duplicate_files takes options {job_id, keep}"
job_id, keep = options["job_id"], options["keep"]
if not isinstance(job_id, str) or not 0 < len(job_id) <= MAX_JOB_ID:
return None, None, "job_id must be a duplicate search's job id"
if not isinstance(keep, list) or not 0 < len(keep) <= MAX_KEEP:
return None, None, f"keep must list 1 to {MAX_KEEP:,} files to keep"
parsed = []
for entry in keep:
key, message = _keep_entry(entry)
if message:
return None, None, message
parsed.append(key)
return job_id, parsed, None
def _select(held, keep):
"""``({group_id: (group, [keeper records])}, None)`` or ``(None, message)``."""
groups = {group["group_id"]: group for group in held.groups}
selection = {}
for group_id, index, path in keep:
group = groups.get(group_id)
if group is None:
return None, "A kept file names a group this search did not find"
match = [
record
for record in group["files"]
if record["root_index"] == index and record["rel"] == path
]
if not match:
return None, "A kept file is not in its group"
_group, keepers = selection.setdefault(group_id, (group, []))
if match[0] not in keepers:
keepers.append(match[0])
return selection, None
SOURCES_CHANGED_REVIEW = (
"The sources changed since the search. Search for duplicates again."
)
def _fingerprint(targets):
digest = hashlib.sha256()
for target in sorted(targets, key=lambda t: _identity(t["abs"])):
part = [
os.path.normcase(target["abs"]),
target["hash"],
target["size"],
os.path.normcase(target["keeper_abs"]),
]
digest.update(json.dumps(part).encode("utf-8") + b"\n")
return "sha256:" + digest.hexdigest()
def _record_parts(records, target):
"""What moves with one removal: its links, tags and Needs review items."""
found = records.get(target["abs"], {})
links, tags, review = (found.get(key, 0) for key in ("links", "tags", "review"))
parts = [
(
(
f"{_plural(links, 'database link')} "
f"{'moves' if links == 1 else 'move'} to the kept copy"
)
if links
else "no database links"
)
]
if tags:
parts.append("its image tags join the kept copy's")
if review:
item = "Needs review item is" if review == 1 else "Needs review items are"
parts.append(f"{review:,} {item} settled")
return parts
def _public_plan(targets, protected, records):
count = len(targets)
groups = len({t["group_id"] for t in targets})
size = sum(t["size"] for t in targets)
reasons = [
f"{n:,} {PROTECTED_REASONS[code][0 if n == 1 else 1]}"
for code, n in sorted(protected.items())
]
details = [
{
"label": target["rel"],
"detail": f"{target['root_label']} · a copy of {target['keeper_rel']}"
+ "".join(f" · {part}" for part in _record_parts(records, target)),
}
for target in targets[:MAX_DETAILS]
]
label = f"{_plural(count, 'duplicate file')} in {_plural(groups, 'group')}"
media = (
f"Deletes {_plural(count, 'file')} ({human_bytes(size)}) from disk; the "
"kept copies stay. The safety backup covers the database, not these files."
)
tagged = sum(1 for t in targets if records.get(t["abs"], {}).get("tags"))
review = sum(records.get(t["abs"], {}).get("review", 0) for t in targets)
if tagged:
media += (
f" The image tags of {_plural(tagged, 'removed copy', 'removed copies')}"
" join their kept copy's."
)
if review:
media += (
f" {_plural(review, 'pending Needs review item')} of the removed "
f"{'copy is' if review == 1 else 'copies are'} settled."
)
return {
"operation": DUPLICATE_FILES,
"scope": {"count": count, "label": label},
"protected": {"count": sum(protected.values()), "reasons": reasons},
"media": {"deletes_files": True, "label": media},
"details": details,
}
+133
View File
@@ -0,0 +1,133 @@
"""Every path a duplicate copy is recorded under (#249 System S6, Task 17).
The search walks each source's real path, but a record names a file the way
it was seen when it was written: under the folder the image monitor watches,
the configured spelling, ComfyUI's output folder, a symlink to it, or a
parent of a source listed later. ``Spellings`` finds them by containment:
for each known folder spelling ``S`` (resolving to the folder ``R``) and a
file ``P`` inside ``R``, the file is also recorded as ``S + relpath(P, R)``.
Folders are compared by ``os.stat`` identity (``st_dev``, ``st_ino``), not
by text, so a case variant on a case-insensitive filesystem (macOS, Windows)
or a symlinked spelling matches. A volume that reports no file id falls
back to ``normcase(realpath)``. Favorites are keyed by a listed root and the
path below it, so ``favorites`` lists every listed root that holds a file.
Blocking: each lookup stats the file's folders once (cached).
"""
import os
from pathlib import Path
from .images import _favorite_relative_key, _favorite_root_key
def _identity(path):
try:
info = os.stat(path)
except OSError:
return None
if info.st_ino:
return ("inode", info.st_dev, info.st_ino)
return ("path", os.path.normcase(os.path.realpath(str(path))))
class Spellings:
"""Folder spellings and listed roots, matched to files by containment."""
def __init__(self, folders, listed_roots):
self._cache = {}
self._folders = []
for spelled in dict.fromkeys(str(f) for f in folders if f):
identity = self._identity(spelled)
if identity is not None:
self._folders.append((spelled, identity))
self._listed = []
for root in dict.fromkeys(str(r) for r in listed_roots if r):
identity = self._identity(root)
if identity is not None:
self._listed.append((root, identity))
def _identity(self, path):
if path not in self._cache:
self._cache[path] = _identity(path)
return self._cache[path]
def _ancestors(self, path):
"""``{identity: folder}`` of every folder above ``path``, nearest first."""
found = {}
folder = os.path.dirname(os.path.abspath(str(path)))
while True:
identity = self._identity(folder)
if identity is not None:
found.setdefault(identity, folder)
parent = os.path.dirname(folder)
if parent == folder:
return found
folder = parent
def of(self, path):
"""``{folder spelling: path spelled under it}``; ``""`` is ``path`` itself."""
path = os.path.abspath(str(path))
spelled = {"": path}
ancestors = self._ancestors(path)
for folder, identity in self._folders:
inside = ancestors.get(identity)
if inside is not None:
spelled.setdefault(
folder, os.path.join(folder, os.path.relpath(path, inside))
)
return spelled
def favorites(self, path):
"""``{listed root: path below it (POSIX)}`` for each listed root holding it."""
path = os.path.abspath(str(path))
ancestors = self._ancestors(path)
held = {}
for root, identity in self._listed:
inside = ancestors.get(identity)
if inside is not None:
held[root] = Path(os.path.relpath(path, inside)).as_posix()
return held
def copy_move(target, kept_stat, spellings):
"""What ``move_copy_records`` moves for one removal.
A record found under a folder spelling moves to the kept copy under
the same spelling when that folder holds it too, else to its real
path. Its favorite moves under every listed root holding the file.
"""
removed = spellings.of(target["abs"])
kept = spellings.of(target["keeper_abs"])
targets = {
spelled: kept.get(folder, target["keeper_abs"])
for folder, spelled in removed.items()
}
kept_favorites = spellings.favorites(target["keeper_abs"])
favorites = []
for root, relative in spellings.favorites(target["abs"]).items():
kept_root, kept_rel = (
(root, kept_favorites[root])
if root in kept_favorites
else (target["keeper_root"], target["keeper_rel"])
)
favorites.append(
(
_favorite_root_key(root),
_favorite_relative_key(relative),
_favorite_root_key(kept_root),
_favorite_relative_key(kept_rel),
)
)
return {
"removed": list(removed.values()),
"kept": list(kept.values()),
"targets": targets,
"kept_path": target["keeper_abs"],
"favorites": favorites,
"removed_locator": target["locator"],
"kept_locator": target["keeper_locator"],
"sha256": target["hash"],
"kept_size": kept_stat.st_size,
"kept_mtime_ns": kept_stat.st_mtime_ns,
}
+16 -5
View File
@@ -19,7 +19,11 @@ ones. One job of a kind runs at a time, and each job has an admission class
- ``guarded`` must not overlap record changes but may overlap other guarded
work (thumbnail generation and clearing, later backup). It is refused while
a ``records`` job runs.
- ``read`` only reads; it is refused only by a running job of its own kind.
- ``read`` only reads; it is refused only by a running job of its own kind
or an ``exclusive`` one.
- ``exclusive`` runs alone (Optimize's VACUUM, which rewrites the whole
database under the write lock): it is refused while any job runs, and every
job is refused while it runs.
A short legacy-page operation takes a ``hold``: a running, non-cancellable
job that lasts as long as the request, so the Activity page lists it and a
@@ -75,12 +79,14 @@ IMAGE_TAG_GROUP_NAMES = ("tagged", "unchanged", "skipped", "errors")
# Thumbnail generation (#249 System S5): previews written, already fresh, failed.
THUMBNAIL_GROUP_NAMES = ("generated", "skipped", "errors")
ADMISSIONS = ("read", "guarded", "records")
ADMISSIONS = ("read", "guarded", "records", "exclusive")
# What each admission class is refused beside (besides its own kind).
# ``exclusive`` (Optimize's VACUUM) runs alone: beside no job of any class.
_REFUSED_BESIDE = {
"read": frozenset(),
"guarded": frozenset({"records"}),
"records": frozenset({"records", "guarded"}),
"read": frozenset({"exclusive"}),
"guarded": frozenset({"records", "exclusive"}),
"records": frozenset({"records", "guarded", "exclusive"}),
"exclusive": frozenset(ADMISSIONS),
}
# Kinds of the legacy-page operations that hold the registry (D80; titles in
@@ -91,6 +97,11 @@ DUPLICATE_CLEANUP_JOB_KIND = "duplicate_cleanup"
THUMBNAILS_JOB_KIND = "thumbnails"
LORA_IMPORT_JOB_KIND = "lora_import"
AUTOTAG_BATCH_JOB_KIND = "autotag"
# Duplicate image files (#249 System S6, py/api/duplicate_files.py): the
# find is a cancellable ``read`` job; the reviewed removal is ``records``.
DUPLICATES_JOB_KIND = "duplicates"
DUPLICATE_GROUP_NAMES = ("groups_found",)
DUPLICATE_CLEANUP_GROUP_NAMES = ("removed", "kept", "errors")
# The System health check (py/api/system.py): a read-only ``read`` job with
# no file outcomes, so it may run beside a scan.
HEALTH_JOB_KIND = "health"
+773
View File
@@ -0,0 +1,773 @@
"""Maintenance routes: reviewed, read-only plans (#249 System S6).
``POST /prompt_manager/maintenance/preview {operation, options:{}}`` reviews
one operation (``duplicate_prompts``, ``optimize``, ``unused_components`` or
``missing_links``; ``database/maintenance.py``) off the event loop and
answers a ``Plan``: ``{success, token, operation, expires_at, scope:{count,
label}, protected:{count, reasons}, backup:{policy, destination_label},
media:{deletes_files, label}, details:[≤50 {label, detail}]}``. The preview
writes nothing. The plan's fingerprint, its targets and the sources it judged
stay in the server's ``PlanStore`` under the token, bound to the database the
preview read (a restore that replaces it invalidates the plan).
The operation enum is strict, ``options`` must be an empty object and the
body is capped at ``PREVIEW_BODY_MAX`` bytes. Errors are ``{success:false,
code, error}`` with ``invalid-request`` (400), ``too-large`` (413: the body,
or a plan over the store's budget), ``storage-changed`` (409: the database
was replaced while the preview read it) or a null code (500); none carries a
server path or exception text.
``POST /prompt_manager/maintenance/apply {token}`` applies a reviewed plan as
an Activity job (202 ``{job}``): kind ``maintenance``, titled after the
operation, admission ``records`` (refused while a scan, restore, LoRA import
or other record-changing job runs, and refusing them meanwhile), not
cancellable, with outcome groups ``changed``, ``protected`` and ``errors``.
The token is consumed when the job starts (on the same loop turn, before its
runner runs); a refused start keeps it. The runner:
1. takes a safety backup when the plan was reviewed with the policy on or the
policy (``maintenance.safety_backup``, default on) is on now: SQLite's
backup API into a temp file beside the database, ``verify_database_file``,
fsync, ``0o600`` on POSIX, then ``os.replace`` onto a unique
``<db>.backup_<YYYYmmdd_HHMMSS_ffffff>[_n]`` name. Any failure removes
what it wrote and ends the job ``backup-failed`` with nothing changed;
2. applies the plan, which first rechecks the reviewed fingerprint inside
``BEGIN IMMEDIATE`` (``plan-changed`` writes nothing);
3. counts the changed and protected records, and prunes old safety backups.
The job log names the backup by its file name only, and the job says "No
safety backup" when none was taken. Apply errors: ``invalid-request`` (400),
``plan-expired`` (410: unknown, expired or forgotten after a restart),
``plan-used`` (409, with the job that applied it), ``storage-changed`` (409:
a restore replaced the database since the review) and ``busy`` (409).
"""
import datetime
import json
import os
import shutil
import sqlite3
import tempfile
from contextlib import closing
from aiohttp import web
from .admin import _prune_safety_backups, _scan_location, _scan_roots
from .job_routes import busy_response
from .jobs import MAINTENANCE_JOB_KIND, JobBusy
from .maintenance_plans import PlanExpired, PlanUsed, ScopeTooLarge
try:
from ...database.health import readonly_uri
from ...database.maintenance import MAINTENANCE_OPERATIONS, PlanChanged
except ImportError:
from database.health import readonly_uri
from database.maintenance import MAINTENANCE_OPERATIONS, PlanChanged
PREVIEW_BODY_MAX = 4096
# Duplicate-file previews carry the reviewed keepers (py/api/duplicate_files.py).
DUPLICATE_PREVIEW_BODY_MAX = 1024 * 1024
DUPLICATE_FILES = "duplicate_files"
PREVIEW_OPERATIONS = MAINTENANCE_OPERATIONS + (DUPLICATE_FILES,)
# Room left over beyond what a step writes, so a disk is never filled to 0.
DISK_MARGIN_BYTES = 16 * 2**20
# The concept's Maintenance row titles, which name each operation's job.
MAINTENANCE_TITLES = {
"duplicate_prompts": "Duplicate prompt text",
"optimize": "Optimize database",
"unused_components": "Unused prompt components",
"missing_links": "Missing-file links",
}
_OPERATION_ERROR = "operation must be one of: " + ", ".join(PREVIEW_OPERATIONS)
MAINTENANCE_GROUPS = ("changed", "protected", "errors")
MAX_TOKEN_LENGTH = 128
BACKUP_FAILED = (
"The safety backup could not be taken and verified; nothing was changed."
)
PLAN_CHANGED = (
"The library changed since the review; nothing was changed. "
"Review the plan again."
)
APPLY_FAILED = "Maintenance failed; nothing was changed."
NO_BACKUP = "No safety backup"
_NAME_ATTEMPTS = 1000
class SafetyBackupFailed(Exception):
"""The safety backup could not be written, verified or put in place."""
def _now():
return datetime.datetime.now()
def _free_bytes(path):
return shutil.disk_usage(path).free
def _device(path):
try:
return os.stat(path).st_dev
except OSError:
return path
def sqlite_temp_dir(environ=None, *, windows=None, usable=None):
"""The folder SQLite builds its temporary files in (VACUUM's copy).
Mirrors SQLite's own search, which differs from ``tempfile``'s. On Unix
(``unixTempFileDir``): ``SQLITE_TMPDIR``, ``TMPDIR``, then ``/var/tmp``,
``/usr/tmp`` and ``/tmp``, the first that is a writable directory, else
the current directory. On Windows (``GetTempPathW``): ``TMP``, ``TEMP``,
then ``USERPROFILE``, else the current directory.
"""
environ = os.environ if environ is None else environ
if windows is None:
windows = os.name == "nt"
if windows:
for name in ("TMP", "TEMP", "USERPROFILE"):
if environ.get(name):
return environ[name]
return os.path.abspath(".")
if usable is None:
def usable(path):
return os.path.isdir(path) and os.access(path, os.W_OK | os.X_OK)
candidates = [environ.get("SQLITE_TMPDIR"), environ.get("TMPDIR")]
candidates += ["/var/tmp", "/usr/tmp", "/tmp"]
for candidate in candidates:
if candidate and usable(candidate):
return candidate
return os.path.abspath(".")
def database_bytes(db_path):
"""The size a safety backup of the database will be (blocking).
SQLite's backup API copies every page, committed WAL content included, so
the copy is ``page_count * page_size`` read on one read-only snapshot.
"""
conn = sqlite3.connect(readonly_uri(db_path), uri=True)
try:
page_size = conn.execute("PRAGMA page_size").fetchone()[0]
return page_size * conn.execute("PRAGMA page_count").fetchone()[0]
finally:
conn.close()
def _megabytes(size):
return f"{size / 2**20:,.1f} MB"
def disk_plan(db_path, *, backup, vacuum):
"""Whether the disk has room for an apply's writes (blocking).
``backup``: the safety backup, one database-sized file beside the
database. ``vacuum``: VACUUM writes the rebuilt database through the WAL
beside it and builds it in a temporary copy in SQLite's temporary folder
(``sqlite_temp_dir``; SQLite's documentation: up to twice the database
in free space). Needs
on one device add up. Returns ``(public, short)``: ``public`` is
``{free_bytes, required_bytes, sufficient}`` for the database folder,
``short`` None or ``(place, what, required, free)`` of the first place
that is short, and ``size`` the database's size (0 when nothing is
written).
"""
db_path = os.path.abspath(db_path)
folder = os.path.dirname(db_path)
size = database_bytes(db_path) if backup or vacuum else 0
needs = [] # (place, folder, bytes, what)
if backup:
needs.append(("database folder", folder, size, "the safety backup"))
if vacuum:
needs.append(("database folder", folder, size, "Optimize"))
needs.append(("temporary folder", sqlite_temp_dir(), size, "Optimize"))
devices = {}
for place, path, needed, what in needs:
entry = devices.setdefault(_device(path), {"place": place, "path": path})
entry["bytes"] = entry.get("bytes", 0) + needed
entry.setdefault("what", [])
if what not in entry["what"]:
entry["what"].append(what)
short = None
public = {"free_bytes": _free_bytes(folder), "required_bytes": 0}
for device, entry in devices.items():
required = entry["bytes"] + DISK_MARGIN_BYTES
free = (
public["free_bytes"]
if entry["path"] == folder
else _free_bytes(entry["path"])
)
if device == _device(folder):
public["required_bytes"] = required
if free < required and short is None:
short = (entry["place"], " and ".join(entry["what"]), required, free)
public["sufficient"] = short is None
return public, short, size
# A restore upload's hard ceiling; below it, free disk space decides.
RESTORE_MAX_BYTES = 64 * 2**30
def size_label(size):
"""``1.6 GB`` or ``120.0 MB``: a backup's size for a message."""
if size >= 2**30:
return f"{size / 2**30:,.1f} GB"
return _megabytes(size)
def restore_room(db_path):
"""The largest restore upload the disks have room for (blocking).
The upload is written to a temp file and staged once more in the
temporary folder (2x), then copied over the live database beside it,
whose current contents are first kept as a safety backup (1x the upload
for the WAL, plus the current database). Needs on one device add up,
each with ``DISK_MARGIN_BYTES``.
"""
db_path = os.path.abspath(db_path)
folder = os.path.dirname(db_path)
current = database_bytes(db_path) if os.path.exists(db_path) else 0
devices = {}
for path, per_upload, fixed in (
(tempfile.gettempdir(), 2, 0),
(folder, 1, current),
):
entry = devices.setdefault(_device(path), {"path": path, "per": 0, "fixed": 0})
entry["per"] += per_upload
entry["fixed"] += fixed
room = min(
(_free_bytes(e["path"]) - e["fixed"] - DISK_MARGIN_BYTES) // e["per"]
for e in devices.values()
)
return max(0, room)
def restore_space_message(size, room, *, known):
"""The truthful refusal of a restore the disks have no room for."""
if known:
return (
f"Not enough free disk space to restore a {size_label(size)} backup; "
f"there is room for about {size_label(room)}. Nothing was changed."
)
return (
"Not enough free disk space to restore this backup: it is larger than "
f"the {size_label(room)} there is room for. Nothing was changed."
)
def short_space_message(short):
"""The truthful refusal for ``disk_plan``'s ``short``."""
place, what, required, free = short
verb = "need" if " and " in what else "needs"
return (
f"Not enough free disk space: {what} {verb} about {_megabytes(required)} "
f"in the {place} and {_megabytes(free)} is free. Nothing was changed."
)
def _changed_line(operation, count):
if operation == "optimize":
return "Optimized the database"
if operation == "duplicate_prompts":
noun = "duplicate prompt" if count == 1 else "duplicate prompts"
return f"Merged {count:,} {noun}"
if operation == "unused_components":
noun = "unused prompt component" if count == 1 else "unused prompt components"
return f"Removed {count:,} {noun}"
noun = "link to a missing file" if count == 1 else "links to missing files"
return f"Removed {count:,} {noun}"
def _remove_quietly(*paths):
for path in paths:
try:
if path and os.path.exists(path):
os.remove(path)
except OSError:
pass
def _reserve_backup_name(db_path):
"""Create and return an unused ``<db>.backup_<stamp>[_n]`` (O_EXCL)."""
for candidate in _backup_names(db_path):
try:
fd = os.open(candidate, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
except FileExistsError:
continue
os.close(fd)
return candidate
raise SafetyBackupFailed("No free backup name")
def _backup_names(db_path):
base = f"{db_path}.backup_{_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):
"""Give the verified ``temp`` its unique backup name; return that path.
``os.link`` is atomic and refuses an existing name, so no empty or half
file ever carries a backup name, even across a crash. A filesystem
without hard links falls back to reserving the name (``O_EXCL``) and
``os.replace``. The temp file is removed by the caller.
"""
for candidate in _backup_names(db_path):
try:
os.link(temp, candidate)
except FileExistsError:
continue
except (OSError, NotImplementedError):
break # no hard links here
else:
return candidate
else:
raise SafetyBackupFailed("No free backup name")
final = _reserve_backup_name(db_path)
try:
os.replace(temp, final)
except OSError:
_remove_quietly(final)
raise
return final
SAFETY_TEMP_PREFIX = ".pm-safety-"
def sweep_safety_temps(db_path):
"""Remove the temp copies an interrupted safety backup left (blocking).
Called once at startup, before any job can take a backup. Returns the
names removed.
"""
folder = os.path.dirname(os.path.abspath(db_path))
removed = []
try:
names = os.listdir(folder)
except OSError:
return removed
for name in names:
if not (name.startswith(SAFETY_TEMP_PREFIX) and ".tmp" in name):
continue
try:
os.remove(os.path.join(folder, name))
removed.append(name)
except OSError:
continue
return removed
def _standalone(path):
"""Make the copy one self-contained file (rollback journal, not WAL).
The backup API copies the live database's WAL flag; a WAL file opened
read-only later (a verify, a restore preview) would leave ``-wal`` and
``-shm`` sidecars beside it.
"""
try:
with closing(sqlite3.connect(path)) as conn:
conn.execute("PRAGMA journal_mode=DELETE").fetchone()
except sqlite3.Error as exc:
raise SafetyBackupFailed("The backup could not be finished") from exc
def _fsync(path):
with open(path, "rb+") as handle:
os.fsync(handle.fileno())
def take_safety_backup(model):
"""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.
"""
db_path = os.path.abspath(model.db_path)
fd, temp = tempfile.mkstemp(
dir=os.path.dirname(db_path), prefix=SAFETY_TEMP_PREFIX, suffix=".tmp"
)
os.close(fd)
final = None
try:
if not model.backup_database(temp):
raise SafetyBackupFailed("The backup could not be written")
_standalone(temp)
ok, _reason = model.verify_database_file(temp)
if not ok:
raise SafetyBackupFailed("The backup did not verify")
_fsync(temp)
if os.name != "nt":
os.chmod(temp, 0o600)
final = _place_backup(temp, db_path)
return os.path.basename(final)
except OSError as exc:
_remove_quietly(final)
raise SafetyBackupFailed("The backup could not be put in place") from exc
except SafetyBackupFailed:
_remove_quietly(final)
raise
finally:
_remove_quietly(temp, temp + "-journal", temp + "-wal", temp + "-shm")
def _failure(message, status, code=None, **extra):
return web.json_response(
{"success": False, "code": code, "error": message, **extra}, status=status
)
def _invalid(message):
return _failure(message, 400, "invalid-request")
def _iso(timestamp):
return datetime.datetime.fromtimestamp(timestamp, datetime.timezone.utc).isoformat()
class MaintenanceRoutesMixin:
"""Mixin providing the /prompt_manager/maintenance plan endpoints."""
def _register_maintenance_routes(self, routes):
@routes.post("/prompt_manager/maintenance/preview")
async def maintenance_preview_route(request):
return await self.preview_maintenance(request)
@routes.post("/prompt_manager/maintenance/apply")
async def maintenance_apply_route(request):
return await self.apply_maintenance_plan(request)
def _database_identity(self, db):
"""What binds a plan to one database: the swap generation and its file."""
path = os.path.normcase(os.path.realpath(db.model.db_path))
return (self._db_generation, path)
@staticmethod
def _safety_backup_policy():
"""``on`` or ``off``: whether an apply takes a safety backup first."""
from ..config import PromptManagerConfig
return "on" if PromptManagerConfig.SAFETY_BACKUP else "off"
async def _read_maintenance_body(self, request, allowed, cap=None):
"""``(data, None)`` for a capped JSON object with only ``allowed`` keys.
``cap(data)``, when given, is the body's byte limit for what it asks
(read up to ``DUPLICATE_PREVIEW_BODY_MAX``); else PREVIEW_BODY_MAX.
"""
limit = PREVIEW_BODY_MAX if cap is None else DUPLICATE_PREVIEW_BODY_MAX
data, error = await self._read_json_body(request, max_bytes=limit)
if error is not None:
if error.status == 413:
return None, _failure("Request body too large", 413, "too-large")
return None, _invalid("Request body must be a JSON object")
if cap is not None:
size = len(json.dumps(data, separators=(",", ":")).encode("utf-8"))
if size > cap(data):
return None, _failure("Request body too large", 413, "too-large")
unknown = sorted(set(data) - set(allowed))
if unknown:
return None, _invalid("Unknown field: " + ", ".join(unknown)[:200])
return data, None
async def preview_maintenance(self, request):
data, error = await self._read_maintenance_body(
request,
("operation", "options"),
cap=lambda body: (
DUPLICATE_PREVIEW_BODY_MAX
if body.get("operation") == DUPLICATE_FILES
else PREVIEW_BODY_MAX
),
)
if error is not None:
return error
operation = data.get("operation")
options = data.get("options", {})
if not isinstance(operation, str) or operation not in PREVIEW_OPERATIONS:
return _invalid(_OPERATION_ERROR)
if operation == DUPLICATE_FILES:
return await self.preview_duplicate_files(options)
if not isinstance(options, dict) or options:
return _invalid(f"{operation} takes no options")
db = self.db
identity = self._database_identity(db)
try:
roots = (
await self._run_in_executor(self._health_roots)
if operation == "missing_links"
else []
)
draft = await self._run_in_executor(
db.preview_maintenance, operation, roots
)
except Exception:
self.logger.exception(f"Maintenance preview of {operation} failed")
return _failure("Failed to review maintenance", 500)
return await self._plan_response(
db,
identity,
draft.operation,
draft.fingerprint,
draft.public(),
{
"targets": draft.targets,
"roots": roots,
"protected": draft.protected["count"],
},
)
async def _plan_response(
self, db, identity, operation, fingerprint, public, payload
):
"""Issue the reviewed plan's token and answer the ``Plan``.
Refused with ``storage-changed`` when the database was replaced while
it was reviewed. The plan reports the safety backup's expected size
and the disk space the apply needs (``disk_plan``).
"""
if identity != self._database_identity(self.db):
return _failure(
"The database was replaced while it was reviewed. Review again.",
409,
"storage-changed",
)
policy = self._safety_backup_policy()
try:
disk, _short, size = await self._run_in_executor(
disk_plan,
db.model.db_path,
backup=policy == "on",
vacuum=operation == "optimize",
)
except (OSError, sqlite3.Error):
self.logger.exception("Could not measure the disk space for a plan")
return _failure("Failed to review maintenance", 500)
store = self._maintenance_plans
try:
token = store.issue(
operation, fingerprint, {**payload, "backup": policy}, identity
)
except ScopeTooLarge:
return _failure(
"This plan is too large to review at once.", 413, "too-large"
)
plan = store.peek(token, operation)
name = os.path.basename(db.model.db_path)
return web.json_response(
{
"success": True,
"token": token,
"expires_at": _iso(plan.expires_at),
**public,
"backup": {
"policy": policy,
"destination_label": f"{name}.backup_<date and time>",
"expected_bytes": size if policy == "on" else None,
},
"disk": disk,
},
headers={"Cache-Control": "no-store"},
)
# ── apply ─────────────────────────────────────────────────────────
async def apply_maintenance_plan(self, request):
data, error = await self._read_maintenance_body(request, ("token",))
if error is not None:
return error
token = data.get("token")
if not isinstance(token, str) or not 0 < len(token) <= MAX_TOKEN_LENGTH:
return _invalid("token must be the plan's token")
plan, refused = self._peek_plan(token)
if refused is not None:
return refused
store = self._maintenance_plans
db = self.db
if plan.db_identity != self._database_identity(db):
store.discard(token)
return _failure(
"The database was replaced since the review. Review again.",
409,
"storage-changed",
)
backup = self._backup_promised(plan)
try:
_disk, short, _size = await self._run_in_executor(
disk_plan,
db.model.db_path,
backup=backup,
vacuum=plan.operation == "optimize",
)
except (OSError, sqlite3.Error):
self.logger.exception("Could not measure the disk space for an apply")
return _failure(APPLY_FAILED, 500)
if short is not None:
# Refused before any job starts: nothing fails halfway for space.
return _failure(short_space_message(short), 507, "insufficient-space")
# The token may have been used, expired or dropped while the disk was
# measured: look again with no await before the job starts.
plan, refused = self._peek_plan(token)
if refused is not None:
return refused
try:
if plan.operation == DUPLICATE_FILES:
job = self._start_duplicate_cleanup(db, plan, backup)
else:
job = self._jobs.start(
MAINTENANCE_JOB_KIND,
MAINTENANCE_TITLES[plan.operation],
# VACUUM rewrites the whole file under the write lock: it
# runs alone, beside no other job of any class.
admission=(
"exclusive" if plan.operation == "optimize" else "records"
),
cancellable=False,
scope={"operation": plan.operation},
groups=MAINTENANCE_GROUPS,
runner=lambda job: self._run_maintenance_plan(
job, db, plan, backup
),
)
except JobBusy as busy:
return busy_response(busy) # the token stays usable
# Same loop turn as the peek and the start: the runner has not run.
store.consume(token, plan.operation, job_id=job.id)
return web.json_response({"success": True, "job": job.summary()}, status=202)
def _peek_plan(self, token):
"""``(plan, None)`` for a usable token, else ``(None, refusal)``."""
try:
return self._maintenance_plans.peek(token), None
except PlanUsed as used:
job = self._jobs.get(used.job_id)
return None, _failure(
"This plan was already applied.",
409,
"plan-used",
job={"id": used.job_id, "title": job.title if job else None},
)
except PlanExpired:
return None, _failure(
"This review expired. Review the plan again.", 410, "plan-expired"
)
def _backup_promised(self, plan):
"""A plan reviewed with the policy on keeps that promise; so does on now.
Decided once, when the apply starts (the disk check measured it), and
handed to the job.
"""
return plan.payload["backup"] == "on" or self._safety_backup_policy() == "on"
async def _safety_backup_step(self, job, db, backup_on):
"""``(ok, name)``: take the promised safety backup, or report none.
``name`` is the backup's file name, or None when no backup was
promised. A failed backup ends the job ``backup-failed`` (``ok``
False).
"""
if not backup_on:
job.log_line("No safety backup (the policy is off)")
return True, None
job.update(status_line="Taking a safety backup")
job.log_line("Taking a safety backup of the database")
try:
backup = await self._run_in_executor(take_safety_backup, db.model)
except SafetyBackupFailed:
self.logger.error("Safety backup failed", exc_info=True)
_fail_job(job, "backup-failed", BACKUP_FAILED)
return False, None
job.log_line(f"Safety backup {backup} verified")
return True, backup
async def _prune_backup_pool(self, job, db):
"""Keep the one ``<db>.backup_*`` pool (restore's too) at its cap."""
removed = await self._run_in_executor(_prune_safety_backups, db.model.db_path)
if removed:
job.log_line(f"Removed {len(removed)} older safety backups")
async def _run_maintenance_plan(self, job, db, plan, backup_on):
ok, backup = await self._safety_backup_step(job, db, backup_on)
if not ok:
return
try:
await self._apply_maintenance_plan(job, db, plan, backup)
finally:
if backup:
await self._prune_backup_pool(job, db)
async def _apply_maintenance_plan(self, job, db, plan, backup):
payload = plan.payload
job.update(status_line=f"Applying: {job.title}")
try:
result = await self._run_in_executor(
db.apply_maintenance,
plan.operation,
plan.fingerprint,
payload["targets"],
payload["roots"],
)
except PlanChanged:
return _fail_job(job, "plan-changed", PLAN_CHANGED)
except Exception:
self.logger.exception(f"Maintenance {plan.operation} failed")
return _fail_job(job, None, APPLY_FAILED)
rows = await self._run_in_executor(
self._maintenance_rows, plan.operation, result["rows"]
)
job.record_counted("changed", result["changed"], rows)
job.record_counted("protected", payload["protected"])
line = _changed_line(plan.operation, result["changed"])
job.log_line(line)
job.update(progress=100, status_line=f"{line} · {backup_note(backup)}")
def _maintenance_rows(self, operation, described):
"""Outcome rows for the changed records (root-relative paths only)."""
if operation == "missing_links":
roots = _scan_roots(self._get_all_output_dirs())
rows = []
for item in described:
path, label, index = _scan_location(item["image_path"], roots)
rows.append(
{
"path": path,
"root": label,
"root_index": index,
"detail": "Link removed; the file was already missing",
}
)
return rows
if operation == "duplicate_prompts":
return [
{
"path": "",
"detail": (
f"Prompt {item['prompt_id']} merged into "
f"prompt {item['kept_id']}"
),
}
for item in described
]
return [
{"path": "", "detail": f"Removed: {item.get('text') or item['prompt_id']}"}
for item in described
]
def backup_note(backup):
"""The status line's backup part: the file name, or that there is none."""
return f"Safety backup {backup}" if backup else NO_BACKUP
def _fail_job(job, code, message):
"""End ``job`` failed with ``message`` (one error row) and its code."""
job.record("errors", {"path": "", "detail": message}, f"! {message}")
job.update(status_line=message)
job.publish({"type": "error", "code": code, "message": message})
job.finish("failed")
+138
View File
@@ -0,0 +1,138 @@
"""The reviewed-plan store behind maintenance previews (#249 System S6).
A preview issues a plan: its operation, the fingerprint of exactly what was
reviewed, the payload an apply needs (target ids, the sources it judged) and
the identity of the database it describes. Everything is held here, in
server memory, under a single-use token (``secrets.token_urlsafe(32)``); the
client only ever sees the token and the public plan. Plans expire after
``ttl`` seconds and are gone after a restart.
The store keeps at most ``max_plans`` plans and ``max_bytes`` of payload
(measured as compact JSON). A plan larger than ``max_bytes`` is refused with
``ScopeTooLarge``: it is never cut down to fit, since an apply must do
exactly what was reviewed. A new plan that fits makes room by dropping the
oldest ones, which then read as expired.
A consumed token is remembered (the last ``MAX_USED``) with the job it
started, so a second apply of it is told which job ran it.
"""
import json
import secrets
import time
from collections import OrderedDict
from dataclasses import dataclass, field
class PlanExpired(Exception):
"""No live plan has the token: unknown, expired, dropped or another operation's."""
class PlanUsed(Exception):
"""The token was already applied; ``job_id`` is the job that applied it."""
def __init__(self, job_id):
super().__init__("This plan was already applied")
self.job_id = job_id
class ScopeTooLarge(Exception):
"""The plan's payload exceeds the store's byte budget."""
@dataclass
class StoredPlan:
"""One reviewed plan, as only the server sees it."""
operation: str
fingerprint: str
payload: dict
db_identity: tuple
expires_at: float # wall-clock seconds since the epoch, for the response
deadline: float = field(repr=False) # on the store's monotonic clock
size: int = field(repr=False)
def _payload_size(payload):
return len(json.dumps(payload, separators=(",", ":"), default=str))
class PlanStore:
"""Reviewed maintenance plans under single-use tokens."""
MAX_USED = 64 # consumed tokens remembered for plan-used answers
def __init__(
self,
ttl=900,
max_plans=8,
max_bytes=16 * 2**20,
*,
clock=time.monotonic,
wall_clock=time.time,
):
self.ttl = ttl
self.max_plans = max_plans
self.max_bytes = max_bytes
self._clock = clock
self._wall_clock = wall_clock
self._plans = OrderedDict() # token -> StoredPlan, oldest first
self._used = OrderedDict() # token -> job id, oldest first
def issue(self, operation, fingerprint, payload, db_identity) -> str:
"""Hold a plan and return its token; ScopeTooLarge when it cannot fit."""
size = _payload_size(payload)
if size > self.max_bytes:
raise ScopeTooLarge(f"The plan needs {size} bytes")
self._drop_expired()
while self._plans and (
len(self._plans) >= self.max_plans
or self._held_bytes() + size > self.max_bytes
):
self._plans.popitem(last=False)
token = secrets.token_urlsafe(32)
self._plans[token] = StoredPlan(
operation=operation,
fingerprint=fingerprint,
payload=payload,
db_identity=tuple(db_identity),
expires_at=self._wall_clock() + self.ttl,
deadline=self._clock() + self.ttl,
size=size,
)
return token
def peek(self, token, operation=None) -> StoredPlan:
"""The live plan of ``token``, left in place; raises PlanUsed or PlanExpired."""
if not isinstance(token, str):
raise PlanExpired("No such plan")
if token in self._used:
raise PlanUsed(self._used[token])
plan = self._plans.get(token)
if plan is None or (operation is not None and plan.operation != operation):
raise PlanExpired("No such plan")
if self._clock() > plan.deadline:
del self._plans[token]
raise PlanExpired("The plan expired")
return plan
def consume(self, token, operation=None, *, job_id=None) -> StoredPlan:
"""Take the plan of ``token`` once; later calls raise PlanUsed(job_id)."""
plan = self.peek(token, operation)
del self._plans[token]
self._used[token] = job_id
while len(self._used) > self.MAX_USED:
self._used.popitem(last=False)
return plan
def discard(self, token):
"""Forget a plan that can no longer apply (its database was replaced)."""
self._plans.pop(token, None)
def _held_bytes(self):
return sum(plan.size for plan in self._plans.values())
def _drop_expired(self):
now = self._clock()
for token in [t for t, plan in self._plans.items() if now > plan.deadline]:
del self._plans[token]
+6
View File
@@ -40,6 +40,7 @@ _SETTINGS_FILE_KEYS = {
"INFINITE_SCROLL": ("web_ui", "infinite_scroll"),
"WORKER_THREADS": ("performance", "worker_threads"),
"MONITORING_ENABLED": ("gallery", "monitoring", "enabled"),
"SAFETY_BACKUP": ("maintenance", "safety_backup"),
}
# Attributes that live on GalleryConfig; the others on PromptManagerConfig.
_GALLERY_SETTINGS = frozenset({"MONITORING_ENABLED"})
@@ -187,6 +188,7 @@ class SettingsRoutesMixin:
"webui_display_mode": PromptManagerConfig.WEBUI_DISPLAY_MODE,
"infinite_scroll": PromptManagerConfig.INFINITE_SCROLL,
"worker_threads": PromptManagerConfig.WORKER_THREADS,
"safety_backup": PromptManagerConfig.SAFETY_BACKUP,
"cpu_count": PromptManagerConfig.max_worker_threads(),
"gallery_root_paths": root_paths,
"gallery_root_path": root_paths[0] if root_paths else "",
@@ -481,6 +483,10 @@ class SettingsRoutesMixin:
if not isinstance(data["monitoring_enabled"], bool):
raise ValueError("monitoring_enabled must be true or false")
updates["MONITORING_ENABLED"] = data["monitoring_enabled"]
if "safety_backup" in data:
if not isinstance(data["safety_backup"], bool):
raise ValueError("safety_backup must be true or false")
updates["SAFETY_BACKUP"] = data["safety_backup"]
if "worker_threads" in data:
updates["WORKER_THREADS"] = validate_worker_threads(
data["worker_threads"], PromptManagerConfig.max_worker_threads()
+11
View File
@@ -659,6 +659,10 @@ class PromptManagerConfig:
DEFAULT_WORKER_THREADS = max(1, (os.cpu_count() or 1) // 2)
WORKER_THREADS = DEFAULT_WORKER_THREADS
# Maintenance: a reviewed apply takes a verified safety backup of the
# database first (System › Backup & export › Recovery policy, #249 S6).
SAFETY_BACKUP = True
@classmethod
def max_worker_threads(cls) -> int:
"""Largest accepted worker-thread count: the detected core count."""
@@ -713,6 +717,7 @@ class PromptManagerConfig:
"auto_backup_interval": cls.AUTO_BACKUP_INTERVAL,
"worker_threads": cls.WORKER_THREADS,
},
"maintenance": {"safety_backup": cls.SAFETY_BACKUP},
"gallery": GalleryConfig.get_config(),
"integrations": IntegrationConfig.get_config(),
"autotag": AutotagConfig.get_config(),
@@ -848,6 +853,12 @@ class PromptManagerConfig:
performance["worker_threads"], cls.WORKER_THREADS
)
maintenance = new_config.get("maintenance", {})
if isinstance(maintenance, dict) and isinstance(
maintenance.get("safety_backup"), bool
):
cls.SAFETY_BACKUP = maintenance["safety_backup"]
# Update gallery config
if "gallery" in new_config:
GalleryConfig.update_config(new_config["gallery"])
+141 -1
View File
@@ -47,11 +47,13 @@ def _texts(path):
class BackupRestoreTestCase(AioHTTPTestCase):
app_options = {}
async def get_application(self):
self.tmpdir = tempfile.mkdtemp()
self.db_path = os.path.join(self.tmpdir, "prompts.db")
app = web.Application()
app = web.Application(**self.app_options)
routes = web.RouteTableDef()
self.api = PromptManagerAPI(db=PromptDatabase(self.db_path))
self.api.add_routes(routes)
@@ -327,6 +329,15 @@ class TestRestoreRoute(BackupRestoreTestCase):
self.assertIn("too large", data["error"].lower())
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
async def test_restore_cap_keeps_the_legacy_shape(self):
self.api.restore_max_bytes = 1024
upload = self._write("big.db", b"SQLite format 3\x00" + b"\x00" * 8192)
resp = await self._restore(upload)
data = await resp.json()
self.assertEqual(resp.status, 400)
self.assertEqual(data["code"], "too-large")
self.assertNotIn(self.tmpdir, data["error"])
async def test_non_multipart_request_is_rejected(self):
resp = await self.client.post("/prompt_manager/restore", data=b"raw bytes")
self.assertIn(resp.status, (400, 500))
@@ -336,3 +347,132 @@ class TestRestoreRoute(BackupRestoreTestCase):
if __name__ == "__main__":
unittest.main()
MIB = 2**20
class TestLargeRestores(BackupRestoreTestCase):
"""A real library's backup restores: room on disk decides, not 100 MB."""
# As ComfyUI builds its app: client_max_size is --max-upload-size (100
# MB by default). The restore streams with read_chunk, which it does not
# cap.
app_options = {"client_max_size": 100 * 2**20}
def _large_db_file(self, texts, pad_bytes):
path = self._make_db_file(texts, name="large.db")
conn = sqlite3.connect(path)
try:
conn.execute("CREATE TABLE padding(b BLOB)")
conn.execute("INSERT INTO padding VALUES (zeroblob(?))", (pad_bytes,))
conn.commit()
finally:
conn.close()
return path
async def _stream_restore(self, path, chunked=False):
with open(path, "rb") as fh:
data = fh.read()
form = aiohttp.FormData()
form.add_field(
"database_file",
data,
filename=os.path.basename(path),
content_type="application/octet-stream",
)
# chunked=False must not be passed: aiohttp 3.14 then frames the body
# as chunked while also sending a Content-Length.
options = {"chunked": True} if chunked else {}
return await self.client.post("/prompt_manager/restore", data=form, **options)
def _spy_uploads(self):
"""Patch the upload's mkstemp; the list collects each temp path made."""
created = []
real_mkstemp = tempfile.mkstemp
def spy(*args, **kwargs):
result = real_mkstemp(*args, **kwargs)
created.append(result[1])
return result
return patch("py.api.admin.tempfile.mkstemp", spy), created
def _assert_removed(self, created):
for path in created:
for suffix in ("", "-wal", "-shm", "-journal"):
self.assertFalse(os.path.exists(path + suffix), path + suffix)
async def test_a_150_mb_backup_restores(self):
self._save("before")
upload = self._large_db_file(["restored"], 150 * MIB)
self.assertGreater(os.path.getsize(upload), 150 * MIB)
spy, created = self._spy_uploads()
with spy:
resp = await self._stream_restore(upload)
data = await resp.json()
self.assertEqual(resp.status, 200, data)
self.assertEqual(
[p["text"] for p in self.api.db.search_prompts()], ["restored"]
)
self.assertEqual(len(created), 1)
self._assert_removed(created)
async def test_a_short_disk_is_refused_up_front(self):
from py.api import maintenance
self._save("keep me")
upload = self._make_db_file(["restored"])
spy, created = self._spy_uploads()
with spy, patch.object(maintenance, "_free_bytes", return_value=MIB):
resp = await self._restore(upload)
data = await resp.json()
self.assertEqual(resp.status, 507)
self.assertEqual(data["code"], "insufficient-space")
self.assertTrue(
data["error"].startswith("Not enough free disk space to restore a "),
data["error"],
)
self.assertTrue(data["error"].endswith("Nothing was changed."))
self.assertNotIn(self.tmpdir, data["error"])
self.assertNotIn(tempfile.gettempdir(), data["error"])
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
self.assertEqual(created, []) # refused before anything was written
async def test_a_short_disk_is_refused_while_streaming_without_a_length(self):
from py.api import maintenance
self._save("keep me")
upload = self._large_db_file(["restored"], 4 * MIB)
spy, created = self._spy_uploads()
# Room for 1 MiB only, and no Content-Length to refuse up front.
with spy, patch.object(maintenance, "restore_room", return_value=MIB):
resp = await self._stream_restore(upload, chunked=True)
data = await resp.json()
self.assertEqual(resp.status, 507)
self.assertEqual(data["code"], "insufficient-space")
self.assertTrue(data["error"].endswith("Nothing was changed."))
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
self.assertEqual(len(created), 1)
self._assert_removed(created)
class TestRestoreRoom(unittest.TestCase):
def test_room_counts_the_upload_its_staged_copy_and_the_swap(self):
from py.api import maintenance
tmp = tempfile.mkdtemp()
try:
db_path = os.path.join(tmp, "prompts.db")
free = 10 * 2**30
with (
patch.object(maintenance, "_free_bytes", return_value=free),
patch.object(maintenance, "_device", return_value="one disk"),
):
room = maintenance.restore_room(db_path)
# Upload + staged copy + WAL growth on one disk, no database yet.
self.assertEqual(room, (free - maintenance.DISK_MARGIN_BYTES) // 3)
self.assertEqual(maintenance.size_label(int(1.6 * 2**30)), "1.6 GB")
self.assertEqual(maintenance.size_label(120 * 2**20), "120.0 MB")
finally:
os.rmdir(tmp)
+986
View File
@@ -0,0 +1,986 @@
"""Duplicate image files on the job registry (#249 System S6, Task 17).
``POST /prompt_manager/jobs/duplicates`` finds content-identical media files
as a ``read`` Activity job; ``GET /prompt_manager/maintenance/duplicates/{id}``
pages the groups it holds; ``duplicate_files`` previews a removal bound to the
reviewed keepers, and its apply job removes only the reviewed, revalidated
non-keepers, moving their database links to the kept copy. The legacy gallery
routes keep their shapes on the same registry kinds.
"""
import asyncio
import hashlib
import json
import os
import sqlite3
import threading
from pathlib import Path
from unittest.mock import patch
try:
from tests.open_handles import open_handles_under
from tests.test_api_maintenance import ApplyTestCase, _digest
except ImportError: # discovered with tests/ as the top-level directory
from open_handles import open_handles_under
from test_api_maintenance import ApplyTestCase, _digest
from py.api import duplicate_files # noqa: E402
from py.config import GalleryConfig, PromptManagerConfig # noqa: E402
WAIT = 10
FIND = "/prompt_manager/jobs/duplicates"
GROUPS = "/prompt_manager/maintenance/duplicates/"
PREVIEW = "/prompt_manager/maintenance/preview"
STREAM = "/prompt_manager/scan_duplicates/stream"
STATUS = "/prompt_manager/scan_duplicates/status"
LEGACY_DELETE = "/prompt_manager/delete_duplicate_images"
def _sse_payloads(body):
return [
json.loads(line[len("data: ") :])
for line in body.splitlines()
if line.startswith("data: ")
]
def _sha(path):
with open(path, "rb") as handle:
return hashlib.sha256(handle.read()).hexdigest()
class DuplicateFilesTestCase(ApplyTestCase):
async def get_application(self):
app = await super().get_application()
GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]
return app
def write(self, name, data, *, folder=None, mtime=None):
path = Path(folder or self.output_dir) / name
path.parent.mkdir(parents=True, exist_ok=True)
path.write_bytes(data)
if mtime is not None:
os.utime(path, (mtime, mtime))
return path
def second_source(self):
"""A second valid source outside ComfyUI (an extra gallery root)."""
second = self.outside_dir / "renders"
second.mkdir(parents=True, exist_ok=True)
os.environ["PROMPT_MANAGER_EXTRA_GALLERY_ROOTS"] = str(second)
return second
def trio(self):
"""a.png (oldest), b.png and c.png share bytes; d.png is unique."""
self.write("a.png", b"same bytes", mtime=1_000_000)
self.write("b.png", b"same bytes", mtime=2_000_000)
self.write("sub/c.png", b"same bytes", mtime=3_000_000)
self.write("d.png", b"unique", mtime=4_000_000)
async def find(self):
response = await self.client.request("POST", FIND)
text = await response.text()
self.assertEqual(response.status, 202, text)
self.assertNotIn(self.tmpdir, text)
job = self.api._jobs.get(json.loads(text)["job"]["id"])
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
return job
async def groups(self, job_id, status=200, **query):
params = "&".join(f"{key}={value}" for key, value in query.items())
response = await self.client.request("GET", f"{GROUPS}{job_id}?{params}")
text = await response.text()
self.assertEqual(response.status, status, text)
self.assertNotIn(self.tmpdir, text)
return json.loads(text)
async def review(self, job_id, keep, status=200):
return await self.post(
PREVIEW,
{
"operation": "duplicate_files",
"options": {"job_id": job_id, "keep": keep},
},
status,
)
async def found_group(self):
job = await self.find()
[group] = (await self.groups(job.id))["groups"]
return job, group
@staticmethod
def keep_of(group, path, root_index=None):
[entry] = [
f
for f in group["files"]
if f["path"] == path and root_index in (None, f["root_index"])
]
return {
"group_id": group["group_id"],
"root_index": entry["root_index"],
"path": entry["path"],
}
async def remove_all_but(self, path):
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, path)])
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
return plan, cleanup
async def until(self, predicate):
"""Wait for a server-side state (a frame at a time), never a fixed timer."""
deadline = asyncio.get_running_loop().time() + WAIT
while not predicate():
if asyncio.get_running_loop().time() > deadline:
self.fail("the expected server state never appeared")
await asyncio.sleep(0)
def gate_hashing(self):
gate = threading.Event()
calls = []
original = self.api._calculate_file_hash
def gated(path):
calls.append(path)
gate.wait(WAIT)
return original(path)
self.api._calculate_file_hash = gated
self.addCleanup(gate.set)
return gate, calls
class TestFindJob(DuplicateFilesTestCase):
async def test_find_is_listed_in_activity_and_legacy_stream_attaches(self):
self.trio()
gate, calls = self.gate_hashing()
response = await self.client.request("POST", FIND)
self.assertEqual(response.status, 202)
summary = (await response.json())["job"]
listed = (
await (await self.client.request("GET", "/prompt_manager/jobs")).json()
)["jobs"]
self.assertEqual([j["id"] for j in listed], [summary["id"]])
self.assertEqual(summary["kind"], "duplicates")
self.assertEqual(summary["title"], "Find duplicate image files")
self.assertEqual(summary["admission"], "read")
self.assertTrue(summary["cancellable"])
self.assertEqual(summary["counters"], {"groups_found": 0})
# The legacy gallery attaches to the same job, and its status says so.
legacy = asyncio.ensure_future(self.client.request("POST", STREAM))
job = self.api._jobs.get(summary["id"])
await self.until(lambda: job.subscriber_count > 0)
status = await (await self.client.request("GET", STATUS)).json()
self.assertTrue(status["running"])
gate.set()
body = await (await legacy).text()
job = self.api._jobs.get(summary["id"])
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
events = _sse_payloads(body)
done = events[-1]
self.assertEqual(done["type"], "complete")
self.assertEqual(done["found"], 1)
self.assertEqual(
sorted(i["filename"] for i in done["duplicates"][0]["images"]),
["a.png", "b.png", "c.png"],
)
self.assertEqual(len(calls), 4) # one pass: each file hashed once
self.assertEqual(job.state, "finished")
self.assertEqual(job.counters, {"groups_found": 1})
[row] = job.outcomes["groups_found"]
self.assertEqual(row["path"], "a.png")
self.assertEqual(row["root_index"], 0)
self.assertEqual(row["detail"], "3 identical files · 10 B each")
self.assertEqual(job.status_line, "Found 1 group of duplicate files")
finished = await (await self.client.request("GET", STATUS)).json()
self.assertFalse(finished["running"])
self.assertEqual(finished["last_event"]["found"], 1)
async def test_legacy_stream_starts_a_registry_job(self):
self.trio()
body = await (await self.client.request("POST", STREAM)).text()
self.assertEqual(_sse_payloads(body)[-1]["type"], "complete")
[job] = self.api._jobs.list()
self.assertEqual(job.kind, "duplicates")
self.assertEqual(job.title, "Find duplicate image files (gallery page)")
self.assertNotIn(self.tmpdir, body)
# Its groups are reviewable in the workspace too.
page = await self.groups(job.id)
self.assertEqual(page["total"], 1)
async def test_find_covers_every_source_with_its_index(self):
second = self.second_source()
self.write("x.png", b"twin", folder=second)
GalleryConfig.MONITORING_DIRECTORIES = [
self._root(self.output_dir),
self._root(second),
]
self.write("x.png", b"twin")
job = await self.find()
[group] = (await self.groups(job.id))["groups"]
self.assertEqual(
sorted((f["root_index"], f["path"]) for f in group["files"]),
[(0, "x.png"), (1, "x.png")],
)
self.assertEqual({f["root"] for f in group["files"]}, {"output", "renders"})
async def test_find_can_be_cancelled(self):
self.trio()
gate, _calls = self.gate_hashing()
response = await self.client.request("POST", FIND)
job_id = (await response.json())["job"]["id"]
stream = asyncio.ensure_future(self.client.request("POST", STREAM))
found = self.api._jobs.get(job_id)
await self.until(lambda: found.subscriber_count > 0)
cancel = await self.client.request(
"POST", f"/prompt_manager/jobs/{job_id}/cancel"
)
self.assertTrue((await cancel.json())["cancellation_requested"])
gate.set()
job = self.api._jobs.get(job_id)
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
self.assertEqual(job.state, "cancelled")
last = _sse_payloads(await (await stream).text())[-1]
self.assertEqual(last, {"type": "error", "message": "Duplicate scan cancelled"})
await self.groups(job_id, 409)
async def test_second_find_is_busy_and_names_the_running_job(self):
self.trio()
gate, _calls = self.gate_hashing()
first = (await (await self.client.request("POST", FIND)).json())["job"]
second = await self.client.request("POST", FIND)
self.assertEqual(second.status, 409)
body = await second.json()
self.assertEqual(body["code"], "busy")
self.assertEqual(body["running"]["id"], first["id"])
gate.set()
async def test_failed_find_is_coded_and_releases_the_kind(self):
self.write("a.png", b"x")
with patch.object(
self.api, "_hash_media_groups", side_effect=OSError(f"boom {self.tmpdir}")
):
with self.assertLogs("prompt_manager", "ERROR"):
job = await self.find()
self.assertEqual(job.state, "failed")
self.assertEqual(job.status_line, "The duplicate search failed.")
self.assertNotIn("boom", json.dumps(job.report()))
await self.groups(job.id, 409)
self.assertEqual((await self.find()).state, "finished")
class TestGroupPages(DuplicateFilesTestCase):
async def test_groups_are_paged_ten_at_a_time(self):
for index in range(12):
for copy in range(2):
self.write(f"g{index:02d}_{copy}.png", f"group {index}".encode())
job = await self.find()
first = await self.groups(job.id)
self.assertEqual((first["total"], first["offset"], first["limit"]), (12, 0, 10))
self.assertEqual(len(first["groups"]), 10)
rest = await self.groups(job.id, offset=10, limit=10)
self.assertEqual(len(rest["groups"]), 2)
ids = [g["group_id"] for g in first["groups"] + rest["groups"]]
self.assertEqual(len(set(ids)), 12)
entry = first["groups"][0]["files"][0]
for key in ("path", "root", "root_index", "url", "size", "media_type"):
self.assertIn(key, entry)
self.assertTrue(entry["url"].endswith("?root=0"))
for query in ({"limit": 11}, {"limit": 0}, {"offset": -1}, {"offset": "x"}):
with self.subTest(query=query):
reply = await self.groups(job.id, 400, **query)
self.assertEqual(reply["code"], "invalid-request")
async def test_unknown_or_running_job_has_no_groups(self):
reply = await self.groups("f" * 32, 404)
self.assertEqual(reply["code"], "not-found")
self.trio()
gate, _calls = self.gate_hashing()
job_id = (await (await self.client.request("POST", FIND)).json())["job"]["id"]
reply = await self.groups(job_id, 409)
self.assertEqual(reply["code"], "job-running")
gate.set()
class TestReview(DuplicateFilesTestCase):
async def test_preview_lists_only_non_keepers_and_writes_nothing(self):
self.trio()
prompt = self.save("Linked")
self.api.db.link_image_to_prompt(prompt, str(self.output_dir / "b.png"))
job, group = await self.found_group()
db_path = self.api.db.model.db_path
files = [db_path, db_path + "-wal"]
before = [_digest(path) for path in files]
media = {p: _digest(p) for p in self.output_dir.rglob("*") if p.is_file()}
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
self.assertEqual([_digest(path) for path in files], before)
self.assertEqual(
{p: _digest(p) for p in self.output_dir.rglob("*") if p.is_file()},
media,
)
self.assertEqual(plan["operation"], "duplicate_files")
self.assertEqual(
plan["scope"], {"count": 2, "label": "2 duplicate files in 1 group"}
)
self.assertTrue(plan["media"]["deletes_files"])
self.assertEqual(
plan["media"]["label"],
"Deletes 2 files (20 B) from disk; the kept copies stay. "
"The safety backup covers the database, not these files.",
)
self.assertEqual(
[row["label"] for row in plan["details"]], ["b.png", "sub/c.png"]
)
self.assertEqual(
plan["details"][0]["detail"],
"output · a copy of a.png · 1 database link moves to the kept copy",
)
self.assertEqual(plan["protected"], {"count": 0, "reasons": []})
async def test_cannot_delete_keeper_or_outside_group(self):
self.trio()
self.write("e.png", b"other pair")
self.write("f.png", b"other pair")
job = await self.find()
groups = (await self.groups(job.id))["groups"]
trio = next(g for g in groups if len(g["files"]) == 3)
pair = next(g for g in groups if len(g["files"]) == 2)
bad = [
# a file of another group named as this group's keeper
[{**self.keep_of(trio, "a.png"), "path": "e.png"}],
# a file that is in no group
[{**self.keep_of(trio, "a.png"), "path": "d.png"}],
# the right path under another source index
[{**self.keep_of(trio, "a.png"), "root_index": 1}],
[{**self.keep_of(trio, "a.png"), "group_id": "0" * 64}],
[{**self.keep_of(trio, "a.png"), "extra": 1}],
[],
"a.png",
]
for keep in bad:
with self.subTest(keep=keep):
reply = await self.review(job.id, keep, 400)
self.assertEqual(reply["code"], "invalid-request")
# Two keepers of one group: only the third copy goes.
plan = await self.review(
job.id,
[self.keep_of(trio, "a.png"), self.keep_of(trio, "sub/c.png")],
)
self.assertEqual([row["label"] for row in plan["details"]], ["b.png"])
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
self.assertEqual(cleanup.counters["removed"], 1)
self.assertTrue((self.output_dir / "a.png").is_file())
self.assertTrue((self.output_dir / "sub" / "c.png").is_file())
self.assertFalse((self.output_dir / "b.png").exists())
# The untouched group is untouched.
self.assertTrue((self.output_dir / "e.png").is_file())
self.assertTrue((self.output_dir / "f.png").is_file())
self.assertEqual(len(pair["files"]), 2)
async def test_a_source_removed_since_the_find_is_refused(self):
self.trio()
job, group = await self.found_group()
GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.second_source())]
reply = await self.review(job.id, [self.keep_of(group, "a.png")], 409)
self.assertEqual(reply["code"], "sources-changed")
def test_containment_is_decided_as_text_before_the_filesystem(self):
root = str(self.output_dir)
outside = os.path.join(str(self.outside_dir), "x.png")
touched = []
real_lstat = os.lstat
def lstat(path, *args, **kwargs):
touched.append(str(path))
return real_lstat(path, *args, **kwargs)
with patch.object(duplicate_files.os, "lstat", side_effect=lstat):
for candidate in (
outside,
os.path.join(root, "..", "outside", "x.png"),
os.path.join(root, "thumbnails", "a_thumb.png"),
os.path.join(root, "notes.txt"),
):
with self.subTest(candidate=candidate):
reason = duplicate_files.file_problem(candidate, root, "0" * 64, 1)
self.assertEqual(reason, duplicate_files.OUTSIDE_SOURCE)
self.assertEqual(touched, [])
@staticmethod
def _can_symlink(directory):
probe = Path(directory) / "probe-link"
try:
os.symlink(str(Path(directory) / "x"), str(probe))
except (OSError, NotImplementedError):
return False
probe.unlink()
return True
async def test_symlink_rejected(self):
if not self._can_symlink(self.tmpdir): # pragma: no cover - Windows w/o rights
self.skipTest("symlinks are not available")
self.trio()
job, group = await self.found_group()
# b.png becomes a link to the keeper after the find: never removed.
(self.output_dir / "b.png").unlink()
os.symlink(str(self.output_dir / "a.png"), str(self.output_dir / "b.png"))
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
self.assertEqual([row["label"] for row in plan["details"]], ["sub/c.png"])
self.assertEqual(
plan["protected"],
{"count": 1, "reasons": ["1 file is a link, not a regular file"]},
)
# A keeper that became a link protects its whole group.
(self.output_dir / "a.png").rename(self.output_dir / "real.png")
os.symlink(str(self.output_dir / "real.png"), str(self.output_dir / "a.png"))
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
self.assertEqual(plan["scope"]["count"], 0)
self.assertEqual(
plan["protected"]["reasons"],
["2 files kept: their kept copy changed or is gone"],
)
class TestCleanupJob(DuplicateFilesTestCase):
async def test_cleanup_is_a_records_job_with_a_safety_backup(self):
self.trio()
plan, cleanup = await self.remove_all_but("a.png")
summary = cleanup.summary()
self.assertEqual(summary["kind"], "duplicate_cleanup")
self.assertEqual(summary["title"], "Remove duplicate image files")
self.assertEqual(summary["admission"], "records")
self.assertFalse(summary["cancellable"])
self.assertEqual(summary["scope"], {"operation": "duplicate_files"})
self.assertEqual(cleanup.state, "finished", cleanup.log)
self.assertEqual(cleanup.counters, {"removed": 2, "kept": 1, "errors": 0})
self.assertEqual(
sorted(row["path"] for row in cleanup.outcomes["removed"]),
["b.png", "sub/c.png"],
)
self.assertEqual(cleanup.outcomes["kept"][0]["path"], "a.png")
[backup] = self.backups()
self.assertEqual(
cleanup.status_line,
f"Removed 2 duplicate files · Safety backup {backup}",
)
self.assertEqual(
_sha(self.output_dir / "a.png"), hashlib.sha256(b"same bytes").hexdigest()
)
self.assertFalse((self.output_dir / "b.png").exists())
self.assertFalse((self.output_dir / "sub" / "c.png").exists())
self.assertNotIn(self.tmpdir, json.dumps(cleanup.report()))
# Single use, like every plan.
again = await self.apply(plan["token"], 409)
self.assertEqual(again["code"], "plan-used")
async def test_thumbnails_of_removed_files_are_removed(self):
self.trio()
self.write("clip.mp4", b"video")
self.write("copy.mp4", b"video")
kept = self.write("thumbnails/a_thumb.png", b"t")
gone = self.write("thumbnails/b_thumb.png", b"t")
video_thumb = self.write("thumbnails/copy_thumb.jpg", b"t")
job = await self.find()
groups = (await self.groups(job.id))["groups"]
keep = []
for group in groups:
paths = sorted(f["path"] for f in group["files"])
keep.append(
self.keep_of(group, "a.png" if "a.png" in paths else "clip.mp4")
)
plan = await self.review(job.id, keep)
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
self.assertTrue(kept.is_file())
self.assertFalse(gone.exists())
self.assertFalse(video_thumb.exists())
async def test_changed_content_skips_file(self):
self.trio()
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
self.write("b.png", b"new bytes") # changed since the review
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
self.assertEqual(cleanup.counters, {"removed": 1, "kept": 2, "errors": 0})
self.assertEqual((self.output_dir / "b.png").read_bytes(), b"new bytes")
self.assertFalse((self.output_dir / "sub" / "c.png").exists())
rows = {row["path"]: row["detail"] for row in cleanup.outcomes["kept"]}
self.assertEqual(rows["b.png"], duplicate_files.CHANGED)
async def test_a_changed_keeper_keeps_its_whole_group(self):
self.trio()
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
self.write("a.png", b"edited keeper")
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
self.assertEqual(cleanup.counters["removed"], 0)
self.assertTrue((self.output_dir / "b.png").is_file())
self.assertTrue((self.output_dir / "sub" / "c.png").is_file())
self.assertTrue(
cleanup.status_line.startswith(
"Removed 0 duplicate files, 2 not removed · "
)
)
rows = {row["path"]: row["detail"] for row in cleanup.outcomes["kept"]}
self.assertEqual(rows["b.png"], duplicate_files.KEEPER_CHANGED)
async def test_partial_failure_reports_each_path(self):
self.trio()
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
real = duplicate_files._unlink
def unlink(path):
if os.path.basename(path) == "c.png":
raise PermissionError(13, f"Permission denied: {path}")
real(path)
self.write("z.png", b"same bytes") # never reviewed: not touched
with patch.object(duplicate_files, "_unlink", side_effect=unlink):
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
self.assertEqual(cleanup.state, "failed")
self.assertEqual(cleanup.counters, {"removed": 1, "kept": 1, "errors": 1})
self.assertEqual(
[row["path"] for row in cleanup.outcomes["removed"]], ["b.png"]
)
[error] = cleanup.outcomes["errors"]
self.assertEqual(
(error["path"], error["detail"]),
("sub/c.png", duplicate_files.LOCKED_FILE),
)
self.assertEqual(
cleanup.status_line,
"Stopped after an error · Removed 1 duplicate file, 1 not removed · "
"Safety backup " + self.backups()[0],
)
report = json.dumps(cleanup.report())
self.assertNotIn("Permission denied", report)
self.assertNotIn("Errno", report)
self.assertTrue((self.output_dir / "z.png").is_file())
async def test_stop_on_failure_leaves_later_files_untouched(self):
self.write("a.png", b"same", mtime=1)
for name in ("b.png", "c.png", "e.png"):
self.write(name, b"same", mtime=2)
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
with (
patch.object(
duplicate_files, "_unlink", side_effect=OSError(5, "I/O error")
),
self.assertLogs("prompt_manager", "WARNING"),
):
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
self.assertEqual(cleanup.counters, {"removed": 0, "kept": 3, "errors": 1})
self.assertEqual(
cleanup.outcomes["errors"][0]["detail"], duplicate_files.DELETE_FAILED
)
not_attempted = [
row
for row in cleanup.outcomes["kept"]
if row["detail"] == duplicate_files.NOT_ATTEMPTED
]
self.assertEqual(len(not_attempted), 2)
for name in ("b.png", "c.png", "e.png"):
self.assertTrue((self.output_dir / name).is_file())
async def test_locked_file_gives_a_coded_message(self):
self.trio()
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
target = self.output_dir / "b.png"
handle = None
if os.name == "nt": # pragma: no cover - the native sharing violation
handle = open(target, "rb")
context = patch.object(duplicate_files, "_unlink", duplicate_files._unlink)
else:
context = patch.object(
duplicate_files,
"_unlink",
side_effect=PermissionError(
32, "The process cannot access the file", str(target)
),
)
try:
with context:
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
finally:
if handle is not None: # pragma: no cover - Windows
handle.close()
[error] = cleanup.outcomes["errors"]
self.assertEqual(error["detail"], duplicate_files.LOCKED_FILE)
self.assertNotIn("WinError", json.dumps(cleanup.report()))
self.assertTrue(target.is_file())
async def test_handles_closed_before_unlink(self):
self.trio()
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
real = duplicate_files._unlink
open_at_unlink = []
def unlink(path):
open_at_unlink.extend(open_handles_under(path))
real(path)
with patch.object(duplicate_files, "_unlink", side_effect=unlink):
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
self.assertEqual(cleanup.counters["removed"], 2)
self.assertEqual(open_at_unlink, [])
async def test_db_links_reconciled_only_for_removed_files(self):
self.trio()
db = self.api.db
moved = self.save("Moves to the keeper")
merged = self.save("Already on the keeper")
db.link_image_to_prompt(moved, str(self.output_dir / "b.png"))
db.link_image_to_prompt(moved, str(self.output_dir / "d.png"))
db.link_image_to_prompt(merged, str(self.output_dir / "a.png"))
db.link_image_to_prompt(merged, str(self.output_dir / "sub" / "c.png"))
db.set_image_favorite(
duplicate_files.favorite_root_key(self.output_dir), "b.png"
)
plan, cleanup = await self.remove_all_but("a.png")
self.assertEqual(
plan["details"][0]["detail"],
"output · a copy of a.png · 1 database link moves to the kept copy",
)
def files_of(prompt_id):
return sorted(
os.path.relpath(image["image_path"], self.output_dir)
for image in db.get_prompt_images(prompt_id)
)
self.assertEqual(files_of(moved), ["a.png", "d.png"])
self.assertEqual(files_of(merged), ["a.png"])
root_key = duplicate_files.favorite_root_key(self.output_dir)
self.assertEqual(db.get_root_favorites(root_key), {"a.png"})
self.assertIn(
"Moved 2 database links to the kept copies", "\n".join(cleanup.log)
)
async def test_a_failed_removal_leaves_its_links_on_the_kept_copy(self):
self.trio()
prompt = self.save("Stays on c")
self.api.db.link_image_to_prompt(prompt, str(self.output_dir / "sub" / "c.png"))
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
real = duplicate_files._unlink
def unlink(path):
if os.path.basename(path) == "c.png":
raise PermissionError(13, "denied")
real(path)
with patch.object(duplicate_files, "_unlink", side_effect=unlink):
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
# Its records moved to the identical kept copy before the unlink failed.
[image] = self.api.db.get_prompt_images(prompt)
self.assertTrue(image["image_path"].endswith("a.png"))
self.assertTrue((self.output_dir / "sub" / "c.png").is_file())
async def test_cleanup_and_find_exclude_each_other(self):
self.trio()
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
gate, _calls = self.gate_hashing()
find = (await (await self.client.request("POST", FIND)).json())["job"]
busy = await self.apply(plan["token"], 409)
self.assertEqual((busy["code"], busy["running"]["id"]), ("busy", find["id"]))
gate.set()
await asyncio.wait_for(
asyncio.shield(self.api._jobs.get(find["id"]).task), WAIT
)
# The token survived the refusal; the cleanup holds off a new find.
release = asyncio.Event()
real = self.api._run_duplicate_cleanup
async def held(*args):
await release.wait()
return await real(*args)
with patch.object(self.api, "_run_duplicate_cleanup", side_effect=held):
body = await self.apply(plan["token"])
refused = await self.client.request("POST", FIND)
self.assertEqual(refused.status, 409)
legacy = await self.client.request("POST", STREAM)
self.assertEqual(legacy.status, 409)
release.set()
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
async def test_cleanup_refused_up_front_when_the_backup_does_not_fit(self):
from py.api import maintenance
self.trio()
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
with patch.object(maintenance, "_free_bytes", return_value=0):
reply = await self.apply(plan["token"], 507)
self.assertEqual(reply["code"], "insufficient-space")
self.assertTrue((self.output_dir / "b.png").is_file())
async def test_policy_off_removes_without_a_backup(self):
PromptManagerConfig.SAFETY_BACKUP = False
self.trio()
_plan, cleanup = await self.remove_all_but("a.png")
self.assertEqual(self.backups(), [])
self.assertEqual(
cleanup.status_line, "Removed 2 duplicate files · No safety backup"
)
class TestLegacyDuplicateRoutes(DuplicateFilesTestCase):
async def test_legacy_delete_duplicates_unchanged(self):
self.trio()
response = await self.client.request(
"POST", LEGACY_DELETE, json={"image_paths": ["b.png", "nope.png"]}
)
body = await response.json()
self.assertEqual(response.status, 200)
self.assertEqual(
set(body),
{"success", "deleted_count", "failed_count", "message", "failed_files"},
)
self.assertEqual((body["deleted_count"], body["failed_count"]), (1, 1))
self.assertEqual(body["failed_files"], ["nope.png (file not found)"])
[hold] = self.api._jobs.list()
self.assertEqual(
(hold.kind, hold.title, hold.admission),
("duplicate_cleanup", "Duplicate image delete (gallery page)", "records"),
)
async def test_legacy_delete_refused_while_a_find_runs(self):
self.trio()
gate, _calls = self.gate_hashing()
await self.client.request("POST", FIND)
response = await self.client.request(
"POST", LEGACY_DELETE, json={"image_paths": ["b.png"]}
)
self.assertEqual(response.status, 409)
self.assertEqual(
await response.json(),
{"success": False, "error": "Find duplicate image files already running"},
)
self.assertTrue((self.output_dir / "b.png").is_file())
gate.set()
async def test_legacy_scan_holds_the_duplicates_kind(self):
self.trio()
gate, _calls = self.gate_hashing()
legacy = asyncio.ensure_future(
self.client.request("GET", "/prompt_manager/scan_duplicates")
)
await self.until(lambda: self.api._jobs.active("duplicates") is not None)
[hold] = self.api._jobs.list()
self.assertEqual(
(hold.kind, hold.title, hold.admission, hold.cancellable),
("duplicates", "Find duplicate image files (gallery page)", "read", False),
)
refused = await self.client.request("POST", FIND)
self.assertEqual(refused.status, 409)
# A held legacy scan publishes no events: the stream does not attach.
stream = await self.client.request("POST", STREAM)
self.assertEqual(stream.status, 409)
gate.set()
body = await (await legacy).json()
self.assertEqual(set(body), {"success", "duplicates", "message"})
self.assertEqual(body["duplicates"][0]["count"], 3)
async def test_legacy_stream_does_not_attach_to_a_find_missing_its_folder(self):
other = self.second_source()
self.write("x.png", b"x", folder=other)
GalleryConfig.MONITORING_DIRECTORIES = [self._root(other)]
self.api._find_comfyui_output_dir = lambda: str(self.output_dir)
gate, _calls = self.gate_hashing()
await self.client.request("POST", FIND)
stream = await self.client.request("POST", STREAM)
self.assertEqual(stream.status, 409)
gate.set()
class TestLinkQueries(DuplicateFilesTestCase):
def test_link_queries_use_the_path_indexes(self):
conn = sqlite3.connect(self.api.db.model.db_path)
try:
for sql in duplicate_files.query_plan_sql():
plan = " ".join(
row[3]
for row in conn.execute(
"EXPLAIN QUERY PLAN " + sql, [None] * sql.count("?")
)
)
with self.subTest(sql=sql):
self.assertNotIn("SCAN generated_images", plan)
self.assertNotIn("SCAN image_favorites", plan)
finally:
conn.close()
class TestReviewHelpers(DuplicateFilesTestCase):
"""The pure checks of py/api/duplicate_review.py."""
def test_file_problem_on_disk(self):
from py.api import duplicate_review as review
root = str(self.output_dir)
good = self.write("ok.png", b"bytes")
digest = hashlib.sha256(b"bytes").hexdigest()
self.assertIsNone(review.file_problem(str(good), root, digest, 5))
self.assertEqual(
review.file_problem(str(good), root, digest, 6), review.CHANGED_CONTENT
)
self.assertEqual(
review.file_problem(str(good), root, "0" * 64, 5), review.CHANGED_CONTENT
)
self.assertEqual(
review.file_problem(os.path.join(root, "gone.png"), root, digest, 5),
review.GONE,
)
folder = self.output_dir / "folder.png"
folder.mkdir()
self.assertEqual(review.file_problem(str(folder), root, digest, 5), review.LINK)
with patch.object(review.os, "lstat", side_effect=OSError("io")):
self.assertEqual(
review.file_problem(str(good), root, digest, 5), review.UNREADABLE
)
def unreadable(_path):
raise OSError("io")
self.assertEqual(
review.file_problem(str(good), root, digest, 5, unreadable),
review.UNREADABLE,
)
def vanished(_path):
raise FileNotFoundError
self.assertEqual(
review.file_problem(str(good), root, digest, 5, vanished), review.GONE
)
def test_a_linked_folder_is_outside_its_source(self):
from py.api import duplicate_review as review
if not TestReview._can_symlink(self.tmpdir): # pragma: no cover
self.skipTest("symlinks are not available")
target = self.write("x.png", b"bytes", folder=self.outside_dir)
os.symlink(str(self.outside_dir), str(self.output_dir / "linked"))
digest = hashlib.sha256(b"bytes").hexdigest()
path = os.path.join(str(self.output_dir), "linked", "x.png")
self.assertTrue(os.path.isfile(path))
self.assertEqual(
review.file_problem(path, str(self.output_dir), digest, 5),
review.OUTSIDE_SOURCE,
)
self.assertTrue(target.is_file())
def test_options_are_validated(self):
from py.api import duplicate_review as review
keep = {"group_id": "g", "root_index": 0, "path": "a.png"}
cases = [
None,
{"job_id": "x"},
{"job_id": "", "keep": [keep]},
{"job_id": "x" * 65, "keep": [keep]},
{"job_id": "x", "keep": {}},
{"job_id": "x", "keep": [keep] * (review.MAX_KEEP + 1)},
{"job_id": "x", "keep": [{**keep, "group_id": 1}]},
{"job_id": "x", "keep": [{**keep, "root_index": True}]},
{"job_id": "x", "keep": [{**keep, "root_index": -1}]},
{"job_id": "x", "keep": [{**keep, "path": ""}]},
{"job_id": "x", "keep": [["g", 0, "a.png"]]},
]
for options in cases:
with self.subTest(options=options):
self.assertIsNotNone(review._parse_options(options)[2])
job_id, parsed, message = review._parse_options(
{"job_id": "x", "keep": [{**keep, "root_index": None}]}
)
self.assertEqual((job_id, parsed, message), ("x", [("g", None, "a.png")], None))
def test_sizes_read_naturally(self):
from py.api.duplicate_review import human_bytes
self.assertEqual(human_bytes(10), "10 B")
self.assertEqual(human_bytes(1536), "1.5 KB")
self.assertEqual(human_bytes(3 * 2**20), "3.0 MB")
self.assertEqual(human_bytes(2**40 * 5000), "5,000.0 TB")
async def test_a_kept_copy_listed_twice_counts_once(self):
self.trio()
job, group = await self.found_group()
keep = self.keep_of(group, "a.png")
plan = await self.review(job.id, [keep, keep])
self.assertEqual(plan["scope"]["count"], 2)
async def test_preview_failure_is_coded(self):
self.trio()
job, group = await self.found_group()
with (
patch.object(
self.api, "_review_duplicate_files", side_effect=OSError("boom")
),
self.assertLogs("prompt_manager", "ERROR"),
):
reply = await self.review(job.id, [self.keep_of(group, "a.png")], 500)
self.assertEqual(reply["error"], "Failed to review maintenance")
async def test_a_source_removed_before_the_apply_deletes_nothing(self):
self.trio()
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.second_source())]
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
self.assertEqual(cleanup.state, "failed")
self.assertEqual(cleanup.last_event["code"], "sources-changed")
self.assertEqual(self.backups(), [])
self.assertTrue((self.output_dir / "b.png").is_file())
async def test_a_records_move_failure_deletes_nothing(self):
self.trio()
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
with (
patch.object(
self.api.db,
"move_copy_records",
side_effect=sqlite3.OperationalError("x"),
),
self.assertLogs("prompt_manager", "ERROR"),
):
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
self.assertEqual(cleanup.state, "failed")
self.assertEqual(
cleanup.outcomes["errors"][-1]["detail"], duplicate_files.RECORDS_FAILED
)
self.assertEqual(cleanup.counters["removed"], 0)
self.assertTrue((self.output_dir / "b.png").is_file())
self.assertTrue((self.output_dir / "sub" / "c.png").is_file())
+558
View File
@@ -0,0 +1,558 @@
"""Duplicate-file removal safety (#249 System S6 final review: C1, I1, I2, I5, M9).
- C1: one physical file reached through two listed sources, a symlinked
source or a hardlink is never offered or removed as a copy of itself.
- I1: each file's records move to its verified kept copy before that file
is unlinked, so no database row ever names a deleted file.
- I2: links stored under another spelling of the source move too.
- I5: image tags move to the kept copy (their union) and the removed copy's
pending Needs review items are settled.
- M9: the preview binds files by size, mtime and inode without hashing.
"""
import asyncio
import hashlib
import os
import sqlite3
import types
from pathlib import Path
from unittest.mock import patch
try:
from tests.open_handles import open_handles_under
from tests import test_api_duplicate_files as dup_tests
except ImportError: # discovered with tests/ as the top-level directory
import test_api_duplicate_files as dup_tests
from open_handles import open_handles_under
from py.api import duplicate_files # noqa: E402
from py.config import GalleryConfig # noqa: E402
WAIT = 10
def _locator(path):
return os.path.normcase(os.path.realpath(str(path)))
class SafetyTestCase(dup_tests.DuplicateFilesTestCase):
def need_symlinks(self):
if not dup_tests.TestReview._can_symlink(self.tmpdir): # pragma: no cover
self.skipTest("symlinks are not available")
def linked_files(self):
"""Every image record's path that names a file that does not exist."""
with self.api.db.model.get_connection() as conn:
rows = conn.execute("SELECT image_path FROM generated_images").fetchall()
return [row[0] for row in rows if not os.path.exists(row[0])]
async def apply_plan(self, plan):
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
return cleanup
class TestOnePhysicalFile(SafetyTestCase):
"""C1: the kept copy survives every alias of itself."""
async def test_a_folder_listed_twice_never_loses_its_only_copy(self):
self.need_symlinks()
only = self.write("a.png", b"the only copy")
alias = Path(self.tmpdir) / "Output-alias"
os.symlink(str(self.output_dir), str(alias))
with patch.object(
self.api, "_get_all_output_dirs", return_value=[self.output_dir, alias]
):
job = await self.find()
page = await self.groups(job.id)
self.assertEqual(page["total"], 0)
self.assertEqual(only.read_bytes(), b"the only copy")
async def test_a_symlinked_source_inside_another_is_one_file(self):
self.need_symlinks()
only = self.write("sub/c.png", b"one file")
linked = self.output_dir / "linked" # walked by neither as a link
os.symlink(str(self.output_dir / "sub"), str(linked))
GalleryConfig.MONITORING_DIRECTORIES = [
self._root(self.output_dir),
str(linked),
]
self.assertEqual(len(self.api._get_all_output_dirs()), 2)
job = await self.find()
self.assertEqual((await self.groups(job.id))["total"], 0)
self.assertTrue(only.is_file())
async def test_a_hardlink_is_never_a_copy_to_remove(self):
first = self.write("a.png", b"same bytes", mtime=1)
os.link(str(first), str(self.output_dir / "b.png"))
self.write("c.png", b"same bytes", mtime=3)
job, group = await self.found_group()
self.assertEqual(len(group["files"]), 2) # a.png and b.png are one file
# c.png becomes a hardlink of the keeper after the find: protected.
(self.output_dir / "c.png").unlink()
os.link(str(first), str(self.output_dir / "c.png"))
kept = group["files"][0]["path"]
plan = await self.review(job.id, [self.keep_of(group, kept)])
self.assertEqual(plan["scope"]["count"], 0)
self.assertEqual(
plan["protected"]["reasons"], ["1 file is the same file as its kept copy"]
)
async def test_the_job_refuses_a_file_that_became_its_kept_copy(self):
first = self.write("a.png", b"same bytes", mtime=1)
self.write("b.png", b"same bytes", mtime=2)
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
(self.output_dir / "b.png").unlink()
os.link(str(first), str(self.output_dir / "b.png"))
cleanup = await self.apply_plan(plan)
self.assertEqual(cleanup.counters["removed"], 0)
rows = {row["path"]: row["detail"] for row in cleanup.outcomes["kept"]}
self.assertEqual(rows["b.png"], duplicate_files.SAME_FILE)
self.assertTrue((self.output_dir / "b.png").is_file())
class TestRecordsMoveBeforeTheUnlink(SafetyTestCase):
"""I1: a deleted file is never named by a database row."""
async def test_a_failure_between_two_files_leaves_no_row_on_a_deleted_file(self):
self.write("a.png", b"same", mtime=1)
self.write("b.png", b"same", mtime=2)
self.write("c.png", b"same", mtime=3)
prompt = self.save("Linked to every copy")
for name in ("b.png", "c.png"):
self.api.db.link_image_to_prompt(prompt, str(self.output_dir / name))
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
real = self.api.db.move_copy_records
calls = []
def move(*args, **kwargs):
calls.append(args)
if len(calls) == 2:
raise sqlite3.OperationalError("database is locked")
return real(*args, **kwargs)
with patch.object(self.api.db, "move_copy_records", side_effect=move):
with self.assertLogs("prompt_manager", "ERROR"):
cleanup = await self.apply_plan(plan)
self.assertEqual(self.linked_files(), [])
self.assertEqual(cleanup.counters["removed"], 1)
self.assertEqual(cleanup.state, "failed")
[error] = cleanup.outcomes["errors"]
self.assertEqual(error["detail"], duplicate_files.RECORDS_FAILED)
self.assertNotIn("Missing-file", cleanup.status_line)
self.assertEqual(
cleanup.status_line.split(" · ")[:2],
["Stopped after an error", "Removed 1 duplicate file, 1 not removed"],
)
async def test_a_failed_unlink_leaves_its_links_on_the_kept_copy(self):
self.trio()
prompt = self.save("On c")
self.api.db.link_image_to_prompt(prompt, str(self.output_dir / "sub" / "c.png"))
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
real = duplicate_files._unlink
def unlink(path):
if os.path.basename(path) == "c.png":
raise PermissionError(13, "denied")
real(path)
with patch.object(duplicate_files, "_unlink", side_effect=unlink):
with self.assertLogs("prompt_manager", "WARNING"):
cleanup = await self.apply_plan(plan)
[image] = self.api.db.get_prompt_images(prompt)
self.assertEqual(os.path.basename(image["image_path"]), "a.png")
self.assertTrue((self.output_dir / "sub" / "c.png").is_file())
self.assertEqual(self.linked_files(), [])
self.assertEqual(
cleanup.outcomes["errors"][0]["detail"], duplicate_files.LOCKED_FILE
)
class TestEverySpelling(SafetyTestCase):
"""I2: links under the monitor's spelling of a source move too."""
async def test_links_under_a_symlinked_source_spelling_move(self):
self.need_symlinks()
self.trio()
spelled = Path(self.comfy_dir) / "out-link"
os.symlink(str(self.output_dir), str(spelled))
GalleryConfig.MONITORING_DIRECTORIES = [str(spelled)]
prompt = self.save("Linked through the symlink")
self.api.db.link_image_to_prompt(prompt, str(spelled / "b.png"))
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
self.assertIn("1 database link moves", plan["details"][0]["detail"])
await self.apply_plan(plan)
[image] = self.api.db.get_prompt_images(prompt)
self.assertEqual(
_locator(image["image_path"]), _locator(self.output_dir / "a.png")
)
self.assertEqual(self.linked_files(), [])
class TestTagsAndReview(SafetyTestCase):
"""I5: tags join the kept copy; pending review items of the copy settle."""
def tag(self, path, tags):
data = Path(path).read_bytes()
info = os.stat(path)
return self.api.db.replace_image_tags(
locator=_locator(path),
sha256=hashlib.sha256(data).hexdigest(),
size=info.st_size,
mtime_ns=info.st_mtime_ns,
expected_epoch=self.api.db.get_execution_epoch(),
expected_tag_revision=0,
tags=tags,
)
def tags_of(self, path):
data = Path(path).read_bytes()
return self.api.db.read_image_tags(
_locator(path), hashlib.sha256(data).hexdigest()
)["tags"]
def queue(self, path):
data = Path(path).read_bytes()
info = os.stat(path)
return self.api.db.record_review_item(
locator=_locator(path),
spelled_path=os.path.realpath(str(path)),
sha256=hashlib.sha256(data).hexdigest(),
size=info.st_size,
mtime_ns=info.st_mtime_ns,
epoch=self.api.db.get_execution_epoch(),
resolver_version=1,
evidence={"candidates": []},
candidate_count=0,
limited=True,
)
def pending(self):
return self.api.db.review_queue_page(limit=50)["items"]
async def test_tags_move_as_a_union_and_pending_items_settle(self):
self.write("a.png", b"same", mtime=1)
self.write("b.png", b"same", mtime=2)
keeper, removed = self.output_dir / "a.png", self.output_dir / "b.png"
self.tag(keeper, ["kept"])
self.tag(removed, ["moved", "kept"])
self.assertTrue(self.queue(removed))
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
self.assertIn("image tags join the kept copy", plan["details"][0]["detail"])
self.assertIn("1 Needs review item is settled", plan["details"][0]["detail"])
await self.apply_plan(plan)
self.assertEqual(self.tags_of(keeper), ["kept", "moved"])
self.assertEqual(self.pending(), [])
class TestPreviewDoesNotHash(SafetyTestCase):
"""M9: the preview binds by size, mtime and inode; the job hashes."""
async def test_preview_reads_no_file_and_spots_a_touched_one(self):
self.trio()
job, group = await self.found_group()
calls = []
real = self.api._calculate_file_hash
def spy(path):
calls.append(path)
return real(path)
self.api._calculate_file_hash = spy
os.utime(self.output_dir / "b.png", (5_000_000, 5_000_000))
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
self.assertEqual(calls, [])
self.assertEqual(plan["scope"]["count"], 1)
self.assertEqual(
plan["protected"]["reasons"], ["1 file changed since the search"]
)
cleanup = await self.apply_plan(plan)
self.assertEqual(cleanup.counters["removed"], 1)
self.assertGreaterEqual(len(calls), 2) # the job hashes keeper and file
def _fake_stat(real, **changes):
"""A stat result with some fields changed (st_dev, st_ino, st_mtime_ns...)."""
fields = {name: getattr(real, name) for name in dir(real) if name.startswith("st_")}
fields.update(changes)
return types.SimpleNamespace(**fields)
class TestAnotherMount(SafetyTestCase):
"""Round 2, R1: one file reached through two mounts is never removed.
The second source stands for the same folder seen through another
filesystem: each test reports its file on another ``st_dev``.
"""
MAYBE = "1 file may be the same file reached through another mount or share"
def setUp(self):
super().setUp()
self.real_lstat, self.real_stat = os.lstat, os.stat
def mount(self, rel, keeper_rel, *, ino="same", mtime_shift=0):
"""Copies at ``output/keeper_rel`` and ``renders/rel``; the patch."""
second = self.second_source()
GalleryConfig.MONITORING_DIRECTORIES = [
self._root(self.output_dir),
self._root(second),
]
stamp = 1_700_000_000_123_456_789
keeper = self.write(keeper_rel, b"one file")
other = self.write(rel, b"one file", folder=second)
os.utime(keeper, ns=(stamp, stamp))
os.utime(other, ns=(stamp, stamp))
target = os.path.normcase(str(other))
keeper_ino = self.real_stat(keeper).st_ino
def changed(real):
return _fake_stat(
real,
st_dev=real.st_dev + 7,
st_ino=keeper_ino if ino == "same" else real.st_ino + 1000,
st_mtime_ns=real.st_mtime_ns + mtime_shift,
)
def lstat(path, *args, **kwargs):
real = self.real_lstat(path, *args, **kwargs)
return changed(real) if os.path.normcase(str(path)) == target else real
def stat(path, *args, **kwargs):
real = self.real_stat(path, *args, **kwargs)
return changed(real) if os.path.normcase(str(path)) == target else real
patches = (
patch("os.lstat", side_effect=lstat),
patch("os.stat", side_effect=stat),
)
return keeper, other, patches
async def review_with(self, patches, keeper_rel):
# The mount reports its device and file id to the find too.
with patches[0], patches[1]:
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, keeper_rel, 0)])
return job, group, plan
async def test_nfs_and_smb_spellings_of_one_share_are_protected(self):
# NFS and SMB report their own device and file id; SMB rounds mtime.
keeper, other, patches = self.mount(
"x.png", "x.png", ino="other", mtime_shift=-89
)
_job, _group, plan = await self.review_with(patches, "x.png")
self.assertEqual(plan["scope"]["count"], 0)
self.assertEqual(plan["protected"]["reasons"], [self.MAYBE])
self.assertTrue(other.is_file() and keeper.is_file())
async def test_a_loopback_mount_with_the_same_inode_is_protected(self):
keeper, other, patches = self.mount("y.png", "sub/y.png", ino="same")
_job, _group, plan = await self.review_with(patches, "sub/y.png")
self.assertEqual(plan["protected"]["reasons"], [self.MAYBE])
async def test_a_fuse_pool_path_is_protected(self):
# mergerfs, rclone or sshfs synthesize both the device and the inode.
keeper, other, patches = self.mount("x.png", "x.png", ino="other")
_job, _group, plan = await self.review_with(patches, "x.png")
self.assertEqual(plan["protected"]["reasons"], [self.MAYBE])
async def test_a_true_copy_on_another_disk_is_still_offered(self):
keeper, other, patches = self.mount("elsewhere.png", "x.png", ino="other")
_job, _group, plan = await self.review_with(patches, "x.png")
self.assertEqual(plan["scope"]["count"], 1)
async def test_the_job_checks_again_before_the_unlink(self):
keeper, other, patches = self.mount("x.png", "x.png", ino="other")
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "x.png", 0)])
self.assertEqual(plan["scope"]["count"], 1) # looked separate then
with patches[0], patches[1]:
cleanup = await self.apply_plan(plan)
self.assertEqual(cleanup.counters["removed"], 0)
rows = {row["path"]: row["detail"] for row in cleanup.outcomes["kept"]}
self.assertEqual(rows["x.png"], duplicate_files.MAYBE_SAME_FILE)
self.assertTrue(other.is_file() and keeper.is_file())
async def test_a_volume_without_file_ids_is_named_truthfully(self):
keeper, other, patches = self.mount("x.png", "x.png")
def zero(path, *args, **kwargs):
return _fake_stat(self.real_stat(path), st_dev=0, st_ino=0)
with patch("os.lstat", side_effect=zero), patch("os.stat", side_effect=zero):
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "x.png", 0)])
self.assertEqual(
plan["protected"]["reasons"],
["1 file: can't confirm this is a separate file on this drive"],
)
async def test_the_kept_copy_is_held_open_during_the_unlink(self):
self.trio()
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
real = duplicate_files._unlink
held = []
def unlink(path):
held.append(bool(open_handles_under(self.output_dir / "a.png")))
real(path)
if not os.path.isdir("/proc/self/fd"): # pragma: no cover - no /proc
self.skipTest("open handles are read from /proc")
with patch.object(duplicate_files, "_unlink", side_effect=unlink):
await self.apply_plan(plan)
self.assertEqual(held, [True, True])
class TestNestedSpellings(SafetyTestCase):
"""Round 2, R2: spellings and favorites by containment, not equality."""
async def test_a_parent_spelling_of_a_nested_source(self):
# The reviewer's probe: sources [output/sub, out-link -> output].
self.need_symlinks()
self.write("sub/a.png", b"same", mtime=1)
self.write("sub/b.png", b"same", mtime=2)
out_link = Path(self.comfy_dir) / "out-link"
os.symlink(str(self.output_dir), str(out_link))
GalleryConfig.MONITORING_DIRECTORIES = [
self._root(self.output_dir / "sub"),
str(out_link),
]
prompt = self.save("Linked before sub was a source")
self.api.db.link_image_to_prompt(prompt, str(out_link / "sub" / "b.png"))
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
self.assertIn("1 database link moves", plan["details"][0]["detail"])
await self.apply_plan(plan)
[image] = self.api.db.get_prompt_images(prompt)
self.assertEqual(
_locator(image["image_path"]), _locator(self.output_dir / "sub" / "a.png")
)
self.assertEqual(self.linked_files(), [])
async def test_nested_sources_subfolder_first_with_a_parent_favorite(self):
self.write("sub/a.png", b"same", mtime=1)
self.write("sub/b.png", b"same", mtime=2)
GalleryConfig.MONITORING_DIRECTORIES = [
self._root(self.output_dir / "sub"),
self._root(self.output_dir),
]
parent_key = duplicate_files.favorite_root_key(self.output_dir)
self.api.db.set_image_favorite(parent_key, "sub/b.png")
prompt = self.save("Linked under the parent")
self.api.db.link_image_to_prompt(prompt, str(self.output_dir / "sub" / "b.png"))
job, group = await self.found_group()
self.assertEqual({f["root_index"] for f in group["files"]}, {0})
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
await self.apply_plan(plan)
self.assertEqual(self.api.db.get_root_favorites(parent_key), {"sub/a.png"})
self.assertEqual(self.linked_files(), [])
class TestLockedRowSaysWhereLinksAre(SafetyTestCase):
"""Round 2, m1."""
async def test_the_locked_row_names_the_kept_copy(self):
self.trio()
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
with (
patch.object(
duplicate_files, "_unlink", side_effect=PermissionError(13, "denied")
),
self.assertLogs("prompt_manager", "WARNING"),
):
cleanup = await self.apply_plan(plan)
self.assertIn(
"its links now name the kept copy", cleanup.outcomes["errors"][0]["detail"]
)
class TestRound3(SafetyTestCase):
"""Re-review 2: n1 (sources at different depths) and n2 (the kept copy
must still be the held file right before the unlink)."""
MAYBE = TestAnotherMount.MAYBE
def setUp(self):
super().setUp()
self.real_lstat, self.real_stat = os.lstat, os.stat
mount = TestAnotherMount.mount
review_with = TestAnotherMount.review_with
async def test_a_deeper_source_on_another_mount_is_protected(self):
# The reviewer's probe: ComfyUI/output -> /mnt/pool/output and
# /mnt/disk1/output/2024 as an extra root, FUSE-style inodes.
_keeper, other, patches = self.mount("x.png", "2024/x.png", ino="other")
_job, _group, plan = await self.review_with(patches, "2024/x.png")
self.assertEqual(plan["protected"]["reasons"], [self.MAYBE])
self.assertTrue(other.is_file())
async def test_a_shallower_source_on_another_mount_is_protected(self):
_keeper, _other, patches = self.mount("2024/x.png", "x.png", ino="other")
_job, _group, plan = await self.review_with(patches, "x.png")
self.assertEqual(plan["protected"]["reasons"], [self.MAYBE])
async def test_only_whole_folder_names_match(self):
_keeper, _other, patches = self.mount("photo-x.png", "x.png", ino="other")
_job, _group, plan = await self.review_with(patches, "x.png")
self.assertEqual(plan["scope"]["count"], 1)
async def removal_while_the_keeper_goes(self, change):
self.write("a.png", b"same", mtime=1)
self.write("b.png", b"same", mtime=2)
prompt = self.save("On b")
self.api.db.link_image_to_prompt(prompt, str(self.output_dir / "b.png"))
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
real = self.api.db.move_copy_records
refused = []
def move(*args, **kwargs):
counts = real(*args, **kwargs)
if args[0]["kept_path"].endswith("a.png"):
try:
change(self.output_dir / "a.png")
except PermissionError:
# Windows: the job holds the kept copy open (no
# FILE_SHARE_DELETE), so it cannot be deleted or replaced.
if os.name != "nt":
raise
refused.append(True)
return counts
with patch.object(self.api.db, "move_copy_records", side_effect=move):
cleanup = await self.apply_plan(plan)
if refused: # pragma: no cover - native Windows
self.assertEqual((self.output_dir / "a.png").read_bytes(), b"same")
self.assertEqual(cleanup.counters["removed"], 1)
self.assertEqual(self.linked_files(), [])
return
self.assertTrue((self.output_dir / "b.png").is_file())
self.assertEqual(cleanup.counters["removed"], 0)
rows = {row["path"]: row["detail"] for row in cleanup.outcomes["kept"]}
self.assertEqual(rows["b.png"], duplicate_files.KEEPER_WENT)
self.assertEqual(self.linked_files(), [])
[image] = self.api.db.get_prompt_images(prompt)
self.assertEqual(os.path.basename(image["image_path"]), "b.png")
self.assertIn("1 not removed", cleanup.status_line)
async def test_a_kept_copy_deleted_during_the_move_keeps_the_file(self):
await self.removal_while_the_keeper_goes(lambda path: path.unlink())
async def test_a_kept_copy_replaced_during_the_move_keeps_the_file(self):
def replace(path):
spare = path.with_name("spare.tmp")
spare.write_bytes(b"other bytes")
os.replace(spare, path)
await self.removal_while_the_keeper_goes(replace)
File diff suppressed because it is too large Load Diff
+6 -2
View File
@@ -1374,13 +1374,17 @@ class TestModelHousekeeping(DatabaseTestCase):
):
self.assertEqual(self.db.model.get_database_info(), {})
def test_vacuum_logs_errors_instead_of_raising(self):
def test_vacuum_logs_and_raises_its_failure(self):
# S6: a failed VACUUM is reported (legacy /maintenance success:false),
# never logged as done.
self._save("x")
self.db.model.vacuum_database()
with patch(
"database.models.sqlite3.connect", side_effect=sqlite3.OperationalError
):
self.db.model.vacuum_database() # logged, not raised
with self.assertLogs("prompt_manager", "ERROR"):
with self.assertRaises(sqlite3.OperationalError):
self.db.model.vacuum_database()
self.assertEqual(len(self.db.search_prompts()), 1)
def test_close_tolerates_a_connection_that_fails_to_close(self):
+2 -2
View File
@@ -63,7 +63,7 @@ class TestAdmissionClasses(unittest.IsolatedAsyncioTestCase):
)
async def test_admission_names_and_explicit_read(self):
self.assertEqual(ADMISSIONS, ("read", "guarded", "records"))
self.assertEqual(ADMISSIONS, ("read", "guarded", "records", "exclusive"))
registry = JobRegistry()
job = registry.start(
"report", "Report", admission="read", scope={}, runner=_noop
@@ -72,7 +72,7 @@ class TestAdmissionClasses(unittest.IsolatedAsyncioTestCase):
self.assertEqual(job.admission, "read")
self.assertFalse(job.record_changing)
with self.assertRaises(ValueError):
registry.start("x", "X", admission="exclusive", scope={}, runner=_noop)
registry.start("x", "X", admission="bogus", scope={}, runner=_noop)
with self.assertRaises(ValueError):
registry.start(
"x",
+546
View File
@@ -0,0 +1,546 @@
"""Maintenance plans (#249 System S6): the plan store and the read-only previews.
``PlanStore`` holds each reviewed plan server-side under a single-use token
for ``ttl`` seconds, within a count and byte budget it never truncates to fit.
``PromptDatabase.preview_maintenance`` builds a plan on its own ``mode=ro``
connection: it writes nothing (the TESTING.md read-only standard) and reuses
the legacy maintenance queries, so a preview counts what the legacy cleanup
would change, except that a missing-file link is only counted under a source
that can be read now.
"""
import hashlib
import os
import sqlite3
import sys
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from database import maintenance # noqa: E402
from database.operations import PromptDatabase # noqa: E402
from py.api.maintenance_plans import ( # noqa: E402
PlanExpired,
PlanStore,
PlanUsed,
ScopeTooLarge,
)
from utils.hashing import generate_prompt_hash # noqa: E402
try:
from tests import test_execution_history_database as history_tests # noqa: E402
from tests.open_handles import assert_closed # noqa: E402
from tests.test_execution_context_database import bundle # noqa: E402
except ImportError: # discovered with tests/ as the top-level directory
import test_execution_history_database as history_tests # noqa: E402
from open_handles import assert_closed # noqa: E402
from test_execution_context_database import bundle # noqa: E402
OPERATIONS = ("duplicate_prompts", "optimize", "unused_components", "missing_links")
def _digest(path):
if not os.path.exists(path):
return None
with open(path, "rb") as handle:
return hashlib.sha256(handle.read()).hexdigest()
class FakeClock:
def __init__(self):
self.now = 1000.0
def __call__(self):
return self.now
class TestPlanStore(unittest.TestCase):
def test_plan_expiry_budget_and_single_use(self):
clock = FakeClock()
store = PlanStore(ttl=900, max_plans=2, max_bytes=200, clock=clock)
first = store.issue("optimize", "sha256:a", {"targets": []}, ("db", 0))
self.assertGreaterEqual(len(first), 40) # token_urlsafe(32)
plan = store.peek(first, "optimize")
self.assertEqual(
(plan.operation, plan.fingerprint, plan.db_identity),
("optimize", "sha256:a", ("db", 0)),
)
# Single use: the second consume names the job the first one started.
self.assertIs(store.consume(first, "optimize", job_id="job-1"), plan)
with self.assertRaises(PlanUsed) as used:
store.consume(first, "optimize", job_id="job-2")
self.assertEqual(used.exception.job_id, "job-1")
with self.assertRaises(PlanUsed):
store.peek(first, "optimize")
# Another operation's token, an unknown token and a non-string are expired.
second = store.issue("missing_links", "sha256:b", {"targets": [1]}, ("db", 0))
for token, operation in (
(second, "optimize"),
("nope", "missing_links"),
(None, "missing_links"),
):
with self.subTest(token=token), self.assertRaises(PlanExpired):
store.peek(token, operation)
# Expiry: valid up to ttl, gone after it.
clock.now += 900
self.assertEqual(store.peek(second, "missing_links").fingerprint, "sha256:b")
clock.now += 0.001
with self.assertRaises(PlanExpired):
store.consume(second, "missing_links", job_id="job-3")
# Budget: a plan over max_bytes is refused whole, never truncated.
with self.assertRaises(ScopeTooLarge):
store.issue("unused_components", "x", {"targets": list(range(100))}, ())
# Count budget: the oldest plan makes room for a new one and is expired.
a = store.issue("optimize", "a", {}, ())
b = store.issue("optimize", "b", {}, ())
c = store.issue("optimize", "c", {}, ())
with self.assertRaises(PlanExpired):
store.peek(a, "optimize")
self.assertEqual(
[store.peek(t, "optimize").fingerprint for t in (b, c)], ["b", "c"]
)
# Byte budget across plans: older plans are dropped until the new one fits.
sized = PlanStore(max_plans=8, max_bytes=100, clock=clock)
pad = {"pad": "x" * 40} # 50 bytes as compact JSON
p1, p2, p3 = (sized.issue("optimize", str(n), pad, ()) for n in range(3))
with self.assertRaises(PlanExpired):
sized.peek(p1, "optimize")
self.assertEqual([sized.peek(t).fingerprint for t in (p2, p3)], ["1", "2"])
def test_expires_at_is_wall_clock(self):
store = PlanStore(ttl=60, clock=FakeClock(), wall_clock=lambda: 5000.0)
token = store.issue("optimize", "f", {}, ())
self.assertEqual(store.peek(token, "optimize").expires_at, 5060.0)
def test_discard_forgets_a_plan(self):
store = PlanStore()
token = store.issue("optimize", "f", {}, ())
store.discard(token)
store.discard("unknown")
with self.assertRaises(PlanExpired):
store.peek(token, "optimize")
def test_used_tokens_are_remembered_within_a_bound(self):
store = PlanStore(max_plans=2)
tokens = []
for index in range(store.MAX_USED + 1):
token = store.issue("optimize", str(index), {}, ())
store.consume(token, "optimize", job_id=f"job-{index}")
tokens.append(token)
with self.assertRaises(PlanExpired): # the oldest record was dropped
store.peek(tokens[0], "optimize")
with self.assertRaises(PlanUsed):
store.peek(tokens[-1], "optimize")
class MaintenanceDatabaseCase(unittest.TestCase):
def setUp(self):
self.temp = tempfile.TemporaryDirectory()
self.root = Path(os.path.realpath(self.temp.name))
self.output = self.root / "output"
self.output.mkdir()
self.db_path = str(self.root / "prompts.db")
self.db = PromptDatabase(self.db_path)
def tearDown(self):
self.db.close_all()
assert_closed(self, self.root)
self.temp.cleanup()
def save(self, text, **metadata):
return self.db.save_prompt(
text, prompt_hash=generate_prompt_hash(text), **metadata
)
def image(self, prompt_id, name, *, exists=True, folder=None):
path = Path(folder or self.output) / name
path.parent.mkdir(parents=True, exist_ok=True)
if exists:
path.write_bytes(b"png")
return self.db.link_image_to_prompt(prompt_id, str(path))
def duplicate(self, text):
"""A second prompt row whose text differs only in case (legacy duplicate)."""
with self.db.model.get_connection() as conn:
cursor = conn.execute(
"INSERT INTO prompts(text, created_at) VALUES(?, '2999-01-01')",
(text,),
)
return cursor.lastrowid
def preview(self, operation, roots=None):
roots = [str(self.output)] if roots is None else roots
return self.db.preview_maintenance(operation, roots)
def ids(self, sql, *args):
with self.db.model.get_connection() as conn:
return [row[0] for row in conn.execute(sql, args)]
class TestPreviews(MaintenanceDatabaseCase):
def populate(self):
kept = self.save("Kept prompt")
self.image(kept, "kept.png")
self.image(kept, "gone.png", exists=False)
self.save("Unused prompt")
self.duplicate("KEPT PROMPT")
return kept
def test_preview_writes_nothing(self):
self.populate()
# The app's own connection stays open, as in the running server.
self.db.get_statistics()
files = [self.db_path, self.db_path + "-wal"]
real_connect = sqlite3.connect
class Recording(sqlite3.Connection):
changes_at_close = []
def close(self):
Recording.changes_at_close.append(self.total_changes)
super().close()
opened = []
def connect(*args, **kwargs):
conn = real_connect(*args, factory=Recording, **kwargs)
opened.append((args, kwargs, conn))
return conn
for operation in OPERATIONS:
with self.subTest(operation=operation):
before = [_digest(path) for path in files]
shm_size = os.path.getsize(self.db_path + "-shm")
media = {
name: _digest(self.output / name)
for name in os.listdir(self.output)
}
prompts = self.ids("SELECT id FROM prompts ORDER BY id")
images = self.ids("SELECT id FROM generated_images ORDER BY id")
opened.clear()
Recording.changes_at_close.clear()
with patch.object(maintenance.sqlite3, "connect", side_effect=connect):
draft = self.preview(operation)
self.assertTrue(draft.fingerprint.startswith("sha256:"))
self.assertEqual(len(opened), 1)
args, kwargs, conn = opened[0]
self.assertTrue(str(args[0]).endswith("?mode=ro"), args)
self.assertTrue(kwargs.get("uri"))
self.assertEqual(Recording.changes_at_close, [0])
with self.assertRaises(sqlite3.ProgrammingError): # closed
conn.execute("SELECT 1")
self.assertEqual([_digest(path) for path in files], before)
self.assertEqual(os.path.getsize(self.db_path + "-shm"), shm_size)
self.assertEqual(
{
name: _digest(self.output / name)
for name in os.listdir(self.output)
},
media,
)
self.assertEqual(
self.ids("SELECT id FROM prompts ORDER BY id"), prompts
)
self.assertEqual(
self.ids("SELECT id FROM generated_images ORDER BY id"), images
)
def test_unknown_operation_raises(self):
with self.assertRaises(ValueError):
self.db.preview_maintenance("reboot", [])
def test_same_library_gives_the_same_fingerprint(self):
self.populate()
for operation in OPERATIONS:
with self.subTest(operation=operation):
self.assertEqual(
self.preview(operation).fingerprint,
self.preview(operation).fingerprint,
)
def test_preview_details_are_capped_and_carry_no_server_path(self):
prompt = self.save("Many")
for index in range(60):
self.image(prompt, f"sub/gone-{index}.png", exists=False)
draft = self.preview("missing_links")
self.assertEqual(draft.scope["count"], 60)
self.assertEqual(len(draft.details), maintenance.MAX_DETAILS)
self.assertEqual(draft.details[0]["label"], "sub/gone-0.png")
self.assertEqual(draft.details[0]["detail"], "output · the file is missing")
self.assertNotIn(str(self.root), str(draft.public()))
self.assertEqual(len(draft.targets), 60)
def test_duplicate_prompts_plan_counts_the_legacy_groups(self):
kept = self.save("A cat")
self.duplicate("A CAT")
self.duplicate("a cat")
paired = self.db.get_prompt_pair(self.db.save_prompt_pair("a dog", "blur"))
self.duplicate("A DOG")
draft = self.preview("duplicate_prompts")
self.assertEqual(draft.scope["count"], 2)
self.assertEqual(draft.scope["label"], "2 duplicate prompts in 1 group")
self.assertEqual(draft.protected["count"], 2)
self.assertEqual(
draft.protected["reasons"], ["1 group belongs to a saved pair"]
)
self.assertEqual(draft.targets[0][0], kept)
self.assertEqual(len(draft.targets), 1)
self.assertEqual(
[row["detail"] for row in draft.details],
["Keeps the oldest of 3 copies", "Kept: belongs to a saved pair"],
)
self.assertEqual(draft.details[0]["label"], "A cat")
self.assertIsNotNone(paired)
def test_unused_components_plan_protects_pairs_history_and_titled_parents(self):
orphan = self.save("Orphan")
pair = self.db.get_prompt_pair(self.db.save_prompt_pair("paired", "neg"))
titled = self.save("Titled parent")
untitled = self.save("Explicit parent")
protected = self.save("Protected", tags=["__protected__"])
with self.db.model.get_connection() as conn:
conn.execute(
"INSERT INTO prompt_library_state(prompt_id,title) "
"VALUES(?,?),(?,NULL)",
(titled, "My recipe", untitled),
)
hinted = self.save("Recorded run")
draft_before_hint = self.preview("unused_components")
self.assertEqual(sorted(draft_before_hint.targets), sorted([orphan, hinted]))
self.record_run_hint(hinted)
draft = self.preview("unused_components")
self.assertEqual(draft.targets, [orphan])
self.assertEqual(draft.scope["count"], 1)
self.assertEqual(draft.scope["label"], "1 unused prompt component")
self.assertEqual(draft.protected["count"], 6)
self.assertEqual(
draft.protected["reasons"],
[
"2 in saved pairs",
"1 tagged __protected__",
"2 with Library state (a title or an explicit parent)",
"1 used by a recorded run",
],
)
self.assertEqual(
draft.details,
[
{
"label": "Orphan",
"detail": maintenance.UNUSED_DETAIL,
}
],
)
# The plan counts exactly what the legacy prune removes.
self.assertIsNotNone(pair)
self.assertEqual(protected > 0, True)
self.assertEqual(self.db.prune_orphaned_prompts(), 1)
self.assertEqual(self.preview("unused_components").targets, [])
def record_run_hint(self, prompt_id, run=0):
"""Make ``prompt_id`` the live positive hint of a recorded run."""
text = self.db.get_prompt_by_id(prompt_id)["text"]
history = history_tests.HistoryDatabaseTests
start = history.start(self, self.db, run)
self.db.write_execution_start(start)
self.db.write_execution_finish(history.finish(self, start))
self.assertTrue(
self.db.write_execution_contexts(
start.run_id, start.storage_epoch, bundle(positive=text, negative="")
)
)
def test_missing_links_plan_skips_unavailable_roots(self):
prompt = self.save("Links")
self.image(prompt, "here.png")
gone = self.image(prompt, "gone.png", exists=False)
unmounted = self.root / "unmounted"
self.image(prompt, "away.png", folder=unmounted)
os.remove(unmounted / "away.png")
os.rmdir(unmounted) # the drive went away: its links look missing
self.image(
prompt, "elsewhere.png", folder=self.root / "elsewhere", exists=False
)
draft = self.preview("missing_links", [str(self.output), str(unmounted)])
self.assertEqual(draft.targets, [gone])
self.assertEqual(draft.scope["count"], 1)
self.assertEqual(draft.scope["label"], "1 link to a missing file")
self.assertEqual(draft.protected["count"], 2)
self.assertEqual(
draft.protected["reasons"],
[
"1 link under an unavailable source",
"1 link outside the image sources",
],
)
self.assertFalse(draft.media["deletes_files"])
# No readable source: nothing is counted, every link is protected.
none = self.preview("missing_links", [str(unmounted)])
self.assertEqual((none.targets, none.scope["count"]), ([], 0))
self.assertEqual(none.protected["count"], 4)
def test_missing_links_plan_matches_the_health_check(self):
prompt = self.save("Links")
self.image(prompt, "a.png", exists=False)
self.image(prompt, "b.png", exists=False)
self.image(prompt, "c.png")
report = self.db.run_health_checks([str(self.output)])
missing = next(c for c in report["checks"] if c["id"] == "missing_files")
self.assertEqual(self.preview("missing_links").scope["count"], missing["count"])
def test_optimize_plan_reports_sizes_without_paths(self):
self.populate()
draft = self.preview("optimize")
self.assertEqual(draft.targets, [])
self.assertEqual(draft.scope, {"count": 1, "label": "The whole database"})
self.assertEqual(
[row["label"] for row in draft.details],
["Database file", "Free pages", "Lock", "Free space needed"],
)
self.assertNotIn(str(self.root), str(draft.public()))
def test_legacy_cleanup_missing_images_skips_unavailable_roots(self):
prompt = self.save("Links")
gone = self.image(prompt, "gone.png", exists=False)
unmounted = self.root / "unmounted"
away = self.image(prompt, "away.png", folder=unmounted, exists=False)
outside = self.image(prompt, "x.png", folder=self.root / "other", exists=False)
unmounted.rmdir()
self.assertEqual(
self.db.cleanup_missing_images(roots=[str(self.output), str(unmounted)]), 1
)
remaining = self.ids("SELECT id FROM generated_images ORDER BY id")
self.assertEqual(remaining, sorted([away, outside]))
self.assertNotIn(gone, remaining)
# Without roots it keeps the legacy rule: every link whose file is gone.
self.assertEqual(self.db.cleanup_missing_images(), 2)
def test_reader_closes_its_connection_when_a_read_raises(self):
self.populate()
with patch.object(
maintenance, "_orphaned_ids", side_effect=sqlite3.OperationalError("boom")
):
with self.assertRaises(sqlite3.OperationalError):
self.preview("unused_components")
# The teardown's assert_closed proves the read-only handle was closed.
class TestApply(MaintenanceDatabaseCase):
"""``apply_maintenance`` rechecks the fingerprint inside BEGIN IMMEDIATE."""
def apply(self, draft, roots=None):
roots = [str(self.output)] if roots is None else roots
return self.db.apply_maintenance(
draft.operation, draft.fingerprint, draft.targets, roots
)
def test_unused_components_apply_removes_exactly_the_reviewed_prompts(self):
first, second = self.save("One"), self.save("Two")
draft = self.preview("unused_components")
with patch.object(maintenance, "ID_BATCH", 1): # one id per statement
result = self.apply(draft)
self.assertEqual(result["changed"], 2)
self.assertEqual(
[row["prompt_id"] for row in result["rows"]], sorted([first, second])
)
self.assertIsNone(self.db.get_prompt_by_id(first))
self.assertIsNone(self.db.get_prompt_by_id(second))
def test_changed_targets_raise_before_anything_is_written(self):
orphan, protected = self.save("Orphan"), self.save("Soon protected")
prompt = self.save("Links")
gone = self.image(prompt, "gone.png", exists=False)
kept = self.save("Kept")
self.duplicate("KEPT")
cases = {
"unused_components": lambda: self.db.model.get_connection().execute(
"INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)",
(protected, "Now a recipe"),
),
"missing_links": lambda: (self.output / "gone.png").write_bytes(b"back"),
"duplicate_prompts": lambda: self.db.update_prompt_metadata(kept, rating=4),
}
for operation, change in cases.items():
with self.subTest(operation=operation):
draft = self.preview(operation)
self.assertTrue(draft.targets)
change()
self.db.model.get_connection().commit()
with self.assertRaises(maintenance.PlanChanged):
self.apply(draft)
self.assertIsNotNone(self.db.get_prompt_by_id(orphan))
self.assertIsNotNone(self.db.get_image_by_id(gone))
self.assertEqual(len(self.db.search_prompts(text="kept")), 2)
def test_a_link_whose_source_went_away_is_never_removed(self):
folder = self.root / "drive"
prompt = self.save("Links")
link = self.image(prompt, "gone.png", folder=folder, exists=False)
roots = [str(folder)]
draft = self.preview("missing_links", roots)
self.assertEqual(draft.targets, [link])
folder.rmdir() # unmounted between the review and the apply
with self.assertRaises(maintenance.PlanChanged):
self.apply(draft, roots)
self.assertIsNotNone(self.db.get_image_by_id(link))
def test_missing_links_apply_removes_rows_only(self):
prompt = self.save("Links")
kept = self.image(prompt, "kept.png")
gone = [self.image(prompt, f"gone-{n}.png", exists=False) for n in range(3)]
media = {p.name: _digest(p) for p in self.output.iterdir()}
draft = self.preview("missing_links")
with patch.object(maintenance, "ID_BATCH", 2):
result = self.apply(draft)
self.assertEqual(result["changed"], 3)
self.assertEqual(
[os.path.basename(row["image_path"]) for row in result["rows"]],
[f"gone-{n}.png" for n in range(3)],
)
self.assertEqual({p.name: _digest(p) for p in self.output.iterdir()}, media)
self.assertIsNotNone(self.db.get_image_by_id(kept))
for image_id in gone:
self.assertIsNone(self.db.get_image_by_id(image_id))
def test_duplicate_prompts_apply_merges_into_the_oldest(self):
kept = self.save("A cat")
copy = self.duplicate("A CAT")
image = self.image(copy, "cat.png")
draft = self.preview("duplicate_prompts")
result = self.apply(draft)
self.assertEqual(result["changed"], 1)
self.assertEqual(result["rows"], [{"prompt_id": copy, "kept_id": kept}])
self.assertIsNone(self.db.get_prompt_by_id(copy))
self.assertEqual(self.db.get_image_by_id(image)["prompt_id"], kept)
def test_optimize_apply_vacuums(self):
self.save("Some")
draft = self.preview("optimize")
with patch.object(self.db.model, "vacuum_database") as vacuum:
self.assertEqual(self.apply(draft), {"changed": 1, "rows": []})
vacuum.assert_called_once_with()
self.assertEqual(self.apply(self.preview("optimize"))["changed"], 1)
with self.assertRaises(maintenance.PlanChanged):
self.db.apply_maintenance("optimize", "sha256:other", [], [])
def test_apply_of_an_unknown_operation_raises(self):
with self.assertRaises(ValueError):
self.db.apply_maintenance("reboot", "x", [], [])
class TestVacuum(MaintenanceDatabaseCase):
def test_vacuum_failure_raises(self):
with patch(
"database.models.sqlite3.connect",
side_effect=sqlite3.OperationalError("disk I/O error"),
):
with self.assertLogs("prompt_manager", "ERROR"):
with self.assertRaises(sqlite3.OperationalError):
self.db.model.vacuum_database()
if __name__ == "__main__":
unittest.main()