"""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 download_backup(self): """GET /backup, read every byte, then wait for its hold to end. The client can hold the last byte before the handler has finished: it records the download, then closes and removes its snapshot. So the state is read only once the snapshot file is gone. """ resp = await self.client.get("/prompt_manager/backup") data = await resp.read() while self.api._jobs.active("backup") is not None or any( name.startswith(".pm-safety-") for name in os.listdir(self.tmpdir) ): await asyncio.sleep(0) return resp, data 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_a_restore_drops_previews_of_the_old_library(self): first = await self.preview(self.upload(["first"])) other = await self.preview(self.upload(["other"], "other.db")) job = await self.finish((await self.apply(first["token"]))["job"]["id"]) self.assertEqual(job.state, "finished", job.log) self.assertEqual(self.stages(), []) refused = await self.apply(other["token"], status=410) self.assertEqual(refused["code"], "plan-expired") 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): # A corrupt snapshot is refused (a semantic-only failure still # downloads: tests/test_restore_review_fixes.py, #335 review M5). self.save("keep") with patch( "py.api.backup_restore.snapshot_corruption", return_value="Integrity check failed: page 3", ): 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, data = await self.download_backup() 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 self.download_backup() 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") await self.download_backup() [job] = self.api._jobs.list() self.assertEqual( (job.kind, job.title, job.admission, job.state), ("backup", "Database backup", "guarded", "finished"), )