diff --git a/README.md b/README.md index e259362..b60585e 100644 --- a/README.md +++ b/README.md @@ -970,6 +970,11 @@ MIT License - see LICENSE file for details. ## Changelog +### Unreleased (Workflows stored once) + +- **🗜️ Smaller database**: each distinct workflow and prompt graph is stored once, compressed, instead of again for every image. Images recorded by this version use it from the first start; the maintenance job "Store each workflow once" rewrites the existing records and then reclaims the space. A 20,000-image test library went from 1.4 GB to 60 MB. The job always takes a verified backup first: the first time, `prompts.db.pre-workflow-storage_` beside the database, kept until you delete it; on later runs, a normal pooled safety backup. No image or video file is opened or changed. +- **⚠️ Downgrading**: images recorded by this version (and, after the job, all images) are stored in the new form. An older Prompt Manager version still opens and even restores such a database, but shows those images without workflows or prompt graphs until you update again. To go back to an older version, stop ComfyUI, delete `prompts.db-wal` and `prompts.db-shm`, and copy the **earliest** `prompts.db.pre-workflow-storage_` over `prompts.db`, or, if the job never ran, a backup taken before this update (records added since that backup are lost). Otherwise stay on this version. + ### v3.3.1 (Proxy Routing Fix) - **🌐 Proxy support**: the dashboard, gallery, tag page and metadata viewer now work when ComfyUI is reached through a proxy mount such as SwarmUI's `/ComfyBackendDirect/`. Scripts, images, API calls, progress streams and gallery folder filters stay under the mount, and the canvas node's buttons use ComfyUI's `fileURL()`. Direct ComfyUI URLs are unchanged. Thanks to @ghosttigerllc-bit ([#164](https://github.com/ComfyAssets/ComfyUI_PromptManager/pull/164)) diff --git a/TESTING.md b/TESTING.md index 8267f6a..213d3eb 100644 --- a/TESTING.md +++ b/TESTING.md @@ -2708,6 +2708,73 @@ runs: every per-prompt protection and fingerprint read is an index search, the duplicate groups and the link rows are one scan each, and a delete batch (400 ids) searches by rowid with its cascades by index. +**Workflow storage** (`database/workflow_documents.py`, `database/workflow_storage.py`, +`py/api/workflow_storage.py`; `tests/test_workflow_documents.py`, +`tests/test_workflow_migration.py`, `tests/test_api_workflow_storage.py`). Each distinct +workflow and prompt graph text is one `workflow_documents` row: `sha256` of its exact +UTF-8 bytes (unique), `byte_length` and `data` = `zlib.compress` of those bytes. Image rows +reference it through `workflow_document_id` / `prompt_document_id` (foreign keys, one +index each). The text is stored and hashed exactly as written, never parsed and +re-serialized. Two forms are readable forever: old (text inline, no id) and new (inline +NULL, id set); inline text wins when both are present. Every reader goes through +`read_stored_text` / `resolve_image_row` (image lists, image by id, prompt info for the +inspector and metadata page, diagnostics, pair ingestion's comparison), which also drops +the id keys so payloads keep their shape; a document that fails its length or hash check +reads as unreadable (None), never as other text. Round trips are proven byte for byte +over edge texts (spacing, key order, duplicate keys, number spellings, `NaN`/`Infinity` +tokens, NUL, BOM, astral and combining characters, 500-deep nesting, a 1.2 MB+ +workflow, `""` and NULL) and seeded random text and workflows +(`TestRoundTrip`). New writes (`link_image_to_prompt`, pair ingestion for scan, monitor +and Needs review) store the new form. No trigger is installed (restore validators refuse +them): a document goes in the transaction that drops its last reference (image and +prompt deletes, missing-link cleanup, duplicate merges, duplicate-file link merges, +ingestion updates: `TestGarbage`, `TestMoreGarbagePaths`), and startup, the migration's +end and Reclaim space sweep anything a path missed. `verify_database_file` checks the +table definition, the id columns, their foreign keys and indexes, the stored values and +dangling ids (`TestSchemaErrors`); a pre-dedupe file is accepted and migrated on open +and on restore. The `workflow_storage` maintenance plan counts the old-form records +(one pass over record headers) and its apply starts "Store each workflow once" +(`records`, cancellable): first, always and whatever `maintenance.safety_backup` says, +a verified pre-migration backup named `.pre-workflow-storage_`, +outside the pruned `.backup_*` pool and kept until the user removes it (the preview +reports `backup.policy: "required"`; `test_the_backup_is_taken_even_with_the_policy_off`, +`test_the_backup_is_outside_the_pool_never_pruned_and_restores`); then batches of at most +200 rows, 16 MB of text or 0.5 s, each one `BEGIN IMMEDIATE` transaction found by a rowid +range search and read one row at a time, followed by 0.15 s with the write lock free so +a writer on another connection gets in (`test_a_writer_on_another_connection_gets_in_between_batches`). +Each text is read back through the accessor before its row changes. Each row has its own +savepoint: a row that cannot be rewritten (text SQLite cannot decode, a text that would +not come back) stays inline, is listed in `errors` and named in the job's last line, and +the rest go on, Reclaim space included (`TestRowErrors`); any other failure (a full disk, +an I/O error, a malformed database, a bug) rolls the batch back whole and fails the job +(`TestSystemErrorsFailTheBatch`, `TestSystemErrorsInTheJob`). With nothing to rewrite the +plan says "Nothing to rewrite" and Apply answers 409 `nothing-to-do` with no backup; only +the first apply keeps a `pre-workflow-storage` backup, and a later one (after a Cancel, a +rerun) takes a pooled safety backup (`TestOneKeptBackup`). The job loop's pause is tested +with a second-connection writer (`TestLiveLoopLetsWritersIn`). A reused document is decompressed and checked first and repaired from the text +when damaged, and `verify_database_file` (restore and backup verification) reads every +document and refuses a damaged one (`TestDamagedDocuments`). The legacy table rebuilds +refuse a file that already has the document id columns (`TestRebuildGuards`). +`tools/workflow_text_digests.py ` writes the sha256 of every record's stored +texts, so a `diff` of its output before and after the rewrite proves byte identity. It +loads only `database/workflow_documents.py` (no package start-up, no log handler) and +opens the database `mode=ro`, with `immutable=1` when no `-wal`/`-shm` exists, so it +creates none (`TestDigestToolSideEffects`). A damaged document's verify reason names it +and says what to do, and a backup that fails on it says so in the job line. A crash or Cancel leaves +committed batches new form and the rest old form, and the next apply continues +(`test_a_crash_mid_migration_leaves_every_row_readable_and_resumes`, +`test_cancel_stops_after_the_current_batch_and_a_rerun_continues`). Values that are not +text stay inline (protected). "Reclaim space" follows when every row was looked at +(`exclusive`, not cancellable): a disk check for VACUUM's copies of the live pages, +then sweep, VACUUM and `wal_checkpoint(TRUNCATE)`; when another job holds the registry it +does not start and the finished migration names Optimize database. No image or video +file is opened (`test_the_jobs_never_open_a_media_file` tracks every Python open under +the output folder and compares bytes and mtimes). Memory stays bounded by one row plus +the batch's distinct documents (`test_memory_stays_bounded_by_one_row`, tracemalloc). +Restoring a deduped backup on a version before this one is not supported: that +version's restore accepts the file and then reads rewritten images without their +workflow (README changelog). + ## Workspace data fields for the concept pages (#249) These fields are additive. The legacy `/prompt_manager/images/output` listing (no diff --git a/database/duplicate_copies.py b/database/duplicate_copies.py index fb33d3c..2baea72 100644 --- a/database/duplicate_copies.py +++ b/database/duplicate_copies.py @@ -25,6 +25,7 @@ from .annotations import _BUMP_INDEX from .models import normalize_image_path, utc_now_iso from .review import _BUMP_REVIEW_INDEX, _PENDING_OF_LOCATOR, _SETTLE_ITEM from .review_changes import notify_review_change +from .workflow_documents import documents_of_rows, release_documents LINKS_BY_FILE_SQL = "SELECT id, prompt_id FROM generated_images WHERE file_path = ?" LEGACY_LINKS_SQL = ( @@ -155,7 +156,9 @@ class DuplicateCopyOperations: break if existing is not None: conn.execute(MOVE_PAIR_LINKS_SQL, (existing[0], image_id)) + documents = documents_of_rows(conn, "id", [image_id]) conn.execute(DELETE_LINK_SQL, (image_id,)) + release_documents(conn, documents) counts["merged"] += 1 continue conn.execute( diff --git a/database/maintenance.py b/database/maintenance.py index 7b0c19c..76a7abd 100644 --- a/database/maintenance.py +++ b/database/maintenance.py @@ -42,12 +42,15 @@ from pathlib import Path from .execution_contexts import LIVE_HINT_CTES from .health import BATCH_ROWS, link_states, readable_roots, readonly_uri from .pairs import PairDependencyError +from .workflow_documents import documents_of_rows, release_documents MAINTENANCE_OPERATIONS = ( "duplicate_prompts", "optimize", "unused_components", "missing_links", + # Applied as its own batched job (database/workflow_storage.py). + "workflow_storage", ) MAX_DETAILS = 50 # detail rows a plan lists MAX_RESULT_ROWS = 200 # changed records an apply describes (the job lists 200) @@ -260,16 +263,24 @@ def _orphaned_ids(conn): def _delete_ids(conn, table, ids): - """Delete ``table`` rows by id in batches; returns the rows removed.""" + """Delete ``table`` rows by id in batches; returns the rows removed. + + The workflow documents only those rows referenced go with them (a + prompt's images go by cascade), in the same transaction. + """ if table not in ("prompts", "generated_images"): raise ValueError(f"Unsupported table {table!r}") removed = 0 for batch in _batches(ids): + documents = documents_of_rows( + conn, "id" if table == "generated_images" else "prompt_id", batch + ) cursor = conn.execute( f"DELETE FROM {table} WHERE id IN ({_marks(batch)})", # nosec B608 batch, ) removed += cursor.rowcount + release_documents(conn, documents) return removed @@ -379,10 +390,13 @@ class MaintenanceOperations: conn, primary_id, merged_metadata, strict=strict ) removed = 0 + # Images the lenient merge could not move go with their prompt. + documents = documents_of_rows(conn, "prompt_id", duplicate_ids) for duplicate_id in duplicate_ids: conn.execute("DELETE FROM prompts WHERE id = ?", (duplicate_id,)) removed += 1 self.logger.debug(f"Removed duplicate prompt {duplicate_id}") + release_documents(conn, documents) return removed, moved @staticmethod @@ -408,6 +422,8 @@ class MaintenanceOperations: roots = [str(root) for root in roots] if operation == "optimize": return self._plan_optimize() + if operation == "workflow_storage": + return self._plan_workflow_storage() if operation == "missing_links": rows = _read_snapshot(self.model.db_path, self._link_rows) return self._plan_missing_links(rows, roots) @@ -656,7 +672,7 @@ class MaintenanceOperations: records changed and up to ``MAX_RESULT_ROWS`` descriptions of them (``{prompt_id[, kept_id]}`` or ``{image_path}``). """ - if operation not in MAINTENANCE_OPERATIONS: + if operation not in MAINTENANCE_OPERATIONS or operation == "workflow_storage": raise ValueError(f"Unknown maintenance operation {operation!r}") if operation == "optimize": # VACUUM cannot run in a transaction; the plan binds the database. diff --git a/database/models.py b/database/models.py index 07c2e4f..f36fdbc 100644 --- a/database/models.py +++ b/database/models.py @@ -24,6 +24,13 @@ from .review_schema import ( review_without_annotations_error, ) 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, @@ -50,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): @@ -316,6 +336,8 @@ class PromptModel: raise sqlite3.DatabaseError(review_error) migrate_annotation_state(conn) migrate_review_state(conn) + # Last: the older generated_images rebuilds run only before file_path. + migrate_workflow_documents(conn) def _migrate_prompt_pairs(self, conn: sqlite3.Connection) -> None: conn.execute("SAVEPOINT prompt_pair_schema") @@ -711,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): @@ -876,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): @@ -1117,6 +1141,11 @@ class PromptModel: ) if review_error: return False, review_error + documents_error = workflow_documents_schema_error( + conn, allow_missing=not require_library_state, documents=True + ) + if documents_error: + return False, documents_error return True, "ok" except sqlite3.Error as e: return False, f"Invalid SQLite database: {e}" diff --git a/database/operations.py b/database/operations.py index 54a4a79..aaaf2bd 100644 --- a/database/operations.py +++ b/database/operations.py @@ -25,6 +25,14 @@ from .execution_outputs import ExecutionOutputOperations from .history_reads import HistoryReadOperations from .health import HealthOperations, iter_consistency_issues from .duplicate_copies import DuplicateCopyOperations +from .workflow_storage import WorkflowStorageOperations +from .workflow_documents import ( + DocumentReader, + documents_of_rows, + release_documents, + resolve_image_row, + stored_values, +) from .maintenance import ( # noqa: F401 (ORPHANED_PROMPTS_SQL is re-exported) ORPHANED_PROMPTS_SQL, _LINK_ROWS_SQL, @@ -339,6 +347,7 @@ class PromptDatabase( HealthOperations, MaintenanceOperations, DuplicateCopyOperations, + WorkflowStorageOperations, ): """Database operations class for managing prompts.""" @@ -776,10 +785,12 @@ class PromptDatabase( """ with self._pair_transaction() as conn: self._ensure_prompts_unpaired(conn, [prompt_id]) + documents = documents_of_rows(conn, "prompt_id", [prompt_id]) # First delete related images to avoid foreign key constraint conn.execute( "DELETE FROM generated_images WHERE prompt_id = ?", (prompt_id,) ) + release_documents(conn, documents) # Then delete the prompt cursor = conn.execute("DELETE FROM prompts WHERE id = ?", (prompt_id,)) conn.commit() @@ -803,10 +814,12 @@ class PromptDatabase( return 0 self._ensure_prompts_unpaired(conn, ids) placeholders = ",".join("?" * len(ids)) + documents = documents_of_rows(conn, "prompt_id", ids) conn.execute( f"DELETE FROM generated_images WHERE prompt_id IN ({placeholders})", ids, ) + release_documents(conn, documents) conn.execute( f"DELETE FROM prompt_tags WHERE prompt_id IN ({placeholders})", ids, @@ -1607,6 +1620,14 @@ class PromptDatabase( filename = os.path.basename(image_path) file_path = normalize_image_path(image_path) file_size, width, height, image_format = _coerce_file_info(metadata) + # Each workflow and prompt graph is stored once (new form). + documents = stored_values( + conn, + { + "workflow_data": dumps_strict(metadata.get("workflow", {})), + "prompt_metadata": dumps_strict(metadata.get("prompt", {})), + }, + ) # INSERT OR IGNORE skips a file already linked to this prompt # (same prompt_id + normalised full path) @@ -1614,8 +1635,9 @@ class PromptDatabase( """ INSERT OR IGNORE INTO generated_images (prompt_id, image_path, filename, file_path, file_size, width, - height, format, workflow_data, prompt_metadata, parameters) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + height, format, workflow_data, prompt_metadata, parameters, + workflow_document_id, prompt_document_id) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( prompt_id_int, @@ -1626,11 +1648,22 @@ class PromptDatabase( width, height, image_format, - dumps_strict(metadata.get("workflow", {})), - dumps_strict(metadata.get("prompt", {})), + documents["workflow_data"], + documents["prompt_metadata"], dumps_strict(metadata.get("parameters", {})), + documents["workflow_document_id"], + documents["prompt_document_id"], ), ) + if cursor.rowcount == 0: + # Already linked: a document made only for this row goes. + release_documents( + conn, + ( + documents["workflow_document_id"], + documents["prompt_document_id"], + ), + ) conn.commit() if cursor.rowcount == 0: @@ -1664,7 +1697,8 @@ class PromptDatabase( """, (prompt_id,), ) - return [self._image_row_to_dict(row) for row in cursor.fetchall()] + reader = DocumentReader(conn) + return [self._image_row_to_dict(row, reader) for row in cursor] def get_recent_images( self, limit: int = 50, offset: int = 0 @@ -1690,7 +1724,8 @@ class PromptDatabase( """, (int(limit), max(0, int(offset))), ) - return [self._image_row_to_dict(row) for row in cursor.fetchall()] + reader = DocumentReader(conn) + return [self._image_row_to_dict(row, reader) for row in cursor] def get_all_images( self, @@ -1724,10 +1759,10 @@ class PromptDatabase( with self.model.get_connection() as conn: cursor = conn.execute(sql, params) - rows = cursor.fetchall() + reader = DocumentReader(conn) result = [] - for row in rows: - data = self._image_row_to_dict(row) + for row in cursor: + data = self._image_row_to_dict(row, reader) if row["_prompt_tags_list"]: data["prompt_tags"] = row["_prompt_tags_list"].split("|||") else: @@ -1765,7 +1800,8 @@ class PromptDatabase( max(0, int(offset)), ), ) - return [self._image_row_to_dict(row) for row in cursor.fetchall()] + reader = DocumentReader(conn) + return [self._image_row_to_dict(row, reader) for row in cursor] def get_image_by_id(self, image_id: int) -> Optional[Dict[str, Any]]: """ @@ -1782,7 +1818,7 @@ class PromptDatabase( "SELECT * FROM generated_images WHERE id = ?", (image_id,) ) row = cursor.fetchone() - return self._image_row_to_dict(row) if row else None + return self._image_row_to_dict(row, DocumentReader(conn)) if row else None def delete_image(self, image_id: int) -> bool: """ @@ -1795,9 +1831,11 @@ class PromptDatabase( bool: True if deletion was successful """ with self.model.get_connection() as conn: + documents = documents_of_rows(conn, "id", [image_id]) cursor = conn.execute( "DELETE FROM generated_images WHERE id = ?", (image_id,) ) + release_documents(conn, documents) conn.commit() return cursor.rowcount > 0 @@ -1831,12 +1869,16 @@ class PromptDatabase( cursor = conn.execute("SELECT id, image_path FROM generated_images") images = cursor.fetchall() - for image in images: - if not os.path.exists(image["image_path"]): - conn.execute( - "DELETE FROM generated_images WHERE id = ?", (image["id"],) - ) - removed_count += 1 + removed = [ + image["id"] + for image in images + if not os.path.exists(image["image_path"]) + ] + documents = documents_of_rows(conn, "id", removed) + for image_id in removed: + conn.execute("DELETE FROM generated_images WHERE id = ?", (image_id,)) + removed_count += 1 + release_documents(conn, documents) conn.commit() @@ -2127,11 +2169,13 @@ class PromptDatabase( count = 0 with self._pair_transaction() as conn: self._ensure_prompts_unpaired(conn, prompt_ids) + documents = documents_of_rows(conn, "prompt_id", prompt_ids) for pid in prompt_ids: conn.execute("DELETE FROM generated_images WHERE prompt_id = ?", (pid,)) cursor = conn.execute("DELETE FROM prompts WHERE id = ?", (pid,)) if cursor.rowcount > 0: count += 1 + release_documents(conn, documents) conn.commit() return count @@ -2205,7 +2249,8 @@ class PromptDatabase( cursor = conn.execute( """SELECT gi.prompt_id, p.text, p.category, p.rating, p.notes, gi.workflow_data, gi.prompt_metadata, gi.generation_time, - gi.parameters, gi.id, gi.image_path, gi.file_path + gi.parameters, gi.id, gi.image_path, gi.file_path, + gi.workflow_document_id, gi.prompt_document_id FROM generated_images gi JOIN prompts p ON gi.prompt_id = p.id WHERE gi.image_path = ? OR gi.file_path = ? @@ -2285,12 +2330,21 @@ class PromptDatabase( "text_source": source, } tags = self._get_prompt_tags(conn, row[0]) + texts = resolve_image_row(DocumentReader(conn), row) try: - workflow = json.loads(row[5]) if row[5] else None + workflow = ( + json.loads(texts["workflow_data"]) + if texts["workflow_data"] + else None + ) except (json.JSONDecodeError, TypeError): workflow = None try: - prompt_meta = json.loads(row[6]) if row[6] else None + prompt_meta = ( + json.loads(texts["prompt_metadata"]) + if texts["prompt_metadata"] + else None + ) except (json.JSONDecodeError, TypeError): prompt_meta = None return { @@ -2540,17 +2594,20 @@ class PromptDatabase( with self.model.get_connection() as conn: return list(iter_consistency_issues(conn)) - def _image_row_to_dict(self, row: sqlite3.Row) -> Dict[str, Any]: + def _image_row_to_dict( + self, row: sqlite3.Row, reader: DocumentReader + ) -> Dict[str, Any]: """ Convert an image database row to a dictionary with parsed JSON fields. Args: row: SQLite row object + reader: The workflow document reader of the row's connection Returns: Dictionary representation of the row """ - data = dict(row) + data = resolve_image_row(reader, row) # Parse JSON fields for field in ["workflow_data", "prompt_metadata", "parameters"]: diff --git a/database/pair_ingestion.py b/database/pair_ingestion.py index 23c7044..da1db08 100644 --- a/database/pair_ingestion.py +++ b/database/pair_ingestion.py @@ -5,6 +5,13 @@ from typing import Any, Dict from .models import normalize_image_path from .operations import _coerce_file_info +from .workflow_documents import ( + TEXT_COLUMNS, + DocumentReader, + release_documents, + resolve_image_row, + stored_values, +) try: from ..utils.hashing import generate_prompt_hash @@ -116,25 +123,35 @@ def _ingest_resolved_pair_in( changed = False if existing: image_id = existing["id"] + # Compared as stored text, whichever form the row is in. + current = resolve_image_row(DocumentReader(conn), existing) updates = { - key: value for key, value in snapshots.items() if existing[key] != value + key: value for key, value in snapshots.items() if current[key] != value } if updates: + replaced = [ + existing[reference] + for column, reference in TEXT_COLUMNS + if column in updates + ] + updates = _document_form(conn, updates) conn.execute( "UPDATE generated_images SET " + ",".join(f"{key}=?" for key in updates) + " WHERE id=?", [*updates.values(), image_id], ) + release_documents(conn, replaced) changed = True else: - columns = ["prompt_id", "image_path", "filename", "file_path", *snapshots] + stored = _document_form(conn, snapshots) + columns = ["prompt_id", "image_path", "filename", "file_path", *stored] values = [ pair["positive_prompt_id"], image_path, os.path.basename(image_path), normalized, - *snapshots.values(), + *stored.values(), ] image_id = conn.execute( "INSERT INTO generated_images(" @@ -162,3 +179,18 @@ def _ingest_resolved_pair_in( "pair_id": pair_id, "image_id": image_id, } + + +def _document_form(conn: Any, values: Dict[str, Any]) -> Dict[str, Any]: + """``values`` with the workflow and prompt graph texts stored once. + + Each of those columns becomes NULL beside its document id (new form); + every other column is unchanged. + """ + texts = { + column: values[column] + for column, _reference in TEXT_COLUMNS + if column in values + } + rest = {key: value for key, value in values.items() if key not in texts} + return {**rest, **stored_values(conn, texts)} diff --git a/database/workflow_documents.py b/database/workflow_documents.py new file mode 100644 index 0000000..0f6173c --- /dev/null +++ b/database/workflow_documents.py @@ -0,0 +1,477 @@ +"""Each workflow and prompt graph stored once, compressed (#249 S6 carry-over). + +``generated_images.workflow_data`` (the UI workflow) and ``prompt_metadata`` +(the API prompt graph) used to hold the same large JSON text again for every +image of a run. Now each distinct text is one ``workflow_documents`` row: + +- ``sha256``: the SHA-256 of the text's exact UTF-8 bytes (unique key); +- ``byte_length``: the length of those bytes; +- ``data``: ``zlib.compress`` of those bytes. + +An image row references its texts through ``workflow_document_id`` and +``prompt_document_id``. Two forms coexist until the migration has rewritten +every row, and both stay readable forever: + +- old form: the text inline in ``workflow_data`` / ``prompt_metadata``, no id; +- new form: the inline column NULL and the document id set. + +The text is stored and hashed exactly as written, never parsed and +re-serialized, because a JSON round trip can change bytes (key order, +spacing, number spelling, duplicate keys, ComfyUI's ``NaN`` tokens). + +``read_stored_text`` is the one decoder. Inline text, when present, wins (a +row an older app version rewrote inline stays right); otherwise the document +is read by its rowid, decompressed within its recorded length and checked +against its length and hash. A document that fails a check reads as +unreadable (None, logged), never as other text. ``resolve_image_row`` puts +the texts back into a ``generated_images`` row and drops the id keys, so +every payload keeps its shape. + +No trigger: every restore validator refuses triggers in a restored file. A +document is removed in code: ``release_documents`` in the same transaction +that drops a row's last reference, and ``sweep_documents`` (at startup, after +the migration and before Reclaim space) for anything a path missed. +""" + +import hashlib +import logging +import re +import sqlite3 +import zlib +from collections import OrderedDict +from typing import Any, Dict, Iterable, Optional, Set + +logger = logging.getLogger("prompt_manager.database.workflow_documents") + +TABLE = "workflow_documents" +# (inline text column, document id column) on generated_images. +TEXT_COLUMNS = ( + ("workflow_data", "workflow_document_id"), + ("prompt_metadata", "prompt_document_id"), +) +DOCUMENT_ID_COLUMNS = tuple(reference for _column, reference in TEXT_COLUMNS) +INDEXES = { + "workflow_document_id": "idx_image_workflow_document", + "prompt_document_id": "idx_image_prompt_document", +} +COMPRESSION_LEVEL = 6 +# A text larger than this stays inline: a document is decompressed in one +# piece, and its recorded length bounds that (a restored file cannot make a +# small row expand without limit). +MAX_DOCUMENT_BYTES = 256 * 2**20 +ID_BATCH = 400 # ids bound per statement (under SQLite's old 999 limit) + +_DDL = ( + f"CREATE TABLE {TABLE} (id INTEGER PRIMARY KEY, sha256 BLOB NOT NULL UNIQUE, " + "byte_length INTEGER NOT NULL, data BLOB NOT NULL)" +) +_AUTOINDEX = f"sqlite_autoindex_{TABLE}_1" +_REFERENCE_TYPE = f"INTEGER REFERENCES {TABLE}(id)" + +_UNREFERENCED = ( + "NOT EXISTS (SELECT 1 FROM generated_images g " + "WHERE g.workflow_document_id=workflow_documents.id) " + "AND NOT EXISTS (SELECT 1 FROM generated_images g " + "WHERE g.prompt_document_id=workflow_documents.id)" +) +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 (?, ?, ?)" +) +# Built from constants only; the ids are bound as "?" marks. +RELEASE_SQL = "DELETE FROM workflow_documents WHERE id IN ({marks}) AND " +RELEASE_SQL += _UNREFERENCED +# The sweep walks the ids through the small unique index, not the blobs. +SWEEP_SQL = ( + "DELETE FROM workflow_documents WHERE id IN (SELECT id FROM workflow_documents " + "INDEXED BY sqlite_autoindex_workflow_documents_1 WHERE " +) +SWEEP_SQL += _UNREFERENCED + ")" +IMAGE_DOCUMENTS_SQL = ( + "SELECT workflow_document_id, prompt_document_id FROM generated_images WHERE " +) + +# Statements whose query plans tests/test_workflow_documents.py records. +QUERY_PLAN_STATEMENTS = { + "lookup by hash": (LOOKUP_SQL, (b"\0" * 32,)), + "read by id": (READ_SQL, (1,)), + "release": (RELEASE_SQL.format(marks="?,?"), (1, 2)), + "sweep": (SWEEP_SQL, ()), + "documents of images": (IMAGE_DOCUMENTS_SQL + "id IN (?,?)", (1, 2)), + "documents of prompts": (IMAGE_DOCUMENTS_SQL + "prompt_id IN (?,?)", (1, 2)), +} + + +class DocumentMismatch(Exception): + """A document would not give back exactly the text it was made from.""" + + +def _compact(sql: str) -> str: + return re.sub(r"\s+", "", sql or "").lower() + + +# ── schema ──────────────────────────────────────────────────────────── + + +def _reference_columns(conn: sqlite3.Connection) -> Dict[str, tuple]: + return { + row[1]: tuple(row) + for row in conn.execute("PRAGMA table_info(generated_images)") + if row[1] in DOCUMENT_ID_COLUMNS + } + + +def _reference_error(conn: sqlite3.Connection, columns: Dict[str, tuple]): + for name, info in columns.items(): + # (cid, name, type, notnull, default, pk) + if (str(info[2]).upper(), info[3], info[4], info[5]) != ("INTEGER", 0, None, 0): + return "Workflow document references have an invalid column" + keys = { + (row[3], row[2], row[4]) + for row in conn.execute("PRAGMA foreign_key_list(generated_images)") + if row[3] in DOCUMENT_ID_COLUMNS + } + if keys != {(name, TABLE, "id") for name in DOCUMENT_ID_COLUMNS}: + return "Workflow document references have invalid foreign keys" + for column, index in INDEXES.items(): + info = conn.execute(f'PRAGMA index_info("{index}")').fetchall() + owner = conn.execute( + "SELECT tbl_name FROM sqlite_master WHERE type='index' AND name=?", + (index,), + ).fetchone() + if [row[2] for row in info] != [column] or not owner: + return "Workflow document references are missing their index" + if owner[0] != "generated_images": + return "Workflow document references are missing their index" + return None + + +def _value_error(conn: sqlite3.Connection): + if conn.execute( + "SELECT 1 FROM workflow_documents WHERE typeof(sha256)!='blob' " + "OR length(sha256)!=32 OR typeof(byte_length)!='integer' OR byte_length<0 " + "OR byte_length>? OR typeof(data)!='blob' LIMIT 1", + (MAX_DOCUMENT_BYTES,), + ).fetchone(): + return "Workflow documents contain invalid values" + for column in DOCUMENT_ID_COLUMNS: + # typeof() reads no overflow page; a dangling id is a rowid probe. + # ``column`` is one of DOCUMENT_ID_COLUMNS, never caller data. + sql = ( + f"SELECT 1 FROM generated_images g WHERE typeof(g.{column}) " # nosec B608 + f"NOT IN ('null','integer') OR (g.{column} IS NOT NULL AND NOT EXISTS " + f"(SELECT 1 FROM workflow_documents d WHERE d.id=g.{column})) LIMIT 1" + ) + if conn.execute(sql).fetchone(): + return "Workflow document references point at missing documents" + return None + + +def damaged_message(document_id: int) -> str: + """The verify reason for a damaged document: which one, and what to do.""" + return ( + f"Workflow document {document_id} is damaged: its text no longer matches " + "its checksum, and the images that use it show no workflow. To fix it, " + "restore a backup from before the damage, or store a new image with the " + "same workflow (that repairs it)." + ) + + +def _damaged_document(conn: sqlite3.Connection) -> Optional[int]: + """The first document failing its length or hash check, or None.""" + for row in conn.execute(ALL_DOCUMENTS_SQL): + if _decode(row[:3], row[3], quiet=True) is None: + return row[3] + return None + + +def workflow_documents_schema_error( + conn: sqlite3.Connection, *, allow_missing: bool = True, documents: bool = False +) -> Optional[str]: + """Reject an unsupported document schema or invalid stored values (reads only). + + A database from before the dedupe has neither the table nor the id + columns; that is acceptable unless ``allow_missing`` is false. Half of + the schema never is. + """ + objects = conn.execute( + "SELECT type,name,sql FROM sqlite_master WHERE name=? " + "OR (type IN ('index','trigger') AND tbl_name=?)", + (TABLE, TABLE), + ).fetchall() + has_images = conn.execute( + "SELECT 1 FROM sqlite_master WHERE type='table' AND name='generated_images'" + ).fetchone() + columns = _reference_columns(conn) if has_images else {} + if not objects and not columns: + if allow_missing or not has_images: + return None + return "Missing workflow document schema" + if not objects or set(columns) != set(DOCUMENT_ID_COLUMNS): + return "Incomplete workflow document schema" + tables = [sql for kind, _name, sql in objects if kind == "table"] + others = {(kind, name, sql) for kind, name, sql in objects if kind != "table"} + if len(tables) != 1 or others != {("index", _AUTOINDEX, None)}: + return "Workflow documents name is reserved by an unsupported schema object" + # The exact definition fixes the columns, their types and constraints. + if _compact(tables[0]) != _compact(_DDL): + return "Workflow documents have an invalid table definition" + error = _reference_error(conn, columns) or _value_error(conn) + if error is None and documents: + # Restore and backup verification only: every document is read. + damaged = _damaged_document(conn) + if damaged is not None: + return damaged_message(damaged) + return error + + +def migrate_workflow_documents(conn: sqlite3.Connection) -> None: + """Install the table, the id columns and their indexes atomically. + + Runs after every older migration (the table rebuilds only run on + databases without ``file_path``, before this). Validates before and + after, then sweeps documents no row references. + """ + conn.execute("SAVEPOINT workflow_documents_schema") + try: + error = workflow_documents_schema_error(conn) + if error: + raise sqlite3.DatabaseError(error) + conn.execute(_DDL.replace("CREATE TABLE", "CREATE TABLE IF NOT EXISTS", 1)) + present = _reference_columns(conn) + for column in DOCUMENT_ID_COLUMNS: + if column not in present: + conn.execute( + f"ALTER TABLE generated_images ADD COLUMN {column} " + f"{_REFERENCE_TYPE}" + ) + for column, index in INDEXES.items(): + conn.execute( + f"CREATE INDEX IF NOT EXISTS {index} ON generated_images({column})" + ) + error = workflow_documents_schema_error(conn, allow_missing=False) + if error: + raise sqlite3.DatabaseError(error) + sweep_documents(conn) + conn.execute("RELEASE workflow_documents_schema") + except BaseException: + conn.execute("ROLLBACK TO workflow_documents_schema") + conn.execute("RELEASE workflow_documents_schema") + raise + + +# ── write ───────────────────────────────────────────────────────────── + + +def storable(text: Any) -> bool: + """Whether ``text`` can be a document (text within MAX_DOCUMENT_BYTES).""" + if not isinstance(text, str): + return False + # UTF-8 takes at most 4 bytes a character: most texts need no encoding. + if len(text) * 4 <= MAX_DOCUMENT_BYTES: + return True + return len(text.encode("utf-8")) <= MAX_DOCUMENT_BYTES + + +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, 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 + if not isinstance(text, str): + raise TypeError("Only text is stored as a workflow document") + raw = text.encode("utf-8") + if len(raw) > MAX_DOCUMENT_BYTES: + raise ValueError("Text is too large for a workflow document") + digest = hashlib.sha256(raw).digest() + if not conn.in_transaction: + conn.execute("BEGIN IMMEDIATE") + row = conn.execute(LOOKUP_SQL, (digest,)).fetchone() + 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") + 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( + conn: sqlite3.Connection, texts: Dict[str, Optional[str]] +) -> Dict[str, Any]: + """New-form column values for ``{text column: text}`` (writers). + + Each text becomes ``{column: None, id column: document id}``; None + stays ``{column: None, id column: None}``; a value that cannot be a + document (not text, or too large) stays inline with no id. + """ + references = dict(TEXT_COLUMNS) + values: Dict[str, Any] = {} + for column, text in texts.items(): + reference = references[column] + if text is None or not storable(text): + values[column], values[reference] = text, None + else: + values[column], values[reference] = None, store_text(conn, text) + return values + + +# ── read ────────────────────────────────────────────────────────────── + + +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 + digest, length, data = bytes(row[0]), row[1], row[2] + try: + if not isinstance(length, int) or not 0 <= length <= MAX_DOCUMENT_BYTES: + raise ValueError("invalid length") + inflater = zlib.decompressobj() + raw = inflater.decompress(bytes(data), length + 1) + if len(raw) != length or inflater.unconsumed_tail or not inflater.eof: + raise ValueError("length differs") + if hashlib.sha256(raw).digest() != digest: + raise ValueError("hash differs") + return raw.decode("utf-8") + except (zlib.error, ValueError, TypeError) as error: + if not quiet: + logger.warning(f"Workflow document {document_id} is unreadable: {error}") + return None + + +def read_stored_text( + conn: sqlite3.Connection, inline: Any, document_id: Optional[int] +) -> Any: + """The stored text of one column: inline when present, else the document. + + Returns None when there is neither, or when the document is missing or + fails its checks (logged). An inline value is returned unchanged. + """ + if inline is not None: + return inline + if document_id is None: + return None + return _decode(conn.execute(READ_SQL, (document_id,)).fetchone(), document_id) + + +class DocumentReader: + """``read_stored_text`` with a small memo for list reads on one connection.""" + + MEMO = 32 + + def __init__(self, conn: sqlite3.Connection): + self.conn = conn + self._memo: "OrderedDict[int, Optional[str]]" = OrderedDict() + + def text(self, inline: Any, document_id: Optional[int]) -> Any: + if inline is not None or document_id is None: + return read_stored_text(self.conn, inline, document_id) + if document_id in self._memo: + self._memo.move_to_end(document_id) + return self._memo[document_id] + text = read_stored_text(self.conn, None, document_id) + self._memo[document_id] = text + if len(self._memo) > self.MEMO: + self._memo.popitem(last=False) + return text + + +def resolve_image_row(reader: DocumentReader, row) -> Dict[str, Any]: + """A ``generated_images`` row as a dict with its texts read back. + + The id columns are dropped, so a payload built from it has the shape it + had before the dedupe. Columns the row did not select stay absent. + """ + data = dict(row) + for column, reference in TEXT_COLUMNS: + has_reference = reference in data + document_id = data.pop(reference, None) + if column in data or has_reference: + data[column] = reader.text(data.get(column), document_id) + return data + + +# ── garbage ─────────────────────────────────────────────────────────── + + +def _chunks(values): + values = list(values) + for start in range(0, len(values), ID_BATCH): + yield values[start : start + ID_BATCH] + + +def image_documents(conn: sqlite3.Connection, where: str, params) -> Set[int]: + """The document ids the image rows matching ``where`` reference. + + ``where`` is a fixed condition written in this package (``id IN (...)``, + ``prompt_id IN (...)``), never caller data; values go in ``params``. + """ + ids: Set[int] = set() + sql = IMAGE_DOCUMENTS_SQL + where + for row in conn.execute(sql, list(params)): # nosec B608 + ids.update(value for value in row if value is not None) + return ids + + +def documents_of_rows( + conn: sqlite3.Connection, column: str, values: Iterable[Any] +) -> Set[int]: + """Document ids of the image rows whose ``column`` is in ``values``. + + ``column`` is ``id`` or ``prompt_id`` (the rows a prompt delete cascades to). + """ + if column not in ("id", "prompt_id"): + raise ValueError(f"Unsupported column {column!r}") + ids: Set[int] = set() + for chunk in _chunks(values): + marks = ",".join("?" for _ in chunk) + ids |= image_documents(conn, f"{column} IN ({marks})", chunk) + return ids + + +def release_documents(conn: sqlite3.Connection, ids: Iterable[Optional[int]]) -> int: + """Delete the documents in ``ids`` that no image row references now. + + Runs in the caller's transaction, right after the delete or update that + dropped the references. Returns how many documents were removed. + """ + removed = 0 + for chunk in _chunks(sorted({i for i in ids if i is not None})): + marks = ",".join("?" for _ in chunk) + # Only "?" marks are interpolated. + sql = RELEASE_SQL.format(marks=marks) + removed += conn.execute(sql, chunk).rowcount # nosec B608 + return removed + + +def sweep_documents(conn: sqlite3.Connection) -> int: + """Delete every document no image row references; returns the count.""" + return conn.execute(SWEEP_SQL).rowcount diff --git a/database/workflow_storage.py b/database/workflow_storage.py new file mode 100644 index 0000000..be3a32f --- /dev/null +++ b/database/workflow_storage.py @@ -0,0 +1,455 @@ +"""The workflow storage migration and Reclaim space (#249 S6 carry-over). + +``migrate_workflow_storage`` rewrites old-form image rows (workflow and +prompt graph text inline) to the new form (a ``workflow_documents`` id, +``database/workflow_documents.py``) in batches: + +- each batch is one ``BEGIN IMMEDIATE`` transaction over at most + ``BATCH_ROWS`` rows or about ``BATCH_BYTES`` of text, found by a rowid + range search (``NEXT_IDS_SQL``) and read one row at a time, so memory is + bounded by one row plus the batch's distinct documents; +- each text is stored (or its document reused) and read back through the + 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; +- 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 +from the first old-form row. Only database rows change: no image or video +file is opened. + +``reclaim_space`` is the VACUUM afterwards: it sweeps unreferenced +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 ( + TEXT_COLUMNS, + DocumentMismatch, + read_stored_text, + release_documents, + storable, + store_text, + sweep_documents, +) + +OPERATION = "workflow_storage" +# The first apply's backup: ``.pre-workflow-storage_``, outside the +# pruned ``.backup_*`` pool. Only one is ever made; later applies take a +# pooled safety backup. +PRE_MIGRATION_LABEL = "pre-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 + + +def _row_error(error) -> bool: + """Whether ``error`` is one row's own cause, which keeps only that row inline. + + Text SQLite cannot decode for Python, and a text that would not come + back identical. Anything else (a full disk, an I/O error, a malformed + database, a bug) fails the whole batch. The decode error is matched by + its message: ``sqlite_errorcode`` needs Python 3.11. + """ + if isinstance(error, (DocumentMismatch, UnicodeError)): + return True + return isinstance(error, sqlite3.OperationalError) and str(error).startswith( + "Could not decode" + ) + + +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 = ( + "SELECT id FROM generated_images WHERE id>? AND " + "(typeof(workflow_data)!='null' OR typeof(prompt_metadata)!='null') " + "ORDER BY id LIMIT ?" +) +ROW_SQL = ( + "SELECT workflow_data, prompt_metadata, workflow_document_id, " + "prompt_document_id FROM generated_images WHERE id=?" +) +UPDATE_SQL = ( + "UPDATE generated_images SET {assignments} WHERE id=?" # columns are constants +) +# One pass over the record headers: rows with inline text, and inline +# values that are not text (they stay inline). +INLINE_COUNTS_SQL = ( + "SELECT " + "SUM(typeof(workflow_data)='text' OR typeof(prompt_metadata)='text'), " + "SUM(typeof(workflow_data) NOT IN ('null','text')) " + "+ SUM(typeof(prompt_metadata) NOT IN ('null','text')) " + "FROM generated_images" +) +# octet_length() (SQLite 3.43+) reads no overflow page either. +INLINE_BYTES_SQL = ( + "SELECT SUM(octet_length(workflow_data)), SUM(octet_length(prompt_metadata)) " + "FROM generated_images" +) + +QUERY_PLAN_STATEMENTS = { + "next old-form rows": (NEXT_IDS_SQL, (0, BATCH_ROWS)), + "one row": (ROW_SQL, (1,)), +} + + +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, 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 + for index, (column, reference) in enumerate(TEXT_COLUMNS): + value, current = row[index], row[index + len(TEXT_COLUMNS)] + if value is None: + continue + if not storable(value): + protected += 1 + continue + 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: + raise DocumentMismatch(f"Image {image_id}: {column} would change") + changes[column], changes[reference] = None, document_id + if current is not None and current != document_id: + replaced.append(current) + size += len(value) + if changes: + assignments = ",".join(f"{key}=?" for key in changes) + conn.execute( + UPDATE_SQL.format(assignments=assignments), + [*changes.values(), image_id], + ) + release_documents(conn, replaced) + return bool(changes), protected, size + + +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. + + 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_error``) 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, + "errors": [], + "done": not ids, + } + verified, checked, size = {}, set(), 0 + for image_id in ids: + 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 Exception as error: + if not _row_error(error): + raise + 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 or _clock() - started >= max_seconds: + break + conn.commit() + return result + except BaseException: + conn.rollback() + raise + + +def pre_migration_backups(db_path): + """The kept pre-migration backups beside ``db_path``, earliest first.""" + db_path = os.path.abspath(db_path) + prefix = os.path.basename(db_path) + "." + PRE_MIGRATION_LABEL + "_" + folder = os.path.dirname(db_path) + try: + names = os.listdir(folder) + except OSError: + return [] + return sorted( + name + for name in names + if name.startswith(prefix) and os.path.isfile(os.path.join(folder, name)) + ) + + +def _file_bytes(db_path): + return sum( + os.path.getsize(path) + for path in (db_path, db_path + "-wal") + if os.path.exists(path) + ) + + +class WorkflowStorageOperations: + """PromptDatabase mixin: the workflow storage migration and Reclaim space.""" + + def migrate_workflow_storage( + self, + *, + on_batch=None, + should_stop=None, + max_rows=BATCH_ROWS, + max_bytes=BATCH_BYTES, + ) -> dict: + """Rewrite every old-form row, batch by batch (blocking). + + ``on_batch(progress)`` runs after each committed batch with the + running totals ``{rows, protected, batches, last_id}``; + ``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(): + stopped = True + break + 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 + totals["last_id"] = batch["last_id"] + if on_batch is not None: + on_batch(dict(totals)) + removed = 0 if stopped else self.sweep_workflow_documents() + return { + "rows": totals["rows"], + "protected": totals["protected"], + "batches": totals["batches"], + "stopped": stopped, + "errors": errors, + "documents_removed": removed, + } + + def migrate_workflow_batch( + self, after_id, max_rows=BATCH_ROWS, max_bytes=BATCH_BYTES + ) -> dict: + """One batch on the calling thread's connection (``migrate_batch``).""" + return migrate_batch(self.model.get_connection(), after_id, max_rows, max_bytes) + + def sweep_workflow_documents(self) -> int: + """Remove every document no image row references (one transaction).""" + conn = self.model.get_connection() + with conn: + return sweep_documents(conn) + + def workflow_storage_status(self) -> dict: + """``workflow_storage_counts`` on one read-only snapshot (blocking).""" + from .maintenance import _read_snapshot + + return _read_snapshot(self.model.db_path, self.workflow_storage_counts) + + def workflow_storage_counts(self, conn) -> dict: + """``{rows, protected, documents, inline_bytes}`` on ``conn`` (reads only). + + ``inline_bytes`` is None where SQLite has no ``octet_length``. + """ + rows, protected = conn.execute(INLINE_COUNTS_SQL).fetchone() + documents = conn.execute("SELECT COUNT(*) FROM workflow_documents").fetchone() + inline_bytes = None + if sqlite3.sqlite_version_info >= (3, 43, 0): + row = conn.execute(INLINE_BYTES_SQL).fetchone() + inline_bytes = sum(value or 0 for value in row) + return { + "rows": rows or 0, + "protected": protected or 0, + "documents": documents[0], + "inline_bytes": inline_bytes, + } + + def reclaim_space(self) -> dict: + """Sweep, VACUUM and truncate the WAL (blocking); file bytes before/after. + + The caller holds the registry's ``exclusive`` slot. Raises whatever + stopped VACUUM; the database is then unchanged. + """ + db_path = self.model.db_path + removed = self.sweep_workflow_documents() + before = _file_bytes(db_path) + self.model.vacuum_database() + with closing(sqlite3.connect(db_path)) as checkpoint: + busy = checkpoint.execute("PRAGMA wal_checkpoint(TRUNCATE)").fetchone()[0] + return { + "before_bytes": before, + "after_bytes": _file_bytes(db_path), + "documents_removed": removed, + "checkpoint_busy": bool(busy), + } + + # ── preview ─────────────────────────────────────────────────────── + + def _plan_workflow_storage(self): + """The reviewed plan of the migration (one read-only snapshot).""" + from .maintenance import ( + MaintenanceDraft, + _Fingerprint, + _megabytes, + _plural, + _read_snapshot, + ) + + def read(conn): + sizes = [ + conn.execute(f"PRAGMA {name}").fetchone()[0] + for name in ("page_size", "page_count", "freelist_count") + ] + return sizes, self.workflow_storage_counts(conn) + + (page_size, page_count, free_pages), counts = _read_snapshot( + self.model.db_path, read + ) + size = page_size * page_count + details = [ + { + "label": "Database file", + "detail": f"{_megabytes(size)}, of which " + f"{_megabytes(free_pages * page_size)} is free pages", + }, + { + "label": "Image records", + "detail": f"{_plural(counts['rows'], 'record')} keep their workflow " + "or prompt graph inline; each distinct text is stored once, " + "compressed", + }, + ] + if not counts["rows"]: + details[1]["detail"] = NOTHING_TO_DO + if counts["inline_bytes"] is not None: + details.append( + { + "label": "Inline text", + "detail": f"{_megabytes(counts['inline_bytes'])} now", + } + ) + if counts["documents"]: + details.append( + { + "label": "Already stored once", + "detail": _plural(counts["documents"], "workflow document"), + } + ) + name = os.path.basename(self.model.db_path) + kept = pre_migration_backups(self.model.db_path) + backup = ( + POOLED_BACKUP_NOTE.format( + name=f"{name}.backup_", kept=kept[0] + ) + if kept + else BACKUP_NOTE.format( + name=f"{name}.{PRE_MIGRATION_LABEL}_" + ) + ) + details += [ + {"label": "Backup", "detail": backup}, + {"label": "Reclaim space", "detail": RECLAIM_NOTE}, + {"label": "Older versions", "detail": OLDER_VERSIONS_NOTE}, + ] + protected = counts["protected"] + return MaintenanceDraft( + operation=OPERATION, + fingerprint=_Fingerprint(OPERATION).hexdigest(), + targets=[], + scope={ + "count": counts["rows"], + "label": ( + _plural(counts["rows"], "image record") + " to rewrite" + if counts["rows"] + else "Nothing to rewrite" + ), + }, + protected={ + "count": protected, + "reasons": ( + [f"{_plural(protected, 'stored value')} not text; kept as is"] + if protected + else [] + ), + }, + media={"deletes_files": False, "label": MEDIA_NOTE}, + details=details, + ) + + +MEDIA_NOTE = "No image or video file is opened or changed; only the database." +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)." +) +POOLED_BACKUP_NOTE = ( + "A safety backup is taken first, whatever the safety backup setting: {name} " + "among the pooled backups; the pre-migration backup {kept} is already kept." +) +NOTHING_TO_DO = "Every image record already stores its workflow once." +OLDER_VERSIONS_NOTE = ( + "This version stores the workflows of images it records in the new form, " + "whether or not this job runs, and this job rewrites the rest. An older " + "Prompt Manager version shows those images without their workflow until you " + "update again. The way back to an older version is the earliest " + "pre-workflow-storage backup." +) diff --git a/py/api/__init__.py b/py/api/__init__.py index bc4ac04..62551a9 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -14,6 +14,7 @@ Split into domain-specific mixins for maintainability: - ReviewEventsMixin: the Needs review change stream an open page re-reads on - SystemRoutesMixin: System health check, Overview and usage statistics - MaintenanceRoutesMixin: reviewed maintenance plans (preview, apply) +- WorkflowStorageRoutesMixin: store each workflow once, then Reclaim space - DuplicateFilesMixin: duplicate image files as find and removal jobs """ @@ -49,6 +50,7 @@ from .image_tag_jobs import ImageTagJobsMixin from .thumbnail_jobs import ThumbnailJobsMixin from .system import SystemRoutesMixin from .maintenance import MaintenanceRoutesMixin +from .workflow_storage import WorkflowStorageRoutesMixin from .duplicate_files import DuplicateFilesMixin from .maintenance_plans import PlanStore from .jobs import JobRegistry @@ -291,6 +293,7 @@ class PromptManagerAPI( ThumbnailJobsMixin, SystemRoutesMixin, MaintenanceRoutesMixin, + WorkflowStorageRoutesMixin, DuplicateFilesMixin, LoggingRoutesMixin, AutotagRoutesMixin, diff --git a/py/api/maintenance.py b/py/api/maintenance.py index 031b3d9..dddc114 100644 --- a/py/api/maintenance.py +++ b/py/api/maintenance.py @@ -58,9 +58,11 @@ from .jobs import MAINTENANCE_JOB_KIND, JobBusy from .maintenance_plans import PlanExpired, PlanUsed, ScopeTooLarge try: + 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.workflow_storage import PRE_MIGRATION_LABEL, pre_migration_backups from database.health import readonly_uri from database.maintenance import MAINTENANCE_OPERATIONS, PlanChanged @@ -68,6 +70,7 @@ PREVIEW_BODY_MAX = 4096 # Duplicate-file previews carry the reviewed keepers (py/api/duplicate_files.py). DUPLICATE_PREVIEW_BODY_MAX = 1024 * 1024 DUPLICATE_FILES = "duplicate_files" +WORKFLOW_STORAGE = "workflow_storage" # applied by py/api/workflow_storage.py PREVIEW_OPERATIONS = MAINTENANCE_OPERATIONS + (DUPLICATE_FILES,) # Room left over beyond what a step writes, so a disk is never filled to 0. DISK_MARGIN_BYTES = 16 * 2**20 @@ -77,6 +80,7 @@ MAINTENANCE_TITLES = { "optimize": "Optimize database", "unused_components": "Unused prompt components", "missing_links": "Missing-file links", + "workflow_storage": "Store each workflow once", } _OPERATION_ERROR = "operation must be one of: " + ", ".join(PREVIEW_OPERATIONS) MAINTENANCE_GROUPS = ("changed", "protected", "errors") @@ -89,12 +93,28 @@ PLAN_CHANGED = ( "Review the plan again." ) 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.""" + """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): + """The job line for a failed backup, naming a damaged workflow document.""" + reason = getattr(error, "reason", None) or "" + if reason.startswith("Workflow document "): + return f"{BACKUP_FAILED} {reason}" + return BACKUP_FAILED def _now(): @@ -160,7 +180,7 @@ def _megabytes(size): return f"{size / 2**20:,.1f} MB" -def disk_plan(db_path, *, backup, vacuum): +def disk_plan(db_path, *, backup, vacuum, vacuum_bytes=None): """Whether the disk has room for an apply's writes (blocking). ``backup``: the safety backup, one database-sized file beside the @@ -168,7 +188,9 @@ def disk_plan(db_path, *, backup, vacuum): beside it and builds it in a temporary copy in SQLite's temporary folder (``sqlite_temp_dir``; SQLite's documentation: up to twice the database in free space). Needs - on one device add up. Returns ``(public, short)``: ``public`` is + on one device add up. ``vacuum_bytes``, when given, is what VACUUM + copies (the live pages, for Reclaim space) instead of the whole file. + Returns ``(public, short)``: ``public`` is ``{free_bytes, required_bytes, sufficient}`` for the database folder, ``short`` None or ``(place, what, required, free)`` of the first place that is short, and ``size`` the database's size (0 when nothing is @@ -181,8 +203,10 @@ def disk_plan(db_path, *, backup, vacuum): if backup: needs.append(("database folder", folder, size, "the safety backup")) if vacuum: - needs.append(("database folder", folder, size, "Optimize")) - needs.append(("temporary folder", sqlite_temp_dir(), size, "Optimize")) + copy = size if vacuum_bytes is None else vacuum_bytes + what = "Optimize" if vacuum_bytes is None else "Reclaim space" + needs.append(("database folder", folder, copy, what)) + needs.append(("temporary folder", sqlite_temp_dir(), copy, what)) devices = {} for place, path, needed, what in needs: entry = devices.setdefault(_device(path), {"place": place, "path": path}) @@ -290,9 +314,9 @@ def _remove_quietly(*paths): pass -def _reserve_backup_name(db_path): - """Create and return an unused ``.backup_[_n]`` (O_EXCL).""" - for candidate in _backup_names(db_path): +def _reserve_backup_name(db_path, label="backup"): + """Create and return an unused ``.