feat(database): keep the migration safe row by row and beside other writers

Final review fixes M1, M2, M3 and M5 for the workflow storage dedupe:

- M1: a batch also ends after 0.5 s, and the loop leaves the write lock
  free for 0.15 s after each committed batch (longer than SQLite's longest
  busy poll), so the image monitor and run counting on another connection
  get in between. RED: a threaded writer with a 2 s busy timeout hit
  "database is locked" (4 of 5 links landed).
- M2: each row runs in its own savepoint. A row that cannot be rewritten
  (undecodable text, a text that would not come back) is rolled back alone,
  stays inline and is listed in errors; the rest of the batch goes on.
- M3: a reused document is decompressed and checked before its id is handed
  out, and repaired from the text (whose hash is its key) when damaged;
  verify_database_file reads every document and refuses a damaged one.
- M5: the two legacy generated_images rebuilds refuse a file that already
  has the document id columns instead of dropping them.
- The preview names the pre-migration backup and says it is always taken.

database/models.py

@@ -25,9 +25,12 @@ from .review_schema import (
 )
 from .saved_views import migrate_saved_views, saved_views_schema_error
 from .workflow_documents import (
+    DOCUMENT_ID_COLUMNS,
     migrate_workflow_documents,
     workflow_documents_schema_error,
 )
+
+
 from .execution_history_schema import (
     history_schema_error,
     migrate_execution_history,
@@ -54,6 +57,19 @@ MAX_PAIR_PROVENANCE_LENGTH = 4_000_000
 MAX_PAIR_PROVENANCE_DEPTH = 64
 
 
+def _refuse_rebuild_with_documents(columns) -> None:
+    """The legacy generated_images rebuilds copy a fixed column list.
+
+    A file that needs one (``prompt_id`` TEXT, no ``file_path``) but already
+    has the workflow document ids would lose them, so it is refused.
+    """
+    if any(column in columns for column in DOCUMENT_ID_COLUMNS):
+        raise sqlite3.DatabaseError(
+            "This database needs an old image table rebuild but already has "
+            "workflow document references; the rebuild would drop them"
+        )
+
+
 def pair_provenance_error(encoded: object) -> Optional[str]:
     """Validate the same bounded JSON object for live writes and restores."""
     if not isinstance(encoded, str):
@@ -717,6 +733,7 @@ class PromptModel:
             columns = {column[1]: column[2] for column in cursor.fetchall()}
             if columns.get("prompt_id") != "TEXT":
                 return
+            _refuse_rebuild_with_documents(columns)
 
             self.logger.info("Migrating foreign key types: prompt_id TEXT -> INTEGER")
             with self._table_rebuild(conn):
@@ -882,6 +899,7 @@ class PromptModel:
             }
             if "file_path" in columns:
                 return
+            _refuse_rebuild_with_documents(columns)
 
             self.logger.info("Migrating database: image uniqueness by full path")
             with self._table_rebuild(conn):
@@ -1124,7 +1142,7 @@ class PromptModel:
             if review_error:
                 return False, review_error
             documents_error = workflow_documents_schema_error(
-                conn, allow_missing=not require_library_state
+                conn, allow_missing=not require_library_state, documents=True
             )
             if documents_error:
                 return False, documents_error

database/workflow_documents.py

@@ -74,7 +74,11 @@ _UNREFERENCED = (
     "AND NOT EXISTS (SELECT 1 FROM generated_images g "
     "WHERE g.prompt_document_id=workflow_documents.id)"
 )
