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

Duplicate image files move onto the job registry:

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

Coordinator rulings on Tasks 15-16:

- Safety backups share the restore pool of 10 `<db>.backup_*` files; the
  pool is now also pruned when an apply fails after its backup.
- Plans report the safety backup's size and the free disk space, and
  apply refuses up front (507 insufficient-space) when the backup or
  VACUUM does not fit, so no job fails halfway for space.
This commit is contained in:
Vito Sansevero
2026-10-23 18:16:58 -07:00
parent 7df509e6fb
commit 237a97d7cb
10 changed files with 2628 additions and 182 deletions
+63 -1
View File
@@ -2518,7 +2518,8 @@ 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. A preview meets the
strict, `options` must be `{}` and the body is capped at 4 KiB (512 KiB for
`duplicate_files`, 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
@@ -2565,6 +2566,67 @@ replaced the database since the review; the plan is dropped). The job:
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".
Safety backups and restore backups are one pool: every `<db>.backup_*` file beside the
database counts toward the same `MAX_SAFETY_BACKUPS` (10), whichever wrote it, and the
pool is pruned after every backup a job took, also when the apply then failed
(`test_safety_backups_share_the_restore_pool`,
`test_the_pool_is_pruned_when_the_apply_fails_after_the_backup`). The legacy
`GET /backup` download writes no file beside the database, so it adds nothing to the
pool. Each plan reports `backup.expected_bytes` (the backup's size, `page_count *
page_size`, or null with the policy off) and `disk: {free_bytes, required_bytes,
sufficient}` for the database folder: the backup needs its size there, and Optimize
needs the database's size twice (the WAL beside it and SQLite's temporary copy, added
up when the temporary folder is on the same disk), each plus a 16 MiB margin. 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
server holds the groups (content hash, root-relative paths, source index) 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 512 KiB, at most 5,000 keep entries) 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, the
same size and the same sha256. 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 and the file is removed (every hash
handle closed first: `test_handles_closed_before_unlink`), with its thumbnail (a
video's `.jpg` too). 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 never the OS text, and the
files after it are listed as not attempted (`test_partial_failure_reports_each_path`).
Only the files actually removed have their links moved: each image record moves to the
kept copy, or, when its prompt already links the kept copy, hands its saved-pair links
to that record and is deleted; a favorite moves with it
(`test_db_links_reconciled_only_for_removed_files`). 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
+105
View File
@@ -41,6 +41,7 @@ from pathlib import Path
from .execution_contexts import LIVE_HINT_CTES
from .health import BATCH_ROWS, link_states, readable_roots, readonly_uri
from .models import normalize_image_path
from .pairs import PairDependencyError
MAINTENANCE_OPERATIONS = (
@@ -627,3 +628,107 @@ class MaintenanceOperations:
"changed": changed,
"rows": [{"image_path": row[1]} for row in rows[:MAX_RESULT_ROWS]],
}
# ── duplicate image files (py/api/duplicate_files.py) ────────────
def count_file_links(self, file_paths) -> dict:
"""``{path: n}``: the image records linking each file (one snapshot).
A record links a file by its normalized ``file_path``, or, in rows
that predate that column, by ``image_path``.
"""
def read(conn):
return {
path: sum(
1
for _row in _file_link_rows(conn, normalize_image_path(path), path)
)
for path in file_paths
}
return _read_snapshot(self.model.db_path, read)
def move_file_links(self, moves, favorites=()) -> dict:
"""Point the links of removed duplicate files at their kept copy.
``moves`` is ``[(removed path, kept path)]`` for files that were
removed; nothing else is touched. Each image record of a removed file
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.
``favorites`` is ``[(root key, removed key, kept root key, kept key)]``:
a removed file's favorite moves to the kept copy. One ``BEGIN
IMMEDIATE`` transaction. Returns ``{moved, merged, favorites}``.
"""
counts = {"moved": 0, "merged": 0, "favorites": 0}
with self._pair_transaction() as conn:
for removed, kept in moves:
kept_key = normalize_image_path(kept)
rows = list(
_file_link_rows(conn, normalize_image_path(removed), removed)
)
for image_id, prompt_id in rows:
existing = conn.execute(KEEPER_LINK_SQL, (prompt_id, kept_key))
existing = existing.fetchone()
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,
(kept, os.path.basename(kept), kept_key, image_id),
)
counts["moved"] += 1
for root_key, removed_key, kept_root, kept_key in favorites:
row = conn.execute(FAVORITE_SQL, (root_key, removed_key)).fetchone()
if row is None:
continue
conn.execute(INSERT_FAVORITE_SQL, (kept_root, kept_key, row[0]))
conn.execute(DELETE_FAVORITE_SQL, (root_key, removed_key))
counts["favorites"] += 1
return counts
# Link reconciliation of removed duplicate files: every statement is an
# index SEARCH (idx_image_file_path, idx_image_path, the (prompt_id,
# file_path) unique index, idx_image_prompt_pairs_image, the favorites key).
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 = ?"
)
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,
)
def _file_link_rows(conn, file_key, raw_path):
"""``(id, prompt_id)`` of the image records linking one file."""
yield from (tuple(row) for row in conn.execute(LINKS_BY_FILE_SQL, (file_key,)))
yield from (tuple(row) for row in conn.execute(LEGACY_LINKS_SQL, (raw_path,)))
+7 -2
View File
@@ -14,6 +14,7 @@ Split into domain-specific mixins for maintainability:
- 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
@@ -48,6 +49,7 @@ 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
@@ -289,6 +291,7 @@ class PromptManagerAPI(
ThumbnailJobsMixin,
SystemRoutesMixin,
MaintenanceRoutesMixin,
DuplicateFilesMixin,
LoggingRoutesMixin,
AutotagRoutesMixin,
LoraIntegrationMixin,
@@ -345,8 +348,9 @@ 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).
@@ -516,6 +520,7 @@ class PromptManagerAPI(
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)
+166 -126
View File
@@ -24,9 +24,9 @@ from .jobs import (
SCAN_GROUPS,
SCAN_SCOPES,
DUPLICATE_CLEANUP_JOB_KIND,
DUPLICATES_JOB_KIND,
MAINTENANCE_JOB_KIND,
RESTORE_JOB_KIND,
Job,
JobBusy,
)
@@ -69,7 +69,9 @@ except ImportError:
IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp", ".tiff")
VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv", ".webm", ".m4v", ".wmv")
SCAN_BATCH_SIZE = 50 # Files per executor call during the output scan
MAX_SAFETY_BACKUPS = 10 # '<db>.backup_*' files kept next to the live database
# '<db>.backup_*' files kept next to the live database: one pool for the
# restore backups and the maintenance safety backups (py/api/maintenance.py).
MAX_SAFETY_BACKUPS = 10
MAX_SCAN_FILES = 50_000 # Upper bound on media files one scan will look at
SCAN_COUNT_KEYS = (
"processed",
@@ -90,6 +92,7 @@ MAX_SCAN_DEPTH = 12 # Directory levels below an output root that are scanned
# the common single-f TIFF spelling. Anything else inside the output
# directory (databases, scripts, notes) is never a "duplicate image".
DELETABLE_MEDIA_EXTENSIONS = frozenset(IMAGE_EXTENSIONS + VIDEO_EXTENSIONS + (".tif",))
DUPLICATE_FIND_LEGACY_TITLE = "Find duplicate image files (gallery page)"
def _collect_media_files(output_dirs, extensions):
@@ -137,6 +140,36 @@ def _thumbnail_rel_path(rel_path, thumbnail_ext):
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 +594,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 +630,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 +661,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 +697,95 @@ 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 file is listed once, under the first
root that holds it, and 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}]}]`` for every hash shared by two or more files, ``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):
key = os.path.normcase(str(path))
if key not in seen:
seen.add(key)
media_files.append((slot, root, path))
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 = item
info = None
if stop is None or not stop():
try:
stat = media_path.stat()
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,
"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 +799,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 +835,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 +849,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 +862,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 +949,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):
+664
View File
@@ -0,0 +1,664 @@
"""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,
KEPT_COPY,
LINKS_FAILED,
LOCKED_FILE,
NOT_ATTEMPTED,
OUTSIDE_SOURCE,
ROW_DETAILS,
SOURCES_CHANGED,
SOURCES_CHANGED_REVIEW,
HeldFind,
SourcesChanged,
_fingerprint,
_group_order,
_identity,
_parse_options,
_plural,
_public_plan,
_select,
file_problem,
human_bytes,
)
from .images import _favorite_relative_key, _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.maintenance import FILE_LINK_SQL
except ImportError:
from database.maintenance 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."""
listed = self._listed_sources()
targets, kept, protected = [], [], {}
hasher = self._calculate_file_hash
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 file_problem(
k["abs"], k["root"], group["hash"], k["size"], hasher
)
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 = file_problem(
record["abs"],
record["root"],
group["hash"],
record["size"],
hasher,
)
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))
links = db.count_file_links([t["abs"] for t in targets]) if targets else {}
return {
"targets": targets,
"kept": kept,
"fingerprint": _fingerprint(targets),
"public": _public_plan(targets, protected, links),
}
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"],
}
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):
"""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),
)
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):
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, plan)
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")
removed, failure, keepers = [], None, {}
for done, target in enumerate(targets, start=1):
if failure is not None:
self._record_target(job, "kept", target, NOT_ATTEMPTED)
continue
outcome, code = await self._run_in_executor(
self._remove_one_duplicate, target, keepers
)
if outcome == "removed":
removed.append(target)
detail = f"A copy of {target['keeper_rel']}"
self._record_target(job, "removed", target, detail, "-")
elif outcome == "kept":
self._record_target(job, "kept", target, ROW_DETAILS[code])
else:
failure = code
self._record_target(job, "errors", target, code, "!")
job.update(processed=done, progress=int(done / len(targets) * 100))
links_ok = await self._move_links(job, db, removed)
if removed:
self.invalidate_gallery_cache()
line = f"Removed {_plural(len(removed), 'duplicate file')}"
job.log_line(line)
if failure is not None or not links_ok:
line = "Stopped after an error · " + line if failure else line
job.update(progress=100, status_line=f"{line} · {backup_note(backup)}")
job.publish({"type": "error", "code": "delete-failed", "message": line})
job.finish("failed")
return
job.update(progress=100, status_line=f"{line} · {backup_note(backup)}")
@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, target, keepers):
"""``(outcome, detail or code)`` for one reviewed removal (blocking).
``keepers`` caches each group's kept-copy check for this job.
"""
hasher = self._calculate_file_hash
group_id = target["group_id"]
if group_id not in keepers:
keepers[group_id] = (
file_problem(
target["keeper_abs"],
target["keeper_root"],
target["hash"],
target["keeper_size"],
hasher,
)
is None
)
if not keepers[group_id]:
return "kept", KEEPER
problem = file_problem(
target["abs"], target["root"], target["hash"], target["size"], hasher
)
if problem:
return "kept", problem
try:
_unlink(target["abs"])
except FileNotFoundError:
return "kept", GONE
except PermissionError:
self.logger.warning(f"Duplicate file is locked: {target['rel']}")
return "error", LOCKED_FILE
except OSError:
self.logger.warning(
f"Could not delete duplicate file {target['rel']}", exc_info=True
)
return "error", DELETE_FAILED
self.logger.info(f"Deleted duplicate file: {target['rel']}")
_remove_thumbnail(Path(target["root"]), Path(target["rel"]), self.logger)
return "removed", None
async def _move_links(self, job, db, removed):
"""Move the removed files' links to their kept copy; False on failure."""
if not removed:
return True
moves = [(t["abs"], t["keeper_abs"]) for t in removed]
favorites = [
(
_favorite_root_key(t["root"]),
_favorite_relative_key(t["rel"]),
_favorite_root_key(t["keeper_root"]),
_favorite_relative_key(t["keeper_rel"]),
)
for t in removed
]
try:
counts = await self._run_in_executor(db.move_file_links, moves, favorites)
except Exception:
self.logger.exception("Moving duplicate file links failed")
job.record(
"errors", {"path": "", "detail": LINKS_FAILED}, f"! {LINKS_FAILED}"
)
return False
links = counts["moved"] + counts["merged"]
job.log_line(f"Moved {_plural(links, 'database link')} to the kept copies")
if counts["favorites"]:
job.log_line(
f"Moved {_plural(counts['favorites'], 'favorite')} to the kept copies"
)
return True
+285
View File
@@ -0,0 +1,285 @@
"""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
MAX_KEEP = 5000 # keep entries one review takes
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"
PROTECTED_REASONS = {
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."
ROW_DETAILS = {
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,
}
LOCKED_FILE = "The file is open in another program or read-only; it was not deleted."
DELETE_FAILED = "The file could not be deleted."
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."
)
LINKS_FAILED = (
"The files were removed, but their database links could not be moved. "
"Review Missing-file links to tidy them."
)
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."""
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):
"""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, of ``size`` bytes and with content hash ``digest``. 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 info.st_size != size:
return CHANGED_CONTENT
if (hasher or _hash_file)(path) != digest:
return CHANGED_CONTENT
except FileNotFoundError:
return GONE
except OSError:
return UNREADABLE
return None
@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 _public_plan(targets, protected, links):
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 = []
for target in targets[:MAX_DETAILS]:
n = links.get(target["abs"], 0)
moves = (
f"{_plural(n, 'database link')} "
f"{'moves' if n == 1 else 'move'} to the kept copy"
if n
else "no database links"
)
details.append(
{
"label": target["rel"],
"detail": f"{target['root_label']} · a copy of {target['keeper_rel']}"
f" · {moves}",
}
)
label = f"{_plural(count, 'duplicate file')} in {_plural(groups, 'group')}"
return {
"operation": DUPLICATE_FILES,
"scope": {"count": count, "label": label},
"protected": {"count": sum(protected.values()), "reasons": reasons},
"media": {
"deletes_files": True,
"label": f"Deletes {_plural(count, 'file')} ({human_bytes(size)}) from "
"disk; the kept copies stay. The safety backup covers the database, "
"not these files.",
},
"details": details,
}
+5
View File
@@ -91,6 +91,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"
+226 -48
View File
@@ -43,7 +43,9 @@ 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
@@ -56,11 +58,19 @@ 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 = 512 * 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",
@@ -68,7 +78,7 @@ MAINTENANCE_TITLES = {
"unused_components": "Unused prompt components",
"missing_links": "Missing-file links",
}
_OPERATION_ERROR = "operation must be one of: " + ", ".join(MAINTENANCE_OPERATIONS)
_OPERATION_ERROR = "operation must be one of: " + ", ".join(PREVIEW_OPERATIONS)
MAINTENANCE_GROUPS = ("changed", "protected", "errors")
MAX_TOKEN_LENGTH = 128
BACKUP_FAILED = (
@@ -91,6 +101,91 @@ 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 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 the temporary folder
(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", tempfile.gettempdir(), 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
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"
@@ -220,13 +315,22 @@ class MaintenanceRoutesMixin:
return "on" if PromptManagerConfig.SAFETY_BACKUP else "off"
async def _read_maintenance_body(self, request, allowed):
"""``(data, None)`` for a capped JSON object with only ``allowed`` keys."""
data, error = await self._read_json_body(request, max_bytes=PREVIEW_BODY_MAX)
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])
@@ -234,14 +338,22 @@ class MaintenanceRoutesMixin:
async def preview_maintenance(self, request):
data, error = await self._read_maintenance_body(
request, ("operation", "options")
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 MAINTENANCE_OPERATIONS:
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
@@ -258,6 +370,28 @@ class MaintenanceRoutesMixin:
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.",
@@ -265,18 +399,20 @@ class MaintenanceRoutesMixin:
"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,
draft.fingerprint,
{
"targets": draft.targets,
"roots": roots,
"protected": draft.protected["count"],
"backup": policy,
},
identity,
operation, fingerprint, {**payload, "backup": policy}, identity
)
except ScopeTooLarge:
return _failure(
@@ -289,11 +425,13 @@ class MaintenanceRoutesMixin:
"success": True,
"token": token,
"expires_at": _iso(plan.expires_at),
**draft.public(),
**public,
"backup": {
"policy": policy,
"destination_label": f"{name}.backup_<date and time>",
"expected_bytes": size if policy == "on" else None,
},
"disk": disk,
},
headers={"Cache-Control": "no-store"},
)
@@ -330,36 +468,81 @@ class MaintenanceRoutesMixin:
409,
"storage-changed",
)
backup = self._backup_promised(plan)
try:
job = self._jobs.start(
MAINTENANCE_JOB_KIND,
MAINTENANCE_TITLES[plan.operation],
admission="records",
cancellable=False,
scope={"operation": plan.operation},
groups=MAINTENANCE_GROUPS,
runner=lambda job: self._run_maintenance_plan(job, db, plan),
_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")
try:
if plan.operation == DUPLICATE_FILES:
job = self._start_duplicate_cleanup(db, plan)
else:
job = self._jobs.start(
MAINTENANCE_JOB_KIND,
MAINTENANCE_TITLES[plan.operation],
admission="records",
cancellable=False,
scope={"operation": plan.operation},
groups=MAINTENANCE_GROUPS,
runner=lambda job: self._run_maintenance_plan(job, db, plan),
)
except JobBusy as busy:
return busy_response(busy) # the token stays usable
# Same loop turn as the start: the runner has not run yet.
store.consume(token, plan.operation, job_id=job.id)
return web.json_response({"success": True, "job": job.summary()}, status=202)
async def _run_maintenance_plan(self, job, db, plan):
payload = plan.payload
backup = None
if payload["backup"] == "on" or self._safety_backup_policy() == "on":
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)
return _fail_job(job, "backup-failed", BACKUP_FAILED)
job.log_line(f"Safety backup {backup} verified")
else:
def _backup_promised(self, plan):
"""A plan reviewed with the policy on keeps that promise; so does on now."""
return plan.payload["backup"] == "on" or self._safety_backup_policy() == "on"
async def _safety_backup_step(self, job, db, plan):
"""``(ok, name)``: take the promised safety backup, or report none.
``name`` is the backup's file name, or None when the policy is off.
A failed backup ends the job ``backup-failed`` (``ok`` False).
"""
if not self._backup_promised(plan):
job.log_line("No safety backup (the policy is off)")
return True, None
job.update(status_line="Taking a safety backup")
job.log_line("Taking a safety backup of the database")
try:
backup = await self._run_in_executor(take_safety_backup, db.model)
except SafetyBackupFailed:
self.logger.error("Safety backup failed", exc_info=True)
_fail_job(job, "backup-failed", BACKUP_FAILED)
return False, None
job.log_line(f"Safety backup {backup} verified")
return True, backup
async def _prune_backup_pool(self, job, db):
"""Keep the one ``<db>.backup_*`` pool (restore's too) at its cap."""
removed = await self._run_in_executor(_prune_safety_backups, db.model.db_path)
if removed:
job.log_line(f"Removed {len(removed)} older safety backups")
async def _run_maintenance_plan(self, job, db, plan):
ok, backup = await self._safety_backup_step(job, db, plan)
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(
@@ -381,17 +564,7 @@ class MaintenanceRoutesMixin:
job.record_counted("protected", payload["protected"])
line = _changed_line(plan.operation, result["changed"])
job.log_line(line)
if backup:
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")
job.update(
progress=100,
status_line=f"{line} · "
+ (f"Safety backup {backup}" if backup else NO_BACKUP),
)
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)."""
@@ -426,6 +599,11 @@ class MaintenanceRoutesMixin:
]
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}")
+964
View File
@@ -0,0 +1,964 @@
"""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):
[entry] = [f for f in group["files"] if f["path"] == path]
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
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))
await asyncio.sleep(0.1)
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))
await asyncio.sleep(0.1)
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 · "))
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 · 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_keeps_its_links(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)
[image] = self.api.db.get_prompt_images(prompt)
self.assertTrue(image["image_path"].endswith("c.png"))
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 asyncio.sleep(0.1)
[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_link_move_failure_is_reported(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_file_links",
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.LINKS_FAILED
)
self.assertEqual(cleanup.counters["removed"], 2)
+143 -5
View File
@@ -26,7 +26,7 @@ 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 # noqa: E402
from py.config import GalleryConfig, PromptManagerConfig # noqa: E402
from utils.hashing import generate_prompt_hash # noqa: E402
try:
@@ -47,7 +47,9 @@ PLAN_KEYS = {
"backup",
"media",
"details",
"disk",
}
DISK_KEYS = {"free_bytes", "required_bytes", "sufficient"}
def _digest(path):
@@ -77,6 +79,14 @@ class MaintenanceAPITestCase(AdminAPITestCase):
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()
@@ -102,6 +112,7 @@ class TestPreviewRoute(MaintenanceAPITestCase):
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"],
{
@@ -109,6 +120,8 @@ class TestPreviewRoute(MaintenanceAPITestCase):
"destination_label": "prompts.db.backup_<date and time>",
},
)
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"))
@@ -377,7 +390,6 @@ class TestApplyRoute(ApplyTestCase):
self.assertEqual(self.strays(), [])
async def test_policy_off_reports_no_backup(self):
from py.config import PromptManagerConfig
PromptManagerConfig.SAFETY_BACKUP = False
orphan = self.save("Orphan")
@@ -392,7 +404,6 @@ class TestApplyRoute(ApplyTestCase):
self.assertIsNone(self.api.db.get_prompt_by_id(orphan))
async def test_a_plan_reviewed_with_the_backup_on_keeps_that_promise(self):
from py.config import PromptManagerConfig
self.save("Orphan")
plan = await self.preview("unused_components")
@@ -624,6 +635,134 @@ class TestSafetyBackup(ApplyTestCase):
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")
with patch.object(maintenance, "_free_bytes", return_value=10 * 2**30):
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")
@@ -655,9 +794,8 @@ class TestSafetyBackupPolicy(MaintenanceAPITestCase):
self.assertFalse(bad["success"])
async def test_config_file_value_is_read_at_load(self):
from py.config import PromptManagerConfig
PromptManagerConfig.update_config({"maintenance": {"safety_backup": False}})
PromptManagerConfig.SAFETY_BACKUP = False
self.assertFalse(PromptManagerConfig.SAFETY_BACKUP)
PromptManagerConfig.update_config({"maintenance": {"safety_backup": "no"}})
self.assertFalse(PromptManagerConfig.SAFETY_BACKUP) # non-bool ignored