diff --git a/TESTING.md b/TESTING.md index 7c421961..fd335557 100644 --- a/TESTING.md +++ b/TESTING.md @@ -3199,8 +3199,11 @@ the library's own folder (never the system temp folder); free space there must h the upload three times (upload, stage, WAL growth) plus the current database (`restore_room`). The stage checks the upload without reading every workflow document, copies it with the backup API (rollback-journal and WAL files alike), migrates an older -schema, repairs damaged workflow documents (re-stored from an inline copy with the same -SHA-256, else dropped with their images' references cleared, both counted), rotates the +schema, repairs damaged workflow documents (read one row at a time: a body that inflates +completely while only its recorded hash or only its length disagrees is re-keyed or merged +into the document holding that text; any other damaged body is re-stored from an inline +copy with the same SHA-256, else dropped with its images' references cleared; repaired, +cleared and distinct images are counted), rotates the history epoch, makes the copy standalone, verifies it with every document read and records its SHA-256. Publishing rechecks that digest (`StageChanged`), keeps the current database as `.backup_[_n]` (`database/safety_backup.py`: @@ -3211,27 +3214,40 @@ logging one that stays (`tests/test_restore_stage.py`). `POST /prompt_manager/restore/preview` answers `{token, expires_at, summary, schema, integrity, documents, live, size, safety_backup}` and writes nothing live; at most two -stages are held (429 `too-large`) for 15 minutes; `DELETE .../preview/{token}` removes -one. `POST /prompt_manager/restore/apply {token}` rechecks the live fingerprint (409 +stages are held (429 `too-large`) for 15 minutes, and a timer removes an expired stage +without another request; `DELETE .../preview/{token}` removes one. A preview that is +refused, fails or is cancelled (even while the worker still stages) removes its stage. The +store of reviewed previews changes only on the event loop; stage folders are removed in +the executor. `POST /prompt_manager/restore/apply {token}` rechecks the live fingerprint (409 `plan-changed`, stage dropped) and starts "Database restore": kind `restore`, admission `exclusive` (it runs alone, like Move library and Optimize), not cancellable, -indeterminate. Its runner waits for request leases, rechecks the fingerprint, +indeterminate. Its runner waits for request leases (when they outlast the wait the job +fails `busy` and the review is kept under its token), rechecks the fingerprint, publishes, swaps `self.db`, prunes the pool and logs the safety backup by name, any document repair and the restored counts; a failed page copy keeps the safety backup and -says so. Tokens are single use (409 `plan-used` with the job), expire (410 +says so. Once the copy replaced the library the restore succeeded: pruning, reopening +`self.db`, dropping the old library's previews and closing the previous object only log +their failures, in the job and the legacy route alike. The fingerprint covers row counts +and the newest row and edit of prompts, images, pairs, runs and saved views, not every +edit (a tag or favorite change on an older row); those survive in the safety backup. Tokens are single use (409 `plan-used` with the job), expire (410 `plan-expired`) and are bound to the database identity (409 `storage-changed`). The legacy `/restore` runs the same stage and publish under the same `exclusive` kind with its old response; a repair adds `workflow_documents` and a sentence to `message` (`tests/test_restore_plans.py`, `tests/test_api_backup_restore.py`). Short request-scoped work takes a registry lease (`JobRegistry.lease`): a Needs review -decision, a maintenance preview, a restore preview and the apply's check. A lease is not +decision, a maintenance preview, the restore preview's read of the live library (not its +upload or staging) and the apply's check. A lease is not listed, never refuses another lease or a job, and is refused (409 `busy`) while an `exclusive` job runs; the restore waits for running leases before it switches the file (`tests/test_job_leases.py`). `GET /prompt_manager/backup` holds `backup` ("Database -backup", `guarded`), writes a standalone copy beside the database, verifies it -(integrity and schema, not every document) before a byte is sent (500 and no bytes -otherwise), streams it and closes the file before removing it. `GET +backup", `guarded`) only while it writes a standalone copy beside the database, so a +stalled download never blocks other work. A copy that is not SQLite or fails `PRAGMA +integrity_check` is a 500 with no bytes. A copy that fails only a semantic check +(`verify_database_file`, workflow documents not read) is still sent, so a user can always +get a backup: the reason is logged, sent as `X-PromptManager-Verify` (`ok` otherwise) and +recorded as `verified:false`. The file is streamed outside the hold, then closed and +removed. `GET /prompt_manager/backup/status` lists up to ten `.backup_*` files newest first (name, bytes, modified time; no paths) and this session's last verified download, or null. diff --git a/database/restore_stage.py b/database/restore_stage.py index 923d429b..75d8883a 100644 --- a/database/restore_stage.py +++ b/database/restore_stage.py @@ -186,15 +186,20 @@ _FINGERPRINT_SQL = ( def live_fingerprint(conn: sqlite3.Connection): """``(summary, fingerprint)`` of a live library, read in one snapshot. - The fingerprint changes when anything a restore would overwrite was - added or edited since (counts plus the newest row and edit of each - table); the apply refuses with ``plan-changed`` then. + The fingerprint changes when rows are added to or removed from the + tables a restore replaces, or when a prompt, pair, run or saved view is + the newest edited (counts plus the newest row and edit of each table); + the apply refuses with ``plan-changed`` then. It does not see every + edit (a tag or favorite change on an older row); those are kept in the + restore's safety backup. A read transaction the caller already has is + used and left open; one opened here is ended here. """ import json - with conn: - if not conn.in_transaction: - conn.execute("BEGIN") # one read snapshot for every count + opened = not conn.in_transaction + if opened: + conn.execute("BEGIN") # one read snapshot for every count + try: summary = library_summary(conn) newest = [] for table, sql in _FINGERPRINT_SQL: @@ -202,6 +207,9 @@ def live_fingerprint(conn: sqlite3.Connection): newest.append([table, list(conn.execute(sql).fetchone())]) except sqlite3.OperationalError: newest.append([table, None]) # no such table or column + finally: + if opened: + conn.rollback() # it only read text = json.dumps([summary, newest], sort_keys=True, default=str) return summary, "sha256:" + hashlib.sha256(text.encode("utf-8")).hexdigest() diff --git a/database/workflow_documents.py b/database/workflow_documents.py index a8979898..3c8434a2 100644 --- a/database/workflow_documents.py +++ b/database/workflow_documents.py @@ -483,12 +483,57 @@ INLINE_TEXTS_SQL = ( "SELECT workflow_data, prompt_metadata FROM generated_images " "WHERE workflow_data IS NOT NULL OR prompt_metadata IS NOT NULL" ) +REKEY_SQL = "UPDATE workflow_documents SET sha256=?, byte_length=? WHERE id=?" + + +def _rows(conn: sqlite3.Connection, sql: str, params=()): + """Each row of ``sql``, read one at a time. + + Iterating a managed connection's cursor prefetches ``ITER_CHUNK`` (128) + rows, and a document or inline text can be megabytes. + """ + cursor = conn.execute(sql, params) + while (row := cursor.fetchone()) is not None: + yield row + + +def _inflated(data) -> Optional[bytes]: + """The complete text bytes ``data`` inflates to, or None. + + Complete means the zlib stream ends (its Adler-32 checked) within + MAX_DOCUMENT_BYTES with nothing after it: the body is intact even when + the row's recorded hash or length is not. + """ + try: + inflater = zlib.decompressobj() + raw = inflater.decompress(bytes(data), MAX_DOCUMENT_BYTES + 1) + if not inflater.eof or inflater.unconsumed_tail or inflater.unused_data: + return None + if len(raw) > MAX_DOCUMENT_BYTES: + return None + raw.decode("utf-8") + return raw + except (zlib.error, TypeError, UnicodeDecodeError): + return None + + +def _one_column_off(row, raw: bytes) -> bool: + """Whether only the recorded hash or only the recorded length is wrong. + + The two columns check the body independently. When the body inflates + completely and one of them still agrees with it, that column vouches for + the body and the other one is the damaged value. When both disagree, the + body itself is the suspect (another valid stream), so it is not trusted. + """ + hash_ok = hashlib.sha256(raw).digest() == bytes(row[0]) + length_ok = row[1] == len(raw) + return hash_ok != length_ok def _inline_copies(conn: sqlite3.Connection, wanted: Dict[bytes, int]): """``{document id: text}`` for each wanted hash an inline row still holds.""" found: Dict[int, str] = {} - for row in conn.execute(INLINE_TEXTS_SQL): + for row in _rows(conn, INLINE_TEXTS_SQL): for text in row: if not isinstance(text, str) or not storable(text): continue @@ -500,40 +545,76 @@ def _inline_copies(conn: sqlite3.Connection, wanted: Dict[bytes, int]): return found +def _rekey(conn: sqlite3.Connection, document_id: int, raw: bytes) -> None: + """Give an intact body its true key, or merge it into the document that has it.""" + digest = hashlib.sha256(raw).digest() + holder = conn.execute(LOOKUP_SQL, (digest,)).fetchone() + if holder is None or holder[0] == document_id: + conn.execute(REKEY_SQL, (digest, len(raw), document_id)) + return + for column in DOCUMENT_ID_COLUMNS: + # ``column`` is one of DOCUMENT_ID_COLUMNS, never caller data. + conn.execute( + f"UPDATE generated_images SET {column}=? WHERE {column}=?", # nosec B608 + (holder[0], document_id), + ) + conn.execute("DELETE FROM workflow_documents WHERE id=?", (document_id,)) + if _decode(holder[1:], holder[0], quiet=True) is None: + conn.execute( + REPAIR_SQL, (len(raw), zlib.compress(raw, COMPRESSION_LEVEL), holder[0]) + ) + + def repair_damaged_documents(conn: sqlite3.Connection) -> Dict[str, int]: """Repair or clear every document failing its checks; commit (writes). - For a restore's stage only, never the live library. A damaged document - whose exact text an inline row still holds (the same SHA-256, its key) is - stored again from that text. Any other damaged document is unreadable for - good: its images' references are cleared (they already showed no - workflow) and the document is removed, so the stage verifies. Returns - ``{repaired, cleared, images}``: documents repaired, documents cleared - and the image rows whose reference was cleared. + For a restore's stage only, never the live library. In order: + + 1. A document whose body inflates completely and whose recorded hash or + length (not both) disagrees is re-keyed from its text, or merged + into the document that already holds that text (``_one_column_off``). + 2. A damaged body whose exact text an inline row still holds (the same + SHA-256, its key) is stored again from that text. + 3. Any other damaged document is unreadable for good: its images' + references are cleared (they already showed no workflow) and the + document is removed, so the stage verifies. + + Returns ``{repaired, cleared, images}``: documents repaired (1 and 2), + documents cleared and the distinct image rows whose reference was + cleared. Rows are read one at a time. """ damaged: Dict[bytes, int] = {} - for row in conn.execute(ALL_DOCUMENTS_SQL): + intact: Dict[int, bytes] = {} + for row in _rows(conn, ALL_DOCUMENTS_SQL): if _decode(row[:3], row[3], quiet=True) is None: - damaged[bytes(row[0])] = row[3] + raw = _inflated(row[2]) + if raw is not None and _one_column_off(row, raw): + intact[row[3]] = raw + else: + damaged[bytes(row[0])] = row[3] result = {"repaired": 0, "cleared": 0, "images": 0} - if not damaged: + if not damaged and not intact: return result with conn: - copies = _inline_copies(conn, damaged) + for document_id, raw in intact.items(): + _rekey(conn, document_id, raw) + copies = _inline_copies(conn, damaged) if damaged else {} for document_id, text in copies.items(): raw = text.encode("utf-8") data = zlib.compress(raw, COMPRESSION_LEVEL) conn.execute(REPAIR_SQL, (len(raw), data, document_id)) lost = sorted(set(damaged.values()) - set(copies)) - images = 0 + images: Set[int] = set() for chunk in _chunks(lost): marks = ",".join("?" * len(chunk)) # ``marks`` is only "?" placeholders; the ids are bound. - images += conn.execute( - "SELECT COUNT(*) FROM generated_images WHERE " # nosec B608 + for (image_id,) in _rows( + conn, + "SELECT id FROM generated_images WHERE " # nosec B608 f"workflow_document_id IN ({marks}) OR prompt_document_id IN ({marks})", chunk + chunk, - ).fetchone()[0] + ): + images.add(image_id) for column in DOCUMENT_ID_COLUMNS: conn.execute( f"UPDATE generated_images SET {column}=NULL " # nosec B608 @@ -544,12 +625,14 @@ def repair_damaged_documents(conn: sqlite3.Connection) -> Dict[str, int]: f"DELETE FROM workflow_documents WHERE id IN ({marks})", # nosec B608 chunk, ) - for document_id in copies: + for document_id in list(intact) + list(copies): logger.warning(f"Workflow document {document_id} was damaged; repaired") if lost: logger.warning( f"{len(lost)} damaged workflow document(s) could not be repaired; " - f"{images} image(s) show no workflow" + f"{len(images)} image(s) show no workflow" ) - result.update(repaired=len(copies), cleared=len(lost), images=images) + result.update( + repaired=len(intact) + len(copies), cleared=len(lost), images=len(images) + ) return result diff --git a/py/api/backup_restore.py b/py/api/backup_restore.py index bf35a3d2..493904dd 100644 --- a/py/api/backup_restore.py +++ b/py/api/backup_restore.py @@ -3,12 +3,16 @@ 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. + While the copy is made 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 and made one standalone file; the hold ends there, so a slow or + stalled download never keeps other work waiting. A copy that fails its + integrity check is a 500 with no bytes. One that fails only a semantic + check (``verify_database_file``) is still sent: a user can always get a + backup. The reason is logged and sent as ``X-PromptManager-Verify`` + (``ok`` otherwise), and the status records it as not verified. The file + is streamed, 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``). @@ -28,16 +32,20 @@ Reviewed restore (new routes answer errors as ``{success:false, code, error}``): 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. + backups and logs what was restored. Once the copy replaced the library, + the restore is done: the later steps only log their failures. - ``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 asyncio import datetime import os import sqlite3 import tempfile +from contextlib import closing +from pathlib import Path from aiohttp import web @@ -82,6 +90,12 @@ REQUESTS_BUSY = ( "Other requests are still reading the library; nothing was changed. " "Try again in a moment." ) +REVIEW_KEPT = ( + "Other requests are still reading the library; nothing was changed. " + "The review is kept: start the restore again in a moment." +) +VERIFY_HEADER = "X-PromptManager-Verify" +SQLITE_HEADER = b"SQLite format 3\x00" PLAN_CHANGED = ( "The library changed since the review; nothing was changed. " "Review the restore again." @@ -158,27 +172,55 @@ def safety_backups(db_path, limit=MAX_LISTED_BACKUPS): ] +def snapshot_corruption(path): + """Why the snapshot file is corrupt, or None (blocking). + + Only physical damage refuses a download: not a SQLite file, or a failed + ``PRAGMA integrity_check``. + """ + try: + with open(path, "rb") as handle: + if handle.read(len(SQLITE_HEADER)) != SQLITE_HEADER: + return "Not a SQLite database file" + uri = Path(os.path.abspath(path)).as_uri() + "?mode=ro" + with closing(sqlite3.connect(uri, uri=True)) as conn: + row = conn.execute("PRAGMA integrity_check").fetchone() + except (OSError, sqlite3.Error) as error: + return f"Invalid SQLite database: {error}" + verdict = row[0] if row else "no result" + return None if verdict == "ok" else f"Integrity check failed: {verdict}" + + def write_backup_snapshot(model): - """A verified, standalone copy of the library beside it; its path (blocking). + """A standalone copy of the library beside it: ``(path, warning)`` (blocking). Raises SafetyBackupFailed with nothing left on disk when the copy cannot - be written or does not verify. The caller removes the returned file. + be written or is corrupt (``snapshot_corruption``). ``warning`` is the + reason a semantic check (``verify_database_file``, workflow documents + not read) refused the copy, or None. The caller removes the 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) + corrupt = snapshot_corruption(temp) + if corrupt is not None: + raise SafetyBackupFailed("The backup is corrupt", corrupt) 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 + return temp, (None if ok else reason) except BaseException: remove_quietly(*(temp + suffix for suffix in SIDECARS)) raise +def _header_text(text): + """A reason as a one-line, latin-1 header value.""" + flat = " ".join(str(text).split())[:300] + return flat.encode("latin-1", "replace").decode("latin-1") + + class BackupRestoreRoutesMixin: """Mixin providing the backup and restore endpoints.""" @@ -209,13 +251,8 @@ class BackupRestoreRoutesMixin: # ── 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), - ) + """Download a copy of the prompts database (WAL included).""" + return await self._stream_backup(request) async def _stream_backup(self, request): model = self.db.model @@ -223,8 +260,21 @@ class BackupRestoreRoutesMixin: return web.json_response( {"success": False, "error": "Database file not found"}, status=404 ) + made = await self._run_held( + BACKUP_JOB_KIND, + BACKUP_HOLD_TITLE, + "guarded", + lambda: self._make_backup_snapshot(model), + ) + if isinstance(made, web.StreamResponse): + return made # refused, or the snapshot failed + temp, warning = made + return await self._send_snapshot(request, temp, warning) + + async def _make_backup_snapshot(self, model): + """``(path, warning)``, or the error response (inside the hold).""" try: - temp = await self._run_in_executor(write_backup_snapshot, model) + temp, warning = 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( @@ -240,6 +290,12 @@ class BackupRestoreRoutesMixin: }, status=500, ) + if warning is not None: + self.logger.warning(f"Backup sent with a failed check: {warning}") + return temp, warning + + async def _send_snapshot(self, request, temp, warning): + """Stream the snapshot (outside the hold), then close and remove it.""" timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") handle = None try: @@ -252,6 +308,7 @@ class BackupRestoreRoutesMixin: f'attachment; filename="prompts_backup_{timestamp}.db"' ), "Content-Length": str(size), + VERIFY_HEADER: "ok" if warning is None else _header_text(warning), } ) await response.prepare(request) @@ -264,7 +321,7 @@ class BackupRestoreRoutesMixin: self._last_backup_download = { "at": datetime.datetime.now(datetime.timezone.utc).isoformat(), "bytes": size, - "verified": True, + "verified": warning is None, } return response finally: @@ -386,12 +443,19 @@ class BackupRestoreRoutesMixin: ) async def _prune_after(self, db_path, *, failed=False): + """Prune old safety backups; a failure is only logged (never raises). + + It runs after a restore replaced the library, or after one failed: + either way the restore's outcome stands. + """ 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") + except Exception as error: + when = "failed" if failed else "completed" + self.logger.warning( + f"Could not prune safety backups after a {when} restore: " + f"{type(error).__name__}" + ) return if removed: self.logger.info(f"Removed {len(removed)} old safety backups") @@ -405,18 +469,75 @@ class BackupRestoreRoutesMixin: await self._prune_after(model.db_path, failed=True) raise - async def _switch_to_restored(self, db_path): - """Open the restored library as ``self.db`` and close the previous one. + async def _after_publish(self, db_path): + """Everything after the copy replaced the library: best effort, logged. - Its connections would otherwise linger (a file lock on Windows). The - health report and every held restore preview described the old - library, so they are dropped (the stage files off the event loop). + The restore is done once the copy succeeded, so nothing here fails it + (#335 review I1). The library is adopted as ``self.db`` (the previous + object, which points at the same file, stays when a new one cannot + be opened), the generation is bumped (so the health report and + reviewed plans of the old library are dropped), held restore + previews are dropped and the previous object's connections closed. """ + await self._prune_after(db_path) previous = self.db - self.db = PromptDatabase(db_path) - self._database_replaced() - await self._run_in_executor(self._restore_plans.discard_all) - await self._run_in_executor(previous.close_all) + try: + self.db = PromptDatabase(db_path) + except Exception: + self.logger.exception("Could not reopen the restored library") + try: + self._database_replaced() + except Exception: + self.logger.exception("Could not drop the old library's reports") + try: + await self._discard_held(self._restore_plans.pop_all()) + except Exception: + self.logger.exception("Could not remove the old restore previews") + if previous is not self.db: + try: + # Its connections would otherwise linger (a lock on Windows). + await self._run_in_executor(previous.close_all) + except Exception: + self.logger.exception("Could not close the previous library") + + async def _discard_held(self, held): + """Remove the stage folders of dropped store entries (off the loop).""" + for entry in held: + try: + await self._run_in_executor(entry.staged.discard) + except Exception: + self.logger.exception("Could not remove a restore stage") + + async def _drop_expired_restores(self): + """Drop expired previews from the store (on the loop), then their files.""" + await self._discard_held(self._restore_plans.pop_expired()) + + def _schedule_restore_expiry(self): + """Remove each preview's stage when it expires, without a request. + + A loop timer (``call_later``), so a server that stops with previews + held leaves nothing pending; the startup sweep removes the folders. + """ + store = self._restore_plans + expires = store.next_expiry() + if expires is None: + return + handle = getattr(self, "_restore_expiry_handle", None) + if handle is not None: + handle.cancel() + loop = asyncio.get_running_loop() + + def fire(): + self._restore_expiry_handle = None + task = loop.create_task(self._expire_restores()) + self._restore_expiry_task = task + + delay = max(0.0, expires - store.now()) + self._restore_expiry_handle = loop.call_later(delay, fire) + + async def _expire_restores(self): + await self._drop_expired_restores() + self._schedule_restore_expiry() # ── legacy one-step restore ─────────────────────────────────────── @@ -461,9 +582,15 @@ class BackupRestoreRoutesMixin: backup_path = await self._publish(model, staged) finally: await self._run_in_executor(staged.discard) - await self._prune_after(model.db_path) - await self._switch_to_restored(model.db_path) - prompt_count = self.db.model.get_database_info().get("total_prompts", 0) + # The library is replaced: from here on the restore succeeded. + await self._after_publish(model.db_path) + try: + prompt_count = await self._run_in_executor( + lambda: self.db.model.get_database_info().get("total_prompts", 0) + ) + except Exception: + self.logger.exception("Could not count the restored prompts") + prompt_count = staged.summary.get("prompts", 0) message = f"Database restored successfully. Found {prompt_count} prompts." body = { "success": True, @@ -497,62 +624,100 @@ class BackupRestoreRoutesMixin: 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", - ) + await self._drop_expired_restores() + if self._restore_plans.full(): + return _too_many_reviews() 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) + 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 + return await self._review_upload(db, identity, temp_path, uploads) 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", + + async def _stage_for_review(self, db, temp_path, uploads): + """Stage ``temp_path`` in the executor; a cancel discards the result. + + A cancelled handler (shutdown, a dropped client) cannot stop the + worker thread: the stage it finishes is removed when it is done, and + so is the upload it reads (``_orphan_stage_cleanup``). + """ + loop = asyncio.get_running_loop() + future = loop.run_in_executor(None, db.model.stage_restore, temp_path) + try: + return await asyncio.shield(future) + except asyncio.CancelledError: + orphans = getattr(self, "_orphan_stages", None) + if orphans is None: + orphans = self._orphan_stages = set() + entry = (future, tuple(uploads)) + orphans.add(entry) + uploads.clear() # still being read: removed with the stage + future.add_done_callback( + lambda _f: loop.create_task(self._orphan_stage_cleanup()) ) - token = store.issue(staged, fingerprint, identity, live) - held = store.peek(token) + raise + + async def _orphan_stage_cleanup(self): + """Remove what cancelled previews staged, once their workers are done.""" + for entry in list(getattr(self, "_orphan_stages", ()) or ()): + future, uploads = entry + if not future.done(): + continue + self._orphan_stages.discard(entry) + if not future.cancelled() and future.exception() is None: + await self._discard_held([_Orphan(future.result())]) + await self._remove_uploads(uploads) + + async def _review_upload(self, db, identity, temp_path, uploads=None): + """Check, stage and compare an upload; the preview response. + + Only the read of the live library takes a lease (#335 review M2): + the upload and the staging never touch it, and a lease held for a + long staging would make a restore's drain fail. A stage that is not + handed to the store (a refusal, an error, a cancel) is removed. + """ + uploads = [] if uploads is None else uploads # the caller removes them + reason = await self._upload_problem(temp_path) + if reason is not None: + return _failure(reason, 400, "invalid-backup") + staged = None + issued = False + try: + try: + staged = await self._stage_for_review(db, temp_path, uploads) + except ValueError as invalid: + return _failure(str(invalid), 400, "invalid-backup") + try: + async with self._jobs.lease("Restore preview"): + live, fingerprint = await self._run_in_executor( + self._live_fingerprint, db + ) + except JobBusy as busy: + return busy_response(busy) + store = self._restore_plans + if identity != self._database_identity(self.db): + return _failure( + "The database was replaced while the restore was reviewed. " + "Review again.", + 409, + "storage-changed", + ) + if store.full(): + return _too_many_reviews() + token = store.issue(staged, fingerprint, identity, live) + issued = True + self._schedule_restore_expiry() + finally: + if staged is not None and not issued: + await self._discard_held([_Orphan(staged)]) + held = store.entry(token) name = os.path.basename(db.model.db_path) return web.json_response( { @@ -576,9 +741,11 @@ class BackupRestoreRoutesMixin: 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: + held = self._restore_plans.pop(token) + if held is None: + await self._drop_expired_restores() return _failure("Unknown restore review.", 410, "plan-expired") + await self._discard_held([held]) return web.json_response({"success": True}) def _peek_restore(self, token): @@ -612,13 +779,14 @@ class BackupRestoreRoutesMixin: 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) + await self._drop_expired_restores() held, refused = self._peek_restore(token) if refused is not None: + await self._drop_expired_restores() # one peek found expired return refused db = self.db if held.db_identity != self._database_identity(db): - await self._run_in_executor(self._restore_plans.discard, token) + await self._discard_held([self._restore_plans.pop(token) or held]) return _failure( "The database was replaced since the review. Review again.", 409, @@ -633,11 +801,12 @@ class BackupRestoreRoutesMixin: 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) + await self._discard_held([self._restore_plans.pop(token) or held]) 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: + await self._drop_expired_restores() return refused if db is not self.db: return _failure( @@ -655,7 +824,7 @@ class BackupRestoreRoutesMixin: cancellable=False, scope={}, groups=RESTORE_GROUPS, - runner=lambda job: self._run_restore(job, db, held), + runner=lambda job: self._run_restore(job, db, held, token), ) except JobBusy as busy: return busy_response(busy) # the token stays usable @@ -663,13 +832,20 @@ class BackupRestoreRoutesMixin: 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): + async def _run_restore(self, job, db, held, token=None): from .maintenance import _fail_job staged = held.staged + keep_stage = False try: job.update(status_line="Waiting for other requests to finish") if not await self._jobs.drain_leases(): + if token is not None: + # Nothing ran: the reviewed stage stays usable (#335 M2). + self._restore_plans.give_back(token, held) + keep_stage = True + self._schedule_restore_expiry() + return _fail_job(job, "busy", REVIEW_KEPT) 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) @@ -697,20 +873,37 @@ class BackupRestoreRoutesMixin: ) self.logger.error(f"Restore failed: {error}") return _fail_job(job, "restore-failed", NOT_RESTORED + kept) - await self._prune_after(model.db_path) - await self._switch_to_restored(model.db_path) - 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) + return _fail_job(job, None, NOT_RESTORED) finally: - await self._run_in_executor(staged.discard) + if not keep_stage: + await self._discard_held([_Orphan(staged)]) + # The copy replaced the library: the restore is done (#335 review I1). + # Every later step only logs a failure. + await self._after_publish(model.db_path) + 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) + + +class _Orphan: + """A stage with no store entry, in the shape ``_discard_held`` takes.""" + + def __init__(self, staged): + self.staged = staged + + +def _too_many_reviews(): + return _failure( + "Two restores are already being reviewed. Discard one first.", + 429, + "too-large", + ) diff --git a/py/api/restore_plans.py b/py/api/restore_plans.py index e15b9de8..7181c423 100644 --- a/py/api/restore_plans.py +++ b/py/api/restore_plans.py @@ -12,12 +12,17 @@ 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. + +The store is only read and changed on the event loop (it has no lock). Its +entries own stage folders, whose removal is blocking: every method that +drops an entry returns it (or keeps it in ``take_doomed``), and the caller +removes the folder in the executor (``HeldRestore.staged.discard``). """ import secrets import time from dataclasses import dataclass -from typing import Any, Callable, Dict, Optional +from typing import Any, Callable, Dict, List, Optional from .maintenance_plans import PlanExpired, PlanUsed @@ -42,15 +47,23 @@ class HeldRestore: class RestorePlanStore: def __init__( self, + clock: Optional[Callable[[], float]] = None, 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] = {} + self._doomed: List[HeldRestore] = [] + + @property + def ttl(self) -> float: + return self._ttl + + def now(self) -> float: + return self._clock() def __len__(self): return len(self._held) @@ -79,7 +92,7 @@ class RestorePlanStore: if held is None: raise PlanExpired() if held.expires_at <= self._clock(): - self.discard(token) + self._doomed.append(self._held.pop(token)) raise PlanExpired() return held @@ -92,23 +105,35 @@ class RestorePlanStore: 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 give_back(self, token, held) -> None: + """Hold a taken stage again under its token (its job did not apply it).""" + self._used.pop(token, None) + self._held[token] = held - def drop_expired(self) -> int: - """Remove every expired stage (blocking file removal); the count.""" + def entry(self, token) -> Optional[HeldRestore]: + """The held entry, expired or not (for the response that issued it).""" + return self._held.get(token) + + def pop(self, token) -> Optional[HeldRestore]: + """Drop the token; its entry (to discard) or None when it was not held.""" + return self._held.pop(token, None) + + def pop_expired(self) -> List[HeldRestore]: + """Drop every expired entry; them, with any ``peek`` found expired.""" 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) + return [self._held.pop(token) for token in expired] + self.take_doomed() - def discard_all(self) -> None: - """Remove every held stage (the database was replaced or shutdown).""" - for token in list(self._held): - self.discard(token) + def pop_all(self) -> List[HeldRestore]: + """Drop every entry (the library was replaced); them, to discard.""" + held = list(self._held.values()) + self.take_doomed() + self._held.clear() + return held + + def take_doomed(self) -> List[HeldRestore]: + doomed, self._doomed = self._doomed, [] + return doomed + + def next_expiry(self) -> Optional[float]: + """When the soonest held entry expires (clock time), or None.""" + return min((held.expires_at for held in self._held.values()), default=None) diff --git a/tests/test_api_backup_restore.py b/tests/test_api_backup_restore.py index 5317594d..e9b451cf 100644 --- a/tests/test_api_backup_restore.py +++ b/tests/test_api_backup_restore.py @@ -132,9 +132,9 @@ class TestBackupRoute(BackupRestoreTestCase): self.assertTrue(disposition.rstrip('"').endswith(".db")) body = await resp.read() - # The handler removes its temp file as it leaves the registry hold, - # which can be just after the client holds the last byte. - while self.api._jobs.active("backup") is not None: + # The handler removes its temp file after the last byte is sent, + # which can be just after the client holds it: wait on that state. + while self.api._jobs.active("backup") is not None or os.path.exists(created[0]): await asyncio.sleep(0) downloaded = self._write("downloaded.db", body) self.assertEqual( diff --git a/tests/test_diagnostics_config.py b/tests/test_diagnostics_config.py index 405ddc16..2fca7ac4 100644 --- a/tests/test_diagnostics_config.py +++ b/tests/test_diagnostics_config.py @@ -177,8 +177,9 @@ class TestAdminEndpointsUseConfigPath(unittest.TestCase): try: stub = self._make_api_stub(db_file) stub.db.model = real_db.model - snapshot = write_backup_snapshot(stub.db.model) + snapshot, warning = write_backup_snapshot(stub.db.model) try: + self.assertIsNone(warning) self.assertEqual(os.path.dirname(snapshot), folder) with open(snapshot, "rb") as handle: self.assertTrue(handle.read(16) == b"SQLite format 3\x00") diff --git a/tests/test_restore_plans.py b/tests/test_restore_plans.py index 9530d984..5ad20eaa 100644 --- a/tests/test_restore_plans.py +++ b/tests/test_restore_plans.py @@ -136,12 +136,15 @@ class RestoreTestCase(AdminAPITestCase): 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 left its - hold, so the registry is read only once the hold says it is done. + 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: + 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 @@ -502,9 +505,12 @@ class TestLegacyRestore(RestoreTestCase): 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.object( - PromptModel, "verify_database_file", return_value=(False, "bad") + 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) diff --git a/tests/test_restore_review_fixes.py b/tests/test_restore_review_fixes.py new file mode 100644 index 00000000..a5f7014b --- /dev/null +++ b/tests/test_restore_review_fixes.py @@ -0,0 +1,446 @@ +"""PR #335 final review fixes (#249 System S8): I1, M1 to M7, N1 and N3. + +- I1: once the stage replaced the library, the restore is done: later steps + (pruning, the switch) only log, and the job and legacy route say so. +- M1: a backup download holds the registry only while its snapshot is made. +- M2: a restore preview leases only its live read; a busy drain keeps the + reviewed stage. +- M3: the reviewed-preview store changes only on the event loop. +- M4: a document whose body is intact but whose key or length is damaged is + re-keyed, not dropped. +- M5: a backup download is refused only for a corrupt snapshot; a semantic + check failure downloads, logged and reported. +- M6: a cancelled preview removes what it staged. +- M7 / N1 / N3: repair reads one row at a time; images counted once; the + fingerprint never commits a caller's transaction. + +Ordering is driven by events and job state, never by timers. +""" + +import asyncio +import hashlib +import os +import sqlite3 +import threading +import zlib +from contextlib import closing +from unittest.mock import patch + +from aiohttp import web + +from database import restore_stage +from database import workflow_documents as wd +from database.managed_connection import ManagedCursor +from database.models import PromptModel +from database.operations import PromptDatabase +from py.api import admin as admin_mod +from py.api import backup_restore, jobs + +try: + from tests.test_restore_plans import RestoreTestCase, _texts +except ImportError: # discovered with tests/ as the top-level directory + from test_restore_plans import RestoreTestCase, _texts + +WAIT = 60 + + +class TestPostPublishSteps(RestoreTestCase): + """I1: a failure after the copy never says the library was not changed.""" + + async def test_job_finishes_when_pruning_fails_after_the_copy(self): + self.save("old") + old_db = self.api.db + generation = self.api._db_generation + body = await self.preview(self.upload(["restored"])) + with ( + patch.object(admin_mod, "_prune_safety_backups", side_effect=OSError("x")), + self.assertLogs(self.api.logger, "WARNING") as logs, + ): + started = await self.apply(body["token"]) + job = await self.finish(started["job"]["id"]) + self.assertEqual(job.state, "finished", job.log) + self.assertIn("Restored 1 prompt", job.status_line) + self.assertNotIn("not changed", "\n".join(job.log)) + self.assertIsNot(self.api.db, old_db) + self.assertEqual(self.api._db_generation, generation + 1) + self.assertEqual(_texts(self.db_path), ["restored"]) + self.assertIn("prune", "\n".join(logs.output).lower()) + + async def test_job_adopts_the_library_when_closing_the_old_one_fails(self): + self.save("old") + old_db = self.api.db + body = await self.preview(self.upload(["restored"])) + with patch.object(old_db, "close_all", side_effect=RuntimeError("close")): + started = await self.apply(body["token"]) + job = await self.finish(started["job"]["id"]) + self.assertEqual(job.state, "finished", job.log) + self.assertIsNot(self.api.db, old_db) + self.assertEqual( + [p["text"] for p in self.api.db.search_prompts()], ["restored"] + ) + + async def test_job_finishes_when_opening_the_new_object_fails(self): + self.save("old") + generation = self.api._db_generation + body = await self.preview(self.upload(["restored"])) + with patch.object( + backup_restore, "PromptDatabase", side_effect=RuntimeError("open") + ): + started = await self.apply(body["token"]) + job = await self.finish(started["job"]["id"]) + self.assertEqual(job.state, "finished", job.log) + # The old object still points at the same, restored file. + self.assertEqual( + [p["text"] for p in self.api.db.search_prompts()], ["restored"] + ) + self.assertEqual(self.api._db_generation, generation + 1) + + async def test_legacy_restore_succeeds_when_pruning_fails_after_the_copy(self): + self.save("old") + old_db = self.api.db + with patch.object(admin_mod, "_prune_safety_backups", side_effect=OSError("x")): + resp = await self.client.post( + "/prompt_manager/restore", data=self.form(self.upload(["restored"])) + ) + body = await resp.json() + self.assertEqual(resp.status, 200, body) + self.assertEqual( + body["message"], "Database restored successfully. Found 1 prompts." + ) + self.assertIsNot(self.api.db, old_db) + + +class TestBackupHoldScope(RestoreTestCase): + """M1 and M5: the hold covers the snapshot only; semantic issues warn.""" + + async def test_hold_ends_before_the_first_byte_is_sent(self): + self.save("keep") + seen = [] + real = web.StreamResponse.prepare + + async def prepare(response, request): + seen.append(self.api._jobs.active("backup")) + return await real(response, request) + + with patch.object(web.StreamResponse, "prepare", prepare): + resp, data = await self.download_backup() + self.assertEqual(resp.status, 200) + self.assertTrue(seen) + self.assertEqual(set(seen), {None}) # no hold while bytes are sent + [job] = self.api._jobs.list() + self.assertEqual((job.kind, job.state), ("backup", "finished")) + self.assertTrue(data.startswith(b"SQLite format 3\x00")) + + async def test_a_corrupt_snapshot_is_refused(self): + self.save("keep") + with patch.object( + backup_restore, "snapshot_corruption", return_value="Integrity check failed" + ): + 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-")], [] + ) + [job] = self.api._jobs.list() + self.assertEqual(job.state, "failed") + + async def test_a_semantic_check_failure_still_downloads_and_says_so(self): + self.save("keep") + reason = "Image pair association has a different positive prompt owner" + with ( + patch.object( + PromptModel, "verify_database_file", return_value=(False, reason) + ), + self.assertLogs(self.api.logger, "WARNING") as logs, + ): + resp, data = await self.download_backup() + self.assertEqual(resp.status, 200) + self.assertEqual(resp.headers["X-PromptManager-Verify"], reason) + self.assertTrue(data.startswith(b"SQLite format 3\x00")) + self.assertIn(reason, "\n".join(logs.output)) + status = await (await self.client.get("/prompt_manager/backup/status")).json() + self.assertFalse(status["last_download"]["verified"]) + + async def test_a_verified_download_says_ok(self): + self.save("keep") + resp, _data = await self.download_backup() + self.assertEqual(resp.headers["X-PromptManager-Verify"], "ok") + + +class TestPreviewLease(RestoreTestCase): + """M2: upload and staging take no lease; a busy drain keeps the stage.""" + + async def test_staging_holds_no_lease(self): + leases = [] + real = PromptModel.stage_restore + + def stage(model, path): + leases.append(dict(self.api._jobs.leases)) + return real(model, path) + + with patch.object(PromptModel, "stage_restore", stage): + await self.preview(self.upload(["restored"])) + self.assertEqual(leases, [{}]) + + async def test_a_busy_drain_keeps_the_reviewed_stage(self): + self.save("live") + body = await self.preview(self.upload(["restored"])) + lease = self.api._jobs.lease("Stuck preview") + await lease.__aenter__() + try: + with patch.object(jobs, "LEASE_DRAIN_SECONDS", 0): + started = await self.apply(body["token"]) + job = await self.finish(started["job"]["id"]) + finally: + await lease.__aexit__(None, None, None) + self.assertEqual(job.state, "failed") + self.assertEqual(job.last_event["code"], "busy") + self.assertIn("review is kept", job.status_line) + self.assertEqual(len(self.stages()), 1) + # The same review applies once the requests are done. + job = await self.finish((await self.apply(body["token"]))["job"]["id"]) + self.assertEqual(job.state, "finished", job.log) + self.assertEqual(_texts(self.db_path), ["restored"]) + + +class TestStoreOnTheLoop(RestoreTestCase): + """M3: entries leave the store on the loop; only file removal is threaded.""" + + async def test_expired_entries_leave_the_store_before_their_files_go(self): + clock = [1000.0] + store = self.api._restore_plan_store = backup_restore.RestorePlanStore( + clock=lambda: clock[0] + ) + body = await self.preview(self.upload(["restored"])) + held = store.peek(body["token"]) + entered, release = threading.Event(), threading.Event() + real = held.staged.discard + + def slow_discard(): + entered.set() + release.wait(WAIT) + real() + + held.staged.discard = slow_discard + clock[0] += 10_000 + cleanup = asyncio.ensure_future(self.api._drop_expired_restores()) + loop = asyncio.get_running_loop() + self.assertTrue(await loop.run_in_executor(None, entered.wait, WAIT)) + # The removal runs on a worker; the store already forgot the entry. + self.assertEqual(len(store), 0) + second = await self.preview(self.upload(["other"], "other.db")) + self.assertEqual(len(store), 1) + release.set() + await cleanup + self.assertEqual(len(self.stages()), 1) + store.peek(second["token"]) + + async def test_an_expired_preview_is_removed_without_another_request(self): + clock = [1000.0] + self.api._restore_plan_store = backup_restore.RestorePlanStore( + clock=lambda: clock[0], ttl=0 + ) + await self.preview(self.upload(["restored"])) + # The expiry timer fires without another request: wait on its state. + while self.stages() or len(self.api._restore_plans): + await asyncio.sleep(0) + self.assertEqual(self.stages(), []) + + +class TestCancelledPreview(RestoreTestCase): + """M6: a cancelled preview leaves no stage behind.""" + + async def test_cancel_while_reading_the_live_library(self): + path = os.path.join(self.tmpdir, "uploads", "up.db") + self.upload(["restored"], "up.db") + entered, release = threading.Event(), threading.Event() + real = self.api._live_fingerprint + + def slow(db): + entered.set() + release.wait(WAIT) + return real(db) + + with patch.object(self.api, "_live_fingerprint", side_effect=slow): + task = asyncio.ensure_future( + self.api._review_upload( + self.api.db, self.api._database_identity(self.api.db), path + ) + ) + loop = asyncio.get_running_loop() + self.assertTrue(await loop.run_in_executor(None, entered.wait, WAIT)) + self.assertEqual(len(self.stages()), 1) + task.cancel() + release.set() + with self.assertRaises(asyncio.CancelledError): + await task + self.assertEqual(self.stages(), []) + self.assertEqual(len(self.api._restore_plans), 0) + + async def test_cancel_while_staging(self): + path = os.path.join(self.tmpdir, "uploads", "up.db") + self.upload(["restored"], "up.db") + entered, release = threading.Event(), threading.Event() + real = PromptModel.stage_restore + done = threading.Event() + + def slow(model, src): + entered.set() + release.wait(WAIT) + try: + return real(model, src) + finally: + done.set() + + with patch.object(PromptModel, "stage_restore", slow): + task = asyncio.ensure_future( + self.api._review_upload( + self.api.db, self.api._database_identity(self.api.db), path + ) + ) + loop = asyncio.get_running_loop() + self.assertTrue(await loop.run_in_executor(None, entered.wait, WAIT)) + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + release.set() + self.assertTrue(await loop.run_in_executor(None, done.wait, WAIT)) + # The worker's finished stage is removed without another request. + while self.api._orphan_stages or self.stages(): + await asyncio.sleep(0) + self.assertEqual(self.stages(), []) + + +class TestRekeyDocuments(RestoreTestCase): + """M4: an intact body with a damaged key or length is re-keyed.""" + + def library(self, damage): + path = os.path.join(self.tmpdir, "uploads", "d.db") + self.upload(["x"], "d.db") + raw = b'{"nodes": [1, 2, 3]}' + digest = hashlib.sha256(raw).digest() + with closing(sqlite3.connect(path)) as conn: + sha, length = damage(digest, len(raw)) + doc = conn.execute( + "INSERT INTO workflow_documents (sha256, byte_length, data) " + "VALUES (?,?,?)", + (sha, length, zlib.compress(raw)), + ).lastrowid + prompt_id = conn.execute("SELECT id FROM prompts").fetchone()[0] + conn.execute( + "INSERT INTO generated_images (prompt_id, image_path, filename, " + "workflow_document_id) VALUES (?, '/x/a.png', 'a.png', ?)", + (prompt_id, doc), + ) + conn.commit() + return path, raw.decode() + + def staged_texts(self, staged): + with closing(sqlite3.connect(staged.path)) as conn: + return [ + wd.read_stored_text(conn, None, row[0]) + for row in conn.execute( + "SELECT workflow_document_id FROM generated_images" + ) + ] + + def test_a_damaged_key_is_recomputed(self): + def flip(digest, length): + bad = bytearray(digest) + bad[0] ^= 1 + return bytes(bad), length + + path, text = self.library(flip) + staged = self.api.db.model.stage_restore(path) + try: + self.assertEqual( + staged.documents, {"repaired": 1, "cleared": 0, "images": 0} + ) + self.assertEqual(self.staged_texts(staged), [text]) + finally: + staged.discard() + + def test_a_damaged_length_is_recomputed(self): + path, text = self.library(lambda digest, length: (digest, length + 7)) + staged = self.api.db.model.stage_restore(path) + try: + self.assertEqual(staged.documents["repaired"], 1) + self.assertEqual(self.staged_texts(staged), [text]) + finally: + staged.discard() + + def test_a_key_that_another_document_holds_is_merged_into_it(self): + path, text = self.library(lambda digest, length: (b"\0" * 32, length)) + with closing(sqlite3.connect(path)) as conn: + raw = text.encode() + conn.execute( + "INSERT INTO workflow_documents (sha256, byte_length, data) " + "VALUES (?,?,?)", + (hashlib.sha256(raw).digest(), len(raw), zlib.compress(raw)), + ) + conn.commit() + staged = self.api.db.model.stage_restore(path) + try: + self.assertEqual(staged.documents["repaired"], 1) + self.assertEqual(self.staged_texts(staged), [text]) + with closing(sqlite3.connect(staged.path)) as conn: + self.assertEqual( + conn.execute("SELECT COUNT(*) FROM workflow_documents").fetchone(), + (1,), + ) + finally: + staged.discard() + + +class TestRepairReadsOneRowAtATime(RestoreTestCase): + """M7: repair never iterates a managed cursor (128-row prefetch).""" + + def test_repair_uses_fetchone(self): + upload = os.path.join(self.tmpdir, "uploads", "dmg.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: + text = wd.read_stored_text( + conn, + None, + conn.execute( + "SELECT workflow_document_id FROM generated_images WHERE id=?", + (image,), + ).fetchone()[0], + ) + conn.execute( + "INSERT INTO generated_images (prompt_id, image_path, filename, " + "workflow_data) VALUES (?, '/x/b.png', 'b.png', ?)", + (prompt_id, text), + ) + conn.execute("UPDATE workflow_documents SET data=x'00'") + try: + with patch.object( + ManagedCursor, "fetchmany", side_effect=AssertionError("prefetch") + ): + result = wd.repair_damaged_documents(conn) + self.assertEqual(result["repaired"], 1) + finally: + src.close_all() + + +class TestFingerprintTransaction(RestoreTestCase): + """N3: the fingerprint never commits a transaction it did not open.""" + + def test_a_callers_transaction_is_left_open(self): + conn = self.api.db.model.get_connection() + conn.execute("BEGIN") + conn.execute("INSERT INTO prompts (text, hash) VALUES ('pending', 'p')") + restore_stage.live_fingerprint(conn) + self.assertTrue(conn.in_transaction) + conn.rollback() + self.assertEqual([p["text"] for p in self.api.db.search_prompts()], [])