feat(maintenance): reviewed, read-only maintenance plans
POST /prompt_manager/maintenance/preview reviews duplicate_prompts, optimize, unused_components or missing_links on its own mode=ro snapshot and answers a Plan (scope, protected reasons, backup policy, media, at most 50 detail rows). The fingerprint, targets and judged sources stay in a server-side PlanStore under a single-use token (15 min, 8 plans, 16 MiB, never truncated; a restart forgets them), bound to the database it read. The previews reuse the legacy queries: duplicate groups and the merge are shared with cleanup_duplicates, ORPHANED_PROMPTS_SQL is built from the #239 protection predicates the plan reports, and missing links are judged by the health check's link_states. Legacy /maintenance keeps its shape; vacuum now reports success:false when VACUUM fails, and cleanup_orphaned_images skips links under unavailable sources.
This commit is contained in:
+37
@@ -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
|
||||
|
||||
+32
-6
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
},
|
||||
],
|
||||
)
|
||||
@@ -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:
|
||||
"""
|
||||
|
||||
+37
-104
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
|
||||
+3
-1
@@ -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()
|
||||
|
||||
@@ -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_<date and time>",
|
||||
},
|
||||
},
|
||||
headers={"Cache-Control": "no-store"},
|
||||
)
|
||||
@@ -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]
|
||||
@@ -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_<date and time>",
|
||||
},
|
||||
)
|
||||
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()
|
||||
@@ -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):
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user