Merge origin/main (System S8) into feat/327-system-s10-integrations
This commit is contained in:
Vito Sansevero
+52
-62
@@ -9,7 +9,6 @@ import math
|
||||
import sqlite3
|
||||
import os
|
||||
import threading
|
||||
import tempfile
|
||||
import weakref
|
||||
from contextlib import closing
|
||||
from pathlib import Path
|
||||
@@ -20,6 +19,7 @@ from .managed_connection import (
|
||||
await_closed,
|
||||
retire_connections,
|
||||
)
|
||||
from . import restore_stage
|
||||
from .library_schema import library_schema_error, migrate_library_state
|
||||
from .image_favorites import image_favorites_schema_error, migrate_image_favorites
|
||||
from .annotation_schema import annotation_schema_error, migrate_annotation_state
|
||||
@@ -39,7 +39,6 @@ from .workflow_documents import (
|
||||
from .execution_history_schema import (
|
||||
history_schema_error,
|
||||
migrate_execution_history,
|
||||
rotate_history_epoch,
|
||||
)
|
||||
|
||||
# Import logging system
|
||||
@@ -555,7 +554,10 @@ class PromptModel:
|
||||
"WHERE pp.positive_prompt_id IS NOT gi.prompt_id LIMIT 1"
|
||||
).fetchone():
|
||||
return "Image pair association has a different positive prompt owner"
|
||||
for row in conn.execute("SELECT provenance FROM image_prompt_pairs"):
|
||||
# fetchone, not iteration: a managed connection's cursor prefetches
|
||||
# ITER_CHUNK rows when iterated, and each provenance may be ~4 MB.
|
||||
cursor = conn.execute("SELECT provenance FROM image_prompt_pairs")
|
||||
while (row := cursor.fetchone()) is not None:
|
||||
error = pair_provenance_error(row[0])
|
||||
if error:
|
||||
return error
|
||||
@@ -1132,7 +1134,7 @@ class PromptModel:
|
||||
|
||||
@staticmethod
|
||||
def verify_database_file(
|
||||
path: str, *, require_library_state: bool = False
|
||||
path: str, *, require_library_state: bool = False, documents: bool = True
|
||||
) -> Tuple[bool, str]:
|
||||
"""
|
||||
Check that a file is a healthy prompts database before restoring it.
|
||||
@@ -1146,6 +1148,9 @@ class PromptModel:
|
||||
database has (Library state, image favorites, saved views,
|
||||
image annotations and the review queue), as a migrated
|
||||
restore stage must; old backups may lack them.
|
||||
documents: Read every workflow document and refuse a damaged
|
||||
one. A restore checks its upload without it and repairs
|
||||
damaged documents on the stage (``database/restore_stage.py``).
|
||||
|
||||
Returns:
|
||||
(ok, reason): reason is "ok" when the file is acceptable, otherwise
|
||||
@@ -1221,7 +1226,7 @@ class PromptModel:
|
||||
if review_error:
|
||||
return False, review_error
|
||||
documents_error = workflow_documents_schema_error(
|
||||
conn, allow_missing=not require_library_state, documents=True
|
||||
conn, allow_missing=not require_library_state, documents=documents
|
||||
)
|
||||
if documents_error:
|
||||
return False, documents_error
|
||||
@@ -1255,17 +1260,43 @@ class PromptModel:
|
||||
finally:
|
||||
source.close()
|
||||
|
||||
def stage_restore(self, src_path: str):
|
||||
"""Stage ``src_path`` beside this database for a restore (blocking).
|
||||
|
||||
Returns a ``StagedRestore`` (``database/restore_stage.py``): a
|
||||
migrated, repaired and verified copy in the library's folder, its
|
||||
digest and what it holds. Neither the live database nor the source
|
||||
is touched. Raises ValueError when either does not verify.
|
||||
"""
|
||||
return restore_stage.stage_restore(PromptModel, self.db_path, src_path)
|
||||
|
||||
def publish_staged(self, staged) -> str:
|
||||
"""Replace the live database with ``staged``; return the safety backup.
|
||||
|
||||
See ``restore_stage.publish_staged``: the staged bytes are rechecked,
|
||||
the current database is kept as ``<db_path>.backup_<stamp>`` and the
|
||||
WAL is truncated afterwards.
|
||||
"""
|
||||
backup_path = restore_stage.publish_staged(self, staged)
|
||||
self.logger.info("Database restored from a staged copy")
|
||||
from .execution_history import notify_history_restore
|
||||
|
||||
notify_history_restore(self.db_path)
|
||||
return backup_path
|
||||
|
||||
def restore_from_file(self, src_path: str) -> str:
|
||||
"""
|
||||
Replace the live database with the file at src_path.
|
||||
|
||||
A disposable SQLite backup of the source is migrated and verified first.
|
||||
Only then is the current database backed up next to itself as
|
||||
``<db_path>.backup_<YYYYmmdd_HHMMSS>`` (WAL content included), and the
|
||||
staged pages replace the live pages through SQLite's backup API.
|
||||
Migration failures never replace live content or modify the source.
|
||||
An online replacement failure leaves the live database intact and
|
||||
propagates; the completed safety backup remains available.
|
||||
``stage_restore`` then ``publish_staged``: a disposable copy of the
|
||||
source is migrated, repaired and verified in the library's folder
|
||||
first. Only then is the current database backed up next to itself as
|
||||
``<db_path>.backup_<YYYYmmdd_HHMMSS_ffffff>`` (WAL content included),
|
||||
and the staged pages replace the live pages through SQLite's backup
|
||||
API. Staging failures never replace live content or modify the
|
||||
source. An online replacement failure leaves the live database intact
|
||||
and propagates; the completed safety backup remains available. The
|
||||
stage is removed either way.
|
||||
|
||||
Args:
|
||||
src_path: SQLite file to restore
|
||||
@@ -1275,55 +1306,14 @@ class PromptModel:
|
||||
there was no database to back up.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the previous database could not be backed up
|
||||
ValueError: If the source or its migrated copy does not verify
|
||||
SafetyBackupFailed: If the previous database could not be backed up
|
||||
sqlite3.OperationalError: If SQLite refuses online replacement
|
||||
"""
|
||||
ok, reason = self.verify_database_file(src_path)
|
||||
if not ok:
|
||||
raise ValueError(reason)
|
||||
with tempfile.TemporaryDirectory(prefix="prompt-manager-restore-") as directory:
|
||||
staged_path = str(Path(directory) / "staged.db")
|
||||
source = sqlite3.connect(
|
||||
Path(src_path).resolve().as_uri() + "?mode=ro", uri=True
|
||||
)
|
||||
try:
|
||||
with closing(sqlite3.connect(staged_path)) as destination:
|
||||
source.backup(destination)
|
||||
finally:
|
||||
source.close()
|
||||
staged = None
|
||||
try:
|
||||
staged = PromptModel(staged_path)
|
||||
rotate_history_epoch(staged.get_connection())
|
||||
ok, reason = self.verify_database_file(
|
||||
staged_path, require_library_state=True
|
||||
)
|
||||
if not ok:
|
||||
raise ValueError(reason)
|
||||
staged.close_all()
|
||||
backup_path = ""
|
||||
if os.path.exists(self.db_path):
|
||||
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
backup_path = f"{self.db_path}.backup_{timestamp}"
|
||||
if not self.backup_database(backup_path):
|
||||
raise RuntimeError("Could not back up the current database")
|
||||
self.logger.info(f"Current database backed up to: {backup_path}")
|
||||
try:
|
||||
self._replace_contents_from(staged_path)
|
||||
except sqlite3.OperationalError as exc:
|
||||
exc._prompt_manager_restore_backup_created = bool(backup_path)
|
||||
raise
|
||||
with PromptModel._init_lock:
|
||||
PromptModel._initialized_paths.add(self._schema_key())
|
||||
self.logger.info(f"Database restored from {src_path}")
|
||||
from .execution_history import notify_history_restore
|
||||
|
||||
notify_history_restore(self.db_path)
|
||||
return backup_path
|
||||
finally:
|
||||
if staged is not None:
|
||||
staged.close_all()
|
||||
with PromptModel._init_lock:
|
||||
PromptModel._initialized_paths.discard(
|
||||
os.path.normcase(os.path.realpath(staged_path))
|
||||
)
|
||||
staged = self.stage_restore(src_path)
|
||||
try:
|
||||
backup_path = self.publish_staged(staged)
|
||||
finally:
|
||||
staged.discard()
|
||||
self.logger.info(f"Database restored from {src_path}")
|
||||
return backup_path
|
||||
|
||||
@@ -0,0 +1,365 @@
|
||||
"""A restore in two steps: stage the upload, then publish it (#249 System S8).
|
||||
|
||||
**Stage** (``stage_restore``) never touches the live database or the source.
|
||||
In a new folder ``.pm-restore-*`` in the library's own folder (#318: never
|
||||
the system temp folder, which may be another disk), it:
|
||||
|
||||
1. checks the source (``verify_database_file`` without reading every
|
||||
workflow document: a damaged one is repaired below, not refused);
|
||||
2. copies it with SQLite's backup API (a rollback-journal file and a WAL file
|
||||
alike) and opens the copy as a ``PromptModel``, which migrates an older
|
||||
schema;
|
||||
3. repairs damaged workflow documents (``repair_damaged_documents``),
|
||||
rotates the history epoch and counts what the copy holds;
|
||||
4. makes the copy one self-contained file and verifies it fully (every
|
||||
current table, every document);
|
||||
5. records the copy's SHA-256, so publishing binds exactly these bytes.
|
||||
|
||||
**Publish** (``publish_staged``) re-hashes the stage (``StageChanged`` if it
|
||||
differs), keeps the current database as a uniquely named, standalone safety
|
||||
backup beside it (``database/safety_backup.py``; not verified, because the
|
||||
database being replaced may be the damaged one), copies the stage over the
|
||||
live pages through the backup API (no file is swapped under open handles; an
|
||||
``OperationalError`` there leaves the live data intact, #226), then
|
||||
checkpoints the WAL with TRUNCATE so the copied pages do not leave a
|
||||
database-sized ``-wal`` behind. A reader that keeps a snapshot past the busy
|
||||
wait is logged, not fatal: the restore itself is done.
|
||||
|
||||
Every handle on the stage is closed before its folder is removed (Windows).
|
||||
``sweep_restore_temps`` removes what an interrupted restore left (startup).
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import os
|
||||
import sqlite3
|
||||
import tempfile
|
||||
from contextlib import closing
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional
|
||||
|
||||
from . import workflow_documents as wd
|
||||
from .execution_history_schema import rotate_history_epoch
|
||||
from .safety_backup import take_safety_backup
|
||||
|
||||
logger = logging.getLogger("prompt_manager.database.restore_stage")
|
||||
|
||||
STAGE_PREFIX = ".pm-restore-"
|
||||
UPLOAD_PREFIX = ".pm-restore-upload-"
|
||||
STAGED_NAME = "staged.db"
|
||||
# How long the post-restore TRUNCATE checkpoint waits for readers to leave.
|
||||
CHECKPOINT_BUSY_MS = 30_000
|
||||
_HASH_CHUNK = 1024 * 1024
|
||||
SUMMARY_KEYS = (
|
||||
"prompts",
|
||||
"positives",
|
||||
"pairs",
|
||||
"archived_pairs",
|
||||
"image_links",
|
||||
"runs",
|
||||
"saved_views",
|
||||
)
|
||||
|
||||
|
||||
class StageChanged(ValueError):
|
||||
"""The staged file is no longer the bytes that were reviewed."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class StagedRestore:
|
||||
"""A verified, migrated copy of an upload, ready to publish."""
|
||||
|
||||
path: str
|
||||
digest: str
|
||||
size: int
|
||||
summary: Dict[str, int]
|
||||
schema: Dict[str, bool]
|
||||
documents: Dict[str, int] = field(
|
||||
default_factory=lambda: {"repaired": 0, "cleared": 0, "images": 0}
|
||||
)
|
||||
|
||||
@property
|
||||
def directory(self) -> str:
|
||||
return os.path.dirname(self.path)
|
||||
|
||||
def discard(self) -> None:
|
||||
"""Remove the stage folder (idempotent; its handles are all closed)."""
|
||||
remove_stage_folder(self.directory)
|
||||
|
||||
|
||||
def remove_stage_folder(directory: str) -> None:
|
||||
"""Remove a stage folder and its files, logging what cannot be removed."""
|
||||
if not directory or not os.path.isdir(directory):
|
||||
return
|
||||
for name in os.listdir(directory):
|
||||
try:
|
||||
os.remove(os.path.join(directory, name))
|
||||
except OSError as error:
|
||||
logger.warning(f"Could not remove restore stage file {name}: {error}")
|
||||
try:
|
||||
os.rmdir(directory)
|
||||
except OSError as error:
|
||||
logger.warning(f"Could not remove the restore stage folder: {error}")
|
||||
|
||||
|
||||
def file_digest(path: str) -> str:
|
||||
"""``sha256:<hex>`` of the file's bytes (blocking)."""
|
||||
digest = hashlib.sha256()
|
||||
with open(path, "rb") as handle:
|
||||
for chunk in iter(lambda: handle.read(_HASH_CHUNK), b""):
|
||||
digest.update(chunk)
|
||||
return "sha256:" + digest.hexdigest()
|
||||
|
||||
|
||||
def _readonly(path: str) -> sqlite3.Connection:
|
||||
return sqlite3.connect(Path(path).resolve().as_uri() + "?mode=ro", uri=True)
|
||||
|
||||
|
||||
def _schema(path: str) -> list:
|
||||
"""The file's schema objects, to tell whether staging migrated it."""
|
||||
with closing(_readonly(path)) as conn:
|
||||
return sorted(
|
||||
conn.execute(
|
||||
"SELECT type, name, tbl_name, sql FROM sqlite_master "
|
||||
"WHERE name NOT LIKE 'sqlite_%'"
|
||||
).fetchall(),
|
||||
key=repr,
|
||||
)
|
||||
|
||||
|
||||
def _count(conn: sqlite3.Connection, table: str, where: str = "") -> int:
|
||||
exists = conn.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", (table,)
|
||||
).fetchone()
|
||||
if not exists:
|
||||
return 0
|
||||
# ``table`` and ``where`` are constants of library_summary, never input.
|
||||
return conn.execute(f"SELECT COUNT(*) FROM {table} {where}").fetchone()[0] # nosec
|
||||
|
||||
|
||||
def library_summary(conn: sqlite3.Connection) -> Dict[str, int]:
|
||||
"""What a library holds, for a restore's review (reads only).
|
||||
|
||||
``positives`` are the Library's prompts (``library_counts``'s
|
||||
``positive_prompts``); ``pairs`` the active saved pairs and
|
||||
``archived_pairs`` the archived ones; ``image_links`` the indexed image
|
||||
rows; ``runs`` the recorded executions.
|
||||
"""
|
||||
from .library import _ELIGIBLE
|
||||
|
||||
has_state = conn.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE type='table' AND "
|
||||
"name='prompt_library_state'"
|
||||
).fetchone()
|
||||
has_pairs = conn.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE type='table' AND name='prompt_pairs'"
|
||||
).fetchone()
|
||||
if has_state and has_pairs:
|
||||
positives = conn.execute(
|
||||
"SELECT COUNT(*) FROM prompts p WHERE " + _ELIGIBLE
|
||||
).fetchone()[0]
|
||||
else:
|
||||
positives = _count(conn, "prompts")
|
||||
return {
|
||||
"prompts": _count(conn, "prompts"),
|
||||
"positives": positives,
|
||||
"pairs": _count(conn, "prompt_pairs", "WHERE archived=0"),
|
||||
"archived_pairs": _count(conn, "prompt_pairs", "WHERE archived<>0"),
|
||||
"image_links": _count(conn, "generated_images"),
|
||||
"runs": _count(conn, "execution_runs"),
|
||||
"saved_views": _count(conn, "saved_views"),
|
||||
}
|
||||
|
||||
|
||||
# The newest row of each table a restore would overwrite: with the counts,
|
||||
# a write after the review (a new run, image, prompt or edit) changes it.
|
||||
_FINGERPRINT_SQL = (
|
||||
("prompts", "SELECT MAX(id), MAX(updated_at), MAX(last_used_at) FROM prompts"),
|
||||
("generated_images", "SELECT MAX(id) FROM generated_images"),
|
||||
("prompt_pairs", "SELECT MAX(id), MAX(updated_at) FROM prompt_pairs"),
|
||||
("execution_runs", "SELECT COUNT(*), MAX(rowid) FROM execution_runs"),
|
||||
("saved_views", "SELECT MAX(rowid) FROM saved_views"),
|
||||
)
|
||||
|
||||
|
||||
def live_fingerprint(conn: sqlite3.Connection):
|
||||
"""``(summary, fingerprint)`` of a live library, read in one snapshot.
|
||||
|
||||
The fingerprint changes when rows are added to or removed from the
|
||||
tables a restore replaces, or when a prompt, pair, run or saved view is
|
||||
the newest edited (counts plus the newest row and edit of each table);
|
||||
the apply refuses with ``plan-changed`` then. It does not see every
|
||||
edit (a tag or favorite change on an older row); those are kept in the
|
||||
restore's safety backup. A read transaction the caller already has is
|
||||
used and left open; one opened here is ended here.
|
||||
"""
|
||||
import json
|
||||
|
||||
opened = not conn.in_transaction
|
||||
if opened:
|
||||
conn.execute("BEGIN") # one read snapshot for every count
|
||||
try:
|
||||
summary = library_summary(conn)
|
||||
newest = []
|
||||
for table, sql in _FINGERPRINT_SQL:
|
||||
try:
|
||||
newest.append([table, list(conn.execute(sql).fetchone())])
|
||||
except sqlite3.OperationalError:
|
||||
newest.append([table, None]) # no such table or column
|
||||
finally:
|
||||
if opened:
|
||||
conn.rollback() # it only read
|
||||
text = json.dumps([summary, newest], sort_keys=True, default=str)
|
||||
return summary, "sha256:" + hashlib.sha256(text.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def _copy(src_path: str, staged_path: str) -> None:
|
||||
source = _readonly(src_path)
|
||||
try:
|
||||
with closing(sqlite3.connect(staged_path)) as destination:
|
||||
source.backup(destination)
|
||||
finally:
|
||||
source.close()
|
||||
|
||||
|
||||
def _standalone(path: str) -> None:
|
||||
with closing(sqlite3.connect(path)) as conn:
|
||||
conn.execute("PRAGMA journal_mode=DELETE").fetchone()
|
||||
|
||||
|
||||
def stage_restore(model_cls, db_path: str, src_path: str) -> StagedRestore:
|
||||
"""Stage ``src_path`` beside ``db_path`` (blocking); see the module docstring.
|
||||
|
||||
Raises ValueError with a readable reason when the source or the migrated
|
||||
stage does not verify; the stage folder is then already removed.
|
||||
"""
|
||||
ok, reason = model_cls.verify_database_file(src_path, documents=False)
|
||||
if not ok:
|
||||
raise ValueError(reason)
|
||||
before = _schema(src_path)
|
||||
with closing(_readonly(src_path)) as conn:
|
||||
library_state = bool(
|
||||
conn.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE type='table' AND "
|
||||
"name='prompt_library_state'"
|
||||
).fetchone()
|
||||
)
|
||||
folder = os.path.dirname(os.path.abspath(db_path))
|
||||
directory = tempfile.mkdtemp(prefix=STAGE_PREFIX, dir=folder)
|
||||
staged_path = os.path.join(directory, STAGED_NAME)
|
||||
staged_key = os.path.normcase(os.path.realpath(staged_path))
|
||||
try:
|
||||
_copy(src_path, staged_path)
|
||||
staged = None
|
||||
try:
|
||||
staged = model_cls(staged_path)
|
||||
conn = staged.get_connection()
|
||||
documents = wd.repair_damaged_documents(conn)
|
||||
rotate_history_epoch(conn)
|
||||
summary = library_summary(conn)
|
||||
finally:
|
||||
if staged is not None:
|
||||
staged.close_all()
|
||||
_standalone(staged_path)
|
||||
ok, reason = model_cls.verify_database_file(
|
||||
staged_path, require_library_state=True
|
||||
)
|
||||
if not ok:
|
||||
raise ValueError(reason)
|
||||
result = StagedRestore(
|
||||
path=staged_path,
|
||||
digest=file_digest(staged_path),
|
||||
size=os.path.getsize(staged_path),
|
||||
summary=summary,
|
||||
schema={
|
||||
"migrated": _schema(staged_path) != before,
|
||||
"library_state": library_state,
|
||||
},
|
||||
documents=documents,
|
||||
)
|
||||
except BaseException:
|
||||
remove_stage_folder(directory)
|
||||
raise
|
||||
finally:
|
||||
with model_cls._init_lock:
|
||||
model_cls._initialized_paths.discard(staged_key)
|
||||
return result
|
||||
|
||||
|
||||
def checkpoint_truncate(conn: sqlite3.Connection) -> bool:
|
||||
"""Checkpoint the WAL and truncate it to 0 bytes; False when a reader held it."""
|
||||
previous = conn.execute("PRAGMA busy_timeout").fetchone()[0]
|
||||
try:
|
||||
conn.execute(f"PRAGMA busy_timeout = {int(CHECKPOINT_BUSY_MS)}")
|
||||
busy = conn.execute("PRAGMA wal_checkpoint(TRUNCATE)").fetchone()[0]
|
||||
finally:
|
||||
conn.execute(f"PRAGMA busy_timeout = {int(previous)}")
|
||||
return not busy
|
||||
|
||||
|
||||
def publish_staged(model, staged: StagedRestore) -> str:
|
||||
"""Publish ``staged`` over ``model``'s database; return the safety backup path.
|
||||
|
||||
"" when there was no database to back up. Raises StageChanged before
|
||||
anything is written when the stage's bytes changed since staging, and
|
||||
SafetyBackupFailed when the current database could not be kept. An
|
||||
``OperationalError`` from the page copy propagates with
|
||||
``_prompt_manager_restore_backup_created`` set; the live data is intact.
|
||||
"""
|
||||
if not os.path.isfile(staged.path) or file_digest(staged.path) != staged.digest:
|
||||
raise StageChanged("The staged restore changed since it was reviewed")
|
||||
backup_path = ""
|
||||
if os.path.exists(model.db_path):
|
||||
backup_path = take_safety_backup(model, verify=False)
|
||||
logger.info(f"Current database backed up to: {os.path.basename(backup_path)}")
|
||||
try:
|
||||
model._replace_contents_from(staged.path)
|
||||
except sqlite3.OperationalError as exc:
|
||||
exc._prompt_manager_restore_backup_created = bool(backup_path)
|
||||
raise
|
||||
with model._init_lock:
|
||||
model._initialized_paths.add(model._schema_key())
|
||||
_truncate_wal(model)
|
||||
return backup_path
|
||||
|
||||
|
||||
def _truncate_wal(model) -> Optional[bool]:
|
||||
"""TRUNCATE-checkpoint the live WAL after a publish (never raises)."""
|
||||
try:
|
||||
conn = model.get_connection()
|
||||
with conn.using():
|
||||
truncated = checkpoint_truncate(conn)
|
||||
except sqlite3.Error as error:
|
||||
logger.warning(f"The database log could not be truncated: {error}")
|
||||
return None
|
||||
if not truncated:
|
||||
logger.warning(
|
||||
"The database log could not be truncated after the restore: a "
|
||||
"reader kept it open. It shrinks at the next checkpoint with no "
|
||||
"reader."
|
||||
)
|
||||
return truncated
|
||||
|
||||
|
||||
def sweep_restore_temps(db_path: str) -> list:
|
||||
"""Remove the stages and uploads an interrupted restore left (startup)."""
|
||||
folder = os.path.dirname(os.path.abspath(db_path))
|
||||
removed = []
|
||||
try:
|
||||
names = os.listdir(folder)
|
||||
except OSError:
|
||||
return removed
|
||||
for name in names:
|
||||
path = os.path.join(folder, name)
|
||||
if name.startswith(UPLOAD_PREFIX) and os.path.isfile(path):
|
||||
try:
|
||||
os.remove(path)
|
||||
removed.append(name)
|
||||
except OSError:
|
||||
continue
|
||||
elif name.startswith(STAGE_PREFIX) and os.path.isdir(path):
|
||||
remove_stage_folder(path)
|
||||
if not os.path.exists(path):
|
||||
removed.append(name)
|
||||
return removed
|
||||
@@ -380,15 +380,21 @@ def _structure_error(conn: sqlite3.Connection) -> Optional[str]:
|
||||
|
||||
|
||||
def _evidence_error(conn: sqlite3.Connection) -> Optional[str]:
|
||||
"""Stream every evidence document once; one row is held at a time."""
|
||||
for evidence_hash, text in conn.execute(
|
||||
"""Stream every evidence document once; one row is held at a time.
|
||||
|
||||
Rows are read with ``fetchone``: iterating a managed connection's cursor
|
||||
would prefetch ``ITER_CHUNK`` (128) rows (``database/managed_connection.py``).
|
||||
"""
|
||||
items = conn.execute(
|
||||
f"SELECT evidence_hash, evidence FROM {ITEMS_TABLE}" # nosec B608
|
||||
):
|
||||
)
|
||||
while (row := items.fetchone()) is not None:
|
||||
evidence_hash, text = row[0], row[1]
|
||||
if not _is_canonical(text) or evidence_digest(text) != evidence_hash:
|
||||
return f"Review queue table {ITEMS_TABLE} contains invalid evidence"
|
||||
for (text,) in conn.execute(
|
||||
f"SELECT evidence FROM {DECISIONS_TABLE}" # nosec B608
|
||||
):
|
||||
decisions = conn.execute(f"SELECT evidence FROM {DECISIONS_TABLE}") # nosec B608
|
||||
while (row := decisions.fetchone()) is not None:
|
||||
text = row[0]
|
||||
if not _is_canonical(text):
|
||||
return f"Review queue table {DECISIONS_TABLE} contains invalid evidence"
|
||||
return None
|
||||
|
||||
@@ -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")
|
||||
@@ -475,3 +475,178 @@ def release_documents(conn: sqlite3.Connection, ids: Iterable[Optional[int]]) ->
|
||||
def sweep_documents(conn: sqlite3.Connection) -> int:
|
||||
"""Delete every document no image row references; returns the count."""
|
||||
return conn.execute(SWEEP_SQL).rowcount
|
||||
|
||||
|
||||
# ── repair (a restore's stage) ────────────────────────────────────────
|
||||
|
||||
INLINE_TEXTS_SQL = (
|
||||
"SELECT workflow_data, prompt_metadata FROM generated_images "
|
||||
"WHERE workflow_data IS NOT NULL OR prompt_metadata IS NOT NULL"
|
||||
)
|
||||
REKEY_SQL = "UPDATE workflow_documents SET sha256=?, byte_length=? WHERE id=?"
|
||||
|
||||
|
||||
def _rows(conn: sqlite3.Connection, sql: str, params=()):
|
||||
"""Each row of ``sql``, read one at a time.
|
||||
|
||||
Iterating a managed connection's cursor prefetches ``ITER_CHUNK`` (128)
|
||||
rows, and a document or inline text can be megabytes.
|
||||
"""
|
||||
cursor = conn.execute(sql, params)
|
||||
while (row := cursor.fetchone()) is not None:
|
||||
yield row
|
||||
|
||||
|
||||
def _inflated(data) -> Optional[bytes]:
|
||||
"""The complete text bytes ``data`` inflates to, or None.
|
||||
|
||||
Complete means the zlib stream ends (its Adler-32 checked) within
|
||||
MAX_DOCUMENT_BYTES with nothing after it: the body is intact even when
|
||||
the row's recorded hash or length is not.
|
||||
"""
|
||||
try:
|
||||
inflater = zlib.decompressobj()
|
||||
raw = inflater.decompress(bytes(data), MAX_DOCUMENT_BYTES + 1)
|
||||
if not inflater.eof or inflater.unconsumed_tail or inflater.unused_data:
|
||||
return None
|
||||
if len(raw) > MAX_DOCUMENT_BYTES:
|
||||
return None
|
||||
raw.decode("utf-8")
|
||||
return raw
|
||||
except (zlib.error, TypeError, UnicodeDecodeError):
|
||||
return None
|
||||
|
||||
|
||||
def _one_column_off(row, raw: bytes) -> bool:
|
||||
"""Whether only the recorded hash or only the recorded length is wrong.
|
||||
|
||||
The two columns check the body independently. When the body inflates
|
||||
completely and one of them still agrees with it, that column vouches for
|
||||
the body and the other one is the damaged value. When both disagree, the
|
||||
body itself is the suspect (another valid stream), so it is not trusted.
|
||||
"""
|
||||
hash_ok = hashlib.sha256(raw).digest() == bytes(row[0])
|
||||
length_ok = row[1] == len(raw)
|
||||
return hash_ok != length_ok
|
||||
|
||||
|
||||
def _inline_copies(conn: sqlite3.Connection, wanted: Dict[bytes, int]):
|
||||
"""``{document id: text}`` for each wanted hash an inline row still holds."""
|
||||
found: Dict[int, str] = {}
|
||||
for row in _rows(conn, INLINE_TEXTS_SQL):
|
||||
for text in row:
|
||||
if not isinstance(text, str) or not storable(text):
|
||||
continue
|
||||
document_id = wanted.get(hashlib.sha256(text.encode("utf-8")).digest())
|
||||
if document_id is not None and document_id not in found:
|
||||
found[document_id] = text
|
||||
if len(found) == len(wanted):
|
||||
break
|
||||
return found
|
||||
|
||||
|
||||
def _rekey(conn: sqlite3.Connection, document_id: int, raw: bytes) -> Optional[int]:
|
||||
"""Give an intact body its true key, or merge it into the document that has it.
|
||||
|
||||
Returns the holder's id when its own body was damaged and has just been
|
||||
rewritten from ``raw`` (it is repaired, never to be cleared), else None.
|
||||
"""
|
||||
digest = hashlib.sha256(raw).digest()
|
||||
holder = conn.execute(LOOKUP_SQL, (digest,)).fetchone()
|
||||
if holder is None or holder[0] == document_id:
|
||||
conn.execute(REKEY_SQL, (digest, len(raw), document_id))
|
||||
return None
|
||||
for column in DOCUMENT_ID_COLUMNS:
|
||||
# ``column`` is one of DOCUMENT_ID_COLUMNS, never caller data.
|
||||
conn.execute(
|
||||
f"UPDATE generated_images SET {column}=? WHERE {column}=?", # nosec B608
|
||||
(holder[0], document_id),
|
||||
)
|
||||
conn.execute("DELETE FROM workflow_documents WHERE id=?", (document_id,))
|
||||
if _decode(holder[1:], holder[0], quiet=True) is None:
|
||||
conn.execute(
|
||||
REPAIR_SQL, (len(raw), zlib.compress(raw, COMPRESSION_LEVEL), holder[0])
|
||||
)
|
||||
return holder[0]
|
||||
return None
|
||||
|
||||
|
||||
def repair_damaged_documents(conn: sqlite3.Connection) -> Dict[str, int]:
|
||||
"""Repair or clear every document failing its checks; commit (writes).
|
||||
|
||||
For a restore's stage only, never the live library. In order:
|
||||
|
||||
1. A document whose body inflates completely and whose recorded hash or
|
||||
length (not both) disagrees is re-keyed from its text, or merged
|
||||
into the document that already holds that text (``_one_column_off``).
|
||||
2. A damaged body whose exact text an inline row still holds (the same
|
||||
SHA-256, its key) is stored again from that text.
|
||||
3. Any other damaged document is unreadable for good: its images'
|
||||
references are cleared (they already showed no workflow) and the
|
||||
document is removed, so the stage verifies.
|
||||
|
||||
Returns ``{repaired, cleared, images}``: documents repaired (1 and 2),
|
||||
documents cleared and the distinct image rows whose reference was
|
||||
cleared. Rows are read one at a time.
|
||||
"""
|
||||
damaged: Dict[bytes, int] = {}
|
||||
intact: Dict[int, bytes] = {}
|
||||
for row in _rows(conn, ALL_DOCUMENTS_SQL):
|
||||
if _decode(row[:3], row[3], quiet=True) is None:
|
||||
raw = _inflated(row[2])
|
||||
if raw is not None and _one_column_off(row, raw):
|
||||
intact[row[3]] = raw
|
||||
else:
|
||||
damaged[bytes(row[0])] = row[3]
|
||||
result = {"repaired": 0, "cleared": 0, "images": 0}
|
||||
if not damaged and not intact:
|
||||
return result
|
||||
with conn:
|
||||
healed = []
|
||||
for document_id, raw in intact.items():
|
||||
holder = _rekey(conn, document_id, raw)
|
||||
if holder is not None and holder in damaged.values():
|
||||
# Rewritten from the merged body: repaired, not to be cleared
|
||||
# (it now also holds the merged images' references).
|
||||
damaged = {k: v for k, v in damaged.items() if v != holder}
|
||||
healed.append(holder)
|
||||
copies = _inline_copies(conn, damaged) if damaged else {}
|
||||
for document_id, text in copies.items():
|
||||
raw = text.encode("utf-8")
|
||||
data = zlib.compress(raw, COMPRESSION_LEVEL)
|
||||
conn.execute(REPAIR_SQL, (len(raw), data, document_id))
|
||||
lost = sorted(set(damaged.values()) - set(copies))
|
||||
images: Set[int] = set()
|
||||
for chunk in _chunks(lost):
|
||||
marks = ",".join("?" * len(chunk))
|
||||
# ``marks`` is only "?" placeholders; the ids are bound.
|
||||
for (image_id,) in _rows(
|
||||
conn,
|
||||
"SELECT id FROM generated_images WHERE " # nosec B608
|
||||
f"workflow_document_id IN ({marks}) OR prompt_document_id IN ({marks})",
|
||||
chunk + chunk,
|
||||
):
|
||||
images.add(image_id)
|
||||
for column in DOCUMENT_ID_COLUMNS:
|
||||
conn.execute(
|
||||
f"UPDATE generated_images SET {column}=NULL " # nosec B608
|
||||
f"WHERE {column} IN ({marks})",
|
||||
chunk,
|
||||
)
|
||||
conn.execute(
|
||||
f"DELETE FROM workflow_documents WHERE id IN ({marks})", # nosec B608
|
||||
chunk,
|
||||
)
|
||||
for document_id in list(intact) + healed + list(copies):
|
||||
logger.warning(f"Workflow document {document_id} was damaged; repaired")
|
||||
if lost:
|
||||
logger.warning(
|
||||
f"{len(lost)} damaged workflow document(s) could not be repaired; "
|
||||
f"{len(images)} image(s) show no workflow"
|
||||
)
|
||||
result.update(
|
||||
repaired=len(intact) + len(healed) + len(copies),
|
||||
cleared=len(lost),
|
||||
images=len(images),
|
||||
)
|
||||
return result
|
||||
|
||||
Reference in New Issue
Block a user