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

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

Browse Source 
- 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
2026-10-25 08:11:27 -07:00
parent a8ce107f1f
commit 5ce3349d64
9 changed files with 947 additions and 169 deletions
Download Patch File Download Diff File Expand all files Collapse all files
database/workflow_documents.py
+102 -19
View File
@@ -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