diff --git a/database/models.py b/database/models.py index 94773c4e..f36fdbc9 100644 --- a/database/models.py +++ b/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 diff --git a/database/workflow_documents.py b/database/workflow_documents.py index 7272b682..d53737b2 100644 --- a/database/workflow_documents.py +++ b/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 diff --git a/database/workflow_storage.py b/database/workflow_storage.py index 3e1ba3de..c9cbe48a 100644 --- a/database/workflow_storage.py +++ b/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_" + ), + }, {"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." ) diff --git a/tests/test_workflow_documents.py b/tests/test_workflow_documents.py index bd60a5c0..40cfdbae 100644 --- a/tests/test_workflow_documents.py +++ b/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( diff --git a/tests/test_workflow_migration.py b/tests/test_workflow_migration.py index 51f8d922..3f98ef40 100644 --- a/tests/test_workflow_migration.py +++ b/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)