diff --git a/py/api/__init__.py b/py/api/__init__.py index 6b2748f3..97504b7b 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -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) diff --git a/py/api/admin.py b/py/api/admin.py index 1e54e991..eaf62124 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -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. diff --git a/py/api/backup_restore.py b/py/api/backup_restore.py new file mode 100644 index 00000000..5b177e10 --- /dev/null +++ b/py/api/backup_restore.py @@ -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 + (``.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): + """``.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_" + }, + }, + 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) diff --git a/py/api/jobs.py b/py/api/jobs.py index d27426c6..9e5c463b 100644 --- a/py/api/jobs.py +++ b/py/api/jobs.py @@ -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 diff --git a/py/api/maintenance.py b/py/api/maintenance.py index cc881348..f53b502e 100644 --- a/py/api/maintenance.py +++ b/py/api/maintenance.py @@ -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"), diff --git a/py/api/restore_plans.py b/py/api/restore_plans.py new file mode 100644 index 00000000..e15b9de8 --- /dev/null +++ b/py/api/restore_plans.py @@ -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) diff --git a/py/api/review.py b/py/api/review.py index 9490ae34..6f81e545 100644 --- a/py/api/review.py +++ b/py/api/review.py @@ -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 diff --git a/tests/test_api_admin.py b/tests/test_api_admin.py index 39ec1ffd..2fb4305c 100644 --- a/tests/test_api_admin.py +++ b/tests/test_api_admin.py @@ -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) diff --git a/tests/test_api_backup_restore.py b/tests/test_api_backup_restore.py index 96b1c4da..19b4d6d2 100644 --- a/tests/test_api_backup_restore.py +++ b/tests/test_api_backup_restore.py @@ -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: diff --git a/tests/test_job_leases.py b/tests/test_job_leases.py new file mode 100644 index 00000000..18ca2b5c --- /dev/null +++ b/tests/test_job_leases.py @@ -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() diff --git a/tests/test_restore_plans.py b/tests/test_restore_plans.py new file mode 100644 index 00000000..73415d7a --- /dev/null +++ b/tests/test_restore_plans.py @@ -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_"}, + ) + 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"), + )