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()