diff --git a/TESTING.md b/TESTING.md index 1409f3c..fd33555 100644 --- a/TESTING.md +++ b/TESTING.md @@ -3191,6 +3191,66 @@ are closed before `os.replace`; renames stay inside one folder; folder identity folder are refused. The browser test server closes its own log handlers before its temp root (which now holds the logs) is removed. +### Backup and restore contract (#249 System S8) + +A restore is staged, then published (`database/restore_stage.py`). The upload is +streamed to `.pm-restore-upload-*.tmp` and staged in a `.pm-restore-*` folder, both in +the library's own folder (never the system temp folder); free space there must hold +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 (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`: +standalone, fsynced, 0600, not verified because the database being replaced may be the +damaged one), copies the stage over the live pages through the backup API (#226: no +file copy) and checkpoints the WAL with TRUNCATE, waiting up to 30 s for readers and +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, 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 (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. 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, 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`) 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. + ## Workspace data fields for the concept pages (#249) These fields are additive. The legacy `/prompt_manager/images/output` listing (no diff --git a/database/models.py b/database/models.py index 78c0fc9..e49ff94 100644 --- a/database/models.py +++ b/database/models.py @@ -9,7 +9,6 @@ import math import sqlite3 import os import threading -import tempfile import weakref from contextlib import closing from pathlib import Path @@ -20,6 +19,7 @@ from .managed_connection import ( await_closed, retire_connections, ) +from . import restore_stage from .library_schema import library_schema_error, migrate_library_state from .image_favorites import image_favorites_schema_error, migrate_image_favorites from .annotation_schema import annotation_schema_error, migrate_annotation_state @@ -39,7 +39,6 @@ from .workflow_documents import ( from .execution_history_schema import ( history_schema_error, migrate_execution_history, - rotate_history_epoch, ) # Import logging system @@ -555,7 +554,10 @@ class PromptModel: "WHERE pp.positive_prompt_id IS NOT gi.prompt_id LIMIT 1" ).fetchone(): return "Image pair association has a different positive prompt owner" - for row in conn.execute("SELECT provenance FROM image_prompt_pairs"): + # fetchone, not iteration: a managed connection's cursor prefetches + # ITER_CHUNK rows when iterated, and each provenance may be ~4 MB. + cursor = conn.execute("SELECT provenance FROM image_prompt_pairs") + while (row := cursor.fetchone()) is not None: error = pair_provenance_error(row[0]) if error: return error @@ -1132,7 +1134,7 @@ class PromptModel: @staticmethod def verify_database_file( - path: str, *, require_library_state: bool = False + path: str, *, require_library_state: bool = False, documents: bool = True ) -> Tuple[bool, str]: """ Check that a file is a healthy prompts database before restoring it. @@ -1146,6 +1148,9 @@ class PromptModel: database has (Library state, image favorites, saved views, image annotations and the review queue), as a migrated restore stage must; old backups may lack them. + documents: Read every workflow document and refuse a damaged + one. A restore checks its upload without it and repairs + damaged documents on the stage (``database/restore_stage.py``). Returns: (ok, reason): reason is "ok" when the file is acceptable, otherwise @@ -1221,7 +1226,7 @@ class PromptModel: if review_error: return False, review_error documents_error = workflow_documents_schema_error( - conn, allow_missing=not require_library_state, documents=True + conn, allow_missing=not require_library_state, documents=documents ) if documents_error: return False, documents_error @@ -1255,17 +1260,43 @@ class PromptModel: finally: source.close() + def stage_restore(self, src_path: str): + """Stage ``src_path`` beside this database for a restore (blocking). + + Returns a ``StagedRestore`` (``database/restore_stage.py``): a + migrated, repaired and verified copy in the library's folder, its + digest and what it holds. Neither the live database nor the source + is touched. Raises ValueError when either does not verify. + """ + return restore_stage.stage_restore(PromptModel, self.db_path, src_path) + + def publish_staged(self, staged) -> str: + """Replace the live database with ``staged``; return the safety backup. + + See ``restore_stage.publish_staged``: the staged bytes are rechecked, + the current database is kept as ``.backup_`` and the + WAL is truncated afterwards. + """ + backup_path = restore_stage.publish_staged(self, staged) + self.logger.info("Database restored from a staged copy") + from .execution_history import notify_history_restore + + notify_history_restore(self.db_path) + return backup_path + def restore_from_file(self, src_path: str) -> str: """ Replace the live database with the file at src_path. - A disposable SQLite backup of the source is migrated and verified first. - Only then is the current database backed up next to itself as - ``.backup_`` (WAL content included), and the - staged pages replace the live pages through SQLite's backup API. - Migration failures never replace live content or modify the source. - An online replacement failure leaves the live database intact and - propagates; the completed safety backup remains available. + ``stage_restore`` then ``publish_staged``: a disposable copy of the + source is migrated, repaired and verified in the library's folder + first. Only then is the current database backed up next to itself as + ``.backup_`` (WAL content included), + and the staged pages replace the live pages through SQLite's backup + API. Staging failures never replace live content or modify the + source. An online replacement failure leaves the live database intact + and propagates; the completed safety backup remains available. The + stage is removed either way. Args: src_path: SQLite file to restore @@ -1275,55 +1306,14 @@ class PromptModel: there was no database to back up. Raises: - RuntimeError: If the previous database could not be backed up + ValueError: If the source or its migrated copy does not verify + SafetyBackupFailed: If the previous database could not be backed up sqlite3.OperationalError: If SQLite refuses online replacement """ - ok, reason = self.verify_database_file(src_path) - if not ok: - raise ValueError(reason) - with tempfile.TemporaryDirectory(prefix="prompt-manager-restore-") as directory: - staged_path = str(Path(directory) / "staged.db") - source = sqlite3.connect( - Path(src_path).resolve().as_uri() + "?mode=ro", uri=True - ) - try: - with closing(sqlite3.connect(staged_path)) as destination: - source.backup(destination) - finally: - source.close() - staged = None - try: - staged = PromptModel(staged_path) - rotate_history_epoch(staged.get_connection()) - ok, reason = self.verify_database_file( - staged_path, require_library_state=True - ) - if not ok: - raise ValueError(reason) - staged.close_all() - backup_path = "" - if os.path.exists(self.db_path): - timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") - backup_path = f"{self.db_path}.backup_{timestamp}" - if not self.backup_database(backup_path): - raise RuntimeError("Could not back up the current database") - self.logger.info(f"Current database backed up to: {backup_path}") - try: - self._replace_contents_from(staged_path) - except sqlite3.OperationalError as exc: - exc._prompt_manager_restore_backup_created = bool(backup_path) - raise - with PromptModel._init_lock: - PromptModel._initialized_paths.add(self._schema_key()) - self.logger.info(f"Database restored from {src_path}") - from .execution_history import notify_history_restore - - notify_history_restore(self.db_path) - return backup_path - finally: - if staged is not None: - staged.close_all() - with PromptModel._init_lock: - PromptModel._initialized_paths.discard( - os.path.normcase(os.path.realpath(staged_path)) - ) + staged = self.stage_restore(src_path) + try: + backup_path = self.publish_staged(staged) + finally: + staged.discard() + self.logger.info(f"Database restored from {src_path}") + return backup_path diff --git a/database/restore_stage.py b/database/restore_stage.py new file mode 100644 index 0000000..75d8883 --- /dev/null +++ b/database/restore_stage.py @@ -0,0 +1,365 @@ +"""A restore in two steps: stage the upload, then publish it (#249 System S8). + +**Stage** (``stage_restore``) never touches the live database or the source. +In a new folder ``.pm-restore-*`` in the library's own folder (#318: never +the system temp folder, which may be another disk), it: + +1. checks the source (``verify_database_file`` without reading every + workflow document: a damaged one is repaired below, not refused); +2. copies it with SQLite's backup API (a rollback-journal file and a WAL file + alike) and opens the copy as a ``PromptModel``, which migrates an older + schema; +3. repairs damaged workflow documents (``repair_damaged_documents``), + rotates the history epoch and counts what the copy holds; +4. makes the copy one self-contained file and verifies it fully (every + current table, every document); +5. records the copy's SHA-256, so publishing binds exactly these bytes. + +**Publish** (``publish_staged``) re-hashes the stage (``StageChanged`` if it +differs), keeps the current database as a uniquely named, standalone safety +backup beside it (``database/safety_backup.py``; not verified, because the +database being replaced may be the damaged one), copies the stage over the +live pages through the backup API (no file is swapped under open handles; an +``OperationalError`` there leaves the live data intact, #226), then +checkpoints the WAL with TRUNCATE so the copied pages do not leave a +database-sized ``-wal`` behind. A reader that keeps a snapshot past the busy +wait is logged, not fatal: the restore itself is done. + +Every handle on the stage is closed before its folder is removed (Windows). +``sweep_restore_temps`` removes what an interrupted restore left (startup). +""" + +import hashlib +import logging +import os +import sqlite3 +import tempfile +from contextlib import closing +from dataclasses import dataclass, field +from pathlib import Path +from typing import Dict, Optional + +from . import workflow_documents as wd +from .execution_history_schema import rotate_history_epoch +from .safety_backup import take_safety_backup + +logger = logging.getLogger("prompt_manager.database.restore_stage") + +STAGE_PREFIX = ".pm-restore-" +UPLOAD_PREFIX = ".pm-restore-upload-" +STAGED_NAME = "staged.db" +# How long the post-restore TRUNCATE checkpoint waits for readers to leave. +CHECKPOINT_BUSY_MS = 30_000 +_HASH_CHUNK = 1024 * 1024 +SUMMARY_KEYS = ( + "prompts", + "positives", + "pairs", + "archived_pairs", + "image_links", + "runs", + "saved_views", +) + + +class StageChanged(ValueError): + """The staged file is no longer the bytes that were reviewed.""" + + +@dataclass +class StagedRestore: + """A verified, migrated copy of an upload, ready to publish.""" + + path: str + digest: str + size: int + summary: Dict[str, int] + schema: Dict[str, bool] + documents: Dict[str, int] = field( + default_factory=lambda: {"repaired": 0, "cleared": 0, "images": 0} + ) + + @property + def directory(self) -> str: + return os.path.dirname(self.path) + + def discard(self) -> None: + """Remove the stage folder (idempotent; its handles are all closed).""" + remove_stage_folder(self.directory) + + +def remove_stage_folder(directory: str) -> None: + """Remove a stage folder and its files, logging what cannot be removed.""" + if not directory or not os.path.isdir(directory): + return + for name in os.listdir(directory): + try: + os.remove(os.path.join(directory, name)) + except OSError as error: + logger.warning(f"Could not remove restore stage file {name}: {error}") + try: + os.rmdir(directory) + except OSError as error: + logger.warning(f"Could not remove the restore stage folder: {error}") + + +def file_digest(path: str) -> str: + """``sha256:`` of the file's bytes (blocking).""" + digest = hashlib.sha256() + with open(path, "rb") as handle: + for chunk in iter(lambda: handle.read(_HASH_CHUNK), b""): + digest.update(chunk) + return "sha256:" + digest.hexdigest() + + +def _readonly(path: str) -> sqlite3.Connection: + return sqlite3.connect(Path(path).resolve().as_uri() + "?mode=ro", uri=True) + + +def _schema(path: str) -> list: + """The file's schema objects, to tell whether staging migrated it.""" + with closing(_readonly(path)) as conn: + return sorted( + conn.execute( + "SELECT type, name, tbl_name, sql FROM sqlite_master " + "WHERE name NOT LIKE 'sqlite_%'" + ).fetchall(), + key=repr, + ) + + +def _count(conn: sqlite3.Connection, table: str, where: str = "") -> int: + exists = conn.execute( + "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", (table,) + ).fetchone() + if not exists: + return 0 + # ``table`` and ``where`` are constants of library_summary, never input. + return conn.execute(f"SELECT COUNT(*) FROM {table} {where}").fetchone()[0] # nosec + + +def library_summary(conn: sqlite3.Connection) -> Dict[str, int]: + """What a library holds, for a restore's review (reads only). + + ``positives`` are the Library's prompts (``library_counts``'s + ``positive_prompts``); ``pairs`` the active saved pairs and + ``archived_pairs`` the archived ones; ``image_links`` the indexed image + rows; ``runs`` the recorded executions. + """ + from .library import _ELIGIBLE + + has_state = conn.execute( + "SELECT 1 FROM sqlite_master WHERE type='table' AND " + "name='prompt_library_state'" + ).fetchone() + has_pairs = conn.execute( + "SELECT 1 FROM sqlite_master WHERE type='table' AND name='prompt_pairs'" + ).fetchone() + if has_state and has_pairs: + positives = conn.execute( + "SELECT COUNT(*) FROM prompts p WHERE " + _ELIGIBLE + ).fetchone()[0] + else: + positives = _count(conn, "prompts") + return { + "prompts": _count(conn, "prompts"), + "positives": positives, + "pairs": _count(conn, "prompt_pairs", "WHERE archived=0"), + "archived_pairs": _count(conn, "prompt_pairs", "WHERE archived<>0"), + "image_links": _count(conn, "generated_images"), + "runs": _count(conn, "execution_runs"), + "saved_views": _count(conn, "saved_views"), + } + + +# The newest row of each table a restore would overwrite: with the counts, +# a write after the review (a new run, image, prompt or edit) changes it. +_FINGERPRINT_SQL = ( + ("prompts", "SELECT MAX(id), MAX(updated_at), MAX(last_used_at) FROM prompts"), + ("generated_images", "SELECT MAX(id) FROM generated_images"), + ("prompt_pairs", "SELECT MAX(id), MAX(updated_at) FROM prompt_pairs"), + ("execution_runs", "SELECT COUNT(*), MAX(rowid) FROM execution_runs"), + ("saved_views", "SELECT MAX(rowid) FROM saved_views"), +) + + +def live_fingerprint(conn: sqlite3.Connection): + """``(summary, fingerprint)`` of a live library, read in one snapshot. + + 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 + + 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: + try: + 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() + + +def _copy(src_path: str, staged_path: str) -> None: + source = _readonly(src_path) + try: + with closing(sqlite3.connect(staged_path)) as destination: + source.backup(destination) + finally: + source.close() + + +def _standalone(path: str) -> None: + with closing(sqlite3.connect(path)) as conn: + conn.execute("PRAGMA journal_mode=DELETE").fetchone() + + +def stage_restore(model_cls, db_path: str, src_path: str) -> StagedRestore: + """Stage ``src_path`` beside ``db_path`` (blocking); see the module docstring. + + Raises ValueError with a readable reason when the source or the migrated + stage does not verify; the stage folder is then already removed. + """ + ok, reason = model_cls.verify_database_file(src_path, documents=False) + if not ok: + raise ValueError(reason) + before = _schema(src_path) + with closing(_readonly(src_path)) as conn: + library_state = bool( + conn.execute( + "SELECT 1 FROM sqlite_master WHERE type='table' AND " + "name='prompt_library_state'" + ).fetchone() + ) + folder = os.path.dirname(os.path.abspath(db_path)) + directory = tempfile.mkdtemp(prefix=STAGE_PREFIX, dir=folder) + staged_path = os.path.join(directory, STAGED_NAME) + staged_key = os.path.normcase(os.path.realpath(staged_path)) + try: + _copy(src_path, staged_path) + staged = None + try: + staged = model_cls(staged_path) + conn = staged.get_connection() + documents = wd.repair_damaged_documents(conn) + rotate_history_epoch(conn) + summary = library_summary(conn) + finally: + if staged is not None: + staged.close_all() + _standalone(staged_path) + ok, reason = model_cls.verify_database_file( + staged_path, require_library_state=True + ) + if not ok: + raise ValueError(reason) + result = StagedRestore( + path=staged_path, + digest=file_digest(staged_path), + size=os.path.getsize(staged_path), + summary=summary, + schema={ + "migrated": _schema(staged_path) != before, + "library_state": library_state, + }, + documents=documents, + ) + except BaseException: + remove_stage_folder(directory) + raise + finally: + with model_cls._init_lock: + model_cls._initialized_paths.discard(staged_key) + return result + + +def checkpoint_truncate(conn: sqlite3.Connection) -> bool: + """Checkpoint the WAL and truncate it to 0 bytes; False when a reader held it.""" + previous = conn.execute("PRAGMA busy_timeout").fetchone()[0] + try: + conn.execute(f"PRAGMA busy_timeout = {int(CHECKPOINT_BUSY_MS)}") + busy = conn.execute("PRAGMA wal_checkpoint(TRUNCATE)").fetchone()[0] + finally: + conn.execute(f"PRAGMA busy_timeout = {int(previous)}") + return not busy + + +def publish_staged(model, staged: StagedRestore) -> str: + """Publish ``staged`` over ``model``'s database; return the safety backup path. + + "" when there was no database to back up. Raises StageChanged before + anything is written when the stage's bytes changed since staging, and + SafetyBackupFailed when the current database could not be kept. An + ``OperationalError`` from the page copy propagates with + ``_prompt_manager_restore_backup_created`` set; the live data is intact. + """ + if not os.path.isfile(staged.path) or file_digest(staged.path) != staged.digest: + raise StageChanged("The staged restore changed since it was reviewed") + backup_path = "" + if os.path.exists(model.db_path): + backup_path = take_safety_backup(model, verify=False) + logger.info(f"Current database backed up to: {os.path.basename(backup_path)}") + try: + model._replace_contents_from(staged.path) + except sqlite3.OperationalError as exc: + exc._prompt_manager_restore_backup_created = bool(backup_path) + raise + with model._init_lock: + model._initialized_paths.add(model._schema_key()) + _truncate_wal(model) + return backup_path + + +def _truncate_wal(model) -> Optional[bool]: + """TRUNCATE-checkpoint the live WAL after a publish (never raises).""" + try: + conn = model.get_connection() + with conn.using(): + truncated = checkpoint_truncate(conn) + except sqlite3.Error as error: + logger.warning(f"The database log could not be truncated: {error}") + return None + if not truncated: + logger.warning( + "The database log could not be truncated after the restore: a " + "reader kept it open. It shrinks at the next checkpoint with no " + "reader." + ) + return truncated + + +def sweep_restore_temps(db_path: str) -> list: + """Remove the stages and uploads an interrupted restore left (startup).""" + folder = os.path.dirname(os.path.abspath(db_path)) + removed = [] + try: + names = os.listdir(folder) + except OSError: + return removed + for name in names: + path = os.path.join(folder, name) + if name.startswith(UPLOAD_PREFIX) and os.path.isfile(path): + try: + os.remove(path) + removed.append(name) + except OSError: + continue + elif name.startswith(STAGE_PREFIX) and os.path.isdir(path): + remove_stage_folder(path) + if not os.path.exists(path): + removed.append(name) + return removed diff --git a/database/review_schema.py b/database/review_schema.py index 0370a5e..0aa8318 100644 --- a/database/review_schema.py +++ b/database/review_schema.py @@ -380,15 +380,21 @@ def _structure_error(conn: sqlite3.Connection) -> Optional[str]: def _evidence_error(conn: sqlite3.Connection) -> Optional[str]: - """Stream every evidence document once; one row is held at a time.""" - for evidence_hash, text in conn.execute( + """Stream every evidence document once; one row is held at a time. + + Rows are read with ``fetchone``: iterating a managed connection's cursor + would prefetch ``ITER_CHUNK`` (128) rows (``database/managed_connection.py``). + """ + items = conn.execute( f"SELECT evidence_hash, evidence FROM {ITEMS_TABLE}" # nosec B608 - ): + ) + while (row := items.fetchone()) is not None: + evidence_hash, text = row[0], row[1] if not _is_canonical(text) or evidence_digest(text) != evidence_hash: return f"Review queue table {ITEMS_TABLE} contains invalid evidence" - for (text,) in conn.execute( - f"SELECT evidence FROM {DECISIONS_TABLE}" # nosec B608 - ): + decisions = conn.execute(f"SELECT evidence FROM {DECISIONS_TABLE}") # nosec B608 + while (row := decisions.fetchone()) is not None: + text = row[0] if not _is_canonical(text): return f"Review queue table {DECISIONS_TABLE} contains invalid evidence" return None diff --git a/database/safety_backup.py b/database/safety_backup.py new file mode 100644 index 0000000..ad212a1 --- /dev/null +++ b/database/safety_backup.py @@ -0,0 +1,185 @@ +"""Safety backups beside the live database: one writer for every caller. + +A safety backup is ``.