From 02b0e6813f46504e633cb94932aba4e29423d9bb Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Sat, 24 Oct 2026 04:34:09 -0700 Subject: [PATCH] feat(database): migrate old workflow rows in batches, then reclaim space The workflow storage migration rewrites old-form image rows (text inline) to document ids, batch by batch: - Each batch is one BEGIN IMMEDIATE transaction of at most 200 rows or 16 MB of text, found by a rowid range search and read one row at a time. - Each text is stored or reused, then read back through the one accessor before its row changes; a text that would not come back identical fails and rolls back the whole batch. - A value that is not text stays inline and is counted as protected. - A crash or a stop leaves committed batches new form and the rest old form; both read the same, and a rerun continues. No media file is opened. - Unreferenced documents are swept at the end. The maintenance preview gains a read-only workflow_storage plan (records to rewrite, inline bytes where SQLite has octet_length, the Reclaim space and older-version notes). reclaim_space sweeps, runs VACUUM and truncates the WAL so the freed pages go back to the disk. --- database/maintenance.py | 6 +- database/operations.py | 2 + database/workflow_storage.py | 321 +++++++++++++++++++++ tests/test_workflow_documents.py | 5 +- tests/test_workflow_migration.py | 459 +++++++++++++++++++++++++++++++ 5 files changed, 789 insertions(+), 4 deletions(-) create mode 100644 database/workflow_storage.py create mode 100644 tests/test_workflow_migration.py diff --git a/database/maintenance.py b/database/maintenance.py index 2e53fbab..76a7abdc 100644 --- a/database/maintenance.py +++ b/database/maintenance.py @@ -49,6 +49,8 @@ MAINTENANCE_OPERATIONS = ( "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) @@ -420,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) @@ -668,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/operations.py b/database/operations.py index 74608a19..aaaf2bd3 100644 --- a/database/operations.py +++ b/database/operations.py @@ -25,6 +25,7 @@ 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, @@ -346,6 +347,7 @@ class PromptDatabase( HealthOperations, MaintenanceOperations, DuplicateCopyOperations, + WorkflowStorageOperations, ): """Database operations class for managing prompts.""" diff --git a/database/workflow_storage.py b/database/workflow_storage.py new file mode 100644 index 00000000..758bf794 --- /dev/null +++ b/database/workflow_storage.py @@ -0,0 +1,321 @@ +"""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. + +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 os +import sqlite3 +from contextlib import closing + +from .workflow_documents import ( + TEXT_COLUMNS, + DocumentMismatch, + read_stored_text, + release_documents, + storable, + store_text, + sweep_documents, +) + +OPERATION = "workflow_storage" +BATCH_ROWS = 200 +BATCH_BYTES = 16 * 2**20 # characters of text rewritten per transaction + +# 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): + """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. + """ + 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) + 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): + """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. + """ + if conn.in_transaction: + raise RuntimeError("The migration needs its own transaction") + 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 + for image_id in ids: + row = conn.execute(ROW_SQL, (image_id,)).fetchone() + changed, protected, row_size = _rewrite_row(conn, image_id, row, verified) + result["rows"] += int(changed) + result["protected"] += protected + result["last_id"] = image_id + size += row_size + if size >= max_bytes: + break + conn.commit() + return result + except BaseException: + conn.rollback() + raise + + +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. Returns ``{rows, + protected, batches, stopped, documents_removed}``. Unreferenced + documents are swept once every row was looked at. + """ + conn = self.model.get_connection() + totals = {"rows": 0, "protected": 0, "batches": 0, "last_id": 0} + stopped = False + while True: + if should_stop is not None and should_stop(): + stopped = True + break + batch = migrate_batch(conn, totals["last_id"], max_rows, max_bytes) + if batch["done"]: + break + 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 not stopped: + with conn: + removed = sweep_documents(conn) + return { + "rows": totals["rows"], + "protected": totals["protected"], + "batches": totals["batches"], + "stopped": stopped, + "documents_removed": removed, + } + + 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 + conn = self.model.get_connection() + with conn: + removed = sweep_documents(conn) + 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 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"), + } + ) + details += [ + {"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", + }, + 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." +) +OLDER_VERSIONS_NOTE = ( + "An older Prompt Manager version shows rewritten images without their " + "workflow. Restoring the safety backup brings back the old form." +) diff --git a/tests/test_workflow_documents.py b/tests/test_workflow_documents.py index e137bb76..011dfbdb 100644 --- a/tests/test_workflow_documents.py +++ b/tests/test_workflow_documents.py @@ -11,6 +11,7 @@ reads both forms, and gives back the exact text that was stored. import json import os import random +import shutil import sqlite3 import sys import tempfile @@ -134,9 +135,7 @@ class DocumentsTestCase(unittest.TestCase): PromptModel.close_all_instances() PromptModel.reset_schema_cache() assert_closed(self, self.tmpdir) - for name in os.listdir(self.tmpdir): - os.unlink(os.path.join(self.tmpdir, name)) - os.rmdir(self.tmpdir) + shutil.rmtree(self.tmpdir) def conn(self): return self.db.model.get_connection() diff --git a/tests/test_workflow_migration.py b/tests/test_workflow_migration.py new file mode 100644 index 00000000..32d51df0 --- /dev/null +++ b/tests/test_workflow_migration.py @@ -0,0 +1,459 @@ +"""The workflow storage migration (database/workflow_storage.py). + +Rewrites old-form rows (text inline) to the new form (a document id) in +batches. Each batch is one transaction, so a crash leaves every row in one +form or the other and readable either way; a rerun continues where the last +committed batch stopped. Rows are read one at a time, the stored text comes +back byte for byte, and no image or video file is ever opened. +""" + +import builtins +import hashlib +import json +import os +import random +import sqlite3 +import tracemalloc +import unittest +from unittest.mock import patch + +from PIL import Image, PngImagePlugin + +from database import workflow_documents as wd +from database import workflow_storage as ws +from database.models import PromptModel + +try: + from tests.test_workflow_documents import ( + DocumentsTestCase, + edge_texts, + random_text, + sample_prompt_graph, + sample_workflow, + ) +except ImportError: # discovered with tests/ as the top-level directory + from test_workflow_documents import ( + DocumentsTestCase, + edge_texts, + random_text, + sample_prompt_graph, + sample_workflow, + ) + + +def _digest(path): + if not os.path.exists(path): + return None + with open(path, "rb") as handle: + return hashlib.sha256(handle.read()).hexdigest() + + +class MigrationTestCase(DocumentsTestCase): + def library(self, count=40, seed=11): + """Old-form rows: duplicated runs, edge texts and random text.""" + rng = random.Random(seed) + prompt_id = self.prompt() + texts = edge_texts() + expected = {} + for index in range(count): + run = index // 4 # four images per run share their texts + if index % 10 == 9: + workflow = texts[index % len(texts)] + graph = random_text(rng) + else: + workflow = json.dumps(sample_workflow(run, nodes=20)) + graph = json.dumps(sample_prompt_graph(run)) + if index % 13 == 5: + graph = None + image_id = self.old_form_image(prompt_id, f"{index}.png", workflow, graph) + expected[image_id] = (workflow, graph) + return prompt_id, expected + + def texts(self): + conn = self.conn() + reader = wd.DocumentReader(conn) + rows = conn.execute( + "SELECT id, workflow_data, prompt_metadata, workflow_document_id, " + "prompt_document_id FROM generated_images ORDER BY id" + ) + result = {} + for row in rows: + data = wd.resolve_image_row(reader, row) + result[row["id"]] = (data["workflow_data"], data["prompt_metadata"]) + return result + + def inline_ids(self): + return { + row[0] + for row in self.conn().execute( + "SELECT id FROM generated_images WHERE workflow_data IS NOT NULL " + "OR prompt_metadata IS NOT NULL" + ) + } + + def payloads(self, prompt_id): + images = sorted(self.db.get_prompt_images(prompt_id), key=lambda i: i["id"]) + for image in images: + image.pop("generation_time", None) + return images + + def migrate(self, **kwargs): + kwargs.setdefault("max_rows", 7) + return self.db.migrate_workflow_storage(**kwargs) + + +class TestMigration(MigrationTestCase): + def test_every_row_is_rewritten_and_reads_back_identical(self): + prompt_id, expected = self.library() + before = self.payloads(prompt_id) + result = self.migrate() + self.assertEqual(self.texts(), expected) + self.assertEqual(self.inline_ids(), set()) + self.assertEqual(self.payloads(prompt_id), before) + distinct = {t for pair in expected.values() for t in pair if t is not None} + self.assertEqual(self.document_count(), len(distinct)) + self.assertEqual(result["rows"], len(expected)) + self.assertGreater(result["batches"], 1) + self.assertFalse(result["stopped"]) + + def test_a_second_run_changes_nothing(self): + _prompt_id, expected = self.library(count=12) + self.migrate() + conn = self.conn() + documents = list(conn.execute("SELECT * FROM workflow_documents ORDER BY id")) + result = self.migrate() + self.assertEqual(result["rows"], 0) + self.assertEqual( + [tuple(r) for r in documents], + [tuple(r) for r in conn.execute("SELECT * FROM workflow_documents")], + ) + self.assertEqual(self.texts(), expected) + + def test_a_crash_mid_migration_leaves_every_row_readable_and_resumes(self): + prompt_id, expected = self.library() + before = self.payloads(prompt_id) + real = wd.store_text + calls = [] + + def crashing(conn, text): + calls.append(text) + if len(calls) == 23: + raise RuntimeError("power cut") + return real(conn, text) + + with patch.object(ws, "store_text", side_effect=crashing): + with self.assertRaises(RuntimeError): + self.migrate() + self.assertFalse(self.conn().in_transaction) + self.assertEqual(self.texts(), expected) + self.assertEqual(self.payloads(prompt_id), before) + remaining = self.inline_ids() + self.assertTrue(0 < len(remaining) < len(expected)) + # The failed batch rolled back whole: what is left is a suffix by id. + self.assertEqual(remaining, {i for i in expected if i >= min(remaining)}) + self.migrate() + self.assertEqual(self.inline_ids(), set()) + self.assertEqual(self.texts(), expected) + + def test_a_crash_inside_a_batch_writes_nothing_of_that_batch(self): + _prompt_id, expected = self.library(count=10) + with patch.object(ws, "store_text", side_effect=RuntimeError("full")): + with self.assertRaises(RuntimeError): + self.migrate(max_rows=50) + 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): + _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) + self.assertEqual(self.inline_ids(), set(expected)) + self.assertEqual(self.texts(), expected) + + def test_stop_between_batches_then_continue(self): + _prompt_id, expected = self.library(count=30) + seen = [] + result = self.migrate(on_batch=seen.append, should_stop=lambda: len(seen) >= 2) + self.assertTrue(result["stopped"]) + self.assertEqual(result["rows"], 14) + self.assertEqual(len(self.inline_ids()), len(expected) - 14) + self.assertEqual(seen[-1]["rows"], 14) + self.migrate() + self.assertEqual(self.inline_ids(), set()) + self.assertEqual(self.texts(), expected) + + def test_values_that_are_not_text_stay_inline(self): + # A TEXT column turns numbers into text; only a blob stays other. + prompt_id = self.prompt() + blob = self.old_form_image(prompt_id, "blob.png", b"\x00raw", "{}") + other = self.old_form_image(prompt_id, "other.png", "42", b"\x01") + result = self.migrate() + self.assertEqual(result["protected"], 2) + conn = self.conn() + rows = { + row[0]: tuple(row[1:]) + for row in conn.execute( + "SELECT id, workflow_data, prompt_metadata, workflow_document_id, " + "prompt_document_id FROM generated_images" + ) + } + self.assertEqual(rows[blob][:2], (b"\x00raw", None)) + self.assertIsNone(rows[blob][2]) + self.assertIsNotNone(rows[blob][3]) + self.assertEqual(rows[other][:2], (None, b"\x01")) + self.assertEqual(self.texts()[other], ("42", b"\x01")) + self.assertEqual(self.migrate()["protected"], 2) # still there, still kept + + def test_a_row_an_older_version_rewrote_inline_takes_the_inline_text(self): + prompt_id = self.prompt() + image = self.db.link_image_to_prompt( + prompt_id, os.path.join(self.tmpdir, "a.png"), {"workflow": {"v": 1}} + ) + conn = self.conn() + with conn: # an older version updates the inline column only + conn.execute( + "UPDATE generated_images SET workflow_data=? WHERE id=?", + ('{"v": 2}', image), + ) + self.assertEqual(self.texts()[image][0], '{"v": 2}') + self.migrate() + self.assertEqual(self.texts()[image][0], '{"v": 2}') + self.assertEqual(self.inline_ids(), set()) + stored = {row[0] for row in conn.execute("SELECT data FROM workflow_documents")} + self.assertNotIn(b'{"v": 1}', {bytes(wd_data) for wd_data in stored}) + texts = [ + wd.read_stored_text(conn, None, row[0]) + for row in conn.execute("SELECT id FROM workflow_documents") + ] + self.assertNotIn('{"v": 1}', texts) + + def test_new_writes_during_the_migration_are_kept(self): + prompt_id, expected = self.library(count=20) + added = [] + + def between(_progress): + if not added: + added.append( + self.db.link_image_to_prompt( + prompt_id, + os.path.join(self.tmpdir, "new.png"), + {"workflow": {"new": True}}, + ) + ) + added.append( + self.old_form_image(prompt_id, "late.png", '{"late": 1}', None) + ) + + self.migrate(on_batch=between) + texts = self.texts() + self.assertEqual(texts[added[0]][0], '{"new": true}') + self.assertEqual(texts[added[1]], ('{"late": 1}', None)) + self.assertEqual(self.inline_ids(), set()) + for image_id, pair in expected.items(): + self.assertEqual(texts[image_id], pair) + + def test_memory_stays_bounded_by_one_row(self): + prompt_id = self.prompt() + big = json.dumps(sample_workflow(1, nodes=3000)) # about 0.9 MB + for index in range(40): + self.old_form_image( + prompt_id, f"{index}.png", big[:-1] + f', "n": {index}}}', None + ) + tracemalloc.start() + try: + self.migrate(max_rows=200, max_bytes=4 * 2**20) + _current, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + self.assertLess(peak, 12 * 2**20, f"peak {peak / 2**20:.1f} MB") + self.assertEqual(self.inline_ids(), set()) + + def test_batches_are_bounded_by_bytes(self): + prompt_id = self.prompt() + text = "x" * 100_000 + for index in range(10): + self.old_form_image(prompt_id, f"{index}.png", text + str(index), None) + seen = [] + self.migrate(max_rows=200, max_bytes=250_000, on_batch=seen.append) + self.assertEqual([batch["rows"] for batch in seen], [3, 6, 9, 10]) + + def test_unreferenced_documents_are_swept_at_the_end(self): + prompt_id = self.prompt() + image = self.db.link_image_to_prompt( + prompt_id, os.path.join(self.tmpdir, "a.png"), {"workflow": {"v": 1}} + ) + conn = self.conn() + with conn: + conn.execute("DELETE FROM generated_images WHERE id=?", (image,)) + result = self.migrate() + self.assertEqual(result["documents_removed"], 2) + self.assertEqual(self.document_count(), 0) + + def test_no_media_file_is_opened_or_changed(self): + media = os.path.join(self.tmpdir, "output") + os.makedirs(media) + prompt_id = self.prompt() + for index in range(5): + path = os.path.join(media, f"{index}.png") + info = PngImagePlugin.PngInfo() + info.add_text("workflow", json.dumps(sample_workflow(index))) + with Image.new("RGB", (4, 4)) as image: + image.save(path, pnginfo=info) + conn = self.conn() + with conn: + conn.execute( + "INSERT INTO generated_images (prompt_id, image_path, filename, " + "file_path, workflow_data) VALUES (?, ?, ?, ?, ?)", + (prompt_id, path, f"{index}.png", path, '{"n": 1}'), + ) + before = { + name: ( + _digest(os.path.join(media, name)), + os.stat(os.path.join(media, name)), + ) + for name in os.listdir(media) + } + opened = [] + real_open, real_os_open = builtins.open, os.open + + def tracking_open(file, *args, **kwargs): + opened.append(str(file)) + return real_open(file, *args, **kwargs) + + def tracking_os_open(path, *args, **kwargs): + opened.append(str(path)) + return real_os_open(path, *args, **kwargs) + + with patch("builtins.open", tracking_open), patch("os.open", tracking_os_open): + with patch.object(Image, "open", side_effect=AssertionError("opened")): + self.migrate() + self.assertEqual([p for p in opened if p.startswith(media)], []) + for name, (digest, stat) in before.items(): + now = os.stat(os.path.join(media, name)) + self.assertEqual(_digest(os.path.join(media, name)), digest) + self.assertEqual(now.st_mtime_ns, stat.st_mtime_ns) + self.assertEqual(self.inline_ids(), set()) + + +class TestPreview(MigrationTestCase): + def test_preview_counts_and_writes_nothing(self): + prompt_id, expected = self.library(count=20) + self.old_form_image(prompt_id, "blob.png", b"raw", None) + files = [self.db_path + suffix for suffix in ("", "-wal")] + before = [_digest(path) for path in files] + draft = self.db.preview_maintenance("workflow_storage") + self.assertEqual([_digest(path) for path in files], before) + self.assertEqual(draft.operation, "workflow_storage") + self.assertEqual(draft.scope["count"], len(expected)) + self.assertEqual(draft.protected["count"], 1) + self.assertFalse(draft.media["deletes_files"]) + self.assertIn("No image or video file", draft.media["label"]) + labels = [row["label"] for row in draft.details] + self.assertIn("Database file", labels) + self.assertIn("Older versions", labels) + self.assertIn("Reclaim space", labels) + + def test_preview_of_a_migrated_library_has_nothing_to_do(self): + self.library(count=5) + self.migrate() + draft = self.db.preview_maintenance("workflow_storage") + self.assertEqual(draft.scope["count"], 0) + + def test_the_batch_reads_are_index_searches(self): + conn = self.conn() + for name, (sql, params) in ws.QUERY_PLAN_STATEMENTS.items(): + plan = " | ".join( + row[3] for row in conn.execute("EXPLAIN QUERY PLAN " + sql, params) + ) + with self.subTest(name=name): + self.assertIn("USING INTEGER PRIMARY KEY", plan, plan) + + +class TestReclaimSpace(MigrationTestCase): + def test_reclaim_space_gives_the_freed_pages_back(self): + prompt_id = self.prompt() + big = json.dumps(sample_workflow(1, nodes=600)) + for index in range(60): + self.old_form_image(prompt_id, f"{index}.png", big, big) + expected = self.texts() + self.migrate() + conn = self.conn() + with conn: # a document a path forgot to release + conn.execute( + "INSERT INTO workflow_documents (sha256, byte_length, data) " + "VALUES (zeroblob(32), 0, x'789c030000000001')" + ) + result = self.db.reclaim_space() + self.assertEqual(result["documents_removed"], 1) + self.assertFalse(result["checkpoint_busy"]) + self.assertLess(result["after_bytes"], result["before_bytes"] / 4) + wal = self.db_path + "-wal" + self.assertEqual(os.path.getsize(wal) if os.path.exists(wal) else 0, 0) + self.assertEqual(self.texts(), expected) + + +class TestRestore(MigrationTestCase): + def test_restoring_a_pre_dedupe_backup_reads_and_migrates(self): + prompt_id, expected = self.library(count=8) + backup = os.path.join(self.tmpdir, "old-backup.db") + conn = sqlite3.connect(backup) + try: # as a version before the dedupe wrote it: no table, no columns + conn.executescript( + "CREATE TABLE prompts (id INTEGER PRIMARY KEY AUTOINCREMENT, " + "text TEXT NOT NULL, created_at TIMESTAMP, updated_at TIMESTAMP, " + "category TEXT, tags TEXT, rating INTEGER, notes TEXT, " + "hash TEXT UNIQUE);" + "CREATE TABLE generated_images (id INTEGER PRIMARY KEY " + "AUTOINCREMENT, prompt_id INTEGER NOT NULL, image_path TEXT NOT " + "NULL, filename TEXT NOT NULL, generation_time TIMESTAMP, " + "file_size INTEGER, width INTEGER, height INTEGER, format TEXT, " + "workflow_data TEXT, prompt_metadata TEXT, parameters TEXT, " + "file_path TEXT, FOREIGN KEY (prompt_id) REFERENCES prompts(id) " + "ON DELETE CASCADE, UNIQUE(prompt_id, file_path));" + ) + conn.execute( + "INSERT INTO prompts (id, text, hash) VALUES (?, 'a cat', 'h')", + (prompt_id,), + ) + for image_id, (workflow, graph) in expected.items(): + conn.execute( + "INSERT INTO generated_images (id, prompt_id, image_path, " + "filename, file_path, workflow_data, prompt_metadata) " + "VALUES (?, ?, ?, ?, ?, ?, ?)", + ( + image_id, + prompt_id, + f"/x/{image_id}", + "x", + f"/x/{image_id}", + workflow, + graph, + ), + ) + conn.commit() + finally: + conn.close() + self.migrate() # the live library is new form now + self.db.model.restore_from_file(backup) + self.assertEqual(self.texts(), expected) + self.assertEqual(self.inline_ids(), set(expected)) + self.migrate() + self.assertEqual(self.texts(), expected) + self.assertEqual(self.inline_ids(), set()) + + def test_restoring_a_deduped_backup_reads_the_same(self): + prompt_id, expected = self.library(count=8) + self.migrate() + backup = os.path.join(self.tmpdir, "new-backup.db") + self.assertTrue(self.db.model.backup_database(backup)) + ok, reason = PromptModel.verify_database_file(backup) + self.assertTrue(ok, reason) + self.db.delete_prompt(prompt_id) + self.assertEqual(self.document_count(), 0) + self.db.model.restore_from_file(backup) + self.assertEqual(self.texts(), expected) + + +if __name__ == "__main__": + unittest.main()