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

refactor(database): stage restores before publishing

Browse Source 
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.
This commit is contained in:
Vito Sansevero
2026-10-25 05:54:20 -07:00
parent 2964a3b632
commit 901d40d007
10 changed files with 1144 additions and 219 deletions
Download Patch File Download Diff File Expand all files Collapse all files
database/models.py
+48 -61
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
@@ -1132,7 +1131,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 +1145,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 +1223,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 +1257,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 +1303,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
+357
View File
@@ -0,0 +1,357 @@
"""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 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.
"""
import json
with conn:
if not conn.in_transaction:
conn.execute("BEGIN") # one read snapshot for every count
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
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/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
+78
View File
@@ -475,3 +475,81 @@ 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"
)
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 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 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.
"""
damaged: Dict[bytes, int] = {}
for row in conn.execute(ALL_DOCUMENTS_SQL):
if _decode(row[:3], row[3], quiet=True) is None:
damaged[bytes(row[0])] = row[3]
result = {"repaired": 0, "cleared": 0, "images": 0}
if not damaged:
return result
with conn:
copies = _inline_copies(conn, damaged)
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
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
f"workflow_document_id IN ({marks}) OR prompt_document_id IN ({marks})",
chunk + chunk,
).fetchone()[0]
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 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"
)
result.update(repaired=len(copies), cleared=len(lost), images=images)
return result
py/api/maintenance.py
+15 -149
View File
@@ -48,7 +48,6 @@ import os
import shutil
import sqlite3
import tempfile
from contextlib import closing
from aiohttp import web
@@ -58,10 +57,22 @@ 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
@@ -95,18 +106,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):
@@ -117,10 +116,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
@@ -305,142 +300,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):
tests/test_api_admin.py
+1 -1
View File
@@ -1963,7 +1963,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)
tests/test_api_backup_restore.py
+16 -5
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:
@@ -267,9 +274,13 @@ class TestRestoreRoute(BackupRestoreTestCase):
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)
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")
@@ -417,7 +428,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 +466,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,
@@ -987,7 +988,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_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_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()