ComfyAssets-ComfyUI_PromptManager/py/api/library_delete.py

"""Delete or archive selected Library prompts, and restore archived ones (#443).

The workspace Prompt library's Delete… (D152). The legacy ``POST
/prompt_manager/bulk/delete`` and ``DELETE /prompt_manager/delete/{id}`` are
separate and unchanged.

- ``POST /prompt_manager/library/delete/preview {prompt_ids, expected_epoch}``
  classifies the selection (reads only, under a lease): the ids to ``delete``
  (in no saved pair), to ``archive`` (paired), the ``protected`` ones (tagged
  ``__protected__``, kept), counts of ``already_archived`` and ``not_found``,
  the ``image_links`` a delete removes, the safety ``backup`` policy and,
  when a backup would be taken, the ``disk`` it needs.
- ``POST /prompt_manager/library/delete {archive_ids, delete_ids,
  expected_epoch, backup}`` executes that review under a ``records`` hold
  (refused 409 ``busy`` beside a scan, restore, maintenance, move or
  Optimize). When a delete remains and a backup was promised (the reviewed
  ``backup`` was "on" or the policy is on now) it checks the disk (507
  ``insufficient-space``) and takes a verified safety backup (500
  ``backup-failed``); then one transaction deletes the reviewed deletes whose
  fingerprint is unchanged since the backup and archives the paired ones.
- ``POST /prompt_manager/library/restore {prompt_ids, expected_epoch}``
  removes the archive mark (under a lease).

A changed storage epoch answers 409 ``storage-changed``; nothing is written.
"""

import sqlite3

from aiohttp import web

try:
    from ...database.library_commands import LibraryCommandError
    from ...database.safety_backup import SafetyBackupFailed
except ImportError:
    from database.library_commands import LibraryCommandError
    from database.safety_backup import SafetyBackupFailed

from . import maintenance as _maintenance
from .admin import _prune_safety_backups
from .job_routes import busy_response
from .jobs import PROMPT_DELETE_JOB_KIND, JobBusy
from .prompts import parse_bulk_prompt_ids

PREVIEW_URL = "/prompt_manager/library/delete/preview"
DELETE_URL = "/prompt_manager/library/delete"
RESTORE_URL = "/prompt_manager/library/restore"
DELETE_TITLE = "Delete prompts (Library)"
_SELECTION_KEYS = frozenset({"prompt_ids", "expected_epoch"})
_DELETE_KEYS = frozenset({"archive_ids", "delete_ids", "expected_epoch", "backup"})
# #482: Delete permanently, with its digests; both or neither (a #443 body).
_PAIRED_KEYS = frozenset({"delete_paired_ids", "paired_digests"})
# #492: the shelf the Delete… runs on, optional on its own (default "active").
_SHELF_KEYS = frozenset({"shelf"})
STORAGE_CHANGED = "Prompt storage was restored; refresh the library"


def _failure(status, code, error, **extra):
    return web.json_response(
        {"success": False, "code": code, "error": error, **extra}, status=status
    )


def _invalid(message):
    return _failure(400, "invalid-request", message)


def _storage_changed(exc):
    return _failure(409, exc.code, STORAGE_CHANGED, library_epoch=exc.epoch)


async def _body(request, keys, *groups):
    """The JSON object with exactly ``keys``, plus all or none of each group."""
    try:
        data = await request.json()
    except ValueError:
        raise ValueError("Expected a JSON object") from None
    if not isinstance(data, dict):
        raise ValueError("Expected a JSON object")
    extra = set(data) - keys
    partial = False
    for group in groups:
        if extra & group:
            partial = partial or not group <= extra
            extra -= group
    if not keys <= set(data) or extra or partial:
        raise ValueError("Request must contain exactly " + ", ".join(sorted(keys)))
    return data


def _shelf(data):
    """#492: the shelf a Delete… was reviewed on."""
    shelf = data.get("shelf", "active")
    if shelf not in ("active", "archived"):
        raise ValueError("shelf must be active or archived")
    return shelf


