diff --git a/TESTING.md b/TESTING.md index 7881dc3b..5dff959d 100644 --- a/TESTING.md +++ b/TESTING.md @@ -2504,6 +2504,43 @@ npm run --prefix tools/vendor test:workflows -- system-concept.spec.mjs workspac --project chromium --project touch --project firefox --project webkit ``` +### Maintenance contract (#249 System S6) + +`POST /prompt_manager/maintenance/preview {operation, options:{}}` reviews one +operation and answers a `Plan` (`py/api/maintenance.py`, `database/maintenance.py`): +`duplicate_prompts` (legacy `cleanup_duplicates`: equal text ignoring case and outer +spaces, the oldest prompt of each group kept; a group with a member in a saved pair is +kept whole), `optimize` (VACUUM, with the database size, its free pages and the free +space it needs), `unused_components` (legacy `prune_orphaned_prompts`; the #239 +protections are the predicates `ORPHANED_PROMPTS_SQL` is built from, so the plan's +protected reasons are exactly the prune's rules) and `missing_links` (legacy +`cleanup_orphaned_images`, judged as the health check's `missing_files` through +`health.link_states`: only a link under a source that can be read now counts; one under +an unavailable source or outside every source is protected). The operation enum is +strict, `options` must be `{}` and the body is capped at 4 KiB. A preview meets the +read-only standard above on its own `mode=ro` snapshot (`test_preview_writes_nothing`, +per operation, plus the media bytes) and stats files after the snapshot ended. The +plan lists at most 50 detail rows and no server path. Its fingerprint (the ids, text +hashes, metadata, tags, Library state and pair refs of every prompt involved, or each +link's row and source), its targets and the sources it judged stay in the server's +`PlanStore` under a single-use `secrets.token_urlsafe(32)` token: 15 minutes, at most 8 +plans and 16 MiB of payload. A plan over that budget is refused with `too-large`, never +truncated; a new plan drops the oldest ones, which then read as expired; a restart +forgets them all. A plan is bound to the database it read: a restore during the preview +answers `storage-changed`. Errors are coded (`invalid-request`, `too-large`, +`storage-changed`, or a null code with "Failed to review maintenance") and carry no +path or exception text. + +The legacy `POST /prompt_manager/maintenance` keeps its shape. Its `vacuum` result is +`success:false` with the error when VACUUM fails (`vacuum_database` now raises after +logging), and `cleanup_orphaned_images` removes only links under a source that can be +read now (the health check's sources), never every link of an unmounted drive. +`cleanup_missing_images()` without sources keeps the old rule for its other callers. +Query plans on a fixture of 5,100 prompts, 8,000 images, 300 pairs and 200 recorded +runs: every per-prompt protection and fingerprint read is an index search, the +duplicate groups and the link rows are one scan each, and a delete batch (400 ids) +searches by rowid with its cascades by index. + ## Workspace data fields for the concept pages (#249) These fields are additive. The legacy `/prompt_manager/images/output` listing (no diff --git a/database/health.py b/database/health.py index 7feedda3..e17a2022 100644 --- a/database/health.py +++ b/database/health.py @@ -189,6 +189,33 @@ def _image_owner(image_path, file_path, prefixes): return owner +def readable_roots(roots): + """Whether each source can be read now (an unmounted drive cannot).""" + return [_readable(root) for root in roots] + + +def link_states(rows, roots, readable): + """``(row, owner, state)`` for each image row ``(id, image_path, filename, + file_path)``: ``owner`` is the index of the source it lies under (None + outside every source) and ``state`` is ``missing``, ``present``, + ``unavailable`` (its source cannot be read, so the file is not looked at) + or ``outside``. ``readable`` is ``readable_roots(roots)``. The health + check's ``missing_files`` and the maintenance missing-link plan both + judge links this way, so their counts agree. + """ + prefixes = [_root_prefix(root) for root in roots] + for row in rows: + owner = _image_owner(row[1], row[3], prefixes) + if owner is None: + yield row, None, "outside" + elif not readable[owner]: + yield row, owner, "unavailable" + elif os.path.exists(row[1]): + yield row, owner, "present" + else: + yield row, owner, "missing" + + def _unavailable_line(unavailable, *, partial): noun = "source is" if unavailable == 1 else "sources are" effect = "its files were" if unavailable == 1 else "their files were" @@ -224,7 +251,7 @@ def _missing_files(conn, roots): def _check_files(rows, roots): """The ``missing_files`` result of the image *rows*, from the file system.""" - readable = [_readable(root) for root in roots] + readable = readable_roots(roots) sources = [ { "name": _source_name(root, index), @@ -239,13 +266,12 @@ def _check_files(rows, roots): None, [_unavailable_line(unavailable, partial=False)], status="unavailable" ) return {**result, "partial": False, "sources": sources} - prefixes = [_root_prefix(root) for root in roots] outside, missing = 0, [] - for image_id, image_path, filename, file_path in rows: - owner = _image_owner(image_path, file_path, prefixes) - if owner is None: + for row, owner, state in link_states(rows, roots, readable): + if state == "outside": outside += 1 - elif readable[owner] and not os.path.exists(image_path): + elif state == "missing": + image_id, image_path, filename, _file_path = row sources[owner]["count"] += 1 if len(missing) < MAX_DETAIL: name = filename or os.path.basename(image_path) diff --git a/database/maintenance.py b/database/maintenance.py new file mode 100644 index 00000000..0abcd16d --- /dev/null +++ b/database/maintenance.py @@ -0,0 +1,550 @@ +"""Reviewed maintenance plans: read-only previews and their exact apply (#249 S6). + +Four operations, each the native owner of a legacy ``/maintenance`` one and +built on the same queries, so a preview counts what the legacy cleanup would +change: + +- ``duplicate_prompts``: prompts whose text is equal ignoring case and outer + spaces (``DUPLICATE_PROMPTS_SQL``, legacy ``cleanup_duplicates``); each + group keeps its oldest prompt, which takes over the others' images and + metadata. A group with a member in a saved pair is kept whole. +- ``optimize``: VACUUM (legacy ``vacuum``). +- ``unused_components``: prompts nothing references (``ORPHANED_PROMPTS_SQL``, + legacy ``prune_orphaned_prompts``, with the #239 protections). +- ``missing_links``: image rows whose file is gone (legacy + ``cleanup_orphaned_images``), counted only under a source that can be read + now, as the health check's ``missing_files`` judges them + (``health.link_states``). It deletes database rows only, never a file. + +``preview_maintenance`` reads one snapshot on its own ``mode=ro`` connection +(``health.readonly_uri``) and closes it in ``finally``; files are looked at +after the snapshot ended. It returns a ``MaintenanceDraft``: the public plan +fields, the targets an apply needs and the fingerprint of exactly what was +reviewed (ids, text hashes, metadata, tags, Library state and pair refs of +every prompt involved, or each link's row and source). The fingerprint never +leaves the server. + +``apply_maintenance`` recomputes that fingerprint on the write connection +inside ``BEGIN IMMEDIATE`` and raises ``PlanChanged`` before writing anything +when it differs; otherwise it applies the reviewed targets in batches of +``ID_BATCH`` in that one transaction. VACUUM cannot run in a transaction, so +``optimize`` binds to its database only (the caller checks that). +""" + +import hashlib +import json +import os +import shutil +import sqlite3 +from dataclasses import dataclass, field +from pathlib import Path + +from .execution_contexts import LIVE_HINT_CTES +from .health import BATCH_ROWS, link_states, readable_roots, readonly_uri +from .pairs import PairDependencyError + +MAINTENANCE_OPERATIONS = ( + "duplicate_prompts", + "optimize", + "unused_components", + "missing_links", +) +MAX_DETAILS = 50 # detail rows a plan lists +EXCERPT_CHARS = 80 +# Ids bound per statement: the pair-reference read binds each id twice, which +# stays under SQLite's historical 999-variable limit. +ID_BATCH = 400 + +# Why a prompt no image links to is kept (#239): predicates on ``p``, in the +# order the plan names them. ORPHANED_PROMPTS_SQL is built from them, so the +# reasons a plan lists are exactly the rules the prune applies. +PROTECTIONS = ( + ( + "EXISTS (SELECT 1 FROM prompt_pairs pp " + "WHERE pp.positive_prompt_id = p.id OR pp.negative_prompt_id = p.id)", + "in saved pairs", + "in a saved pair", + ), + ( + "EXISTS (SELECT 1 FROM prompt_tags pt JOIN tags t ON pt.tag_id = t.id " + "WHERE pt.prompt_id = p.id AND t.name = '__protected__')", + "tagged __protected__", + "tagged __protected__", + ), + ( + "EXISTS (SELECT 1 FROM prompt_library_state ls WHERE ls.prompt_id = p.id)", + "with Library state (a title or an explicit parent)", + "with Library state (a title or an explicit parent)", + ), + ( + "p.id IN (SELECT prompt_id FROM live_hint_prompts)", + "used by recorded runs", + "used by a recorded run", + ), +) + +# Exempt references: images, pairs, the __protected__ tag, Library parent +# state, and live recorded-run hints (computed once, see LIVE_HINT_CTES). +ORPHANED_PROMPTS_SQL = ( + "WITH " + + LIVE_HINT_CTES + + " SELECT p.id FROM prompts p" + + " LEFT JOIN generated_images gi ON p.id = gi.prompt_id" + + " WHERE gi.prompt_id IS NULL" + + "".join(f" AND NOT ({predicate})" for predicate, _many, _one in PROTECTIONS) +) +# The prompts no image links to, with which protections keep each one. +_UNLINKED_PROTECTIONS_SQL = ( + "WITH " + + LIVE_HINT_CTES + + " SELECT " + + ", ".join( + f"COALESCE(SUM({predicate}), 0)" for predicate, _many, _one in PROTECTIONS + ) + + ", COALESCE(SUM(" + + " OR ".join(f"({predicate})" for predicate, _many, _one in PROTECTIONS) + + "), 0) FROM prompts p WHERE NOT EXISTS " + "(SELECT 1 FROM generated_images gi WHERE gi.prompt_id = p.id)" +) + +# Duplicate prompt groups, as the legacy cleanup finds them (case-insensitive +# text; GROUP_CONCAT order is sorted by created_at afterwards). +DUPLICATE_PROMPTS_SQL = """ + SELECT LOWER(TRIM(text)) as normalized_text, COUNT(*) as count, + GROUP_CONCAT(id) as ids, + GROUP_CONCAT(created_at) as created_dates + FROM prompts + GROUP BY LOWER(TRIM(text)) + HAVING COUNT(*) > 1 +""" + +_LINK_ROWS_SQL = ( + "SELECT id, image_path, filename, file_path FROM generated_images ORDER BY id" +) + +# What a prompt's fingerprint covers, per batch of ids. +_PROMPT_STATE_SQL = ( + "SELECT id, hash, text, category, rating, notes, created_at, updated_at " + "FROM prompts WHERE id IN ({marks}) ORDER BY id", + "SELECT prompt_id, tag_id FROM prompt_tags WHERE prompt_id IN ({marks}) " + "ORDER BY prompt_id, tag_id", + "SELECT * FROM prompt_library_state WHERE prompt_id IN ({marks}) " + "ORDER BY prompt_id", + "SELECT id, positive_prompt_id, negative_prompt_id FROM prompt_pairs " + "WHERE positive_prompt_id IN ({marks}) OR negative_prompt_id IN ({marks}) " + "ORDER BY id", +) + +NO_FILES_DELETED = "No image file is deleted." +LINKS_ONLY = "Database links only; no image file is deleted." +UNUSED_DETAIL = "No image, saved pair, Library state or recorded run uses it" + + +class PlanChanged(Exception): + """What a plan was reviewed on changed before it was applied.""" + + +@dataclass +class MaintenanceDraft: + """A preview: the public plan fields plus what only the server keeps.""" + + operation: str + fingerprint: str + targets: list + scope: dict + protected: dict + media: dict + details: list = field(default_factory=list) + + def public(self): + """The plan fields a response may carry (no fingerprint, no targets).""" + return { + "operation": self.operation, + "scope": dict(self.scope), + "protected": { + "count": self.protected["count"], + "reasons": list(self.protected["reasons"]), + }, + "media": dict(self.media), + "details": [dict(row) for row in self.details], + } + + +def _plural(count, one, many=None): + return f"{count:,} {one if count == 1 else (many or one + 's')}" + + +def _excerpt(text): + text = " ".join(str(text or "").split()) + if len(text) <= EXCERPT_CHARS: + return text + return text[: EXCERPT_CHARS - 1].rstrip() + "…" + + +def _megabytes(size): + return f"{size / 2**20:,.1f} MB" + + +class _Fingerprint: + """A streaming sha256 over JSON-encoded parts.""" + + def __init__(self, operation): + self._hash = hashlib.sha256() + self.add(operation) + + def add(self, part): + encoded = json.dumps(part, separators=(",", ":"), default=str) + self._hash.update(encoded.encode("utf-8") + b"\n") + + def hexdigest(self): + return "sha256:" + self._hash.hexdigest() + + +def _batches(ids): + ids = list(ids) + for start in range(0, len(ids), ID_BATCH): + yield ids[start : start + ID_BATCH] + + +def _marks(batch): + return ",".join("?" for _ in batch) + + +def _add_prompt_state(fingerprint, conn, ids): + """Fold every reviewed fact about the prompts ``ids`` into ``fingerprint``.""" + for batch in _batches(sorted(ids)): + marks = _marks(batch) + for sql in _PROMPT_STATE_SQL: + params = batch * sql.count("{marks}") + for row in conn.execute(sql.format(marks=marks), params): + fingerprint.add(list(row)) + + +def _texts(conn, ids): + """``{id: text}`` of at most MAX_DETAILS prompts, for detail rows.""" + ids = list(ids)[:MAX_DETAILS] + if not ids: + return {} + rows = conn.execute( + f"SELECT id, text FROM prompts WHERE id IN ({_marks(ids)})", ids + ) + return {row[0]: row[1] for row in rows} + + +def _orphaned_ids(conn): + return sorted(row[0] for row in conn.execute(ORPHANED_PROMPTS_SQL)) + + +def _delete_ids(conn, table, ids): + """Delete ``table`` rows by id in batches; returns the rows removed.""" + if table not in ("prompts", "generated_images"): + raise ValueError(f"Unsupported table {table!r}") + removed = 0 + for batch in _batches(ids): + cursor = conn.execute( + f"DELETE FROM {table} WHERE id IN ({_marks(batch)})", # nosec B608 + batch, + ) + removed += cursor.rowcount + return removed + + +def _read_snapshot(db_path, read): + """``read(conn)`` on one ``mode=ro`` snapshot; the connection always closes.""" + conn = sqlite3.connect(readonly_uri(db_path), uri=True) + try: + conn.row_factory = sqlite3.Row + conn.execute("BEGIN") # one read snapshot for the whole plan + return read(conn) + finally: + conn.close() + + +def _relative(image_path, root): + try: + return Path(os.path.relpath(image_path, root)).as_posix() + except ValueError: # another drive (Windows) + return os.path.basename(image_path) + + +def _source_name(root, index): + return os.path.basename(str(root).rstrip("\\/")) or f"Source {index + 1}" + + +class MaintenanceOperations: + """Mixin for PromptDatabase: maintenance previews and their apply.""" + + # ── shared with the legacy cleanup ─────────────────────────────── + + def _duplicate_groups(self, conn): + """``(groups, paired)``: duplicate id groups, oldest first. + + ``paired`` are the groups with a member in a saved pair, which the + cleanup keeps whole. + """ + groups, paired = [], [] + for duplicate in conn.execute(DUPLICATE_PROMPTS_SQL).fetchall(): + ids = duplicate["ids"].split(",") + created_dates = duplicate["created_dates"].split(",") + # Sort IDs by created_at date to keep the oldest + id_date_pairs = list(zip(ids, created_dates)) + id_date_pairs.sort(key=lambda x: x[1]) + sorted_ids = [int(pair[0]) for pair in id_date_pairs] + try: + self._ensure_prompts_unpaired(conn, sorted_ids) + except PairDependencyError: + paired.append(sorted_ids) + continue + groups.append(sorted_ids) + return groups, paired + + def _merge_duplicate_group(self, conn, sorted_ids): + """Keep the oldest prompt of a group; returns (removed, images moved).""" + primary_id, duplicate_ids = sorted_ids[0], sorted_ids[1:] + self.logger.debug( + f"Merging duplicates: keeping {primary_id}, removing {duplicate_ids}" + ) + primary = conn.execute( + "SELECT * FROM prompts WHERE id = ?", (primary_id,) + ).fetchone() + if not primary: + return 0, 0 + merged_metadata = self._merge_duplicate_metadata( + conn, primary_id, duplicate_ids + ) + moved = self._transfer_images_to_primary(conn, primary_id, duplicate_ids) + if merged_metadata: + self._update_primary_with_merged_metadata(conn, primary_id, merged_metadata) + removed = 0 + for duplicate_id in duplicate_ids: + conn.execute("DELETE FROM prompts WHERE id = ?", (duplicate_id,)) + removed += 1 + self.logger.debug(f"Removed duplicate prompt {duplicate_id}") + return removed, moved + + @staticmethod + def _missing_link_ids(rows, roots): + """Ids of the ``rows`` whose file is gone under a readable source.""" + readable = readable_roots(roots) + return [ + row[0] + for row, _owner, state in link_states(rows, roots, readable) + if state == "missing" + ] + + # ── previews ────────────────────────────────────────────────────── + + def preview_maintenance(self, operation, roots=()) -> MaintenanceDraft: + """The reviewed plan of ``operation``; writes nothing (blocking). + + ``roots`` are the sources missing links are judged under (the health + check's: the effective sources plus configured ones that are gone). + """ + if operation not in MAINTENANCE_OPERATIONS: + raise ValueError(f"Unknown maintenance operation {operation!r}") + roots = [str(root) for root in roots] + if operation == "optimize": + return self._plan_optimize() + if operation == "missing_links": + rows = _read_snapshot(self.model.db_path, self._link_rows) + return self._plan_missing_links(rows, roots) + reader = ( + self._plan_duplicate_prompts + if operation == "duplicate_prompts" + else self._plan_unused_components + ) + return _read_snapshot(self.model.db_path, reader) + + @staticmethod + def _link_rows(conn): + rows = [] + cursor = conn.execute(_LINK_ROWS_SQL) + while batch := cursor.fetchmany(BATCH_ROWS): + rows.extend(tuple(row) for row in batch) + return rows + + def _duplicates_fingerprint(self, conn): + groups, paired = self._duplicate_groups(conn) + fingerprint = _Fingerprint("duplicate_prompts") + fingerprint.add({"groups": groups, "paired": paired}) + _add_prompt_state(fingerprint, conn, [i for group in groups for i in group]) + return groups, paired, fingerprint.hexdigest() + + def _plan_duplicate_prompts(self, conn): + groups, paired, fingerprint = self._duplicates_fingerprint(conn) + removed = sum(len(group) - 1 for group in groups) + texts = _texts(conn, [group[0] for group in (groups + paired)]) + details = [ + { + "label": _excerpt(texts.get(group[0])), + "detail": f"Keeps the oldest of {len(group)} copies", + } + for group in groups + ] + [ + { + "label": _excerpt(texts.get(group[0])), + "detail": "Kept: belongs to a saved pair", + } + for group in paired + ] + if groups: + label = ( + f"{_plural(removed, 'duplicate prompt')} in " + f"{_plural(len(groups), 'group')}" + ) + else: + label = "No duplicate prompt text" + reasons = [] + if len(paired) == 1: + reasons.append("1 group belongs to a saved pair") + elif paired: + reasons.append(f"{len(paired):,} groups belong to saved pairs") + return MaintenanceDraft( + operation="duplicate_prompts", + fingerprint=fingerprint, + targets=groups, + scope={"count": removed, "label": label}, + protected={ + "count": sum(len(group) for group in paired), + "reasons": reasons, + }, + media={"deletes_files": False, "label": NO_FILES_DELETED}, + details=details[:MAX_DETAILS], + ) + + def _unused_fingerprint(self, conn): + ids = _orphaned_ids(conn) + fingerprint = _Fingerprint("unused_components") + fingerprint.add(ids) + _add_prompt_state(fingerprint, conn, ids) + return ids, fingerprint.hexdigest() + + def _plan_unused_components(self, conn): + ids, fingerprint = self._unused_fingerprint(conn) + *by_reason, kept = conn.execute(_UNLINKED_PROTECTIONS_SQL).fetchone() + reasons = [ + f"{count:,} {many if count != 1 else one}" + for count, (_predicate, many, one) in zip(by_reason, PROTECTIONS) + if count + ] + texts = _texts(conn, ids) + count = len(ids) + return MaintenanceDraft( + operation="unused_components", + fingerprint=fingerprint, + targets=ids, + scope={ + "count": count, + "label": ( + _plural(count, "unused prompt component") + if count + else "No unused prompt components" + ), + }, + protected={"count": kept, "reasons": reasons}, + media={"deletes_files": False, "label": NO_FILES_DELETED}, + details=[ + {"label": _excerpt(texts.get(prompt_id)), "detail": UNUSED_DETAIL} + for prompt_id in ids[:MAX_DETAILS] + ], + ) + + @staticmethod + def _links_fingerprint(candidates): + fingerprint = _Fingerprint("missing_links") + for row, owner in candidates: + fingerprint.add([row[0], row[1], row[3], owner]) + return fingerprint.hexdigest() + + def _plan_missing_links(self, rows, roots): + readable = readable_roots(roots) + candidates, unavailable, outside = [], 0, 0 + for row, owner, state in link_states(rows, roots, readable): + if state == "missing": + candidates.append((row, owner)) + elif state == "unavailable": + unavailable += 1 + elif state == "outside": + outside += 1 + reasons = [] + if unavailable: + reasons.append( + _plural( + unavailable, + "link under an unavailable source", + "links under unavailable sources", + ) + ) + if outside: + reasons.append( + _plural( + outside, + "link outside the image sources", + "links outside the image sources", + ) + ) + count = len(candidates) + return MaintenanceDraft( + operation="missing_links", + fingerprint=self._links_fingerprint(candidates), + targets=[row[0] for row, _owner in candidates], + scope={ + "count": count, + "label": ( + _plural(count, "link to a missing file", "links to missing files") + if count + else "No links to missing files" + ), + }, + protected={"count": unavailable + outside, "reasons": reasons}, + media={"deletes_files": False, "label": LINKS_ONLY}, + details=[ + { + "label": _relative(row[1], roots[owner]), + "detail": ( + f"{_source_name(roots[owner], owner)} · the file is missing" + ), + } + for row, owner in candidates[:MAX_DETAILS] + ], + ) + + def _plan_optimize(self): + def read(conn): + return [ + conn.execute(f"PRAGMA {name}").fetchone()[0] + for name in ("page_size", "page_count", "freelist_count") + ] + + page_size, page_count, free_pages = _read_snapshot(self.model.db_path, read) + db_path = self.model.db_path + wal = db_path + "-wal" + size = page_size * page_count + wal_size = os.path.getsize(wal) if os.path.exists(wal) else 0 + free = shutil.disk_usage(os.path.dirname(os.path.abspath(db_path))).free + # VACUUM writes the rebuilt database through the WAL beside it. + return MaintenanceDraft( + operation="optimize", + fingerprint=_Fingerprint("optimize").hexdigest(), + targets=[], + scope={"count": 1, "label": "The whole database"}, + protected={"count": 0, "reasons": []}, + media={"deletes_files": False, "label": NO_FILES_DELETED}, + details=[ + { + "label": "Database file", + "detail": f"{_megabytes(size)}, plus {_megabytes(wal_size)} " + "not yet checkpointed", + }, + { + "label": "Free pages", + "detail": f"About {_megabytes(free_pages * page_size)} " + "can be reclaimed", + }, + { + "label": "Free space needed", + "detail": f"About {_megabytes(size)} in the database folder; " + f"{_megabytes(free)} available", + }, + ], + ) diff --git a/database/models.py b/database/models.py index 37244656..07c2e4f1 100644 --- a/database/models.py +++ b/database/models.py @@ -941,6 +941,10 @@ class PromptModel: Reclaims unused space and defragments the database file, improving query performance and reducing file size. + + Raises: + Exception: whatever stopped VACUUM (logged first); the database is + then unchanged. """ try: with closing(sqlite3.connect(self.db_path)) as conn: @@ -948,6 +952,7 @@ class PromptModel: conn.commit() except Exception as e: self.logger.error(f"Error vacuuming database: {e}") + raise def get_database_info(self) -> dict: """ diff --git a/database/operations.py b/database/operations.py index cf9cec13..6eaa557a 100644 --- a/database/operations.py +++ b/database/operations.py @@ -10,7 +10,7 @@ import os from typing import Optional, List, Dict, Any, Set, Tuple from .models import PromptModel, normalize_image_path, utc_now_iso -from .pairs import PairDependencyError, PairOperations +from .pairs import PairDependencyError, PairOperations # noqa: F401 (re-exported) from .library import LibraryOperations, titled from .image_favorites import ImageFavoriteOperations from .annotations import AnnotationOperations @@ -20,10 +20,17 @@ from .saved_views import SavedViewOperations from .library_commands import LibraryCommandOperations from .vocabulary import VocabularyOperations from .execution_history import ExecutionHistoryOperations -from .execution_contexts import ExecutionContextOperations, LIVE_HINT_CTES +from .execution_contexts import ExecutionContextOperations from .execution_outputs import ExecutionOutputOperations from .history_reads import HistoryReadOperations from .health import HealthOperations, iter_consistency_issues +from .maintenance import ( # noqa: F401 (ORPHANED_PROMPTS_SQL is re-exported) + ORPHANED_PROMPTS_SQL, + _LINK_ROWS_SQL, + MaintenanceOperations, + _delete_ids, + _orphaned_ids, +) # Import logging system try: @@ -166,28 +173,6 @@ def _coerce_file_info(metadata: Any) -> tuple: ) -# Exempt references: images, pairs, the __protected__ tag, Library parent -# state, and live recorded-run hints (computed once, see LIVE_HINT_CTES). -ORPHANED_PROMPTS_SQL = "WITH " + LIVE_HINT_CTES + """ -SELECT p.id FROM prompts p -LEFT JOIN generated_images gi ON p.id = gi.prompt_id -WHERE gi.prompt_id IS NULL -AND NOT EXISTS ( - SELECT 1 FROM prompt_pairs pp - WHERE pp.positive_prompt_id = p.id OR pp.negative_prompt_id = p.id -) -AND NOT EXISTS ( - SELECT 1 FROM prompt_tags pt - JOIN tags t ON pt.tag_id = t.id - WHERE pt.prompt_id = p.id AND t.name = '__protected__' -) -AND NOT EXISTS ( - SELECT 1 FROM prompt_library_state ls WHERE ls.prompt_id = p.id -) -AND p.id NOT IN (SELECT prompt_id FROM live_hint_prompts) -""" - - # The scan's known set (#249 Activity-2): every image row, with its stored file # size and whether its pair link is resolved: exactly one distinct pair, the rule # _LINK_SUMMARY_SQL and get_image_prompt_info use. A file is known and resolved @@ -351,6 +336,7 @@ class PromptDatabase( ReviewDecisionOperations, SavedViewOperations, HealthOperations, + MaintenanceOperations, ): """Database operations class for managing prompts.""" @@ -1541,80 +1527,14 @@ class PromptDatabase( self.logger.info("Starting duplicate cleanup process with image preservation") try: with self._pair_transaction() as conn: - # Find duplicates by text content (case-insensitive) - # Note: Removed ORDER BY from GROUP_CONCAT for SQLite compatibility - # We'll sort the IDs manually after fetching - cursor = conn.execute(""" - SELECT LOWER(TRIM(text)) as normalized_text, COUNT(*) as count, - GROUP_CONCAT(id) as ids, - GROUP_CONCAT(created_at) as created_dates - FROM prompts - GROUP BY LOWER(TRIM(text)) - HAVING COUNT(*) > 1 - """) - - duplicates = cursor.fetchall() - self.logger.debug( - f"Found {len(duplicates)} groups of duplicate prompts" - ) + groups, _paired = self._duplicate_groups(conn) + self.logger.debug(f"Found {len(groups)} groups of duplicate prompts") total_removed = 0 total_images_transferred = 0 - - for duplicate in duplicates: - ids = duplicate["ids"].split(",") - created_dates = duplicate["created_dates"].split(",") - - # Sort IDs by created_at date to keep the oldest - id_date_pairs = list(zip(ids, created_dates)) - id_date_pairs.sort(key=lambda x: x[1]) # Sort by date - sorted_ids = [int(pair[0]) for pair in id_date_pairs] - - try: - self._ensure_prompts_unpaired(conn, sorted_ids) - except PairDependencyError: - continue - - # Keep the oldest one (first), merge and delete the rest - primary_id = sorted_ids[0] # Keep the oldest - duplicate_ids = sorted_ids[1:] - - self.logger.debug( - f"Merging duplicates: keeping {primary_id}, removing " - f"{duplicate_ids}" - ) - - # Get primary prompt details - cursor = conn.execute( - "SELECT * FROM prompts WHERE id = ?", (primary_id,) - ) - primary_prompt = cursor.fetchone() - if not primary_prompt: - continue - - # Collect and merge metadata from all duplicates - merged_metadata = self._merge_duplicate_metadata( - conn, primary_id, duplicate_ids - ) - - # Transfer all images from duplicates to primary prompt - images_transferred = self._transfer_images_to_primary( - conn, primary_id, duplicate_ids - ) - total_images_transferred += images_transferred - - # Update primary prompt with merged metadata - if merged_metadata: - self._update_primary_with_merged_metadata( - conn, primary_id, merged_metadata - ) - - # Delete duplicate prompts (images already transferred) - for duplicate_id in duplicate_ids: - conn.execute( - "DELETE FROM prompts WHERE id = ?", (duplicate_id,) - ) - total_removed += 1 - self.logger.debug(f"Removed duplicate prompt {duplicate_id}") + for sorted_ids in groups: + removed, moved = self._merge_duplicate_group(conn, sorted_ids) + total_removed += removed + total_images_transferred += moved conn.commit() @@ -1879,13 +1799,30 @@ class PromptDatabase( conn.commit() return cursor.rowcount > 0 - def cleanup_missing_images(self) -> int: + def cleanup_missing_images(self, roots=None) -> int: """ Remove image records where the actual file no longer exists. + Args: + roots: The image sources links are judged under (the health + check's). When given, only a link under a source that can be + read now is removed: a link under an unavailable source (an + unmounted drive) or outside every source is kept. None keeps + the old rule: every link whose file does not exist. + Returns: int: Number of orphaned records removed """ + if roots is not None: + with self.model.get_connection() as conn: + rows = [tuple(row) for row in conn.execute(_LINK_ROWS_SQL)] + # Files are looked at before the write transaction starts. + ids = self._missing_link_ids(rows, [str(root) for root in roots]) + if not ids: + return 0 + with self._pair_transaction() as conn: + return _delete_ids(conn, "generated_images", ids) + removed_count = 0 with self.model.get_connection() as conn: @@ -2567,16 +2504,12 @@ class PromptDatabase( Returns number of prompts removed. """ with self._pair_transaction() as conn: - cursor = conn.execute(ORPHANED_PROMPTS_SQL) - orphaned = [row["id"] for row in cursor.fetchall()] + orphaned = _orphaned_ids(conn) if not orphaned: return 0 - placeholders = ",".join(["?"] * len(orphaned)) - cursor = conn.execute( - f"DELETE FROM prompts WHERE id IN ({placeholders})", orphaned - ) + removed = _delete_ids(conn, "prompts", orphaned) conn.commit() - return cursor.rowcount + return removed def check_consistency(self) -> List[str]: """Run consistency checks and return a list of issue descriptions. diff --git a/py/api/__init__.py b/py/api/__init__.py index f929b2d1..8a3e6e21 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -13,6 +13,7 @@ Split into domain-specific mixins for maintainability: - ReviewRecordingMixin: Needs review recording for the scan and the monitor - ReviewEventsMixin: the Needs review change stream an open page re-reads on - SystemRoutesMixin: System health check, Overview and usage statistics +- MaintenanceRoutesMixin: reviewed maintenance plans (preview, apply) """ import asyncio @@ -46,6 +47,8 @@ from .job_routes import JobRoutesMixin from .image_tag_jobs import ImageTagJobsMixin from .thumbnail_jobs import ThumbnailJobsMixin from .system import SystemRoutesMixin +from .maintenance import MaintenanceRoutesMixin +from .maintenance_plans import PlanStore from .jobs import JobRegistry from .logging_routes import LoggingRoutesMixin from .autotag_routes import AutotagRoutesMixin @@ -285,6 +288,7 @@ class PromptManagerAPI( ImageTagJobsMixin, ThumbnailJobsMixin, SystemRoutesMixin, + MaintenanceRoutesMixin, LoggingRoutesMixin, AutotagRoutesMixin, LoraIntegrationMixin, @@ -348,6 +352,9 @@ class PromptManagerAPI( # (py/api/system.py). self._health_report = None self._db_generation = 0 + # Reviewed maintenance plans under single-use tokens; a restart + # forgets them (py/api/maintenance.py). + self._maintenance_plans = PlanStore() # Run cleanup on initialization to remove any existing duplicates try: @@ -508,6 +515,7 @@ class PromptManagerAPI( self._register_settings_routes(routes) self._register_job_routes(routes) self._register_system_routes(routes) + self._register_maintenance_routes(routes) self._register_logging_routes(routes) self._register_autotag_routes(routes) self._register_lora_routes(routes) diff --git a/py/api/admin.py b/py/api/admin.py index 0b9b89cb..0da16fa3 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -1112,8 +1112,10 @@ class AdminRoutesMixin: db.cleanup_duplicates(), "duplicate prompts" ), "vacuum": vacuum, + # Only under sources that can be read now (never an unmounted drive). "cleanup_orphaned_images": lambda: count_result( - db.cleanup_missing_images(), "orphaned image records" + db.cleanup_missing_images(self._health_roots()), + "orphaned image records", ), "check_hash_duplicates": lambda: self._hash_duplicates_result( db.check_hash_duplicates() diff --git a/py/api/maintenance.py b/py/api/maintenance.py new file mode 100644 index 00000000..ca2df99b --- /dev/null +++ b/py/api/maintenance.py @@ -0,0 +1,152 @@ +"""Maintenance routes: reviewed, read-only plans (#249 System S6). + +``POST /prompt_manager/maintenance/preview {operation, options:{}}`` reviews +one operation (``duplicate_prompts``, ``optimize``, ``unused_components`` or +``missing_links``; ``database/maintenance.py``) off the event loop and +answers a ``Plan``: ``{success, token, operation, expires_at, scope:{count, +label}, protected:{count, reasons}, backup:{policy, destination_label}, +media:{deletes_files, label}, details:[≤50 {label, detail}]}``. The preview +writes nothing. The plan's fingerprint, its targets and the sources it judged +stay in the server's ``PlanStore`` under the token, bound to the database the +preview read (a restore that replaces it invalidates the plan). + +The operation enum is strict, ``options`` must be an empty object and the +body is capped at ``PREVIEW_BODY_MAX`` bytes. Errors are ``{success:false, +code, error}`` with ``invalid-request`` (400), ``too-large`` (413: the body, +or a plan over the store's budget), ``storage-changed`` (409: the database +was replaced while the preview read it) or a null code (500); none carries a +server path or exception text. +""" + +import datetime +import os + +from aiohttp import web + +from .maintenance_plans import ScopeTooLarge + +try: + from ...database.maintenance import MAINTENANCE_OPERATIONS +except ImportError: + from database.maintenance import MAINTENANCE_OPERATIONS + +PREVIEW_BODY_MAX = 4096 +# The concept's Maintenance row titles, which name each operation's job. +MAINTENANCE_TITLES = { + "duplicate_prompts": "Duplicate prompt text", + "optimize": "Optimize database", + "unused_components": "Unused prompt components", + "missing_links": "Missing-file links", +} +_OPERATION_ERROR = "operation must be one of: " + ", ".join(MAINTENANCE_OPERATIONS) + + +def _failure(message, status, code=None, **extra): + return web.json_response( + {"success": False, "code": code, "error": message, **extra}, status=status + ) + + +def _invalid(message): + return _failure(message, 400, "invalid-request") + + +def _iso(timestamp): + return datetime.datetime.fromtimestamp(timestamp, datetime.timezone.utc).isoformat() + + +class MaintenanceRoutesMixin: + """Mixin providing the /prompt_manager/maintenance plan endpoints.""" + + def _register_maintenance_routes(self, routes): + @routes.post("/prompt_manager/maintenance/preview") + async def maintenance_preview_route(request): + return await self.preview_maintenance(request) + + def _database_identity(self, db): + """What binds a plan to one database: the swap generation and its file.""" + path = os.path.normcase(os.path.realpath(db.model.db_path)) + return (self._db_generation, path) + + @staticmethod + def _safety_backup_policy(): + """``on`` or ``off``: whether an apply takes a safety backup first.""" + return "on" + + 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) + 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") + unknown = sorted(set(data) - set(allowed)) + if unknown: + return None, _invalid("Unknown field: " + ", ".join(unknown)[:200]) + return data, None + + async def preview_maintenance(self, request): + data, error = await self._read_maintenance_body( + request, ("operation", "options") + ) + 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: + return _invalid(_OPERATION_ERROR) + if not isinstance(options, dict) or options: + return _invalid(f"{operation} takes no options") + db = self.db + identity = self._database_identity(db) + try: + roots = ( + await self._run_in_executor(self._health_roots) + if operation == "missing_links" + else [] + ) + draft = await self._run_in_executor( + db.preview_maintenance, operation, roots + ) + except Exception: + self.logger.exception(f"Maintenance preview of {operation} failed") + return _failure("Failed to review maintenance", 500) + if identity != self._database_identity(self.db): + return _failure( + "The database was replaced while it was reviewed. Review again.", + 409, + "storage-changed", + ) + policy = self._safety_backup_policy() + store = self._maintenance_plans + try: + token = store.issue( + operation, + draft.fingerprint, + { + "targets": draft.targets, + "roots": roots, + "protected": draft.protected["count"], + "backup": policy, + }, + identity, + ) + except ScopeTooLarge: + return _failure( + "This plan is too large to review at once.", 413, "too-large" + ) + plan = store.peek(token, operation) + name = os.path.basename(db.model.db_path) + return web.json_response( + { + "success": True, + "token": token, + "expires_at": _iso(plan.expires_at), + **draft.public(), + "backup": { + "policy": policy, + "destination_label": f"{name}.backup_", + }, + }, + headers={"Cache-Control": "no-store"}, + ) diff --git a/py/api/maintenance_plans.py b/py/api/maintenance_plans.py new file mode 100644 index 00000000..aa257d9a --- /dev/null +++ b/py/api/maintenance_plans.py @@ -0,0 +1,138 @@ +"""The reviewed-plan store behind maintenance previews (#249 System S6). + +A preview issues a plan: its operation, the fingerprint of exactly what was +reviewed, the payload an apply needs (target ids, the sources it judged) and +the identity of the database it describes. Everything is held here, in +server memory, under a single-use token (``secrets.token_urlsafe(32)``); the +client only ever sees the token and the public plan. Plans expire after +``ttl`` seconds and are gone after a restart. + +The store keeps at most ``max_plans`` plans and ``max_bytes`` of payload +(measured as compact JSON). A plan larger than ``max_bytes`` is refused with +``ScopeTooLarge``: it is never cut down to fit, since an apply must do +exactly what was reviewed. A new plan that fits makes room by dropping the +oldest ones, which then read as expired. + +A consumed token is remembered (the last ``MAX_USED``) with the job it +started, so a second apply of it is told which job ran it. +""" + +import json +import secrets +import time +from collections import OrderedDict +from dataclasses import dataclass, field + + +class PlanExpired(Exception): + """No live plan has the token: unknown, expired, dropped or another operation's.""" + + +class PlanUsed(Exception): + """The token was already applied; ``job_id`` is the job that applied it.""" + + def __init__(self, job_id): + super().__init__("This plan was already applied") + self.job_id = job_id + + +class ScopeTooLarge(Exception): + """The plan's payload exceeds the store's byte budget.""" + + +@dataclass +class StoredPlan: + """One reviewed plan, as only the server sees it.""" + + operation: str + fingerprint: str + payload: dict + db_identity: tuple + expires_at: float # wall-clock seconds since the epoch, for the response + deadline: float = field(repr=False) # on the store's monotonic clock + size: int = field(repr=False) + + +def _payload_size(payload): + return len(json.dumps(payload, separators=(",", ":"), default=str)) + + +class PlanStore: + """Reviewed maintenance plans under single-use tokens.""" + + MAX_USED = 64 # consumed tokens remembered for plan-used answers + + def __init__( + self, + ttl=900, + max_plans=8, + max_bytes=16 * 2**20, + *, + clock=time.monotonic, + wall_clock=time.time, + ): + self.ttl = ttl + self.max_plans = max_plans + self.max_bytes = max_bytes + self._clock = clock + self._wall_clock = wall_clock + self._plans = OrderedDict() # token -> StoredPlan, oldest first + self._used = OrderedDict() # token -> job id, oldest first + + def issue(self, operation, fingerprint, payload, db_identity) -> str: + """Hold a plan and return its token; ScopeTooLarge when it cannot fit.""" + size = _payload_size(payload) + if size > self.max_bytes: + raise ScopeTooLarge(f"The plan needs {size} bytes") + self._drop_expired() + while self._plans and ( + len(self._plans) >= self.max_plans + or self._held_bytes() + size > self.max_bytes + ): + self._plans.popitem(last=False) + token = secrets.token_urlsafe(32) + self._plans[token] = StoredPlan( + operation=operation, + fingerprint=fingerprint, + payload=payload, + db_identity=tuple(db_identity), + expires_at=self._wall_clock() + self.ttl, + deadline=self._clock() + self.ttl, + size=size, + ) + return token + + def peek(self, token, operation=None) -> StoredPlan: + """The live plan of ``token``, left in place; raises PlanUsed or PlanExpired.""" + if not isinstance(token, str): + raise PlanExpired("No such plan") + if token in self._used: + raise PlanUsed(self._used[token]) + plan = self._plans.get(token) + if plan is None or (operation is not None and plan.operation != operation): + raise PlanExpired("No such plan") + if self._clock() > plan.deadline: + del self._plans[token] + raise PlanExpired("The plan expired") + return plan + + def consume(self, token, operation=None, *, job_id=None) -> StoredPlan: + """Take the plan of ``token`` once; later calls raise PlanUsed(job_id).""" + plan = self.peek(token, operation) + del self._plans[token] + self._used[token] = job_id + while len(self._used) > self.MAX_USED: + self._used.popitem(last=False) + return plan + + def discard(self, token): + """Forget a plan that can no longer apply (its database was replaced).""" + self._plans.pop(token, None) + + def _held_bytes(self): + return sum(plan.size for plan in self._plans.values()) + + def _drop_expired(self): + now = self._clock() + for token in [t for t, plan in self._plans.items() if now > plan.deadline]: + del self._plans[token] diff --git a/tests/test_api_maintenance.py b/tests/test_api_maintenance.py new file mode 100644 index 00000000..50adf750 --- /dev/null +++ b/tests/test_api_maintenance.py @@ -0,0 +1,273 @@ +"""Maintenance routes (py/api/maintenance.py, #249 System S6). + +``POST /prompt_manager/maintenance/preview`` answers a reviewed ``Plan`` for +one of four operations and writes nothing; its token binds a later apply to +exactly that plan. Errors are coded and never carry a server path or an +exception text. +""" + +import hashlib +import json +import os +import sqlite3 +import sys +from pathlib import Path +from unittest.mock import MagicMock, patch + +_mock_server = MagicMock() +_mock_server.PromptServer.instance.routes = MagicMock() +sys.modules.setdefault("server", _mock_server) + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from database.operations import PromptDatabase # noqa: E402 +from py.api import PromptManagerAPI # noqa: E402 +from py.api.maintenance_plans import PlanExpired # noqa: E402 +from py.config import GalleryConfig # noqa: E402 +from utils.hashing import generate_prompt_hash # noqa: E402 + +try: + from tests.test_api_admin import AdminAPITestCase # noqa: E402 +except ImportError: # discovered with tests/ as the top-level directory + from test_api_admin import AdminAPITestCase # noqa: E402 + +WAIT = 10 +PREVIEW = "/prompt_manager/maintenance/preview" +OPERATIONS = ("duplicate_prompts", "optimize", "unused_components", "missing_links") +PLAN_KEYS = { + "success", + "token", + "operation", + "expires_at", + "scope", + "protected", + "backup", + "media", + "details", +} + + +def _digest(path): + if not os.path.exists(path): + return None + with open(path, "rb") as handle: + return hashlib.sha256(handle.read()).hexdigest() + + +class MaintenanceAPITestCase(AdminAPITestCase): + async def get_application(self): + app = await super().get_application() + self.api._find_comfyui_output_dir = lambda: ( + str(self.output_dir) if self.output_dir.is_dir() else None + ) + return app + + def save(self, text, **metadata): + return self.api.db.save_prompt( + text, prompt_hash=generate_prompt_hash(text), **metadata + ) + + def image(self, prompt_id, name, *, exists=True, folder=None): + path = Path(folder or self.output_dir) / name + path.parent.mkdir(parents=True, exist_ok=True) + if exists: + path.write_bytes(b"png") + return self.api.db.link_image_to_prompt(prompt_id, str(path)) + + async def post(self, url, body, status): + response = await self.client.request("POST", url, json=body) + text = await response.text() + self.assertEqual(response.status, status, text) + self.assertNotIn(self.tmpdir, text) + return json.loads(text) + + async def preview(self, operation, status=200, **body): + return await self.post(PREVIEW, {"operation": operation, **body}, status) + + +class TestPreviewRoute(MaintenanceAPITestCase): + async def test_preview_answers_a_plan_for_each_operation(self): + prompt = self.save("Kept") + self.image(prompt, "gone.png", exists=False) + self.save("Unused") + for operation in OPERATIONS: + with self.subTest(operation=operation): + plan = await self.preview(operation, options={}) + self.assertEqual(set(plan), PLAN_KEYS) + self.assertTrue(plan["success"]) + self.assertEqual(plan["operation"], operation) + self.assertIsInstance(plan["token"], str) + self.assertEqual(set(plan["scope"]), {"count", "label"}) + self.assertEqual(set(plan["protected"]), {"count", "reasons"}) + self.assertEqual( + plan["backup"], + { + "policy": "on", + "destination_label": "prompts.db.backup_", + }, + ) + self.assertFalse(plan["media"]["deletes_files"]) + self.assertLessEqual(len(plan["details"]), 50) + self.assertTrue(plan["expires_at"].endswith("+00:00")) + self.assertNotIn("fingerprint", json.dumps(plan)) + missing = await self.preview("missing_links") + self.assertEqual( + missing["scope"], {"count": 1, "label": "1 link to a missing file"} + ) + unused = await self.preview("unused_components") + self.assertEqual(unused["scope"]["count"], 1) + + async def test_preview_writes_nothing(self): + prompt = self.save("Kept") + self.image(prompt, "kept.png") + self.image(prompt, "gone.png", exists=False) + self.save("Unused") + db_path = self.api.db.model.db_path + self.api.db.get_statistics() + files = [db_path, db_path + "-wal"] + for operation in OPERATIONS: + with self.subTest(operation=operation): + before = [_digest(path) for path in files] + shm = os.path.getsize(db_path + "-shm") + media = {p.name: _digest(p) for p in self.output_dir.iterdir()} + await self.preview(operation) + self.assertEqual([_digest(path) for path in files], before) + self.assertEqual(os.path.getsize(db_path + "-shm"), shm) + self.assertEqual( + {p.name: _digest(p) for p in self.output_dir.iterdir()}, media + ) + self.assertEqual( + [p["text"] for p in self.api.db.search_prompts()], ["Unused", "Kept"] + ) + + async def test_unknown_operation_is_400(self): + cases = [ + {"operation": "reboot"}, + {"operation": "duplicate_files"}, # S6 Task 17, not yet + {"operation": 3}, + {}, + {"operation": "optimize", "options": {"all": True}}, + {"operation": "optimize", "options": []}, + {"operation": "optimize", "extra": 1}, + ] + for body in cases: + with self.subTest(body=body): + reply = await self.post(PREVIEW, body, 400) + self.assertEqual(reply["success"], False) + self.assertEqual(reply["code"], "invalid-request") + response = await self.client.request( + "POST", PREVIEW, data=b"[1]", headers={"Content-Type": "application/json"} + ) + self.assertEqual(response.status, 400) + self.assertEqual((await response.json())["code"], "invalid-request") + big = json.dumps({"operation": "optimize", "pad": "x" * 8192}) + response = await self.client.request( + "POST", PREVIEW, data=big, headers={"Content-Type": "application/json"} + ) + self.assertEqual(response.status, 413) + self.assertEqual((await response.json())["code"], "too-large") + + async def test_preview_failure_is_coded_without_exception_text(self): + with patch.object( + self.api.db, + "preview_maintenance", + side_effect=sqlite3.OperationalError(f"boom {self.tmpdir}"), + ): + with self.assertLogs("prompt_manager", "ERROR"): + reply = await self.preview("optimize", status=500) + self.assertEqual( + reply, + {"success": False, "code": None, "error": "Failed to review maintenance"}, + ) + + async def test_scope_too_large_is_413(self): + from py.api.maintenance_plans import PlanStore + + self.api._maintenance_plans = PlanStore(max_bytes=10) + self.save("Unused") + reply = await self.preview("unused_components", status=413) + self.assertEqual(reply["code"], "too-large") + + async def test_preview_during_a_restore_swap_is_storage_changed(self): + real = self.api.db.preview_maintenance + + def swapped(*args, **kwargs): + draft = real(*args, **kwargs) + self.api._db_generation += 1 # a restore replaced the database + return draft + + with patch.object(self.api.db, "preview_maintenance", side_effect=swapped): + reply = await self.preview("optimize", status=409) + self.assertEqual(reply["code"], "storage-changed") + + async def test_missing_links_preview_uses_the_health_roots(self): + unmounted = Path(self.tmpdir) / "unmounted" + unmounted.mkdir() + GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir), str(unmounted)] + prompt = self.save("Links") + self.image(prompt, "gone.png", exists=False) + self.image(prompt, "away.png", folder=unmounted, exists=False) + unmounted.rmdir() + plan = await self.preview("missing_links") + self.assertEqual(plan["scope"]["count"], 1) + self.assertEqual( + plan["protected"], + {"count": 1, "reasons": ["1 link under an unavailable source"]}, + ) + + async def test_restart_loses_plans(self): + plan = await self.preview("optimize") + self.assertEqual( + self.api._maintenance_plans.peek(plan["token"], "optimize").operation, + "optimize", + ) + restarted = PromptManagerAPI(db=PromptDatabase(self.api.db.model.db_path)) + try: + with self.assertRaises(PlanExpired): + restarted._maintenance_plans.peek(plan["token"], "optimize") + finally: + restarted.db.close_all() + + +class TestLegacyMaintenance(MaintenanceAPITestCase): + async def test_vacuum_failure_reported_not_success(self): + with patch( + "database.models.sqlite3.connect", + side_effect=sqlite3.OperationalError("database is locked"), + ): + with self.assertLogs("prompt_manager", "ERROR"): + reply = await self.post( + "/prompt_manager/maintenance", {"operations": ["vacuum"]}, 200 + ) + self.assertFalse(reply["all_successful"]) + self.assertFalse(reply["results"]["vacuum"]["success"]) + self.assertEqual(reply["results"]["vacuum"]["error"], "database is locked") + + async def test_legacy_orphaned_image_cleanup_skips_unavailable_sources(self): + unmounted = Path(self.tmpdir) / "unmounted" + unmounted.mkdir() + GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir), str(unmounted)] + prompt = self.save("Links") + self.image(prompt, "gone.png", exists=False) + away = self.image(prompt, "away.png", folder=unmounted, exists=False) + unmounted.rmdir() + reply = await self.post( + "/prompt_manager/maintenance", + {"operations": ["cleanup_orphaned_images"]}, + 200, + ) + self.assertEqual( + reply["results"]["cleanup_orphaned_images"], + { + "success": True, + "removed_count": 1, + "message": "Removed 1 orphaned image records", + }, + ) + self.assertIsNotNone(self.api.db.get_image_by_id(away)) + + +if __name__ == "__main__": # pragma: no cover + import unittest + + unittest.main() diff --git a/tests/test_database.py b/tests/test_database.py index bcc8498d..92935ce1 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -1374,13 +1374,17 @@ class TestModelHousekeeping(DatabaseTestCase): ): self.assertEqual(self.db.model.get_database_info(), {}) - def test_vacuum_logs_errors_instead_of_raising(self): + def test_vacuum_logs_and_raises_its_failure(self): + # S6: a failed VACUUM is reported (legacy /maintenance success:false), + # never logged as done. self._save("x") self.db.model.vacuum_database() with patch( "database.models.sqlite3.connect", side_effect=sqlite3.OperationalError ): - self.db.model.vacuum_database() # logged, not raised + with self.assertLogs("prompt_manager", "ERROR"): + with self.assertRaises(sqlite3.OperationalError): + self.db.model.vacuum_database() self.assertEqual(len(self.db.search_prompts()), 1) def test_close_tolerates_a_connection_that_fails_to_close(self): diff --git a/tests/test_maintenance_plans.py b/tests/test_maintenance_plans.py new file mode 100644 index 00000000..b5efe347 --- /dev/null +++ b/tests/test_maintenance_plans.py @@ -0,0 +1,443 @@ +"""Maintenance plans (#249 System S6): the plan store and the read-only previews. + +``PlanStore`` holds each reviewed plan server-side under a single-use token +for ``ttl`` seconds, within a count and byte budget it never truncates to fit. +``PromptDatabase.preview_maintenance`` builds a plan on its own ``mode=ro`` +connection: it writes nothing (the TESTING.md read-only standard) and reuses +the legacy maintenance queries, so a preview counts what the legacy cleanup +would change, except that a missing-file link is only counted under a source +that can be read now. +""" + +import hashlib +import os +import sqlite3 +import sys +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from database import maintenance # noqa: E402 +from database.operations import PromptDatabase # noqa: E402 +from py.api.maintenance_plans import ( # noqa: E402 + PlanExpired, + PlanStore, + PlanUsed, + ScopeTooLarge, +) +from utils.hashing import generate_prompt_hash # noqa: E402 + +try: + from tests import test_execution_history_database as history_tests # noqa: E402 + from tests.open_handles import assert_closed # noqa: E402 + from tests.test_execution_context_database import bundle # noqa: E402 +except ImportError: # discovered with tests/ as the top-level directory + import test_execution_history_database as history_tests # noqa: E402 + from open_handles import assert_closed # noqa: E402 + from test_execution_context_database import bundle # noqa: E402 + +OPERATIONS = ("duplicate_prompts", "optimize", "unused_components", "missing_links") + + +def _digest(path): + if not os.path.exists(path): + return None + with open(path, "rb") as handle: + return hashlib.sha256(handle.read()).hexdigest() + + +class FakeClock: + def __init__(self): + self.now = 1000.0 + + def __call__(self): + return self.now + + +class TestPlanStore(unittest.TestCase): + def test_plan_expiry_budget_and_single_use(self): + clock = FakeClock() + store = PlanStore(ttl=900, max_plans=2, max_bytes=200, clock=clock) + first = store.issue("optimize", "sha256:a", {"targets": []}, ("db", 0)) + self.assertGreaterEqual(len(first), 40) # token_urlsafe(32) + plan = store.peek(first, "optimize") + self.assertEqual( + (plan.operation, plan.fingerprint, plan.db_identity), + ("optimize", "sha256:a", ("db", 0)), + ) + # Single use: the second consume names the job the first one started. + self.assertIs(store.consume(first, "optimize", job_id="job-1"), plan) + with self.assertRaises(PlanUsed) as used: + store.consume(first, "optimize", job_id="job-2") + self.assertEqual(used.exception.job_id, "job-1") + with self.assertRaises(PlanUsed): + store.peek(first, "optimize") + # Another operation's token, an unknown token and a non-string are expired. + second = store.issue("missing_links", "sha256:b", {"targets": [1]}, ("db", 0)) + for token, operation in ( + (second, "optimize"), + ("nope", "missing_links"), + (None, "missing_links"), + ): + with self.subTest(token=token), self.assertRaises(PlanExpired): + store.peek(token, operation) + # Expiry: valid up to ttl, gone after it. + clock.now += 900 + self.assertEqual(store.peek(second, "missing_links").fingerprint, "sha256:b") + clock.now += 0.001 + with self.assertRaises(PlanExpired): + store.consume(second, "missing_links", job_id="job-3") + # Budget: a plan over max_bytes is refused whole, never truncated. + with self.assertRaises(ScopeTooLarge): + store.issue("unused_components", "x", {"targets": list(range(100))}, ()) + # Count budget: the oldest plan makes room for a new one and is expired. + a = store.issue("optimize", "a", {}, ()) + b = store.issue("optimize", "b", {}, ()) + c = store.issue("optimize", "c", {}, ()) + with self.assertRaises(PlanExpired): + store.peek(a, "optimize") + self.assertEqual( + [store.peek(t, "optimize").fingerprint for t in (b, c)], ["b", "c"] + ) + # Byte budget across plans: older plans are dropped until the new one fits. + sized = PlanStore(max_plans=8, max_bytes=100, clock=clock) + pad = {"pad": "x" * 40} # 50 bytes as compact JSON + p1, p2, p3 = (sized.issue("optimize", str(n), pad, ()) for n in range(3)) + with self.assertRaises(PlanExpired): + sized.peek(p1, "optimize") + self.assertEqual([sized.peek(t).fingerprint for t in (p2, p3)], ["1", "2"]) + + def test_expires_at_is_wall_clock(self): + store = PlanStore(ttl=60, clock=FakeClock(), wall_clock=lambda: 5000.0) + token = store.issue("optimize", "f", {}, ()) + self.assertEqual(store.peek(token, "optimize").expires_at, 5060.0) + + def test_discard_forgets_a_plan(self): + store = PlanStore() + token = store.issue("optimize", "f", {}, ()) + store.discard(token) + store.discard("unknown") + with self.assertRaises(PlanExpired): + store.peek(token, "optimize") + + def test_used_tokens_are_remembered_within_a_bound(self): + store = PlanStore(max_plans=2) + tokens = [] + for index in range(store.MAX_USED + 1): + token = store.issue("optimize", str(index), {}, ()) + store.consume(token, "optimize", job_id=f"job-{index}") + tokens.append(token) + with self.assertRaises(PlanExpired): # the oldest record was dropped + store.peek(tokens[0], "optimize") + with self.assertRaises(PlanUsed): + store.peek(tokens[-1], "optimize") + + +class MaintenanceDatabaseCase(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.root = Path(os.path.realpath(self.temp.name)) + self.output = self.root / "output" + self.output.mkdir() + self.db_path = str(self.root / "prompts.db") + self.db = PromptDatabase(self.db_path) + + def tearDown(self): + self.db.close_all() + assert_closed(self, self.root) + self.temp.cleanup() + + def save(self, text, **metadata): + return self.db.save_prompt( + text, prompt_hash=generate_prompt_hash(text), **metadata + ) + + def image(self, prompt_id, name, *, exists=True, folder=None): + path = Path(folder or self.output) / name + path.parent.mkdir(parents=True, exist_ok=True) + if exists: + path.write_bytes(b"png") + return self.db.link_image_to_prompt(prompt_id, str(path)) + + def duplicate(self, text): + """A second prompt row whose text differs only in case (legacy duplicate).""" + with self.db.model.get_connection() as conn: + cursor = conn.execute( + "INSERT INTO prompts(text, created_at) VALUES(?, '2999-01-01')", + (text,), + ) + return cursor.lastrowid + + def preview(self, operation, roots=None): + roots = [str(self.output)] if roots is None else roots + return self.db.preview_maintenance(operation, roots) + + def ids(self, sql, *args): + with self.db.model.get_connection() as conn: + return [row[0] for row in conn.execute(sql, args)] + + +class TestPreviews(MaintenanceDatabaseCase): + def populate(self): + kept = self.save("Kept prompt") + self.image(kept, "kept.png") + self.image(kept, "gone.png", exists=False) + self.save("Unused prompt") + self.duplicate("KEPT PROMPT") + return kept + + def test_preview_writes_nothing(self): + self.populate() + # The app's own connection stays open, as in the running server. + self.db.get_statistics() + files = [self.db_path, self.db_path + "-wal"] + real_connect = sqlite3.connect + + class Recording(sqlite3.Connection): + changes_at_close = [] + + def close(self): + Recording.changes_at_close.append(self.total_changes) + super().close() + + opened = [] + + def connect(*args, **kwargs): + conn = real_connect(*args, factory=Recording, **kwargs) + opened.append((args, kwargs, conn)) + return conn + + for operation in OPERATIONS: + with self.subTest(operation=operation): + before = [_digest(path) for path in files] + shm_size = os.path.getsize(self.db_path + "-shm") + media = { + name: _digest(self.output / name) + for name in os.listdir(self.output) + } + prompts = self.ids("SELECT id FROM prompts ORDER BY id") + images = self.ids("SELECT id FROM generated_images ORDER BY id") + opened.clear() + Recording.changes_at_close.clear() + with patch.object(maintenance.sqlite3, "connect", side_effect=connect): + draft = self.preview(operation) + self.assertTrue(draft.fingerprint.startswith("sha256:")) + self.assertEqual(len(opened), 1) + args, kwargs, conn = opened[0] + self.assertTrue(str(args[0]).endswith("?mode=ro"), args) + self.assertTrue(kwargs.get("uri")) + self.assertEqual(Recording.changes_at_close, [0]) + with self.assertRaises(sqlite3.ProgrammingError): # closed + conn.execute("SELECT 1") + self.assertEqual([_digest(path) for path in files], before) + self.assertEqual(os.path.getsize(self.db_path + "-shm"), shm_size) + self.assertEqual( + { + name: _digest(self.output / name) + for name in os.listdir(self.output) + }, + media, + ) + self.assertEqual( + self.ids("SELECT id FROM prompts ORDER BY id"), prompts + ) + self.assertEqual( + self.ids("SELECT id FROM generated_images ORDER BY id"), images + ) + + def test_unknown_operation_raises(self): + with self.assertRaises(ValueError): + self.db.preview_maintenance("reboot", []) + + def test_same_library_gives_the_same_fingerprint(self): + self.populate() + for operation in OPERATIONS: + with self.subTest(operation=operation): + self.assertEqual( + self.preview(operation).fingerprint, + self.preview(operation).fingerprint, + ) + + def test_preview_details_are_capped_and_carry_no_server_path(self): + prompt = self.save("Many") + for index in range(60): + self.image(prompt, f"sub/gone-{index}.png", exists=False) + draft = self.preview("missing_links") + self.assertEqual(draft.scope["count"], 60) + self.assertEqual(len(draft.details), maintenance.MAX_DETAILS) + self.assertEqual(draft.details[0]["label"], "sub/gone-0.png") + self.assertEqual(draft.details[0]["detail"], "output · the file is missing") + self.assertNotIn(str(self.root), str(draft.public())) + self.assertEqual(len(draft.targets), 60) + + def test_duplicate_prompts_plan_counts_the_legacy_groups(self): + kept = self.save("A cat") + self.duplicate("A CAT") + self.duplicate("a cat") + paired = self.db.get_prompt_pair(self.db.save_prompt_pair("a dog", "blur")) + self.duplicate("A DOG") + draft = self.preview("duplicate_prompts") + self.assertEqual(draft.scope["count"], 2) + self.assertEqual(draft.scope["label"], "2 duplicate prompts in 1 group") + self.assertEqual(draft.protected["count"], 2) + self.assertEqual( + draft.protected["reasons"], ["1 group belongs to a saved pair"] + ) + self.assertEqual(draft.targets[0][0], kept) + self.assertEqual(len(draft.targets), 1) + self.assertEqual( + [row["detail"] for row in draft.details], + ["Keeps the oldest of 3 copies", "Kept: belongs to a saved pair"], + ) + self.assertEqual(draft.details[0]["label"], "A cat") + self.assertIsNotNone(paired) + + def test_unused_components_plan_protects_pairs_history_and_titled_parents(self): + orphan = self.save("Orphan") + pair = self.db.get_prompt_pair(self.db.save_prompt_pair("paired", "neg")) + titled = self.save("Titled parent") + untitled = self.save("Explicit parent") + protected = self.save("Protected", tags=["__protected__"]) + with self.db.model.get_connection() as conn: + conn.execute( + "INSERT INTO prompt_library_state(prompt_id,title) " + "VALUES(?,?),(?,NULL)", + (titled, "My recipe", untitled), + ) + hinted = self.save("Recorded run") + draft_before_hint = self.preview("unused_components") + self.assertEqual(sorted(draft_before_hint.targets), sorted([orphan, hinted])) + self.record_run_hint(hinted) + draft = self.preview("unused_components") + self.assertEqual(draft.targets, [orphan]) + self.assertEqual(draft.scope["count"], 1) + self.assertEqual(draft.scope["label"], "1 unused prompt component") + self.assertEqual(draft.protected["count"], 6) + self.assertEqual( + draft.protected["reasons"], + [ + "2 in saved pairs", + "1 tagged __protected__", + "2 with Library state (a title or an explicit parent)", + "1 used by a recorded run", + ], + ) + self.assertEqual( + draft.details, + [ + { + "label": "Orphan", + "detail": maintenance.UNUSED_DETAIL, + } + ], + ) + # The plan counts exactly what the legacy prune removes. + self.assertIsNotNone(pair) + self.assertEqual(protected > 0, True) + self.assertEqual(self.db.prune_orphaned_prompts(), 1) + self.assertEqual(self.preview("unused_components").targets, []) + + def record_run_hint(self, prompt_id, run=0): + """Make ``prompt_id`` the live positive hint of a recorded run.""" + text = self.db.get_prompt_by_id(prompt_id)["text"] + history = history_tests.HistoryDatabaseTests + start = history.start(self, self.db, run) + self.db.write_execution_start(start) + self.db.write_execution_finish(history.finish(self, start)) + self.assertTrue( + self.db.write_execution_contexts( + start.run_id, start.storage_epoch, bundle(positive=text, negative="") + ) + ) + + def test_missing_links_plan_skips_unavailable_roots(self): + prompt = self.save("Links") + self.image(prompt, "here.png") + gone = self.image(prompt, "gone.png", exists=False) + unmounted = self.root / "unmounted" + self.image(prompt, "away.png", folder=unmounted) + os.remove(unmounted / "away.png") + os.rmdir(unmounted) # the drive went away: its links look missing + self.image( + prompt, "elsewhere.png", folder=self.root / "elsewhere", exists=False + ) + draft = self.preview("missing_links", [str(self.output), str(unmounted)]) + self.assertEqual(draft.targets, [gone]) + self.assertEqual(draft.scope["count"], 1) + self.assertEqual(draft.scope["label"], "1 link to a missing file") + self.assertEqual(draft.protected["count"], 2) + self.assertEqual( + draft.protected["reasons"], + [ + "1 link under an unavailable source", + "1 link outside the image sources", + ], + ) + self.assertFalse(draft.media["deletes_files"]) + # No readable source: nothing is counted, every link is protected. + none = self.preview("missing_links", [str(unmounted)]) + self.assertEqual((none.targets, none.scope["count"]), ([], 0)) + self.assertEqual(none.protected["count"], 4) + + def test_missing_links_plan_matches_the_health_check(self): + prompt = self.save("Links") + self.image(prompt, "a.png", exists=False) + self.image(prompt, "b.png", exists=False) + self.image(prompt, "c.png") + report = self.db.run_health_checks([str(self.output)]) + missing = next(c for c in report["checks"] if c["id"] == "missing_files") + self.assertEqual(self.preview("missing_links").scope["count"], missing["count"]) + + def test_optimize_plan_reports_sizes_without_paths(self): + self.populate() + draft = self.preview("optimize") + self.assertEqual(draft.targets, []) + self.assertEqual(draft.scope, {"count": 1, "label": "The whole database"}) + self.assertEqual( + [row["label"] for row in draft.details], + ["Database file", "Free pages", "Free space needed"], + ) + self.assertNotIn(str(self.root), str(draft.public())) + + def test_legacy_cleanup_missing_images_skips_unavailable_roots(self): + prompt = self.save("Links") + gone = self.image(prompt, "gone.png", exists=False) + unmounted = self.root / "unmounted" + away = self.image(prompt, "away.png", folder=unmounted, exists=False) + outside = self.image(prompt, "x.png", folder=self.root / "other", exists=False) + unmounted.rmdir() + self.assertEqual( + self.db.cleanup_missing_images(roots=[str(self.output), str(unmounted)]), 1 + ) + remaining = self.ids("SELECT id FROM generated_images ORDER BY id") + self.assertEqual(remaining, sorted([away, outside])) + self.assertNotIn(gone, remaining) + # Without roots it keeps the legacy rule: every link whose file is gone. + self.assertEqual(self.db.cleanup_missing_images(), 2) + + def test_reader_closes_its_connection_when_a_read_raises(self): + self.populate() + with patch.object( + maintenance, "_orphaned_ids", side_effect=sqlite3.OperationalError("boom") + ): + with self.assertRaises(sqlite3.OperationalError): + self.preview("unused_components") + # The teardown's assert_closed proves the read-only handle was closed. + + +class TestVacuum(MaintenanceDatabaseCase): + def test_vacuum_failure_raises(self): + with patch( + "database.models.sqlite3.connect", + side_effect=sqlite3.OperationalError("disk I/O error"), + ): + with self.assertLogs("prompt_manager", "ERROR"): + with self.assertRaises(sqlite3.OperationalError): + self.db.model.vacuum_database() + + +if __name__ == "__main__": + unittest.main()