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"})
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):
    """The JSON object with exactly ``keys``, or raise ValueError."""
    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")
    if set(data) != keys:
        raise ValueError("Request must contain exactly " + ", ".join(sorted(keys)))
    return data


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)
            ids = parse_bulk_prompt_ids(data["prompt_ids"])
        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"],
                )
                disk = None
                if review["delete"] 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_found": len(review["not_found"]),
                "image_links": review["image_links"],
                "backup": policy,
                "disk": disk,
            }
        )

    async def delete_library_prompts(self, request):
        try:
            data = await _body(request, _DELETE_KEYS)
            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")
        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
                )
                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):
        epoch = data["expected_epoch"]
        promised = data["backup"] == "on" or self._safety_backup_policy() == "on"
        backup = None
        try:
            candidates = None
            if delete_ids:
                review = await self._run_in_executor(
                    db.classify_library_delete, delete_ids, expected_epoch=epoch
                )
                candidates = review["fingerprints"]
                if candidates 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 {},
            )
        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})