-LOOKUP_SQL = "SELECT id, byte_length FROM workflow_documents WHERE sha256=?"
+LOOKUP_SQL = (
+    "SELECT id, sha256, byte_length, data FROM workflow_documents WHERE sha256=?"
+)
+REPAIR_SQL = "UPDATE workflow_documents SET byte_length=?, data=? WHERE id=?"
+ALL_DOCUMENTS_SQL = "SELECT sha256, byte_length, data, id FROM workflow_documents"
 READ_SQL = "SELECT sha256, byte_length, data FROM workflow_documents WHERE id=?"
 INSERT_SQL = (
     "INSERT INTO workflow_documents (sha256, byte_length, data) VALUES (?, ?, ?)"
@@ -168,8 +172,16 @@ def _value_error(conn: sqlite3.Connection):
     return None
 
 
+def _damaged_document(conn: sqlite3.Connection) -> bool:
+    """Whether a document fails its length or hash check (one row at a time)."""
+    for row in conn.execute(ALL_DOCUMENTS_SQL):
+        if _decode(row[:3], row[3], quiet=True) is None:
+            return True
+    return False
+
+
 def workflow_documents_schema_error(
-    conn: sqlite3.Connection, *, allow_missing: bool = True
+    conn: sqlite3.Connection, *, allow_missing: bool = True, documents: bool = False
 ) -> Optional[str]:
     """Reject an unsupported document schema or invalid stored values (reads only).
 
@@ -199,7 +211,11 @@ def workflow_documents_schema_error(
     # The exact definition fixes the columns, their types and constraints.
     if _compact(tables[0]) != _compact(_DDL):
         return "Workflow documents have an invalid table definition"
-    return _reference_error(conn, columns) or _value_error(conn)
+    error = _reference_error(conn, columns) or _value_error(conn)
+    if error is None and documents and _damaged_document(conn):
+        # Restore and backup verification only: every document is read.
+        return "Workflow documents contain a damaged document"
+    return error
 
 
 def migrate_workflow_documents(conn: sqlite3.Connection) -> None:
@@ -250,15 +266,21 @@ def storable(text: Any) -> bool:
     return len(text.encode("utf-8")) <= MAX_DOCUMENT_BYTES
 
 
-def store_text(conn: sqlite3.Connection, text: Optional[str]) -> Optional[int]:
+def store_text(
+    conn: sqlite3.Connection, text: Optional[str], verified: Optional[set] = None
+) -> Optional[int]:
     """The id of the document holding exactly ``text``; None for None.
 
-    Reuses the document with the same hash and length, or inserts a new one
-    whose compressed bytes were checked to decompress to ``text``'s bytes.
-    Opens a write transaction first when none is open, so a concurrent
-    ``release_documents`` cannot remove a reused document before the caller
-    references it. Raises DocumentMismatch, TypeError for a non-text value
-    and ValueError for a text over MAX_DOCUMENT_BYTES.
+    Reuses the document with the same hash, or inserts a new one whose
+    compressed bytes were checked to decompress to ``text``'s bytes. A reused
+    document is decompressed and checked first; a damaged one (a restored
+    file's corruption under a valid hash) is repaired from ``text``, whose
+    hash is its key, so every row that references it reads right again.
+    ``verified``, when given, holds the ids already checked (a migration
+    batch), which are then not read again. Opens a write transaction first
+    when none is open, so a concurrent ``release_documents`` cannot remove a
+    reused document before the caller references it. Raises TypeError for a
+    non-text value and ValueError for a text over MAX_DOCUMENT_BYTES.
     """
     if text is None:
         return None
@@ -271,14 +293,21 @@ def store_text(conn: sqlite3.Connection, text: Optional[str]) -> Optional[int]:
     if not conn.in_transaction:
         conn.execute("BEGIN IMMEDIATE")
     row = conn.execute(LOOKUP_SQL, (digest,)).fetchone()
-    if row is not None:
-        if row[1] != len(raw):
-            raise DocumentMismatch("A stored document has another length")
+    if row is not None and verified is not None and row[0] in verified:
         return row[0]
     data = zlib.compress(raw, COMPRESSION_LEVEL)
     if zlib.decompress(data) != raw:  # pragma: no cover - zlib is exact
         raise DocumentMismatch("Compression did not round-trip")
-    return conn.execute(INSERT_SQL, (digest, len(raw), data)).lastrowid
+    if row is None:
+        document_id = conn.execute(INSERT_SQL, (digest, len(raw), data)).lastrowid
+    else:
+        document_id = row[0]
+        if _decode(row[1:], document_id, quiet=True) != text:
+            logger.warning(f"Workflow document {document_id} was damaged; repaired")
+            conn.execute(REPAIR_SQL, (len(raw), data, document_id))
+    if verified is not None:
+        verified.add(document_id)
+    return document_id
 
 
 def stored_values(
@@ -304,7 +333,8 @@ def stored_values(
 # ── read ──────────────────────────────────────────────────────────────
 
 
-def _decode(row, document_id) -> Optional[str]:
+def _decode(row, document_id, quiet=False) -> Optional[str]:
+    """``(sha256, byte_length, data)`` as checked text, or None (logged)."""
     if row is None:
         logger.warning(f"Workflow document {document_id} is missing")
         return None
@@ -320,7 +350,8 @@ def _decode(row, document_id) -> Optional[str]:
             raise ValueError("hash differs")
         return raw.decode("utf-8")
     except (zlib.error, ValueError, TypeError) as error:
-        logger.warning(f"Workflow document {document_id} is unreadable: {error}")
+        if not quiet:
+            logger.warning(f"Workflow document {document_id} is unreadable: {error}")
         return None
 
 

database/workflow_storage.py

@@ -12,7 +12,12 @@ prompt graph text inline) to the new form (a ``workflow_documents`` id,
   one accessor before its row changes; a text that would not come back
   identical fails the batch (``DocumentMismatch``) and rolls it back whole;
 - a value that is not text (or is over ``MAX_DOCUMENT_BYTES``) stays inline
-  and is counted as protected.
+  and is counted as protected;
+- each row runs in its own savepoint: one that cannot be rewritten (text
+  SQLite cannot decode, a text that would not come back) stays inline and
+  is listed, and the rest of the batch goes on;
+- a batch ends after ``BATCH_SECONDS``, and the write lock is then left
+  free for ``BATCH_PAUSE``, so other writers are never starved.
 
 A crash or a stop between batches leaves committed batches in the new form
 and the rest in the old form; both read the same, and a rerun continues
@@ -24,8 +29,10 @@ documents, rebuilds the file and checkpoints the WAL with TRUNCATE, so the
 freed pages go back to the disk.
 """
 
+import logging
 import os
 import sqlite3
+import time
 from contextlib import closing
 
 from .workflow_documents import (
@@ -41,6 +48,18 @@ from .workflow_documents import (
 OPERATION = "workflow_storage"
 BATCH_ROWS = 200
 BATCH_BYTES = 16 * 2**20  # characters of text rewritten per transaction
+# A batch holds the write lock: it ends after this long, and the loop then
+# leaves the lock free for BATCH_PAUSE, longer than SQLite's longest busy
+# poll (100 ms), so the image monitor and the run counting get in between.
+BATCH_SECONDS = 0.5
+BATCH_PAUSE = 0.15
+# What one row's rewrite may raise without failing its batch: the row stays
+# inline and is reported (text SQLite cannot decode, a text that would not
+# come back identical).
+ROW_ERRORS = (sqlite3.Error, DocumentMismatch, UnicodeError, ValueError, TypeError)
+
+logger = logging.getLogger("prompt_manager.database.workflow_storage")
+_clock = time.monotonic
 
 # typeof() reads a column's type from the record header: no overflow page.
 NEXT_IDS_SQL = (
@@ -76,11 +95,12 @@ QUERY_PLAN_STATEMENTS = {
 }
 
 
-def _rewrite_row(conn, image_id, row, verified):
+def _rewrite_row(conn, image_id, row, verified, checked):
     """Store the row's inline texts as documents; returns (changed, protected, size).
 
     ``verified`` maps document ids already read back in this batch to their
-    text, so a document shared by many rows is decompressed once.
+    text, and ``checked`` holds the reused ids ``store_text`` already
+    checked, so a document shared by many rows is decompressed once.
     """
     changes, replaced = {}, []
     protected = size = 0
@@ -91,7 +111,7 @@ def _rewrite_row(conn, image_id, row, verified):
         if not storable(value):
             protected += 1
             continue
-        document_id = store_text(conn, value)
+        document_id = store_text(conn, value, checked)
         if document_id not in verified:
             verified[document_id] = read_stored_text(conn, None, document_id)
         if verified[document_id] != value:
@@ -110,28 +130,60 @@ def _rewrite_row(conn, image_id, row, verified):
     return bool(changes), protected, size
 
 
-def migrate_batch(conn, after_id, max_rows=BATCH_ROWS, max_bytes=BATCH_BYTES):
+def migrate_batch(
+    conn,
+    after_id,
+    max_rows=BATCH_ROWS,
+    max_bytes=BATCH_BYTES,
+    max_seconds=BATCH_SECONDS,
+):
     """Rewrite the next batch after ``after_id`` in one transaction.
 
-    Returns ``{last_id, rows, protected, done}``: the last row looked at,
-    rows rewritten, inline values kept (not text) and whether there was
-    nothing left to look at.
+    The batch ends after ``max_rows`` rows, about ``max_bytes`` of text or
+    ``max_seconds``, whichever comes first. Each row runs in its own
+    savepoint: a row that cannot be rewritten (``ROW_ERRORS``) is rolled
+    back alone, stays inline and is listed in ``errors``; anything else rolls
+    the whole batch back and propagates. Returns ``{last_id, rows,
+    protected, errors, done}``: the last row looked at, rows rewritten,
+    inline values kept (not text), the ids kept by an error and whether
+    there was nothing left to look at.
     """
     if conn.in_transaction:
         raise RuntimeError("The migration needs its own transaction")
+    started = _clock()
     conn.execute("BEGIN IMMEDIATE")
     try:
         ids = [row[0] for row in conn.execute(NEXT_IDS_SQL, (after_id, max_rows))]
-        result = {"last_id": after_id, "rows": 0, "protected": 0, "done": not ids}
-        verified, size = {}, 0
+        result = {
+            "last_id": after_id,
+            "rows": 0,
+            "protected": 0,
+            "errors": [],
+            "done": not ids,
+        }
+        verified, checked, size = {}, set(), 0
         for image_id in ids:
-            row = conn.execute(ROW_SQL, (image_id,)).fetchone()
-            changed, protected, row_size = _rewrite_row(conn, image_id, row, verified)
+            conn.execute("SAVEPOINT workflow_row")
+            try:
+                row = conn.execute(ROW_SQL, (image_id,)).fetchone()
+                changed, protected, row_size = _rewrite_row(
+                    conn, image_id, row, verified, checked
+                )
+            except ROW_ERRORS as error:
+                conn.execute("ROLLBACK TO workflow_row")
+                conn.execute("RELEASE workflow_row")
+                verified.clear()  # the row's documents were rolled back too
+                checked.clear()
+                logger.warning(f"Image record {image_id} kept inline: {error}")
+                result["errors"].append(image_id)
+                changed, protected, row_size = False, 0, 0
+            else:
+                conn.execute("RELEASE workflow_row")
             result["rows"] += int(changed)
             result["protected"] += protected
             result["last_id"] = image_id
             size += row_size
-            if size >= max_bytes:
+            if size >= max_bytes or _clock() - started >= max_seconds:
                 break
         conn.commit()
         return result
@@ -163,11 +215,14 @@ class WorkflowStorageOperations:
 
         ``on_batch(progress)`` runs after each committed batch with the
         running totals ``{rows, protected, batches, last_id}``;
-        ``should_stop()`` is asked before each batch. Returns ``{rows,
-        protected, batches, stopped, documents_removed}``. Unreferenced
+        ``should_stop()`` is asked before each batch. Each committed batch is
+        followed by ``BATCH_PAUSE`` with the write lock free. Returns
+        ``{rows, protected, batches, stopped, errors, documents_removed}``
+        (``errors``: the ids of rows kept inline by an error). Unreferenced
         documents are swept once every row was looked at.
         """
         totals = {"rows": 0, "protected": 0, "batches": 0, "last_id": 0}
+        errors = []
         stopped = False
         while True:
             if should_stop is not None and should_stop():
@@ -176,6 +231,8 @@ class WorkflowStorageOperations:
             batch = self.migrate_workflow_batch(totals["last_id"], max_rows, max_bytes)
             if batch["done"]:
                 break
+            time.sleep(BATCH_PAUSE)  # leave the write lock to other writers
+            errors += batch["errors"]
             totals["rows"] += batch["rows"]
             totals["protected"] += batch["protected"]
             totals["batches"] += 1
@@ -188,6 +245,7 @@ class WorkflowStorageOperations:
             "protected": totals["protected"],
             "batches": totals["batches"],
             "stopped": stopped,
+            "errors": errors,
             "documents_removed": removed,
         }
 
@@ -296,7 +354,14 @@ class WorkflowStorageOperations:
                     "detail": _plural(counts["documents"], "workflow document"),
                 }
             )
+        name = os.path.basename(self.model.db_path)
         details += [
+            {
+                "label": "Backup",
+                "detail": BACKUP_NOTE.format(
+                    name=f"{name}.pre-workflow-storage_<date and time>"
+                ),
+            },
             {"label": "Reclaim space", "detail": RECLAIM_NOTE},
             {"label": "Older versions", "detail": OLDER_VERSIONS_NOTE},
         ]
@@ -327,7 +392,11 @@ RECLAIM_NOTE = (
     "When every record is rewritten, Reclaim space runs alone to give the freed "
     "space back to the disk (VACUUM). The database is locked while it runs."
 )
+BACKUP_NOTE = (
+    "Always taken first, whatever the safety backup setting: {name} beside "
+    "the database, verified, and kept until you remove it (it is never pruned)."
+)
 OLDER_VERSIONS_NOTE = (
     "An older Prompt Manager version shows rewritten images without their "
-    "workflow. Restoring the safety backup brings back the old form."
+    "workflow. Restoring the pre-migration backup brings back the old form."
 )

tests/test_workflow_documents.py

@@ -517,9 +517,9 @@ class TestEdgeErrors(DocumentsTestCase):
                 "UPDATE workflow_documents SET byte_length=5 WHERE id=?",
                 (document_id,),
             )
-        with self.assertRaises(wd.DocumentMismatch):
-            wd.store_text(conn, "abc")
-        conn.rollback()
+        with conn, self.assertLogs(wd.logger, "WARNING"):  # repaired from the text
+            self.assertEqual(wd.store_text(conn, "abc"), document_id)
+        self.assertEqual(wd.read_stored_text(conn, None, document_id), "abc")
         for length in (-1, 2, 4):
             with conn:
                 conn.execute(

tests/test_workflow_migration.py

@@ -135,11 +135,11 @@ class TestMigration(MigrationTestCase):
         real = wd.store_text
         calls = []
 
-        def crashing(conn, text):
+        def crashing(conn, text, *args):
             calls.append(text)
             if len(calls) == 23:
                 raise RuntimeError("power cut")
-            return real(conn, text)
+            return real(conn, text, *args)
 
         with patch.object(ws, "store_text", side_effect=crashing):
             with self.assertRaises(RuntimeError):
@@ -163,13 +163,14 @@ class TestMigration(MigrationTestCase):
         self.assertEqual(self.inline_ids(), set(expected))
         self.assertEqual(self.document_count(), 0)
 
-    def test_a_document_that_reads_back_differently_fails_the_batch(self):
+    def test_a_document_that_reads_back_differently_keeps_its_row(self):
         _prompt_id, expected = self.library(count=6)
         with patch.object(ws, "read_stored_text", return_value="something else"):
-            with self.assertRaises(wd.DocumentMismatch):
-                self.migrate(max_rows=50)
+            result = self.migrate(max_rows=50)
+        self.assertEqual(result["errors"], sorted(expected))
         self.assertEqual(self.inline_ids(), set(expected))
         self.assertEqual(self.texts(), expected)
+        self.assertEqual(self.document_count(), 0)
 
     def test_stop_between_batches_then_continue(self):
         _prompt_id, expected = self.library(count=30)