nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(backup): review a restore, run it as an Activity job, verify backups

Browse Source 
- POST /restore/preview stages an upload in the library's folder and
  answers what it holds beside the live library, writing nothing live;
  DELETE discards it; at most two stages are held (429 too-large).
- POST /restore/apply rechecks the live fingerprint (409 plan-changed) and
  starts "Database restore": exclusive, not cancellable, indeterminate. It
  waits for request leases, publishes the exact staged bytes, swaps the
  database, prunes safety backups and logs counts, repairs and the backup.
- The legacy /restore runs the same stage and publish as an exclusive hold
  with its old response, so a restore and the health check, Move library
  and every other job refuse each other.
- JobRegistry.lease: Needs review decisions and maintenance and restore
  previews are refused while an exclusive job runs, and a restore waits
  for them before it switches the file.
- GET /backup holds a guarded slot, verifies a standalone snapshot before
  streaming it and removes the closed temp file; GET /backup/status lists
  the safety backups on disk and this session's verified download.
This commit is contained in:
Vito Sansevero
2026-10-25 06:20:12 -07:00
parent b9bafbbfa0
commit f193a3a1ce
11 changed files with 1656 additions and 241 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/api/__init__.py
+6
View File
@@ -56,6 +56,7 @@ from .thumbnail_auto import AutoThumbnailMixin
from .system import SystemRoutesMixin
from .library_location import LibraryLocationRoutesMixin
from .maintenance import MaintenanceRoutesMixin
from .backup_restore import BackupRestoreRoutesMixin
from .workflow_storage import WorkflowStorageRoutesMixin
from .duplicate_files import DuplicateFilesMixin
from .maintenance_plans import PlanStore
@@ -302,6 +303,7 @@ class PromptManagerAPI(
SystemRoutesMixin,
LibraryLocationRoutesMixin,
MaintenanceRoutesMixin,
BackupRestoreRoutesMixin,
WorkflowStorageRoutesMixin,
DuplicateFilesMixin,
LoggingRoutesMixin,
@@ -376,7 +378,10 @@ class PromptManagerAPI(
try:
from .maintenance import sweep_safety_temps
from .backup_restore import restore_stage
swept = sweep_safety_temps(self.db.model.db_path)
swept += restore_stage.sweep_restore_temps(self.db.model.db_path)
if swept:
self.logger.info(f"Removed {len(swept)} interrupted backup copies")
except Exception as e:
@@ -547,6 +552,7 @@ class PromptManagerAPI(
self._register_system_routes(routes)
self._register_library_location_routes(routes)
self._register_maintenance_routes(routes)
self._register_backup_restore_routes(routes)
self._register_duplicate_routes(routes)
self._register_logging_routes(routes)
self._register_autotag_routes(routes)
py/api/admin.py
-215
View File
@@ -3,14 +3,12 @@
import asyncio
import contextlib
from contextlib import closing
import datetime
import functools
import hashlib
import json
import math
import os
import sqlite3
import tempfile
import threading
import time
from pathlib import Path
@@ -27,7 +25,6 @@ from .jobs import (
DUPLICATE_CLEANUP_JOB_KIND,
DUPLICATES_JOB_KIND,
MAINTENANCE_JOB_KIND,
RESTORE_JOB_KIND,
JobBusy,
)
@@ -37,7 +34,6 @@ except ImportError:
from utils.parallel import map_parallel
try:
from ...database.operations import PromptDatabase
from ...database.pair_ingestion import ingest_resolved_pair
from ...utils.generation_params import (
PLACEHOLDERS,
@@ -48,7 +44,6 @@ try:
from ...utils.prompt_pair_resolver import resolve_prompt_pair
from ...utils.review_capture import reviewable
except ImportError:
from database.operations import PromptDatabase
from database.pair_ingestion import ingest_resolved_pair
from utils.generation_params import (
PLACEHOLDERS,
@@ -1236,216 +1231,6 @@ class AdminRoutesMixin:
"message": f"Found {len(issues)} consistency issues",
}
async def backup_database(self, request):
"""Download a consistent copy of the prompts database (WAL included)."""
try:
model = self.db.model
db_path = model.db_path
if not os.path.exists(db_path):
return web.json_response(
{"success": False, "error": "Database file not found"}, status=404
)
fd, temp_path = tempfile.mkstemp(suffix=".db")
os.close(fd)
try:
ok = await self._run_in_executor(model.backup_database, temp_path)
if not ok:
return web.json_response(
{"success": False, "error": "Failed to create database backup"},
status=500,
)
# Read once (the file is deleted right after) rather than
# streaming so the temp file never outlives the request.
with open(temp_path, "rb") as fh:
file_data = fh.read()
finally:
if os.path.exists(temp_path):
os.unlink(temp_path)
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
filename = f"prompts_backup_{timestamp}.db"
return web.Response(
body=file_data,
content_type="application/octet-stream",
headers={
"Content-Disposition": f'attachment; filename="{filename}"',
"Content-Length": str(len(file_data)),
},
)
except Exception as e:
self.logger.error(f"Backup error: {e}", exc_info=True)
return web.json_response(
{
"success": False,
"error": f"Failed to backup database: {self._public_error(e)}",
},
status=500,
)
async def restore_database(self, request):
"""Restore the prompts database from an uploaded SQLite file.
The upload is streamed to a temp file. Free disk space decides how
large it may be (``maintenance.restore_room``), under a 64 GiB ceiling
(``restore_max_bytes``): a request whose Content-Length does not fit
is refused before anything is written, and one without is refused as
soon as its running size does not. The file is
verified with PromptModel.verify_database_file and only then copied
over the live database, which is backed up first. The whole request
holds the registry (records): it is refused while a scan or another
record-changing job runs, and those are refused while it runs.
"""
return await self._run_held(
RESTORE_JOB_KIND,
"Database restore (admin page)",
"records",
lambda: self._restore_database(request),
)
def _restore_refusal(self, size, ceiling, room, *, known):
"""``400 too-large`` past the ceiling, else ``507 insufficient-space``."""
from .maintenance import restore_space_message
if size > ceiling:
return web.json_response(
{
"success": False,
"code": "too-large",
"error": (
"File too large. Maximum size is "
f"{ceiling // (1024 * 1024)}MB"
),
},
status=400,
)
return web.json_response(
{
"success": False,
"code": "insufficient-space",
"error": restore_space_message(size, room, known=known),
},
status=507,
)
async def _restore_database(self, request):
from .maintenance import RESTORE_MAX_BYTES, restore_room
ceiling = getattr(self, "restore_max_bytes", RESTORE_MAX_BYTES)
temp_path = None
try:
# Free disk space decides how large a backup may be (#313): the
# upload, its staged copy and the swap must all fit.
room = await self._run_in_executor(restore_room, self.db.model.db_path)
max_bytes = min(ceiling, room)
length = request.content_length
if length is not None and length > max_bytes:
return self._restore_refusal(length, ceiling, room, known=True)
reader = await request.multipart()
field = await reader.next()
if not field or field.name != "database_file":
return web.json_response(
{
"success": False,
"error": (
"No database file uploaded. "
"Expected field name: database_file"
),
},
status=400,
)
fd, temp_path = tempfile.mkstemp(suffix=".db")
os.close(fd)
size = 0
with open(temp_path, "wb") as out:
while True:
chunk = await field.read_chunk()
if not chunk:
break
size += len(chunk)
if size > max_bytes:
return self._restore_refusal(size, ceiling, room, known=False)
out.write(chunk)
if size == 0:
return web.json_response(
{"success": False, "error": "Uploaded file is empty"}, status=400
)
model = self.db.model
ok, reason = await self._run_in_executor(
model.verify_database_file, temp_path
)
if ok:
reason = await self._run_in_executor(_inspect_restore_upload, temp_path)
ok = reason is None
if not ok:
return web.json_response(
{"success": False, "error": reason}, status=400
)
try:
backup_path = await self._run_in_executor(
model.restore_from_file, temp_path
)
except sqlite3.OperationalError as exc:
if (
getattr(exc, "_prompt_manager_restore_backup_created", False)
is True
):
try:
await self._run_in_executor(
_prune_safety_backups, model.db_path
)
except Exception:
self.logger.warning(
"Could not prune safety backups after failed restore"
)
raise
removed = await self._run_in_executor(_prune_safety_backups, model.db_path)
if removed:
self.logger.info(f"Removed {len(removed)} old safety backups")
previous = self.db
self.db = PromptDatabase(model.db_path)
self._database_replaced() # a health report described the old one
# Its connections would otherwise linger (a file lock on Windows)
previous.close_all()
prompt_count = self.db.model.get_database_info().get("total_prompts", 0)
return web.json_response(
{
"success": True,
"message": (
"Database restored successfully. "
f"Found {prompt_count} prompts."
),
"prompt_count": prompt_count,
"backup_created": (
self._public_path(backup_path) if backup_path else None
),
}
)
except Exception as e:
self.logger.error(f"Restore error: {e}", exc_info=True)
return web.json_response(
{
"success": False,
"error": f"Failed to restore database: {self._public_error(e)}",
},
status=500,
)
finally:
if temp_path:
for suffix in ("", "-wal", "-shm", "-journal"):
if os.path.exists(temp_path + suffix):
os.unlink(temp_path + suffix)
async def scan_images(self, request):
"""Scan ComfyUI output images for prompt metadata and add them to the database.
py/api/backup_restore.py
+714
View File
@@ -0,0 +1,714 @@
"""Database backup and restore routes (#249 System S8).
Legacy routes, responses unchanged:
- ``GET /prompt_manager/backup`` downloads a consistent copy of the library.
It holds the registry as ``backup`` ("Database backup", ``guarded``): it is
refused while a restore, maintenance apply or other record-changing job
runs. The copy is written in the library's folder, made one standalone
file and verified (integrity and schema) before a byte is sent; a copy that
does not verify is a 500 with no bytes. It is streamed from the open file,
then closed and removed.
- ``POST /prompt_manager/restore`` (multipart ``database_file``) restores in
one request through the same stage and publish as the restore job, holding
the registry as ``restore`` (``exclusive``).
Reviewed restore (new routes answer errors as ``{success:false, code, error}``):
- ``POST /prompt_manager/restore/preview`` (multipart ``database_file``, the
same disk-space cap) stages the upload in the library's folder
(``database/restore_stage.py``) and answers what it holds beside the live
library: ``{token, expires_at, summary, schema:{migrated, library_state},
integrity:"ok", documents:{repaired, cleared, images}, live, size,
safety_backup:{destination_label}}``. Nothing live is written or backed
up. At most two previews are held (``too-large``, 429).
- ``DELETE /prompt_manager/restore/preview/{token}`` discards a stage.
- ``POST /prompt_manager/restore/apply {token}`` starts the restore job (202
``{job}``): kind ``restore``, "Database restore", ``exclusive`` (it runs
alone, like Move library), not cancellable, indeterminate. The runner waits
for request leases to end, rechecks the live fingerprint (``plan-changed``),
publishes the exact staged bytes, swaps ``self.db``, prunes old safety
backups and logs what was restored.
- ``GET /prompt_manager/backup/status``: the safety backups on disk
(``<db>.backup_*``, newest first, names only) and this session's last
verified download (Q9: no sidecar file).
"""
import datetime
import os
import sqlite3
import tempfile
from aiohttp import web
from . import admin as _admin
from .job_routes import busy_response
from .jobs import RESTORE_JOB_KIND, JobBusy
from .maintenance_plans import PlanExpired, PlanUsed
from .restore_plans import RestorePlanStore
try:
from ...database import restore_stage
from ...database.operations import PromptDatabase
from ...database.safety_backup import (
SafetyBackupFailed,
fsync_file,
make_standalone,
remove_quietly,
safety_temp,
)
except ImportError:
from database import restore_stage
from database.operations import PromptDatabase
from database.safety_backup import (
SafetyBackupFailed,
fsync_file,
make_standalone,
remove_quietly,
safety_temp,
)
BACKUP_JOB_KIND = "backup"
BACKUP_HOLD_TITLE = "Database backup"
LEGACY_RESTORE_TITLE = "Database restore (admin page)"
RESTORE_JOB_TITLE = "Database restore"
RESTORE_GROUPS = ("errors",)
MAX_TOKEN_LENGTH = 128
APPLY_BODY_MAX = 1024
STREAM_CHUNK = 1024 * 1024
MAX_LISTED_BACKUPS = 10
SIDECARS = ("", "-wal", "-shm", "-journal")
REQUESTS_BUSY = (
"Other requests are still reading the library; nothing was changed. "
"Try again in a moment."
)
PLAN_CHANGED = (
"The library changed since the review; nothing was changed. "
"Review the restore again."
)
NOT_RESTORED = "The restore failed; the library was not changed."
def _failure(message, status, code=None, **extra):
return web.json_response(
{"success": False, "code": code, "error": message, **extra}, status=status
)
def _iso(timestamp):
return datetime.datetime.fromtimestamp(timestamp, datetime.timezone.utc).isoformat()
def _plural(count, word):
return f"{count:,} {word}" + ("" if count == 1 else "s")
def documents_line(documents):
"""The repair report's sentence, or "" when nothing was damaged."""
repaired, cleared = documents.get("repaired", 0), documents.get("cleared", 0)
parts = []
if repaired:
parts.append(f"repaired {_plural(repaired, 'damaged workflow document')}")
if cleared:
images = documents.get("images", 0)
parts.append(
f"could not repair {_plural(cleared, 'workflow document')}: "
f"{_plural(images, 'image')} {'shows' if images == 1 else 'show'} "
"no workflow"
)
if not parts:
return ""
text = "; ".join(parts)
return text[0].upper() + text[1:] + "."
def restored_line(summary):
return (
f"Restored {_plural(summary['prompts'], 'prompt')}, "
f"{_plural(summary['image_links'], 'image link')}, "
f"{_plural(summary['pairs'], 'saved pair')} and "
f"{_plural(summary['runs'], 'run')}"
)
def safety_backups(db_path, limit=MAX_LISTED_BACKUPS):
"""``<db>.backup_*`` files beside the database, newest first (blocking)."""
db_path = os.path.abspath(db_path)
folder = os.path.dirname(db_path)
prefix = os.path.basename(db_path) + ".backup_"
found = []
try:
names = os.listdir(folder)
except OSError:
return []
for name in names:
if not name.startswith(prefix):
continue
try:
stat = os.stat(os.path.join(folder, name))
except OSError:
continue
if not os.path.isfile(os.path.join(folder, name)):
continue
found.append((stat.st_mtime, name, stat.st_size))
found.sort(key=lambda item: (item[0], item[1]), reverse=True)
return [
{"name": name, "bytes": size, "modified_at": _iso(mtime)}
for mtime, name, size in found[:limit]
]
def write_backup_snapshot(model):
"""A verified, standalone copy of the library beside it; its path (blocking).
Raises SafetyBackupFailed with nothing left on disk when the copy cannot
be written or does not verify. The caller removes the returned file.
"""
temp = safety_temp(model.db_path)
try:
if not model.backup_database(temp):
raise SafetyBackupFailed("The backup could not be written")
make_standalone(temp)
ok, reason = model.verify_database_file(temp, documents=False)
if not ok:
raise SafetyBackupFailed("The backup did not verify", reason)
fsync_file(temp)
return temp
except BaseException:
remove_quietly(*(temp + suffix for suffix in SIDECARS))
raise
class BackupRestoreRoutesMixin:
"""Mixin providing the backup and restore endpoints."""
def _register_backup_restore_routes(self, routes):
@routes.post("/prompt_manager/restore/preview")
async def restore_preview_route(request):
return await self.preview_restore(request)
@routes.delete("/prompt_manager/restore/preview/{token}")
async def restore_discard_route(request):
return await self.discard_restore_preview(request)
@routes.post("/prompt_manager/restore/apply")
async def restore_apply_route(request):
return await self.apply_restore(request)
@routes.get("/prompt_manager/backup/status")
async def backup_status_route(request):
return await self.backup_status(request)
@property
def _restore_plans(self):
store = getattr(self, "_restore_plan_store", None)
if store is None:
store = self._restore_plan_store = RestorePlanStore()
return store
# ── backup ────────────────────────────────────────────────────────
async def backup_database(self, request):
"""Download a verified copy of the prompts database (WAL included)."""
return await self._run_held(
BACKUP_JOB_KIND,
BACKUP_HOLD_TITLE,
"guarded",
lambda: self._stream_backup(request),
)
async def _stream_backup(self, request):
model = self.db.model
if not os.path.exists(model.db_path):
return web.json_response(
{"success": False, "error": "Database file not found"}, status=404
)
try:
temp = await self._run_in_executor(write_backup_snapshot, model)
except SafetyBackupFailed as failed:
self.logger.error(f"Backup failed: {failed} {failed.reason or ''}")
return web.json_response(
{"success": False, "error": "Failed to create database backup"},
status=500,
)
except Exception as e:
self.logger.error(f"Backup error: {e}", exc_info=True)
return web.json_response(
{
"success": False,
"error": f"Failed to backup database: {self._public_error(e)}",
},
status=500,
)
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
handle = None
try:
size = os.path.getsize(temp)
handle = open(temp, "rb")
response = web.StreamResponse(
headers={
"Content-Type": "application/octet-stream",
"Content-Disposition": (
f'attachment; filename="prompts_backup_{timestamp}.db"'
),
"Content-Length": str(size),
}
)
await response.prepare(request)
while True:
chunk = await self._run_in_executor(handle.read, STREAM_CHUNK)
if not chunk:
break
await response.write(chunk)
await response.write_eof()
self._last_backup_download = {
"at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
"bytes": size,
"verified": True,
}
return response
finally:
if handle is not None:
handle.close() # before the unlink (Windows)
await self._run_in_executor(
remove_quietly, *(temp + suffix for suffix in SIDECARS)
)
async def backup_status(self, request):
try:
backups = await self._run_in_executor(safety_backups, self.db.model.db_path)
except Exception:
self.logger.exception("Backup status failed")
return _failure("Failed to read the backup status", 500)
return web.json_response(
{
"success": True,
"safety_backups": backups,
"last_download": getattr(self, "_last_backup_download", None),
},
headers={"Cache-Control": "no-store"},
)
# ── the upload ────────────────────────────────────────────────────
def _restore_refusal(self, size, ceiling, room, *, known):
"""``400 too-large`` past the ceiling, else ``507 insufficient-space``."""
from .maintenance import restore_space_message
if size > ceiling:
return web.json_response(
{
"success": False,
"code": "too-large",
"error": (
"File too large. Maximum size is "
f"{ceiling // (1024 * 1024)}MB"
),
},
status=400,
)
return web.json_response(
{
"success": False,
"code": "insufficient-space",
"error": restore_space_message(size, room, known=known),
},
status=507,
)
async def _receive_restore_upload(self, request, uploads):
"""Stream ``database_file`` into the library's folder.
Returns ``(path, None)`` or ``(None, refusal)``. Each temp path is
appended to ``uploads`` as soon as it exists, so the caller removes
it (and its sidecars) whatever happens.
"""
from .maintenance import RESTORE_MAX_BYTES, restore_room
ceiling = getattr(self, "restore_max_bytes", RESTORE_MAX_BYTES)
db_path = self.db.model.db_path
# Free disk space decides how large a backup may be (#313): the
# upload, its staged copy and the swap must all fit.
room = await self._run_in_executor(restore_room, db_path)
max_bytes = min(ceiling, room)
length = request.content_length
if length is not None and length > max_bytes:
return None, self._restore_refusal(length, ceiling, room, known=True)
reader = await request.multipart()
field = await reader.next()
if not field or field.name != "database_file":
return None, web.json_response(
{
"success": False,
"error": (
"No database file uploaded. "
"Expected field name: database_file"
),
},
status=400,
)
fd, temp_path = tempfile.mkstemp(
prefix=restore_stage.UPLOAD_PREFIX,
suffix=".tmp",
dir=os.path.dirname(os.path.abspath(db_path)),
)
uploads.append(temp_path)
size = 0
with os.fdopen(fd, "wb") as out:
while True:
chunk = await field.read_chunk()
if not chunk:
break
size += len(chunk)
if size > max_bytes:
return None, self._restore_refusal(size, ceiling, room, known=False)
out.write(chunk)
if size == 0:
return None, web.json_response(
{"success": False, "error": "Uploaded file is empty"}, status=400
)
return temp_path, None
async def _upload_problem(self, path):
"""None when the upload may be staged, else a readable reason."""
ok, reason = await self._run_in_executor(
self.db.model.verify_database_file, path, documents=False
)
if ok:
reason = await self._run_in_executor(_admin._inspect_restore_upload, path)
ok = reason is None
return None if ok else reason
async def _remove_uploads(self, uploads):
for path in uploads:
await self._run_in_executor(
remove_quietly, *(path + suffix for suffix in SIDECARS)
)
async def _prune_after(self, db_path, *, failed=False):
try:
removed = await self._run_in_executor(_admin._prune_safety_backups, db_path)
except Exception:
if not failed:
raise
self.logger.warning("Could not prune safety backups after failed restore")
return
if removed:
self.logger.info(f"Removed {len(removed)} old safety backups")
async def _publish(self, model, staged):
"""Publish ``staged``; on a page-copy failure prune and re-raise."""
try:
return await self._run_in_executor(model.publish_staged, staged)
except sqlite3.OperationalError as exc:
if getattr(exc, "_prompt_manager_restore_backup_created", False) is True:
await self._prune_after(model.db_path, failed=True)
raise
def _swap_database(self, db_path):
"""Open the restored library as ``self.db``; the previous one to close."""
previous = self.db
self.db = PromptDatabase(db_path)
self._database_replaced() # a health report described the old one
self._restore_plans.discard_all() # stages reviewed against the old one
return previous
# ── legacy one-step restore ───────────────────────────────────────
async def restore_database(self, request):
"""Restore the prompts database from an uploaded SQLite file.
The upload is streamed to a temp file in the library's folder. Free
disk space decides how large it may be (``maintenance.restore_room``),
under a 64 GiB ceiling (``restore_max_bytes``): a request whose
Content-Length does not fit is refused before anything is written,
and one without is refused as soon as its running size does not. The
file is checked, staged, repaired and verified, then published over
the live database, which is backed up first. The whole request holds
the registry as ``restore`` (``exclusive``): it is refused while any
job runs, and every job and request lease is refused while it runs.
"""
return await self._run_held(
RESTORE_JOB_KIND,
LEGACY_RESTORE_TITLE,
"exclusive",
lambda: self._restore_database(request),
)
async def _restore_database(self, request):
uploads = []
try:
temp_path, refused = await self._receive_restore_upload(request, uploads)
if refused is not None:
return refused
reason = await self._upload_problem(temp_path)
if reason is not None:
return web.json_response(
{"success": False, "error": reason}, status=400
)
if not await self._jobs.drain_leases():
return web.json_response(
{"success": False, "error": REQUESTS_BUSY}, status=409
)
model = self.db.model
staged = await self._run_in_executor(model.stage_restore, temp_path)
try:
backup_path = await self._publish(model, staged)
finally:
await self._run_in_executor(staged.discard)
await self._prune_after(model.db_path)
previous = self._swap_database(model.db_path)
# Its connections would otherwise linger (a file lock on Windows)
await self._run_in_executor(previous.close_all)
prompt_count = self.db.model.get_database_info().get("total_prompts", 0)
message = f"Database restored successfully. Found {prompt_count} prompts."
body = {
"success": True,
"message": message,
"prompt_count": prompt_count,
"backup_created": (
self._public_path(backup_path) if backup_path else None
),
}
repaired = documents_line(staged.documents)
if repaired:
body["message"] = f"{message} {repaired}"
body["workflow_documents"] = dict(staged.documents)
return web.json_response(body)
except Exception as e:
self.logger.error(f"Restore error: {e}", exc_info=True)
return web.json_response(
{
"success": False,
"error": f"Failed to restore database: {self._public_error(e)}",
},
status=500,
)
finally:
await self._remove_uploads(uploads)
# ── reviewed restore ──────────────────────────────────────────────
def _live_fingerprint(self, db):
return restore_stage.live_fingerprint(db.model.get_connection())
async def preview_restore(self, request):
"""Stage an upload for review; nothing live is written (see module)."""
store = self._restore_plans
await self._run_in_executor(store.drop_expired)
if store.full():
return _failure(
"Two restores are already being reviewed. Discard one first.",
429,
"too-large",
)
uploads = []
staged = None
try:
async with self._jobs.lease("Restore preview"):
db = self.db
identity = self._database_identity(db)
temp_path, refused = await self._receive_restore_upload(
request, uploads
)
if refused is not None:
return refused
reason = await self._upload_problem(temp_path)
if reason is not None:
return _failure(reason, 400, "invalid-backup")
try:
staged = await self._run_in_executor(
db.model.stage_restore, temp_path
)
except ValueError as invalid:
return _failure(str(invalid), 400, "invalid-backup")
live, fingerprint = await self._run_in_executor(
self._live_fingerprint, db
)
except JobBusy as busy:
return busy_response(busy)
except Exception:
self.logger.exception("Restore preview failed")
if staged is not None:
await self._run_in_executor(staged.discard)
return _failure("Failed to review the restore", 500)
finally:
await self._remove_uploads(uploads)
if identity != self._database_identity(self.db) or store.full():
await self._run_in_executor(staged.discard)
if store.full():
return _failure(
"Two restores are already being reviewed. Discard one first.",
429,
"too-large",
)
return _failure(
"The database was replaced while the restore was reviewed. "
"Review again.",
409,
"storage-changed",
)
token = store.issue(staged, fingerprint, identity, live)
held = store.peek(token)
name = os.path.basename(db.model.db_path)
return web.json_response(
{
"success": True,
"token": token,
"expires_at": _iso(held.expires_at),
"summary": dict(staged.summary),
"schema": dict(staged.schema),
"integrity": "ok",
"documents": dict(staged.documents),
"live": dict(live),
"size": staged.size,
"safety_backup": {
"destination_label": f"{name}.backup_<date and time>"
},
},
headers={"Cache-Control": "no-store"},
)
async def discard_restore_preview(self, request):
token = request.match_info.get("token", "")
if not 0 < len(token) <= MAX_TOKEN_LENGTH:
return _failure("Unknown restore review.", 410, "plan-expired")
discarded = await self._run_in_executor(self._restore_plans.discard, token)
if not discarded:
return _failure("Unknown restore review.", 410, "plan-expired")
return web.json_response({"success": True})
def _peek_restore(self, token):
try:
return self._restore_plans.peek(token), None
except PlanUsed as used:
job = self._jobs.get(used.job_id)
return None, _failure(
"This restore was already started.",
409,
"plan-used",
job={"id": used.job_id, "title": job.title if job else None},
)
except PlanExpired:
return None, _failure(
"This review expired. Review the restore again.",
410,
"plan-expired",
)
async def apply_restore(self, request):
data, error = await self._read_json_body(request, max_bytes=APPLY_BODY_MAX)
if error is not None:
if error.status == 413:
return _failure("Request body too large", 413, "too-large")
return _failure(
"Request body must be a JSON object", 400, "invalid-request"
)
token = data.get("token")
if set(data) != {"token"} or not isinstance(token, str):
return _failure("token must be the review's token", 400, "invalid-request")
if not 0 < len(token) <= MAX_TOKEN_LENGTH:
return _failure("token must be the review's token", 400, "invalid-request")
await self._run_in_executor(self._restore_plans.drop_expired)
held, refused = self._peek_restore(token)
if refused is not None:
return refused
db = self.db
if held.db_identity != self._database_identity(db):
await self._run_in_executor(self._restore_plans.discard, token)
return _failure(
"The database was replaced since the review. Review again.",
409,
"storage-changed",
)
try:
async with self._jobs.lease("Restore review check"):
_live, fingerprint = await self._run_in_executor(
self._live_fingerprint, db
)
except JobBusy as busy:
return busy_response(busy)
if fingerprint != held.fingerprint:
# It can never match again: free the stage for a new review.
await self._run_in_executor(self._restore_plans.discard, token)
return _failure(PLAN_CHANGED, 409, "plan-changed")
# The token may have been used, expired or discarded meanwhile.
held, refused = self._peek_restore(token)
if refused is not None:
return refused
if db is not self.db:
return _failure(
"The database was replaced since the review. Review again.",
409,
"storage-changed",
)
try:
job = self._jobs.start(
RESTORE_JOB_KIND,
RESTORE_JOB_TITLE,
# It replaces the whole database: it runs alone, beside no
# job of any class, like Move library.
admission="exclusive",
cancellable=False,
scope={},
groups=RESTORE_GROUPS,
runner=lambda job: self._run_restore(job, db, held),
)
except JobBusy as busy:
return busy_response(busy) # the token stays usable
# Same loop turn as the peek and the start: the runner has not run.
self._restore_plans.take(token, job.id)
return web.json_response({"success": True, "job": job.summary()}, status=202)
async def _run_restore(self, job, db, held):
from .maintenance import _fail_job
staged = held.staged
try:
job.update(status_line="Waiting for other requests to finish")
if not await self._jobs.drain_leases():
return _fail_job(job, "busy", REQUESTS_BUSY)
job.update(status_line="Checking the library")
_live, fingerprint = await self._run_in_executor(self._live_fingerprint, db)
if db is not self.db or fingerprint != held.fingerprint:
return _fail_job(job, "plan-changed", PLAN_CHANGED)
job.update(status_line="Restoring the database")
job.log_line("Restoring the reviewed backup")
model = db.model
try:
backup_path = await self._publish(model, staged)
except restore_stage.StageChanged:
return _fail_job(job, "plan-changed", PLAN_CHANGED)
except SafetyBackupFailed:
return _fail_job(
job,
"backup-failed",
"The current library could not be backed up; nothing was "
"changed.",
)
except sqlite3.OperationalError as error:
kept = (
" The safety backup was kept."
if getattr(error, "_prompt_manager_restore_backup_created", False)
else ""
)
self.logger.error(f"Restore failed: {error}")
return _fail_job(job, "restore-failed", NOT_RESTORED + kept)
await self._prune_after(model.db_path)
previous = self._swap_database(model.db_path)
await self._run_in_executor(previous.close_all)
if backup_path:
job.log_line(f"Safety backup {os.path.basename(backup_path)}")
else:
job.log_line("No safety backup: there was no library to keep")
repaired = documents_line(staged.documents)
if repaired:
job.log_line(repaired)
line = restored_line(staged.summary)
job.log_line(line)
job.update(progress=100, status_line=line)
except Exception:
self.logger.exception("Restore job failed")
_fail_job(job, None, NOT_RESTORED)
finally:
await self._run_in_executor(staged.discard)
py/api/jobs.py
+100
View File
@@ -28,6 +28,15 @@ ones. One job of a kind runs at a time, and each job has an admission class
A short legacy-page operation takes a ``hold``: a running, non-cancellable
job that lasts as long as the request, so the Activity page lists it and a
workspace job is refused while it runs (and the other way round).
Short request-scoped work that reads or writes the library beside any
ordinary job (a Needs review decision, a maintenance or restore preview)
takes a ``lease`` instead: it is not listed (so quick requests never push
real jobs out of the finished history, and two of them never refuse each
other), but it is refused while an ``exclusive`` job runs, and an
``exclusive`` job that replaces the database (a restore) awaits
``drain_leases`` before it switches the file, so no lease's connection is
open across the switch (#249 S8).
"""
import asyncio
@@ -48,6 +57,9 @@ MAX_OUTCOME_ROWS = 200 # rows listed per outcome group; counters stay exact
MAX_LOG_LINES = 2000
LOG_CAP_LINE = "[Log capped at 2,000 lines]"
MAX_FINISHED_JOBS = 20 # finished jobs this session keeps
# The longest an exclusive job waits for running leases to end (a ceiling a
# hung request reaches, not a pacing timer: leases end with their request).
LEASE_DRAIN_SECONDS = 60
# The metadata scan's outcome groups, in counter-tile order of the spec's
# mapping, and the per-file scan status each group collects. Files without
@@ -403,6 +415,74 @@ class JobRegistry:
def __init__(self, logger_=None):
self._logger = logger_ or logger
self._jobs = {}
self._leases = {} # title -> running count
self._drain_waiters = []
# ── leases ────────────────────────────────────────────────────────
def lease(self, title):
"""An async context manager for short request-scoped library work.
Entering raises JobBusy, naming the job, while an ``exclusive`` job
runs; otherwise the lease is counted until the block exits. Leases
are never listed and never refuse one another or any job.
"""
return _Lease(self, title)
def lease_blocker(self):
"""The running ``exclusive`` job a new lease is refused for, or None."""
for job in self.running():
if job.admission == "exclusive":
return job
return None
@property
def leases(self):
"""Running leases by title (a copy)."""
return dict(self._leases)
def _take_lease(self, title):
blocker = self.lease_blocker()
if blocker is not None:
raise JobBusy(blocker)
self._leases[title] = self._leases.get(title, 0) + 1
def _release_lease(self, title):
count = self._leases.get(title, 0) - 1
if count > 0:
self._leases[title] = count
else:
self._leases.pop(title, None)
if not self._leases:
waiters, self._drain_waiters = self._drain_waiters, []
for waiter in waiters:
if not waiter.done():
waiter.set_result(True)
async def drain_leases(self, timeout=None):
"""Wait until no lease runs; False when ``timeout`` ended the wait.
For an ``exclusive`` job about to switch the database: once it runs
no new lease is admitted, so the count only goes down.
"""
if not self._leases:
return True
waiter = asyncio.get_running_loop().create_future()
self._drain_waiters.append(waiter)
try:
await asyncio.wait_for(
asyncio.shield(waiter),
LEASE_DRAIN_SECONDS if timeout is None else timeout,
)
return True
except asyncio.TimeoutError:
self._logger.warning(
"Still running after the wait: " + ", ".join(sorted(self._leases))
)
return False
finally:
if waiter in self._drain_waiters:
self._drain_waiters.remove(waiter)
def start(
self,
@@ -565,3 +645,23 @@ class _Hold:
if self._job.running:
self._job.finish("finished" if exc_type is None else "failed")
return False
class _Lease:
"""The async context manager ``JobRegistry.lease`` returns."""
def __init__(self, registry, title):
self._registry = registry
self._title = title
self._held = False
async def __aenter__(self):
self._registry._take_lease(self._title)
self._held = True
return self
async def __aexit__(self, exc_type, exc, tb):
if self._held:
self._held = False
self._registry._release_lease(self._title)
return False
py/api/maintenance.py
+21 -19
View File
@@ -47,7 +47,6 @@ import json
import os
import shutil
import sqlite3
import tempfile
from aiohttp import web
@@ -78,6 +77,7 @@ except ImportError:
from database.maintenance import MAINTENANCE_OPERATIONS, PlanChanged
PREVIEW_BODY_MAX = 4096
MAINTENANCE_PREVIEW_LEASE = "Maintenance preview"
# Duplicate-file previews carry the reviewed keepers (py/api/duplicate_files.py).
DUPLICATE_PREVIEW_BODY_MAX = 1024 * 1024
DUPLICATE_FILES = "duplicate_files"
@@ -238,29 +238,18 @@ def size_label(size):
def restore_room(db_path):
"""The largest restore upload the disks have room for (blocking).
"""The largest restore upload the disk has room for (blocking).
The upload is written to a temp file and staged once more in the
temporary folder (2x), then copied over the live database beside it,
whose current contents are first kept as a safety backup (1x the upload
for the WAL, plus the current database). Needs on one device add up,
each with ``DISK_MARGIN_BYTES``.
Everything a restore writes is in the library's folder (#318): the
upload, its staged copy (``database/restore_stage.py``) and the WAL
growth while the stage is copied over the live database (3x the upload),
plus the safety backup of the current database (1x its size), with
``DISK_MARGIN_BYTES`` to spare.
"""
db_path = os.path.abspath(db_path)
folder = os.path.dirname(db_path)
current = database_bytes(db_path) if os.path.exists(db_path) else 0
devices = {}
for path, per_upload, fixed in (
(tempfile.gettempdir(), 2, 0),
(folder, 1, current),
):
entry = devices.setdefault(_device(path), {"path": path, "per": 0, "fixed": 0})
entry["per"] += per_upload
entry["fixed"] += fixed
room = min(
(_free_bytes(e["path"]) - e["fixed"] - DISK_MARGIN_BYTES) // e["per"]
for e in devices.values()
)
room = (_free_bytes(folder) - current - DISK_MARGIN_BYTES) // 3
return max(0, room)
@@ -369,6 +358,19 @@ class MaintenanceRoutesMixin:
return data, None
async def preview_maintenance(self, request):
"""Review one operation under a registry lease (#249 S8).
The preview's read-only connections must not be open while a restore
switches the database: it is refused (409 ``busy``) while an
``exclusive`` job runs, and a restore waits for running previews.
"""
try:
async with self._jobs.lease(MAINTENANCE_PREVIEW_LEASE):
return await self._preview_maintenance(request)
except JobBusy as busy:
return busy_response(busy)
async def _preview_maintenance(self, request):
data, error = await self._read_maintenance_body(
request,
("operation", "options"),
py/api/restore_plans.py
+114
View File
@@ -0,0 +1,114 @@
"""Reviewed restores: staged uploads held under single-use tokens (#249 S8).
A restore preview stages the upload (``database/restore_stage.py``) and
keeps it here under a ``secrets.token_urlsafe(32)`` token with the live
library's fingerprint and the database identity it was reviewed against.
``take`` hands the stage to the restore job exactly once; a second apply of
the token names that job (``PlanUsed``). An expired, discarded or unknown
token raises ``PlanExpired``; its stage folder is removed. A restart forgets
every token (the startup sweep removes the folders).
At most ``max_stages`` stages are held. Bytes are bounded by free disk space
(the upload is refused when the disk has no room, and held stages use that
disk), not by a fixed budget. A stage being applied belongs to its job and is
never evicted or counted.
"""
import secrets
import time
from dataclasses import dataclass
from typing import Any, Callable, Dict, Optional
from .maintenance_plans import PlanExpired, PlanUsed
STAGE_TTL_SECONDS = 900
MAX_STAGES = 2
_USED_KEPT = 32 # applied tokens remembered for plan-used answers
class TooManyStages(Exception):
"""Every stage slot is held by a reviewed restore."""
@dataclass
class HeldRestore:
staged: Any # database.restore_stage.StagedRestore
fingerprint: str
db_identity: Any
live: Dict[str, int]
expires_at: float
class RestorePlanStore:
def __init__(
self,
ttl: float = STAGE_TTL_SECONDS,
max_stages: int = MAX_STAGES,
clock: Optional[Callable[[], float]] = None,
):
self._ttl = ttl
self._max = max_stages
self._clock = clock or time.time
self._held: Dict[str, HeldRestore] = {}
self._used: Dict[str, str] = {}
def __len__(self):
return len(self._held)
def full(self) -> bool:
return len(self._held) >= self._max
def issue(self, staged, fingerprint, db_identity, live) -> str:
"""Hold ``staged`` and return its token; TooManyStages when full."""
if self.full():
raise TooManyStages()
token = secrets.token_urlsafe(32)
self._held[token] = HeldRestore(
staged=staged,
fingerprint=fingerprint,
db_identity=db_identity,
live=dict(live),
expires_at=self._clock() + self._ttl,
)
return token
def peek(self, token) -> HeldRestore:
if token in self._used:
raise PlanUsed(self._used[token])
held = self._held.get(token)
if held is None:
raise PlanExpired()
if held.expires_at <= self._clock():
self.discard(token)
raise PlanExpired()
return held
def take(self, token, job_id) -> HeldRestore:
"""Hand the held stage to the job applying it (single use)."""
held = self.peek(token)
del self._held[token]
self._used[token] = job_id
while len(self._used) > _USED_KEPT:
self._used.pop(next(iter(self._used)))
return held
def discard(self, token) -> bool:
"""Drop the token and remove its stage; False when it was not held."""
held = self._held.pop(token, None)
if held is None:
return False
held.staged.discard()
return True
def drop_expired(self) -> int:
"""Remove every expired stage (blocking file removal); the count."""
now = self._clock()
expired = [t for t, held in self._held.items() if held.expires_at <= now]
for token in expired:
self.discard(token)
return len(expired)
def discard_all(self) -> None:
"""Remove every held stage (the database was replaced or shutdown)."""
for token in list(self._held):
self.discard(token)
py/api/review.py
+19 -1
View File
@@ -31,6 +31,8 @@ from aiohttp import web
from .annotations import _HEX64, _TOKEN, _UUID, _capture_error, _error
from .images import media_thumbnail_url
from .inspection_service import VIDEO_EXTENSIONS
from .job_routes import busy_response
from .jobs import JobBusy
from .prompts import bad_request
try:
@@ -239,6 +241,9 @@ def _side_label(titles: dict, side, empty: str) -> str:
return display_title(titles.get(stored, ()), text) or empty
REVIEW_DECISION_LEASE = "Needs review decision"
class ReviewRoutesMixin:
"""The Needs review queue over E1's schema and E2's recorded items."""
@@ -446,7 +451,20 @@ class ReviewRoutesMixin:
# -- decision --------------------------------------------------------------
async def decide_review_item(self, request: web.Request) -> web.Response:
"""POST one explicit decision on one pending item."""
"""POST one explicit decision on one pending item.
The decision runs under a registry lease (#249 S8): it is refused
(409 ``busy``) while a restore or another ``exclusive`` job runs, and
a restore waits for decisions already running before it switches the
database. Decisions never refuse one another or a scan.
"""
try:
async with self._jobs.lease(REVIEW_DECISION_LEASE):
return await self._decide_review_item(request)
except JobBusy as busy:
return busy_response(busy)
async def _decide_review_item(self, request: web.Request) -> web.Response:
data, error = await self._read_json_body(request, max_bytes=MAX_BODY_BYTES)
if error is not None:
return error
tests/test_api_admin.py
+4 -3
View File
@@ -2057,7 +2057,8 @@ class TestBackupRestoreContract(AdminAPITestCase):
resp = await self._restore("database_file", upload)
self.assertEqual(resp.status, 200)
self.assertIn("verify_database_file", calls)
self.assertIn("restore_from_file", calls)
self.assertIn("stage_restore", calls)
self.assertIn("publish_staged", calls)
async def test_restore_keeps_at_most_ten_safety_backups(self):
import time
@@ -2098,12 +2099,12 @@ class TestBackupRestoreContract(AdminAPITestCase):
self.api.db.save_prompt(text="restored", prompt_hash="h1")
upload = self._snapshot_upload()
def boom(src_path):
def boom(staged):
raise PermissionError(
13, "Permission denied", os.path.join(self.tmpdir, "prompts.db")
)
self.api.db.model.restore_from_file = boom
self.api.db.model.publish_staged = boom
resp = await self._restore("database_file", upload)
body = await resp.text()
self.assertEqual(resp.status, 500)
tests/test_api_backup_restore.py
+4 -3
View File
@@ -121,7 +121,7 @@ class TestBackupRoute(BackupRestoreTestCase):
created.append(result[1])
return result
with patch("py.api.admin.tempfile.mkstemp", spy_mkstemp):
with patch("tempfile.mkstemp", spy_mkstemp):
resp = await self.client.get("/prompt_manager/backup")
self.assertEqual(resp.status, 200)
@@ -270,7 +270,7 @@ class TestRestoreRoute(BackupRestoreTestCase):
created.append(result[1])
return result
with patch("py.api.admin.tempfile.mkstemp", spy_mkstemp):
with patch("tempfile.mkstemp", spy_mkstemp):
resp = await self._restore(upload)
self.assertEqual(resp.status, 200)
@@ -278,6 +278,7 @@ class TestRestoreRoute(BackupRestoreTestCase):
# temp copy are gone afterwards.
uploads = [p for p in created if _is_upload(p)]
self.assertEqual(len(uploads), 1)
self.assertEqual(os.path.dirname(uploads[0]), self.tmpdir)
for path in created:
for suffix in ("", "-wal", "-shm", "-journal"):
self.assertFalse(os.path.exists(path + suffix))
@@ -409,7 +410,7 @@ class TestLargeRestores(BackupRestoreTestCase):
created.append(result[1])
return result
return patch("py.api.admin.tempfile.mkstemp", spy), created
return patch("tempfile.mkstemp", spy), created
def _assert_removed(self, created):
for path in created:
tests/test_job_leases.py
+100
View File
@@ -0,0 +1,100 @@
"""Registry leases: short request-scoped work versus a database switch (#249 S8).
A lease is refused while an ``exclusive`` job (restore, Move library,
Optimize) runs, never refuses another lease or any job, is not listed, and an
exclusive job can await the leases already running. Ordering is driven by
events only.
"""
import asyncio
import os
import sys
import unittest
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from py.api.jobs import JobBusy, JobRegistry # noqa: E402
class TestLeases(unittest.IsolatedAsyncioTestCase):
def setUp(self):
self.registry = JobRegistry()
def start(self, kind, admission):
gate = asyncio.Event()
async def runner(job):
await gate.wait()
job = self.registry.start(
kind, kind.title(), admission=admission, scope={}, runner=runner
)
return job, gate
async def test_leases_share_and_are_not_listed(self):
scan, gate = self.start("scan", "records")
async with self.registry.lease("Needs review decision"):
async with self.registry.lease("Needs review decision"):
self.assertEqual(self.registry.leases, {"Needs review decision": 2})
# A lease never refuses a job, even an exclusive one.
self.assertIsNone(self.registry.blocker("health", admission="read"))
self.assertEqual(self.registry.leases, {})
self.assertEqual([job.kind for job in self.registry.list()], ["scan"])
gate.set()
await scan.task
async def test_lease_refused_while_exclusive_job_runs(self):
restore, gate = self.start("restore", "exclusive")
with self.assertRaises(JobBusy) as refused:
async with self.registry.lease("Maintenance preview"):
self.fail("admitted beside a restore")
self.assertIs(refused.exception.running, restore)
self.assertEqual(self.registry.leases, {})
gate.set()
await restore.task
async with self.registry.lease("Maintenance preview"):
pass
async def test_exclusive_job_awaits_running_leases(self):
entered, release = asyncio.Event(), asyncio.Event()
order = []
async def decision():
async with self.registry.lease("Needs review decision"):
entered.set()
await release.wait()
order.append("decision done")
task = asyncio.ensure_future(decision())
await entered.wait()
drained = asyncio.ensure_future(self.registry.drain_leases())
await asyncio.sleep(0)
self.assertFalse(drained.done())
release.set()
self.assertTrue(await drained)
order.append("drained")
await task
self.assertEqual(order, ["decision done", "drained"])
self.assertTrue(await self.registry.drain_leases()) # none: at once
async def test_drain_reports_a_lease_that_outlives_the_ceiling(self):
hold = self.registry.lease("Stuck preview")
await hold.__aenter__()
try:
with self.assertLogs("py.api.jobs", "WARNING") as logs:
self.assertFalse(await self.registry.drain_leases(timeout=0))
self.assertIn("Stuck preview", "\n".join(logs.output))
self.assertEqual(self.registry._drain_waiters, [])
finally:
await hold.__aexit__(None, None, None)
self.assertEqual(self.registry.leases, {})
async def test_a_failed_body_still_releases_the_lease(self):
with self.assertRaises(RuntimeError):
async with self.registry.lease("Needs review decision"):
raise RuntimeError("boom")
self.assertEqual(self.registry.leases, {})
if __name__ == "__main__":
unittest.main()
tests/test_restore_plans.py
+574
View File
@@ -0,0 +1,574 @@
"""Reviewed restore and verified backup routes (#249 System S8).
- ``POST /prompt_manager/restore/preview`` stages an upload and writes
nothing live; ``DELETE .../preview/{token}`` discards it;
``POST /prompt_manager/restore/apply`` starts the ``exclusive`` restore
job that publishes exactly the reviewed bytes.
- A restore and the health check, maintenance, Move library, a Needs review
decision and a maintenance preview refuse each other or wait (carry-overs
1 and 2).
- ``GET /prompt_manager/backup`` verifies its snapshot before sending a byte;
``GET /prompt_manager/backup/status`` lists the safety backups on disk and
this session's verified download.
Ordering is driven by events and job tasks, never by timers.
"""
import asyncio
import glob
import hashlib
import os
import sqlite3
import threading
from contextlib import closing
from pathlib import Path
from unittest.mock import patch
import aiohttp
from database import restore_stage
from database.models import PromptModel
from database.operations import PromptDatabase
from py.api import restore_plans
from py.api.backup_restore import RESTORE_JOB_TITLE
from utils.hashing import generate_prompt_hash
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 = 60 # a ceiling only a hung test reaches
PREVIEW = "/prompt_manager/restore/preview"
APPLY = "/prompt_manager/restore/apply"
SUMMARY_KEYS = {
"prompts",
"positives",
"pairs",
"archived_pairs",
"image_links",
"runs",
"saved_views",
}
def _texts(path):
with closing(sqlite3.connect(path)) as conn:
return sorted(r[0] for r in conn.execute("SELECT text FROM prompts"))
class RestoreTestCase(AdminAPITestCase):
@property
def db_path(self):
return os.path.join(self.tmpdir, "prompts.db")
def save(self, text, db=None):
return (db or self.api.db).save_prompt(
text=text, prompt_hash=generate_prompt_hash(text)
)
def upload(self, texts, name="upload.db"):
folder = os.path.join(self.tmpdir, "uploads")
os.makedirs(folder, exist_ok=True)
path = os.path.join(folder, name)
db = PromptDatabase(path)
for text in texts:
self.save(text, db)
db.close_all()
return Path(path).read_bytes()
def form(self, data):
form = aiohttp.FormData()
form.add_field(
"database_file",
data,
filename="backup.db",
content_type="application/octet-stream",
)
return form
async def preview(self, data, status=200):
resp = await self.client.post(PREVIEW, data=self.form(data))
body = await resp.json()
self.assertEqual(resp.status, status, body)
return body
async def apply(self, token, status=202):
resp = await self.client.post(APPLY, json={"token": token})
body = await resp.json()
self.assertEqual(resp.status, status, body)
return body
async def finish(self, job_id):
job = self.api._jobs.get(job_id)
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
return job
def live_digest(self):
return {
suffix: hashlib.sha256(Path(self.db_path + suffix).read_bytes()).hexdigest()
for suffix in ("", "-wal")
if os.path.exists(self.db_path + suffix)
}
def backups(self):
return sorted(glob.glob(self.db_path + ".backup_*"))
def stages(self):
return sorted(
name
for name in os.listdir(self.tmpdir)
if name.startswith(restore_stage.STAGE_PREFIX)
)
def gated_job(self, kind, admission, title=None):
gate = asyncio.Event()
async def runner(job):
await gate.wait()
job = self.api._jobs.start(
kind, title or kind.title(), admission=admission, scope={}, runner=runner
)
self.addAsyncCleanup(self._release, job, gate)
return job, gate
async def _release(self, job, gate):
gate.set()
await asyncio.gather(job.task, return_exceptions=True)
class TestPreview(RestoreTestCase):
async def test_preview_never_replaces_or_backs_up(self):
self.save("live one")
self.api.db.model.get_connection()
before = self.live_digest()
body = await self.preview(self.upload(["restored a", "restored b"]))
self.assertEqual(self.live_digest(), before)
self.assertEqual(self.backups(), [])
self.assertEqual(
[p["text"] for p in self.api.db.search_prompts()], ["live one"]
)
self.assertEqual(set(body["summary"]), SUMMARY_KEYS)
self.assertEqual(set(body["live"]), SUMMARY_KEYS)
self.assertEqual((body["summary"]["prompts"], body["live"]["prompts"]), (2, 1))
self.assertEqual(body["schema"], {"migrated": False, "library_state": True})
self.assertEqual(body["integrity"], "ok")
self.assertEqual(body["documents"], {"repaired": 0, "cleared": 0, "images": 0})
self.assertEqual(
body["safety_backup"],
{"destination_label": "prompts.db.backup_<date and time>"},
)
self.assertNotIn(self.tmpdir, str(body))
# Staged in the library's folder; the upload itself is gone.
self.assertEqual(len(self.stages()), 1)
self.assertEqual(
[n for n in os.listdir(self.tmpdir) if "upload" in n and "pm-" in n], []
)
self.assertEqual(self.api._jobs.list(), [])
async def test_invalid_upload_is_400_and_leaves_nothing(self):
resp = await self.client.post(PREVIEW, data=self.form(b"not a database"))
body = await resp.json()
self.assertEqual((resp.status, body["code"]), (400, "invalid-backup"))
self.assertEqual(self.stages(), [])
async def test_expired_or_discarded_token_is_410(self):
body = await self.preview(self.upload(["restored"]))
resp = await self.client.delete(f"{PREVIEW}/{body['token']}")
self.assertEqual(resp.status, 200)
self.assertEqual(self.stages(), [])
refused = await self.apply(body["token"], status=410)
self.assertEqual(refused["code"], "plan-expired")
resp = await self.client.delete(f"{PREVIEW}/{body['token']}")
self.assertEqual(resp.status, 410)
clock = [1000.0]
self.api._restore_plan_store = restore_plans.RestorePlanStore(
clock=lambda: clock[0]
)
body = await self.preview(self.upload(["restored"], "second.db"))
clock[0] += restore_plans.STAGE_TTL_SECONDS + 1
refused = await self.apply(body["token"], status=410)
self.assertEqual(refused["code"], "plan-expired")
self.assertEqual(self.stages(), [])
async def test_stage_budget_429(self):
data = self.upload(["restored"])
first = await self.preview(data)
await self.preview(data)
refused = await self.preview(data, status=429)
self.assertEqual(refused["code"], "too-large")
self.assertEqual(len(self.stages()), 2)
await self.client.delete(f"{PREVIEW}/{first['token']}")
await self.preview(data)
async def test_preview_reports_repaired_documents(self):
upload = os.path.join(self.tmpdir, "uploads", "damaged.db")
os.makedirs(os.path.dirname(upload), exist_ok=True)
src = PromptDatabase(upload)
prompt_id = self.save("with a workflow", src)
image = src.link_image_to_prompt(
prompt_id, os.path.join(self.tmpdir, "a.png"), {"workflow": {"v": 1}}
)
conn = src.model.get_connection()
with conn:
conn.execute(
"UPDATE workflow_documents SET data=x'00' WHERE id=(SELECT "
"workflow_document_id FROM generated_images WHERE id=?)",
(image,),
)
src.close_all()
body = await self.preview(Path(upload).read_bytes())
self.assertEqual(body["documents"], {"repaired": 0, "cleared": 1, "images": 1})
class TestApply(RestoreTestCase):
async def test_apply_restores_and_logs(self):
self.save("old")
old_db = self.api.db
body = await self.preview(self.upload(["restored a", "restored b"]))
started = await self.apply(body["token"])
summary = started["job"]
self.assertEqual(
(summary["kind"], summary["title"], summary["cancellable"]),
("restore", RESTORE_JOB_TITLE, False),
)
self.assertEqual(summary["admission"], "exclusive")
job = await self.finish(summary["id"])
self.assertEqual(job.state, "finished", job.log)
self.assertIsNot(self.api.db, old_db)
self.assertEqual(_texts(self.db_path), ["restored a", "restored b"])
[backup] = self.backups()
self.assertEqual(_texts(backup), ["old"])
log = "\n".join(job.log)
self.assertIn(f"Safety backup {os.path.basename(backup)}", log)
self.assertIn("Restored 2 prompts", log)
self.assertNotIn(self.tmpdir, log)
self.assertEqual(self.stages(), [])
wal = self.db_path + "-wal"
self.assertEqual(os.path.getsize(wal) if os.path.exists(wal) else 0, 0)
used = await self.apply(body["token"], status=409)
self.assertEqual(used["code"], "plan-used")
self.assertEqual(used["job"]["id"], summary["id"])
async def test_apply_binds_exact_staged_bytes(self):
self.save("live")
body = await self.preview(self.upload(["restored"]))
held = self.api._restore_plans.peek(body["token"])
with closing(sqlite3.connect(held.staged.path)) as conn:
conn.execute("INSERT INTO prompts (text, hash) VALUES ('sneaked', 'x')")
conn.commit()
started = await self.apply(body["token"])
job = await self.finish(started["job"]["id"])
self.assertEqual(job.state, "failed")
self.assertEqual(job.last_event["code"], "plan-changed")
self.assertEqual(_texts(self.db_path), ["live"])
self.assertEqual(self.backups(), [])
self.assertEqual(self.stages(), [])
async def test_live_change_after_preview_is_409(self):
self.save("live")
body = await self.preview(self.upload(["restored"]))
self.save("written after the review")
refused = await self.apply(body["token"], status=409)
self.assertEqual(refused["code"], "plan-changed")
self.assertEqual(
sorted(p["text"] for p in self.api.db.search_prompts()),
["live", "written after the review"],
)
self.assertEqual(self.stages(), [])
self.assertEqual(self.api._jobs.list(), [])
async def test_publish_failure_keeps_safety_backup_and_reports(self):
self.save("live")
body = await self.preview(self.upload(["restored"]))
failure = sqlite3.OperationalError("injected page copy failure")
with patch.object(PromptModel, "_replace_contents_from", side_effect=failure):
started = await self.apply(body["token"])
job = await self.finish(started["job"]["id"])
self.assertEqual(job.state, "failed")
self.assertEqual(job.last_event["code"], "restore-failed")
self.assertIn("The safety backup was kept.", job.status_line)
[backup] = self.backups()
self.assertEqual(_texts(backup), ["live"])
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["live"])
self.assertEqual(self.stages(), [])
async def test_jobs_survive_db_swap(self):
earlier, gate = self.gated_job("health", "read", "Health check")
gate.set()
await earlier.task
body = await self.preview(self.upload(["restored"]))
job = await self.finish((await self.apply(body["token"]))["job"]["id"])
self.assertEqual(job.state, "finished")
listed = (await (await self.client.get("/prompt_manager/jobs")).json())["jobs"]
self.assertEqual(
[(j["kind"], j["state"]) for j in listed],
[("restore", "finished"), ("health", "finished")],
)
class TestRestoreAdmission(RestoreTestCase):
"""A restore never interleaves with other library work (carry-overs 1, 2)."""
async def test_restore_job_refused_while_maintenance_runs(self):
self.gated_job("maintenance", "records", "Duplicate prompt text")
body = await self.preview(self.upload(["restored"]))
refused = await self.apply(body["token"], status=409)
self.assertEqual(refused["code"], "busy")
self.assertEqual(refused["running"]["title"], "Duplicate prompt text")
self.api._restore_plans.peek(body["token"]) # the token stays usable
async def test_restore_and_health_check_refuse_each_other(self):
health, gate = self.gated_job("health", "read", "Health check")
body = await self.preview(self.upload(["restored"]))
refused = await self.apply(body["token"], status=409)
self.assertEqual(refused["running"]["title"], "Health check")
legacy = await self.client.post(
"/prompt_manager/restore", data=self.form(self.upload(["x"], "x.db"))
)
self.assertEqual(legacy.status, 409)
gate.set()
await health.task
restore, release = self.gated_job("restore", "exclusive", RESTORE_JOB_TITLE)
resp = await self.client.post("/prompt_manager/system/health")
refusal = await resp.json()
self.assertEqual((resp.status, refusal["code"]), (409, "busy"))
self.assertEqual(refusal["running"]["title"], RESTORE_JOB_TITLE)
async def test_restore_and_move_library_refuse_each_other(self):
self.gated_job("library_move", "exclusive", "Move library")
body = await self.preview(self.upload(["restored"]), status=409)
self.assertEqual(body["running"]["title"], "Move library")
self.api._jobs.list()[0].finish()
body = await self.preview(self.upload(["restored"]))
self.gated_job("restore", "exclusive", RESTORE_JOB_TITLE)
resp = await self.client.post("/prompt_manager/system/library/move")
self.assertEqual(resp.status, 409)
self.assertEqual((await resp.json())["running"]["title"], RESTORE_JOB_TITLE)
async def test_decision_and_preview_refused_while_restore_runs(self):
self.gated_job("restore", "exclusive", RESTORE_JOB_TITLE)
for url, payload in (
("/prompt_manager/review/items/1/decision", {"choice": "unresolved"}),
("/prompt_manager/maintenance/preview", {"operation": "optimize"}),
):
resp = await self.client.post(url, json=payload)
body = await resp.json()
self.assertEqual((resp.status, body["code"]), (409, "busy"), url)
self.assertEqual(body["running"]["title"], RESTORE_JOB_TITLE)
resp = await self.client.post(PREVIEW, data=self.form(self.upload(["r"])))
self.assertEqual(resp.status, 409)
async def test_restore_waits_for_a_running_decision(self):
self.save("live")
body = await self.preview(self.upload(["restored"]))
entered, release = threading.Event(), threading.Event()
loop = asyncio.get_running_loop()
def slow_item(item_id):
entered.set()
release.wait(WAIT)
return None
with patch.object(self.api.db, "review_item", side_effect=slow_item):
decision = asyncio.ensure_future(
self.client.post(
"/prompt_manager/review/items/1/decision",
json={
"expected_epoch": "00000000-0000-4000-8000-000000000000",
"expected_revision": 0,
"capture_token": None,
"sha256": "0" * 64,
"dependency_fingerprint": "0" * 64,
"choice": "unresolved",
},
)
)
self.assertTrue(await loop.run_in_executor(None, entered.wait, WAIT))
self.assertEqual(self.api._jobs.leases, {"Needs review decision": 1})
# The apply's own check takes a lease too: it shares with the decision.
started = await self.apply(body["token"])
job = self.api._jobs.get(started["job"]["id"])
while job.status_line != "Waiting for other requests to finish":
await asyncio.sleep(0)
self.assertEqual(_texts(self.db_path), ["live"])
release.set()
answered = await decision
self.assertNotEqual(answered.status, 409)
await self.finish(job.id)
self.assertEqual(job.state, "finished", job.log)
self.assertEqual(_texts(self.db_path), ["restored"])
class TestWorkerHandles(RestoreTestCase):
"""Carry-over 3: no worker holds the library open across a switch."""
async def test_preview_workers_pause_and_library_connections_are_released(self):
restore, _gate = self.gated_job("restore", "exclusive", RESTORE_JOB_TITLE)
# Automatic previews (thumbnail workers) write nothing while it runs;
# they never open the library file at all.
self.assertIs(self.api._auto_thumbnail_blocker(), restore)
async def test_a_restore_leaves_no_handle_on_the_stage_or_upload(self):
from tests.open_handles import open_handles_under
self.save("live")
body = await self.preview(self.upload(["restored"]))
job = await self.finish((await self.apply(body["token"]))["job"]["id"])
self.assertEqual(job.state, "finished", job.log)
leftovers = [
hit
for hit in open_handles_under(self.tmpdir)
if hit.startswith(".pm-") or hit.startswith("staged")
]
self.assertEqual(leftovers, [])
self.assertEqual(
[n for n in os.listdir(self.tmpdir) if n.startswith(".pm-")], []
)
class TestLegacyRestore(RestoreTestCase):
async def test_legacy_restore_unchanged(self):
self.save("old")
resp = await self.client.post(
"/prompt_manager/restore", data=self.form(self.upload(["a", "b"]))
)
body = await resp.json()
self.assertEqual(resp.status, 200, body)
self.assertEqual(
set(body), {"success", "message", "prompt_count", "backup_created"}
)
self.assertEqual(
body["message"], "Database restored successfully. Found 2 prompts."
)
self.assertEqual(self.stages(), [])
[job] = self.api._jobs.list()
self.assertEqual((job.kind, job.admission), ("restore", "exclusive"))
async def test_legacy_restore_reports_damaged_documents(self):
upload = os.path.join(self.tmpdir, "uploads", "damaged.db")
os.makedirs(os.path.dirname(upload), exist_ok=True)
src = PromptDatabase(upload)
prompt_id = self.save("with a workflow", src)
src.link_image_to_prompt(
prompt_id, os.path.join(self.tmpdir, "a.png"), {"workflow": {"v": 1}}
)
conn = src.model.get_connection()
with conn:
conn.execute("UPDATE workflow_documents SET data=x'00'")
src.close_all()
resp = await self.client.post(
"/prompt_manager/restore", data=self.form(Path(upload).read_bytes())
)
body = await resp.json()
self.assertEqual(resp.status, 200, body)
self.assertEqual(
body["workflow_documents"], {"repaired": 0, "cleared": 2, "images": 1}
) # the image's workflow and its prompt graph
self.assertTrue(
body["message"].endswith(
"Could not repair 2 workflow documents: 1 image shows no workflow."
),
body["message"],
)
ok, reason = PromptModel.verify_database_file(self.db_path)
self.assertTrue(ok, reason)
class TestBackupDownload(RestoreTestCase):
async def test_backup_verifies_snapshot_before_streaming(self):
self.save("keep")
with patch.object(
PromptModel, "verify_database_file", return_value=(False, "bad")
):
resp = await self.client.get("/prompt_manager/backup")
self.assertEqual(resp.status, 500)
self.assertEqual(
await resp.json(),
{"success": False, "error": "Failed to create database backup"},
)
self.assertEqual(
[n for n in os.listdir(self.tmpdir) if n.startswith(".pm-")], []
)
status = await (await self.client.get("/prompt_manager/backup/status")).json()
self.assertIsNone(status["last_download"])
async def test_temp_snapshot_removed_and_closed(self):
self.save("keep")
created = []
from database import safety_backup
original = safety_backup.safety_temp
def spy(db_path):
path = original(db_path)
created.append(path)
return path
with patch("py.api.backup_restore.safety_temp", spy):
resp = await self.client.get("/prompt_manager/backup")
data = await resp.read()
self.assertEqual(resp.status, 200)
[temp] = created
self.assertEqual(os.path.dirname(temp), self.tmpdir)
self.assertFalse(os.path.exists(temp))
downloaded = os.path.join(self.tmpdir, "downloaded.db")
Path(downloaded).write_bytes(data)
self.assertEqual(_texts(downloaded), ["keep"])
with closing(sqlite3.connect(downloaded)) as conn:
self.assertEqual(
conn.execute("PRAGMA journal_mode").fetchone()[0], "delete"
)
async def test_status_null_download_on_fresh_server(self):
resp = await self.client.get("/prompt_manager/backup/status")
body = await resp.json()
self.assertEqual(
body, {"success": True, "safety_backups": [], "last_download": None}
)
await (await self.client.get("/prompt_manager/backup")).read()
body = await (await self.client.get("/prompt_manager/backup/status")).json()
download = body["last_download"]
self.assertEqual(set(download), {"at", "bytes", "verified"})
self.assertTrue(download["verified"])
self.assertGreater(download["bytes"], 0)
async def test_status_lists_safety_backups_from_disk(self):
for index in range(12):
path = Path(f"{self.db_path}.backup_2026010{index % 10}_00000{index}")
path.write_bytes(b"x" * (index + 1))
os.utime(path, (1_000_000 + index, 1_000_000 + index))
Path(self.tmpdir, "other.db.backup_1").write_bytes(b"x")
body = await (await self.client.get("/prompt_manager/backup/status")).json()
listed = body["safety_backups"]
self.assertEqual(len(listed), 10)
self.assertEqual(listed[0]["bytes"], 12)
self.assertEqual(listed[-1]["bytes"], 3)
for row in listed:
self.assertEqual(set(row), {"name", "bytes", "modified_at"})
self.assertTrue(row["name"].startswith("prompts.db.backup_"))
self.assertNotIn(os.sep, row["name"])
self.assertNotIn(self.tmpdir, str(body))
async def test_backup_refused_while_restore_runs(self):
self.gated_job("restore", "exclusive", RESTORE_JOB_TITLE)
resp = await self.client.get("/prompt_manager/backup")
self.assertEqual(resp.status, 409)
self.assertEqual(
await resp.json(),
{"success": False, "error": f"{RESTORE_JOB_TITLE} already running"},
)
async def test_backup_is_listed_while_it_streams(self):
self.save("keep")
resp = await self.client.get("/prompt_manager/backup")
await resp.read()
[job] = self.api._jobs.list()
self.assertEqual(
(job.kind, job.title, job.admission, job.state),
("backup", "Database backup", "guarded", "finished"),
)