From f68e92b3b6b526b34fe1d99d228833579c4d2d7a Mon Sep 17 00:00:00 2001 From: Vito Date: Sat, 24 Oct 2026 05:28:44 -0700 Subject: [PATCH] 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 .backup_[_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 `.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. --- TESTING.md | 204 ++++++ database/duplicate_copies.py | 212 ++++++ database/health.py | 38 +- database/maintenance.py | 724 +++++++++++++++++++ database/models.py | 5 + database/operations.py | 167 ++--- py/api/__init__.py | 27 +- py/api/admin.py | 374 ++++++---- py/api/duplicate_files.py | 779 ++++++++++++++++++++ py/api/duplicate_review.py | 451 ++++++++++++ py/api/duplicate_spellings.py | 133 ++++ py/api/jobs.py | 21 +- py/api/maintenance.py | 773 ++++++++++++++++++++ py/api/maintenance_plans.py | 138 ++++ py/api/settings_routes.py | 6 + py/config.py | 11 + tests/test_api_backup_restore.py | 142 +++- tests/test_api_duplicate_files.py | 986 +++++++++++++++++++++++++ tests/test_api_duplicate_safety.py | 558 +++++++++++++++ tests/test_api_maintenance.py | 1067 ++++++++++++++++++++++++++++ tests/test_database.py | 8 +- tests/test_job_admission.py | 4 +- tests/test_maintenance_plans.py | 546 ++++++++++++++ 23 files changed, 7110 insertions(+), 264 deletions(-) create mode 100644 database/duplicate_copies.py create mode 100644 database/maintenance.py create mode 100644 py/api/duplicate_files.py create mode 100644 py/api/duplicate_review.py create mode 100644 py/api/duplicate_spellings.py create mode 100644 py/api/maintenance.py create mode 100644 py/api/maintenance_plans.py create mode 100644 tests/test_api_duplicate_files.py create mode 100644 tests/test_api_duplicate_safety.py create mode 100644 tests/test_api_maintenance.py create mode 100644 tests/test_maintenance_plans.py diff --git a/TESTING.md b/TESTING.md index 52667c4..8267f6a 100644 --- a/TESTING.md +++ b/TESTING.md @@ -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 `.backup_`, 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 `.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 diff --git a/database/duplicate_copies.py b/database/duplicate_copies.py new file mode 100644 index 0000000..fb33d3c --- /dev/null +++ b/database/duplicate_copies.py @@ -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 diff --git a/database/health.py b/database/health.py index 7feedda..e17a202 100644 --- a/database/health.py +++ b/database/health.py @@ -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) diff --git a/database/maintenance.py b/database/maintenance.py new file mode 100644 index 0000000..7b0c19c --- /dev/null +++ b/database/maintenance.py @@ -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]], + } diff --git a/database/models.py b/database/models.py index 3724465..07c2e4f 100644 --- a/database/models.py +++ b/database/models.py @@ -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: """ diff --git a/database/operations.py b/database/operations.py index cf9cec1..54a4a79 100644 --- a/database/operations.py +++ b/database/operations.py @@ -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. diff --git a/py/api/__init__.py b/py/api/__init__.py index f929b2d..bc4ac04 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -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) diff --git a/py/api/admin.py b/py/api/admin.py index 4b2e9c7..7326e40 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -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 # '.backup_*' files kept next to the live database +# '.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: diff --git a/py/api/duplicate_files.py b/py/api/duplicate_files.py new file mode 100644 index 0000000..6edf265 --- /dev/null +++ b/py/api/duplicate_files.py @@ -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 diff --git a/py/api/duplicate_review.py b/py/api/duplicate_review.py new file mode 100644 index 0000000..79f64ba --- /dev/null +++ b/py/api/duplicate_review.py @@ -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, + } diff --git a/py/api/duplicate_spellings.py b/py/api/duplicate_spellings.py new file mode 100644 index 0000000..31e0f0b --- /dev/null +++ b/py/api/duplicate_spellings.py @@ -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, + } diff --git a/py/api/jobs.py b/py/api/jobs.py index aec92eb..eace0b9 100644 --- a/py/api/jobs.py +++ b/py/api/jobs.py @@ -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" diff --git a/py/api/maintenance.py b/py/api/maintenance.py new file mode 100644 index 0000000..031b3d9 --- /dev/null +++ b/py/api/maintenance.py @@ -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 + ``.backup_[_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 ``.backup_[_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_", + "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 ``.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") diff --git a/py/api/maintenance_plans.py b/py/api/maintenance_plans.py new file mode 100644 index 0000000..aa257d9 --- /dev/null +++ b/py/api/maintenance_plans.py @@ -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] diff --git a/py/api/settings_routes.py b/py/api/settings_routes.py index 4760e08..17637f7 100644 --- a/py/api/settings_routes.py +++ b/py/api/settings_routes.py @@ -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() diff --git a/py/config.py b/py/config.py index c6fa719..18f94f1 100644 --- a/py/config.py +++ b/py/config.py @@ -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"]) diff --git a/tests/test_api_backup_restore.py b/tests/test_api_backup_restore.py index aa211ff..2f6acc7 100644 --- a/tests/test_api_backup_restore.py +++ b/tests/test_api_backup_restore.py @@ -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) diff --git a/tests/test_api_duplicate_files.py b/tests/test_api_duplicate_files.py new file mode 100644 index 0000000..ee33e9b --- /dev/null +++ b/tests/test_api_duplicate_files.py @@ -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()) diff --git a/tests/test_api_duplicate_safety.py b/tests/test_api_duplicate_safety.py new file mode 100644 index 0000000..0211514 --- /dev/null +++ b/tests/test_api_duplicate_safety.py @@ -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) diff --git a/tests/test_api_maintenance.py b/tests/test_api_maintenance.py new file mode 100644 index 0000000..2066d2a --- /dev/null +++ b/tests/test_api_maintenance.py @@ -0,0 +1,1067 @@ +"""Maintenance routes (py/api/maintenance.py, #249 System S6). + +``POST /prompt_manager/maintenance/preview`` answers a reviewed ``Plan`` for +one of four operations and writes nothing; its token binds a later apply to +exactly that plan. Errors are coded and never carry a server path or an +exception text. +""" + +import asyncio +import datetime +import hashlib +import json +import re +import os +import sqlite3 +import sys +from pathlib import Path +from unittest.mock import MagicMock, patch + +_mock_server = MagicMock() +_mock_server.PromptServer.instance.routes = MagicMock() +sys.modules.setdefault("server", _mock_server) + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from database.operations import PromptDatabase # noqa: E402 +from py.api import PromptManagerAPI # noqa: E402 +from py.api.maintenance_plans import PlanExpired # noqa: E402 +from py.config import GalleryConfig, PromptManagerConfig # noqa: E402 +from utils.hashing import generate_prompt_hash # noqa: E402 + +try: + from tests.test_api_admin import AdminAPITestCase # noqa: E402 +except ImportError: # discovered with tests/ as the top-level directory + from test_api_admin import AdminAPITestCase # noqa: E402 + +WAIT = 10 +PREVIEW = "/prompt_manager/maintenance/preview" +OPERATIONS = ("duplicate_prompts", "optimize", "unused_components", "missing_links") +PLAN_KEYS = { + "success", + "token", + "operation", + "expires_at", + "scope", + "protected", + "backup", + "media", + "details", + "disk", +} +DISK_KEYS = {"free_bytes", "required_bytes", "sufficient"} + + +def _digest(path): + if not os.path.exists(path): + return None + with open(path, "rb") as handle: + return hashlib.sha256(handle.read()).hexdigest() + + +class MaintenanceAPITestCase(AdminAPITestCase): + async def get_application(self): + app = await super().get_application() + self.api._find_comfyui_output_dir = lambda: ( + str(self.output_dir) if self.output_dir.is_dir() else None + ) + return app + + def save(self, text, **metadata): + return self.api.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_dir) / name + path.parent.mkdir(parents=True, exist_ok=True) + if exists: + path.write_bytes(b"png") + return self.api.db.link_image_to_prompt(prompt_id, str(path)) + + def database_bytes(self): + conn = sqlite3.connect(self.api.db.model.db_path) + try: + size = conn.execute("PRAGMA page_size").fetchone()[0] + return size * conn.execute("PRAGMA page_count").fetchone()[0] + finally: + conn.close() + + async def post(self, url, body, status): + response = await self.client.request("POST", url, json=body) + text = await response.text() + self.assertEqual(response.status, status, text) + self.assertNotIn(self.tmpdir, text) + return json.loads(text) + + async def preview(self, operation, status=200, **body): + return await self.post(PREVIEW, {"operation": operation, **body}, status) + + +class TestPreviewRoute(MaintenanceAPITestCase): + async def test_preview_answers_a_plan_for_each_operation(self): + prompt = self.save("Kept") + self.image(prompt, "gone.png", exists=False) + self.save("Unused") + for operation in OPERATIONS: + with self.subTest(operation=operation): + plan = await self.preview(operation, options={}) + self.assertEqual(set(plan), PLAN_KEYS) + self.assertTrue(plan["success"]) + self.assertEqual(plan["operation"], operation) + self.assertIsInstance(plan["token"], str) + self.assertEqual(set(plan["scope"]), {"count", "label"}) + self.assertEqual(set(plan["protected"]), {"count", "reasons"}) + expected = plan["backup"].pop("expected_bytes") + self.assertEqual( + plan["backup"], + { + "policy": "on", + "destination_label": "prompts.db.backup_", + }, + ) + self.assertEqual(expected, self.database_bytes()) + self.assertEqual(set(plan["disk"]), DISK_KEYS) + self.assertFalse(plan["media"]["deletes_files"]) + self.assertLessEqual(len(plan["details"]), 50) + self.assertTrue(plan["expires_at"].endswith("+00:00")) + self.assertNotIn("fingerprint", json.dumps(plan)) + missing = await self.preview("missing_links") + self.assertEqual( + missing["scope"], {"count": 1, "label": "1 link to a missing file"} + ) + unused = await self.preview("unused_components") + self.assertEqual(unused["scope"]["count"], 1) + + async def test_preview_writes_nothing(self): + prompt = self.save("Kept") + self.image(prompt, "kept.png") + self.image(prompt, "gone.png", exists=False) + self.save("Unused") + db_path = self.api.db.model.db_path + self.api.db.get_statistics() + files = [db_path, db_path + "-wal"] + for operation in OPERATIONS: + with self.subTest(operation=operation): + before = [_digest(path) for path in files] + shm = os.path.getsize(db_path + "-shm") + media = {p.name: _digest(p) for p in self.output_dir.iterdir()} + await self.preview(operation) + self.assertEqual([_digest(path) for path in files], before) + self.assertEqual(os.path.getsize(db_path + "-shm"), shm) + self.assertEqual( + {p.name: _digest(p) for p in self.output_dir.iterdir()}, media + ) + self.assertEqual( + [p["text"] for p in self.api.db.search_prompts()], ["Unused", "Kept"] + ) + + async def test_unknown_operation_is_400(self): + cases = [ + {"operation": "reboot"}, + {"operation": "duplicate_files"}, # S6 Task 17, not yet + {"operation": 3}, + {}, + {"operation": "optimize", "options": {"all": True}}, + {"operation": "optimize", "options": []}, + {"operation": "optimize", "extra": 1}, + ] + for body in cases: + with self.subTest(body=body): + reply = await self.post(PREVIEW, body, 400) + self.assertEqual(reply["success"], False) + self.assertEqual(reply["code"], "invalid-request") + response = await self.client.request( + "POST", PREVIEW, data=b"[1]", headers={"Content-Type": "application/json"} + ) + self.assertEqual(response.status, 400) + self.assertEqual((await response.json())["code"], "invalid-request") + big = json.dumps({"operation": "optimize", "pad": "x" * 8192}) + response = await self.client.request( + "POST", PREVIEW, data=big, headers={"Content-Type": "application/json"} + ) + self.assertEqual(response.status, 413) + self.assertEqual((await response.json())["code"], "too-large") + + async def test_preview_failure_is_coded_without_exception_text(self): + with patch.object( + self.api.db, + "preview_maintenance", + side_effect=sqlite3.OperationalError(f"boom {self.tmpdir}"), + ): + with self.assertLogs("prompt_manager", "ERROR"): + reply = await self.preview("optimize", status=500) + self.assertEqual( + reply, + {"success": False, "code": None, "error": "Failed to review maintenance"}, + ) + + async def test_scope_too_large_is_413(self): + from py.api.maintenance_plans import PlanStore + + self.api._maintenance_plans = PlanStore(max_bytes=10) + self.save("Unused") + reply = await self.preview("unused_components", status=413) + self.assertEqual(reply["code"], "too-large") + + async def test_preview_during_a_restore_swap_is_storage_changed(self): + real = self.api.db.preview_maintenance + + def swapped(*args, **kwargs): + draft = real(*args, **kwargs) + self.api._db_generation += 1 # a restore replaced the database + return draft + + with patch.object(self.api.db, "preview_maintenance", side_effect=swapped): + reply = await self.preview("optimize", status=409) + self.assertEqual(reply["code"], "storage-changed") + + async def test_missing_links_preview_uses_the_health_roots(self): + unmounted = Path(self.tmpdir) / "unmounted" + unmounted.mkdir() + GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir), str(unmounted)] + prompt = self.save("Links") + self.image(prompt, "gone.png", exists=False) + self.image(prompt, "away.png", folder=unmounted, exists=False) + unmounted.rmdir() + plan = await self.preview("missing_links") + self.assertEqual(plan["scope"]["count"], 1) + self.assertEqual( + plan["protected"], + {"count": 1, "reasons": ["1 link under an unavailable source"]}, + ) + + async def test_restart_loses_plans(self): + plan = await self.preview("optimize") + self.assertEqual( + self.api._maintenance_plans.peek(plan["token"], "optimize").operation, + "optimize", + ) + restarted = PromptManagerAPI(db=PromptDatabase(self.api.db.model.db_path)) + try: + with self.assertRaises(PlanExpired): + restarted._maintenance_plans.peek(plan["token"], "optimize") + finally: + restarted.db.close_all() + + +APPLY = "/prompt_manager/maintenance/apply" +BACKUP_NAME = re.compile(r"^prompts\.db\.backup_\d{8}_\d{6}_\d{6}(_\d+)?$") + + +class ApplyTestCase(MaintenanceAPITestCase): + def backups(self): + return sorted( + name + for name in os.listdir(self.tmpdir) + if name.startswith("prompts.db.backup_") + ) + + def strays(self): + """Files beside the database that are neither it nor a safety backup.""" + known = {"prompts.db", "prompts.db-wal", "prompts.db-shm", "config.json"} + return sorted( + name + for name in os.listdir(self.tmpdir) + if os.path.isfile(os.path.join(self.tmpdir, name)) + and name not in known + and not BACKUP_NAME.match(name) + ) + + async def apply(self, token, status=202): + return await self.post(APPLY, {"token": token}, status) + + async def run_apply(self, operation): + plan = await self.preview(operation) + body = await self.apply(plan["token"]) + job = self.api._jobs.get(body["job"]["id"]) + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + return plan, body["job"], job + + +class TestApplyRoute(ApplyTestCase): + async def test_apply_runs_as_a_records_job_with_a_verified_backup(self): + orphan = self.save("Orphan") + plan, summary, job = await self.run_apply("unused_components") + self.assertEqual(summary["kind"], "maintenance") + self.assertEqual(summary["title"], "Unused prompt components") + self.assertEqual(summary["admission"], "records") + self.assertFalse(summary["cancellable"]) + self.assertEqual( + summary["counters"], {"changed": 0, "protected": 0, "errors": 0} + ) + self.assertEqual(summary["scope"], {"operation": "unused_components"}) + self.assertEqual(job.state, "finished", job.log) + self.assertEqual(job.counters, {"changed": 1, "protected": 0, "errors": 0}) + self.assertIsNone(self.api.db.get_prompt_by_id(orphan)) + [backup] = self.backups() + self.assertRegex(backup, BACKUP_NAME) + ok, _reason = self.api.db.model.verify_database_file( + os.path.join(self.tmpdir, backup) + ) + self.assertTrue(ok) + self.assertEqual(self.strays(), []) + self.assertIn(f"Safety backup {backup} verified", "\n".join(job.log)) + self.assertEqual( + job.status_line, + f"Removed 1 unused prompt component · Safety backup {backup}", + ) + self.assertEqual(job.outcomes["changed"][0]["detail"], "Removed: Orphan") + report = json.dumps(job.report()) + self.assertNotIn(self.tmpdir, report) + if os.name == "nt": # pragma: no cover - native path spelling + self.assertNotIn(self.tmpdir.replace("\\", "/"), report) + + async def test_apply_rejects_changed_target(self): + orphan, soon = self.save("Orphan"), self.save("Soon a recipe") + plan = await self.preview("unused_components") + with self.api.db.model.get_connection() as conn: + conn.execute( + "INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)", + (soon, "Recipe"), + ) + body = await self.apply(plan["token"]) + job = self.api._jobs.get(body["job"]["id"]) + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + self.assertEqual(job.state, "failed") + self.assertEqual(job.last_event["code"], "plan-changed") + self.assertEqual( + job.status_line, + "The library changed since the review; nothing was changed. " + "Review the plan again.", + ) + self.assertEqual(job.counters["errors"], 1) + self.assertIsNotNone(self.api.db.get_prompt_by_id(orphan)) + self.assertIsNotNone(self.api.db.get_prompt_by_id(soon)) + + async def test_token_is_single_use(self): + self.save("Orphan") + plan, summary, job = await self.run_apply("unused_components") + again = await self.apply(plan["token"], 409) + self.assertEqual(again["code"], "plan-used") + self.assertEqual( + again["job"], {"id": summary["id"], "title": "Unused prompt components"} + ) + kinds = [j.kind for j in self.api._jobs.list()] + self.assertEqual(kinds.count("maintenance"), 1) + + async def test_apply_refused_while_restore_runs(self): + orphan = self.save("Orphan") + plan = await self.preview("unused_components") + async with self.api._jobs.hold( + "restore", "Database restore (admin page)", admission="records" + ) as restore: + busy = await self.apply(plan["token"], 409) + self.assertEqual(busy["code"], "busy") + self.assertEqual( + busy["running"], + {"id": restore.id, "title": "Database restore (admin page)"}, + ) + self.assertIsNotNone(self.api.db.get_prompt_by_id(orphan)) + # The refused apply kept the token: it applies once the restore ended. + body = await self.apply(plan["token"]) + job = self.api._jobs.get(body["job"]["id"]) + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + self.assertIsNone(self.api.db.get_prompt_by_id(orphan)) + + async def test_backup_failure_aborts_before_write(self): + orphan = self.save("Orphan") + model = self.api.db.model + failures = { + "backup": patch.object(model, "backup_database", return_value=False), + "verify": patch.object( + model, "verify_database_file", return_value=(False, "broken") + ), + } + for name, failure in failures.items(): + with self.subTest(failure=name): + with failure, patch.object(self.api.db, "apply_maintenance") as applied: + _plan, _summary, job = await self.run_apply("unused_components") + applied.assert_not_called() + self.assertEqual(job.state, "failed") + self.assertEqual(job.last_event["code"], "backup-failed") + self.assertEqual( + job.status_line, + "The safety backup could not be taken and verified; " + "nothing was changed.", + ) + self.assertIsNotNone(self.api.db.get_prompt_by_id(orphan)) + self.assertEqual(self.backups(), []) + self.assertEqual(self.strays(), []) + + async def test_policy_off_reports_no_backup(self): + + PromptManagerConfig.SAFETY_BACKUP = False + orphan = self.save("Orphan") + plan, _summary, job = await self.run_apply("unused_components") + self.assertEqual(plan["backup"]["policy"], "off") + self.assertEqual(job.state, "finished") + self.assertIn("No safety backup (the policy is off)", "\n".join(job.log)) + self.assertEqual( + job.status_line, "Removed 1 unused prompt component · No safety backup" + ) + self.assertEqual(self.backups(), []) + self.assertIsNone(self.api.db.get_prompt_by_id(orphan)) + + async def test_a_plan_reviewed_with_the_backup_on_keeps_that_promise(self): + + self.save("Orphan") + plan = await self.preview("unused_components") + PromptManagerConfig.SAFETY_BACKUP = False # turned off after the review + body = await self.apply(plan["token"]) + job = self.api._jobs.get(body["job"]["id"]) + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + self.assertEqual(len(self.backups()), 1) + + async def test_backup_verified_before_write(self): + self.save("Orphan") + model, db = self.api.db.model, self.api.db + order = [] + real_backup, real_verify = model.backup_database, model.verify_database_file + real_apply = db.apply_maintenance + + def backup(path): + order.append("backup") + return real_backup(path) + + def verify(path, **kwargs): + order.append("verify") + return real_verify(path, **kwargs) + + def apply(*args): + order.append("apply") + self.assertEqual(len(self.backups()), 1) # in place before the write + return real_apply(*args) + + with ( + patch.object(model, "backup_database", side_effect=backup), + patch.object(model, "verify_database_file", side_effect=verify), + patch.object(db, "apply_maintenance", side_effect=apply), + ): + _plan, _summary, job = await self.run_apply("unused_components") + self.assertEqual(job.state, "finished") + self.assertEqual(order, ["backup", "verify", "apply"]) + + async def test_missing_links_never_deletes_files(self): + prompt = self.save("Links") + kept = self.image(prompt, "kept.png") + gone = self.image(prompt, "sub/gone.png", exists=False) + other = Path(self.tmpdir) / "elsewhere" + self.image(prompt, "x.png", folder=other) + media = { + str(path): _digest(path) + for path in Path(self.tmpdir).rglob("*") + if path.is_file() and "prompts.db" not in path.name + } + _plan, _summary, job = await self.run_apply("missing_links") + self.assertEqual(job.state, "finished", job.log) + self.assertEqual(job.counters["changed"], 1) + self.assertEqual( + job.outcomes["changed"][0], + { + "path": "sub/gone.png", + "root": "output", + "root_index": 0, + "detail": "Link removed; the file was already missing", + }, + ) + after = { + str(path): _digest(path) + for path in Path(self.tmpdir).rglob("*") + if path.is_file() and "prompts.db" not in path.name and str(path) in media + } + self.assertEqual(after, media) + self.assertIsNone(self.api.db.get_image_by_id(gone)) + self.assertIsNotNone(self.api.db.get_image_by_id(kept)) + + async def test_duplicate_prompts_and_optimize_apply(self): + kept = self.save("A cat") + with self.api.db.model.get_connection() as conn: + copy = conn.execute( + "INSERT INTO prompts(text, created_at) VALUES('A CAT', '2999-01-01')" + ).lastrowid + _plan, _summary, job = await self.run_apply("duplicate_prompts") + self.assertEqual(job.state, "finished", job.log) + self.assertEqual(job.status_line.split(" · ")[0], "Merged 1 duplicate prompt") + self.assertEqual( + job.outcomes["changed"][0]["detail"], + f"Prompt {copy} merged into prompt {kept}", + ) + self.assertIsNone(self.api.db.get_prompt_by_id(copy)) + _plan, _summary, job = await self.run_apply("optimize") + self.assertEqual(job.state, "finished", job.log) + self.assertEqual(job.status_line.split(" · ")[0], "Optimized the database") + + async def test_protected_items_are_counted(self): + self.save("Orphan") + self.save("Kept", tags=["__protected__"]) + _plan, _summary, job = await self.run_apply("unused_components") + self.assertEqual(job.counters, {"changed": 1, "protected": 1, "errors": 0}) + + async def test_apply_failure_changes_nothing_and_is_coded(self): + orphan = self.save("Orphan") + with patch.object( + self.api.db, + "apply_maintenance", + side_effect=sqlite3.OperationalError(f"locked {self.tmpdir}"), + ): + with self.assertLogs("prompt_manager", "ERROR"): + _plan, _summary, job = await self.run_apply("unused_components") + self.assertEqual(job.state, "failed") + self.assertEqual(job.last_event["code"], None) + self.assertEqual(job.status_line, "Maintenance failed; nothing was changed.") + self.assertNotIn(self.tmpdir, json.dumps(job.report())) + self.assertIsNotNone(self.api.db.get_prompt_by_id(orphan)) + + async def test_apply_errors_are_coded(self): + missing = await self.post(APPLY, {}, 400) + self.assertEqual(missing["code"], "invalid-request") + wrong = await self.post(APPLY, {"token": 7}, 400) + self.assertEqual(wrong["code"], "invalid-request") + extra = await self.post(APPLY, {"token": "x", "force": True}, 400) + self.assertEqual(extra["code"], "invalid-request") + expired = await self.apply("not-a-token", 410) + self.assertEqual(expired["code"], "plan-expired") + + async def test_a_restore_after_the_review_is_storage_changed(self): + orphan = self.save("Orphan") + plan = await self.preview("unused_components") + self.api._database_replaced() # what the legacy restore calls + reply = await self.apply(plan["token"], 409) + self.assertEqual(reply["code"], "storage-changed") + again = await self.apply(plan["token"], 410) # the plan is forgotten + self.assertEqual(again["code"], "plan-expired") + self.assertIsNotNone(self.api.db.get_prompt_by_id(orphan)) + + async def test_restart_loses_plans_for_apply(self): + plan = await self.preview("optimize") + self.api._maintenance_plans = type(self.api._maintenance_plans)() + reply = await self.apply(plan["token"], 410) + self.assertEqual(reply["code"], "plan-expired") + + async def test_safety_backups_are_pruned(self): + from py.api.admin import MAX_SAFETY_BACKUPS + + for index in range(MAX_SAFETY_BACKUPS): + old = os.path.join( + self.tmpdir, f"prompts.db.backup_20000101_0000{index:02d}" + ) + with open(old, "wb") as handle: + handle.write(b"old") + os.utime(old, (1000 + index, 1000 + index)) + self.save("Orphan") + await self.run_apply("unused_components") + backups = self.backups() + self.assertEqual(len(backups), MAX_SAFETY_BACKUPS) + self.assertNotIn("prompts.db.backup_20000101_000000", backups) + + +class TestSafetyBackup(ApplyTestCase): + async def test_backup_names_unique_within_one_second(self): + from py.api import maintenance + + frozen = datetime.datetime(2026, 10, 23, 12, 0, 0, 123456) + model = self.api.db.model + with patch.object(maintenance, "_now", return_value=frozen): + names = [maintenance.take_safety_backup(model) for _ in range(3)] + self.assertEqual( + names, + [ + "prompts.db.backup_20261023_120000_123456", + "prompts.db.backup_20261023_120000_123456_1", + "prompts.db.backup_20261023_120000_123456_2", + ], + ) + for name in names: + ok, _reason = model.verify_database_file(os.path.join(self.tmpdir, name)) + self.assertTrue(ok, name) + self.assertEqual(self.strays(), []) + + async def test_backup_is_private_on_posix(self): + from py.api import maintenance + + name = maintenance.take_safety_backup(self.api.db.model) + if os.name != "nt": + mode = os.stat(os.path.join(self.tmpdir, name)).st_mode & 0o777 + self.assertEqual(mode, 0o600) + + async def test_a_failed_replace_leaves_no_file(self): + from py.api import maintenance + + # A filesystem without hard links falls back to reserve-and-replace. + with ( + patch.object(maintenance.os, "link", side_effect=OSError("no hard links")), + patch.object(maintenance.os, "replace", side_effect=OSError("busy")), + ): + with self.assertRaises(maintenance.SafetyBackupFailed): + maintenance.take_safety_backup(self.api.db.model) + self.assertEqual(self.backups(), []) + self.assertEqual(self.strays(), []) + with patch.object(maintenance.os, "link", side_effect=FileExistsError): + with self.assertRaises(maintenance.SafetyBackupFailed): + maintenance.take_safety_backup(self.api.db.model) + with patch.object(maintenance.os, "link", side_effect=OSError("no hard links")): + name = maintenance.take_safety_backup(self.api.db.model) + self.assertEqual(self.backups(), [name]) + self.assertEqual(self.strays(), []) + + async def test_a_backup_that_cannot_be_finished_leaves_no_file(self): + from py.api import maintenance + + real_connect = sqlite3.connect + calls = [] + + def connect(*args, **kwargs): + calls.append(args[0]) + if len(calls) == 2: # the backup is written; finishing it fails + raise sqlite3.OperationalError("x") + return real_connect(*args, **kwargs) + + with patch.object(maintenance.sqlite3, "connect", side_effect=connect): + with self.assertRaises(maintenance.SafetyBackupFailed) as failed: + maintenance.take_safety_backup(self.api.db.model) + self.assertEqual(str(failed.exception), "The backup could not be finished") + self.assertEqual(calls[0], calls[1]) # the temp copy, both times + self.assertEqual((self.backups(), self.strays()), ([], [])) + + async def test_a_safety_backup_restores_through_the_legacy_route(self): + from aiohttp import FormData + + from py.api import maintenance + + self.save("Before the backup") + name = maintenance.take_safety_backup(self.api.db.model) + self.save("After the backup") + with open(os.path.join(self.tmpdir, name), "rb") as handle: + upload = handle.read() + form = FormData() + form.add_field("database_file", upload, filename=name) + response = await self.client.request( + "POST", "/prompt_manager/restore", data=form + ) + self.assertEqual(response.status, 200, await response.text()) + texts = [p["text"] for p in self.api.db.search_prompts()] + self.assertEqual(texts, ["Before the backup"]) + mode = self.api.db.model.get_connection().execute("PRAGMA journal_mode") + self.assertEqual(mode.fetchone()[0], "wal") + + +class TestBackupPoolAndDiskSpace(ApplyTestCase): + """Coordinator rulings: one backup pool; Apply refuses a short disk up front.""" + + def old_backup(self, name, mtime): + path = os.path.join(self.tmpdir, name) + with open(path, "wb") as handle: + handle.write(b"old") + os.utime(path, (mtime, mtime)) + return name + + async def test_safety_backups_share_the_restore_pool(self): + from py.api.admin import MAX_SAFETY_BACKUPS, _prune_safety_backups + + restore_style = [ + self.old_backup(f"prompts.db.backup_2000010{i}_000000", 1000 + i) + for i in range(6) + ] + maintenance_style = [ + self.old_backup(f"prompts.db.backup_2001010{i}_000000_000000", 2000 + i) + for i in range(5) + ] + self.save("Orphan") + await self.run_apply("unused_components") + backups = self.backups() + # One pool of ten: the new backup plus the nine newest of the eleven. + self.assertEqual(len(backups), MAX_SAFETY_BACKUPS) + self.assertNotIn(restore_style[0], backups) + self.assertNotIn(restore_style[1], backups) + self.assertTrue(set(maintenance_style) <= set(backups)) + # A restore prunes the same pool, maintenance backups included. + self.old_backup("prompts.db.backup_20020101_000000", 3000) + _prune_safety_backups(self.api.db.model.db_path) + self.assertEqual(len(self.backups()), MAX_SAFETY_BACKUPS) + self.assertNotIn(restore_style[2], self.backups()) + + async def test_the_pool_is_pruned_when_the_apply_fails_after_the_backup(self): + from py.api.admin import MAX_SAFETY_BACKUPS + + for index in range(MAX_SAFETY_BACKUPS): + self.old_backup(f"prompts.db.backup_2000010{index}_00000{index}", 1000) + orphan = self.save("Orphan") + plan = await self.preview("unused_components") + self.save("Changed after the review") + body = await self.apply(plan["token"]) + job = self.api._jobs.get(body["job"]["id"]) + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + self.assertEqual(job.last_event["code"], "plan-changed") + self.assertIsNotNone(self.api.db.get_prompt_by_id(orphan)) + self.assertEqual(len(self.backups()), MAX_SAFETY_BACKUPS) + + async def test_preview_reports_the_backup_size_and_free_space(self): + from py.api import maintenance + + self.save("Orphan") + # SQLite's temporary folder on the database's disk: the needs add up. + with ( + patch.object(maintenance, "_free_bytes", return_value=10 * 2**30), + patch.object(maintenance, "sqlite_temp_dir", return_value=self.tmpdir), + ): + plan = await self.preview("unused_components") + optimize = await self.preview("optimize") + size = self.database_bytes() + self.assertEqual(plan["backup"]["expected_bytes"], size) + self.assertEqual( + plan["disk"], + { + "free_bytes": 10 * 2**30, + "required_bytes": size + maintenance.DISK_MARGIN_BYTES, + "sufficient": True, + }, + ) + # VACUUM rebuilds the database beside it, and the backup comes first. + self.assertEqual( + optimize["disk"]["required_bytes"], + 3 * size + maintenance.DISK_MARGIN_BYTES, + ) + + async def test_policy_off_reports_no_backup_size(self): + from py.api import maintenance + + PromptManagerConfig.SAFETY_BACKUP = False + self.save("Orphan") + with patch.object(maintenance, "_free_bytes", return_value=0): + plan = await self.preview("unused_components") + reply = await self.apply(plan["token"]) + self.assertIsNone(plan["backup"]["expected_bytes"]) + self.assertEqual( + plan["disk"], {"free_bytes": 0, "required_bytes": 0, "sufficient": True} + ) + job = self.api._jobs.get(reply["job"]["id"]) + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + self.assertEqual(job.state, "finished") + + async def test_apply_refuses_up_front_when_the_backup_does_not_fit(self): + from py.api import maintenance + + orphan = self.save("Orphan") + plan = await self.preview("unused_components") + with patch.object(maintenance, "_free_bytes", return_value=1024): + reply = await self.apply(plan["token"], 507) + self.assertEqual(reply["code"], "insufficient-space") + needed = self.database_bytes() + maintenance.DISK_MARGIN_BYTES + size_mb = f"{needed / 2**20:,.1f} MB" + self.assertEqual( + reply["error"], + f"Not enough free disk space: the safety backup needs about {size_mb} " + "in the database folder and 0.0 MB is free. Nothing was changed.", + ) + self.assertEqual(self.api._jobs.list(), []) # no job started + self.assertEqual(self.backups(), []) + self.assertIsNotNone(self.api.db.get_prompt_by_id(orphan)) + # The review is still usable once there is room. + body = await self.apply(plan["token"]) + job = self.api._jobs.get(body["job"]["id"]) + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + self.assertIsNone(self.api.db.get_prompt_by_id(orphan)) + + async def test_apply_refuses_up_front_when_vacuum_does_not_fit(self): + from py.api import maintenance + + PromptManagerConfig.SAFETY_BACKUP = False + size = self.database_bytes() + plan = await self.preview("optimize") + free = size # VACUUM needs about twice the database + with patch.object(maintenance, "_free_bytes", return_value=free): + reply = await self.apply(plan["token"], 507) + self.assertEqual(reply["code"], "insufficient-space") + self.assertIn("Optimize needs about", reply["error"]) + self.assertTrue(reply["error"].endswith("Nothing was changed.")) + self.assertEqual(self.api._jobs.list(), []) + + +class TestSafetyBackupPolicy(MaintenanceAPITestCase): + async def test_policy_round_trips_with_revision(self): + settings = await self.client.request("GET", "/prompt_manager/settings") + body = await settings.json() + self.assertTrue(body["settings"]["safety_backup"]) + saved = await self.post( + "/prompt_manager/settings", + {"safety_backup": False, "expected_revision": body["revision"]}, + 200, + ) + self.assertTrue(saved["persisted"]) + self.assertNotEqual(saved["revision"], body["revision"]) + with open(self.config_path, encoding="utf-8") as handle: + self.assertIs(json.load(handle)["maintenance"]["safety_backup"], False) + again = await ( + await self.client.request("GET", "/prompt_manager/settings") + ).json() + self.assertFalse(again["settings"]["safety_backup"]) + self.assertEqual(again["revision"], saved["revision"]) + plan = await self.preview("optimize") + self.assertEqual(plan["backup"]["policy"], "off") + stale = await self.post( + "/prompt_manager/settings", + {"safety_backup": True, "expected_revision": body["revision"]}, + 409, + ) + self.assertEqual(stale["code"], "settings-changed") + bad = await self.post("/prompt_manager/settings", {"safety_backup": "yes"}, 400) + self.assertFalse(bad["success"]) + + async def test_config_file_value_is_read_at_load(self): + + PromptManagerConfig.SAFETY_BACKUP = False + self.assertFalse(PromptManagerConfig.SAFETY_BACKUP) + PromptManagerConfig.update_config({"maintenance": {"safety_backup": "no"}}) + self.assertFalse(PromptManagerConfig.SAFETY_BACKUP) # non-bool ignored + self.assertEqual( + PromptManagerConfig.get_config()["maintenance"], {"safety_backup": False} + ) + + +class TestLegacyMaintenance(MaintenanceAPITestCase): + async def test_vacuum_failure_reported_not_success(self): + with patch( + "database.models.sqlite3.connect", + side_effect=sqlite3.OperationalError("database is locked"), + ): + with self.assertLogs("prompt_manager", "ERROR"): + reply = await self.post( + "/prompt_manager/maintenance", {"operations": ["vacuum"]}, 200 + ) + self.assertFalse(reply["all_successful"]) + self.assertFalse(reply["results"]["vacuum"]["success"]) + self.assertEqual(reply["results"]["vacuum"]["error"], "database is locked") + + async def test_legacy_orphaned_image_cleanup_skips_unavailable_sources(self): + unmounted = Path(self.tmpdir) / "unmounted" + unmounted.mkdir() + GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir), str(unmounted)] + prompt = self.save("Links") + self.image(prompt, "gone.png", exists=False) + away = self.image(prompt, "away.png", folder=unmounted, exists=False) + unmounted.rmdir() + reply = await self.post( + "/prompt_manager/maintenance", + {"operations": ["cleanup_orphaned_images"]}, + 200, + ) + self.assertEqual( + reply["results"]["cleanup_orphaned_images"], + { + "success": True, + "removed_count": 1, + "message": "Removed 1 orphaned image records", + }, + ) + self.assertIsNotNone(self.api.db.get_image_by_id(away)) + + +if __name__ == "__main__": # pragma: no cover + import unittest + + unittest.main() + + +class TestFinalReviewFixes(ApplyTestCase): + """PR #313 final review: I3, M1-M5.""" + + def duplicate_of(self, kept_text, copy_text): + kept = self.save(kept_text) + with self.api.db.model.get_connection() as conn: + copy = conn.execute( + "INSERT INTO prompts(text, created_at) VALUES(?, '2999-01-01')", + (copy_text,), + ).lastrowid + return kept, copy + + def test_sqlite_temp_dir_follows_sqlite_search_order(self): + from py.api.maintenance import sqlite_temp_dir + + def usable(*good): + return lambda path: path in good + + posix = {"windows": False} + self.assertEqual( + sqlite_temp_dir({}, usable=usable("/var/tmp", "/tmp"), **posix), "/var/tmp" + ) + self.assertEqual( + sqlite_temp_dir({}, usable=usable("/usr/tmp", "/tmp"), **posix), "/usr/tmp" + ) + env = {"SQLITE_TMPDIR": "/fast", "TMPDIR": "/t"} + self.assertEqual( + sqlite_temp_dir(env, usable=usable("/fast", "/t", "/tmp"), **posix), + "/fast", + ) + self.assertEqual( + sqlite_temp_dir(env, usable=usable("/t", "/tmp"), **posix), "/t" + ) + self.assertEqual( + sqlite_temp_dir({}, usable=usable(), **posix), os.path.abspath(".") + ) + windows = {"windows": True, "usable": usable()} + env = {"TMP": "C:\\a", "TEMP": "C:\\b", "USERPROFILE": "C:\\u"} + self.assertEqual(sqlite_temp_dir(env, **windows), "C:\\a") + self.assertEqual(sqlite_temp_dir({"TEMP": "C:\\b"}, **windows), "C:\\b") + self.assertEqual(sqlite_temp_dir({"USERPROFILE": "C:\\u"}, **windows), "C:\\u") + + async def test_optimize_measures_sqlites_temp_folder(self): + from py.api import maintenance + + measured = [] + + def free(path): + measured.append(path) + return 10 * 2**30 + + other = os.path.join(self.tmpdir, "sqlite-temp") + os.mkdir(other) + with ( + patch.object(maintenance, "sqlite_temp_dir", return_value=other), + patch.object(maintenance, "_free_bytes", side_effect=free), + patch.object( + maintenance, "_device", side_effect=lambda path: path # another disk + ), + ): + plan = await self.preview("optimize") + self.assertIn(other, measured) + size = self.database_bytes() + self.assertEqual( + plan["disk"]["required_bytes"], 2 * size + maintenance.DISK_MARGIN_BYTES + ) + + async def test_a_plan_used_during_the_disk_check_is_plan_used(self): + from py.api import maintenance + + self.save("Orphan") + plan = await self.preview("unused_components") + real = maintenance.disk_plan + + def racing(*args, **kwargs): + # A parallel apply consumed the token while this one measured. + self.api._maintenance_plans.consume( + plan["token"], "unused_components", job_id="f" * 32 + ) + return real(*args, **kwargs) + + with patch.object(maintenance, "disk_plan", side_effect=racing): + reply = await self.apply(plan["token"], 409) + self.assertEqual(reply["code"], "plan-used") + self.assertEqual(self.api._jobs.list(), []) + + async def test_the_backup_decision_is_frozen_when_the_job_starts(self): + PromptManagerConfig.SAFETY_BACKUP = False + self.save("Orphan") + plan = await self.preview("unused_components") + real = self.api._jobs.start + + def start(*args, **kwargs): + job = real(*args, **kwargs) + PromptManagerConfig.SAFETY_BACKUP = True # switched on meanwhile + return job + + with patch.object(self.api._jobs, "start", side_effect=start): + body = await self.apply(plan["token"]) + job = self.api._jobs.get(body["job"]["id"]) + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + self.assertEqual(self.backups(), []) + self.assertTrue(job.status_line.endswith("No safety backup")) + + async def test_startup_sweeps_interrupted_backup_copies(self): + stray = os.path.join(self.tmpdir, ".pm-safety-crashed.tmp") + with open(stray, "wb") as handle: + handle.write(b"half a backup") + PromptManagerAPI(db=self.api.db) + self.assertFalse(os.path.exists(stray)) + + async def test_a_backup_is_linked_into_place_without_a_reserved_name(self): + from py.api import maintenance + + linked = [] + real = os.link + + def link(source, target): + linked.append(target) + return real(source, target) + + with ( + patch.object(maintenance.os, "link", side_effect=link), + patch.object( + maintenance, "_reserve_backup_name", side_effect=AssertionError + ), + ): + name = maintenance.take_safety_backup(self.api.db.model) + self.assertEqual([os.path.basename(p) for p in linked], [name]) + self.assertEqual(self.strays(), []) + + async def test_optimize_runs_alone_and_says_the_database_is_locked(self): + from py.api.jobs import JobBusy + + plan = await self.preview("optimize") + self.assertIn( + "Runs alone: the database is locked while it runs", + " ".join(row["detail"] for row in plan["details"]), + ) + gate = asyncio.Event() + + async def held(job): + await gate.wait() + + read = self.api._jobs.start( + "health", "Health check", admission="read", runner=held + ) + busy = await self.apply(plan["token"], 409) + self.assertEqual(busy["running"]["id"], read.id) + gate.set() + await asyncio.wait_for(asyncio.shield(read.task), WAIT) + gate.clear() + body = await self.apply(plan["token"]) + self.assertEqual(body["job"]["admission"], "exclusive") + with self.assertRaises(JobBusy): + self.api._jobs.start( + "health", "Health check", admission="read", runner=held + ) + job = self.api._jobs.get(body["job"]["id"]) + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + + async def test_a_link_collision_is_protected_not_swallowed(self): + kept, copy = self.duplicate_of("A cat", "A CAT") + path = os.path.join(str(self.output_dir), "both.png") + self.api.db.link_image_to_prompt(kept, path) + self.api.db.link_image_to_prompt(copy, path) + plan, _summary, job = await self.run_apply("duplicate_prompts") + self.assertEqual(plan["scope"]["count"], 0) + self.assertEqual( + plan["protected"]["reasons"], ["1 group links one image twice"] + ) + self.assertIsNotNone(self.api.db.get_prompt_by_id(copy)) + self.assertEqual(len(self.api.db.get_prompt_images(copy)), 1) + + async def test_a_duplicate_with_a_library_title_is_protected(self): + kept, copy = self.duplicate_of("A dog", "A DOG") + with self.api.db.model.get_connection() as conn: + conn.execute( + "INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)", + (copy, "My dog recipe"), + ) + plan, _summary, _job = await self.run_apply("duplicate_prompts") + self.assertEqual(plan["scope"]["count"], 0) + self.assertEqual( + plan["protected"]["reasons"], ["1 group has its own Library title"] + ) + self.assertIsNotNone(self.api.db.get_prompt_by_id(copy)) + + def test_the_strict_merge_raises_instead_of_dropping_a_link(self): + from database.maintenance import MergeConflict + + kept, copy = self.duplicate_of("A bird", "A BIRD") + path = os.path.join(str(self.output_dir), "bird.png") + self.api.db.link_image_to_prompt(kept, path) + self.api.db.link_image_to_prompt(copy, path) + with self.assertRaises(MergeConflict): + with self.api.db._pair_transaction() as conn: + self.api.db._merge_duplicate_group(conn, [kept, copy], strict=True) + self.assertEqual(len(self.api.db.get_prompt_images(copy)), 1) + # The legacy cleanup keeps its old behaviour (it logs and goes on). + with self.assertLogs("prompt_manager.database", "ERROR"): + self.assertEqual(self.api.db.cleanup_duplicates(), 1) diff --git a/tests/test_database.py b/tests/test_database.py index bcc8498..92935ce 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -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): diff --git a/tests/test_job_admission.py b/tests/test_job_admission.py index 59e4821..0bfb11c 100644 --- a/tests/test_job_admission.py +++ b/tests/test_job_admission.py @@ -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", diff --git a/tests/test_maintenance_plans.py b/tests/test_maintenance_plans.py new file mode 100644 index 0000000..a8cbbfb --- /dev/null +++ b/tests/test_maintenance_plans.py @@ -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()