diff --git a/database/workflow_documents.py b/database/workflow_documents.py index 09b36937..7272b682 100644 --- a/database/workflow_documents.py +++ b/database/workflow_documents.py @@ -66,12 +66,6 @@ _DDL = ( "byte_length INTEGER NOT NULL, data BLOB NOT NULL)" ) _AUTOINDEX = f"sqlite_autoindex_{TABLE}_1" -_COLUMNS = ( - ("id", "INTEGER", 0, 1), - ("sha256", "BLOB", 1, 0), - ("byte_length", "INTEGER", 1, 0), - ("data", "BLOB", 1, 0), -) _REFERENCE_TYPE = f"INTEGER REFERENCES {TABLE}(id)" _UNREFERENCED = ( @@ -202,14 +196,9 @@ def workflow_documents_schema_error( 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" - shape = tuple( - (row[1], str(row[2]).upper(), row[3], row[5]) - for row in conn.execute(f"PRAGMA table_info({TABLE})") - ) - if shape != _COLUMNS: - return "Workflow documents have invalid columns" return _reference_error(conn, columns) or _value_error(conn) diff --git a/database/workflow_storage.py b/database/workflow_storage.py index 758bf794..3e1ba3de 100644 --- a/database/workflow_storage.py +++ b/database/workflow_storage.py @@ -167,14 +167,13 @@ class WorkflowStorageOperations: 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) + batch = self.migrate_workflow_batch(totals["last_id"], max_rows, max_bytes) if batch["done"]: break totals["rows"] += batch["rows"] @@ -183,10 +182,7 @@ class WorkflowStorageOperations: 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) + removed = 0 if stopped else self.sweep_workflow_documents() return { "rows": totals["rows"], "protected": totals["protected"], @@ -195,6 +191,24 @@ class WorkflowStorageOperations: "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). @@ -220,9 +234,7 @@ class WorkflowStorageOperations: stopped VACUUM; the database is then unchanged. """ db_path = self.model.db_path - conn = self.model.get_connection() - with conn: - removed = sweep_documents(conn) + removed = self.sweep_workflow_documents() before = _file_bytes(db_path) self.model.vacuum_database() with closing(sqlite3.connect(db_path)) as checkpoint: diff --git a/py/api/__init__.py b/py/api/__init__.py index bc4ac049..62551a93 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 62ac5853..dad45dda 100644 --- a/py/api/maintenance.py +++ b/py/api/maintenance.py @@ -68,6 +68,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 +78,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") @@ -160,7 +162,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 +170,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 +185,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}) @@ -567,6 +573,8 @@ class MaintenanceRoutesMixin: try: if plan.operation == DUPLICATE_FILES: job = self._start_duplicate_cleanup(db, plan, backup) + elif plan.operation == WORKFLOW_STORAGE: + job = self._start_workflow_storage(db, backup) else: job = self._jobs.start( MAINTENANCE_JOB_KIND, diff --git a/py/api/workflow_storage.py b/py/api/workflow_storage.py new file mode 100644 index 00000000..9de194a0 --- /dev/null +++ b/py/api/workflow_storage.py @@ -0,0 +1,210 @@ +"""Store each workflow once: the migration and Reclaim space jobs (#249 S6). + +Reviewed like every maintenance operation (``POST /prompt_manager/ +maintenance/preview {operation:"workflow_storage"}``, a ``Plan`` with the +records to rewrite, the backup and the disk check) and applied with its token +(``POST /prompt_manager/maintenance/apply``). The apply starts: + +1. **Store each workflow once**: kind ``maintenance``, admission ``records``, + cancellable, groups ``changed``/``protected``/``errors``. It takes the + promised verified safety backup (``backup-failed`` ends it with nothing + rewritten), then rewrites old-form rows batch by batch + (``database/workflow_storage.py``): one transaction per batch, progress + after each, Cancel honoured between batches. A failed batch rolls back + whole and fails the job; every record still reads as before. Stopped or + failed, the next apply continues from the first old-form row. +2. **Reclaim space**, once every record was looked at: kind ``maintenance``, + admission ``exclusive`` (it runs alone; VACUUM holds the write lock), not + cancellable. It checks the disk for VACUUM's copies of the live pages + (``disk_plan``), then sweeps, runs VACUUM and truncates the WAL. When + another job holds the registry, it does not start and the finished + migration says so and names Optimize database. + +Only the database changes: no image or video file is opened. +""" + +import sqlite3 + +from .jobs import MAINTENANCE_JOB_KIND, JobBusy +from .maintenance import ( + MAINTENANCE_GROUPS, + _fail_job, + _megabytes, + backup_note, + disk_plan, + short_space_message, +) + +try: + from ...database import workflow_storage + from ...database.health import readonly_uri +except ImportError: + from database import workflow_storage + from database.health import readonly_uri + +OPERATION = workflow_storage.OPERATION +TITLE = "Store each workflow once" +RECLAIM_TITLE = "Reclaim space" +REWRITE_FAILED = ( + "Storing workflows once stopped on an error. The batch in progress was " + "rolled back; every record reads as before. Run it again to continue." +) +RECLAIM_FAILED = "Reclaim space failed; the database is unchanged." + + +def _plural(count, one, many=None): + return f"{count:,} {one if count == 1 else (many or one + 's')}" + + +def live_bytes(db_path): + """The pages VACUUM copies: page_count minus the free list (blocking).""" + conn = sqlite3.connect(readonly_uri(db_path), uri=True) + try: + page_size = conn.execute("PRAGMA page_size").fetchone()[0] + pages = conn.execute("PRAGMA page_count").fetchone()[0] + free = conn.execute("PRAGMA freelist_count").fetchone()[0] + return page_size * (pages - free) + finally: + conn.close() + + +class WorkflowStorageRoutesMixin: + """Mixin for the maintenance routes: the workflow storage jobs.""" + + def _start_workflow_storage(self, db, backup): + """Start the migration job (raises JobBusy); the token is the caller's.""" + return self._jobs.start( + MAINTENANCE_JOB_KIND, + TITLE, + admission="records", + cancellable=True, + scope={"operation": OPERATION}, + groups=MAINTENANCE_GROUPS, + runner=lambda job: self._run_workflow_storage(job, db, backup), + ) + + async def _run_workflow_storage(self, job, db, backup_on): + ok, backup = await self._safety_backup_step(job, db, backup_on) + if not ok: + return + try: + done = await self._rewrite_workflows(job, db, backup) + finally: + if backup: + await self._prune_backup_pool(job, db) + if done: + job.finish("finished") + self._start_reclaim_space(job, db) + + async def _rewrite_workflows(self, job, db, backup): + """Rewrite every old-form row; True when every row was looked at.""" + counts = await self._run_in_executor(db.workflow_storage_status) + job.update(total=counts["rows"], processed=0) + job.update(status_line="Rewriting image records") + job.log_line(f"{_plural(counts['rows'], 'image record')} to rewrite") + rows = protected = 0 + last_id = 0 + while True: + if job.cancelled: + return self._stopped(job, rows, backup) + try: + batch = await self._run_in_executor( + db.migrate_workflow_batch, + last_id, + workflow_storage.BATCH_ROWS, + workflow_storage.BATCH_BYTES, + ) + except Exception: + self.logger.exception("Storing workflows once failed") + _fail_job(job, None, REWRITE_FAILED) + return False + if batch["done"]: + break + last_id = batch["last_id"] + rows += batch["rows"] + protected += batch["protected"] + job.record_counted("changed", batch["rows"]) + job.record_counted("protected", batch["protected"]) + total = max(job.total, rows) + job.update( + processed=rows, + total=total, + progress=100 * rows // total if total else 100, + status_line=f"Rewrote {_plural(rows, 'image record')}", + ) + job.publish({"type": "progress", "processed": rows, "total": total}) + removed = await self._run_in_executor(db.sweep_workflow_documents) + line = f"Rewrote {_plural(rows, 'image record')}" + if protected: + line += f"; {_plural(protected, 'stored value')} not text, kept as is" + job.log_line(line) + if removed: + job.log_line(f"Removed {_plural(removed, 'unused workflow document')}") + job.update(progress=100, status_line=f"{line} · {backup_note(backup)}") + return True + + def _stopped(self, job, rows, backup): + line = ( + f"Stopped after {_plural(rows, 'image record')}; " + "run it again to continue" + ) + job.log_line(line) + job.update(status_line=f"{line} · {backup_note(backup)}") + job.finish("cancelled") + return False + + def _start_reclaim_space(self, finished, db): + """Start Reclaim space after the migration ``finished``, or say why not.""" + try: + self._jobs.start( + MAINTENANCE_JOB_KIND, + RECLAIM_TITLE, + admission="exclusive", + cancellable=False, + scope={"operation": OPERATION}, + groups=MAINTENANCE_GROUPS, + runner=lambda job: self._run_reclaim_space(job, db), + ) + except JobBusy as busy: + line = ( + f"Reclaim space did not start: {busy.running.title} is running. " + "Run Optimize database to give the space back." + ) + finished.log_line(line) + finished.update(status_line=f"{finished.status_line} · {line}") + return + finished.log_line("Reclaim space started") + + async def _run_reclaim_space(self, job, db): + job.update(status_line="Checking free disk space") + try: + live = await self._run_in_executor(live_bytes, db.model.db_path) + _disk, short, _size = await self._run_in_executor( + disk_plan, + db.model.db_path, + backup=False, + vacuum=True, + vacuum_bytes=live, + ) + except (OSError, sqlite3.Error): + self.logger.exception("Could not measure the disk space for Reclaim space") + return _fail_job(job, None, RECLAIM_FAILED) + if short is not None: + return _fail_job(job, "insufficient-space", short_space_message(short)) + job.update(status_line="Reclaiming space (the database is locked)") + job.log_line("Rebuilding the database file (VACUUM)") + try: + result = await self._run_in_executor(db.reclaim_space) + except Exception: + self.logger.exception("Reclaim space failed") + return _fail_job(job, None, RECLAIM_FAILED) + saved = max(0, result["before_bytes"] - result["after_bytes"]) + line = ( + f"Reclaimed {_megabytes(saved)} · the database is now " + f"{_megabytes(result['after_bytes'])}" + ) + if result["checkpoint_busy"]: + line += " (the rest after the next restart)" + job.record_counted("changed", 1) + job.log_line(line) + job.update(progress=100, status_line=line) diff --git a/tests/test_api_workflow_storage.py b/tests/test_api_workflow_storage.py new file mode 100644 index 00000000..17338d63 --- /dev/null +++ b/tests/test_api_workflow_storage.py @@ -0,0 +1,413 @@ +"""Store each workflow once, as a maintenance job (py/api/workflow_storage.py). + +``POST /prompt_manager/maintenance/preview {operation:"workflow_storage"}`` +reviews the rewrite; applying it starts a ``records`` job that takes the +verified safety backup first, rewrites old-form rows batch by batch (Cancel +stops after the current batch) and then starts "Reclaim space", an +``exclusive`` job with its own disk check. No image file is ever opened. +""" + +import asyncio +import builtins +import contextlib +import hashlib +import json +import os +import sqlite3 +from unittest.mock import patch + +from PIL import Image, PngImagePlugin + +from database import workflow_storage as ws +from database.workflow_documents import DocumentReader, resolve_image_row +from py.api import maintenance +from py.api.jobs import MAINTENANCE_JOB_KIND + +try: + from tests.test_api_maintenance import ApplyTestCase + from tests.test_workflow_documents import sample_prompt_graph, sample_workflow +except ImportError: # discovered with tests/ as the top-level directory + from test_api_maintenance import ApplyTestCase + from test_workflow_documents import sample_prompt_graph, sample_workflow + +WAIT = 20 +OPERATION = "workflow_storage" + + +def _digest(path): + with open(path, "rb") as handle: + return hashlib.sha256(handle.read()).hexdigest() + + +class WorkflowStorageTestCase(ApplyTestCase): + def old_library(self, count=12, media=False): + """Old-form rows (text inline), four images per run.""" + prompt_id = self.save("a cat") + conn = self.api.db.model.get_connection() + for index in range(count): + run = index // 4 + path = self.output_dir / f"{index}.png" + workflow = json.dumps(sample_workflow(run, nodes=120)) + if media: + info = PngImagePlugin.PngInfo() + info.add_text("workflow", workflow) + with Image.new("RGB", (4, 4)) as image: + image.save(path, pnginfo=info) + with conn: + conn.execute( + "INSERT INTO generated_images (prompt_id, image_path, filename, " + "file_path, workflow_data, prompt_metadata, parameters) " + "VALUES (?, ?, ?, ?, ?, ?, '{}')", + ( + prompt_id, + str(path), + path.name, + str(path), + workflow, + json.dumps(sample_prompt_graph(run)), + ), + ) + return prompt_id + + def texts(self): + conn = self.api.db.model.get_connection() + reader = DocumentReader(conn) + return { + row["id"]: resolve_image_row(reader, row) + for row in conn.execute( + "SELECT id, workflow_data, prompt_metadata, workflow_document_id, " + "prompt_document_id FROM generated_images ORDER BY id" + ) + } + + def inline_count(self): + conn = self.api.db.model.get_connection() + return conn.execute( + "SELECT COUNT(*) FROM generated_images WHERE workflow_data IS NOT NULL" + ).fetchone()[0] + + def reclaim_job(self): + jobs = [job for job in self.api._jobs.list() if job.title == "Reclaim space"] + return jobs[0] if jobs else None + + async def finish(self, job): + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + reclaim = self.reclaim_job() + if reclaim is not None and reclaim.task is not None: + await asyncio.wait_for(asyncio.shield(reclaim.task), WAIT) + return reclaim + + async def start(self): + plan = await self.preview(OPERATION) + body = await self.apply(plan["token"]) + return plan, body["job"], self.api._jobs.get(body["job"]["id"]) + + +class TestWorkflowStorageJob(WorkflowStorageTestCase): + async def test_preview_reviews_the_rewrite_and_writes_nothing(self): + self.old_library() + db_path = self.api.db.model.db_path + before = _digest(db_path), _digest(db_path + "-wal") + plan = await self.preview(OPERATION) + self.assertEqual((_digest(db_path), _digest(db_path + "-wal")), before) + self.assertEqual(plan["operation"], OPERATION) + self.assertEqual(plan["scope"]["count"], 12) + self.assertEqual(plan["scope"]["label"], "12 image records to rewrite") + self.assertFalse(plan["media"]["deletes_files"]) + self.assertEqual( + plan["media"]["label"], + "No image or video file is opened or changed; only the database.", + ) + self.assertEqual(plan["backup"]["policy"], "on") + self.assertGreater(plan["backup"]["expected_bytes"], 0) + self.assertTrue(plan["disk"]["sufficient"]) + + async def test_apply_backs_up_rewrites_then_reclaims_space(self): + prompt_id = self.old_library(count=40) + before = self.texts() + listed = await self.client.request( + "GET", f"/prompt_manager/prompts/{prompt_id}/images" + ) + listed_before = await listed.json() + size_before = self.database_bytes() + plan, summary, job = await self.start() + self.assertEqual(summary["kind"], MAINTENANCE_JOB_KIND) + self.assertEqual(summary["title"], "Store each workflow once") + self.assertEqual(summary["admission"], "records") + self.assertTrue(summary["cancellable"]) + self.assertEqual(summary["scope"], {"operation": OPERATION}) + reclaim = await self.finish(job) + self.assertEqual(job.state, "finished", job.log) + self.assertEqual(job.counters, {"changed": 40, "protected": 0, "errors": 0}) + self.assertEqual(job.total, 40) + self.assertEqual(job.processed, 40) + self.assertEqual(self.inline_count(), 0) + self.assertEqual(self.texts(), before) + [backup] = self.backups() + self.assertIn(f"Safety backup {backup} verified", "\n".join(job.log)) + self.assertTrue(job.status_line.startswith("Rewrote 40 image records")) + # The safety backup holds the old form, verified. + conn = sqlite3.connect(os.path.join(self.tmpdir, backup)) + try: + inline = conn.execute( + "SELECT COUNT(*) FROM generated_images WHERE workflow_data IS NOT NULL" + ).fetchone()[0] + finally: + conn.close() + self.assertEqual(inline, 40) + # Reclaim space ran alone afterwards and gave the space back. + self.assertIsNotNone(reclaim) + self.assertEqual(reclaim.admission, "exclusive") + self.assertFalse(reclaim.cancellable) + self.assertEqual(reclaim.state, "finished", reclaim.log) + self.assertIn("Reclaim space started", "\n".join(job.log)) + self.assertLess(self.database_bytes(), size_before / 3) + wal = self.api.db.model.db_path + "-wal" + self.assertEqual(os.path.getsize(wal) if os.path.exists(wal) else 0, 0) + self.assertRegex(reclaim.status_line, r"^Reclaimed [\d.,]+ MB · ") + listed = await self.client.request( + "GET", f"/prompt_manager/prompts/{prompt_id}/images" + ) + self.assertEqual(await listed.json(), listed_before) + self.assertEqual(self.strays(), []) + self.assertNotIn(self.tmpdir, json.dumps(job.report())) + + async def test_a_backup_failure_rewrites_nothing(self): + self.old_library() + with patch.object( + maintenance, + "take_safety_backup", + side_effect=maintenance.SafetyBackupFailed("no"), + ): + _plan, _summary, job = await self.start() + await self.finish(job) + self.assertEqual(job.state, "failed") + self.assertEqual(job.last_event["code"], "backup-failed") + self.assertEqual(self.inline_count(), 12) + self.assertIsNone(self.reclaim_job()) + + async def test_cancel_stops_after_the_current_batch_and_a_rerun_continues(self): + self.old_library() + before = self.texts() + real = self.api.db.migrate_workflow_batch + holder = {} + + def first_then_cancel(*args, **kwargs): + result = real(*args, **kwargs) + holder["job"].cancelled = True # Cancel pressed during batch one + return result + + with ( + patch.object(ws, "BATCH_ROWS", 5), + patch.object( + self.api.db, "migrate_workflow_batch", side_effect=first_then_cancel + ), + ): + plan = await self.preview(OPERATION) + body = await self.apply(plan["token"]) + holder["job"] = job = self.api._jobs.get(body["job"]["id"]) + await self.finish(job) + self.assertEqual(job.state, "cancelled") + self.assertEqual(job.counters["changed"], 5) + self.assertEqual(self.inline_count(), 7) + self.assertIn("run it again to continue", job.status_line) + self.assertIsNone(self.reclaim_job()) + self.assertEqual(self.texts(), before) + _plan, _summary, again = await self.start() + await self.finish(again) + self.assertEqual(again.state, "finished", again.log) + self.assertEqual(again.counters["changed"], 7) + self.assertEqual(self.inline_count(), 0) + self.assertEqual(self.texts(), before) + + async def test_a_failed_batch_fails_the_job_and_every_record_still_reads(self): + self.old_library() + before = self.texts() + real = self.api.db.migrate_workflow_batch + calls = [] + + def failing(*args, **kwargs): + calls.append(1) + if len(calls) == 2: + raise sqlite3.OperationalError("disk I/O error") + return real(*args, **kwargs) + + with ( + patch.object(ws, "BATCH_ROWS", 5), + patch.object(self.api.db, "migrate_workflow_batch", side_effect=failing), + ): + _plan, _summary, job = await self.start() + await self.finish(job) + self.assertEqual(job.state, "failed") + self.assertIn( + "The batch in progress was rolled back; every record reads as before.", + job.status_line, + ) + self.assertNotIn("disk I/O", job.status_line) + self.assertEqual(job.counters["changed"], 5) + self.assertEqual(job.counters["errors"], 1) + self.assertEqual(self.texts(), before) + self.assertIsNone(self.reclaim_job()) + + async def test_reclaim_space_that_cannot_start_says_why(self): + self.old_library() + release = asyncio.Event() + + async def held(job): + await release.wait() + + reader = self.api._jobs.start( + "health", "Health check", admission="read", runner=held + ) + try: + _plan, _summary, job = await self.start() + await self.finish(job) + finally: + release.set() + await asyncio.wait_for(reader.task, WAIT) + self.assertEqual(job.state, "finished", job.log) + self.assertIsNone(self.reclaim_job()) + self.assertIn( + "Reclaim space did not start: Health check is running. " + "Run Optimize database to give the space back.", + "\n".join(job.log), + ) + self.assertIn("Reclaim space did not start", job.status_line) + + async def test_reclaim_space_checks_the_disk_first(self): + self.old_library() + before = self.texts() + _plan, _summary, job = await self.start() + free = {"bytes": 10**12} + + def free_bytes(_path): + return free["bytes"] + + with patch.object(maintenance, "_free_bytes", side_effect=free_bytes): + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + reclaim = self.reclaim_job() + free["bytes"] = 1024 + # The reclaim job checks when it runs, after the rewrite. + await asyncio.wait_for(asyncio.shield(reclaim.task), WAIT) + self.assertEqual(reclaim.state, "failed") + self.assertEqual(reclaim.last_event["code"], "insufficient-space") + self.assertIn("Not enough free disk space", reclaim.status_line) + self.assertEqual(self.texts(), before) + + async def test_apply_refuses_up_front_when_the_backup_does_not_fit(self): + self.old_library() + plan = await self.preview(OPERATION) + with patch.object(maintenance, "_free_bytes", return_value=1024): + body = await self.apply(plan["token"], status=507) + self.assertEqual(body["code"], "insufficient-space") + self.assertEqual(self.inline_count(), 12) + + async def test_the_jobs_never_open_a_media_file(self): + self.old_library(media=True) + media = str(self.output_dir) + 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")): + _plan, _summary, job = await self.start() + reclaim = await self.finish(job) + self.assertEqual(job.state, "finished", job.log) + self.assertEqual(reclaim.state, "finished", reclaim.log) + self.assertEqual([path for path in opened if path.startswith(media)], []) + for name, (digest, stat) in before.items(): + path = os.path.join(media, name) + self.assertEqual(_digest(path), digest) + self.assertEqual(os.stat(path).st_mtime_ns, stat.st_mtime_ns) + + async def test_a_second_apply_of_the_token_is_plan_used(self): + self.old_library(count=4) + plan = await self.preview(OPERATION) + body = await self.apply(plan["token"]) + await self.finish(self.api._jobs.get(body["job"]["id"])) + again = await self.apply(plan["token"], status=409) + self.assertEqual(again["code"], "plan-used") + + +class TestWorkflowStorageOutcomes(WorkflowStorageTestCase): + async def test_kept_values_and_swept_documents_are_logged(self): + prompt_id = self.old_library(count=4) + conn = self.api.db.model.get_connection() + with conn: + conn.execute( + "INSERT INTO generated_images (prompt_id, image_path, filename, " + "file_path, workflow_data) VALUES (?, '/x/b', 'b', '/x/b', x'00')", + (prompt_id,), + ) + conn.execute( + "INSERT INTO workflow_documents (sha256, byte_length, data) " + "VALUES (zeroblob(32), 0, x'789c030000000001')" + ) + _plan, _summary, job = await self.start() + await self.finish(job) + log = "\n".join(job.log) + self.assertEqual(job.counters["protected"], 1) + self.assertIn( + "Rewrote 4 image records; 1 stored value not text, kept as is", log + ) + self.assertIn("Removed 1 unused workflow document", log) + + async def reclaim_with(self, **patches): + self.old_library(count=4) + before = self.texts() + with contextlib.ExitStack() as stack: + for name, value in patches.items(): + stack.enter_context(patch.object(self.api.db, name, value)) + _plan, _summary, job = await self.start() + reclaim = await self.finish(job) + self.assertEqual(self.texts(), before) + return reclaim + + async def test_a_failed_reclaim_says_the_database_is_unchanged(self): + reclaim = await self.reclaim_with( + reclaim_space=lambda: (_ for _ in ()).throw( + sqlite3.OperationalError("database is locked") + ) + ) + self.assertEqual(reclaim.state, "failed") + self.assertEqual( + reclaim.status_line, "Reclaim space failed; the database is unchanged." + ) + + async def test_a_reclaim_whose_disk_cannot_be_measured_fails_cleanly(self): + from py.api import workflow_storage as api_storage + + with patch.object(api_storage, "live_bytes", side_effect=OSError("gone")): + reclaim = await self.reclaim_with() + self.assertEqual(reclaim.state, "failed") + self.assertNotIn("gone", reclaim.status_line) + + async def test_a_busy_checkpoint_says_the_rest_follows_a_restart(self): + reclaim = await self.reclaim_with( + reclaim_space=lambda: { + "before_bytes": 3 * 2**20, + "after_bytes": 2**20, + "documents_removed": 0, + "checkpoint_busy": True, + } + ) + self.assertEqual( + reclaim.status_line, + "Reclaimed 2.0 MB · the database is now 1.0 MB " + "(the rest after the next restart)", + ) diff --git a/tests/test_workflow_documents.py b/tests/test_workflow_documents.py index 011dfbdb..bd60a5c0 100644 --- a/tests/test_workflow_documents.py +++ b/tests/test_workflow_documents.py @@ -17,6 +17,7 @@ import sys import tempfile import unittest import zlib +from unittest.mock import patch sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) @@ -402,6 +403,152 @@ class TestSchema(DocumentsTestCase): self.assertIn("orkflow", reason) +class TestSchemaErrors(unittest.TestCase): + """Each way a restored file's document schema can be wrong is named.""" + + IMAGES = ( + "CREATE TABLE generated_images (id INTEGER PRIMARY KEY, " + "workflow_data TEXT, prompt_metadata TEXT{extra})" + ) + TABLE = ( + "CREATE TABLE workflow_documents (id INTEGER PRIMARY KEY, sha256 BLOB NOT " + "NULL UNIQUE, byte_length INTEGER NOT NULL, data BLOB NOT NULL{extra})" + ) + REFS = ( + ", workflow_document_id {kind}, prompt_document_id " + "INTEGER REFERENCES workflow_documents(id)" + ) + GOOD_REF = "INTEGER REFERENCES workflow_documents(id)" + + def schema(self, images_extra="", table_extra="", indexes=True, table=True): + conn = sqlite3.connect(":memory:") + self.addCleanup(conn.close) + conn.execute(self.IMAGES.format(extra=images_extra)) + if table: + conn.execute(self.TABLE.format(extra=table_extra)) + if indexes and images_extra: + for column, index in wd.INDEXES.items(): + conn.execute(f"CREATE INDEX {index} ON generated_images({column})") + return conn + + def test_each_error_is_named(self): + refs = self.REFS.format(kind=self.GOOD_REF) + cases = { + "Missing workflow document schema": self.schema(table=False), + "Incomplete workflow document schema": self.schema(), + "Workflow documents have an invalid table definition": self.schema( + refs, table_extra=", extra TEXT" + ), + "Workflow document references have an invalid column": self.schema( + self.REFS.format(kind="TEXT REFERENCES workflow_documents(id)") + ), + "Workflow document references have invalid foreign keys": self.schema( + self.REFS.format(kind="INTEGER") + ), + "Workflow document references are missing their index": self.schema( + refs, indexes=False + ), + } + for message, conn in cases.items(): + with self.subTest(message=message): + self.assertEqual( + wd.workflow_documents_schema_error(conn, allow_missing=False), + message, + ) + self.assertIsNone(wd.workflow_documents_schema_error(self.schema(table=False))) + self.assertIsNone( + wd.workflow_documents_schema_error(self.schema(refs), allow_missing=False) + ) + + def test_an_index_on_another_table_is_not_the_index(self): + conn = self.schema(self.REFS.format(kind=self.GOOD_REF), indexes=False) + conn.execute("CREATE TABLE other (workflow_document_id INTEGER)") + conn.execute( + "CREATE INDEX idx_image_workflow_document ON other(workflow_document_id)" + ) + self.assertEqual( + wd.workflow_documents_schema_error(conn), + "Workflow document references are missing their index", + ) + + def test_migrate_refuses_a_bad_schema_and_leaves_it_alone(self): + conn = self.schema(table_extra="") # the table without the columns + with self.assertRaisesRegex(sqlite3.DatabaseError, "Incomplete"): + wd.migrate_workflow_documents(conn) + columns = { + row[1] for row in conn.execute("PRAGMA table_info(generated_images)") + } + self.assertNotIn("workflow_document_id", columns) + + def test_migrate_validates_what_it_installed(self): + conn = self.schema(table=False) + with patch.object( + wd, + "workflow_documents_schema_error", + side_effect=[None, "Workflow documents broke"], + ): + with self.assertRaisesRegex(sqlite3.DatabaseError, "broke"): + wd.migrate_workflow_documents(conn) + tables = {row[0] for row in conn.execute("SELECT name FROM sqlite_master")} + self.assertNotIn("workflow_documents", tables) + + +class TestEdgeErrors(DocumentsTestCase): + def test_store_refuses_values_that_cannot_be_documents(self): + conn = self.conn() + with self.assertRaises(TypeError): + wd.store_text(conn, b"bytes") + with patch.object(wd, "MAX_DOCUMENT_BYTES", 8): + with self.assertRaises(ValueError): + wd.store_text(conn, "123456789") + self.assertFalse(wd.storable("12345678é")) + self.assertTrue(wd.storable("1234567")) + self.assertEqual( + wd.stored_values(conn, {"workflow_data": "123456789"}), + {"workflow_data": "123456789", "workflow_document_id": None}, + ) + conn.rollback() + + def test_a_hash_match_with_another_length_is_a_mismatch(self): + document_id = self.store("abc") + conn = self.conn() + with conn: + conn.execute( + "UPDATE workflow_documents SET byte_length=5 WHERE id=?", + (document_id,), + ) + with self.assertRaises(wd.DocumentMismatch): + wd.store_text(conn, "abc") + conn.rollback() + for length in (-1, 2, 4): + with conn: + conn.execute( + "UPDATE workflow_documents SET byte_length=? WHERE id=?", + (length, document_id), + ) + with self.subTest(length=length), self.assertLogs(wd.logger, "WARNING"): + self.assertIsNone(wd.read_stored_text(conn, None, document_id)) + + def test_the_reader_memo_is_bounded(self): + ids = [self.store(f"text {index}") for index in range(40)] + reader = wd.DocumentReader(self.conn()) + for index, document_id in enumerate(ids): + self.assertEqual(reader.text(None, document_id), f"text {index}") + self.assertEqual(len(reader._memo), wd.DocumentReader.MEMO) + + def test_a_row_without_text_columns_is_left_alone(self): + reader = wd.DocumentReader(self.conn()) + self.assertEqual(wd.resolve_image_row(reader, {"id": 1}), {"id": 1}) + self.assertEqual( + wd.resolve_image_row(reader, {"workflow_document_id": None}), + {"workflow_data": None}, + ) + + def test_documents_of_rows_takes_only_known_columns(self): + with self.assertRaises(ValueError): + wd.documents_of_rows(self.conn(), "image_path", ["x"]) + + class TestWritersStoreNewForm(DocumentsTestCase): def test_link_image_stores_documents_not_inline_text(self): workflow = sample_workflow(1)