feat(maintenance): store each workflow once (#315)
* feat(maintenance): reviewed, read-only maintenance plans POST /prompt_manager/maintenance/preview reviews duplicate_prompts, optimize, unused_components or missing_links on its own mode=ro snapshot and answers a Plan (scope, protected reasons, backup policy, media, at most 50 detail rows). The fingerprint, targets and judged sources stay in a server-side PlanStore under a single-use token (15 min, 8 plans, 16 MiB, never truncated; a restart forgets them), bound to the database it read. The previews reuse the legacy queries: duplicate groups and the merge are shared with cleanup_duplicates, ORPHANED_PROMPTS_SQL is built from the #239 protection predicates the plan reports, and missing links are judged by the health check's link_states. Legacy /maintenance keeps its shape; vacuum now reports success:false when VACUUM fails, and cleanup_orphaned_images skips links under unavailable sources. * feat(maintenance): apply reviewed plans as Activity jobs POST /prompt_manager/maintenance/apply {token} starts a maintenance job (records admission, not cancellable, groups changed/protected/errors). The token is consumed on the loop turn the job starts: a second apply is 409 plan-used naming the job, a busy refusal keeps the token, an unknown or expired one is 410 plan-expired, and a restore since the review is 409 storage-changed. With the safety backup policy on (maintenance.safety_backup, default true, saved through the settings route with the revision), the job first writes a backup into a temp file beside the database, makes it a standalone rollback-journal file, verifies it, fsyncs it and renames it onto a unique <db>.backup_<YYYYmmdd_HHMMSS_ffffff>[_n] name; any failure leaves no file and nothing changed. The apply then rechecks the reviewed fingerprint inside BEGIN IMMEDIATE (plan-changed writes nothing), deletes in batches of 400 ids, counts changed and protected records, logs the backup by name only and prunes old safety backups. With the policy off the log and the status line say "No safety backup". * feat(maintenance): find and remove duplicate files through Activity Duplicate image files move onto the job registry: - POST /prompt_manager/jobs/duplicates finds content-identical media files over every effective source as a cancellable read job ("Find duplicate image files", group groups_found); the server holds the groups per job and GET /prompt_manager/maintenance/duplicates/{id} pages them ten at a time. - The duplicate_files preview binds a removal to the reviewed keepers. Every keeper and removal is checked again: containment as text before any filesystem access, then not a link, still inside its source, same size and hash. The plan lists the database links that move. - Apply runs duplicate_cleanup (records, not cancellable): the promised safety backup, then per file a recheck, the unlink with every handle closed and its thumbnail. A changed file is kept, the first failed unlink stops the job with a coded message (never the OS text), and only removed files have their image links and favorites moved to the kept copy. - The legacy scan, stream, status and delete routes keep their shapes on the same registry kinds; the hashing, grouping and thumbnail removal are shared with them. Coordinator rulings on Tasks 15-16: - Safety backups share the restore pool of 10 `<db>.backup_*` files; the pool is now also pruned when an apply fails after its backup. - Plans report the safety backup's size and the free disk space, and apply refuses up front (507 insufficient-space) when the backup or VACUUM does not fit, so no job fails halfway for space. * fix(maintenance): never remove the kept copy; move records before unlink Final review of #313: - C1: the duplicate find tells files apart by (st_dev, st_ino), so a source listed twice, a symlinked source inside another and a hardlink are one file. The preview and the job refuse a removal that os.path.samestat matches with its kept copy. - I1: each file's records move to its verified kept copy in one short transaction before that file is unlinked; a failed move deletes nothing and a failed unlink leaves the records on the identical kept copy. The status line counts files removed and not removed. - I2: links are looked up under every spelling of the source (the real path, the configured and monitored folders, ComfyUI's output). - I5: image tags join the kept copy as their union and the removed copy's pending, undecided Needs review items are settled; the plan says so. - I3: the VACUUM disk check measures SQLite's real temporary folder. - M1: the token is looked at again with no await before the job starts. - M2: the backup decision is made once and handed to the job. - M3: backups are linked into place; startup sweeps interrupted temps. - M4: Optimize runs alone (new "exclusive" admission) and its plan says the database is locked while VACUUM runs. - M5: the reviewed duplicate merge is strict: collisions and Library titles are protected, and errors roll back. The legacy cleanup is unchanged. - M7: 2,000 keep entries within a 1 MiB body. M8: event waits in tests. - M9: the preview binds files by size, mtime and inode without hashing; the job hashes right before each unlink. * fix(maintenance): protect one file reached through two mounts Re-review of #313, round 2: - R1: a removal on another device than its kept copy, with the same size and mtimes within 2 s, is protected when its inode is equal or its path within its source is the same ("May be the same file reached through another mount or share"), in the preview and again right before each unlink. The job holds the kept copy open during the unlink, so on Windows deleting the kept copy itself through another path fails. - R2: a file's spellings are found by containment: every configured, monitored, ComfyUI-output or listed folder whose stat identity is a folder above the file. Parent spellings of nested sources and case variants match; favorites move under every listed root holding the file. Lookups stay index searches. - m1: a failed unlink's row says the links now name the kept copy. - m2: a volume with no file ids reads "Can't confirm this is a separate file on this drive". * fix(maintenance): stat duplicate candidates with os.stat Python 3.10's Path.stat uses the os.stat bound when pathlib was imported, so the find and the per-file checks could see different stat calls. The find now calls os.stat directly, like every other check; the mount tests failed on py3.10 only. * fix(maintenance): recheck the kept copy right before each unlink Re-review 2 of #313: - n1: the cross-mount guard also matches when one path within its source ends with the other on a whole folder name (sources listed at different depths, e.g. 2024/x.png and x.png on a FUSE or noserverino mount). - n2: after a file's records move and right before its unlink, the held kept copy must still be the file at its path (fstat of the hold against os.stat) and the cross-mount guard must still hold; otherwise the file is kept, its records move back to it, and it counts as not removed. * test(maintenance): the kept-copy race tests hold on Windows On Windows the removal job holds the kept copy open without FILE_SHARE_DELETE, so the tests' attempt to delete or replace it during the record move is refused. The tests now assert that native outcome: the kept copy stays and the removal goes ahead. * feat(database): store each workflow and prompt graph once generated_images.workflow_data and prompt_metadata repeated the same large JSON text for every image of a run. Each distinct text now lives once in workflow_documents, keyed by the sha256 of its exact UTF-8 bytes and zlib-compressed; image rows reference it through workflow_document_id and prompt_document_id. - New writes (link_image_to_prompt, pair ingestion for scan, monitor and Needs review) store the new form. Old rows keep their inline text and read the same until the migration rewrites them. - Every reader goes through one accessor (read_stored_text, resolve_image_row): gallery and legacy image lists, image by id, prompt info for the inspector and metadata page, diagnostics. Payloads keep their shape. Inline text wins when present; a document that fails its length or hash check reads as unreadable, never as other text. - No trigger (restore validators refuse them): a document goes in the same transaction as its last reference (image and prompt deletes, missing-link cleanup, duplicate merges, duplicate-file link merges, ingestion updates), and a sweep at startup catches anything a path missed. - verify_database_file validates the table, the id columns, their foreign keys and indexes, and the stored values. * 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. * feat(maintenance): store each workflow once as an Activity job The reviewed maintenance plan gains workflow_storage. Applying it starts "Store each workflow once": kind maintenance, records admission, cancellable, groups changed/protected/errors. - The promised verified safety backup comes first; a failed backup ends the job with nothing rewritten. The apply's up-front disk check covers it. - Old-form rows are rewritten batch by batch on the event loop's executor, with progress after each committed batch. Cancel stops between batches; a failed batch rolls back whole and fails the job with a coded message. Either way every record reads as before and the next apply continues. - When every row was looked at, "Reclaim space" starts: exclusive (runs alone), not cancellable. It checks the disk for VACUUM's copies of the live pages, then sweeps, runs VACUUM and truncates the WAL. If another job holds the registry it does not start, and the finished migration says so and names Optimize database. - No image or video file is opened; tests track every open under the output folder and compare the media bytes and mtimes. * feat(maintenance): document the workflow storage contract TESTING.md's maintenance contract gains the workflow storage rules: the documents table and its two forms, the one accessor, garbage in the same transaction plus the sweeps, the batched migration and Reclaim space, the media and memory guarantees, and that a deduped backup cannot be restored into an older version. A test covers the preview on SQLite before 3.43 (no octet_length, as in older Windows Python builds): the inline-bytes detail is left out. * feat(database): keep the migration safe row by row and beside other writers Final review fixes M1, M2, M3 and M5 for the workflow storage dedupe: - M1: a batch also ends after 0.5 s, and the loop leaves the write lock free for 0.15 s after each committed batch (longer than SQLite's longest busy poll), so the image monitor and run counting on another connection get in between. RED: a threaded writer with a 2 s busy timeout hit "database is locked" (4 of 5 links landed). - M2: each row runs in its own savepoint. A row that cannot be rewritten (undecodable text, a text that would not come back) is rolled back alone, stays inline and is listed in errors; the rest of the batch goes on. - M3: a reused document is decompressed and checked before its id is handed out, and repaired from the text (whose hash is its key) when damaged; verify_database_file reads every document and refuses a damaged one. - M5: the two legacy generated_images rebuilds refuse a file that already has the document id columns instead of dropping them. - The preview names the pre-migration backup and says it is always taken. * feat(maintenance): always back up before storing workflows once Final review fixes I1, I2 and I3 for the workflow storage dedupe: - I1: the migration always takes a verified backup, whatever maintenance.safety_backup says; the plan reports backup.policy "required" and the disk check counts it. - I2: that backup is prompts.db.pre-workflow-storage_<stamp>, outside the pruned prompts.db.backup_* pool and kept until the user removes it; the job's last line and the preview name it, and the legacy restore accepts it. README's changelog warns that an older version opens and restores a rewritten database but shows no workflows, and says how to go back. - I3: tools/workflow_text_digests.py writes the sha256 of every record's stored texts (read-only; inline before the rewrite, through the accessor after), so a diff before and after proves byte identity. - A row the migration could not rewrite is listed in the job's errors and named in its last line, and Reclaim space still runs. * feat(database): fail the batch on system errors and name a damaged document Re-review fixes R1 and R6 (database side), plus R3's helpers: - R1: a row is kept inline only for its own cause: text SQLite cannot decode, or a text that would not come back. A full disk, an I/O error, a malformed database or a bug now fails the batch and rolls it back, instead of listing every remaining row as "could not be rewritten". - R6: verify_database_file names the damaged document and says what to do (restore a backup from before the damage, or store an image with the same workflow, which repairs it). - R3: pre_migration_backups lists the kept backups, earliest first; the plan says "Nothing to rewrite" and names an existing kept backup, and the older version note says new images use the new form from the first start and names the earliest backup as the way back. * feat(maintenance): one kept pre-migration backup and nothing-to-do refusal Re-review fixes R2, R3, R4, R5 and R6 (job side): - R3: Apply refuses a plan with nothing to rewrite (409 nothing-to-do, no job, no backup). Only the first apply keeps a pre-workflow-storage backup; later applies (after a Cancel, a rerun) take a pooled safety backup, forced on whatever the policy, so one copy is never pruned. - R4: a second-connection writer test runs against the live job loop; with its pause removed the writer hits "database is locked" (4 of 5 links). - R5: the digest tool loads only database/workflow_documents.py by path (no package start-up, no log handler) and opens a stopped database with immutable=1, so it leaves no -wal/-shm; a live one is read through its WAL. - R6: a backup that fails on a damaged document says which one in the job line. - R2: README's changelog says images recorded by this version use the new form from the first start, and the way back is the earliest pre-workflow-storage backup (or a pre-update backup if the job never ran).
This commit is contained in:
Vito
@@ -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_<date time>` 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_<date time>` 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))
|
||||
|
||||
+67
@@ -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 `<db>.pre-workflow-storage_<YYYYmmdd_HHMMSS_ffffff>`,
|
||||
outside the pruned `<db>.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 <db> <out>` 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
|
||||
|
||||
@@ -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(
|
||||
|
||||
+18
-2
@@ -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.
|
||||
|
||||
@@ -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}"
|
||||
|
||||
+79
-22
@@ -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"]:
|
||||
|
||||
@@ -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)}
|
||||
|
||||
@@ -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
|
||||
@@ -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: ``<db>.pre-workflow-storage_<stamp>``, outside the
|
||||
# pruned ``<db>.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_<date and time>", kept=kept[0]
|
||||
)
|
||||
if kept
|
||||
else BACKUP_NOTE.format(
|
||||
name=f"{name}.{PRE_MIGRATION_LABEL}_<date and time>"
|
||||
)
|
||||
)
|
||||
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."
|
||||
)
|
||||
@@ -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,
|
||||
|
||||
+66
-24
@@ -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 ``<db>.backup_<stamp>[_n]`` (O_EXCL)."""
|
||||
for candidate in _backup_names(db_path):
|
||||
def _reserve_backup_name(db_path, label="backup"):
|
||||
"""Create and return an unused ``<db>.<label>_<stamp>[_n]`` (O_EXCL)."""
|
||||
for candidate in _backup_names(db_path, label):
|
||||
try:
|
||||
fd = os.open(candidate, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
|
||||
except FileExistsError:
|
||||
@@ -302,13 +326,13 @@ def _reserve_backup_name(db_path):
|
||||
raise SafetyBackupFailed("No free backup name")
|
||||
|
||||
|
||||
def _backup_names(db_path):
|
||||
base = f"{db_path}.backup_{_now().strftime('%Y%m%d_%H%M%S_%f')}"
|
||||
def _backup_names(db_path, label="backup"):
|
||||
base = f"{db_path}.{label}_{_now().strftime('%Y%m%d_%H%M%S_%f')}"
|
||||
for attempt in range(_NAME_ATTEMPTS):
|
||||
yield base if attempt == 0 else f"{base}_{attempt}"
|
||||
|
||||
|
||||
def _place_backup(temp, db_path):
|
||||
def _place_backup(temp, db_path, label="backup"):
|
||||
"""Give the verified ``temp`` its unique backup name; return that path.
|
||||
|
||||
``os.link`` is atomic and refuses an existing name, so no empty or half
|
||||
@@ -316,7 +340,7 @@ def _place_backup(temp, db_path):
|
||||
without hard links falls back to reserving the name (``O_EXCL``) and
|
||||
``os.replace``. The temp file is removed by the caller.
|
||||
"""
|
||||
for candidate in _backup_names(db_path):
|
||||
for candidate in _backup_names(db_path, label):
|
||||
try:
|
||||
os.link(temp, candidate)
|
||||
except FileExistsError:
|
||||
@@ -327,7 +351,7 @@ def _place_backup(temp, db_path):
|
||||
return candidate
|
||||
else:
|
||||
raise SafetyBackupFailed("No free backup name")
|
||||
final = _reserve_backup_name(db_path)
|
||||
final = _reserve_backup_name(db_path, label)
|
||||
try:
|
||||
os.replace(temp, final)
|
||||
except OSError:
|
||||
@@ -381,12 +405,15 @@ def _fsync(path):
|
||||
os.fsync(handle.fileno())
|
||||
|
||||
|
||||
def take_safety_backup(model):
|
||||
def take_safety_backup(model, label="backup"):
|
||||
"""Write and verify a safety backup; return its file name (blocking).
|
||||
|
||||
The copy is made in a temp file beside the database and only renamed
|
||||
onto its final, unique name once it verified, so a backup file on disk
|
||||
is always a verified one. Raises SafetyBackupFailed, leaving no file.
|
||||
onto its final, unique name ``<db>.<label>_<stamp>`` once it verified,
|
||||
so a backup file on disk is always a verified one. ``label`` "backup"
|
||||
joins the pruned pool; another label (the workflow storage migration's
|
||||
``pre-workflow-storage``) is kept until the user removes it. Raises
|
||||
SafetyBackupFailed, leaving no file.
|
||||
"""
|
||||
db_path = os.path.abspath(model.db_path)
|
||||
fd, temp = tempfile.mkstemp(
|
||||
@@ -398,13 +425,13 @@ def take_safety_backup(model):
|
||||
if not model.backup_database(temp):
|
||||
raise SafetyBackupFailed("The backup could not be written")
|
||||
_standalone(temp)
|
||||
ok, _reason = model.verify_database_file(temp)
|
||||
ok, reason = model.verify_database_file(temp)
|
||||
if not ok:
|
||||
raise SafetyBackupFailed("The backup did not verify")
|
||||
raise SafetyBackupFailed("The backup did not verify", reason)
|
||||
_fsync(temp)
|
||||
if os.name != "nt":
|
||||
os.chmod(temp, 0o600)
|
||||
final = _place_backup(temp, db_path)
|
||||
final = _place_backup(temp, db_path, label)
|
||||
return os.path.basename(final)
|
||||
except OSError as exc:
|
||||
_remove_quietly(final)
|
||||
@@ -519,6 +546,7 @@ class MaintenanceRoutesMixin:
|
||||
"targets": draft.targets,
|
||||
"roots": roots,
|
||||
"protected": draft.protected["count"],
|
||||
"rows": draft.scope["count"],
|
||||
},
|
||||
)
|
||||
|
||||
@@ -538,11 +566,18 @@ class MaintenanceRoutesMixin:
|
||||
"storage-changed",
|
||||
)
|
||||
policy = self._safety_backup_policy()
|
||||
destination = "backup"
|
||||
if operation == WORKFLOW_STORAGE:
|
||||
# Always backed up first, whatever the policy (the way back);
|
||||
# one kept pre-migration backup ever, pooled ones after it.
|
||||
policy = "required"
|
||||
kept = await self._run_in_executor(pre_migration_backups, db.model.db_path)
|
||||
destination = "backup" if kept else PRE_MIGRATION_LABEL
|
||||
try:
|
||||
disk, _short, size = await self._run_in_executor(
|
||||
disk_plan,
|
||||
db.model.db_path,
|
||||
backup=policy == "on",
|
||||
backup=policy != "off",
|
||||
vacuum=operation == "optimize",
|
||||
)
|
||||
except (OSError, sqlite3.Error):
|
||||
@@ -567,8 +602,8 @@ class MaintenanceRoutesMixin:
|
||||
**public,
|
||||
"backup": {
|
||||
"policy": policy,
|
||||
"destination_label": f"{name}.backup_<date and time>",
|
||||
"expected_bytes": size if policy == "on" else None,
|
||||
"destination_label": f"{name}.{destination}_<date and time>",
|
||||
"expected_bytes": size if policy != "off" else None,
|
||||
},
|
||||
"disk": disk,
|
||||
},
|
||||
@@ -596,6 +631,9 @@ class MaintenanceRoutesMixin:
|
||||
409,
|
||||
"storage-changed",
|
||||
)
|
||||
if plan.operation == WORKFLOW_STORAGE and not plan.payload.get("rows"):
|
||||
store.discard(token) # no job, no backup: nothing to rewrite
|
||||
return _failure(NOTHING_TO_DO, 409, "nothing-to-do")
|
||||
backup = self._backup_promised(plan)
|
||||
try:
|
||||
_disk, short, _size = await self._run_in_executor(
|
||||
@@ -618,6 +656,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,
|
||||
@@ -663,6 +703,8 @@ class MaintenanceRoutesMixin:
|
||||
Decided once, when the apply starts (the disk check measured it), and
|
||||
handed to the job.
|
||||
"""
|
||||
if plan.operation == WORKFLOW_STORAGE:
|
||||
return True # never without its pre-migration backup
|
||||
return plan.payload["backup"] == "on" or self._safety_backup_policy() == "on"
|
||||
|
||||
async def _safety_backup_step(self, job, db, backup_on):
|
||||
@@ -679,9 +721,9 @@ class MaintenanceRoutesMixin:
|
||||
job.log_line("Taking a safety backup of the database")
|
||||
try:
|
||||
backup = await self._run_in_executor(take_safety_backup, db.model)
|
||||
except SafetyBackupFailed:
|
||||
except SafetyBackupFailed as error:
|
||||
self.logger.error("Safety backup failed", exc_info=True)
|
||||
_fail_job(job, "backup-failed", BACKUP_FAILED)
|
||||
_fail_job(job, "backup-failed", backup_failed_message(error))
|
||||
return False, None
|
||||
job.log_line(f"Safety backup {backup} verified")
|
||||
return True, backup
|
||||
|
||||
@@ -0,0 +1,277 @@
|
||||
"""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 always takes
|
||||
a verified pre-migration backup, whatever the safety backup policy, named
|
||||
``<db>.pre-workflow-storage_<stamp>`` outside the pruned ``<db>.backup_*``
|
||||
pool and kept until the user removes it (``backup-failed`` ends the job
|
||||
with nothing rewritten). It then rewrites old-form rows batch by batch
|
||||
(``database/workflow_storage.py``): one transaction per batch of at most
|
||||
``BATCH_SECONDS``, then ``BATCH_PAUSE`` with the write lock free, progress
|
||||
after each, Cancel honoured between batches. A row that cannot be
|
||||
rewritten stays inline and is listed in ``errors``; any other failure
|
||||
rolls its batch back whole and fails the job. Every record still reads as
|
||||
before, and 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 asyncio
|
||||
import sqlite3
|
||||
|
||||
from .jobs import MAINTENANCE_JOB_KIND, JobBusy
|
||||
from .maintenance import (
|
||||
MAINTENANCE_GROUPS,
|
||||
PRE_MIGRATION_LABEL,
|
||||
SafetyBackupFailed,
|
||||
backup_failed_message,
|
||||
_fail_job,
|
||||
_megabytes,
|
||||
disk_plan,
|
||||
short_space_message,
|
||||
take_safety_backup,
|
||||
)
|
||||
|
||||
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."
|
||||
ROW_KEPT = "could not be rewritten; it keeps its inline text and reads as before"
|
||||
|
||||
|
||||
NOTHING_LEFT = "Every image record already stores its workflow once"
|
||||
|
||||
|
||||
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):
|
||||
# Always backed up first, whatever the policy: the way back. The first
|
||||
# apply keeps a pre-migration backup; later ones (after a Cancel, a
|
||||
# rerun) take a pooled safety backup, so only one copy is never pruned.
|
||||
counts = await self._run_in_executor(db.workflow_storage_status)
|
||||
if not counts["rows"]:
|
||||
job.log_line(NOTHING_LEFT)
|
||||
job.update(progress=100, status_line=NOTHING_LEFT)
|
||||
return
|
||||
kept = await self._run_in_executor(
|
||||
workflow_storage.pre_migration_backups, db.model.db_path
|
||||
)
|
||||
label = "backup" if kept else PRE_MIGRATION_LABEL
|
||||
backup = await self._backup_before_rewrite(job, db, label, kept)
|
||||
if backup is None:
|
||||
return
|
||||
try:
|
||||
done = await self._rewrite_workflows(job, db, backup)
|
||||
finally:
|
||||
if kept:
|
||||
await self._prune_backup_pool(job, db)
|
||||
if done:
|
||||
job.finish("finished")
|
||||
self._start_reclaim_space(job, db)
|
||||
|
||||
async def _backup_before_rewrite(self, job, db, label, kept):
|
||||
"""The verified backup's status note, or None (the job failed).
|
||||
|
||||
``label`` ``pre-workflow-storage``: kept until the user removes it.
|
||||
``backup``: a pooled safety backup, while the earliest kept
|
||||
pre-migration backup (``kept[0]``) stays. A failure ends the job
|
||||
``backup-failed``, nothing changed.
|
||||
"""
|
||||
noun = "safety backup" if kept else "pre-migration backup"
|
||||
job.update(status_line=f"Taking the {noun}")
|
||||
job.log_line(f"Taking the {noun} of the database")
|
||||
try:
|
||||
name = await self._run_in_executor(take_safety_backup, db.model, label)
|
||||
except SafetyBackupFailed as error:
|
||||
self.logger.error("Backup before the rewrite failed", exc_info=True)
|
||||
_fail_job(job, "backup-failed", backup_failed_message(error))
|
||||
return None
|
||||
if kept:
|
||||
job.log_line(
|
||||
f"Safety backup {name} verified; the pre-migration backup "
|
||||
f"{kept[0]} is kept"
|
||||
)
|
||||
return f"Safety backup {name}"
|
||||
job.log_line(
|
||||
f"Pre-migration backup {name} verified; it is kept until you remove it"
|
||||
)
|
||||
return f"Pre-migration backup {name}"
|
||||
|
||||
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 = errors = 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"]
|
||||
errors += len(batch["errors"])
|
||||
job.record_counted("changed", batch["rows"])
|
||||
job.record_counted("protected", batch["protected"])
|
||||
for image_id in batch["errors"]:
|
||||
job.record(
|
||||
"errors",
|
||||
{"path": "", "detail": f"Image record {image_id}: {ROW_KEPT}"},
|
||||
f"! Image record {image_id}: {ROW_KEPT}",
|
||||
)
|
||||
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})
|
||||
# Leave the write lock free between batches (monitor, run counts).
|
||||
await asyncio.sleep(workflow_storage.BATCH_PAUSE)
|
||||
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"
|
||||
if errors:
|
||||
kept = "keeps its" if errors == 1 else "keep their"
|
||||
line += (
|
||||
f"; {_plural(errors, 'image record')} could not be rewritten and "
|
||||
f"{kept} inline text"
|
||||
)
|
||||
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}")
|
||||
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}")
|
||||
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)
|
||||
@@ -654,11 +654,20 @@ class TestCleanupJob(DuplicateFilesTestCase):
|
||||
db.link_image_to_prompt(moved, str(self.output_dir / "b.png"))
|
||||
db.link_image_to_prompt(moved, str(self.output_dir / "d.png"))
|
||||
db.link_image_to_prompt(merged, str(self.output_dir / "a.png"))
|
||||
db.link_image_to_prompt(merged, str(self.output_dir / "sub" / "c.png"))
|
||||
db.link_image_to_prompt(
|
||||
merged,
|
||||
str(self.output_dir / "sub" / "c.png"),
|
||||
{"workflow": {"only": "c"}, "prompt": {"only": "c"}},
|
||||
)
|
||||
db.set_image_favorite(
|
||||
duplicate_files.favorite_root_key(self.output_dir), "b.png"
|
||||
)
|
||||
plan, cleanup = await self.remove_all_but("a.png")
|
||||
# The merged link's own workflow document went with it (#249 S6 dedupe).
|
||||
documents = db.model.get_connection().execute(
|
||||
"SELECT COUNT(*) FROM workflow_documents"
|
||||
)
|
||||
self.assertEqual(documents.fetchone()[0], 1) # the shared "{}"
|
||||
self.assertEqual(
|
||||
plan["details"][0]["detail"],
|
||||
"output · a copy of a.png · 1 database link moves to the kept copy",
|
||||
|
||||
@@ -0,0 +1,424 @@
|
||||
"""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, store_text
|
||||
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"], "required")
|
||||
self.assertEqual(
|
||||
plan["backup"]["destination_label"],
|
||||
"prompts.db.pre-workflow-storage_<date and time>",
|
||||
)
|
||||
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)
|
||||
self.assertEqual(self.backups(), []) # never in the pruned pool
|
||||
[backup] = [
|
||||
name
|
||||
for name in os.listdir(self.tmpdir)
|
||||
if name.startswith("prompts.db.pre-workflow-storage_")
|
||||
]
|
||||
self.assertIn(f"Pre-migration backup {backup} verified", "\n".join(job.log))
|
||||
self.assertEqual(
|
||||
job.status_line,
|
||||
f"Rewrote 40 image records · Pre-migration backup {backup}",
|
||||
)
|
||||
# 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(), [backup]) # only the kept backup
|
||||
self.assertNotIn(self.tmpdir, json.dumps(job.report()))
|
||||
|
||||
async def test_a_backup_failure_rewrites_nothing(self):
|
||||
self.old_library()
|
||||
from py.api import workflow_storage as api_storage
|
||||
|
||||
with patch.object(
|
||||
api_storage,
|
||||
"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,),
|
||||
)
|
||||
store_text(conn, "a document no record references")
|
||||
_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)",
|
||||
)
|
||||
@@ -99,11 +99,16 @@ class StrictRows:
|
||||
"""Assertions over the rows a write left behind."""
|
||||
|
||||
def image_rows(self, db):
|
||||
"""The stored texts, read through the one workflow document accessor."""
|
||||
from database.workflow_documents import DocumentReader, resolve_image_row
|
||||
|
||||
with db.model.get_connection() as conn:
|
||||
reader = DocumentReader(conn)
|
||||
return [
|
||||
dict(row)
|
||||
resolve_image_row(reader, row)
|
||||
for row in conn.execute(
|
||||
"SELECT prompt_id, workflow_data, prompt_metadata, parameters "
|
||||
"SELECT prompt_id, workflow_data, prompt_metadata, parameters, "
|
||||
"workflow_document_id, prompt_document_id "
|
||||
"FROM generated_images ORDER BY id"
|
||||
)
|
||||
]
|
||||
|
||||
@@ -0,0 +1,817 @@
|
||||
"""Each workflow and prompt graph stored once (database/workflow_documents.py).
|
||||
|
||||
``generated_images.workflow_data`` and ``prompt_metadata`` held the same large
|
||||
JSON text again for every image of a run. The text now lives once in
|
||||
``workflow_documents`` (keyed by the sha256 of its exact UTF-8 bytes, zlib
|
||||
compressed) and image rows reference it. Old rows keep their inline text until
|
||||
the migration rewrites them; every reader goes through one accessor, which
|
||||
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
|
||||
import unittest
|
||||
import zlib
|
||||
from unittest.mock import patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from database import workflow_documents as wd # noqa: E402
|
||||
from database.models import PromptModel # noqa: E402
|
||||
from database.operations import PromptDatabase # noqa: E402
|
||||
from utils.hashing import generate_prompt_hash # noqa: E402
|
||||
from utils.json_strict import dumps_strict # noqa: E402
|
||||
|
||||
try:
|
||||
from tests.open_handles import assert_closed
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from open_handles import assert_closed
|
||||
|
||||
NAN = float("nan")
|
||||
INF = float("inf")
|
||||
|
||||
|
||||
def sample_workflow(seed=0, nodes=40):
|
||||
"""A ComfyUI-shaped UI workflow: nodes, links, widgets and groups."""
|
||||
rng = random.Random(seed)
|
||||
return {
|
||||
"last_node_id": nodes,
|
||||
"last_link_id": nodes * 2,
|
||||
"nodes": [
|
||||
{
|
||||
"id": index,
|
||||
"type": rng.choice(
|
||||
["KSampler", "CLIPTextEncode", "VAEDecode", "LoraLoader"]
|
||||
),
|
||||
"pos": [rng.uniform(-2000, 2000), rng.uniform(-2000, 2000)],
|
||||
"size": {"0": 315, "1": 262},
|
||||
"flags": {},
|
||||
"order": index,
|
||||
"mode": 0,
|
||||
"inputs": [{"name": "model", "type": "MODEL", "link": index}],
|
||||
"outputs": [{"name": "LATENT", "type": "LATENT", "links": [index]}],
|
||||
"properties": {"Node name for S&R": "KSampler"},
|
||||
"widgets_values": [seed * 1000 + index, "fixed", 20, 7.5, "euler"],
|
||||
}
|
||||
for index in range(nodes)
|
||||
],
|
||||
"links": [[i, i, 0, i + 1, 0, "MODEL"] for i in range(nodes)],
|
||||
"groups": [{"title": "Sampler ✨ 采样", "bounding": [0, 0, 100, 100]}],
|
||||
"config": {},
|
||||
"extra": {"ds": {"scale": 0.5, "offset": [0, 0]}},
|
||||
"version": 0.4,
|
||||
}
|
||||
|
||||
|
||||
def sample_prompt_graph(seed=0):
|
||||
return {
|
||||
"3": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {"seed": seed, "steps": 20, "cfg": 7.0, "model": ["4", 0]},
|
||||
},
|
||||
"6": {"class_type": "CLIPTextEncode", "inputs": {"text": f"a cat {seed}"}},
|
||||
}
|
||||
|
||||
|
||||
def edge_texts():
|
||||
"""Texts whose bytes a parse-and-serialize round trip would change."""
|
||||
nested = "[" * 500 + "]" * 500
|
||||
return [
|
||||
"",
|
||||
"{}",
|
||||
"null",
|
||||
" ",
|
||||
"\n",
|
||||
'{"a": 1}',
|
||||
'{"a":1}',
|
||||
'{"a": 1, "a": 2}', # duplicate key: a parse would keep one
|
||||
'{"b": 1, "a": 2}', # key order
|
||||
'{"v": 1.0}',
|
||||
'{"v": 1e3}', # number spelling
|
||||
'{"v": 1.10}',
|
||||
'{"v": NaN, "w": Infinity, "x": -Infinity}', # ComfyUI's tokens
|
||||
'{"text": "NaN"}',
|
||||
'{"u": "\\u00e9\\ud83d\\ude00"}', # escaped
|
||||
'{"u": "é 😀 中文 עברית é"}', # raw unicode, astral, RTL, combining
|
||||
"trailing space ",
|
||||
'{"nul": "a\u0000b"}',
|
||||
"a\u0000b",
|
||||
"{}", # byte-order mark
|
||||
"
",
|
||||
"not json at all",
|
||||
nested,
|
||||
json.dumps(sample_workflow(1)),
|
||||
json.dumps(sample_workflow(1), indent=2),
|
||||
json.dumps(sample_workflow(1), separators=(",", ":")),
|
||||
json.dumps({"v": NAN, "w": [INF, -INF]}), # Python's lenient output
|
||||
dumps_strict({"v": NAN, "w": [INF, -INF]}),
|
||||
]
|
||||
|
||||
|
||||
def random_text(rng):
|
||||
"""Arbitrary text of valid code points (no surrogates), ASCII to astral."""
|
||||
length = rng.choice([0, 1, 2, 7, 64, 1000, 20000])
|
||||
ranges = [(0x00, 0x7F), (0x80, 0x7FF), (0x800, 0xD7FF), (0xE000, 0xFFFF)]
|
||||
ranges.append((0x10000, 0x10FFFF))
|
||||
chars = []
|
||||
for _ in range(length):
|
||||
low, high = rng.choice(ranges)
|
||||
chars.append(chr(rng.randint(low, high)))
|
||||
return "".join(chars)
|
||||
|
||||
|
||||
class DocumentsTestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmpdir = os.path.realpath(tempfile.mkdtemp())
|
||||
self.db_path = os.path.join(self.tmpdir, "prompts.db")
|
||||
self.db = PromptDatabase(self.db_path)
|
||||
self.addCleanup(self._cleanup)
|
||||
|
||||
def _cleanup(self):
|
||||
PromptModel.close_all_instances()
|
||||
PromptModel.reset_schema_cache()
|
||||
assert_closed(self, self.tmpdir)
|
||||
shutil.rmtree(self.tmpdir)
|
||||
|
||||
def conn(self):
|
||||
return self.db.model.get_connection()
|
||||
|
||||
def prompt(self, text="a cat"):
|
||||
return self.db.save_prompt(text, prompt_hash=generate_prompt_hash(text))
|
||||
|
||||
def old_form_image(self, prompt_id, name, workflow, prompt_graph):
|
||||
"""A row as every version before the dedupe wrote it: text inline."""
|
||||
conn = self.conn()
|
||||
with conn:
|
||||
path = os.path.join(self.tmpdir, name)
|
||||
cursor = conn.execute(
|
||||
"INSERT INTO generated_images (prompt_id, image_path, filename, "
|
||||
"file_path, workflow_data, prompt_metadata, parameters) "
|
||||
"VALUES (?, ?, ?, ?, ?, ?, ?)",
|
||||
(prompt_id, path, name, path, workflow, prompt_graph, "{}"),
|
||||
)
|
||||
return cursor.lastrowid
|
||||
|
||||
def raw(self, image_id):
|
||||
return dict(
|
||||
self.conn()
|
||||
.execute(
|
||||
"SELECT workflow_data, prompt_metadata, workflow_document_id, "
|
||||
"prompt_document_id FROM generated_images WHERE id=?",
|
||||
(image_id,),
|
||||
)
|
||||
.fetchone()
|
||||
)
|
||||
|
||||
def document_count(self):
|
||||
return (
|
||||
self.conn().execute("SELECT COUNT(*) FROM workflow_documents").fetchone()[0]
|
||||
)
|
||||
|
||||
def linked(self, prompt_id, name, seed):
|
||||
return self.db.link_image_to_prompt(
|
||||
prompt_id,
|
||||
os.path.join(self.tmpdir, name),
|
||||
{"workflow": sample_workflow(seed), "prompt": sample_prompt_graph(seed)},
|
||||
)
|
||||
|
||||
def store(self, text):
|
||||
conn = self.conn()
|
||||
with conn:
|
||||
return wd.store_text(conn, text)
|
||||
|
||||
|
||||
class TestRoundTrip(DocumentsTestCase):
|
||||
def assert_round_trip(self, text):
|
||||
document_id = self.store(text)
|
||||
self.assertIsInstance(document_id, int)
|
||||
back = wd.read_stored_text(self.conn(), None, document_id)
|
||||
self.assertIsInstance(back, str)
|
||||
self.assertEqual(back, text)
|
||||
self.assertEqual(back.encode("utf-8"), text.encode("utf-8"))
|
||||
return document_id
|
||||
|
||||
def test_edge_texts_come_back_byte_identical(self):
|
||||
for text in edge_texts():
|
||||
with self.subTest(text=text[:40]):
|
||||
self.assert_round_trip(text)
|
||||
|
||||
def test_texts_that_parse_alike_stay_distinct(self):
|
||||
spaced = self.assert_round_trip('{"a": 1}')
|
||||
tight = self.assert_round_trip('{"a":1}')
|
||||
self.assertNotEqual(spaced, tight)
|
||||
|
||||
def test_random_texts_come_back_byte_identical(self):
|
||||
rng = random.Random(20261024)
|
||||
for case in range(400):
|
||||
text = random_text(rng)
|
||||
with self.subTest(case=case):
|
||||
self.assert_round_trip(text)
|
||||
|
||||
def test_random_workflows_come_back_byte_identical(self):
|
||||
rng = random.Random(7)
|
||||
for case in range(150):
|
||||
value = sample_workflow(rng.randint(0, 10**9), nodes=rng.randint(0, 60))
|
||||
if rng.random() < 0.3:
|
||||
value["nodes"].append({"widgets_values": [NAN, INF, -INF]})
|
||||
for text in (json.dumps(value), dumps_strict(value)):
|
||||
with self.subTest(case=case):
|
||||
self.assert_round_trip(text)
|
||||
|
||||
def test_huge_workflow_round_trips(self):
|
||||
text = json.dumps(sample_workflow(3, nodes=4000))
|
||||
self.assertGreater(len(text), 1_200_000)
|
||||
self.assert_round_trip(text)
|
||||
|
||||
def test_identical_text_is_stored_once(self):
|
||||
text = json.dumps(sample_workflow(5))
|
||||
first = self.store(text)
|
||||
second = self.store(text)
|
||||
self.assertEqual(first, second)
|
||||
self.assertEqual(self.document_count(), 1)
|
||||
|
||||
def test_null_stays_null(self):
|
||||
self.assertIsNone(self.store(None))
|
||||
self.assertIsNone(wd.read_stored_text(self.conn(), None, None))
|
||||
self.assertEqual(self.document_count(), 0)
|
||||
|
||||
def test_document_is_compressed_utf8_keyed_by_sha256(self):
|
||||
text = "é" + json.dumps(sample_workflow(9))
|
||||
document_id = self.store(text)
|
||||
row = (
|
||||
self.conn()
|
||||
.execute(
|
||||
"SELECT sha256, byte_length, data FROM workflow_documents WHERE id=?",
|
||||
(document_id,),
|
||||
)
|
||||
.fetchone()
|
||||
)
|
||||
raw = text.encode("utf-8")
|
||||
import hashlib
|
||||
|
||||
self.assertEqual(bytes(row["sha256"]), hashlib.sha256(raw).digest())
|
||||
self.assertEqual(row["byte_length"], len(raw))
|
||||
self.assertEqual(zlib.decompress(row["data"]), raw)
|
||||
self.assertLess(len(row["data"]), len(raw) / 3)
|
||||
|
||||
def test_inline_text_wins_over_a_document(self):
|
||||
document_id = self.store("from the document")
|
||||
self.assertEqual(
|
||||
wd.read_stored_text(self.conn(), "inline", document_id), "inline"
|
||||
)
|
||||
self.assertEqual(wd.read_stored_text(self.conn(), "", document_id), "")
|
||||
|
||||
def test_a_tampered_document_reads_as_unreadable_never_as_other_text(self):
|
||||
document_id = self.store('{"a": 1}')
|
||||
conn = self.conn()
|
||||
for data in (zlib.compress(b'{"a": 2}'), b"not zlib"):
|
||||
with conn:
|
||||
conn.execute(
|
||||
"UPDATE workflow_documents SET data=? WHERE id=?",
|
||||
(data, document_id),
|
||||
)
|
||||
with self.subTest(data=data), self.assertLogs(wd.logger, level="WARNING"):
|
||||
self.assertIsNone(wd.read_stored_text(conn, None, document_id))
|
||||
|
||||
def test_a_missing_document_reads_as_unreadable(self):
|
||||
with self.assertLogs(wd.logger, level="WARNING"):
|
||||
self.assertIsNone(wd.read_stored_text(self.conn(), None, 424242))
|
||||
|
||||
def test_reader_memo_returns_the_same_text(self):
|
||||
text = json.dumps(sample_workflow(2))
|
||||
document_id = self.store(text)
|
||||
reader = wd.DocumentReader(self.conn())
|
||||
self.assertEqual(reader.text(None, document_id), text)
|
||||
self.assertEqual(reader.text(None, document_id), text)
|
||||
self.assertEqual(reader.text("inline", document_id), "inline")
|
||||
|
||||
|
||||
class TestSchema(DocumentsTestCase):
|
||||
def test_fresh_database_has_the_table_columns_and_indexes(self):
|
||||
conn = self.conn()
|
||||
self.assertIsNone(wd.workflow_documents_schema_error(conn, allow_missing=False))
|
||||
columns = {
|
||||
row[1] for row in conn.execute("PRAGMA table_info(generated_images)")
|
||||
}
|
||||
self.assertTrue(
|
||||
{"workflow_document_id", "prompt_document_id"}.issubset(columns)
|
||||
)
|
||||
indexes = {
|
||||
row[1] for row in conn.execute("PRAGMA index_list(generated_images)")
|
||||
}
|
||||
self.assertIn("idx_image_workflow_document", indexes)
|
||||
self.assertIn("idx_image_prompt_document", indexes)
|
||||
|
||||
def test_no_trigger_is_installed(self):
|
||||
# Every restore validator refuses triggers; GC is done in code.
|
||||
row = (
|
||||
self.conn()
|
||||
.execute("SELECT name FROM sqlite_master WHERE type='trigger'")
|
||||
.fetchone()
|
||||
)
|
||||
self.assertIsNone(row)
|
||||
|
||||
def test_current_database_verifies_and_restores(self):
|
||||
image = self.db.link_image_to_prompt(
|
||||
self.prompt(), os.path.join(self.tmpdir, "a.png"), {"workflow": {"n": 1}}
|
||||
)
|
||||
self.assertTrue(image)
|
||||
ok, reason = self.db.model.verify_database_file(
|
||||
self.db_path, require_library_state=True
|
||||
)
|
||||
self.assertTrue(ok, reason)
|
||||
|
||||
def test_an_old_database_opens_migrates_and_keeps_its_inline_rows(self):
|
||||
PromptModel.close_all_instances()
|
||||
PromptModel.reset_schema_cache()
|
||||
old_path = os.path.join(self.tmpdir, "old.db")
|
||||
conn = sqlite3.connect(old_path)
|
||||
try:
|
||||
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));"
|
||||
"INSERT INTO prompts (id, text, hash) VALUES (1, 'a cat', 'h');"
|
||||
"INSERT INTO generated_images (prompt_id, image_path, filename, "
|
||||
"file_path, workflow_data, prompt_metadata) VALUES (1, '/x/a.png', "
|
||||
"'a.png', '/x/a.png', '{\"n\": NaN}', '{\"3\": {}}');"
|
||||
)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
ok, reason = PromptModel.verify_database_file(old_path)
|
||||
self.assertTrue(ok, reason)
|
||||
old = PromptDatabase(old_path)
|
||||
[image] = old.get_prompt_images(1)
|
||||
self.assertEqual(image["workflow_data"], {"n": None})
|
||||
self.assertEqual(image["prompt_metadata"], {"3": {}})
|
||||
row = (
|
||||
old.model.get_connection()
|
||||
.execute("SELECT workflow_data, workflow_document_id FROM generated_images")
|
||||
.fetchone()
|
||||
)
|
||||
self.assertEqual(tuple(row), ('{"n": NaN}', None))
|
||||
old.close_all()
|
||||
|
||||
def test_verify_refuses_a_malformed_documents_table(self):
|
||||
bad = os.path.join(self.tmpdir, "bad.db")
|
||||
PromptModel.close_all_instances()
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
with sqlite3.connect(bad) as target:
|
||||
conn.backup(target)
|
||||
target.close()
|
||||
finally:
|
||||
conn.close()
|
||||
cases = {
|
||||
"table": "DROP TABLE workflow_documents; CREATE TABLE "
|
||||
"workflow_documents (id INTEGER PRIMARY KEY, sha256 TEXT, data TEXT)",
|
||||
"values": "INSERT INTO workflow_documents (sha256, byte_length, data) "
|
||||
"VALUES (zeroblob(32), -1, x'00')",
|
||||
"reference": "PRAGMA foreign_keys=OFF; INSERT INTO prompts (text, hash) "
|
||||
"VALUES ('p', 'h2'); INSERT INTO generated_images (prompt_id, image_path,"
|
||||
" filename, workflow_document_id) VALUES (last_insert_rowid(), '/a', 'a',"
|
||||
" 99)",
|
||||
"index": "DROP INDEX idx_image_workflow_document",
|
||||
}
|
||||
for name, script in cases.items():
|
||||
copy = os.path.join(self.tmpdir, f"bad-{name}.db")
|
||||
with open(bad, "rb") as source, open(copy, "wb") as target:
|
||||
target.write(source.read())
|
||||
conn = sqlite3.connect(copy)
|
||||
try:
|
||||
conn.executescript(script)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
with self.subTest(name=name):
|
||||
ok, reason = PromptModel.verify_database_file(copy)
|
||||
self.assertFalse(ok)
|
||||
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 conn, self.assertLogs(wd.logger, "WARNING"): # repaired from the text
|
||||
self.assertEqual(wd.store_text(conn, "abc"), document_id)
|
||||
self.assertEqual(wd.read_stored_text(conn, None, document_id), "abc")
|
||||
for length in (-1, 2, 4):
|
||||
with conn:
|
||||
conn.execute(
|
||||
"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)
|
||||
graph = sample_prompt_graph(1)
|
||||
prompt_id = self.prompt()
|
||||
first = self.db.link_image_to_prompt(
|
||||
prompt_id,
|
||||
os.path.join(self.tmpdir, "a.png"),
|
||||
{"workflow": workflow, "prompt": graph},
|
||||
)
|
||||
second = self.db.link_image_to_prompt(
|
||||
prompt_id,
|
||||
os.path.join(self.tmpdir, "b.png"),
|
||||
{"workflow": workflow, "prompt": graph},
|
||||
)
|
||||
a, b = self.raw(first), self.raw(second)
|
||||
self.assertIsNone(a["workflow_data"])
|
||||
self.assertIsNone(a["prompt_metadata"])
|
||||
self.assertIsNotNone(a["workflow_document_id"])
|
||||
self.assertEqual(a["workflow_document_id"], b["workflow_document_id"])
|
||||
self.assertEqual(a["prompt_document_id"], b["prompt_document_id"])
|
||||
self.assertEqual(self.document_count(), 2)
|
||||
conn = self.conn()
|
||||
self.assertEqual(
|
||||
wd.read_stored_text(conn, None, a["workflow_document_id"]),
|
||||
dumps_strict(workflow),
|
||||
)
|
||||
|
||||
def test_readers_answer_the_same_for_both_forms(self):
|
||||
workflow, graph = sample_workflow(4), sample_prompt_graph(4)
|
||||
workflow["nodes"][0]["widgets_values"] = [NAN, INF]
|
||||
prompt_id = self.prompt()
|
||||
new_id = self.db.link_image_to_prompt(
|
||||
prompt_id,
|
||||
os.path.join(self.tmpdir, "new.png"),
|
||||
{"workflow": workflow, "prompt": graph},
|
||||
)
|
||||
old_id = self.old_form_image(
|
||||
prompt_id, "old.png", dumps_strict(workflow), dumps_strict(graph)
|
||||
)
|
||||
new, old = self.db.get_image_by_id(new_id), self.db.get_image_by_id(old_id)
|
||||
for key in ("workflow_data", "prompt_metadata", "parameters"):
|
||||
self.assertEqual(new[key], old[key])
|
||||
self.assertEqual(new["workflow_data"]["nodes"][0]["widgets_values"], [None] * 2)
|
||||
self.assertNotIn("workflow_document_id", new)
|
||||
self.assertNotIn("prompt_document_id", old)
|
||||
listed = {
|
||||
image["id"]: image
|
||||
for reader in (
|
||||
lambda: self.db.get_prompt_images(prompt_id),
|
||||
lambda: self.db.get_recent_images(10),
|
||||
lambda: self.db.get_all_images(),
|
||||
lambda: self.db.search_images_by_prompt("cat"),
|
||||
)
|
||||
for image in reader()
|
||||
}
|
||||
self.assertEqual(listed[new_id]["workflow_data"], new["workflow_data"])
|
||||
self.assertEqual(listed[old_id]["prompt_metadata"], graph)
|
||||
info_new = self.db.get_image_prompt_info(os.path.join(self.tmpdir, "new.png"))
|
||||
info_old = self.db.get_image_prompt_info(os.path.join(self.tmpdir, "old.png"))
|
||||
self.assertEqual(info_new["workflow_data"], info_old["workflow_data"])
|
||||
self.assertEqual(info_new["prompt_metadata"], graph)
|
||||
|
||||
def test_diagnostics_recent_images_resolve_documents(self):
|
||||
prompt_id = self.prompt()
|
||||
self.db.link_image_to_prompt(
|
||||
prompt_id,
|
||||
os.path.join(self.tmpdir, "d.png"),
|
||||
{"workflow": {"n": 1}, "prompt": {"3": {}}},
|
||||
)
|
||||
from utils import diagnostics
|
||||
|
||||
tool = diagnostics.GalleryDiagnostics(self.db_path)
|
||||
result = tool.check_images_table()
|
||||
tool_images = result["recent_images"]
|
||||
self.assertEqual(tool_images[0]["workflow_data"], '{"n": 1}')
|
||||
self.assertNotIn("workflow_document_id", tool_images[0])
|
||||
|
||||
|
||||
class TestGarbage(DocumentsTestCase):
|
||||
"""A document no row references is removed with its last reference."""
|
||||
|
||||
def test_delete_image_releases_only_unshared_documents(self):
|
||||
prompt_id = self.prompt()
|
||||
shared_a = self.linked(prompt_id, "a.png", 1)
|
||||
self.linked(prompt_id, "b.png", 1)
|
||||
alone = self.linked(prompt_id, "c.png", 2)
|
||||
self.assertEqual(self.document_count(), 4)
|
||||
self.db.delete_image(shared_a)
|
||||
self.assertEqual(self.document_count(), 4)
|
||||
self.db.delete_image(alone)
|
||||
self.assertEqual(self.document_count(), 2)
|
||||
|
||||
def test_deleting_prompts_releases_their_images_documents(self):
|
||||
first, second = self.prompt("one"), self.prompt("two")
|
||||
self.linked(first, "a.png", 1)
|
||||
self.linked(second, "b.png", 2)
|
||||
self.db.delete_prompt(first)
|
||||
self.assertEqual(self.document_count(), 2)
|
||||
self.db.update_prompt_metadata(second, category="gone")
|
||||
self.assertEqual(self.db.delete_prompts_by_category("gone"), 1)
|
||||
self.assertEqual(self.document_count(), 0)
|
||||
|
||||
def test_cleanup_missing_images_releases_documents(self):
|
||||
prompt_id = self.prompt()
|
||||
self.linked(prompt_id, "gone.png", 1)
|
||||
self.db.cleanup_missing_images()
|
||||
self.assertEqual(self.document_count(), 0)
|
||||
self.linked(prompt_id, "gone-too.png", 2)
|
||||
self.assertEqual(self.db.cleanup_missing_images(roots=[self.tmpdir]), 1)
|
||||
self.assertEqual(self.document_count(), 0)
|
||||
|
||||
def test_sweep_removes_what_a_path_without_release_left(self):
|
||||
prompt_id = self.prompt()
|
||||
image = self.linked(prompt_id, "a.png", 1)
|
||||
conn = self.conn()
|
||||
with conn:
|
||||
conn.execute("DELETE FROM generated_images WHERE id=?", (image,))
|
||||
self.assertEqual(self.document_count(), 2)
|
||||
with conn:
|
||||
self.assertEqual(wd.sweep_documents(conn), 2)
|
||||
self.assertEqual(self.document_count(), 0)
|
||||
|
||||
def test_startup_sweeps_unreferenced_documents(self):
|
||||
prompt_id = self.prompt()
|
||||
image = self.linked(prompt_id, "a.png", 1)
|
||||
conn = self.conn()
|
||||
with conn:
|
||||
conn.execute("DELETE FROM generated_images WHERE id=?", (image,))
|
||||
PromptModel.close_all_instances()
|
||||
PromptModel.reset_schema_cache()
|
||||
self.db = PromptDatabase(self.db_path)
|
||||
self.assertEqual(self.document_count(), 0)
|
||||
|
||||
def test_release_keeps_a_document_the_other_column_references(self):
|
||||
prompt_id = self.prompt()
|
||||
text = '{"same": true}'
|
||||
first = self.db.link_image_to_prompt(
|
||||
prompt_id, os.path.join(self.tmpdir, "a.png"), {"workflow": {"same": True}}
|
||||
)
|
||||
second = self.db.link_image_to_prompt(
|
||||
prompt_id, os.path.join(self.tmpdir, "b.png"), {"prompt": {"same": True}}
|
||||
)
|
||||
self.assertEqual(
|
||||
self.raw(first)["workflow_document_id"],
|
||||
self.raw(second)["prompt_document_id"],
|
||||
)
|
||||
self.db.delete_image(first)
|
||||
self.assertEqual(
|
||||
wd.read_stored_text(
|
||||
self.conn(), None, self.raw(second)["prompt_document_id"]
|
||||
),
|
||||
text,
|
||||
)
|
||||
|
||||
|
||||
class TestMoreGarbagePaths(DocumentsTestCase):
|
||||
def test_bulk_delete_prompts_releases_documents(self):
|
||||
first, second = self.prompt("one"), self.prompt("two")
|
||||
self.linked(first, "a.png", 1)
|
||||
self.linked(second, "b.png", 2)
|
||||
self.assertEqual(self.db.bulk_delete_prompts([first, second]), 2)
|
||||
self.assertEqual(self.document_count(), 0)
|
||||
|
||||
def test_duplicate_prompt_merge_releases_images_it_could_not_move(self):
|
||||
primary = self.prompt("A cat")
|
||||
copy = self.db.save_prompt("a cat ", prompt_hash="other-hash")
|
||||
self.linked(primary, "same.png", 1)
|
||||
self.linked(copy, "same.png", 2) # collides with the primary's link
|
||||
self.assertEqual(self.document_count(), 4)
|
||||
self.db.cleanup_duplicates()
|
||||
self.assertEqual(self.document_count(), 2)
|
||||
|
||||
def test_maintenance_delete_ids_releases_documents(self):
|
||||
from database.maintenance import _delete_ids
|
||||
|
||||
prompt_id = self.prompt()
|
||||
image = self.linked(prompt_id, "a.png", 1)
|
||||
other = self.prompt("other")
|
||||
self.linked(other, "b.png", 2)
|
||||
conn = self.conn()
|
||||
with conn:
|
||||
_delete_ids(conn, "generated_images", [image])
|
||||
self.assertEqual(self.document_count(), 2)
|
||||
with conn:
|
||||
_delete_ids(conn, "prompts", [other])
|
||||
self.assertEqual(self.document_count(), 0)
|
||||
|
||||
def test_failed_insert_of_a_known_link_leaves_no_new_document(self):
|
||||
prompt_id = self.prompt()
|
||||
path = os.path.join(self.tmpdir, "a.png")
|
||||
self.db.link_image_to_prompt(prompt_id, path, {"workflow": {"v": 1}})
|
||||
before = self.document_count()
|
||||
self.assertEqual(
|
||||
self.db.link_image_to_prompt(prompt_id, path, {"workflow": {"v": 2}}), 0
|
||||
)
|
||||
self.assertEqual(self.document_count(), before)
|
||||
|
||||
|
||||
class TestPairIngestionDocuments(DocumentsTestCase):
|
||||
def resolution(self):
|
||||
return {
|
||||
"state": "resolved",
|
||||
"positive": {"text": "cat", "provenance": {"source": "api"}},
|
||||
"negative": {"text": "blur", "provenance": {"source": "api"}},
|
||||
}
|
||||
|
||||
def test_ingestion_stores_new_form_and_replaces_documents_it_updates(self):
|
||||
from database.pair_ingestion import ingest_resolved_pair
|
||||
|
||||
path = os.path.join(self.tmpdir, "scan.png")
|
||||
first = {"workflow": sample_workflow(1), "prompt": sample_prompt_graph(1)}
|
||||
added = ingest_resolved_pair(self.db, path, first, self.resolution())
|
||||
self.assertEqual(added["status"], "added")
|
||||
row = self.raw(added["image_id"])
|
||||
self.assertIsNone(row["workflow_data"])
|
||||
self.assertIsNotNone(row["workflow_document_id"])
|
||||
again = ingest_resolved_pair(self.db, path, first, self.resolution())
|
||||
self.assertEqual(again["status"], "unchanged")
|
||||
second = {"workflow": sample_workflow(2), "prompt": sample_prompt_graph(1)}
|
||||
updated = ingest_resolved_pair(self.db, path, second, self.resolution())
|
||||
self.assertEqual(updated["status"], "enriched")
|
||||
self.assertEqual(self.document_count(), 2) # the first workflow went
|
||||
image = self.db.get_image_by_id(added["image_id"])
|
||||
self.assertEqual(image["workflow_data"], sample_workflow(2))
|
||||
|
||||
def test_an_old_form_row_with_the_same_text_is_unchanged(self):
|
||||
from database.pair_ingestion import ingest_resolved_pair
|
||||
|
||||
path = os.path.join(self.tmpdir, "old.png")
|
||||
metadata = {"workflow": sample_workflow(3), "prompt": sample_prompt_graph(3)}
|
||||
added = ingest_resolved_pair(self.db, path, metadata, self.resolution())
|
||||
conn = self.conn()
|
||||
with conn: # rewrite it as an older version stored it
|
||||
conn.execute(
|
||||
"UPDATE generated_images SET workflow_data=?, prompt_metadata=?, "
|
||||
"workflow_document_id=NULL, prompt_document_id=NULL WHERE id=?",
|
||||
(
|
||||
dumps_strict(metadata["workflow"], sort_keys=True),
|
||||
dumps_strict(metadata["prompt"], sort_keys=True),
|
||||
added["image_id"],
|
||||
),
|
||||
)
|
||||
wd.sweep_documents(conn)
|
||||
again = ingest_resolved_pair(self.db, path, metadata, self.resolution())
|
||||
self.assertEqual(again["status"], "unchanged")
|
||||
self.assertEqual(self.document_count(), 0)
|
||||
|
||||
|
||||
class TestQueryPlans(DocumentsTestCase):
|
||||
"""Every statement the store and the migration run is an index search."""
|
||||
|
||||
def plan(self, sql, params=()):
|
||||
rows = self.conn().execute("EXPLAIN QUERY PLAN " + sql, params).fetchall()
|
||||
return " | ".join(row[3] for row in rows)
|
||||
|
||||
def test_statements_use_indexes(self):
|
||||
for name, (sql, params) in wd.QUERY_PLAN_STATEMENTS.items():
|
||||
with self.subTest(name=name):
|
||||
plan = self.plan(sql, params)
|
||||
self.assertNotRegex(plan, r"\bSCAN (generated_images|gi)\b(?! USING)")
|
||||
self.assertTrue("USING" in plan or "SEARCH" in plan, f"{name}: {plan}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,469 @@
|
||||
"""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, *args):
|
||||
calls.append(text)
|
||||
if len(calls) == 23:
|
||||
raise RuntimeError("power cut")
|
||||
return real(conn, text, *args)
|
||||
|
||||
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_keeps_its_row(self):
|
||||
_prompt_id, expected = self.library(count=6)
|
||||
with patch.object(ws, "read_stored_text", return_value="something else"):
|
||||
result = self.migrate(max_rows=50)
|
||||
self.assertEqual(result["errors"], sorted(expected))
|
||||
self.assertEqual(self.inline_ids(), set(expected))
|
||||
self.assertEqual(self.texts(), expected)
|
||||
self.assertEqual(self.document_count(), 0)
|
||||
|
||||
def test_stop_between_batches_then_continue(self):
|
||||
_prompt_id, expected = self.library(count=30)
|
||||
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_without_octet_length_omits_the_inline_bytes(self):
|
||||
# SQLite before 3.43 (older Windows Python builds) has no octet_length.
|
||||
self.library(count=4)
|
||||
with patch.object(ws.sqlite3, "sqlite_version_info", (3, 37, 2)):
|
||||
draft = self.db.preview_maintenance("workflow_storage")
|
||||
labels = [row["label"] for row in draft.details]
|
||||
self.assertNotIn("Inline text", labels)
|
||||
self.assertEqual(draft.scope["count"], 4)
|
||||
|
||||
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()
|
||||
@@ -0,0 +1,267 @@
|
||||
"""PR #315 scoped re-review fixes: R1, R3, R4, R5 and R6 (R2 is wording).
|
||||
|
||||
- R1: only a row's own cause (undecodable text, a text that would not come
|
||||
back) keeps that row inline; disk-full, I/O and corruption fail the batch.
|
||||
- R3: nothing to rewrite is refused up front; one kept pre-migration backup
|
||||
ever, later applies take a pooled safety backup instead.
|
||||
- R4: the live job loop lets a writer on another connection in.
|
||||
- R5: the digest tool loads no package, adds no log handler and leaves no
|
||||
``-wal``/``-shm`` beside a stopped database.
|
||||
- R6: a damaged document is named, with what to do, in the verify reason and
|
||||
in the job's backup failure.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sqlite3
|
||||
import subprocess # nosec B404
|
||||
import sys
|
||||
import threading
|
||||
from unittest.mock import patch
|
||||
|
||||
from database import workflow_documents as wd
|
||||
from database import workflow_storage as ws
|
||||
from database.models import PromptModel
|
||||
|
||||
try:
|
||||
from tests.test_api_workflow_storage import OPERATION, WorkflowStorageTestCase
|
||||
from tests.test_workflow_migration import MigrationTestCase
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from test_api_workflow_storage import OPERATION, WorkflowStorageTestCase
|
||||
from test_workflow_migration import MigrationTestCase
|
||||
|
||||
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
TOOL = os.path.join(ROOT, "tools", "workflow_text_digests.py")
|
||||
KEPT = "prompts.db.pre-workflow-storage_"
|
||||
WAIT = 20
|
||||
|
||||
|
||||
def disk_full(*_args, **_kwargs):
|
||||
raise sqlite3.OperationalError("database or disk is full")
|
||||
|
||||
|
||||
class TestSystemErrorsFailTheBatch(MigrationTestCase):
|
||||
def test_a_disk_error_fails_the_batch_instead_of_listing_rows(self):
|
||||
_prompt_id, expected = self.library(count=6)
|
||||
for error in (
|
||||
sqlite3.OperationalError("database or disk is full"),
|
||||
sqlite3.OperationalError("disk I/O error"),
|
||||
sqlite3.DatabaseError("database disk image is malformed"),
|
||||
TypeError("a bug"),
|
||||
):
|
||||
with self.subTest(error=str(error)):
|
||||
with patch.object(ws, "store_text", side_effect=error):
|
||||
with self.assertRaises(type(error)):
|
||||
self.migrate(max_rows=50)
|
||||
self.assertFalse(self.conn().in_transaction)
|
||||
self.assertEqual(self.inline_ids(), set(expected))
|
||||
self.assertEqual(self.texts(), expected)
|
||||
|
||||
|
||||
class TestSystemErrorsInTheJob(WorkflowStorageTestCase):
|
||||
async def test_a_disk_error_ends_the_job_failed_with_a_truthful_line(self):
|
||||
self.old_library(count=4)
|
||||
before = self.texts()
|
||||
with patch.object(ws, "store_text", side_effect=disk_full):
|
||||
_plan, _summary, job = await self.start()
|
||||
await self.finish(job)
|
||||
self.assertEqual(job.state, "failed")
|
||||
self.assertIn("The batch in progress was rolled back", job.status_line)
|
||||
self.assertEqual(job.counters["errors"], 1) # the job's one error row
|
||||
self.assertNotIn("could not be rewritten", job.status_line)
|
||||
self.assertIsNone(self.reclaim_job())
|
||||
self.assertEqual(self.texts(), before)
|
||||
|
||||
|
||||
class TestOneKeptBackup(WorkflowStorageTestCase):
|
||||
def kept(self):
|
||||
return sorted(n for n in os.listdir(self.tmpdir) if n.startswith(KEPT))
|
||||
|
||||
async def test_nothing_to_rewrite_is_said_and_refused_without_a_backup(self):
|
||||
self.save("no images")
|
||||
plan = await self.preview(OPERATION)
|
||||
self.assertEqual(plan["scope"]["count"], 0)
|
||||
self.assertEqual(plan["scope"]["label"], "Nothing to rewrite")
|
||||
reply = await self.apply(plan["token"], status=409)
|
||||
self.assertEqual(reply["code"], "nothing-to-do")
|
||||
self.assertEqual(
|
||||
reply["error"],
|
||||
"Every image record already stores its workflow once; nothing to do.",
|
||||
)
|
||||
self.assertEqual(self.kept(), [])
|
||||
self.assertEqual(self.backups(), [])
|
||||
self.assertEqual(self.api._jobs.list(), [])
|
||||
|
||||
async def test_a_cancel_and_rerun_keep_one_kept_backup_and_pool_the_next(self):
|
||||
self.old_library(count=12)
|
||||
real = self.api.db.migrate_workflow_batch
|
||||
holder = {}
|
||||
|
||||
def first_then_cancel(*args, **kwargs):
|
||||
result = real(*args, **kwargs)
|
||||
holder["job"].cancelled = True
|
||||
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")
|
||||
[kept] = self.kept()
|
||||
plan = await self.preview(OPERATION)
|
||||
self.assertEqual(plan["backup"]["policy"], "required")
|
||||
self.assertEqual(
|
||||
plan["backup"]["destination_label"], "prompts.db.backup_<date and time>"
|
||||
)
|
||||
self.assertIn(
|
||||
f"the pre-migration backup {kept} is already kept",
|
||||
" ".join(row["detail"] for row in plan["details"]),
|
||||
)
|
||||
body = await self.apply(plan["token"])
|
||||
again = self.api._jobs.get(body["job"]["id"])
|
||||
await self.finish(again)
|
||||
self.assertEqual(again.state, "finished", again.log)
|
||||
self.assertEqual(self.kept(), [kept]) # still one never-pruned copy
|
||||
[pooled] = self.backups()
|
||||
self.assertIn(f"Safety backup {pooled} verified", "\n".join(again.log))
|
||||
self.assertIn(f"pre-migration backup {kept} is kept", "\n".join(again.log))
|
||||
|
||||
|
||||
class TestLiveLoopLetsWritersIn(WorkflowStorageTestCase):
|
||||
async def test_a_writer_on_another_connection_gets_in_during_the_job(self):
|
||||
prompt_id = self.save("writer")
|
||||
conn = self.api.db.model.get_connection()
|
||||
with conn:
|
||||
for index in range(200):
|
||||
conn.execute(
|
||||
"INSERT INTO generated_images (prompt_id, image_path, filename, "
|
||||
"file_path, workflow_data) VALUES (?, ?, 'x', ?, ?)",
|
||||
(prompt_id, f"/x/{index}", f"/x/{index}", f'{{"n": {index}}}'),
|
||||
)
|
||||
started, done = threading.Event(), threading.Event()
|
||||
errors, written = [], []
|
||||
real = ws.store_text
|
||||
|
||||
def slow(conn, text, *args, **kwargs):
|
||||
started.set()
|
||||
done.wait(0.01) # a slow disk: each row takes a while
|
||||
return real(conn, text, *args, **kwargs)
|
||||
|
||||
def writer():
|
||||
started.wait(WAIT)
|
||||
own = self.api.db.model.get_connection() # this thread's connection
|
||||
own.execute("PRAGMA busy_timeout = 2000")
|
||||
try:
|
||||
for index in range(5):
|
||||
written.append(
|
||||
self.api.db.link_image_to_prompt(
|
||||
prompt_id,
|
||||
os.path.join(self.tmpdir, f"w{index}.png"),
|
||||
{"workflow": {"w": index}},
|
||||
)
|
||||
)
|
||||
except sqlite3.OperationalError as error: # pragma: no cover - RED
|
||||
errors.append(str(error))
|
||||
finally:
|
||||
self.api.db.model.close()
|
||||
|
||||
thread = threading.Thread(target=writer)
|
||||
thread.start()
|
||||
try:
|
||||
with patch.object(ws, "store_text", side_effect=slow):
|
||||
_plan, _summary, job = await self.start()
|
||||
await asyncio.wait_for(asyncio.shield(job.task), 120)
|
||||
await self.finish(job) # and Reclaim space after it
|
||||
finally:
|
||||
done.set()
|
||||
await asyncio.get_running_loop().run_in_executor(None, thread.join, WAIT)
|
||||
self.assertEqual(job.state, "finished", job.log)
|
||||
self.assertEqual(errors, [])
|
||||
self.assertEqual(len([image for image in written if image]), 5)
|
||||
|
||||
|
||||
class TestDigestToolSideEffects(MigrationTestCase):
|
||||
def run_tool(self, path, cwd):
|
||||
out = os.path.join(cwd, "out.sha")
|
||||
code = (
|
||||
"import logging, runpy, sys\n"
|
||||
"sys.argv = ['tool', sys.argv[1], sys.argv[2]]\n"
|
||||
"try:\n"
|
||||
f" runpy.run_path({TOOL!r}, run_name='__main__')\n"
|
||||
"except SystemExit:\n"
|
||||
" pass\n"
|
||||
"loaded = sorted(m for m in sys.modules\n"
|
||||
" if m.split('.')[0] in ('database', 'utils', 'py'))\n"
|
||||
"handlers = [h for n in ('', 'prompt_manager')\n"
|
||||
" for h in logging.getLogger(n).handlers]\n"
|
||||
"print('LOADED', loaded, 'HANDLERS', len(handlers))\n"
|
||||
)
|
||||
result = subprocess.run( # nosec B603
|
||||
[sys.executable, "-c", code, path, out],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
cwd=cwd,
|
||||
check=False,
|
||||
)
|
||||
self.assertEqual(result.returncode, 0, result.stderr)
|
||||
return result.stdout
|
||||
|
||||
def test_a_stopped_database_gets_no_sidecars_and_no_log(self):
|
||||
self.library(count=3)
|
||||
self.migrate()
|
||||
PromptModel.close_all_instances()
|
||||
for suffix in ("-wal", "-shm"):
|
||||
self.assertFalse(os.path.exists(self.db_path + suffix))
|
||||
cwd = os.path.join(self.tmpdir, "run")
|
||||
os.mkdir(cwd)
|
||||
before = sorted(os.listdir(self.tmpdir))
|
||||
printed = self.run_tool(self.db_path, cwd)
|
||||
self.assertEqual(sorted(os.listdir(self.tmpdir)), before)
|
||||
self.assertEqual(sorted(os.listdir(cwd)), ["out.sha"]) # no logs/ folder
|
||||
self.assertIn("LOADED [] HANDLERS 0", printed)
|
||||
with open(os.path.join(cwd, "out.sha"), encoding="utf-8") as handle:
|
||||
self.assertEqual(len(handle.read().splitlines()), 3)
|
||||
|
||||
def test_a_live_database_is_read_with_its_wal(self):
|
||||
prompt_id = self.prompt()
|
||||
self.old_form_image(prompt_id, "a.png", '{"in": "wal"}', None)
|
||||
# The connection stays open: the row may still be only in the WAL.
|
||||
cwd = os.path.join(self.tmpdir, "run")
|
||||
os.mkdir(cwd)
|
||||
self.run_tool(self.db_path, cwd)
|
||||
with open(os.path.join(cwd, "out.sha"), encoding="utf-8") as handle:
|
||||
self.assertEqual(len(handle.read().splitlines()), 1)
|
||||
|
||||
|
||||
class TestDamagedDocumentIsNamed(WorkflowStorageTestCase):
|
||||
def damage(self):
|
||||
conn = self.api.db.model.get_connection()
|
||||
with conn:
|
||||
document_id = wd.store_text(conn, '{"damaged": true}')
|
||||
conn.execute(
|
||||
"UPDATE workflow_documents SET data=x'00' WHERE id=?", (document_id,)
|
||||
)
|
||||
return document_id
|
||||
|
||||
async def test_verify_names_the_document_and_what_to_do(self):
|
||||
document_id = self.damage()
|
||||
ok, reason = PromptModel.verify_database_file(self.api.db.model.db_path)
|
||||
self.assertFalse(ok)
|
||||
self.assertEqual(reason, wd.damaged_message(document_id))
|
||||
self.assertIn(f"Workflow document {document_id} is damaged", reason)
|
||||
self.assertIn("restore a backup", reason)
|
||||
|
||||
async def test_the_backup_failure_says_which_document(self):
|
||||
self.old_library(count=2)
|
||||
document_id = self.damage()
|
||||
_plan, _summary, job = await self.start()
|
||||
await self.finish(job)
|
||||
self.assertEqual(job.state, "failed")
|
||||
self.assertEqual(job.last_event["code"], "backup-failed")
|
||||
self.assertIn(f"Workflow document {document_id} is damaged", job.status_line)
|
||||
@@ -0,0 +1,374 @@
|
||||
"""PR #315 final review fixes: I1, I2, I3, M1, M2, M3 and M5.
|
||||
|
||||
- I1: the migration always takes a verified backup, whatever the policy.
|
||||
- I2: that backup has its own name outside the pruned pool.
|
||||
- I3: ``tools/workflow_text_digests.py`` proves byte identity across the rewrite.
|
||||
- M1: a writer on another connection is never starved by the batches.
|
||||
- M2: a row that cannot be rewritten stays inline, is counted, and the rest go on.
|
||||
- M3: a damaged document is refused on restore and repaired before reuse.
|
||||
- M5: the legacy table rebuilds refuse a file with document id columns.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import os
|
||||
import re
|
||||
import sqlite3
|
||||
import subprocess # nosec B404
|
||||
import sys
|
||||
import threading
|
||||
import zlib
|
||||
from unittest.mock import patch
|
||||
|
||||
from database import workflow_documents as wd
|
||||
from database import workflow_storage as ws
|
||||
from database.models import PromptModel
|
||||
|
||||
try:
|
||||
from tests.test_api_workflow_storage import OPERATION, WorkflowStorageTestCase
|
||||
from tests.test_workflow_migration import MigrationTestCase
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from test_api_workflow_storage import OPERATION, WorkflowStorageTestCase
|
||||
from test_workflow_migration import MigrationTestCase
|
||||
|
||||
# After the API test modules, which stand in ComfyUI's ``server`` module.
|
||||
from py.api.admin import _prune_safety_backups # noqa: E402
|
||||
from py.config import PromptManagerConfig # noqa: E402
|
||||
|
||||
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
PRE_MIGRATION = re.compile(
|
||||
r"^prompts\.db\.pre-workflow-storage_\d{8}_\d{6}_\d{6}(_\d+)?$"
|
||||
)
|
||||
WAIT = 20
|
||||
|
||||
|
||||
class TestMigrationBackup(WorkflowStorageTestCase):
|
||||
def pre_migration(self):
|
||||
return sorted(
|
||||
name for name in os.listdir(self.tmpdir) if PRE_MIGRATION.match(name)
|
||||
)
|
||||
|
||||
async def test_the_backup_is_taken_even_with_the_policy_off(self):
|
||||
PromptManagerConfig.SAFETY_BACKUP = False
|
||||
self.old_library(count=4)
|
||||
plan = await self.preview(OPERATION)
|
||||
self.assertEqual(plan["backup"]["policy"], "required")
|
||||
self.assertGreater(plan["backup"]["expected_bytes"], 0)
|
||||
self.assertGreater(plan["disk"]["required_bytes"], 0)
|
||||
body = await self.apply(plan["token"])
|
||||
job = self.api._jobs.get(body["job"]["id"])
|
||||
await self.finish(job)
|
||||
self.assertEqual(job.state, "finished", job.log)
|
||||
[backup] = self.pre_migration()
|
||||
self.assertNotIn("No safety backup", "\n".join(job.log))
|
||||
self.assertIn(f"Pre-migration backup {backup} verified", "\n".join(job.log))
|
||||
|
||||
async def test_the_backup_is_outside_the_pool_never_pruned_and_restores(self):
|
||||
self.old_library(count=4)
|
||||
before = self.texts()
|
||||
plan = await self.preview(OPERATION)
|
||||
self.assertIn(
|
||||
"prompts.db.pre-workflow-storage_<date and time>",
|
||||
" ".join(row["detail"] for row in plan["details"]),
|
||||
)
|
||||
_plan, _summary, job = await self.start()
|
||||
await self.finish(job)
|
||||
[backup] = self.pre_migration()
|
||||
self.assertEqual(self.backups(), []) # not in the pruned pool
|
||||
self.assertIn(f"Pre-migration backup {backup}", job.status_line)
|
||||
path = os.path.join(self.tmpdir, backup)
|
||||
ok, reason = PromptModel.verify_database_file(path)
|
||||
self.assertTrue(ok, reason)
|
||||
_prune_safety_backups(self.api.db.model.db_path, keep=0)
|
||||
self.assertEqual(self.pre_migration(), [backup])
|
||||
# The way back: the legacy restore accepts it, old form included.
|
||||
self.api.db.model.restore_from_file(path)
|
||||
self.assertEqual(self.texts(), before)
|
||||
self.assertEqual(self.inline_count(), 4)
|
||||
|
||||
|
||||
class TestTextDigestsTool(MigrationTestCase):
|
||||
def digests(self, path, out):
|
||||
result = subprocess.run( # nosec B603
|
||||
[sys.executable, os.path.join(ROOT, "tools", "workflow_text_digests.py")]
|
||||
+ [path, out],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
cwd=self.tmpdir,
|
||||
check=False,
|
||||
)
|
||||
self.assertEqual(result.returncode, 0, result.stderr)
|
||||
with open(out, encoding="utf-8") as handle:
|
||||
return handle.read()
|
||||
|
||||
def test_digests_match_across_the_rewrite_and_read_only(self):
|
||||
_prompt_id, expected = self.library(count=12)
|
||||
PromptModel.close_all_instances()
|
||||
before = self.digests(self.db_path, os.path.join(self.tmpdir, "a.sha"))
|
||||
self.assertEqual(len(before.splitlines()), len(expected))
|
||||
self.migrate()
|
||||
PromptModel.close_all_instances()
|
||||
with open(self.db_path, "rb") as handle:
|
||||
frozen = handle.read()
|
||||
after = self.digests(self.db_path, os.path.join(self.tmpdir, "b.sha"))
|
||||
with open(self.db_path, "rb") as handle:
|
||||
self.assertEqual(handle.read(), frozen)
|
||||
self.assertEqual(after, before)
|
||||
|
||||
def test_a_changed_text_shows_in_the_diff(self):
|
||||
prompt_id = self.prompt()
|
||||
self.old_form_image(prompt_id, "a.png", '{"a": 1}', None)
|
||||
PromptModel.close_all_instances()
|
||||
before = self.digests(self.db_path, os.path.join(self.tmpdir, "a.sha"))
|
||||
conn = self.conn()
|
||||
with conn:
|
||||
conn.execute("UPDATE generated_images SET workflow_data='{\"a\":1}'")
|
||||
PromptModel.close_all_instances()
|
||||
after = self.digests(self.db_path, os.path.join(self.tmpdir, "b.sha"))
|
||||
self.assertNotEqual(after, before)
|
||||
|
||||
def test_main_in_process_hashes_undecodable_text_too(self):
|
||||
import contextlib
|
||||
import importlib.util
|
||||
import io
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"workflow_text_digests",
|
||||
os.path.join(ROOT, "tools", "workflow_text_digests.py"),
|
||||
)
|
||||
tool = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(tool)
|
||||
prompt_id = self.prompt()
|
||||
conn = self.conn()
|
||||
with conn:
|
||||
conn.execute(
|
||||
"INSERT INTO generated_images (prompt_id, image_path, filename, "
|
||||
"file_path, workflow_data) VALUES (?, '/x', 'x', '/x', "
|
||||
"CAST(x'ff' AS TEXT))",
|
||||
(prompt_id,),
|
||||
)
|
||||
out = os.path.join(self.tmpdir, "c.sha")
|
||||
printed = io.StringIO()
|
||||
with contextlib.redirect_stdout(printed):
|
||||
self.assertEqual(tool.main([self.db_path, out]), 0)
|
||||
self.assertEqual(printed.getvalue(), "1 image records\n")
|
||||
with open(out, encoding="utf-8") as handle:
|
||||
line = handle.read().split()
|
||||
self.assertEqual(line[1], hashlib.sha256(b"\xff").hexdigest())
|
||||
self.assertEqual(line[2], "-")
|
||||
|
||||
|
||||
class TestWriterIsNotStarved(MigrationTestCase):
|
||||
def test_a_writer_on_another_connection_gets_in_between_batches(self):
|
||||
prompt_id = self.prompt()
|
||||
for index in range(200):
|
||||
self.old_form_image(prompt_id, f"{index}.png", f'{{"n": {index}}}', None)
|
||||
started, done = threading.Event(), threading.Event()
|
||||
errors, written = [], []
|
||||
real = ws.store_text
|
||||
|
||||
def slow(conn, text, *args, **kwargs):
|
||||
started.set()
|
||||
done.wait(0.01) # a slow disk: each row takes a while
|
||||
return real(conn, text, *args, **kwargs)
|
||||
|
||||
def writer():
|
||||
started.wait(WAIT)
|
||||
conn = self.db.model.get_connection() # this thread's own connection
|
||||
conn.execute("PRAGMA busy_timeout = 2000")
|
||||
try:
|
||||
for index in range(5):
|
||||
written.append(
|
||||
self.db.link_image_to_prompt(
|
||||
prompt_id,
|
||||
os.path.join(self.tmpdir, f"w{index}.png"),
|
||||
{"workflow": {"w": index}},
|
||||
)
|
||||
)
|
||||
except sqlite3.OperationalError as error: # pragma: no cover - RED
|
||||
errors.append(str(error))
|
||||
finally:
|
||||
self.db.model.close()
|
||||
|
||||
thread = threading.Thread(target=writer)
|
||||
thread.start()
|
||||
with patch.object(ws, "store_text", side_effect=slow):
|
||||
self.db.migrate_workflow_storage()
|
||||
done.set()
|
||||
thread.join(WAIT)
|
||||
self.assertEqual(errors, [])
|
||||
self.assertEqual(len([image for image in written if image]), 5)
|
||||
|
||||
def test_a_batch_has_a_time_budget(self):
|
||||
prompt_id = self.prompt()
|
||||
for index in range(30):
|
||||
self.old_form_image(prompt_id, f"{index}.png", f'{{"n": {index}}}', None)
|
||||
ticks = iter(range(10**6))
|
||||
with patch.object(ws, "_clock", side_effect=lambda: next(ticks)):
|
||||
batch = ws.migrate_batch(self.conn(), 0, max_seconds=5)
|
||||
self.assertLess(batch["rows"], 30)
|
||||
self.assertGreater(batch["rows"], 0)
|
||||
|
||||
|
||||
class TestRowErrors(MigrationTestCase):
|
||||
def test_an_unreadable_row_stays_inline_and_the_rest_are_rewritten(self):
|
||||
prompt_id, expected = self.library(count=10)
|
||||
conn = self.conn()
|
||||
with conn:
|
||||
bad = conn.execute(
|
||||
"INSERT INTO generated_images (prompt_id, image_path, filename, "
|
||||
"file_path, workflow_data) VALUES (?, '/x/bad', 'bad', '/x/bad', "
|
||||
"CAST(x'7b22ff227d' AS TEXT))",
|
||||
(prompt_id,),
|
||||
).lastrowid
|
||||
result = self.migrate()
|
||||
self.assertEqual(result["errors"], [bad])
|
||||
self.assertEqual(result["rows"], len(expected))
|
||||
self.assertEqual(self.inline_ids_without(bad), set())
|
||||
raw = conn.execute(
|
||||
"SELECT CAST(workflow_data AS BLOB), workflow_document_id "
|
||||
"FROM generated_images WHERE id=?",
|
||||
(bad,),
|
||||
).fetchone()
|
||||
self.assertEqual((bytes(raw[0]), raw[1]), (b'{"\xff"}', None))
|
||||
self.assertEqual(result["documents_removed"], 0)
|
||||
|
||||
def inline_ids_without(self, image_id):
|
||||
return {
|
||||
row[0]
|
||||
for row in self.conn().execute(
|
||||
"SELECT id FROM generated_images WHERE id != ? AND "
|
||||
"(workflow_data IS NOT NULL OR prompt_metadata IS NOT NULL)",
|
||||
(image_id,),
|
||||
)
|
||||
}
|
||||
|
||||
def test_a_row_whose_text_will_not_come_back_is_kept_and_counted(self):
|
||||
prompt_id = self.prompt()
|
||||
good = self.old_form_image(prompt_id, "good.png", '{"good": 1}', None)
|
||||
bad = self.old_form_image(prompt_id, "bad.png", '{"bad": 1}', None)
|
||||
real = ws.read_stored_text
|
||||
|
||||
def lying(conn, inline, document_id):
|
||||
text = real(conn, inline, document_id)
|
||||
return "other" if text == '{"bad": 1}' else text
|
||||
|
||||
with patch.object(ws, "read_stored_text", side_effect=lying):
|
||||
result = self.migrate()
|
||||
self.assertEqual(result["errors"], [bad])
|
||||
self.assertEqual(self.inline_ids(), {bad})
|
||||
self.assertEqual(self.texts()[good][0], '{"good": 1}')
|
||||
self.assertEqual(self.texts()[bad][0], '{"bad": 1}')
|
||||
self.assertEqual(self.document_count(), 1) # the bad row's went
|
||||
|
||||
|
||||
class TestRowErrorsInTheJob(WorkflowStorageTestCase):
|
||||
async def test_the_job_reports_kept_rows_and_reclaim_space_still_runs(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/bad', 'bad', '/x/bad', "
|
||||
"CAST(x'7b22ff227d' AS TEXT))",
|
||||
(prompt_id,),
|
||||
)
|
||||
_plan, _summary, job = await self.start()
|
||||
reclaim = await self.finish(job)
|
||||
self.assertEqual(job.state, "finished", job.log)
|
||||
self.assertEqual(job.counters["changed"], 4)
|
||||
self.assertEqual(job.counters["errors"], 1)
|
||||
self.assertIn(
|
||||
"1 image record could not be rewritten and keeps its inline text",
|
||||
job.status_line,
|
||||
)
|
||||
self.assertIsNotNone(reclaim)
|
||||
self.assertEqual(reclaim.state, "finished", reclaim.log)
|
||||
|
||||
|
||||
class TestDamagedDocuments(MigrationTestCase):
|
||||
def damage(self, document_id, data=b"not zlib"):
|
||||
conn = self.conn()
|
||||
with conn:
|
||||
conn.execute(
|
||||
"UPDATE workflow_documents SET data=? WHERE id=?", (data, document_id)
|
||||
)
|
||||
|
||||
def test_verify_refuses_a_damaged_document(self):
|
||||
document_id = self.store('{"a": 1}')
|
||||
ok, _reason = PromptModel.verify_database_file(self.db_path)
|
||||
self.assertTrue(ok)
|
||||
for data in (b"not zlib", zlib.compress(b'{"a": 2}')):
|
||||
self.damage(document_id, data)
|
||||
with self.subTest(data=data):
|
||||
ok, reason = PromptModel.verify_database_file(self.db_path)
|
||||
self.assertFalse(ok)
|
||||
self.assertEqual(reason, wd.damaged_message(document_id))
|
||||
|
||||
def test_a_damaged_document_is_repaired_before_it_is_reused(self):
|
||||
prompt_id = self.prompt()
|
||||
first = self.db.link_image_to_prompt(
|
||||
prompt_id, os.path.join(self.tmpdir, "a.png"), {"workflow": {"v": 1}}
|
||||
)
|
||||
document_id = self.raw(first)["workflow_document_id"]
|
||||
self.damage(document_id)
|
||||
with self.assertLogs(wd.logger, "WARNING"):
|
||||
second = self.db.link_image_to_prompt(
|
||||
prompt_id, os.path.join(self.tmpdir, "b.png"), {"workflow": {"v": 1}}
|
||||
)
|
||||
self.assertEqual(self.raw(second)["workflow_document_id"], document_id)
|
||||
for image_id in (first, second):
|
||||
self.assertEqual(self.texts()[image_id][0], '{"v": 1}')
|
||||
|
||||
def test_store_verifies_a_reused_document_once_per_given_set(self):
|
||||
document_id = self.store("abc")
|
||||
verified = set()
|
||||
conn = self.conn()
|
||||
with conn:
|
||||
self.assertEqual(wd.store_text(conn, "abc", verified), document_id)
|
||||
self.assertEqual(verified, {document_id})
|
||||
self.damage(document_id)
|
||||
with conn: # already verified in this set: not read again
|
||||
self.assertEqual(wd.store_text(conn, "abc", verified), document_id)
|
||||
|
||||
|
||||
class TestRebuildGuards(MigrationTestCase):
|
||||
def test_the_legacy_rebuilds_refuse_document_columns(self):
|
||||
for column, kind in (("prompt_id", "TEXT"), ("file_path", None)):
|
||||
path = os.path.join(self.tmpdir, f"guard-{column}.db")
|
||||
conn = sqlite3.connect(path)
|
||||
try:
|
||||
columns = [
|
||||
"id INTEGER PRIMARY KEY AUTOINCREMENT",
|
||||
f"prompt_id {kind or 'INTEGER'} NOT NULL",
|
||||
"image_path TEXT NOT NULL",
|
||||
"filename TEXT NOT NULL",
|
||||
"generation_time TIMESTAMP",
|
||||
"workflow_data TEXT",
|
||||
"prompt_metadata TEXT",
|
||||
"parameters TEXT",
|
||||
"workflow_document_id INTEGER",
|
||||
"prompt_document_id INTEGER",
|
||||
]
|
||||
if kind is not None:
|
||||
columns.append("file_path TEXT")
|
||||
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);"
|
||||
f"CREATE TABLE generated_images ({', '.join(columns)});"
|
||||
)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
with self.subTest(column=column):
|
||||
with self.assertRaisesRegex(sqlite3.DatabaseError, "workflow document"):
|
||||
PromptModel(path)
|
||||
PromptModel.close_all_instances()
|
||||
PromptModel.reset_schema_cache()
|
||||
|
||||
|
||||
if __name__ == "__main__": # pragma: no cover
|
||||
import unittest
|
||||
|
||||
unittest.main()
|
||||
@@ -0,0 +1,94 @@
|
||||
"""Digests of every image record's stored workflow and prompt graph text.
|
||||
|
||||
Proves the workflow storage rewrite (#249 S6) kept every text byte for byte.
|
||||
Run it from this repository on the database before the rewrite (or on its
|
||||
``prompts.db.pre-workflow-storage_*`` backup) and again after, then compare::
|
||||
|
||||
python3 tools/workflow_text_digests.py prompts.db.pre-workflow-storage_<stamp> \\
|
||||
before.sha
|
||||
python3 tools/workflow_text_digests.py prompts.db after.sha
|
||||
diff before.sha after.sha # prints nothing when every text is identical
|
||||
|
||||
Each line is ``<image id> <sha256 of workflow_data> <sha256 of prompt_metadata>``
|
||||
(``-`` for none), hashing the exact stored bytes: the inline column before the
|
||||
rewrite, the document read through the one accessor after it.
|
||||
|
||||
Read-only, with no side effects: the database is opened with ``mode=ro``,
|
||||
and with ``immutable=1`` when no ``-wal``/``-shm`` file exists beside it (a
|
||||
stopped database or a backup), so SQLite creates none. A live database
|
||||
already has them and is read through its WAL. Only
|
||||
``database/workflow_documents.py`` is loaded (by path; it needs only the
|
||||
standard library), so no package start-up runs and no log file or handler is
|
||||
added. Run it as the user that runs ComfyUI.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import importlib.util
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def _load_documents():
|
||||
"""``database/workflow_documents.py`` alone, without the ``database`` package."""
|
||||
path = _ROOT / "database" / "workflow_documents.py"
|
||||
spec = importlib.util.spec_from_file_location("_pm_workflow_documents", path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def database_uri(db_path):
|
||||
"""``mode=ro``, plus ``immutable=1`` when nothing has the database open."""
|
||||
path = os.path.abspath(db_path)
|
||||
uri = Path(path).as_uri() + "?mode=ro"
|
||||
if not any(os.path.exists(path + suffix) for suffix in ("-wal", "-shm")):
|
||||
uri += "&immutable=1"
|
||||
return uri
|
||||
|
||||
|
||||
def _digest(value):
|
||||
if value is None:
|
||||
return "-"
|
||||
raw = value if isinstance(value, bytes) else value.encode("utf-8")
|
||||
return hashlib.sha256(raw).hexdigest()
|
||||
|
||||
|
||||
def write_digests(db_path, out):
|
||||
"""Write one digest line per image record of ``db_path`` to ``out``."""
|
||||
documents = _load_documents()
|
||||
conn = sqlite3.connect(database_uri(db_path), uri=True)
|
||||
try:
|
||||
conn.row_factory = sqlite3.Row
|
||||
conn.text_factory = bytes # inline text exactly as stored, even undecodable
|
||||
reader = documents.DocumentReader(conn)
|
||||
count = 0
|
||||
with open(out, "w", encoding="utf-8") as handle:
|
||||
for row in conn.execute("SELECT * FROM generated_images ORDER BY id"):
|
||||
texts = documents.resolve_image_row(reader, row)
|
||||
handle.write(
|
||||
f"{row['id']} {_digest(texts.get('workflow_data'))} "
|
||||
f"{_digest(texts.get('prompt_metadata'))}\n"
|
||||
)
|
||||
count += 1
|
||||
return count
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def main(argv=None):
|
||||
parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
|
||||
parser.add_argument("database", help="prompts.db or one of its backups")
|
||||
parser.add_argument("out", help="file to write the digests to")
|
||||
args = parser.parse_args(argv)
|
||||
count = write_digests(args.database, args.out)
|
||||
print(f"{count} image records")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -38,6 +38,11 @@ from typing import Dict, Any
|
||||
|
||||
from .logging_config import get_logger
|
||||
|
||||
try:
|
||||
from ..database.workflow_documents import DocumentReader, resolve_image_row
|
||||
except ImportError: # loaded as a top-level ``utils`` package
|
||||
from database.workflow_documents import DocumentReader, resolve_image_row
|
||||
|
||||
|
||||
class GalleryDiagnostics:
|
||||
"""Diagnostics for the gallery system.
|
||||
@@ -221,7 +226,8 @@ class GalleryDiagnostics:
|
||||
ORDER BY gi.generation_time DESC
|
||||
LIMIT 5
|
||||
""")
|
||||
recent_images = [dict(row) for row in cursor.fetchall()]
|
||||
reader = DocumentReader(conn)
|
||||
recent_images = [resolve_image_row(reader, row) for row in cursor]
|
||||
|
||||
self.logger.info(f" [TIME] Recent images: {len(recent_images)}")
|
||||
for img in recent_images:
|
||||
|
||||
Reference in New Issue
Block a user