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

feat(system): backup and restore backend (System S8) (#335)

Browse Source 
* refactor(database): stage restores before publishing

restore_from_file is now stage_restore + publish_staged (database/
restore_stage.py). The stage lives in a .pm-restore-* folder in the
library's own folder, is migrated, has damaged workflow documents repaired
from an inline copy or cleared, is verified with every document read and is
bound by its SHA-256. Publishing keeps a unique, standalone, 0600 safety
backup (microsecond names), copies the pages through the backup API as
before and then truncates the WAL, which stayed database-sized after a
restore. The safety backup writer moves to database/safety_backup.py so
the database layer and maintenance share it.

* fix(database): read the verify streams one row at a time

The provenance and review-evidence checks iterated their cursors; on a
managed connection iteration prefetches 128 rows, so up to 128 large values
were held at once while the docstrings promised one. Both now use fetchone.

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

- POST /restore/preview stages an upload in the library's folder and
  answers what it holds beside the live library, writing nothing live;
  DELETE discards it; at most two stages are held (429 too-large).
- POST /restore/apply rechecks the live fingerprint (409 plan-changed) and
  starts "Database restore": exclusive, not cancellable, indeterminate. It
  waits for request leases, publishes the exact staged bytes, swaps the
  database, prunes safety backups and logs counts, repairs and the backup.
- The legacy /restore runs the same stage and publish as an exclusive hold
  with its old response, so a restore and the health check, Move library
  and every other job refuse each other.
- JobRegistry.lease: Needs review decisions and maintenance and restore
  previews are refused while an exclusive job runs, and a restore waits
  for them before it switches the file.
- GET /backup holds a guarded slot, verifies a standalone snapshot before
  streaming it and removes the closed temp file; GET /backup/status lists
  the safety backups on disk and this session's verified download.

* docs(testing): backup and restore contract

* test(system): a restore refuses a running health check

The health route test relied on a legacy restore interleaving with a running
check; the restore is now exclusive and refuses it, naming the check. The
diagnostics backup tests follow the backup handler into
py/api/backup_restore.py.

* fix(backup): drop stale restore previews off the event loop

A restore drops the previews reviewed against the old library; their stage
files are now removed in the executor, with the old database's connections.

* test(backup): wait for the backup hold to end before reading its state

CI (py3.11/3.12 on Linux, every macOS) read the registry right after the
client had the last byte, before the handler left its hold and removed its
temp file: the job was still running. The tests now wait on that state.

* test(images): the scan dialog's Back up now leaves a finished backup hold

A backup download now holds a guarded registry slot while it streams
(#249 S8), so the jobs list holds that finished hold after the dialog's
Back up now; the spec still proves Cancel starts no scan.

* fix(backup): final review round for the restore and backup routes

- I1: once the stage's copy replaced the library the restore succeeded:
  pruning, reopening self.db, dropping old previews and closing the old
  object only log, so the job finishes and the legacy /restore answers 200
  instead of "the library was not changed". The generation is always bumped.
- M1: a backup download holds its guarded slot only while the snapshot is
  written; the stream runs outside it, so a stalled download blocks nothing.
- M2: a restore preview leases only its read of the live library, not the
  upload or the staging; a restore whose drain still times out keeps the
  reviewed stage under its token.
- M3: the reviewed-preview store changes only on the event loop; stage
  folders are removed in the executor.
- M4: a document whose body inflates completely while only its recorded
  hash or only its length disagrees is re-keyed (or merged into the
  document holding that text) instead of dropped.
- M5: a backup download is refused only when the snapshot is not SQLite or
  fails integrity_check; a semantic check failure still downloads, logged,
  sent as X-PromptManager-Verify and recorded as not verified.
- M6: a refused, failed or cancelled preview removes its stage, also when
  the cancel lands while a worker still stages; a timer removes expired
  stages without another request.
- M7, N1: repair reads rows one at a time and counts distinct images.
- N3: the live fingerprint never commits a transaction it did not open.

* test(backup): close the old library the close-failure test leaves open

CI (py3.12 Linux and the coverage gate) found the old object's
executor-thread connections still open at teardown: the test patches its
close_all to fail, and once the object is collected close_all_instances no
longer reaches it. The test now closes it itself.

* fix(backup): re-review round for restore repair, apply and backup cleanup

- R1: when re-keying merges an intact body into a holder whose own body was
  damaged, the rewritten holder leaves the damaged list, so the clear pass
  no longer deletes it and unlinks the images just moved to it.
- R2: a losing duplicate apply removes a stage only when the store still
  held it, never one another apply's job has taken.
- R3: a restore that times out waiting for requests keeps its review only
  while it is unexpired and its slot is free; otherwise the stage goes and
  the job says to try again.
- R4: a backup download cancelled while its snapshot is written removes
  the snapshot once the worker finishes.
This commit is contained in:
Vito
2026-10-25 10:17:07 -07:00
committed by GitHub
parent a905cc2316
commit 28ad3a1a1f
24 changed files with 3890 additions and 497 deletions
Download Patch File Download Diff File Expand all files Collapse all files
TESTING.md
+60
View File
@@ -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 `<db>.backup_<YYYYmmdd_HHMMSS_ffffff>[_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 `<db>.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
database/models.py
+52 -62
View File
@@ -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 ``<db_path>.backup_<stamp>`` 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
``<db_path>.backup_<YYYYmmdd_HHMMSS>`` (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
``<db_path>.backup_<YYYYmmdd_HHMMSS_ffffff>`` (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
database/restore_stage.py
+365
View File
@@ -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:<hex>`` 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
database/review_schema.py
+12 -6
View File
@@ -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
database/safety_backup.py
+185
View File
@@ -0,0 +1,185 @@
"""Safety backups beside the live database: one writer for every caller.
A safety backup is ``<db>.<label>_<YYYYmmdd_HHMMSS_ffffff>[_n]`` in the
library's own folder. It is written by SQLite's backup API into a temp file
beside the database (``SAFETY_TEMP_PREFIX``, swept at startup after a crash),
made one self-contained file (rollback journal), optionally verified, fsynced,
made private (``0o600`` on POSIX) and only then given its unique name, so a
backup name never carries a half-written file. The names are unique within
one second (microseconds, then a counter): on Windows two backups in the same
second must not overwrite each other.
Maintenance applies take a verified backup (``py/api/maintenance.py``); a
restore takes an unverified one, because the database it replaces may be the
damaged one the user is recovering from (``database/restore_stage.py``).
"""
import datetime
import os
import sqlite3
import tempfile
from contextlib import closing
SAFETY_TEMP_PREFIX = ".pm-safety-"
_NAME_ATTEMPTS = 1000
class SafetyBackupFailed(Exception):
"""The safety backup could not be written, verified or put in place.
``reason`` is the verify reason when the copy did not verify.
"""
def __init__(self, message, reason=None):
super().__init__(message)
self.reason = reason
def _now():
return datetime.datetime.now()
def remove_quietly(*paths):
"""Remove each existing path, ignoring errors (cleanup after a failure)."""
for path in paths:
try:
if path and os.path.exists(path):
os.remove(path)
except OSError:
pass
def backup_names(db_path, label="backup"):
"""Candidate names, unique within one second (microseconds, then ``_n``)."""
base = f"{db_path}.{label}_{_now().strftime('%Y%m%d_%H%M%S_%f')}"
for attempt in range(_NAME_ATTEMPTS):
yield base if attempt == 0 else f"{base}_{attempt}"
def reserve_backup_name(db_path, label="backup"):
"""Create and return an unused ``<db>.<label>_<stamp>[_n]`` (O_EXCL)."""
for candidate in backup_names(db_path, label):
try:
fd = os.open(candidate, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
except FileExistsError:
continue
os.close(fd)
return candidate
raise SafetyBackupFailed("No free backup name")
def place_backup(temp, db_path, label="backup"):
"""Give the finished ``temp`` its unique backup name; return that path.
``os.link`` is atomic and refuses an existing name, so no empty or half
file ever carries a backup name, even across a crash. A filesystem
without hard links falls back to reserving the name (``O_EXCL``) and
``os.replace``. The temp file is removed by the caller.
"""
for candidate in backup_names(db_path, label):
try:
os.link(temp, candidate)
except FileExistsError:
continue
except (OSError, NotImplementedError):
break # no hard links here
else:
return candidate
else:
raise SafetyBackupFailed("No free backup name")
final = reserve_backup_name(db_path, label)
try:
os.replace(temp, final)
except OSError:
remove_quietly(final)
raise
return final
def sweep_safety_temps(db_path):
"""Remove the temp copies an interrupted safety backup left (blocking).
Called once at startup, before any job can take a backup. Returns the
names removed.
"""
folder = os.path.dirname(os.path.abspath(db_path))
removed = []
try:
names = os.listdir(folder)
except OSError:
return removed
for name in names:
if not (name.startswith(SAFETY_TEMP_PREFIX) and ".tmp" in name):
continue
try:
os.remove(os.path.join(folder, name))
removed.append(name)
except OSError:
continue
return removed
def make_standalone(path):
"""Make the copy one self-contained file (rollback journal, not WAL).
The backup API copies the live database's WAL flag; a WAL file opened
read-only later (a verify, a restore preview) would leave ``-wal`` and
``-shm`` sidecars beside it.
"""
try:
with closing(sqlite3.connect(path)) as conn:
conn.execute("PRAGMA journal_mode=DELETE").fetchone()
except sqlite3.Error as exc:
raise SafetyBackupFailed("The backup could not be finished") from exc
def fsync_file(path):
with open(path, "rb+") as handle:
os.fsync(handle.fileno())
def safety_temp(db_path):
"""A new, closed temp file beside the database (``SAFETY_TEMP_PREFIX``)."""
fd, temp = tempfile.mkstemp(
dir=os.path.dirname(os.path.abspath(db_path)),
prefix=SAFETY_TEMP_PREFIX,
suffix=".tmp",
)
os.close(fd)
return temp
def take_safety_backup(model, label="backup", *, verify=True):
"""Write (and by default verify) a safety backup; return its path (blocking).
The copy is made in a temp file beside the database and only renamed
onto its final, unique name ``<db>.<label>_<stamp>`` once it is complete
(and verified), so a backup file on disk is always a whole one. ``label``
"backup" joins the pruned pool; another label (the workflow storage
migration's ``pre-workflow-storage``) is kept until the user removes it.
Raises SafetyBackupFailed, leaving no file.
"""
db_path = os.path.abspath(model.db_path)
temp = safety_temp(db_path)
final = None
try:
if not model.backup_database(temp):
raise SafetyBackupFailed("The backup could not be written")
make_standalone(temp)
if verify:
ok, reason = model.verify_database_file(temp)
if not ok:
raise SafetyBackupFailed("The backup did not verify", reason)
fsync_file(temp)
if os.name != "nt":
os.chmod(temp, 0o600)
final = place_backup(temp, db_path, label)
return final
except OSError as exc:
remove_quietly(final)
raise SafetyBackupFailed("The backup could not be put in place") from exc
except SafetyBackupFailed:
remove_quietly(final)
raise
finally:
remove_quietly(temp, temp + "-journal", temp + "-wal", temp + "-shm")
database/workflow_documents.py
+175
View File
@@ -475,3 +475,178 @@ def release_documents(conn: sqlite3.Connection, ids: Iterable[Optional[int]]) ->
def sweep_documents(conn: sqlite3.Connection) -> int:
"""Delete every document no image row references; returns the count."""
return conn.execute(SWEEP_SQL).rowcount
# ── repair (a restore's stage) ────────────────────────────────────────
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 _rows(conn, INLINE_TEXTS_SQL):
for text in row:
if not isinstance(text, str) or not storable(text):
continue
document_id = wanted.get(hashlib.sha256(text.encode("utf-8")).digest())
if document_id is not None and document_id not in found:
found[document_id] = text
if len(found) == len(wanted):
break
return found
def _rekey(conn: sqlite3.Connection, document_id: int, raw: bytes) -> Optional[int]:
"""Give an intact body its true key, or merge it into the document that has it.
Returns the holder's id when its own body was damaged and has just been
rewritten from ``raw`` (it is repaired, never to be cleared), else None.
"""
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 None
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])
)
return holder[0]
return None
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. 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] = {}
intact: Dict[int, bytes] = {}
for row in _rows(conn, ALL_DOCUMENTS_SQL):
if _decode(row[:3], row[3], quiet=True) is None:
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 and not intact:
return result
with conn:
healed = []
for document_id, raw in intact.items():
holder = _rekey(conn, document_id, raw)
if holder is not None and holder in damaged.values():
# Rewritten from the merged body: repaired, not to be cleared
# (it now also holds the merged images' references).
damaged = {k: v for k, v in damaged.items() if v != holder}
healed.append(holder)
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: Set[int] = set()
for chunk in _chunks(lost):
marks = ",".join("?" * len(chunk))
# ``marks`` is only "?" placeholders; the ids are bound.
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,
):
images.add(image_id)
for column in DOCUMENT_ID_COLUMNS:
conn.execute(
f"UPDATE generated_images SET {column}=NULL " # nosec B608
f"WHERE {column} IN ({marks})",
chunk,
)
conn.execute(
f"DELETE FROM workflow_documents WHERE id IN ({marks})", # nosec B608
chunk,
)
for document_id in list(intact) + healed + 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"{len(images)} image(s) show no workflow"
)
result.update(
repaired=len(intact) + len(healed) + len(copies),
cleared=len(lost),
images=len(images),
)
return result
py/api/__init__.py
+6
View File
@@ -56,6 +56,7 @@ from .thumbnail_auto import AutoThumbnailMixin
from .system import SystemRoutesMixin
from .library_location import LibraryLocationRoutesMixin
from .maintenance import MaintenanceRoutesMixin
from .backup_restore import BackupRestoreRoutesMixin
from .workflow_storage import WorkflowStorageRoutesMixin
from .duplicate_files import DuplicateFilesMixin
from .maintenance_plans import PlanStore
@@ -302,6 +303,7 @@ class PromptManagerAPI(
SystemRoutesMixin,
LibraryLocationRoutesMixin,
MaintenanceRoutesMixin,
BackupRestoreRoutesMixin,
WorkflowStorageRoutesMixin,
DuplicateFilesMixin,
LoggingRoutesMixin,
@@ -376,7 +378,10 @@ class PromptManagerAPI(
try:
from .maintenance import sweep_safety_temps
from .backup_restore import restore_stage
swept = sweep_safety_temps(self.db.model.db_path)
swept += restore_stage.sweep_restore_temps(self.db.model.db_path)
if swept:
self.logger.info(f"Removed {len(swept)} interrupted backup copies")
except Exception as e:
@@ -547,6 +552,7 @@ class PromptManagerAPI(
self._register_system_routes(routes)
self._register_library_location_routes(routes)
self._register_maintenance_routes(routes)
self._register_backup_restore_routes(routes)
self._register_duplicate_routes(routes)
self._register_logging_routes(routes)
self._register_autotag_routes(routes)
py/api/admin.py
-215
View File
@@ -3,14 +3,12 @@
import asyncio
import contextlib
from contextlib import closing
import datetime
import functools
import hashlib
import json
import math
import os
import sqlite3
import tempfile
import threading
import time
import urllib.parse
@@ -28,7 +26,6 @@ from .jobs import (
DUPLICATE_CLEANUP_JOB_KIND,
DUPLICATES_JOB_KIND,
MAINTENANCE_JOB_KIND,
RESTORE_JOB_KIND,
JobBusy,
)
@@ -38,7 +35,6 @@ except ImportError:
from utils.parallel import map_parallel
try:
from ...database.operations import PromptDatabase
from ...database.pair_ingestion import ingest_resolved_pair
from ...utils.generation_params import (
PLACEHOLDERS,
@@ -49,7 +45,6 @@ try:
from ...utils.prompt_pair_resolver import resolve_prompt_pair
from ...utils.review_capture import reviewable
except ImportError:
from database.operations import PromptDatabase
from database.pair_ingestion import ingest_resolved_pair
from utils.generation_params import (
PLACEHOLDERS,
@@ -1259,216 +1254,6 @@ class AdminRoutesMixin:
"message": f"Found {len(issues)} consistency issues",
}
async def backup_database(self, request):
"""Download a consistent copy of the prompts database (WAL included)."""
try:
model = self.db.model
db_path = model.db_path
if not os.path.exists(db_path):
return web.json_response(
{"success": False, "error": "Database file not found"}, status=404
)
fd, temp_path = tempfile.mkstemp(suffix=".db")
os.close(fd)
try:
ok = await self._run_in_executor(model.backup_database, temp_path)
if not ok:
return web.json_response(
{"success": False, "error": "Failed to create database backup"},
status=500,
)
# Read once (the file is deleted right after) rather than
# streaming so the temp file never outlives the request.
with open(temp_path, "rb") as fh:
file_data = fh.read()
finally:
if os.path.exists(temp_path):
os.unlink(temp_path)
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
filename = f"prompts_backup_{timestamp}.db"
return web.Response(
body=file_data,
content_type="application/octet-stream",
headers={
"Content-Disposition": f'attachment; filename="{filename}"',
"Content-Length": str(len(file_data)),
},
)
except Exception as e:
self.logger.error(f"Backup error: {e}", exc_info=True)
return web.json_response(
{
"success": False,
"error": f"Failed to backup database: {self._public_error(e)}",
},
status=500,
)
async def restore_database(self, request):
"""Restore the prompts database from an uploaded SQLite file.
The upload is streamed to a temp file. Free disk space decides how
large it may be (``maintenance.restore_room``), under a 64 GiB ceiling
(``restore_max_bytes``): a request whose Content-Length does not fit
is refused before anything is written, and one without is refused as
soon as its running size does not. The file is
verified with PromptModel.verify_database_file and only then copied
over the live database, which is backed up first. The whole request
holds the registry (records): it is refused while a scan or another
record-changing job runs, and those are refused while it runs.
"""
return await self._run_held(
RESTORE_JOB_KIND,
"Database restore (admin page)",
"records",
lambda: self._restore_database(request),
)
def _restore_refusal(self, size, ceiling, room, *, known):
"""``400 too-large`` past the ceiling, else ``507 insufficient-space``."""
from .maintenance import restore_space_message
if size > ceiling:
return web.json_response(
{
"success": False,
"code": "too-large",
"error": (
"File too large. Maximum size is "
f"{ceiling // (1024 * 1024)}MB"
),
},
status=400,
)
return web.json_response(
{
"success": False,
"code": "insufficient-space",
"error": restore_space_message(size, room, known=known),
},
status=507,
)
async def _restore_database(self, request):
from .maintenance import RESTORE_MAX_BYTES, restore_room
ceiling = getattr(self, "restore_max_bytes", RESTORE_MAX_BYTES)
temp_path = None
try:
# Free disk space decides how large a backup may be (#313): the
# upload, its staged copy and the swap must all fit.
room = await self._run_in_executor(restore_room, self.db.model.db_path)
max_bytes = min(ceiling, room)
length = request.content_length
if length is not None and length > max_bytes:
return self._restore_refusal(length, ceiling, room, known=True)
reader = await request.multipart()
field = await reader.next()
if not field or field.name != "database_file":
return web.json_response(
{
"success": False,
"error": (
"No database file uploaded. "
"Expected field name: database_file"
),
},
status=400,
)
fd, temp_path = tempfile.mkstemp(suffix=".db")
os.close(fd)
size = 0
with open(temp_path, "wb") as out:
while True:
chunk = await field.read_chunk()
if not chunk:
break
size += len(chunk)
if size > max_bytes:
return self._restore_refusal(size, ceiling, room, known=False)
out.write(chunk)
if size == 0:
return web.json_response(
{"success": False, "error": "Uploaded file is empty"}, status=400
)
model = self.db.model
ok, reason = await self._run_in_executor(
model.verify_database_file, temp_path
)
if ok:
reason = await self._run_in_executor(_inspect_restore_upload, temp_path)
ok = reason is None
if not ok:
return web.json_response(
{"success": False, "error": reason}, status=400
)
try:
backup_path = await self._run_in_executor(
model.restore_from_file, temp_path
)
except sqlite3.OperationalError as exc:
if (
getattr(exc, "_prompt_manager_restore_backup_created", False)
is True
):
try:
await self._run_in_executor(
_prune_safety_backups, model.db_path
)
except Exception:
self.logger.warning(
"Could not prune safety backups after failed restore"
)
raise
removed = await self._run_in_executor(_prune_safety_backups, model.db_path)
if removed:
self.logger.info(f"Removed {len(removed)} old safety backups")
previous = self.db
self.db = PromptDatabase(model.db_path)
self._database_replaced() # a health report described the old one
# Its connections would otherwise linger (a file lock on Windows)
previous.close_all()
prompt_count = self.db.model.get_database_info().get("total_prompts", 0)
return web.json_response(
{
"success": True,
"message": (
"Database restored successfully. "
f"Found {prompt_count} prompts."
),
"prompt_count": prompt_count,
"backup_created": (
self._public_path(backup_path) if backup_path else None
),
}
)
except Exception as e:
self.logger.error(f"Restore error: {e}", exc_info=True)
return web.json_response(
{
"success": False,
"error": f"Failed to restore database: {self._public_error(e)}",
},
status=500,
)
finally:
if temp_path:
for suffix in ("", "-wal", "-shm", "-journal"):
if os.path.exists(temp_path + suffix):
os.unlink(temp_path + suffix)
async def scan_images(self, request):
"""Scan ComfyUI output images for prompt metadata and add them to the database.
py/api/backup_restore.py
+937
View File
@@ -0,0 +1,937 @@
"""Database backup and restore routes (#249 System S8).
Legacy routes, responses unchanged:
- ``GET /prompt_manager/backup`` downloads a consistent copy of the library.
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``).
Reviewed restore (new routes answer errors as ``{success:false, code, error}``):
- ``POST /prompt_manager/restore/preview`` (multipart ``database_file``, the
same disk-space cap) stages the upload in the library's folder
(``database/restore_stage.py``) and answers what it holds beside the live
library: ``{token, expires_at, summary, schema:{migrated, library_state},
integrity:"ok", documents:{repaired, cleared, images}, live, size,
safety_backup:{destination_label}}``. Nothing live is written or backed
up. At most two previews are held (``too-large``, 429).
- ``DELETE /prompt_manager/restore/preview/{token}`` discards a stage.
- ``POST /prompt_manager/restore/apply {token}`` starts the restore job (202
``{job}``): kind ``restore``, "Database restore", ``exclusive`` (it runs
alone, like Move library), not cancellable, indeterminate. The runner waits
for request leases to end, rechecks the live fingerprint (``plan-changed``),
publishes the exact staged bytes, swaps ``self.db``, prunes old safety
backups and logs what was restored. 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
(``<db>.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
from . import admin as _admin
from .job_routes import busy_response
from .jobs import RESTORE_JOB_KIND, JobBusy
from .maintenance_plans import PlanExpired, PlanUsed
from .restore_plans import RestorePlanStore
try:
from ...database import restore_stage
from ...database.operations import PromptDatabase
from ...database.safety_backup import (
SafetyBackupFailed,
fsync_file,
make_standalone,
remove_quietly,
safety_temp,
)
except ImportError:
from database import restore_stage
from database.operations import PromptDatabase
from database.safety_backup import (
SafetyBackupFailed,
fsync_file,
make_standalone,
remove_quietly,
safety_temp,
)
BACKUP_JOB_KIND = "backup"
BACKUP_HOLD_TITLE = "Database backup"
LEGACY_RESTORE_TITLE = "Database restore (admin page)"
RESTORE_JOB_TITLE = "Database restore"
RESTORE_GROUPS = ("errors",)
MAX_TOKEN_LENGTH = 128
APPLY_BODY_MAX = 1024
STREAM_CHUNK = 1024 * 1024
MAX_LISTED_BACKUPS = 10
SIDECARS = ("", "-wal", "-shm", "-journal")
REQUESTS_BUSY = (
"Other requests are still reading the library; nothing was changed. "
"Try again in a moment."
)
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."
)
NOT_RESTORED = "The restore failed; the library was not changed."
def _failure(message, status, code=None, **extra):
return web.json_response(
{"success": False, "code": code, "error": message, **extra}, status=status
)
def _iso(timestamp):
return datetime.datetime.fromtimestamp(timestamp, datetime.timezone.utc).isoformat()
def _plural(count, word):
return f"{count:,} {word}" + ("" if count == 1 else "s")
def documents_line(documents):
"""The repair report's sentence, or "" when nothing was damaged."""
repaired, cleared = documents.get("repaired", 0), documents.get("cleared", 0)
parts = []
if repaired:
parts.append(f"repaired {_plural(repaired, 'damaged workflow document')}")
if cleared:
images = documents.get("images", 0)
parts.append(
f"could not repair {_plural(cleared, 'workflow document')}: "
f"{_plural(images, 'image')} {'shows' if images == 1 else 'show'} "
"no workflow"
)
if not parts:
return ""
text = "; ".join(parts)
return text[0].upper() + text[1:] + "."
def restored_line(summary):
return (
f"Restored {_plural(summary['prompts'], 'prompt')}, "
f"{_plural(summary['image_links'], 'image link')}, "
f"{_plural(summary['pairs'], 'saved pair')} and "
f"{_plural(summary['runs'], 'run')}"
)
def safety_backups(db_path, limit=MAX_LISTED_BACKUPS):
"""``<db>.backup_*`` files beside the database, newest first (blocking)."""
db_path = os.path.abspath(db_path)
folder = os.path.dirname(db_path)
prefix = os.path.basename(db_path) + ".backup_"
found = []
try:
names = os.listdir(folder)
except OSError:
return []
for name in names:
if not name.startswith(prefix):
continue
try:
stat = os.stat(os.path.join(folder, name))
except OSError:
continue
if not os.path.isfile(os.path.join(folder, name)):
continue
found.append((stat.st_mtime, name, stat.st_size))
found.sort(key=lambda item: (item[0], item[1]), reverse=True)
return [
{"name": name, "bytes": size, "modified_at": _iso(mtime)}
for mtime, name, size in found[:limit]
]
def 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 standalone copy of the library beside it: ``(path, warning)`` (blocking).
Raises SafetyBackupFailed with nothing left on disk when the copy cannot
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)
fsync_file(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."""
def _register_backup_restore_routes(self, routes):
@routes.post("/prompt_manager/restore/preview")
async def restore_preview_route(request):
return await self.preview_restore(request)
@routes.delete("/prompt_manager/restore/preview/{token}")
async def restore_discard_route(request):
return await self.discard_restore_preview(request)
@routes.post("/prompt_manager/restore/apply")
async def restore_apply_route(request):
return await self.apply_restore(request)
@routes.get("/prompt_manager/backup/status")
async def backup_status_route(request):
return await self.backup_status(request)
@property
def _restore_plans(self):
store = getattr(self, "_restore_plan_store", None)
if store is None:
store = self._restore_plan_store = RestorePlanStore()
return store
# ── backup ────────────────────────────────────────────────────────
async def backup_database(self, request):
"""Download a copy of the prompts database (WAL included)."""
return await self._stream_backup(request)
async def _stream_backup(self, request):
model = self.db.model
if not os.path.exists(model.db_path):
return web.json_response(
{"success": False, "error": "Database file not found"}, status=404
)
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)."""
loop = asyncio.get_running_loop()
future = loop.run_in_executor(None, write_backup_snapshot, model)
try:
temp, warning = await asyncio.shield(future)
except asyncio.CancelledError:
# The worker cannot be stopped: remove what it writes once it is
# done (#335 re-review R4).
future.add_done_callback(
lambda done: loop.create_task(self._remove_finished_snapshot(done))
)
raise
except SafetyBackupFailed as failed:
self.logger.error(f"Backup failed: {failed} {failed.reason or ''}")
return web.json_response(
{"success": False, "error": "Failed to create database backup"},
status=500,
)
except Exception as e:
self.logger.error(f"Backup error: {e}", exc_info=True)
return web.json_response(
{
"success": False,
"error": f"Failed to backup database: {self._public_error(e)}",
},
status=500,
)
if warning is not None:
self.logger.warning(f"Backup sent with a failed check: {warning}")
return temp, warning
async def _remove_finished_snapshot(self, future):
"""Remove the snapshot a cancelled download's worker finished."""
if future.cancelled() or future.exception() is not None:
return # nothing was left on disk
temp, _warning = future.result()
await self._run_in_executor(
remove_quietly, *(temp + suffix for suffix in SIDECARS)
)
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:
size = os.path.getsize(temp)
handle = open(temp, "rb")
response = web.StreamResponse(
headers={
"Content-Type": "application/octet-stream",
"Content-Disposition": (
f'attachment; filename="prompts_backup_{timestamp}.db"'
),
"Content-Length": str(size),
VERIFY_HEADER: "ok" if warning is None else _header_text(warning),
}
)
await response.prepare(request)
while True:
chunk = await self._run_in_executor(handle.read, STREAM_CHUNK)
if not chunk:
break
await response.write(chunk)
await response.write_eof()
self._last_backup_download = {
"at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
"bytes": size,
"verified": warning is None,
}
return response
finally:
if handle is not None:
handle.close() # before the unlink (Windows)
await self._run_in_executor(
remove_quietly, *(temp + suffix for suffix in SIDECARS)
)
async def backup_status(self, request):
try:
backups = await self._run_in_executor(safety_backups, self.db.model.db_path)
except Exception:
self.logger.exception("Backup status failed")
return _failure("Failed to read the backup status", 500)
return web.json_response(
{
"success": True,
"safety_backups": backups,
"last_download": getattr(self, "_last_backup_download", None),
},
headers={"Cache-Control": "no-store"},
)
# ── the upload ────────────────────────────────────────────────────
def _restore_refusal(self, size, ceiling, room, *, known):
"""``400 too-large`` past the ceiling, else ``507 insufficient-space``."""
from .maintenance import restore_space_message
if size > ceiling:
return web.json_response(
{
"success": False,
"code": "too-large",
"error": (
"File too large. Maximum size is "
f"{ceiling // (1024 * 1024)}MB"
),
},
status=400,
)
return web.json_response(
{
"success": False,
"code": "insufficient-space",
"error": restore_space_message(size, room, known=known),
},
status=507,
)
async def _receive_restore_upload(self, request, uploads):
"""Stream ``database_file`` into the library's folder.
Returns ``(path, None)`` or ``(None, refusal)``. Each temp path is
appended to ``uploads`` as soon as it exists, so the caller removes
it (and its sidecars) whatever happens.
"""
from .maintenance import RESTORE_MAX_BYTES, restore_room
ceiling = getattr(self, "restore_max_bytes", RESTORE_MAX_BYTES)
db_path = self.db.model.db_path
# Free disk space decides how large a backup may be (#313): the
# upload, its staged copy and the swap must all fit.
room = await self._run_in_executor(restore_room, db_path)
max_bytes = min(ceiling, room)
length = request.content_length
if length is not None and length > max_bytes:
return None, self._restore_refusal(length, ceiling, room, known=True)
reader = await request.multipart()
field = await reader.next()
if not field or field.name != "database_file":
return None, web.json_response(
{
"success": False,
"error": (
"No database file uploaded. "
"Expected field name: database_file"
),
},
status=400,
)
fd, temp_path = tempfile.mkstemp(
prefix=restore_stage.UPLOAD_PREFIX,
suffix=".tmp",
dir=os.path.dirname(os.path.abspath(db_path)),
)
uploads.append(temp_path)
size = 0
with os.fdopen(fd, "wb") as out:
while True:
chunk = await field.read_chunk()
if not chunk:
break
size += len(chunk)
if size > max_bytes:
return None, self._restore_refusal(size, ceiling, room, known=False)
out.write(chunk)
if size == 0:
return None, web.json_response(
{"success": False, "error": "Uploaded file is empty"}, status=400
)
return temp_path, None
async def _upload_problem(self, path):
"""None when the upload may be staged, else a readable reason."""
ok, reason = await self._run_in_executor(
self.db.model.verify_database_file, path, documents=False
)
if ok:
reason = await self._run_in_executor(_admin._inspect_restore_upload, path)
ok = reason is None
return None if ok else reason
async def _remove_uploads(self, uploads):
for path in uploads:
await self._run_in_executor(
remove_quietly, *(path + suffix for suffix in SIDECARS)
)
async def _prune_after(self, db_path, *, failed=False):
"""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 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")
async def _publish(self, model, staged):
"""Publish ``staged``; on a page-copy failure prune and re-raise."""
try:
return await self._run_in_executor(model.publish_staged, staged)
except sqlite3.OperationalError as exc:
if getattr(exc, "_prompt_manager_restore_backup_created", False) is True:
await self._prune_after(model.db_path, failed=True)
raise
async def _after_publish(self, db_path):
"""Everything after the copy replaced the library: best effort, logged.
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
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 _discard_popped(self, token):
"""Drop ``token`` and remove its stage, only if the store still held it.
A stage another apply's job has taken is not this request's to remove
(#335 re-review R2).
"""
popped = self._restore_plans.pop(token)
if popped is not None:
await self._discard_held([popped])
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 ───────────────────────────────────────
async def restore_database(self, request):
"""Restore the prompts database from an uploaded SQLite file.
The upload is streamed to a temp file in the library's folder. Free
disk space decides how large it may be (``maintenance.restore_room``),
under a 64 GiB ceiling (``restore_max_bytes``): a request whose
Content-Length does not fit is refused before anything is written,
and one without is refused as soon as its running size does not. The
file is checked, staged, repaired and verified, then published over
the live database, which is backed up first. The whole request holds
the registry as ``restore`` (``exclusive``): it is refused while any
job runs, and every job and request lease is refused while it runs.
"""
return await self._run_held(
RESTORE_JOB_KIND,
LEGACY_RESTORE_TITLE,
"exclusive",
lambda: self._restore_database(request),
)
async def _restore_database(self, request):
uploads = []
try:
temp_path, refused = await self._receive_restore_upload(request, uploads)
if refused is not None:
return refused
reason = await self._upload_problem(temp_path)
if reason is not None:
return web.json_response(
{"success": False, "error": reason}, status=400
)
if not await self._jobs.drain_leases():
return web.json_response(
{"success": False, "error": REQUESTS_BUSY}, status=409
)
model = self.db.model
staged = await self._run_in_executor(model.stage_restore, temp_path)
try:
backup_path = await self._publish(model, staged)
finally:
await self._run_in_executor(staged.discard)
# 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,
"message": message,
"prompt_count": prompt_count,
"backup_created": (
self._public_path(backup_path) if backup_path else None
),
}
repaired = documents_line(staged.documents)
if repaired:
body["message"] = f"{message} {repaired}"
body["workflow_documents"] = dict(staged.documents)
return web.json_response(body)
except Exception as e:
self.logger.error(f"Restore error: {e}", exc_info=True)
return web.json_response(
{
"success": False,
"error": f"Failed to restore database: {self._public_error(e)}",
},
status=500,
)
finally:
await self._remove_uploads(uploads)
# ── reviewed restore ──────────────────────────────────────────────
def _live_fingerprint(self, db):
return restore_stage.live_fingerprint(db.model.get_connection())
async def preview_restore(self, request):
"""Stage an upload for review; nothing live is written (see module)."""
await self._drop_expired_restores()
if self._restore_plans.full():
return _too_many_reviews()
uploads = []
try:
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")
return _failure("Failed to review the restore", 500)
finally:
await self._remove_uploads(uploads)
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())
)
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(
{
"success": True,
"token": token,
"expires_at": _iso(held.expires_at),
"summary": dict(staged.summary),
"schema": dict(staged.schema),
"integrity": "ok",
"documents": dict(staged.documents),
"live": dict(live),
"size": staged.size,
"safety_backup": {
"destination_label": f"{name}.backup_<date and time>"
},
},
headers={"Cache-Control": "no-store"},
)
async def discard_restore_preview(self, request):
token = request.match_info.get("token", "")
if not 0 < len(token) <= MAX_TOKEN_LENGTH:
return _failure("Unknown restore review.", 410, "plan-expired")
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):
try:
return self._restore_plans.peek(token), None
except PlanUsed as used:
job = self._jobs.get(used.job_id)
return None, _failure(
"This restore was already started.",
409,
"plan-used",
job={"id": used.job_id, "title": job.title if job else None},
)
except PlanExpired:
return None, _failure(
"This review expired. Review the restore again.",
410,
"plan-expired",
)
async def apply_restore(self, request):
data, error = await self._read_json_body(request, max_bytes=APPLY_BODY_MAX)
if error is not None:
if error.status == 413:
return _failure("Request body too large", 413, "too-large")
return _failure(
"Request body must be a JSON object", 400, "invalid-request"
)
token = data.get("token")
if set(data) != {"token"} or not isinstance(token, str):
return _failure("token must be the review's token", 400, "invalid-request")
if not 0 < len(token) <= MAX_TOKEN_LENGTH:
return _failure("token must be the review's token", 400, "invalid-request")
await self._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._discard_popped(token)
return _failure(
"The database was replaced since the review. Review again.",
409,
"storage-changed",
)
try:
async with self._jobs.lease("Restore review check"):
_live, fingerprint = await self._run_in_executor(
self._live_fingerprint, db
)
except JobBusy as busy:
return busy_response(busy)
if fingerprint != held.fingerprint:
# It can never match again: free the stage for a new review.
await self._discard_popped(token)
return _failure(PLAN_CHANGED, 409, "plan-changed")
# The token may have been used, expired or discarded meanwhile.
held, refused = self._peek_restore(token)
if refused is not None:
await self._drop_expired_restores()
return refused
if db is not self.db:
return _failure(
"The database was replaced since the review. Review again.",
409,
"storage-changed",
)
try:
job = self._jobs.start(
RESTORE_JOB_KIND,
RESTORE_JOB_TITLE,
# It replaces the whole database: it runs alone, beside no
# job of any class, like Move library.
admission="exclusive",
cancellable=False,
scope={},
groups=RESTORE_GROUPS,
runner=lambda job: self._run_restore(job, db, held, token),
)
except JobBusy as busy:
return busy_response(busy) # the token stays usable
# Same loop turn as the peek and the start: the runner has not run.
self._restore_plans.take(token, job.id)
return web.json_response({"success": True, "job": job.summary()}, status=202)
async def _run_restore(self, job, db, held, 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():
# Nothing ran: the reviewed stage stays usable (#335 M2),
# unless it expired meanwhile or its slot was reused.
if token is not None and 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)
if db is not self.db or fingerprint != held.fingerprint:
return _fail_job(job, "plan-changed", PLAN_CHANGED)
job.update(status_line="Restoring the database")
job.log_line("Restoring the reviewed backup")
model = db.model
try:
backup_path = await self._publish(model, staged)
except restore_stage.StageChanged:
return _fail_job(job, "plan-changed", PLAN_CHANGED)
except SafetyBackupFailed:
return _fail_job(
job,
"backup-failed",
"The current library could not be backed up; nothing was "
"changed.",
)
except sqlite3.OperationalError as error:
kept = (
" The safety backup was kept."
if getattr(error, "_prompt_manager_restore_backup_created", False)
else ""
)
self.logger.error(f"Restore failed: {error}")
return _fail_job(job, "restore-failed", NOT_RESTORED + kept)
except Exception:
self.logger.exception("Restore job failed")
return _fail_job(job, None, NOT_RESTORED)
finally:
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",
)
py/api/jobs.py
+100
View File
@@ -28,6 +28,15 @@ ones. One job of a kind runs at a time, and each job has an admission class
A short legacy-page operation takes a ``hold``: a running, non-cancellable
job that lasts as long as the request, so the Activity page lists it and a
workspace job is refused while it runs (and the other way round).
Short request-scoped work that reads or writes the library beside any
ordinary job (a Needs review decision, a maintenance or restore preview)
takes a ``lease`` instead: it is not listed (so quick requests never push
real jobs out of the finished history, and two of them never refuse each
other), but it is refused while an ``exclusive`` job runs, and an
``exclusive`` job that replaces the database (a restore) awaits
``drain_leases`` before it switches the file, so no lease's connection is
open across the switch (#249 S8).
"""
import asyncio
@@ -48,6 +57,9 @@ MAX_OUTCOME_ROWS = 200 # rows listed per outcome group; counters stay exact
MAX_LOG_LINES = 2000
LOG_CAP_LINE = "[Log capped at 2,000 lines]"
MAX_FINISHED_JOBS = 20 # finished jobs this session keeps
# The longest an exclusive job waits for running leases to end (a ceiling a
# hung request reaches, not a pacing timer: leases end with their request).
LEASE_DRAIN_SECONDS = 60
# The metadata scan's outcome groups, in counter-tile order of the spec's
# mapping, and the per-file scan status each group collects. Files without
@@ -403,6 +415,74 @@ class JobRegistry:
def __init__(self, logger_=None):
self._logger = logger_ or logger
self._jobs = {}
self._leases = {} # title -> running count
self._drain_waiters = []
# ── leases ────────────────────────────────────────────────────────
def lease(self, title):
"""An async context manager for short request-scoped library work.
Entering raises JobBusy, naming the job, while an ``exclusive`` job
runs; otherwise the lease is counted until the block exits. Leases
are never listed and never refuse one another or any job.
"""
return _Lease(self, title)
def lease_blocker(self):
"""The running ``exclusive`` job a new lease is refused for, or None."""
for job in self.running():
if job.admission == "exclusive":
return job
return None
@property
def leases(self):
"""Running leases by title (a copy)."""
return dict(self._leases)
def _take_lease(self, title):
blocker = self.lease_blocker()
if blocker is not None:
raise JobBusy(blocker)
self._leases[title] = self._leases.get(title, 0) + 1
def _release_lease(self, title):
count = self._leases.get(title, 0) - 1
if count > 0:
self._leases[title] = count
else:
self._leases.pop(title, None)
if not self._leases:
waiters, self._drain_waiters = self._drain_waiters, []
for waiter in waiters:
if not waiter.done():
waiter.set_result(True)
async def drain_leases(self, timeout=None):
"""Wait until no lease runs; False when ``timeout`` ended the wait.
For an ``exclusive`` job about to switch the database: once it runs
no new lease is admitted, so the count only goes down.
"""
if not self._leases:
return True
waiter = asyncio.get_running_loop().create_future()
self._drain_waiters.append(waiter)
try:
await asyncio.wait_for(
asyncio.shield(waiter),
LEASE_DRAIN_SECONDS if timeout is None else timeout,
)
return True
except asyncio.TimeoutError:
self._logger.warning(
"Still running after the wait: " + ", ".join(sorted(self._leases))
)
return False
finally:
if waiter in self._drain_waiters:
self._drain_waiters.remove(waiter)
def start(
self,
@@ -565,3 +645,23 @@ class _Hold:
if self._job.running:
self._job.finish("finished" if exc_type is None else "failed")
return False
class _Lease:
"""The async context manager ``JobRegistry.lease`` returns."""
def __init__(self, registry, title):
self._registry = registry
self._title = title
self._held = False
async def __aenter__(self):
self._registry._take_lease(self._title)
self._held = True
return self
async def __aexit__(self, exc_type, exc, tb):
if self._held:
self._held = False
self._registry._release_lease(self._title)
return False
py/api/maintenance.py
+36 -168
View File
@@ -48,8 +48,6 @@ import json
import os
import shutil
import sqlite3
import tempfile
from contextlib import closing
from aiohttp import web
@@ -59,15 +57,28 @@ from .jobs import MAINTENANCE_JOB_KIND, JobBusy
from .maintenance_plans import PlanExpired, PlanUsed, ScopeTooLarge
try:
from ...database import safety_backup as _backups
from ...database.safety_backup import ( # noqa: F401 - re-exported
SAFETY_TEMP_PREFIX,
SafetyBackupFailed,
sweep_safety_temps,
)
from ...database.workflow_storage import PRE_MIGRATION_LABEL, pre_migration_backups
from ...database.health import readonly_uri
from ...database.maintenance import MAINTENANCE_OPERATIONS, PlanChanged
except ImportError:
from database import safety_backup as _backups
from database.safety_backup import ( # noqa: F401 - re-exported
SAFETY_TEMP_PREFIX,
SafetyBackupFailed,
sweep_safety_temps,
)
from database.workflow_storage import PRE_MIGRATION_LABEL, pre_migration_backups
from database.health import readonly_uri
from database.maintenance import MAINTENANCE_OPERATIONS, PlanChanged
PREVIEW_BODY_MAX = 4096
MAINTENANCE_PREVIEW_LEASE = "Maintenance preview"
# Duplicate-file previews carry the reviewed keepers (py/api/duplicate_files.py).
DUPLICATE_PREVIEW_BODY_MAX = 1024 * 1024
DUPLICATE_FILES = "duplicate_files"
@@ -96,18 +107,6 @@ PLAN_CHANGED = (
APPLY_FAILED = "Maintenance failed; nothing was changed."
NOTHING_TO_DO = "Every image record already stores its workflow once; nothing to do."
NO_BACKUP = "No safety backup"
_NAME_ATTEMPTS = 1000
class SafetyBackupFailed(Exception):
"""The safety backup could not be written, verified or put in place.
``reason`` is the verify reason when the copy did not verify.
"""
def __init__(self, message, reason=None):
super().__init__(message)
self.reason = reason
def backup_failed_message(error):
@@ -118,10 +117,6 @@ def backup_failed_message(error):
return BACKUP_FAILED
def _now():
return datetime.datetime.now()
def _free_bytes(path):
return shutil.disk_usage(path).free
@@ -244,29 +239,18 @@ def size_label(size):
def restore_room(db_path):
"""The largest restore upload the disks have room for (blocking).
"""The largest restore upload the disk has room for (blocking).
The upload is written to a temp file and staged once more in the
temporary folder (2x), then copied over the live database beside it,
whose current contents are first kept as a safety backup (1x the upload
for the WAL, plus the current database). Needs on one device add up,
each with ``DISK_MARGIN_BYTES``.
Everything a restore writes is in the library's folder (#318): the
upload, its staged copy (``database/restore_stage.py``) and the WAL
growth while the stage is copied over the live database (3x the upload),
plus the safety backup of the current database (1x its size), with
``DISK_MARGIN_BYTES`` to spare.
"""
db_path = os.path.abspath(db_path)
folder = os.path.dirname(db_path)
current = database_bytes(db_path) if os.path.exists(db_path) else 0
devices = {}
for path, per_upload, fixed in (
(tempfile.gettempdir(), 2, 0),
(folder, 1, current),
):
entry = devices.setdefault(_device(path), {"path": path, "per": 0, "fixed": 0})
entry["per"] += per_upload
entry["fixed"] += fixed
room = min(
(_free_bytes(e["path"]) - e["fixed"] - DISK_MARGIN_BYTES) // e["per"]
for e in devices.values()
)
room = (_free_bytes(folder) - current - DISK_MARGIN_BYTES) // 3
return max(0, room)
@@ -306,142 +290,13 @@ def _changed_line(operation, count):
return f"Removed {count:,} {noun}"
def _remove_quietly(*paths):
for path in paths:
try:
if path and os.path.exists(path):
os.remove(path)
except OSError:
pass
def _reserve_backup_name(db_path, label="backup"):
"""Create and return an unused ``<db>.<label>_<stamp>[_n]`` (O_EXCL)."""
for candidate in _backup_names(db_path, label):
try:
fd = os.open(candidate, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
except FileExistsError:
continue
os.close(fd)
return candidate
raise SafetyBackupFailed("No free backup name")
def _backup_names(db_path, label="backup"):
base = f"{db_path}.{label}_{_now().strftime('%Y%m%d_%H%M%S_%f')}"
for attempt in range(_NAME_ATTEMPTS):
yield base if attempt == 0 else f"{base}_{attempt}"
def _place_backup(temp, db_path, label="backup"):
"""Give the verified ``temp`` its unique backup name; return that path.
``os.link`` is atomic and refuses an existing name, so no empty or half
file ever carries a backup name, even across a crash. A filesystem
without hard links falls back to reserving the name (``O_EXCL``) and
``os.replace``. The temp file is removed by the caller.
"""
for candidate in _backup_names(db_path, label):
try:
os.link(temp, candidate)
except FileExistsError:
continue
except (OSError, NotImplementedError):
break # no hard links here
else:
return candidate
else:
raise SafetyBackupFailed("No free backup name")
final = _reserve_backup_name(db_path, label)
try:
os.replace(temp, final)
except OSError:
_remove_quietly(final)
raise
return final
SAFETY_TEMP_PREFIX = ".pm-safety-"
def sweep_safety_temps(db_path):
"""Remove the temp copies an interrupted safety backup left (blocking).
Called once at startup, before any job can take a backup. Returns the
names removed.
"""
folder = os.path.dirname(os.path.abspath(db_path))
removed = []
try:
names = os.listdir(folder)
except OSError:
return removed
for name in names:
if not (name.startswith(SAFETY_TEMP_PREFIX) and ".tmp" in name):
continue
try:
os.remove(os.path.join(folder, name))
removed.append(name)
except OSError:
continue
return removed
def _standalone(path):
"""Make the copy one self-contained file (rollback journal, not WAL).
The backup API copies the live database's WAL flag; a WAL file opened
read-only later (a verify, a restore preview) would leave ``-wal`` and
``-shm`` sidecars beside it.
"""
try:
with closing(sqlite3.connect(path)) as conn:
conn.execute("PRAGMA journal_mode=DELETE").fetchone()
except sqlite3.Error as exc:
raise SafetyBackupFailed("The backup could not be finished") from exc
def _fsync(path):
with open(path, "rb+") as handle:
os.fsync(handle.fileno())
def take_safety_backup(model, label="backup"):
"""Write and verify a safety backup; return its file name (blocking).
The copy is made in a temp file beside the database and only renamed
onto its final, unique name ``<db>.<label>_<stamp>`` once it verified,
so a backup file on disk is always a verified one. ``label`` "backup"
joins the pruned pool; another label (the workflow storage migration's
``pre-workflow-storage``) is kept until the user removes it. Raises
SafetyBackupFailed, leaving no file.
See ``database/safety_backup.py``: a backup file on disk is always a
verified one. Raises SafetyBackupFailed, leaving no file.
"""
db_path = os.path.abspath(model.db_path)
fd, temp = tempfile.mkstemp(
dir=os.path.dirname(db_path), prefix=SAFETY_TEMP_PREFIX, suffix=".tmp"
)
os.close(fd)
final = None
try:
if not model.backup_database(temp):
raise SafetyBackupFailed("The backup could not be written")
_standalone(temp)
ok, reason = model.verify_database_file(temp)
if not ok:
raise SafetyBackupFailed("The backup did not verify", reason)
_fsync(temp)
if os.name != "nt":
os.chmod(temp, 0o600)
final = _place_backup(temp, db_path, label)
return os.path.basename(final)
except OSError as exc:
_remove_quietly(final)
raise SafetyBackupFailed("The backup could not be put in place") from exc
except SafetyBackupFailed:
_remove_quietly(final)
raise
finally:
_remove_quietly(temp, temp + "-journal", temp + "-wal", temp + "-shm")
return os.path.basename(_backups.take_safety_backup(model, label))
def _failure(message, status, code=None, **extra):
@@ -504,6 +359,19 @@ class MaintenanceRoutesMixin:
return data, None
async def preview_maintenance(self, request):
"""Review one operation under a registry lease (#249 S8).
The preview's read-only connections must not be open while a restore
switches the database: it is refused (409 ``busy``) while an
``exclusive`` job runs, and a restore waits for running previews.
"""
try:
async with self._jobs.lease(MAINTENANCE_PREVIEW_LEASE):
return await self._preview_maintenance(request)
except JobBusy as busy:
return busy_response(busy)
async def _preview_maintenance(self, request):
data, error = await self._read_maintenance_body(
request,
("operation", "options"),
py/api/restore_plans.py
+146
View File
@@ -0,0 +1,146 @@
"""Reviewed restores: staged uploads held under single-use tokens (#249 S8).
A restore preview stages the upload (``database/restore_stage.py``) and
keeps it here under a ``secrets.token_urlsafe(32)`` token with the live
library's fingerprint and the database identity it was reviewed against.
``take`` hands the stage to the restore job exactly once; a second apply of
the token names that job (``PlanUsed``). An expired, discarded or unknown
token raises ``PlanExpired``; its stage folder is removed. A restart forgets
every token (the startup sweep removes the folders).
At most ``max_stages`` stages are held. Bytes are bounded by free disk space
(the upload is refused when the disk has no room, and held stages use that
disk), not by a fixed budget. A stage being applied belongs to its job and is
never evicted or counted.
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, List, Optional
from .maintenance_plans import PlanExpired, PlanUsed
STAGE_TTL_SECONDS = 900
MAX_STAGES = 2
_USED_KEPT = 32 # applied tokens remembered for plan-used answers
class TooManyStages(Exception):
"""Every stage slot is held by a reviewed restore."""
@dataclass
class HeldRestore:
staged: Any # database.restore_stage.StagedRestore
fingerprint: str
db_identity: Any
live: Dict[str, int]
expires_at: float
class RestorePlanStore:
def __init__(
self,
clock: Optional[Callable[[], float]] = None,
ttl: float = STAGE_TTL_SECONDS,
max_stages: int = MAX_STAGES,
):
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)
def full(self) -> bool:
return len(self._held) >= self._max
def issue(self, staged, fingerprint, db_identity, live) -> str:
"""Hold ``staged`` and return its token; TooManyStages when full."""
if self.full():
raise TooManyStages()
token = secrets.token_urlsafe(32)
self._held[token] = HeldRestore(
staged=staged,
fingerprint=fingerprint,
db_identity=db_identity,
live=dict(live),
expires_at=self._clock() + self._ttl,
)
return token
def peek(self, token) -> HeldRestore:
if token in self._used:
raise PlanUsed(self._used[token])
held = self._held.get(token)
if held is None:
raise PlanExpired()
if held.expires_at <= self._clock():
self._doomed.append(self._held.pop(token))
raise PlanExpired()
return held
def take(self, token, job_id) -> HeldRestore:
"""Hand the held stage to the job applying it (single use)."""
held = self.peek(token)
del self._held[token]
self._used[token] = job_id
while len(self._used) > _USED_KEPT:
self._used.pop(next(iter(self._used)))
return held
def give_back(self, token, held) -> bool:
"""Hold a taken stage again under its token (its job did not apply it).
False, holding nothing, when the review expired meanwhile or its slot
was reused (the store is full): the caller then discards the stage.
"""
if held.expires_at <= self._clock() or self.full():
return False
self._used.pop(token, None)
self._held[token] = held
return True
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]
return [self._held.pop(token) for token in expired] + self.take_doomed()
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)
py/api/review.py
+19 -1
View File
@@ -31,6 +31,8 @@ from aiohttp import web
from .annotations import _HEX64, _TOKEN, _UUID, _capture_error, _error
from .images import media_thumbnail_url
from .inspection_service import VIDEO_EXTENSIONS
from .job_routes import busy_response
from .jobs import JobBusy
from .prompts import bad_request
try:
@@ -239,6 +241,9 @@ def _side_label(titles: dict, side, empty: str) -> str:
return display_title(titles.get(stored, ()), text) or empty
REVIEW_DECISION_LEASE = "Needs review decision"
class ReviewRoutesMixin:
"""The Needs review queue over E1's schema and E2's recorded items."""
@@ -446,7 +451,20 @@ class ReviewRoutesMixin:
# -- decision --------------------------------------------------------------
async def decide_review_item(self, request: web.Request) -> web.Response:
"""POST one explicit decision on one pending item."""
"""POST one explicit decision on one pending item.
The decision runs under a registry lease (#249 S8): it is refused
(409 ``busy``) while a restore or another ``exclusive`` job runs, and
a restore waits for decisions already running before it switches the
database. Decisions never refuse one another or a scan.
"""
try:
async with self._jobs.lease(REVIEW_DECISION_LEASE):
return await self._decide_review_item(request)
except JobBusy as busy:
return busy_response(busy)
async def _decide_review_item(self, request: web.Request) -> web.Response:
data, error = await self._read_json_body(request, max_bytes=MAX_BODY_BYTES)
if error is not None:
return error
tests/browser/workspace-images.spec.mjs
+4 -2
View File
@@ -2184,6 +2184,8 @@ test('scan for metadata opens the review dialog from Import', async ({ page, req
await dialog.getByRole('button', { name: 'Cancel', exact: true }).click();
await expect(scanReview(page)).toHaveCount(0);
await expect(page.getByRole('button', { name: 'Import', exact: true })).toBeFocused();
const jobs = await (await request.get('/prompt_manager/jobs')).json();
expect(jobs.jobs).toEqual([]);
// No scan started: the one job is the backup's finished registry hold
// (#249 S8: a backup download holds a guarded slot while it streams).
await expect.poll(async () => (await (await request.get('/prompt_manager/jobs')).json()).jobs
.map(job => [job.kind, job.state])).toEqual([['backup', 'finished']]);
});
tests/test_api_admin.py
+5 -4
View File
@@ -2015,7 +2015,7 @@ class TestBackupRestoreContract(AdminAPITestCase):
self.assertFalse((await invalid.json())["success"])
compatible = self._snapshot_upload()
with patch(
"database.models.rotate_history_epoch",
"database.restore_stage.rotate_history_epoch",
side_effect=sqlite3.OperationalError("staging failed before safety backup"),
):
staged = await self._restore("database_file", compatible)
@@ -2109,7 +2109,8 @@ class TestBackupRestoreContract(AdminAPITestCase):
resp = await self._restore("database_file", upload)
self.assertEqual(resp.status, 200)
self.assertIn("verify_database_file", calls)
self.assertIn("restore_from_file", calls)
self.assertIn("stage_restore", calls)
self.assertIn("publish_staged", calls)
async def test_restore_keeps_at_most_ten_safety_backups(self):
import time
@@ -2150,12 +2151,12 @@ class TestBackupRestoreContract(AdminAPITestCase):
self.api.db.save_prompt(text="restored", prompt_hash="h1")
upload = self._snapshot_upload()
def boom(src_path):
def boom(staged):
raise PermissionError(
13, "Permission denied", os.path.join(self.tmpdir, "prompts.db")
)
self.api.db.model.restore_from_file = boom
self.api.db.model.publish_staged = boom
resp = await self._restore("database_file", upload)
body = await resp.text()
self.assertEqual(resp.status, 500)
tests/test_api_backup_restore.py
+24 -8
View File
@@ -38,6 +38,13 @@ from py.api import PromptManagerAPI, _public_path
from utils.hashing import generate_prompt_hash
def _is_upload(path):
"""Not the safety backup's temp copy (database/safety_backup.py)."""
from database.safety_backup import SAFETY_TEMP_PREFIX
return not os.path.basename(path).startswith(SAFETY_TEMP_PREFIX)
def _texts(path):
conn = sqlite3.connect(path)
try:
@@ -114,7 +121,7 @@ class TestBackupRoute(BackupRestoreTestCase):
created.append(result[1])
return result
with patch("py.api.admin.tempfile.mkstemp", spy_mkstemp):
with patch("tempfile.mkstemp", spy_mkstemp):
resp = await self.client.get("/prompt_manager/backup")
self.assertEqual(resp.status, 200)
@@ -125,6 +132,10 @@ class TestBackupRoute(BackupRestoreTestCase):
self.assertTrue(disposition.rstrip('"').endswith(".db"))
body = await resp.read()
# 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(
_texts(downloaded), ["backup me 0", "backup me 1", "backup me 2"]
@@ -263,13 +274,18 @@ class TestRestoreRoute(BackupRestoreTestCase):
created.append(result[1])
return result
with patch("py.api.admin.tempfile.mkstemp", spy_mkstemp):
with patch("tempfile.mkstemp", spy_mkstemp):
resp = await self._restore(upload)
self.assertEqual(resp.status, 200)
self.assertEqual(len(created), 1)
for suffix in ("", "-wal", "-shm", "-journal"):
self.assertFalse(os.path.exists(created[0] + suffix))
# One upload, in the library's folder; it and the safety backup's
# temp copy are gone afterwards.
uploads = [p for p in created if _is_upload(p)]
self.assertEqual(len(uploads), 1)
self.assertEqual(os.path.dirname(uploads[0]), self.tmpdir)
for path in created:
for suffix in ("", "-wal", "-shm", "-journal"):
self.assertFalse(os.path.exists(path + suffix))
async def test_text_file_is_rejected_and_database_untouched(self):
self._save("keep me")
@@ -398,7 +414,7 @@ class TestLargeRestores(BackupRestoreTestCase):
created.append(result[1])
return result
return patch("py.api.admin.tempfile.mkstemp", spy), created
return patch("tempfile.mkstemp", spy), created
def _assert_removed(self, created):
for path in created:
@@ -417,7 +433,7 @@ class TestLargeRestores(BackupRestoreTestCase):
self.assertEqual(
[p["text"] for p in self.api.db.search_prompts()], ["restored"]
)
self.assertEqual(len(created), 1)
self.assertEqual(len([p for p in created if _is_upload(p)]), 1)
self._assert_removed(created)
async def test_a_short_disk_is_refused_up_front(self):
@@ -455,7 +471,7 @@ class TestLargeRestores(BackupRestoreTestCase):
self.assertEqual(data["code"], "insufficient-space")
self.assertTrue(data["error"].endswith("Nothing was changed."))
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
self.assertEqual(len(created), 1)
self.assertEqual(len([p for p in created if _is_upload(p)]), 1)
self._assert_removed(created)
tests/test_api_maintenance.py
+3 -2
View File
@@ -23,6 +23,7 @@ sys.modules.setdefault("server", _mock_server)
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from database import safety_backup # noqa: E402
from database.operations import PromptDatabase # noqa: E402
from py.api import PromptManagerAPI # noqa: E402
from py.api.maintenance_plans import PlanExpired # noqa: E402
@@ -562,7 +563,7 @@ class TestSafetyBackup(ApplyTestCase):
frozen = datetime.datetime(2026, 10, 23, 12, 0, 0, 123456)
model = self.api.db.model
with patch.object(maintenance, "_now", return_value=frozen):
with patch.object(safety_backup, "_now", return_value=frozen):
names = [maintenance.take_safety_backup(model) for _ in range(3)]
self.assertEqual(
names,
@@ -1018,7 +1019,7 @@ class TestFinalReviewFixes(ApplyTestCase):
with (
patch.object(maintenance.os, "link", side_effect=link),
patch.object(
maintenance, "_reserve_backup_name", side_effect=AssertionError
safety_backup, "reserve_backup_name", side_effect=AssertionError
),
):
name = maintenance.take_safety_backup(self.api.db.model)
tests/test_diagnostics_config.py
+20 -18
View File
@@ -161,46 +161,48 @@ class TestAdminEndpointsUseConfigPath(unittest.TestCase):
self.assertNotIn("/nonexistent/custom", message)
def test_backup_uses_model_db_path(self):
"""backup_database should read from self.db.model.db_path."""
"""The backup snapshot is taken of self.db.model.db_path."""
if "server" not in sys.modules:
sys.modules["server"] = MagicMock()
from py.api.admin import AdminRoutesMixin
from py.api.backup_restore import write_backup_snapshot
from database.operations import PromptDatabase
with tempfile.NamedTemporaryFile(delete=False, suffix=".db") as f:
db_file = f.name
folder = tempfile.mkdtemp()
db_file = os.path.join(folder, "custom.db")
# A WAL-safe backup copies pages through sqlite itself, so the
# configured path has to hold a real database, not arbitrary bytes.
real_db = PromptDatabase(db_file)
try:
stub = self._make_api_stub(db_file)
stub.db.model = real_db.model
result = self._run_async(
AdminRoutesMixin.backup_database(stub, MagicMock())
)
self.assertEqual(result.content_type, "application/octet-stream")
self.assertGreater(len(result.body), 0)
self.assertTrue(result.body.startswith(b"SQLite format 3\x00"))
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")
finally:
os.unlink(snapshot)
finally:
real_db.close_all()
for suffix in ("", "-wal", "-shm"):
if os.path.exists(db_file + suffix):
os.unlink(db_file + suffix)
for name in os.listdir(folder):
os.unlink(os.path.join(folder, name))
os.rmdir(folder)
def test_backup_reports_missing_custom_path(self):
"""backup_database should return 404 when configured path doesn't exist."""
"""The backup answers 404 when the configured path doesn't exist."""
if "server" not in sys.modules:
sys.modules["server"] = MagicMock()
from py.api.admin import AdminRoutesMixin
from py.api.backup_restore import BackupRestoreRoutesMixin
stub = self._make_api_stub("/nonexistent/custom/prompts.db")
result = self._run_async(AdminRoutesMixin.backup_database(stub, MagicMock()))
result = self._run_async(
BackupRestoreRoutesMixin._stream_backup(stub, MagicMock())
)
body = json.loads(result.body)
self.assertFalse(body["success"])
self.assertEqual(result.status, 404)
tests/test_execution_history_restore.py
+1 -1
View File
@@ -512,7 +512,7 @@ class HistoryRestoreTests(unittest.TestCase):
subscribe_history_restore(self.path, lambda: notices.append(True))
)
with patch(
"database.models.rotate_history_epoch",
"database.restore_stage.rotate_history_epoch",
side_effect=sqlite3.OperationalError("injected staging failure"),
):
with self.assertRaisesRegex(sqlite3.OperationalError, "injected staging"):
tests/test_job_leases.py
+100
View File
@@ -0,0 +1,100 @@
"""Registry leases: short request-scoped work versus a database switch (#249 S8).
A lease is refused while an ``exclusive`` job (restore, Move library,
Optimize) runs, never refuses another lease or any job, is not listed, and an
exclusive job can await the leases already running. Ordering is driven by
events only.
"""
import asyncio
import os
import sys
import unittest
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from py.api.jobs import JobBusy, JobRegistry # noqa: E402
class TestLeases(unittest.IsolatedAsyncioTestCase):
def setUp(self):
self.registry = JobRegistry()
def start(self, kind, admission):
gate = asyncio.Event()
async def runner(job):
await gate.wait()
job = self.registry.start(
kind, kind.title(), admission=admission, scope={}, runner=runner
)
return job, gate
async def test_leases_share_and_are_not_listed(self):
scan, gate = self.start("scan", "records")
async with self.registry.lease("Needs review decision"):
async with self.registry.lease("Needs review decision"):
self.assertEqual(self.registry.leases, {"Needs review decision": 2})
# A lease never refuses a job, even an exclusive one.
self.assertIsNone(self.registry.blocker("health", admission="read"))
self.assertEqual(self.registry.leases, {})
self.assertEqual([job.kind for job in self.registry.list()], ["scan"])
gate.set()
await scan.task
async def test_lease_refused_while_exclusive_job_runs(self):
restore, gate = self.start("restore", "exclusive")
with self.assertRaises(JobBusy) as refused:
async with self.registry.lease("Maintenance preview"):
self.fail("admitted beside a restore")
self.assertIs(refused.exception.running, restore)
self.assertEqual(self.registry.leases, {})
gate.set()
await restore.task
async with self.registry.lease("Maintenance preview"):
pass
async def test_exclusive_job_awaits_running_leases(self):
entered, release = asyncio.Event(), asyncio.Event()
order = []
async def decision():
async with self.registry.lease("Needs review decision"):
entered.set()
await release.wait()
order.append("decision done")
task = asyncio.ensure_future(decision())
await entered.wait()
drained = asyncio.ensure_future(self.registry.drain_leases())
await asyncio.sleep(0)
self.assertFalse(drained.done())
release.set()
self.assertTrue(await drained)
order.append("drained")
await task
self.assertEqual(order, ["decision done", "drained"])
self.assertTrue(await self.registry.drain_leases()) # none: at once
async def test_drain_reports_a_lease_that_outlives_the_ceiling(self):
hold = self.registry.lease("Stuck preview")
await hold.__aenter__()
try:
with self.assertLogs("py.api.jobs", "WARNING") as logs:
self.assertFalse(await self.registry.drain_leases(timeout=0))
self.assertIn("Stuck preview", "\n".join(logs.output))
self.assertEqual(self.registry._drain_waiters, [])
finally:
await hold.__aexit__(None, None, None)
self.assertEqual(self.registry.leases, {})
async def test_a_failed_body_still_releases_the_lease(self):
with self.assertRaises(RuntimeError):
async with self.registry.lease("Needs review decision"):
raise RuntimeError("boom")
self.assertEqual(self.registry.leases, {})
if __name__ == "__main__":
unittest.main()
tests/test_restore_plans.py
+599
View File
@@ -0,0 +1,599 @@
"""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_<date and time>"},
)
self.assertNotIn(self.tmpdir, str(body))
# Staged in the library's folder; the upload itself is gone.
self.assertEqual(len(self.stages()), 1)
self.assertEqual(
[n for n in os.listdir(self.tmpdir) if "upload" in n and "pm-" in n], []
)
self.assertEqual(self.api._jobs.list(), [])
async def test_invalid_upload_is_400_and_leaves_nothing(self):
resp = await self.client.post(PREVIEW, data=self.form(b"not a database"))
body = await resp.json()
self.assertEqual((resp.status, body["code"]), (400, "invalid-backup"))
self.assertEqual(self.stages(), [])
async def test_expired_or_discarded_token_is_410(self):
body = await self.preview(self.upload(["restored"]))
resp = await self.client.delete(f"{PREVIEW}/{body['token']}")
self.assertEqual(resp.status, 200)
self.assertEqual(self.stages(), [])
refused = await self.apply(body["token"], status=410)
self.assertEqual(refused["code"], "plan-expired")
resp = await self.client.delete(f"{PREVIEW}/{body['token']}")
self.assertEqual(resp.status, 410)
clock = [1000.0]
self.api._restore_plan_store = restore_plans.RestorePlanStore(
clock=lambda: clock[0]
)
body = await self.preview(self.upload(["restored"], "second.db"))
clock[0] += restore_plans.STAGE_TTL_SECONDS + 1
refused = await self.apply(body["token"], status=410)
self.assertEqual(refused["code"], "plan-expired")
self.assertEqual(self.stages(), [])
async def test_stage_budget_429(self):
data = self.upload(["restored"])
first = await self.preview(data)
await self.preview(data)
refused = await self.preview(data, status=429)
self.assertEqual(refused["code"], "too-large")
self.assertEqual(len(self.stages()), 2)
await self.client.delete(f"{PREVIEW}/{first['token']}")
await self.preview(data)
async def test_preview_reports_repaired_documents(self):
upload = os.path.join(self.tmpdir, "uploads", "damaged.db")
os.makedirs(os.path.dirname(upload), exist_ok=True)
src = PromptDatabase(upload)
prompt_id = self.save("with a workflow", src)
image = src.link_image_to_prompt(
prompt_id, os.path.join(self.tmpdir, "a.png"), {"workflow": {"v": 1}}
)
conn = src.model.get_connection()
with conn:
conn.execute(
"UPDATE workflow_documents SET data=x'00' WHERE id=(SELECT "
"workflow_document_id FROM generated_images WHERE id=?)",
(image,),
)
src.close_all()
body = await self.preview(Path(upload).read_bytes())
self.assertEqual(body["documents"], {"repaired": 0, "cleared": 1, "images": 1})
class TestApply(RestoreTestCase):
async def test_apply_restores_and_logs(self):
self.save("old")
old_db = self.api.db
body = await self.preview(self.upload(["restored a", "restored b"]))
started = await self.apply(body["token"])
summary = started["job"]
self.assertEqual(
(summary["kind"], summary["title"], summary["cancellable"]),
("restore", RESTORE_JOB_TITLE, False),
)
self.assertEqual(summary["admission"], "exclusive")
job = await self.finish(summary["id"])
self.assertEqual(job.state, "finished", job.log)
self.assertIsNot(self.api.db, old_db)
self.assertEqual(_texts(self.db_path), ["restored a", "restored b"])
[backup] = self.backups()
self.assertEqual(_texts(backup), ["old"])
log = "\n".join(job.log)
self.assertIn(f"Safety backup {os.path.basename(backup)}", log)
self.assertIn("Restored 2 prompts", log)
self.assertNotIn(self.tmpdir, log)
self.assertEqual(self.stages(), [])
wal = self.db_path + "-wal"
self.assertEqual(os.path.getsize(wal) if os.path.exists(wal) else 0, 0)
used = await self.apply(body["token"], status=409)
self.assertEqual(used["code"], "plan-used")
self.assertEqual(used["job"]["id"], summary["id"])
async def test_apply_binds_exact_staged_bytes(self):
self.save("live")
body = await self.preview(self.upload(["restored"]))
held = self.api._restore_plans.peek(body["token"])
with closing(sqlite3.connect(held.staged.path)) as conn:
conn.execute("INSERT INTO prompts (text, hash) VALUES ('sneaked', 'x')")
conn.commit()
started = await self.apply(body["token"])
job = await self.finish(started["job"]["id"])
self.assertEqual(job.state, "failed")
self.assertEqual(job.last_event["code"], "plan-changed")
self.assertEqual(_texts(self.db_path), ["live"])
self.assertEqual(self.backups(), [])
self.assertEqual(self.stages(), [])
async def test_live_change_after_preview_is_409(self):
self.save("live")
body = await self.preview(self.upload(["restored"]))
self.save("written after the review")
refused = await self.apply(body["token"], status=409)
self.assertEqual(refused["code"], "plan-changed")
self.assertEqual(
sorted(p["text"] for p in self.api.db.search_prompts()),
["live", "written after the review"],
)
self.assertEqual(self.stages(), [])
self.assertEqual(self.api._jobs.list(), [])
async def test_publish_failure_keeps_safety_backup_and_reports(self):
self.save("live")
body = await self.preview(self.upload(["restored"]))
failure = sqlite3.OperationalError("injected page copy failure")
with patch.object(PromptModel, "_replace_contents_from", side_effect=failure):
started = await self.apply(body["token"])
job = await self.finish(started["job"]["id"])
self.assertEqual(job.state, "failed")
self.assertEqual(job.last_event["code"], "restore-failed")
self.assertIn("The safety backup was kept.", job.status_line)
[backup] = self.backups()
self.assertEqual(_texts(backup), ["live"])
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["live"])
self.assertEqual(self.stages(), [])
async def test_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"),
)
tests/test_restore_review_fixes.py
+579
View File
@@ -0,0 +1,579 @@
"""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"]))
try:
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"]
)
finally:
# The patched close left the old object's executor-thread
# connections open; close them before teardown removes the files
# (once the object is collected, close_all_instances misses it).
await asyncio.get_running_loop().run_in_executor(None, old_db.close_all)
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()], [])
class TestRereviewFixes(RestoreTestCase):
"""PR #335 re-review: R1 to R4."""
def test_r1_a_merge_into_a_damaged_holder_keeps_both_images(self):
rekey = TestRekeyDocuments.library
path, text = rekey(self, lambda digest, length: (b"\0" * 32, length))
raw = text.encode()
with closing(sqlite3.connect(path)) as conn:
holder = conn.execute(
"INSERT INTO workflow_documents (sha256, byte_length, data) "
"VALUES (?,?,?)",
(hashlib.sha256(raw).digest(), len(raw), b"garbage-not-zlib"),
).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/b.png', 'b.png', ?)",
(prompt_id, holder),
)
conn.commit()
staged = self.api.db.model.stage_restore(path)
try:
self.assertEqual(
staged.documents, {"repaired": 2, "cleared": 0, "images": 0}
)
self.assertEqual(
TestRekeyDocuments.staged_texts(self, staged), [text, text]
)
with closing(sqlite3.connect(staged.path)) as conn:
self.assertEqual(
conn.execute("SELECT COUNT(*) FROM workflow_documents").fetchone(),
(1,),
)
finally:
staged.discard()
async def test_r2_a_losing_duplicate_apply_leaves_the_taken_stage(self):
self.save("live")
body = await self.preview(self.upload(["restored"]))
store = self.api._restore_plans
held = store.peek(body["token"])
entered, release = threading.Event(), threading.Event()
real = self.api._live_fingerprint
def gated(db):
entered.set()
release.wait(WAIT)
return real(db)
with patch.object(self.api, "_live_fingerprint", side_effect=gated):
second = asyncio.ensure_future(
self.client.post(
"/prompt_manager/restore/apply", json={"token": body["token"]}
)
)
loop = asyncio.get_running_loop()
self.assertTrue(await loop.run_in_executor(None, entered.wait, WAIT))
# Meanwhile the first apply's job took the stage, and a write landed.
store.take(body["token"], "first-job")
self.save("a run count")
release.set()
resp = await second
self.assertEqual(resp.status, 409)
self.assertTrue(os.path.isdir(held.staged.directory))
await self._discard_quietly(held)
async def _discard_quietly(self, held):
await asyncio.get_running_loop().run_in_executor(None, held.staged.discard)
def test_r3_give_back_refuses_a_full_store_or_an_expired_review(self):
clock = [1000.0]
store = backup_restore.RestorePlanStore(clock=lambda: clock[0], max_stages=1)
first = store.issue(_Stage(), "f", "id", {})
held = store.take(first, "job")
other = store.issue(_Stage(), "f", "id", {})
self.assertFalse(store.give_back(first, held)) # the slot was reused
self.assertEqual(len(store), 1)
store.pop(other)
clock[0] = held.expires_at
self.assertFalse(store.give_back(first, held)) # expired meanwhile
clock[0] -= 1
self.assertTrue(store.give_back(first, held))
self.assertIs(store.peek(first), held)
async def test_r4_a_cancelled_backup_removes_its_snapshot(self):
from aiohttp.test_utils import make_mocked_request
self.save("keep")
entered, release, done = (threading.Event() for _ in range(3))
real = backup_restore.write_backup_snapshot
def slow(model):
entered.set()
release.wait(WAIT)
try:
return real(model)
finally:
done.set()
with patch.object(backup_restore, "write_backup_snapshot", slow):
task = asyncio.ensure_future(
self.api._stream_backup(
make_mocked_request("GET", "/prompt_manager/backup")
)
)
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))
async def snapshot_gone():
while any(n.startswith(".pm-safety-") for n in os.listdir(self.tmpdir)):
await asyncio.sleep(0)
# A ceiling only a regression reaches (the file never goes).
await asyncio.wait_for(snapshot_gone(), WAIT)
self.assertIsNone(self.api._jobs.active("backup"))
class _Stage:
def discard(self):
pass
tests/test_restore_stage.py
+440
View File
@@ -0,0 +1,440 @@
"""A restore is staged beside the library, then published (#249 System S8).
``PromptModel.stage_restore`` copies an upload into a stage folder in the
library's own folder, migrates, repairs and verifies it there, and never
touches the live database or the source. ``publish_staged`` re-checks the
staged bytes, keeps a uniquely named safety backup and copies the stage over
the live pages, then truncates the WAL. ``restore_from_file`` is the two in a
row, with the behavior it always had.
"""
import datetime
import hashlib
import os
import sqlite3
import sys
import tempfile
import threading
import unittest
import zlib
from contextlib import closing
from pathlib import Path
from unittest.mock import patch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from database import restore_stage, safety_backup # noqa: E402
from database import workflow_documents as wd # noqa: E402
from database.models import PromptModel # noqa: E402
from database.operations import PromptDatabase # noqa: E402
from utils.hashing import generate_prompt_hash # noqa: E402
try:
from tests.open_handles import assert_closed # noqa: E402
except ImportError: # discovered with tests/ as the top-level directory
from open_handles import assert_closed # noqa: E402
WAIT = 60 # a bound only a hung test reaches
def _texts(path):
with closing(sqlite3.connect(path)) as conn:
return sorted(r[0] for r in conn.execute("SELECT text FROM prompts"))
def _digests(path):
"""sha256 of the file and each sidecar that exists."""
return {
suffix: hashlib.sha256(Path(path + suffix).read_bytes()).hexdigest()
for suffix in ("", "-wal", "-shm", "-journal")
if os.path.exists(path + suffix)
}
class StageTestCase(unittest.TestCase):
def setUp(self):
self.tmpdir = tempfile.mkdtemp()
self.library = os.path.join(self.tmpdir, "library")
os.mkdir(self.library)
self.uploads = os.path.join(self.tmpdir, "uploads")
os.mkdir(self.uploads)
self.path = os.path.join(self.library, "prompts.db")
self.db = PromptDatabase(self.path)
def tearDown(self):
PromptModel.close_all_instances()
assert_closed(self, self.tmpdir)
for root, dirs, files in os.walk(self.tmpdir, topdown=False):
for name in files:
os.unlink(os.path.join(root, name))
for name in dirs:
os.rmdir(os.path.join(root, name))
os.rmdir(self.tmpdir)
def save(self, db, text):
return db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
def source(self, texts, name="source.db", journal="wal"):
path = os.path.join(self.uploads, name)
src = PromptDatabase(path)
for text in texts:
self.save(src, text)
src.close_all()
if journal != "wal":
with closing(sqlite3.connect(path)) as conn:
conn.execute(f"PRAGMA journal_mode={journal}").fetchone()
return path
def stage_dirs(self):
return sorted(
name
for name in os.listdir(self.library)
if name.startswith(restore_stage.STAGE_PREFIX)
)
def backups(self):
return sorted(
name
for name in os.listdir(self.library)
if name.startswith("prompts.db.backup_")
)
class TestStageAndPublish(StageTestCase):
def test_restore_from_file_unchanged_via_stage_publish(self):
self.save(self.db, "old prompt")
src = self.source(["restored one", "restored two"])
model = self.db.model
calls = []
real_stage, real_publish = model.stage_restore, model.publish_staged
def stage(path):
calls.append("stage")
return real_stage(path)
def publish(staged):
calls.append("publish")
return real_publish(staged)
with (
patch.object(model, "stage_restore", stage),
patch.object(model, "publish_staged", publish),
):
backup_path = model.restore_from_file(src)
self.assertEqual(calls, ["stage", "publish"])
self.assertTrue(backup_path.startswith(self.path + ".backup_"))
self.assertEqual(_texts(backup_path), ["old prompt"])
self.assertEqual(_texts(self.path), ["restored one", "restored two"])
self.assertEqual(self.stage_dirs(), [])
def test_stage_never_touches_live_or_source(self):
self.save(self.db, "live")
self.db.model.get_connection() # the live file stays open, WAL and all
src = self.source(["restored"])
live_before = _digests(self.path)
source_before = _digests(src)[""]
staged = self.db.model.stage_restore(src)
try:
self.assertEqual(_digests(self.path), live_before)
# Read-only opens of a WAL-flagged file add empty sidecars, as
# they always did; the source's bytes never change.
self.assertEqual(_digests(src)[""], source_before)
self.assertEqual(self.backups(), [])
self.assertEqual(_texts(self.path), ["live"])
finally:
staged.discard()
def test_stage_lives_in_db_directory_and_is_removed(self):
src = self.source(["restored"])
staged = self.db.model.stage_restore(src)
stage_dir = os.path.dirname(staged.path)
self.assertEqual(os.path.dirname(stage_dir), self.library)
self.assertEqual(self.stage_dirs(), [os.path.basename(stage_dir)])
# One self-contained file: no sidecar of an open or WAL database.
self.assertEqual(os.listdir(stage_dir), [os.path.basename(staged.path)])
staged.discard()
self.assertEqual(self.stage_dirs(), [])
staged.discard() # a second discard is harmless
# A failed stage removes its folder too.
bad = os.path.join(self.uploads, "notes.db")
Path(bad).write_bytes(b"not a database")
with self.assertRaises(ValueError):
self.db.model.stage_restore(bad)
self.assertEqual(self.stage_dirs(), [])
def test_staged_handles_closed_before_cleanup(self):
src = self.source(["restored"])
staged = self.db.model.stage_restore(src)
assert_closed(self, os.path.dirname(staged.path))
self.db.model.publish_staged(staged)
assert_closed(self, os.path.dirname(staged.path))
staged.discard()
self.assertFalse(os.path.exists(os.path.dirname(staged.path)))
def test_safety_backup_name_unique_within_one_second(self):
self.save(self.db, "first")
frozen = datetime.datetime(2026, 10, 25, 9, 30, 0, 123456)
names = []
with patch.object(safety_backup, "_now", return_value=frozen):
for text in ("second", "third"):
staged = self.db.model.stage_restore(self.source([text], f"{text}.db"))
try:
names.append(os.path.basename(self.db.model.publish_staged(staged)))
finally:
staged.discard()
self.assertEqual(
names,
[
"prompts.db.backup_20261025_093000_123456",
"prompts.db.backup_20261025_093000_123456_1",
],
)
self.assertEqual(_texts(os.path.join(self.library, names[1])), ["second"])
for name in names: # standalone: no WAL flag, no sidecars when read
with closing(sqlite3.connect(os.path.join(self.library, name))) as conn:
self.assertEqual(
conn.execute("PRAGMA journal_mode").fetchone()[0], "delete"
)
if os.name != "nt":
mode = os.stat(os.path.join(self.library, name)).st_mode & 0o777
self.assertEqual(mode, 0o600)
def test_publish_binds_the_staged_bytes(self):
self.save(self.db, "live")
staged = self.db.model.stage_restore(self.source(["restored"]))
try:
with closing(sqlite3.connect(staged.path)) as conn:
conn.execute("INSERT INTO prompts (text, hash) VALUES ('sneaked', 'x')")
conn.commit()
with self.assertRaises(restore_stage.StageChanged):
self.db.model.publish_staged(staged)
finally:
staged.discard()
self.assertEqual(_texts(self.path), ["live"])
self.assertEqual(self.backups(), [])
def test_summary_schema_and_digest(self):
prompt_id = self.save(self.db, "unused")
src = self.source(["a", "b"])
staged = self.db.model.stage_restore(src)
try:
self.assertEqual(staged.summary["prompts"], 2)
self.assertEqual(staged.summary["positives"], 2)
for key in ("pairs", "archived_pairs", "image_links", "runs"):
self.assertEqual(staged.summary[key], 0, key)
self.assertEqual(staged.summary["saved_views"], 0)
self.assertEqual(staged.schema, {"migrated": False, "library_state": True})
self.assertEqual(
staged.digest,
"sha256:" + hashlib.sha256(Path(staged.path).read_bytes()).hexdigest(),
)
self.assertEqual(staged.size, os.path.getsize(staged.path))
live = restore_stage.library_summary(self.db.model.get_connection())
self.assertEqual(live["prompts"], 1)
self.assertTrue(prompt_id)
finally:
staged.discard()
class TestOlderBackups(StageTestCase):
"""Rollback-journal and older-schema backups restore exactly as before."""
def test_rollback_journal_backup_restores(self):
self.save(self.db, "live")
src = self.source(["from a delete-mode file"], journal="delete")
with closing(sqlite3.connect(src)) as conn:
self.assertEqual(
conn.execute("PRAGMA journal_mode").fetchone()[0], "delete"
)
before = _digests(src)
self.db.model.restore_from_file(src)
self.assertEqual(_texts(self.path), ["from a delete-mode file"])
self.assertEqual(_digests(src), before)
def test_older_schema_backup_is_migrated(self):
src = os.path.join(self.uploads, "legacy.db")
with closing(sqlite3.connect(src)) as conn:
conn.execute(
"CREATE TABLE prompts (id INTEGER PRIMARY KEY AUTOINCREMENT,"
" text TEXT NOT NULL, created_at TIMESTAMP, updated_at TIMESTAMP,"
" category TEXT, tags TEXT, rating INTEGER, notes TEXT,"
" hash TEXT UNIQUE, workflow_name TEXT)"
)
conn.execute("INSERT INTO prompts (text) VALUES ('legacy')")
conn.commit()
staged = self.db.model.stage_restore(src)
try:
self.assertEqual(staged.schema, {"migrated": True, "library_state": False})
self.assertEqual(staged.summary["prompts"], 1)
self.db.model.publish_staged(staged)
finally:
staged.discard()
self.assertEqual(_texts(self.path), ["legacy"])
ok, reason = PromptModel.verify_database_file(
self.path, require_library_state=True
)
self.assertTrue(ok, reason)
class TestWalAfterRestore(StageTestCase):
"""Carry-over 4: the live WAL is checkpointed and truncated after a restore."""
def test_wal_is_truncated_after_restore(self):
src = self.source(["restored"])
with closing(sqlite3.connect(src)) as conn:
conn.execute("CREATE TABLE padding(b BLOB)")
conn.execute("INSERT INTO padding VALUES (zeroblob(?))", (4 * 2**20,))
conn.commit()
self.save(self.db, "live")
# Another thread's idle connection (the image monitor, a node).
bystander = PromptDatabase(self.path)
opened = threading.Event()
worker = threading.Thread(
target=lambda: (bystander.search_prompts(), opened.set())
)
worker.start()
worker.join(WAIT)
self.assertTrue(opened.is_set())
try:
self.db.model.restore_from_file(src)
self.assertEqual(_texts(self.path), ["restored"])
wal = self.path + "-wal"
self.assertEqual(os.path.getsize(wal) if os.path.exists(wal) else 0, 0)
self.assertEqual(
[p["text"] for p in bystander.search_prompts()], ["restored"]
)
finally:
bystander.close_all()
def test_a_reader_left_open_is_waited_for_then_logged(self):
src = self.source(["restored"])
self.save(self.db, "live")
stuck = PromptDatabase(self.path)
conn = stuck.model.get_connection()
conn.execute("BEGIN")
conn.execute("SELECT COUNT(*) FROM prompts").fetchone() # holds a snapshot
try:
with (
patch.object(restore_stage, "CHECKPOINT_BUSY_MS", 50),
self.assertLogs(
"prompt_manager.database.restore_stage", "WARNING"
) as logs,
):
self.db.model.restore_from_file(src)
finally:
conn.rollback()
stuck.close_all()
self.assertIn("could not be truncated", "\n".join(logs.output))
self.assertEqual(_texts(self.path), ["restored"])
class TestDamagedDocuments(StageTestCase):
"""Carry-over 5: repair what can be repaired, report what cannot."""
def library_with_documents(self, name, inline_copy):
"""A backup whose two documents are damaged; one has an inline copy."""
path = os.path.join(self.uploads, name)
src = PromptDatabase(path)
prompt_id = self.save(src, "with workflows")
first = src.link_image_to_prompt(
prompt_id, os.path.join(self.uploads, "a.png"), {"workflow": {"v": 1}}
)
second = src.link_image_to_prompt(
prompt_id, os.path.join(self.uploads, "b.png"), {"workflow": {"v": 2}}
)
conn = src.model.get_connection()
ids = [
conn.execute(
"SELECT workflow_document_id FROM generated_images WHERE id=?", (i,)
).fetchone()[0]
for i in (first, second)
]
texts = [wd.read_stored_text(conn, None, document_id) for document_id in ids]
with conn:
if inline_copy:
# An older row that still holds the first text inline.
conn.execute(
"INSERT INTO generated_images (prompt_id, image_path, filename,"
" workflow_data) VALUES (?, ?, 'c.png', ?)",
(prompt_id, os.path.join(self.uploads, "c.png"), texts[0]),
)
for document_id in ids:
conn.execute(
"UPDATE workflow_documents SET data=? WHERE id=?",
(zlib.compress(b"damaged"), document_id),
)
src.close_all()
return path, ids, texts, (first, second)
def test_damaged_documents_are_repaired_or_cleared_and_reported(self):
src, ids, texts, images = self.library_with_documents("damaged.db", True)
ok, reason = PromptModel.verify_database_file(src)
self.assertFalse(ok)
self.assertTrue(reason.startswith("Workflow document "), reason)
before = _digests(src)
staged = self.db.model.stage_restore(src)
try:
self.assertEqual(
staged.documents, {"repaired": 1, "cleared": 1, "images": 1}
)
self.db.model.publish_staged(staged)
finally:
staged.discard()
self.assertEqual(_digests(src), before) # the upload is never changed
ok, reason = PromptModel.verify_database_file(self.path)
self.assertTrue(ok, reason)
conn = self.db.model.get_connection()
rows = {
row[0]: (row[1], row[2])
for row in conn.execute(
"SELECT id, workflow_data, workflow_document_id FROM generated_images"
)
}
repaired = rows[images[0]]
self.assertEqual(repaired[1], ids[0])
self.assertEqual(wd.read_stored_text(conn, None, repaired[1]), texts[0])
self.assertEqual(rows[images[1]], (None, None)) # cleared, nothing readable
def test_restore_from_file_accepts_a_backup_with_damaged_documents(self):
src, _ids, _texts_, _images = self.library_with_documents("damaged.db", False)
self.db.model.restore_from_file(src)
ok, reason = PromptModel.verify_database_file(self.path)
self.assertTrue(ok, reason)
self.assertEqual(_texts(self.path), ["with workflows"])
def test_an_intact_backup_reports_nothing_repaired(self):
staged = self.db.model.stage_restore(self.source(["clean"]))
try:
self.assertEqual(
staged.documents, {"repaired": 0, "cleared": 0, "images": 0}
)
finally:
staged.discard()
class TestSweep(StageTestCase):
def test_interrupted_stages_and_uploads_are_swept(self):
stage = os.path.join(self.library, restore_stage.STAGE_PREFIX + "abc")
os.mkdir(stage)
Path(stage, "staged.db").write_bytes(b"x")
upload = os.path.join(self.library, restore_stage.UPLOAD_PREFIX + "x.tmp")
Path(upload).write_bytes(b"x")
keep = os.path.join(self.library, "notes.txt")
Path(keep).write_bytes(b"x")
swept = restore_stage.sweep_restore_temps(self.path)
self.assertEqual(
sorted(swept), sorted([os.path.basename(stage), os.path.basename(upload)])
)
self.assertEqual(self.stage_dirs(), [])
self.assertTrue(os.path.exists(keep))
if __name__ == "__main__":
unittest.main()
tests/test_system_api.py
+22 -10
View File
@@ -60,8 +60,8 @@ class SystemAPITestCase(AdminAPITestCase):
self.assertEqual(response.status, status, await response.text())
return await response.json()
async def restore_snapshot(self):
"""Restore the live database from a copy of itself through the legacy route."""
def snapshot_form(self):
"""A multipart upload of a copy of the live database."""
from contextlib import closing
from aiohttp import FormData
@@ -74,8 +74,12 @@ class SystemAPITestCase(AdminAPITestCase):
upload = handle.read()
form = FormData()
form.add_field("database_file", upload, filename="upload.db")
return form
async def restore_snapshot(self):
"""Restore the live database from a copy of itself through the legacy route."""
response = await self.client.request(
"POST", "/prompt_manager/restore", data=form
"POST", "/prompt_manager/restore", data=self.snapshot_form()
)
self.assertEqual(response.status, 200, await response.text())
@@ -193,8 +197,10 @@ class TestHealthRoute(SystemAPITestCase):
overview = await self.json("GET", "/prompt_manager/system/overview")
self.assertIsNone(overview["health"])
async def test_check_finishing_after_a_restore_is_not_kept(self):
# A check that started on the old database describes it, not the new one.
async def test_restore_refused_while_a_check_runs(self):
# #249 S8 carry-over: a restore and the health check never interleave
# (the restore is exclusive), so no check outlives the database it
# read; a restore after the check forgets its report.
self.save("Any")
release = threading.Event()
entered = threading.Event()
@@ -209,18 +215,24 @@ class TestHealthRoute(SystemAPITestCase):
body = await self.json("POST", "/prompt_manager/system/health", 202)
loop = asyncio.get_running_loop()
await loop.run_in_executor(None, entered.wait, WAIT)
await self.restore_snapshot()
refused = await self.client.post(
"/prompt_manager/restore", data=self.snapshot_form()
)
self.assertEqual(refused.status, 409)
self.assertEqual(
await refused.json(),
{"success": False, "error": "Health check already running"},
)
release.set()
job = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
self.assertEqual(job.state, "finished")
body = await self.json("GET", "/prompt_manager/system/health")
self.assertEqual(body, {"success": True, "report": None})
# A check started on the restored database is kept again.
await self.run_health()
self.assertIsNotNone(
(await self.json("GET", "/prompt_manager/system/health"))["report"]
)
await self.restore_snapshot()
body = await self.json("GET", "/prompt_manager/system/health")
self.assertEqual(body, {"success": True, "report": None})
async def test_health_sources_resolve_off_the_event_loop(self):
loop_thread = threading.get_ident()