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:
+63
-1
@@ -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
|
||||
|
||||
@@ -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
@@ -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
@@ -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):
|
||||
|
||||
@@ -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
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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
@@ -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}")
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user