def _paired_delete(data, archive_ids):
    """#482: the reviewed archive ids to delete with their pairs -> digest."""
    if "delete_paired_ids" not in data:
        return {}
    ids = _reviewed_ids(data["delete_paired_ids"], "delete_paired_ids")
    digests = data["paired_digests"]
    if not isinstance(digests, dict):
        raise ValueError("paired_digests must map prompt ids to digests")
    if not set(ids) <= set(archive_ids):
        raise ValueError("delete_paired_ids must be reviewed archive ids")
    chosen = {}
    for pid in ids:
        digest = digests.get(str(pid))
        if not isinstance(digest, str) or len(digest) != 64:
            raise ValueError("Every delete_paired_ids entry needs its review digest")
        chosen[pid] = digest
    return chosen


def _reviewed_ids(value, name):
    """A reviewed id list: possibly empty, else as ``parse_bulk_prompt_ids``."""
    if value == []:
        return []
    try:
        return parse_bulk_prompt_ids(value)
    except ValueError as exc:
        raise ValueError(str(exc).replace("prompt_ids", name)) from None


class LibraryDeleteRoutesMixin:
    """Mixin providing the workspace Library's delete, archive and restore."""

    def _register_library_delete_routes(self, routes):
        routes.post(PREVIEW_URL)(self.preview_library_delete)
        routes.post(DELETE_URL)(self.delete_library_prompts)
        routes.post(RESTORE_URL)(self.restore_library_prompts)

    async def preview_library_delete(self, request):
        try:
            data = await _body(request, _SELECTION_KEYS, _SHELF_KEYS)
            ids = parse_bulk_prompt_ids(data["prompt_ids"])
            shelf = _shelf(data)
        except ValueError as exc:
            return _invalid(str(exc))
        db = self.db
        policy = self._safety_backup_policy()
        try:
            async with self._jobs.lease("Library delete review"):
                review = await self._run_in_executor(
                    db.classify_library_delete,
                    ids,
                    expected_epoch=data["expected_epoch"],
                    shelf=shelf,
                )
                disk = None
                if (review["delete"] or review["deletable"]) and policy == "on":
                    public, short, _size = await self._run_in_executor(
                        _maintenance.disk_plan,
                        db.model.db_path,
                        backup=True,
                        vacuum=False,
                    )
                    disk = {
                        **public,
                        "message": (
                            None
                            if short is None
                            else _maintenance.short_space_message(short)
                        ),
                    }
        except JobBusy as busy:
            return busy_response(busy)
        except LibraryCommandError as exc:
            return _storage_changed(exc)
        except ValueError as exc:
            return _invalid(str(exc))
        except (OSError, sqlite3.Error):
            self.logger.exception("Library delete review failed")
            return _failure(500, "review-failed", "The delete could not be reviewed")
        return web.json_response(
            {
                "success": True,
                "library_epoch": review["library_epoch"],
                "delete": review["delete"],
                "archive": review["archive"],
                "protected": review["protected"],
                "already_archived": len(review["already_archived"]),
                "not_archived": len(review["not_archived"]),
                "not_found": len(review["not_found"]),
                "image_links": review["image_links"],
                # #482: the archive ids Delete permanently may delete with pairs.
                "deletable": review["deletable"],
                "review_digests": {
                    str(pid): digest for pid, digest in review["review_digests"].items()
                },
                "paired_counts": review["paired_counts"],
                "paired_image_links": review["paired_image_links"],
                "backup": policy,
                "disk": disk,
            }
        )

    async def delete_library_prompts(self, request):
        try:
            data = await _body(request, _DELETE_KEYS, _PAIRED_KEYS, _SHELF_KEYS)
            shelf = _shelf(data)
            archive_ids = _reviewed_ids(data["archive_ids"], "archive_ids")
            delete_ids = _reviewed_ids(data["delete_ids"], "delete_ids")
            if not archive_ids and not delete_ids:
                raise ValueError("Nothing to delete or archive")
            if len(archive_ids) + len(delete_ids) > 1000:
                raise ValueError("At most 1000 prompts can be changed at once")
            if set(archive_ids) & set(delete_ids):
                raise ValueError("A prompt cannot be both archived and deleted")
            if data["backup"] not in ("on", "off"):
                raise ValueError("backup must be on or off")
            paired = _paired_delete(data, archive_ids)
        except ValueError as exc:
            return _invalid(str(exc))
        busy = self._jobs.blocker(PROMPT_DELETE_JOB_KIND, admission="records")
        if busy is not None:
            return busy_response(JobBusy(busy))
        try:
            async with self._jobs.hold(
                PROMPT_DELETE_JOB_KIND, DELETE_TITLE, admission="records"
            ) as job:
                response = await self._held_library_delete(
                    self.db, data, archive_ids, delete_ids, paired, shelf
                )
                if response.status >= 400:
                    job.finish("failed")
                return response
        except JobBusy as busy:
            return busy_response(busy)

    async def _held_library_delete(
        self, db, data, archive_ids, delete_ids, paired, shelf="active"
    ):
        epoch = data["expected_epoch"]
        promised = data["backup"] == "on" or self._safety_backup_policy() == "on"
        backup = None
        try:
            candidates = None
            if delete_ids or paired:
                review = await self._run_in_executor(
                    db.classify_library_delete,
                    [*delete_ids, *paired],
                    expected_epoch=epoch,
                    # #492: the same shelf, or archived ids would skip the backup.
                    shelf=shelf,
                )
                candidates = review["fingerprints"]
                # #482: a permanent delete of a prompt in a saved pair is backed up too.
                if (candidates or review["deletable"]) and promised:
                    _disk, short, _size = await self._run_in_executor(
                        _maintenance.disk_plan,
                        db.model.db_path,
                        backup=True,
                        vacuum=False,
                    )
                    if short is not None:
                        return _failure(
                            507,
                            "insufficient-space",
                            _maintenance.short_space_message(short),
                        )
                    backup = await self._run_in_executor(
                        _maintenance.take_safety_backup, db.model
                    )
            result = await self._run_in_executor(
                db.archive_or_delete_library_prompts,
                archive_ids,
                delete_ids,
                expected_epoch=epoch,
                candidates=candidates if candidates is not None else {},
                delete_paired=paired,
                shelf=shelf,
            )
        except LibraryCommandError as exc:
            return _storage_changed(exc)
        except SafetyBackupFailed as error:
            self.logger.error("Safety backup before a Library delete failed")
            return _failure(
                500, "backup-failed", _maintenance.backup_failed_message(error)
            )
        except ValueError as exc:
            return _invalid(str(exc))
        except (OSError, sqlite3.Error):
            self.logger.exception("Library delete failed")
            return _failure(500, "delete-failed", "The prompts could not be changed")
        finally:
            if backup is not None:
                await self._prune_after_delete(db)
        return web.json_response({"success": True, **result, "backup": backup})

    async def _prune_after_delete(self, db):
        try:
            await self._run_in_executor(_prune_safety_backups, db.model.db_path)
        except Exception:  # pruning old copies never fails the delete
            self.logger.exception("Could not prune old safety backups")

    async def restore_library_prompts(self, request):
        try:
            data = await _body(request, _SELECTION_KEYS)
            ids = parse_bulk_prompt_ids(data["prompt_ids"])
        except ValueError as exc:
            return _invalid(str(exc))
        try:
            async with self._jobs.lease("Library restore"):
                result = await self._run_in_executor(
                    self.db.restore_library_prompts,
                    ids,
                    expected_epoch=data["expected_epoch"],
                )
        except JobBusy as busy:
            return busy_response(busy)
        except LibraryCommandError as exc:
            return _storage_changed(exc)
        except ValueError as exc:
            return _invalid(str(exc))
        except (OSError, sqlite3.Error):
            self.logger.exception("Library restore failed")
            return _failure(500, "restore-failed", "The prompts could not be restored")
        return web.json_response({"success": True, **result})