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