Compare commits
9
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c1e74bd0f1 | ||
|
|
e43157e684 | ||
|
|
44c2fc4133 | ||
|
|
78e6ed7a2d | ||
|
|
297b2060ee | ||
|
|
06099579e6 | ||
|
|
cfa47ccac8 | ||
|
|
64f5b91eee | ||
|
|
f68e92b3b6 |
@@ -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))
|
||||
|
||||
+124
-3
@@ -1932,8 +1932,8 @@ pending count (`GET /prompt_manager/review/items?limit=1`, its `total`).
|
||||
unavailable file's empty media button keeps the concept's 1.33 media area, so its
|
||||
card lines up with the others. A long file name in the panel's `p.small` wraps inside
|
||||
the panel (`overflow-wrap: anywhere`), where the concept's short sample name never
|
||||
overflows. A queue card shows the Images listing's thumbnail when the server knows one
|
||||
(`thumbnail_url`, the same URL the Images page uses), else the original; a re-read
|
||||
overflows. A queue card shows the Images listing's thumbnail (`thumbnail_url`, the
|
||||
same URL the Images page uses; a preview not made yet is made on first view); a re-read
|
||||
rebuilds a card whose media, target or availability changed (a root came or went), so
|
||||
Open and Tag follow it. When an unanswered decision's item leaves the queue in a
|
||||
re-read, the decision was recorded: the page finishes as a recorded decision does (its
|
||||
@@ -2230,6 +2230,59 @@ through the Activity store, which toasts "Thumbnail generation started."; the jo
|
||||
in Activity. A start refused because another job holds the slot toasts the D57 busy
|
||||
link toast and stays. The legacy gallery's thumbnail tools are unchanged.
|
||||
|
||||
Automatic previews (#249, `py/api/thumbnail_auto.py`). The approved mock has no
|
||||
repeated global thumbnail prompt (`docs/concepts/workspace/README.md`), so the
|
||||
workspace never asks; previews are a cache the server keeps:
|
||||
|
||||
- Every listing gives an image or video a thumbnail URL (`media_thumbnail_url`: Images,
|
||||
saved views and favorites, Library cards and galleries, Needs review, duplicate
|
||||
groups). An existing preview's URL keeps its `v=<mtime>`; a missing one is listed
|
||||
without `v=`.
|
||||
- `GET /images/serve/thumbnails/<rel>_thumb<ext>` for a missing preview is refused as
|
||||
text first (traversal, a backslash, a drive or stream colon, empty or dot segments, so
|
||||
`D:/x`, `\\server\share` and a leading `/` never reach a joined path), then admitted
|
||||
on the event loop (`_PreviewGate`): 3 previews are written at once on the API's own
|
||||
`pm-preview` pool (never the default executor), one flight runs per file (the history
|
||||
route shares the flight key and the admission), and at most 64 requests wait; one
|
||||
more gets its fallback at once without listing a folder. Inside the admission the
|
||||
preview maps back to its media file under the same root (exact names, then one folder
|
||||
listing for a mixed-case extension; a regular file inside the root, never through a
|
||||
link) and is made with the job's defaults (512 px, medium). When it cannot be made an
|
||||
image's original is served and a video's placeholder frame (a video is never sent to
|
||||
an `<img>`), both `Cache-Control: no-store`; cards and the inspector cover also swap a
|
||||
video back to its own element when its preview image fails. Only a file that cannot
|
||||
be decoded is remembered for the session (path, size, mtime; retried once it changes)
|
||||
and counted as `failed` in `/images/thumbnails/status`, which Overview's row adds
|
||||
("N missing thumbnails · M failed"); I/O errors and a preview another program holds
|
||||
open (Windows) are not. A preview folder (or any link segment of it) resolving out of
|
||||
the root is never written. Media inside a nested `thumbnails` folder is listed with no
|
||||
thumbnail and keeps its original. The listing's `thumbnail_ready` says whether the
|
||||
preview file exists (the legacy gallery's dot and "Fast").
|
||||
- The video decoder falls through OpenCV, ffmpeg (a tracked process, killed on
|
||||
shutdown) and the placeholder frame: an OpenCV that cannot open or read a clip (AV1)
|
||||
no longer ends the run.
|
||||
- The image monitor hands every settled output to `queue_capture_thumbnail`, which
|
||||
queues it for one background worker: each wake takes at most Worker threads files and
|
||||
writes them on that many threads; a queued file is not queued twice; the queue holds
|
||||
10,000 and a file left out is made on view.
|
||||
- A finished scan that added files counts the missing previews of the sources it
|
||||
walked: 200 or more start the thumbnail job (missing only, visible and cancellable in
|
||||
Activity, never silent bulk work); fewer go through the worker.
|
||||
- Holds: automatic previews write only `thumbnails/` files and no records, so they run
|
||||
beside a scan, image tagging, the thumbnail job and the Needs review recording (S5
|
||||
M7). They wait (a view gets the original, a queued file is left for its view) while a
|
||||
restore, a duplicate removal, a sources change, the gallery's thumbnail clear or an
|
||||
`exclusive` job runs (`_auto_thumbnail_blocker`), checked when a write starts (after
|
||||
its slot) and again just before its rename.
|
||||
- Shutdown is bounded and logged: the worker stops at a file boundary, ffmpeg is
|
||||
killed, an idle worker leaves (the next output starts a fresh one), the view pool
|
||||
shuts down, and after the app's cleanup a late monitor event
|
||||
starts nothing. A package reload stops the workers the previous import left.
|
||||
- Cards and strips load previews only: Images in every layout and size, Library cards,
|
||||
Needs review cards and the inspector's cover (a video's cover is its preview frame).
|
||||
The media viewer's main media and downloads load the original
|
||||
(`workspace-thumbnails.spec.mjs` records every media request; `tests/test_thumbnail_auto.py`).
|
||||
|
||||
Overview (`system-overview.mjs`) reads `GET /prompt_manager/system/overview` once per
|
||||
visit through `latestOnly()` on the section scope; a late answer after the section left
|
||||
paints nothing. Until it answers, the section shows "Loading the overview…"; a failure
|
||||
@@ -2277,7 +2330,8 @@ Recorded deviations (plan D60–D67, D79–D81 renumbered D92–D99 and D111–D
|
||||
"Backup status" reads "No retained backup record." until verified backups exist (S8).
|
||||
The thumbnail row is "Checking thumbnails…" until its own read answers, then "No
|
||||
missing thumbnails", "N missing thumbnails" (the pending count of
|
||||
`/images/thumbnails/status`) or "Thumbnail status unavailable"; Review thumbnails
|
||||
`/images/thumbnails/status`, followed by " · M failed" when automatic previews
|
||||
failed this session) or "Thumbnail status unavailable"; Review thumbnails
|
||||
opens the thumbnail dialog (D99). "One display configuration" says "Grid
|
||||
and density settings are shared across views." (no paging). The sources row is
|
||||
"N output sources" with "Monitoring and enrichment share one source list.", prefixed
|
||||
@@ -2708,6 +2762,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
|
||||
|
||||
+8
-1
@@ -132,11 +132,18 @@ try:
|
||||
_review_api = globals().get("api")
|
||||
if _review_api is not None:
|
||||
_global_image_monitor.set_review_recorder(_review_api.record_review_file)
|
||||
# Every settled output's preview is queued for the API's background
|
||||
# worker (#249): the workspace never asks to generate thumbnails.
|
||||
_global_image_monitor.set_thumbnail_queue(_review_api.queue_capture_thumbnail)
|
||||
_review_host = getattr(config.server_instance, "app", None)
|
||||
if _review_host is not None and hasattr(_review_host, "on_cleanup"):
|
||||
|
||||
async def _stop_review_recording(app, monitor=_global_image_monitor):
|
||||
async def _stop_review_recording(
|
||||
app, monitor=_global_image_monitor, api=_review_api
|
||||
):
|
||||
monitor.set_review_recorder(None)
|
||||
monitor.set_thumbnail_queue(None)
|
||||
api.close_auto_thumbnails(final=True)
|
||||
|
||||
_review_host.on_cleanup.insert(0, _stop_review_recording)
|
||||
|
||||
|
||||
@@ -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."
|
||||
)
|
||||
@@ -9,11 +9,13 @@ Split into domain-specific mixins for maintainability:
|
||||
- JobRoutesMixin: the session job registry behind Activity (py/api/jobs.py)
|
||||
- ImageTagJobsMixin: bulk Add tags / Auto-tag as image-tags jobs
|
||||
- ThumbnailJobsMixin: thumbnail plan, thumbnail job and the gallery stream on it
|
||||
- AutoThumbnailMixin: previews made on first view, on capture and after a scan
|
||||
- ReviewRoutesMixin: the Needs review queue, item detail and decisions
|
||||
- ReviewRecordingMixin: Needs review recording for the scan and the monitor
|
||||
- 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
|
||||
"""
|
||||
|
||||
@@ -47,8 +49,10 @@ from .settings_routes import SettingsRoutesMixin
|
||||
from .job_routes import JobRoutesMixin
|
||||
from .image_tag_jobs import ImageTagJobsMixin
|
||||
from .thumbnail_jobs import ThumbnailJobsMixin
|
||||
from .thumbnail_auto import AutoThumbnailMixin
|
||||
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
|
||||
@@ -289,8 +293,10 @@ class PromptManagerAPI(
|
||||
JobRoutesMixin,
|
||||
ImageTagJobsMixin,
|
||||
ThumbnailJobsMixin,
|
||||
AutoThumbnailMixin,
|
||||
SystemRoutesMixin,
|
||||
MaintenanceRoutesMixin,
|
||||
WorkflowStorageRoutesMixin,
|
||||
DuplicateFilesMixin,
|
||||
LoggingRoutesMixin,
|
||||
AutotagRoutesMixin,
|
||||
|
||||
@@ -1504,6 +1504,9 @@ class AdminRoutesMixin:
|
||||
verb = "cancelled" if state == "cancelled" else "finished"
|
||||
job.log_line(f"Scan {verb} · {_plural(job.processed, 'file')} processed")
|
||||
job.finish(state)
|
||||
if state == "finished":
|
||||
# After the scan, so the guarded thumbnail job is admitted (#249).
|
||||
await self._hand_off_scan_thumbnails(job, output_dirs)
|
||||
|
||||
async def _stream_scan_job(self, request, job):
|
||||
"""Stream ``job`` as SSE; stop quietly if the client goes away."""
|
||||
|
||||
+42
-20
@@ -734,7 +734,10 @@ class HistoryRoutesMixin:
|
||||
raise HistoryReadError(422, "response-limit") from None
|
||||
else:
|
||||
prepared = _prepare_current_media(
|
||||
resolved, scope, thumbnail=kind == "thumbnail"
|
||||
resolved,
|
||||
scope,
|
||||
thumbnail=kind == "thumbnail",
|
||||
**media_options,
|
||||
)
|
||||
_check_current(scope, self, fence, epoch, generation)
|
||||
return prepared
|
||||
@@ -761,25 +764,44 @@ class HistoryRoutesMixin:
|
||||
_check_current(scope, self, fence, epoch, generation)
|
||||
raise
|
||||
|
||||
pending = asyncio.create_task(self._run_in_executor(project))
|
||||
try:
|
||||
prepared = await asyncio.shield(pending)
|
||||
except asyncio.CancelledError:
|
||||
pending.add_done_callback(_close_prepared_result)
|
||||
raise
|
||||
except (
|
||||
HistoryReadError,
|
||||
HistoryEpochChanged,
|
||||
sqlite3.Error,
|
||||
OSError,
|
||||
ValueError,
|
||||
):
|
||||
raise
|
||||
except Exception as error:
|
||||
self.logger.warning(
|
||||
"Current history media unavailable (%s)", type(error).__name__
|
||||
)
|
||||
raise _current_unavailable("temporary_unavailable") from None
|
||||
async def execute():
|
||||
pending = asyncio.create_task(self._run_in_executor(project))
|
||||
try:
|
||||
return await asyncio.shield(pending)
|
||||
except asyncio.CancelledError:
|
||||
pending.add_done_callback(_close_prepared_result)
|
||||
raise
|
||||
except (
|
||||
HistoryReadError,
|
||||
HistoryEpochChanged,
|
||||
sqlite3.Error,
|
||||
OSError,
|
||||
ValueError,
|
||||
):
|
||||
raise
|
||||
except Exception as error:
|
||||
self.logger.warning(
|
||||
"Current history media unavailable (%s)", type(error).__name__
|
||||
)
|
||||
raise _current_unavailable("temporary_unavailable") from None
|
||||
|
||||
# A missing preview (#249) is noted on the first read, made through
|
||||
# the serve route's admission and single flight, then read again. A
|
||||
# video whose preview cannot be made is the placeholder frame.
|
||||
needed = []
|
||||
media_options = {}
|
||||
if kind == "thumbnail":
|
||||
media_options = {
|
||||
"ensure_preview": lambda *preview: needed.append(preview),
|
||||
"placeholder": self._video_placeholder_bytes(),
|
||||
}
|
||||
prepared = await execute()
|
||||
if needed:
|
||||
if not isinstance(prepared, bytes):
|
||||
prepared.close()
|
||||
await self._admit_preview(*needed[0])
|
||||
media_options["ensure_preview"] = None
|
||||
prepared = await execute()
|
||||
if isinstance(prepared, bytes):
|
||||
return _success_response(prepared)
|
||||
return await _current_media_response(request, prepared)
|
||||
|
||||
@@ -5,6 +5,7 @@ from dataclasses import dataclass
|
||||
import errno
|
||||
import hashlib
|
||||
import hmac
|
||||
import io
|
||||
import json
|
||||
import mimetypes
|
||||
import os
|
||||
@@ -309,17 +310,35 @@ class _OpenedCurrentMedia:
|
||||
self.handle.close()
|
||||
|
||||
|
||||
def _preview_writable(root: Path, preview: Path) -> bool:
|
||||
"""A preview may be written: ``thumbnails`` and every link segment of the
|
||||
preview's folder resolve inside *root* (the serve route's rule, #249)."""
|
||||
folder = root / "thumbnails"
|
||||
return _canonical(folder).is_relative_to(_canonical(root)) and _canonical(
|
||||
preview
|
||||
).is_relative_to(_canonical(folder))
|
||||
|
||||
|
||||
def _prepare_current_media(
|
||||
resolved: ResolvedCurrentCandidate, scope: CurrentOutputScope, *, thumbnail: bool
|
||||
resolved: ResolvedCurrentCandidate,
|
||||
scope: CurrentOutputScope,
|
||||
*,
|
||||
thumbnail: bool,
|
||||
ensure_preview=None,
|
||||
placeholder: bytes | None = None,
|
||||
) -> _OpenedCurrentMedia:
|
||||
"""Open the current file, or its preview for ``thumbnail``.
|
||||
|
||||
A missing preview that may be written is handed to
|
||||
``ensure_preview(src, dst, is_video)`` (the route makes it and reads
|
||||
again, #249). Without one an image's original is opened and a video is
|
||||
the ``placeholder`` JPEG, never the video inside an ``<img>``.
|
||||
"""
|
||||
target = resolved.canonical_target
|
||||
if thumbnail:
|
||||
relative = target.relative_to(resolved.gallery_root)
|
||||
extension = (
|
||||
".jpg"
|
||||
if target.suffix.lower() in VIDEO_EXTENSIONS
|
||||
else target.suffix.lower()
|
||||
)
|
||||
is_video = target.suffix.lower() in VIDEO_EXTENSIONS
|
||||
extension = ".jpg" if is_video else target.suffix.lower()
|
||||
preview = (
|
||||
resolved.gallery_root
|
||||
/ "thumbnails"
|
||||
@@ -328,7 +347,14 @@ def _prepare_current_media(
|
||||
try:
|
||||
preview.stat()
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
if ensure_preview is not None and _preview_writable(
|
||||
resolved.gallery_root, preview
|
||||
):
|
||||
ensure_preview(target, preview, is_video)
|
||||
if is_video and placeholder is not None:
|
||||
return _OpenedCurrentMedia(
|
||||
io.BytesIO(placeholder), len(placeholder), "image/jpeg"
|
||||
)
|
||||
else:
|
||||
target = _canonical(preview)
|
||||
handle = _open_current_target(target, resolved.namespace, scope)
|
||||
|
||||
+259
-137
@@ -110,6 +110,8 @@ def _card_header(path: Path) -> tuple:
|
||||
|
||||
# The legacy gallery page's thumbnail generation, as Activity lists it (D80).
|
||||
THUMBNAIL_GENERATION_TITLE = "Thumbnail generation (gallery page)"
|
||||
# Its clear: automatic previews are not written while it runs.
|
||||
THUMBNAIL_CLEAR_TITLE = "Thumbnail clear (gallery page)"
|
||||
|
||||
|
||||
class _CardFacts(NamedTuple):
|
||||
@@ -261,13 +263,32 @@ def _temp_sibling(dst):
|
||||
return dst.with_name(f".{dst.stem}.part-{token}{dst.suffix}")
|
||||
|
||||
|
||||
def _atomic_save(img, dst, fmt=None, **save_kwargs):
|
||||
class PreviewWriteHeld(Exception):
|
||||
"""A preview write stopped before its rename: ``guard()`` said no (#249)."""
|
||||
|
||||
|
||||
# The result error of a write held before its rename (never remembered).
|
||||
PREVIEW_HELD = "Preview writes are held while another job runs"
|
||||
|
||||
|
||||
def _replace_unless_held(tmp, dst, guard=None):
|
||||
"""Rename *tmp* over *dst*, unless ``guard()`` is false (PreviewWriteHeld)."""
|
||||
if guard is not None and not guard():
|
||||
raise PreviewWriteHeld(PREVIEW_HELD)
|
||||
os.replace(tmp, dst)
|
||||
|
||||
|
||||
def _atomic_save(img, dst, fmt=None, *, guard=None, **save_kwargs):
|
||||
"""Save a PIL image to *dst* without a partially written file ever being
|
||||
visible at that path: write a temp sibling, then rename over it."""
|
||||
visible at that path: write a temp sibling, then rename over it.
|
||||
|
||||
``guard()`` (automatic previews) is asked again just before the rename;
|
||||
false leaves *dst* untouched and raises PreviewWriteHeld.
|
||||
"""
|
||||
tmp = _temp_sibling(dst)
|
||||
try:
|
||||
img.save(tmp, fmt, **save_kwargs)
|
||||
os.replace(tmp, dst)
|
||||
_replace_unless_held(tmp, dst, guard)
|
||||
finally:
|
||||
if tmp.exists():
|
||||
tmp.unlink()
|
||||
@@ -305,11 +326,11 @@ def _ffmpeg_qscale(encoding):
|
||||
return str(_FFMPEG_QSCALE.get(encoding["quality"], 4))
|
||||
|
||||
|
||||
def _write_image_thumbnail(src, dst, thumbnail_size, encoding=None):
|
||||
def _write_image_thumbnail(src, dst, thumbnail_size, encoding=None, guard=None):
|
||||
"""Resize *src* into *dst* (blocking PIL work), atomically.
|
||||
|
||||
``encoding`` is one of THUMBNAIL_ENCODINGS (medium when None); a PNG
|
||||
preview keeps only its ``optimize`` flag.
|
||||
preview keeps only its ``optimize`` flag. ``guard`` see _atomic_save.
|
||||
"""
|
||||
encoding = encoding or DEFAULT_THUMBNAIL_ENCODING
|
||||
with Image.open(src) as img:
|
||||
@@ -319,7 +340,50 @@ def _write_image_thumbnail(src, dst, thumbnail_size, encoding=None):
|
||||
save_kwargs = dict(encoding)
|
||||
if dst.suffix.lower() == ".png":
|
||||
save_kwargs = {"optimize": encoding["optimize"]}
|
||||
_atomic_save(img, dst, **save_kwargs)
|
||||
_atomic_save(img, dst, guard=guard, **save_kwargs)
|
||||
|
||||
|
||||
def _video_placeholder_image(thumbnail_size):
|
||||
"""The placeholder frame for a video no decoder could read: a play mark."""
|
||||
from PIL import ImageDraw, ImageFont
|
||||
|
||||
img = Image.new("RGB", thumbnail_size, color=(50, 50, 50))
|
||||
draw = ImageDraw.Draw(img)
|
||||
center_x, center_y = thumbnail_size[0] // 2, thumbnail_size[1] // 2
|
||||
triangle_size = min(thumbnail_size) // 4
|
||||
points = [
|
||||
(center_x - triangle_size // 2, center_y - triangle_size // 2),
|
||||
(center_x - triangle_size // 2, center_y + triangle_size // 2),
|
||||
(center_x + triangle_size // 2, center_y),
|
||||
]
|
||||
draw.polygon(points, fill=(255, 255, 255))
|
||||
try:
|
||||
font = ImageFont.load_default()
|
||||
text = "VIDEO"
|
||||
bbox = draw.textbbox((0, 0), text, font=font)
|
||||
text_width = bbox[2] - bbox[0]
|
||||
draw.text(
|
||||
(center_x - text_width // 2, center_y + triangle_size // 2 + 10),
|
||||
text,
|
||||
fill=(255, 255, 255),
|
||||
font=font,
|
||||
)
|
||||
except (OSError, AttributeError):
|
||||
pass
|
||||
return img
|
||||
|
||||
|
||||
def _is_decode_error(exc):
|
||||
"""True when *exc* says the file itself cannot be decoded (#249 M4).
|
||||
|
||||
Pillow's decoder errors are OSErrors without an errno; a real I/O error
|
||||
(disk full, a share gone) carries one and is worth another try.
|
||||
"""
|
||||
if isinstance(exc, (UnidentifiedImageError, Image.DecompressionBombError)):
|
||||
return True
|
||||
if isinstance(exc, SyntaxError):
|
||||
return True
|
||||
return isinstance(exc, OSError) and exc.errno is None
|
||||
|
||||
|
||||
def _new_thumbnail_stats():
|
||||
@@ -710,23 +774,51 @@ def _iter_media_files(output_path, max_files):
|
||||
yield root_path / name
|
||||
|
||||
|
||||
def media_thumbnail_url(output_path, rel_path, root_index, *, root_scope=None):
|
||||
"""The Images listing's thumbnail URL for *rel_path* under *output_path*,
|
||||
or None when no thumbnail exists (one blocking stat). The Needs review
|
||||
queue cards use the same URL (#249 F, M10)."""
|
||||
def media_thumbnail_rel(rel_path):
|
||||
"""``thumbnails/<rel>_thumb<ext>`` (POSIX) for the media file *rel_path*.
|
||||
|
||||
The extension is lower case, and a video's preview is a ``.jpg``: the
|
||||
name every thumbnail writer (the plan, the job, automatic previews) uses.
|
||||
"""
|
||||
rel_path = Path(rel_path)
|
||||
extension = rel_path.suffix.lower()
|
||||
thumb_rel = (
|
||||
return (
|
||||
f"thumbnails/{rel_path.with_suffix('').as_posix()}_thumb"
|
||||
f"{'.jpg' if extension in VIDEO_EXTENSIONS else extension}"
|
||||
)
|
||||
|
||||
|
||||
def media_thumbnail(output_path, rel_path, root_index, *, root_scope=None):
|
||||
"""``(thumbnail_url, ready)`` for *rel_path* under *output_path* (one
|
||||
blocking stat). The Images listing, the Needs review queue cards and the
|
||||
Library use the same URL (#249 F, M10).
|
||||
|
||||
An existing preview's URL carries its mtime as ``v=`` and ``ready`` is
|
||||
True. A missing one is still listed (without ``v=``): the serve route
|
||||
makes it on first view and falls back when it cannot
|
||||
(py/api/thumbnail_auto.py), so no card loads the full file because its
|
||||
preview was not made yet. Media inside a nested ``thumbnails`` folder is
|
||||
never previewed: ``(None, False)``, and its card shows the original.
|
||||
"""
|
||||
rel_path = Path(rel_path)
|
||||
if "thumbnails" in rel_path.parts[:-1]:
|
||||
return None, False
|
||||
thumb_rel = media_thumbnail_rel(rel_path)
|
||||
thumbnail_url = thumbnail_url_for(thumb_rel, Path(output_path) / thumb_rel)
|
||||
ready = thumbnail_url is not None
|
||||
if thumbnail_url is None:
|
||||
return None
|
||||
thumbnail_url += f"&root={root_index}"
|
||||
quoted = urllib.parse.quote(thumb_rel, safe="/")
|
||||
thumbnail_url = f"/prompt_manager/images/serve/{quoted}?root={root_index}"
|
||||
else:
|
||||
thumbnail_url += f"&root={root_index}"
|
||||
if root_scope is not None:
|
||||
thumbnail_url += f"&root_scope={root_scope}"
|
||||
return thumbnail_url
|
||||
return thumbnail_url, ready
|
||||
|
||||
|
||||
def media_thumbnail_url(output_path, rel_path, root_index, *, root_scope=None):
|
||||
"""The thumbnail URL of ``media_thumbnail`` (None for nested previews)."""
|
||||
return media_thumbnail(output_path, rel_path, root_index, root_scope=root_scope)[0]
|
||||
|
||||
|
||||
def _output_image_entry(
|
||||
@@ -744,7 +836,7 @@ def _output_image_entry(
|
||||
extension = media_path.suffix.lower()
|
||||
is_video = extension in VIDEO_EXTENSIONS
|
||||
|
||||
thumbnail_url = media_thumbnail_url(
|
||||
thumbnail_url, thumbnail_ready = media_thumbnail(
|
||||
output_path, rel_path, root_index, root_scope=root_scope
|
||||
)
|
||||
encoded_path = urllib.parse.quote(rel_path.as_posix(), safe="/")
|
||||
@@ -763,6 +855,8 @@ def _output_image_entry(
|
||||
"root_index": root_index,
|
||||
"url": url,
|
||||
"thumbnail_url": thumbnail_url,
|
||||
# Whether the preview file exists yet (the legacy gallery's dot, #249).
|
||||
"thumbnail_ready": thumbnail_ready,
|
||||
"size": stat.st_size,
|
||||
"modified_time": stat.st_mtime,
|
||||
"extension": extension,
|
||||
@@ -1682,6 +1776,11 @@ class ImageRoutesMixin:
|
||||
if image_path.is_file():
|
||||
return self._file_response(image_path)
|
||||
|
||||
# A listed preview that was not made yet is made now (#249).
|
||||
preview = await self._serve_missing_thumbnail(output_dirs, filepath)
|
||||
if preview is not None:
|
||||
return preview
|
||||
|
||||
# ComfyUI's cache middleware would otherwise let browsers cache this
|
||||
# miss for an hour, hiding a thumbnail generated a moment later.
|
||||
return web.json_response(
|
||||
@@ -1909,6 +2008,8 @@ class ImageRoutesMixin:
|
||||
"total_images": pending + plan["skipped"],
|
||||
"pending": pending,
|
||||
"ready": plan["skipped"],
|
||||
# Files whose automatic preview failed this session (#249).
|
||||
"failed": self._auto_thumbnail_failures(),
|
||||
"collection_id": hashlib.sha256(collection.encode()).hexdigest()[
|
||||
:16
|
||||
],
|
||||
@@ -2003,6 +2104,8 @@ class ImageRoutesMixin:
|
||||
encoding=None,
|
||||
rebuild=False,
|
||||
stop=None,
|
||||
guard=None,
|
||||
track=None,
|
||||
):
|
||||
"""Create one thumbnail (blocking). Never raises; returns a result dict.
|
||||
|
||||
@@ -2010,7 +2113,11 @@ class ImageRoutesMixin:
|
||||
and touches nothing (a job's file boundary). ``rebuild`` writes a
|
||||
fresh preview again and never replaces an existing video preview with
|
||||
the placeholder. Every write is atomic, so a failure keeps the old
|
||||
preview.
|
||||
preview. ``guard()`` false just before the rename gives ``action:
|
||||
"error"`` with ``held: True`` and writes nothing; ``track(process)``
|
||||
is handed the ffmpeg process (both for automatic previews, #249). An
|
||||
error result carries ``decode: True`` when the file itself cannot be
|
||||
decoded (an I/O error does not).
|
||||
"""
|
||||
result = {
|
||||
"file": src.name,
|
||||
@@ -2027,22 +2134,43 @@ class ImageRoutesMixin:
|
||||
dst.parent.mkdir(parents=True, exist_ok=True)
|
||||
if is_video:
|
||||
made = self._generate_video_thumbnail(
|
||||
src, dst, thumbnail_size, encoding, keep_existing=rebuild
|
||||
src,
|
||||
dst,
|
||||
thumbnail_size,
|
||||
encoding,
|
||||
keep_existing=rebuild,
|
||||
guard=guard,
|
||||
track=track,
|
||||
)
|
||||
if not made:
|
||||
return {**result, "action": "error", "error": "Video decode failed"}
|
||||
return {
|
||||
**result,
|
||||
"action": "error",
|
||||
"error": "Video decode failed",
|
||||
"decode": True,
|
||||
}
|
||||
else:
|
||||
_write_image_thumbnail(src, dst, thumbnail_size, encoding)
|
||||
_write_image_thumbnail(src, dst, thumbnail_size, encoding, guard)
|
||||
return {**result, "action": "generated"}
|
||||
except PreviewWriteHeld:
|
||||
return {**result, "action": "error", "error": PREVIEW_HELD, "held": True}
|
||||
except UnidentifiedImageError:
|
||||
return {**result, "action": "error", "error": "Not a valid image file"}
|
||||
return {
|
||||
**result,
|
||||
"action": "error",
|
||||
"error": "Not a valid image file",
|
||||
"decode": True,
|
||||
}
|
||||
except PermissionError as e:
|
||||
message = _preview_write_error(e)
|
||||
self.logger.warning(f"Thumbnail failed for {src.name}: {message}")
|
||||
return {**result, "action": "error", "error": message}
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Thumbnail failed for {src.name}", exc_info=True)
|
||||
return {**result, "action": "error", "error": _safe_error(e)}
|
||||
failed = {**result, "action": "error", "error": _safe_error(e)}
|
||||
if _is_decode_error(e):
|
||||
failed["decode"] = True
|
||||
return failed
|
||||
|
||||
async def _emit_sse(self, response, event, data) -> bool:
|
||||
"""Write one SSE frame. Returns False once the client has gone away."""
|
||||
@@ -2061,147 +2189,59 @@ class ImageRoutesMixin:
|
||||
thumbnail_size,
|
||||
encoding=None,
|
||||
keep_existing=False,
|
||||
*,
|
||||
guard=None,
|
||||
track=None,
|
||||
):
|
||||
"""Generate thumbnail from video file. Returns True if successful.
|
||||
|
||||
``encoding`` is one of THUMBNAIL_ENCODINGS (medium when None); it
|
||||
encodes the OpenCV frame, ffmpeg's frame (as ``-q:v``) and the
|
||||
placeholder alike. Both decoders fit the frame inside
|
||||
``thumbnail_size`` keeping its aspect ratio. With ``keep_existing`` an
|
||||
existing preview is kept, and the run fails, when no frame could be
|
||||
decoded, rather than replaced by the placeholder. A PermissionError
|
||||
while replacing the preview propagates (``_generate_one`` names it).
|
||||
Tries OpenCV, then ffmpeg, then writes the placeholder frame: a
|
||||
decoder that cannot open or read the clip (an AV1 file OpenCV was
|
||||
built without) falls through to the next one (#249 I1). ``encoding``
|
||||
is one of THUMBNAIL_ENCODINGS (medium when None); it encodes the
|
||||
OpenCV frame, ffmpeg's frame (as ``-q:v``) and the placeholder alike.
|
||||
Both decoders fit the frame inside ``thumbnail_size`` keeping its
|
||||
aspect ratio. With ``keep_existing`` an existing preview is kept, and
|
||||
the run fails, when no frame could be decoded, rather than replaced
|
||||
by the placeholder. A PermissionError while replacing the preview
|
||||
propagates (``_generate_one`` names it), as does PreviewWriteHeld
|
||||
when ``guard()`` refuses the rename. ``track(process)`` is handed the
|
||||
ffmpeg process so a shutdown can kill it.
|
||||
"""
|
||||
encoding = encoding or DEFAULT_THUMBNAIL_ENCODING
|
||||
try:
|
||||
# Try using OpenCV first (most reliable)
|
||||
try:
|
||||
import cv2
|
||||
|
||||
cap = cv2.VideoCapture(str(video_path))
|
||||
if not cap.isOpened():
|
||||
return False
|
||||
|
||||
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
target_frame = max(1, int(frame_count * 0.1))
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, target_frame)
|
||||
|
||||
ret, frame = cap.read()
|
||||
cap.release()
|
||||
|
||||
if not ret or frame is None:
|
||||
return False
|
||||
|
||||
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
|
||||
img = Image.fromarray(frame_rgb)
|
||||
img.thumbnail(thumbnail_size, Image.Resampling.LANCZOS)
|
||||
_atomic_save(img, thumbnail_path, "JPEG", **encoding)
|
||||
frame = self._opencv_frame(video_path, thumbnail_size)
|
||||
if frame is not None:
|
||||
_atomic_save(frame, thumbnail_path, "JPEG", guard=guard, **encoding)
|
||||
self.logger.debug(
|
||||
f"Generated video thumbnail using OpenCV: {thumbnail_path}"
|
||||
)
|
||||
return True
|
||||
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# Fallback to ffmpeg
|
||||
try:
|
||||
import subprocess
|
||||
|
||||
tmp_path = _temp_sibling(thumbnail_path)
|
||||
cmd = [
|
||||
"ffmpeg",
|
||||
"-i",
|
||||
str(video_path),
|
||||
"-ss",
|
||||
"00:00:01",
|
||||
"-vframes",
|
||||
"1",
|
||||
"-vf",
|
||||
f"scale={thumbnail_size[0]}:{thumbnail_size[1]}"
|
||||
":force_original_aspect_ratio=decrease",
|
||||
"-q:v",
|
||||
_ffmpeg_qscale(encoding),
|
||||
"-y",
|
||||
str(tmp_path),
|
||||
]
|
||||
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
|
||||
if result.returncode == 0 and tmp_path.exists():
|
||||
os.replace(tmp_path, thumbnail_path)
|
||||
self.logger.debug(
|
||||
f"Generated video thumbnail using ffmpeg: {thumbnail_path}"
|
||||
)
|
||||
return True
|
||||
else:
|
||||
self.logger.warning(
|
||||
f"ffmpeg failed for {video_path}: {result.stderr}"
|
||||
)
|
||||
|
||||
except (ImportError, subprocess.TimeoutExpired, FileNotFoundError):
|
||||
pass
|
||||
finally:
|
||||
try:
|
||||
tmp_path.unlink()
|
||||
except (NameError, OSError):
|
||||
pass
|
||||
|
||||
if self._ffmpeg_frame(
|
||||
video_path, thumbnail_path, thumbnail_size, encoding, guard, track
|
||||
):
|
||||
return True
|
||||
if keep_existing and thumbnail_path.exists():
|
||||
self.logger.warning(
|
||||
f"No frame decoded for {video_path.name}; keeping its preview"
|
||||
)
|
||||
return False
|
||||
|
||||
# Last resort: create a placeholder thumbnail
|
||||
# Last resort: the placeholder frame.
|
||||
try:
|
||||
from PIL import ImageDraw, ImageFont
|
||||
|
||||
img = Image.new("RGB", thumbnail_size, color=(50, 50, 50))
|
||||
draw = ImageDraw.Draw(img)
|
||||
|
||||
center_x, center_y = thumbnail_size[0] // 2, thumbnail_size[1] // 2
|
||||
triangle_size = min(thumbnail_size) // 4
|
||||
|
||||
points = [
|
||||
(center_x - triangle_size // 2, center_y - triangle_size // 2),
|
||||
(center_x - triangle_size // 2, center_y + triangle_size // 2),
|
||||
(center_x + triangle_size // 2, center_y),
|
||||
]
|
||||
draw.polygon(points, fill=(255, 255, 255))
|
||||
|
||||
try:
|
||||
font = ImageFont.load_default()
|
||||
text = "VIDEO"
|
||||
bbox = draw.textbbox((0, 0), text, font=font)
|
||||
text_width = bbox[2] - bbox[0]
|
||||
draw.text(
|
||||
(
|
||||
center_x - text_width // 2,
|
||||
center_y + triangle_size // 2 + 10,
|
||||
),
|
||||
text,
|
||||
fill=(255, 255, 255),
|
||||
font=font,
|
||||
)
|
||||
except (OSError, AttributeError):
|
||||
pass
|
||||
|
||||
_atomic_save(img, thumbnail_path, "JPEG", **encoding)
|
||||
img = _video_placeholder_image(thumbnail_size)
|
||||
_atomic_save(img, thumbnail_path, "JPEG", guard=guard, **encoding)
|
||||
self.logger.debug(
|
||||
f"Generated placeholder video thumbnail: {thumbnail_path}"
|
||||
)
|
||||
return True
|
||||
|
||||
except PermissionError:
|
||||
except (PermissionError, PreviewWriteHeld):
|
||||
raise
|
||||
except Exception as e:
|
||||
self.logger.warning(
|
||||
f"Failed to create placeholder thumbnail for {video_path}: {e}"
|
||||
)
|
||||
return False
|
||||
|
||||
except PermissionError:
|
||||
except (PermissionError, PreviewWriteHeld):
|
||||
raise
|
||||
except Exception as e:
|
||||
self.logger.error(
|
||||
@@ -2209,6 +2249,88 @@ class ImageRoutesMixin:
|
||||
)
|
||||
return False
|
||||
|
||||
def _opencv_frame(self, video_path, thumbnail_size):
|
||||
"""A frame about 10% in, fitted to *thumbnail_size*, or None (no
|
||||
OpenCV, or it could not open or read the clip)."""
|
||||
try:
|
||||
import cv2
|
||||
except ImportError:
|
||||
return None
|
||||
try:
|
||||
cap = cv2.VideoCapture(str(video_path))
|
||||
try:
|
||||
if not cap.isOpened():
|
||||
return None
|
||||
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, max(1, int(frame_count * 0.1)))
|
||||
ret, frame = cap.read()
|
||||
finally:
|
||||
cap.release()
|
||||
if not ret or frame is None:
|
||||
return None
|
||||
img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||||
img.thumbnail(thumbnail_size, Image.Resampling.LANCZOS)
|
||||
return img
|
||||
except Exception as e:
|
||||
self.logger.debug(f"OpenCV could not read {video_path}: {e}")
|
||||
return None
|
||||
|
||||
def _ffmpeg_frame(
|
||||
self, video_path, thumbnail_path, thumbnail_size, encoding, guard, track
|
||||
):
|
||||
"""Write ffmpeg's frame at 1 s into *thumbnail_path*; True on success."""
|
||||
import subprocess # nosec B404 - fixed argv, no shell
|
||||
|
||||
tmp_path = _temp_sibling(thumbnail_path)
|
||||
cmd = [
|
||||
"ffmpeg",
|
||||
"-i",
|
||||
str(video_path),
|
||||
"-ss",
|
||||
"00:00:01",
|
||||
"-vframes",
|
||||
"1",
|
||||
"-vf",
|
||||
f"scale={thumbnail_size[0]}:{thumbnail_size[1]}"
|
||||
":force_original_aspect_ratio=decrease",
|
||||
"-q:v",
|
||||
_ffmpeg_qscale(encoding),
|
||||
"-y",
|
||||
str(tmp_path),
|
||||
]
|
||||
try:
|
||||
process = subprocess.Popen( # nosec B603 B607 - fixed argv
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
)
|
||||
if track is not None:
|
||||
track(process)
|
||||
try:
|
||||
_out, stderr = process.communicate(timeout=30)
|
||||
except subprocess.TimeoutExpired:
|
||||
process.kill()
|
||||
process.communicate()
|
||||
return False
|
||||
if process.returncode == 0 and tmp_path.exists():
|
||||
_replace_unless_held(tmp_path, thumbnail_path, guard)
|
||||
self.logger.debug(
|
||||
f"Generated video thumbnail using ffmpeg: {thumbnail_path}"
|
||||
)
|
||||
return True
|
||||
self.logger.warning(f"ffmpeg failed for {video_path}: {stderr}")
|
||||
return False
|
||||
except (FileNotFoundError, OSError) as e:
|
||||
if isinstance(e, PermissionError) and tmp_path.exists():
|
||||
raise
|
||||
return False
|
||||
finally:
|
||||
try:
|
||||
tmp_path.unlink()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
async def clear_thumbnails(self, request: web.Request) -> web.Response:
|
||||
"""Clear generated thumbnails across gallery roots, preserving originals.
|
||||
|
||||
@@ -2218,7 +2340,7 @@ class ImageRoutesMixin:
|
||||
"""
|
||||
return await self._run_held(
|
||||
THUMBNAILS_JOB_KIND,
|
||||
"Thumbnail clear (gallery page)",
|
||||
THUMBNAIL_CLEAR_TITLE,
|
||||
"guarded",
|
||||
self._clear_thumbnails,
|
||||
)
|
||||
|
||||
@@ -515,6 +515,12 @@ class JobRegistry:
|
||||
def get(self, job_id):
|
||||
return self._jobs.get(job_id)
|
||||
|
||||
def running(self):
|
||||
"""The running jobs, oldest first (a snapshot, safe off the event loop)."""
|
||||
jobs = [job for job in tuple(self._jobs.values()) if job.running]
|
||||
jobs.sort(key=lambda job: job._started_seq)
|
||||
return jobs
|
||||
|
||||
def active(self, kind):
|
||||
"""The running job of ``kind``, or None (safe off the event loop)."""
|
||||
for job in tuple(self._jobs.values()):
|
||||
|
||||
+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,843 @@
|
||||
"""Automatic thumbnails (#249): previews are a cache the server keeps itself.
|
||||
|
||||
The approved mock has no repeated global thumbnail prompt
|
||||
(``docs/concepts/workspace/README.md``), so the workspace never asks. Every
|
||||
listing gives an image or video a thumbnail URL (``media_thumbnail``), and the
|
||||
previews behind those URLs are made here, in three ways:
|
||||
|
||||
- **On first view.** ``GET /prompt_manager/images/serve/thumbnails/<rel>_thumb
|
||||
<ext>`` for a preview that does not exist yet is admitted on the event loop
|
||||
(``_PreviewGate``): at most ``ON_DEMAND_LIMIT`` previews are written at once,
|
||||
on this API's own small pool (never the default executor), one flight runs
|
||||
per file, and the requests waiting are capped (``MAX_WAITING``); past the
|
||||
cap a request gets its fallback at once. Inside the admission the preview is
|
||||
mapped back to its media file under the same root (the request text is
|
||||
checked first: no traversal, drive, UNC or empty segment; then a regular
|
||||
file inside the root, never through a link), made and served. When it
|
||||
cannot be made an image's original is served and a video's placeholder
|
||||
frame, both ``Cache-Control: no-store``: a video is never sent to an
|
||||
``<img>``. A file that cannot be decoded is remembered for this session
|
||||
(keyed by its path, size and mtime) and counted for the System Overview
|
||||
(``failed`` in ``/images/thumbnails/status``); I/O errors and a preview
|
||||
another program holds open (Windows) are not remembered. The recorded
|
||||
history's thumbnail route goes through the same admission and flight key.
|
||||
- **On capture.** The image monitor hands every new output it settles to
|
||||
``queue_capture_thumbnail``, which queues its preview for one background
|
||||
worker thread (``AutoThumbnailer``). Each wake takes at most Worker threads
|
||||
(``WORKER_THREADS``) files and writes them on that many threads; enqueuing
|
||||
is O(1), so the monitor and ComfyUI's queue never wait. A file already
|
||||
queued is not queued again, the queue is bounded (``MAX_QUEUED``) and a
|
||||
preview left out is made on its first view instead.
|
||||
- **After a scan.** A scan that added files checks the missing previews; from
|
||||
``SCAN_THUMBNAIL_JOB_THRESHOLD`` up it starts the thumbnail Activity job
|
||||
(missing only, visible and cancellable) instead of writing them silently,
|
||||
and fewer go through the worker.
|
||||
|
||||
Every preview is written by ``_generate_one`` with the thumbnail job's default
|
||||
size and quality (``DEFAULT_THUMBNAIL_EDGE``, medium), into a temp sibling
|
||||
renamed over the old one, so a reader never sees a partial file.
|
||||
|
||||
Holds: automatic previews only write files under ``thumbnails/`` and never
|
||||
touch the database, like the thumbnail job, so they run beside a scan, image
|
||||
tagging and the Needs review recording (S5 M7). They are not written while a
|
||||
job runs that replaces the database or removes media and previews: a restore,
|
||||
a duplicate removal, a sources change, the gallery's thumbnail clear or an
|
||||
``exclusive`` job (``_auto_thumbnail_blocker``). The holds are checked when a
|
||||
write starts (after its slot) and again just before its rename.
|
||||
|
||||
Shutdown: ``close_auto_thumbnails`` stops the worker at a file boundary,
|
||||
kills a running ffmpeg, waits a bounded time (and logs when the wait ran
|
||||
out) and shuts the view pool down; after the app's cleanup
|
||||
(``final=True``) nothing starts a worker again. Importing this module again
|
||||
(a package reload) stops the workers the previous import left.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import collections
|
||||
import functools
|
||||
import io
|
||||
import os
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
from .images import (
|
||||
GALLERY_IMAGE_EXTENSIONS,
|
||||
MEDIA_EXTENSIONS,
|
||||
THUMBNAIL_CLEAR_TITLE,
|
||||
THUMBNAIL_ENCODINGS,
|
||||
VIDEO_EXTENSIONS,
|
||||
_has_traversal,
|
||||
_is_within,
|
||||
_thumbnail_is_fresh,
|
||||
_video_placeholder_image,
|
||||
_worker_threads,
|
||||
media_thumbnail_rel,
|
||||
)
|
||||
from .jobs import (
|
||||
DUPLICATE_CLEANUP_JOB_KIND,
|
||||
RESTORE_JOB_KIND,
|
||||
THUMBNAILS_JOB_KIND,
|
||||
JobBusy,
|
||||
)
|
||||
from .settings_routes import SOURCE_CHANGE_JOB_KIND
|
||||
from .thumbnail_jobs import DEFAULT_THUMBNAIL_EDGE, ThumbnailOptions
|
||||
|
||||
try:
|
||||
from ...utils.parallel import map_parallel
|
||||
except ImportError: # tests import the package top-level
|
||||
from utils.parallel import map_parallel
|
||||
|
||||
AUTO_THUMBNAIL_EDGE = DEFAULT_THUMBNAIL_EDGE
|
||||
AUTO_THUMBNAIL_SIZE = (AUTO_THUMBNAIL_EDGE, AUTO_THUMBNAIL_EDGE)
|
||||
AUTO_THUMBNAIL_ENCODING = THUMBNAIL_ENCODINGS["medium"]
|
||||
# Previews views may write at once (the size of the view pool).
|
||||
ON_DEMAND_LIMIT = 3
|
||||
# View requests that may wait on the gate; one more gets its fallback at once.
|
||||
MAX_WAITING = 64
|
||||
# How long a pool thread waits for the capture worker's flight of the same file.
|
||||
FLIGHT_WAIT_SECONDS = 20
|
||||
# How long close() waits for the worker's batch in hand (ffmpeg's own limit).
|
||||
CLOSE_WAIT_SECONDS = 30
|
||||
# Files the capture worker holds; more are made on their first view.
|
||||
MAX_QUEUED = 10_000
|
||||
# Undecodable files remembered per session (oldest forgotten first).
|
||||
MAX_FAILURES_REMEMBERED = 4096
|
||||
# Missing previews after a scan from which the visible thumbnail job starts:
|
||||
# a fresh install or a newly added source. Fewer go through the worker.
|
||||
SCAN_THUMBNAIL_JOB_THRESHOLD = 200
|
||||
# Jobs beside which no automatic preview is written (see the module docstring).
|
||||
AUTO_REFUSED_KINDS = frozenset(
|
||||
{RESTORE_JOB_KIND, DUPLICATE_CLEANUP_JOB_KIND, SOURCE_CHANGE_JOB_KIND}
|
||||
)
|
||||
WORKER_NAME = "PromptManagerThumbnailWorker"
|
||||
POOL_PREFIX = "pm-preview"
|
||||
# Identifies this import's workers: a reload stops the ones it did not start.
|
||||
_MODULE_TOKEN = object()
|
||||
|
||||
READY, FAILED, BUSY, DEFERRED = "ready", "failed", "busy", "deferred"
|
||||
# What the gate answers a request past MAX_WAITING.
|
||||
OVER = object()
|
||||
|
||||
|
||||
def _fingerprint(src):
|
||||
"""``(path, size, mtime_ns)`` of *src*, or None when it cannot be read."""
|
||||
try:
|
||||
stat = os.stat(src)
|
||||
except OSError:
|
||||
return None
|
||||
return (os.path.normcase(str(src)), stat.st_size, stat.st_mtime_ns)
|
||||
|
||||
|
||||
def flight_key(root, rel_parts):
|
||||
"""The single-flight key of ``root/rel_parts``: its folder resolved, case
|
||||
folded. The serve route, the history route and the worker share it."""
|
||||
path = Path(root).joinpath(*rel_parts)
|
||||
return os.path.normcase(os.path.join(os.path.realpath(path.parent), path.name))
|
||||
|
||||
|
||||
def _dst_key(dst):
|
||||
return flight_key(Path(dst).parent, (Path(dst).name,))
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def video_placeholder_jpeg():
|
||||
"""The placeholder frame as JPEG bytes (a video whose preview is missing)."""
|
||||
buffer = io.BytesIO()
|
||||
_video_placeholder_image(AUTO_THUMBNAIL_SIZE).save(
|
||||
buffer, "JPEG", **AUTO_THUMBNAIL_ENCODING
|
||||
)
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
def _is_preview_request(filepath):
|
||||
"""Text check only: ``thumbnails/<...>/<name>_thumb<lower-case ext>``.
|
||||
|
||||
Refuses, as text and before any path is joined, what a joined path could
|
||||
take as an anchor: a backslash, a drive or stream colon, and empty, ``.``
|
||||
or ``..`` segments (so a leading ``/`` and ``//server/share`` too).
|
||||
"""
|
||||
if "\\" in filepath or ":" in filepath:
|
||||
return False
|
||||
parts = filepath.split("/")
|
||||
if len(parts) < 2 or parts[0] != "thumbnails":
|
||||
return False
|
||||
if any(part in ("", ".", "..") for part in parts) or "thumbnails" in parts[1:-1]:
|
||||
return False
|
||||
stem, extension = os.path.splitext(parts[-1])
|
||||
return (
|
||||
stem.endswith("_thumb")
|
||||
and len(stem) > len("_thumb")
|
||||
and extension == extension.lower()
|
||||
and extension in GALLERY_IMAGE_EXTENSIONS
|
||||
)
|
||||
|
||||
|
||||
def _regular_file(path):
|
||||
"""True for a regular file that is not a link (one lstat)."""
|
||||
try:
|
||||
return not path.is_symlink() and path.is_file()
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
|
||||
def _candidates(parts):
|
||||
"""``(folder parts, base, extensions)`` of a checked preview request."""
|
||||
stem, extension = os.path.splitext(parts[-1])
|
||||
wanted = [extension] + (sorted(VIDEO_EXTENSIONS) if extension == ".jpg" else [])
|
||||
return parts[1:-1], stem[: -len("_thumb")], wanted
|
||||
|
||||
|
||||
def _quick_source(root, parts):
|
||||
"""``(src, is_video)`` by exact name only (no folder listing), or None."""
|
||||
root = Path(root)
|
||||
folder_parts, base, wanted = _candidates(parts)
|
||||
folder = root.joinpath(*folder_parts)
|
||||
for extension in wanted:
|
||||
for name in (base + extension, base + extension.upper()):
|
||||
path = folder / name
|
||||
if _regular_file(path) and _is_within(path, root):
|
||||
return path, extension in VIDEO_EXTENSIONS
|
||||
return None
|
||||
|
||||
|
||||
def _thumbnail_source(root, parts):
|
||||
"""``(src, dst, is_video)`` for the preview *parts* under *root*, or None.
|
||||
|
||||
*parts* passed ``_is_preview_request``. The source has the preview's name
|
||||
without ``_thumb`` and an extension that previews as the requested one
|
||||
(any case: exact names first, then one listing of the folder); a ``.jpg``
|
||||
may also be a video's. It must be a regular file inside *root*. ``dst``
|
||||
is None when the preview's folder resolves outside ``root/thumbnails`` (a
|
||||
link): nothing is written there.
|
||||
"""
|
||||
root = Path(root)
|
||||
folder_parts, base, wanted = _candidates(parts)
|
||||
quick = _quick_source(root, parts)
|
||||
if quick is not None:
|
||||
found = quick[0]
|
||||
else:
|
||||
found = _scan_folder(root.joinpath(*folder_parts), base, set(wanted))
|
||||
if found is None or not _is_within(found, root):
|
||||
return None
|
||||
rel = Path(*folder_parts, found.name)
|
||||
dst = root / media_thumbnail_rel(rel)
|
||||
if dst != root.joinpath(*parts):
|
||||
return None
|
||||
if not _is_within(root / "thumbnails", root) or not _is_within(
|
||||
dst, root / "thumbnails"
|
||||
):
|
||||
dst = None
|
||||
return found, dst, found.suffix.lower() in VIDEO_EXTENSIONS
|
||||
|
||||
|
||||
def _scan_folder(folder, base, extensions):
|
||||
"""A mixed-case extension (``.Png``): one listing of *folder*, or None."""
|
||||
try:
|
||||
with os.scandir(folder) as entries:
|
||||
for entry in entries:
|
||||
stem, extension = os.path.splitext(entry.name)
|
||||
if stem == base and extension.lower() in extensions:
|
||||
path = Path(entry.path)
|
||||
if _regular_file(path):
|
||||
return path
|
||||
except OSError:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _stop_stale_workers():
|
||||
"""Stop capture workers a previous import of this module started."""
|
||||
for thread in threading.enumerate():
|
||||
token = getattr(thread, "pm_preview_token", None)
|
||||
close = getattr(thread, "pm_preview_close", None)
|
||||
if token is not None and token is not _MODULE_TOKEN and close is not None:
|
||||
close()
|
||||
|
||||
|
||||
class AutoThumbnailer:
|
||||
"""Single-flight preview writes and the capture worker (thread-safe).
|
||||
|
||||
``generate(src, dst, is_video, *, stop, guard, track)`` writes one
|
||||
preview and returns ``_generate_one``'s result dict; ``may_write()``
|
||||
says whether automatic previews may be written now. Views call
|
||||
``ensure`` on the view pool, the monitor calls ``enqueue`` from its
|
||||
thread.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
generate,
|
||||
may_write,
|
||||
logger,
|
||||
*,
|
||||
workers=_worker_threads,
|
||||
max_queued=MAX_QUEUED,
|
||||
start_worker=True,
|
||||
):
|
||||
self._generate = generate
|
||||
self._may_write = may_write
|
||||
self._logger = logger
|
||||
self._workers = workers
|
||||
self._max_queued = max_queued
|
||||
self._start_worker = start_worker
|
||||
self._changed = threading.Condition()
|
||||
self._inflight = {}
|
||||
self._failures = collections.OrderedDict()
|
||||
self._queue = collections.deque()
|
||||
self._queued = set()
|
||||
self._processes = []
|
||||
self._worker = None
|
||||
# True from a worker's start until it leaves (under the lock), so an
|
||||
# idle worker can exit and the next enqueue starts a fresh one.
|
||||
self._worker_running = False
|
||||
self._closed = False
|
||||
self._busy = False
|
||||
self._batch = 0
|
||||
|
||||
@property
|
||||
def failure_count(self):
|
||||
with self._changed:
|
||||
return len(self._failures)
|
||||
|
||||
@property
|
||||
def closed(self):
|
||||
return self._closed
|
||||
|
||||
def _allowed(self):
|
||||
return not self._closed and self._may_write()
|
||||
|
||||
def _track(self, process):
|
||||
"""Keep an ffmpeg process so close() can kill it."""
|
||||
with self._changed:
|
||||
self._processes = [p for p in self._processes if p.poll() is None]
|
||||
self._processes.append(process)
|
||||
closed = self._closed
|
||||
if closed:
|
||||
process.kill()
|
||||
|
||||
# ── one preview ───────────────────────────────────────────────────
|
||||
|
||||
def ensure(self, src, dst, is_video):
|
||||
"""Make *dst* for *src* unless it is fresh; returns the outcome.
|
||||
|
||||
``ready`` (the preview is there), ``failed`` (the file cannot be
|
||||
decoded, remembered), ``busy`` (an I/O error, the file is gone, or
|
||||
another flight did not finish) or ``deferred`` (a hold refuses
|
||||
automatic writes, checked when the write starts and before its
|
||||
rename). Blocking: call it off the event loop.
|
||||
"""
|
||||
if _thumbnail_is_fresh(src, dst):
|
||||
return READY
|
||||
fingerprint = _fingerprint(src)
|
||||
if fingerprint is None:
|
||||
return BUSY
|
||||
key = _dst_key(dst)
|
||||
with self._changed:
|
||||
if fingerprint in self._failures:
|
||||
return FAILED
|
||||
flight = self._inflight.get(key)
|
||||
if flight is None:
|
||||
self._inflight[key] = done = threading.Event()
|
||||
if flight is not None:
|
||||
flight.wait(FLIGHT_WAIT_SECONDS)
|
||||
if _thumbnail_is_fresh(src, dst):
|
||||
return READY
|
||||
with self._changed:
|
||||
return FAILED if fingerprint in self._failures else BUSY
|
||||
try:
|
||||
return self._lead(src, dst, is_video, fingerprint)
|
||||
finally:
|
||||
with self._changed:
|
||||
self._inflight.pop(key, None)
|
||||
self._changed.notify_all()
|
||||
done.set()
|
||||
|
||||
def _lead(self, src, dst, is_video, fingerprint):
|
||||
if not self._allowed():
|
||||
return DEFERRED
|
||||
if _thumbnail_is_fresh(src, dst):
|
||||
return READY
|
||||
result = self._generate(
|
||||
src,
|
||||
dst,
|
||||
is_video,
|
||||
stop=lambda: self._closed,
|
||||
guard=self._allowed,
|
||||
track=self._track,
|
||||
)
|
||||
if result.get("action") in ("generated", "skipped"):
|
||||
return READY
|
||||
if result.get("held"):
|
||||
return DEFERRED
|
||||
if not result.get("decode"):
|
||||
return BUSY # an I/O error says nothing about the file
|
||||
error = result.get("error") or "failed"
|
||||
with self._changed:
|
||||
self._failures[fingerprint] = error
|
||||
while len(self._failures) > MAX_FAILURES_REMEMBERED:
|
||||
self._failures.popitem(last=False)
|
||||
self._logger.info(
|
||||
f"No preview for {os.path.basename(str(src))} ({error}); "
|
||||
"the original is shown"
|
||||
)
|
||||
return FAILED
|
||||
|
||||
# ── the capture worker ────────────────────────────────────────────
|
||||
|
||||
def enqueue(self, src, dst, is_video):
|
||||
"""Queue *dst*'s preview for the worker; False when already queued,
|
||||
full or closed. Never blocks on a write."""
|
||||
key = _dst_key(dst)
|
||||
with self._changed:
|
||||
if self._closed or key in self._queued:
|
||||
return False
|
||||
if len(self._queue) >= self._max_queued:
|
||||
return False
|
||||
self._queued.add(key)
|
||||
self._queue.append((key, Path(src), Path(dst), bool(is_video)))
|
||||
self._busy = True
|
||||
if self._start_worker and not self._worker_running:
|
||||
self._start_worker_locked()
|
||||
self._changed.notify_all()
|
||||
return True
|
||||
|
||||
def _start_worker_locked(self):
|
||||
"""Start the worker (lock held). A start that fails leaves the file
|
||||
queued and the next enqueue tries again."""
|
||||
worker = threading.Thread(target=self._run, name=WORKER_NAME, daemon=True)
|
||||
worker.pm_preview_token = _MODULE_TOKEN
|
||||
worker.pm_preview_close = self.close
|
||||
self._worker_running = True
|
||||
try:
|
||||
worker.start()
|
||||
except RuntimeError:
|
||||
self._worker_running = False
|
||||
self._busy = False
|
||||
self._logger.warning(
|
||||
"The automatic preview worker could not start; the next output "
|
||||
"tries again",
|
||||
exc_info=True,
|
||||
)
|
||||
return
|
||||
self._worker = worker
|
||||
|
||||
def _run(self):
|
||||
try:
|
||||
self._drain()
|
||||
except Exception:
|
||||
self._logger.warning("The automatic preview worker stopped", exc_info=True)
|
||||
finally:
|
||||
# Whatever ended the loop, the next enqueue starts a fresh worker.
|
||||
with self._changed:
|
||||
self._worker_running = False
|
||||
self._busy = False
|
||||
self._batch = 0
|
||||
self._changed.notify_all()
|
||||
|
||||
def _drain(self):
|
||||
while True:
|
||||
with self._changed:
|
||||
if self._closed or not self._queue:
|
||||
# Idle or closed: leave, so no thread outlives its work.
|
||||
self._busy = False
|
||||
self._worker_running = False
|
||||
self._changed.notify_all()
|
||||
return
|
||||
take = min(len(self._queue), max(1, int(self._workers() or 1)))
|
||||
batch = [self._queue.popleft() for _ in range(take)]
|
||||
self._batch = len(batch)
|
||||
self._changed.notify_all()
|
||||
try:
|
||||
map_parallel(self._write_queued, batch, len(batch))
|
||||
finally:
|
||||
with self._changed:
|
||||
self._batch = 0
|
||||
for key, *_rest in batch:
|
||||
self._queued.discard(key)
|
||||
|
||||
def _write_queued(self, item):
|
||||
_key, src, dst, is_video = item
|
||||
try:
|
||||
self.ensure(src, dst, is_video)
|
||||
except Exception: # one file never stops the worker
|
||||
self._logger.warning(
|
||||
f"Automatic preview failed for {src.name}", exc_info=True
|
||||
)
|
||||
|
||||
def close(self):
|
||||
"""Stop the worker at its next file boundary (queued files are
|
||||
dropped), kill a running ffmpeg and wait up to CLOSE_WAIT_SECONDS;
|
||||
a worker still running after that is logged."""
|
||||
with self._changed:
|
||||
self._closed = True
|
||||
self._queue.clear()
|
||||
self._queued.clear()
|
||||
processes, self._processes = self._processes, []
|
||||
self._changed.notify_all()
|
||||
worker = self._worker
|
||||
for process in processes:
|
||||
try:
|
||||
process.kill()
|
||||
except OSError:
|
||||
pass
|
||||
if worker is None or worker is threading.current_thread():
|
||||
return
|
||||
worker.join(CLOSE_WAIT_SECONDS)
|
||||
if worker.is_alive():
|
||||
self._logger.warning(
|
||||
"The automatic preview worker is still running after "
|
||||
f"{CLOSE_WAIT_SECONDS} s; it stops at its next file"
|
||||
)
|
||||
|
||||
# ── waits (tests and shutdown; never timers) ──────────────────────
|
||||
|
||||
def _wait_for(self, predicate, timeout):
|
||||
with self._changed:
|
||||
return self._changed.wait_for(predicate, timeout)
|
||||
|
||||
def wait_idle(self, timeout=None):
|
||||
"""True once the worker has nothing queued or in hand."""
|
||||
return self._wait_for(lambda: not self._busy, timeout)
|
||||
|
||||
def wait_batch(self, size, timeout=None):
|
||||
"""True once the worker holds a batch of *size* files."""
|
||||
return self._wait_for(lambda: self._batch == size, timeout)
|
||||
|
||||
|
||||
class _PreviewGate:
|
||||
"""Admission of view requests on the event loop (#249 I2).
|
||||
|
||||
Waiters are coroutines: at most ``limit`` previews are written at once
|
||||
on the view pool, one flight runs per key, and past ``MAX_WAITING``
|
||||
waiting requests ``run`` answers ``OVER`` at once.
|
||||
"""
|
||||
|
||||
def __init__(self, limit, logger):
|
||||
self.loop = asyncio.get_running_loop()
|
||||
self._slots = asyncio.Semaphore(limit)
|
||||
self._flights = {}
|
||||
self._tasks = set()
|
||||
self._logger = logger
|
||||
self._tick = asyncio.Event()
|
||||
self.waiting = 0
|
||||
|
||||
def _changed(self):
|
||||
tick, self._tick = self._tick, asyncio.Event()
|
||||
tick.set()
|
||||
|
||||
async def settled(self, predicate):
|
||||
"""Return once ``predicate()`` holds (re-checked at each change)."""
|
||||
while not predicate():
|
||||
await self._tick.wait()
|
||||
|
||||
async def run(self, key, work, pool):
|
||||
"""``work()``'s answer for *key* (shared by every request of the
|
||||
flight), None when it raised, or OVER past the cap."""
|
||||
if self.waiting >= MAX_WAITING:
|
||||
return OVER
|
||||
flight = self._flights.get(key)
|
||||
if flight is None:
|
||||
flight = self.loop.create_future()
|
||||
self._flights[key] = flight
|
||||
task = self.loop.create_task(self._lead(key, flight, work, pool))
|
||||
self._tasks.add(task)
|
||||
task.add_done_callback(self._tasks.discard)
|
||||
self.waiting += 1
|
||||
self._changed()
|
||||
try:
|
||||
return await asyncio.shield(flight)
|
||||
finally:
|
||||
self.waiting -= 1
|
||||
self._changed()
|
||||
|
||||
async def _lead(self, key, flight, work, pool):
|
||||
result = None
|
||||
try:
|
||||
async with self._slots:
|
||||
result = await self.loop.run_in_executor(pool, work)
|
||||
except Exception:
|
||||
self._logger.warning("Automatic preview failed", exc_info=True)
|
||||
finally:
|
||||
self._flights.pop(key, None)
|
||||
if not flight.done():
|
||||
flight.set_result(result)
|
||||
|
||||
|
||||
class AutoThumbnailMixin:
|
||||
"""Automatic previews on the API: views, the monitor's queue, scans."""
|
||||
|
||||
_auto_thumbnailer = None
|
||||
_auto_thumbnails_shut = False
|
||||
_preview_gate_obj = None
|
||||
_preview_pool_obj = None
|
||||
_auto_thumbnailer_lock = threading.Lock()
|
||||
|
||||
def _auto_thumbnails(self):
|
||||
"""This API's AutoThumbnailer, created on first use (no worker after
|
||||
the app's cleanup)."""
|
||||
with AutoThumbnailMixin._auto_thumbnailer_lock:
|
||||
if self._auto_thumbnailer is None:
|
||||
self._auto_thumbnailer = AutoThumbnailer(
|
||||
self._write_auto_thumbnail,
|
||||
lambda: self._auto_thumbnail_blocker() is None,
|
||||
self.logger,
|
||||
start_worker=not self._auto_thumbnails_shut,
|
||||
)
|
||||
return self._auto_thumbnailer
|
||||
|
||||
def close_auto_thumbnails(self, final=False):
|
||||
"""Stop the capture worker and the view pool, and forget them.
|
||||
|
||||
Bounded (see AutoThumbnailer.close). ``final`` (the app's cleanup)
|
||||
also keeps a late monitor event from starting a worker again; a
|
||||
test reset passes False and the next use starts fresh ones.
|
||||
"""
|
||||
with AutoThumbnailMixin._auto_thumbnailer_lock:
|
||||
if final:
|
||||
self._auto_thumbnails_shut = True
|
||||
thumbnailer, self._auto_thumbnailer = self._auto_thumbnailer, None
|
||||
pool, self._preview_pool_obj = self._preview_pool_obj, None
|
||||
self._preview_gate_obj = None
|
||||
if thumbnailer is not None:
|
||||
thumbnailer.close()
|
||||
if pool is not None:
|
||||
pool.shutdown(wait=False, cancel_futures=True)
|
||||
|
||||
def _auto_thumbnail_failures(self):
|
||||
thumbnailer = self._auto_thumbnailer
|
||||
return 0 if thumbnailer is None else thumbnailer.failure_count
|
||||
|
||||
def _write_auto_thumbnail(self, src, dst, is_video, *, stop, guard, track):
|
||||
return self._generate_one(
|
||||
src,
|
||||
dst,
|
||||
AUTO_THUMBNAIL_SIZE,
|
||||
is_video,
|
||||
encoding=AUTO_THUMBNAIL_ENCODING,
|
||||
stop=stop,
|
||||
guard=guard,
|
||||
track=track,
|
||||
)
|
||||
|
||||
def _auto_thumbnail_blocker(self):
|
||||
"""The running job automatic previews wait for, or None (any thread)."""
|
||||
for job in self._jobs.running():
|
||||
if job.admission == "exclusive" or job.kind in AUTO_REFUSED_KINDS:
|
||||
return job
|
||||
if job.kind == THUMBNAILS_JOB_KIND and job.title == THUMBNAIL_CLEAR_TITLE:
|
||||
return job
|
||||
return None
|
||||
|
||||
# ── the view admission ────────────────────────────────────────────
|
||||
|
||||
def _preview_gate(self):
|
||||
"""The gate of the running loop (a new loop gets a new one)."""
|
||||
gate = self._preview_gate_obj
|
||||
if gate is None or gate.loop is not asyncio.get_running_loop():
|
||||
gate = self._preview_gate_obj = _PreviewGate(ON_DEMAND_LIMIT, self.logger)
|
||||
return gate
|
||||
|
||||
def _preview_pool(self):
|
||||
with AutoThumbnailMixin._auto_thumbnailer_lock:
|
||||
if self._preview_pool_obj is None:
|
||||
self._preview_pool_obj = ThreadPoolExecutor(
|
||||
ON_DEMAND_LIMIT, thread_name_prefix=POOL_PREFIX
|
||||
)
|
||||
return self._preview_pool_obj
|
||||
|
||||
def _preview_pool_threads(self):
|
||||
pool = self._preview_pool_obj
|
||||
return [] if pool is None else list(getattr(pool, "_threads", ()))
|
||||
|
||||
def _video_placeholder_bytes(self):
|
||||
return video_placeholder_jpeg()
|
||||
|
||||
async def _admit_preview(self, src, dst, is_video):
|
||||
"""Make one preview through the view admission (the history route).
|
||||
|
||||
Shares the serve route's flight key; every flight answers the
|
||||
``ensure`` outcome (or None, or OVER), whichever route started it.
|
||||
"""
|
||||
return await self._preview_gate().run(
|
||||
flight_key(Path(dst).parent, (Path(dst).name,)),
|
||||
functools.partial(self._ensure_view, src, dst, is_video),
|
||||
self._preview_pool(),
|
||||
)
|
||||
|
||||
def _ensure_view(self, src, dst, is_video):
|
||||
"""``ensure``, reporting a preview removed under us as busy."""
|
||||
outcome = self._auto_thumbnails().ensure(src, dst, is_video)
|
||||
if outcome == READY and not Path(dst).is_file():
|
||||
return BUSY # removed under us: the fallback this once
|
||||
return outcome
|
||||
|
||||
# ── on first view ─────────────────────────────────────────────────
|
||||
|
||||
def _preview_on_demand(self, root, parts):
|
||||
"""The ``ensure`` outcome for the preview *parts* under *root*
|
||||
(``deferred`` when its folder is linked out of the root), or None
|
||||
when no media file previews as *parts*."""
|
||||
located = _thumbnail_source(root, parts)
|
||||
if located is None:
|
||||
return None
|
||||
src, dst, is_video = located
|
||||
if dst is None:
|
||||
return DEFERRED
|
||||
return self._ensure_view(src, dst, is_video)
|
||||
|
||||
def _preview_response(self, outcome, src, dst, is_video):
|
||||
if outcome == READY and dst is not None:
|
||||
return self._file_response(dst)
|
||||
if is_video:
|
||||
return web.Response(
|
||||
body=video_placeholder_jpeg(),
|
||||
content_type="image/jpeg",
|
||||
headers={"Cache-Control": "no-store"},
|
||||
)
|
||||
response = web.FileResponse(src)
|
||||
response.headers["Cache-Control"] = "no-store"
|
||||
return response
|
||||
|
||||
async def _serve_missing_thumbnail(self, output_dirs, filepath):
|
||||
"""The made preview, or its fallback; None when *filepath* names no
|
||||
preview of a media file in *output_dirs*.
|
||||
|
||||
The flight answers only the outcome (a history flight may be the one
|
||||
joined); the file it previews is derived here: the preview path is
|
||||
the request's own, and a fallback looks its source up again on the
|
||||
view pool (bounded by it, never the default executor).
|
||||
"""
|
||||
if _has_traversal(filepath) or not _is_preview_request(filepath):
|
||||
return None
|
||||
parts = tuple(filepath.split("/"))
|
||||
gate = self._preview_gate()
|
||||
pool = self._preview_pool()
|
||||
for root in output_dirs:
|
||||
root = Path(root)
|
||||
outcome = await gate.run(
|
||||
flight_key(root, parts),
|
||||
functools.partial(self._preview_on_demand, root, parts),
|
||||
pool,
|
||||
)
|
||||
if outcome is None:
|
||||
continue
|
||||
if outcome is OVER:
|
||||
# Past the cap: exact names only, no listing, no write.
|
||||
quick = _quick_source(root, parts)
|
||||
if quick is None:
|
||||
continue
|
||||
return self._preview_response(BUSY, quick[0], None, quick[1])
|
||||
dst = root.joinpath(*parts)
|
||||
if outcome == READY:
|
||||
return self._file_response(dst)
|
||||
located = await gate.loop.run_in_executor(
|
||||
pool, _thumbnail_source, root, parts
|
||||
)
|
||||
if located is None:
|
||||
continue
|
||||
src, _dst, is_video = located
|
||||
return self._preview_response(outcome, src, None, is_video)
|
||||
return None
|
||||
|
||||
# ── on capture ────────────────────────────────────────────────────
|
||||
|
||||
def queue_capture_thumbnail(self, image_path):
|
||||
"""Queue the preview of a new output (the monitor's thread).
|
||||
|
||||
Returns ``queued``, ``fresh`` (its preview is current) or ``skipped``
|
||||
(not a media file under a current root, a link, a preview itself,
|
||||
already queued, or after the app's cleanup). Never writes on the
|
||||
caller's thread.
|
||||
"""
|
||||
if self._auto_thumbnails_shut:
|
||||
return "skipped"
|
||||
located = self._capture_target(image_path)
|
||||
if located is None:
|
||||
return "skipped"
|
||||
src, dst, is_video = located
|
||||
if _thumbnail_is_fresh(src, dst):
|
||||
return "fresh"
|
||||
if self._auto_thumbnails().enqueue(src, dst, is_video):
|
||||
return "queued"
|
||||
return "skipped"
|
||||
|
||||
def _capture_target(self, image_path):
|
||||
path = Path(os.path.abspath(str(image_path)))
|
||||
if path.suffix.lower() not in MEDIA_EXTENSIONS or not _regular_file(path):
|
||||
return None
|
||||
folded = os.path.normcase(str(path))
|
||||
for root in self._get_all_output_dirs():
|
||||
root = Path(os.path.abspath(str(root)))
|
||||
base = os.path.normcase(str(root)).rstrip("\\/") + os.sep
|
||||
if not folded.startswith(base) or not _is_within(path, root):
|
||||
continue
|
||||
rel = Path(os.path.relpath(path, root))
|
||||
if "thumbnails" in rel.parts[:-1]:
|
||||
return None
|
||||
dst = root / media_thumbnail_rel(rel)
|
||||
if not _is_within(root / "thumbnails", root) or not _is_within(
|
||||
dst, root / "thumbnails"
|
||||
):
|
||||
return None
|
||||
return path, dst, path.suffix.lower() in VIDEO_EXTENSIONS
|
||||
return None
|
||||
|
||||
# ── after a scan ──────────────────────────────────────────────────
|
||||
|
||||
async def _thumbnails_after_scan(self, job, output_dirs=None):
|
||||
"""Hand a finished scan's missing previews on; returns what it did.
|
||||
|
||||
None (the scan added nothing), ``job`` (the thumbnail job started for
|
||||
``SCAN_THUMBNAIL_JOB_THRESHOLD`` or more missing previews), ``busy``
|
||||
(a job refused it; the previews are made on view), ``queued`` (fewer,
|
||||
through the worker) or ``none`` (nothing missing). It counts the
|
||||
sources the scan walked (``output_dirs``, narrowed to the job's
|
||||
``roots``; every source when None). The server log says which.
|
||||
"""
|
||||
if job.counters.get("added", 0) <= 0:
|
||||
return None
|
||||
if output_dirs is None:
|
||||
output_dirs = await self._run_in_executor(self._get_all_output_dirs)
|
||||
roots = job.scope.get("roots")
|
||||
if roots:
|
||||
output_dirs = [output_dirs[i] for i in roots if 0 <= i < len(output_dirs)]
|
||||
plan = await self._run_in_executor(self._plan_gallery_thumbnails, output_dirs)
|
||||
missing = len(plan["targets"]) + plan["remaining"]
|
||||
if missing >= SCAN_THUMBNAIL_JOB_THRESHOLD:
|
||||
options = ThumbnailOptions("missing", AUTO_THUMBNAIL_EDGE, "medium")
|
||||
try:
|
||||
self._start_thumbnail_job(options)
|
||||
except JobBusy as busy:
|
||||
self.logger.info(
|
||||
f"{missing:,} previews missing after the scan; the thumbnail "
|
||||
f"job did not start while {busy.running.title} runs"
|
||||
)
|
||||
return "busy"
|
||||
self.logger.info(
|
||||
f"{missing:,} previews missing after the scan; thumbnail job started"
|
||||
)
|
||||
return "job"
|
||||
if not missing:
|
||||
return "none"
|
||||
thumbnailer = self._auto_thumbnails()
|
||||
for src, dst, is_video in plan["targets"]:
|
||||
thumbnailer.enqueue(src, dst, is_video)
|
||||
self.logger.info(f"{missing:,} previews missing after the scan; queued")
|
||||
return "queued"
|
||||
|
||||
async def _hand_off_scan_thumbnails(self, job, output_dirs=None):
|
||||
"""``_thumbnails_after_scan`` for the scan runner: never raises."""
|
||||
try:
|
||||
return await self._thumbnails_after_scan(job, output_dirs)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
self.logger.warning("Previews after the scan failed", exc_info=True)
|
||||
return None
|
||||
|
||||
|
||||
# A package reload re-runs this module: stop the workers the old one left.
|
||||
_stop_stale_workers()
|
||||
@@ -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)
|
||||
@@ -386,6 +386,8 @@ def create_app(root: Path) -> web.Application:
|
||||
async def reset(request: web.Request) -> web.Response:
|
||||
nonlocal history_fixture
|
||||
request_holds.release()
|
||||
# No preview queued by the last test may land in the new fixture.
|
||||
await api._run_in_executor(api.close_auto_thumbnails)
|
||||
library_editor_commands.reset()
|
||||
data = await request.json() if request.can_read_body else {}
|
||||
if not isinstance(data, dict):
|
||||
|
||||
@@ -398,18 +398,20 @@ test('preview uses actual Viewer.js and cleans up when closing or leaving Images
|
||||
test('real video media has a usable card and native playback cleans up on close and navigation', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset', { data: { include_video: true } })).toBeOK();
|
||||
const media = await (await request.get('/prompt_manager/images/output?limit=40')).json();
|
||||
expect(media.images.find(image => image.filename === 'clip.webm').thumbnail_url).toBeNull();
|
||||
// #249: a video with no preview yet still lists its preview URL; the server makes the frame on first view.
|
||||
expect(media.images.find(image => image.filename === 'clip.webm').thumbnail_url).toMatch(/\/thumbnails\/clip_thumb\.jpg\?root=0/);
|
||||
await page.goto(workspace);
|
||||
await refine(page, { Media: 'video' });
|
||||
await expect(cards(page)).toHaveCount(1);
|
||||
const card = cards(page).first();
|
||||
await expect(card).toHaveAttribute('data-image-key', '0:clip.webm');
|
||||
await expect(card.locator('img')).toHaveCount(0);
|
||||
await expect(card.getByRole('button', { name: 'Open clip.webm', exact: true })).toBeVisible();
|
||||
// Legacy video art inside the concept card; D11 names the video's own extension.
|
||||
await expect(card.locator('.media-art > .image-open > video')).toHaveCount(1);
|
||||
// The card shows the preview frame, never the video itself; D11 names the video's own extension.
|
||||
await expect(card.locator('video')).toHaveCount(0);
|
||||
await expect(card.locator('.media-art > .image-open > img')).toHaveCount(1);
|
||||
await expect(card.locator('.media-info p')).toHaveText(/(?:^|· )WEBM$/);
|
||||
expect(await card.locator('video').evaluate(node => node.getBoundingClientRect().height)).toBeGreaterThan(20);
|
||||
await expect.poll(() => card.locator('img').evaluate(node => node.complete && node.naturalHeight > 0)).toBe(true);
|
||||
expect(await card.locator('img').evaluate(node => node.getBoundingClientRect().height)).toBeGreaterThan(20);
|
||||
await card.getByRole('button', { name: 'Open clip.webm', exact: true }).click();
|
||||
const video = page.locator('dialog[open] video');
|
||||
await expect(video).toBeVisible();
|
||||
@@ -713,16 +715,45 @@ test('selected layout stays blue with readable text on hover in either theme', a
|
||||
}
|
||||
});
|
||||
|
||||
test('refresh releases the previous inspector video source', async ({ page, request }) => {
|
||||
test('refresh releases the previous inspector cover source', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset', { data: { include_video: true } })).toBeOK();
|
||||
await page.goto(workspace);
|
||||
const card = cards(page).filter({ hasText: 'clip.webm' });
|
||||
await card.getByRole('button', { name: 'Inspect clip.webm', exact: true }).click();
|
||||
// #249: a video's cover is its preview frame, not the video.
|
||||
const cover = details(page).locator('img.detail-cover');
|
||||
await expect(cover).toHaveAttribute('src', /\/thumbnails\/clip_thumb\.jpg\?/);
|
||||
await expect(details(page).locator('video')).toHaveCount(0);
|
||||
await cover.evaluate(element => { window.__oldInspectorCover = element; });
|
||||
await rescan(page);
|
||||
await expect.poll(() => page.evaluate(() => !window.__oldInspectorCover.isConnected)).toBe(true);
|
||||
expect(await page.evaluate(() => window.__oldInspectorCover.getAttribute('src'))).toBeNull();
|
||||
});
|
||||
|
||||
/** A listing that names no preview for its videos (the fallback a card and the cover keep for one). */
|
||||
async function videosWithoutPreview(page) {
|
||||
await page.route(url => url.pathname === '/prompt_manager/images/output', async route => {
|
||||
const response = await route.fetch();
|
||||
const body = await response.json();
|
||||
if (Array.isArray(body.images)) body.images = body.images.map(image => image.is_video ? { ...image, thumbnail_url: null } : image);
|
||||
await route.fulfill({ response, json: body });
|
||||
});
|
||||
}
|
||||
|
||||
test('a video listed without a preview keeps its own element, released on refresh', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset', { data: { include_video: true } })).toBeOK();
|
||||
await videosWithoutPreview(page);
|
||||
await page.goto(workspace);
|
||||
const card = cards(page).filter({ hasText: 'clip.webm' });
|
||||
await expect(card.locator('.media-art > .image-open > video')).toHaveCount(1);
|
||||
await expect(card.locator('img')).toHaveCount(0);
|
||||
await card.getByRole('button', { name: 'Inspect clip.webm', exact: true }).click();
|
||||
await expect(details(page).locator('video')).toBeVisible();
|
||||
await details(page).locator('video').evaluate(element => { window.__oldInspectorVideo = element; });
|
||||
await card.locator('video').evaluate(element => { window.__oldCardVideo = element; });
|
||||
await rescan(page);
|
||||
await expect.poll(() => page.evaluate(() => !window.__oldInspectorVideo.isConnected)).toBe(true);
|
||||
expect(await page.evaluate(() => window.__oldInspectorVideo.getAttribute('src'))).toBeNull();
|
||||
await expect.poll(() => page.evaluate(() => !window.__oldInspectorVideo.isConnected && !window.__oldCardVideo.isConnected)).toBe(true);
|
||||
expect(await page.evaluate(() => [window.__oldInspectorVideo.getAttribute('src'), window.__oldCardVideo.getAttribute('src')])).toEqual([null, null]);
|
||||
});
|
||||
|
||||
test('Import Rescan rescans cached files and preserves the collection when rescan fails', async ({ page, request }) => {
|
||||
|
||||
@@ -0,0 +1,135 @@
|
||||
import { startCoverage, saveCoverage } from '../../tools/vendor/browser-coverage.mjs';
|
||||
import { test, expect } from './fixtures.mjs';
|
||||
|
||||
/**
|
||||
* #249: cards and strips load previews, never the full file. The approved mock has no repeated thumbnail prompt
|
||||
* (docs/concepts/workspace/README.md), so a preview that was not made yet is made by the server on first view:
|
||||
* the listing still names the `_thumb` URL. Only the media viewer's main image (and downloads) load the original.
|
||||
*/
|
||||
const images = '/prompt_manager/workspace#/images';
|
||||
const library = '/prompt_manager/workspace#/library';
|
||||
const cards = page => page.locator('article[data-image-key]');
|
||||
const details = page => page.getByRole('complementary', { name: 'Selected image details', exact: true });
|
||||
const SERVE = '/prompt_manager/images/serve/';
|
||||
|
||||
/** Every media file the page asks the server for, as `{path, preview}` (path relative to its root). */
|
||||
function recordMedia(page) {
|
||||
const seen = [];
|
||||
page.on('request', request => {
|
||||
const url = new URL(request.url());
|
||||
if (!url.pathname.startsWith(SERVE)) return;
|
||||
const path = decodeURIComponent(url.pathname.slice(SERVE.length));
|
||||
seen.push({ path, preview: /^thumbnails\/.+_thumb\.[a-z0-9]+$/.test(path) });
|
||||
});
|
||||
return seen;
|
||||
}
|
||||
/** Every card image has loaded (or failed): what it asked for is all in `seen`. */
|
||||
const settled = page => page.evaluate(() => Promise.all([...document.querySelectorAll('article img')].map(image => image.decode().catch(() => {}))));
|
||||
const originals = seen => seen.filter(item => !item.preview).map(item => item.path);
|
||||
|
||||
test.beforeEach(async ({ page, browserName }) => {
|
||||
await startCoverage(page, browserName);
|
||||
page.thumbErrors = [];
|
||||
page.on('pageerror', error => page.thumbErrors.push(error.message));
|
||||
});
|
||||
test.afterEach(async ({ page, browserName }) => {
|
||||
await saveCoverage(page, browserName);
|
||||
expect(page.thumbErrors).toEqual([]);
|
||||
});
|
||||
|
||||
test('a page of the Images grid requests previews, never originals, in every layout and size', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset')).toBeOK();
|
||||
// Newest outputs with no preview yet: the live finding (844 such files listed with a null thumbnail).
|
||||
await expect(await request.post('/__test__/system/unthumbnailed', { data: { count: 6 } })).toBeOK();
|
||||
const listing = await (await request.get('/prompt_manager/images/output?limit=40')).json();
|
||||
const fresh = listing.images.filter(image => image.path.startsWith('unthumbnailed/'));
|
||||
expect(fresh).toHaveLength(6);
|
||||
for (const image of fresh) expect(image.thumbnail_url).toMatch(/\/thumbnails\/unthumbnailed\/plain_\d{3}_thumb\.png\?root=0/);
|
||||
const seen = recordMedia(page);
|
||||
await page.goto(images);
|
||||
await expect(cards(page)).toHaveCount(8);
|
||||
await settled(page);
|
||||
expect(originals(seen)).toEqual([]);
|
||||
expect(seen.filter(item => item.path.startsWith('thumbnails/unthumbnailed/'))).toHaveLength(6);
|
||||
// The server made them on that first view.
|
||||
await expect.poll(async () => (await (await request.get('/prompt_manager/images/thumbnails/status')).json()).pending).toBe(0);
|
||||
for (const layout of ['List', 'Masonry', 'Grid']) {
|
||||
await page.getByRole('button', { name: layout, exact: true }).click();
|
||||
await expect(page.getByRole('button', { name: layout, exact: true })).toHaveAttribute('aria-pressed', 'true');
|
||||
for (const density of ['large', 'compact', 'comfortable']) {
|
||||
await page.getByRole('combobox', { name: 'Thumbnail size', exact: true }).selectOption(density);
|
||||
await expect(cards(page)).toHaveCount(8);
|
||||
await settled(page);
|
||||
}
|
||||
}
|
||||
expect(originals(seen)).toEqual([]);
|
||||
});
|
||||
|
||||
test('the inspector cover is the preview and only the media viewer loads the original', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset')).toBeOK();
|
||||
await expect(await request.post('/__test__/system/unthumbnailed', { data: { count: 1 } })).toBeOK();
|
||||
const seen = recordMedia(page);
|
||||
await page.goto(images);
|
||||
const card = cards(page).filter({ hasText: 'plain_001.png' });
|
||||
await card.getByRole('button', { name: 'Inspect plain_001.png', exact: true }).click();
|
||||
const cover = details(page).locator('img.detail-cover');
|
||||
await expect(cover).toHaveAttribute('src', /\/thumbnails\/unthumbnailed\/plain_001_thumb\.png\?/);
|
||||
await settled(page);
|
||||
await cover.evaluate(image => image.decode().catch(() => {}));
|
||||
expect(originals(seen)).toEqual([]);
|
||||
await card.getByRole('button', { name: 'Open plain_001.png', exact: true }).click();
|
||||
await expect(page.locator('dialog#viewer-dialog')).toBeVisible();
|
||||
await expect.poll(() => originals(seen)).toEqual(['unthumbnailed/plain_001.png']);
|
||||
await page.getByRole('button', { name: 'Close image viewer', exact: true }).click();
|
||||
});
|
||||
|
||||
test('a video card shows its preview frame instead of loading the video', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset', { data: { include_video: true } })).toBeOK();
|
||||
const seen = recordMedia(page);
|
||||
await page.goto(images);
|
||||
const card = cards(page).filter({ hasText: 'clip.webm' });
|
||||
await expect(card.locator('.media-art > .image-open > img')).toHaveAttribute('src', /\/thumbnails\/clip_thumb\.jpg\?/);
|
||||
await expect(card.locator('video')).toHaveCount(0);
|
||||
await settled(page);
|
||||
expect(originals(seen)).toEqual([]);
|
||||
});
|
||||
|
||||
test('Library cards and the shared gallery request previews only', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset', { data: { native_library: true } })).toBeOK();
|
||||
const seen = recordMedia(page);
|
||||
await page.goto(library);
|
||||
const cat = page.locator('article[data-prompt-id]').filter({ has: page.getByText('Cat portrait', { exact: true }) });
|
||||
await expect(cat.locator('.card-art img')).not.toHaveCount(0);
|
||||
await settled(page);
|
||||
expect(seen.length).toBeGreaterThan(0);
|
||||
expect(originals(seen)).toEqual([]);
|
||||
});
|
||||
|
||||
test('a video whose preview cannot be loaded swaps back to its own element', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset', { data: { include_video: true } })).toBeOK();
|
||||
await page.route(url => url.pathname.endsWith('/thumbnails/clip_thumb.jpg'), route => route.fulfill({ status: 404, body: '' }));
|
||||
await page.goto(images);
|
||||
const card = cards(page).filter({ hasText: 'clip.webm' });
|
||||
await expect(card.locator('.media-art > .image-open > video')).toHaveCount(1);
|
||||
await expect(card.locator('.media-art > .image-open > img')).toHaveCount(0);
|
||||
await card.getByRole('button', { name: 'Inspect clip.webm', exact: true }).click();
|
||||
await expect(details(page).locator('video.detail-cover')).toHaveCount(1);
|
||||
await expect(details(page).locator('img.detail-cover')).toHaveCount(0);
|
||||
});
|
||||
|
||||
test('an image listed without a preview (a nested thumbnails folder) shows its original', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset')).toBeOK();
|
||||
// M10: such files are listed with a null thumbnail_url; the card falls back to the original as before.
|
||||
await page.route(url => url.pathname === '/prompt_manager/images/output', async route => {
|
||||
const response = await route.fetch();
|
||||
const body = await response.json();
|
||||
if (Array.isArray(body.images)) body.images = body.images.map(image => ({ ...image, thumbnail_url: null }));
|
||||
await route.fulfill({ response, json: body });
|
||||
});
|
||||
const seen = recordMedia(page);
|
||||
await page.goto(images);
|
||||
await expect(cards(page)).toHaveCount(2);
|
||||
await settled(page);
|
||||
expect(seen.length).toBeGreaterThan(0);
|
||||
expect(seen.every(item => !item.preview)).toBe(true);
|
||||
});
|
||||
@@ -39,3 +39,4 @@ test('stale inspection echoes use the canonical recovery vocabulary and identity
|
||||
});
|
||||
}
|
||||
});
|
||||
test('the recorded output inspector cover is its preview, not the original (#249)',()=>{const a=adapter(),target=media('2',1);const record=a.toRecord(envelope(target),target,extractGenerationFacts);assert.deepEqual(record.preview,{url:target.thumbnail_url});assert.notEqual(record.preview.url,record.download.url);});
|
||||
|
||||
@@ -895,3 +895,30 @@ test('cache off never rereads a selection whose source changed, and refresh neve
|
||||
assert.equal(deferred.state.deferred, true);
|
||||
deferred.dispose();
|
||||
});
|
||||
|
||||
test('record carries the listing thumbnail as its preview, never another file (#249)', () => {
|
||||
const toRecord = contract('toInspectionRecord'), entry = image();
|
||||
const thumb = `/prompt_manager/images/serve/thumbnails/folder/a%20%23%25_thumb.png?root=0&root_scope=${scopeA}`;
|
||||
assert.deepEqual(toRecord(envelope(entry), { ...entry, thumbnail_url: thumb }, extractGenerationFacts).preview, { url: thumb });
|
||||
const versioned = `/prompt_manager/images/serve/thumbnails/folder/a%20%23%25_thumb.png?v=12&root=0&root_scope=${scopeA}`;
|
||||
assert.deepEqual(toRecord(envelope(entry), { ...entry, thumbnail_url: versioned }, extractGenerationFacts).preview, { url: versioned });
|
||||
// No thumbnail, or one naming another file, root or origin: no preview (the cover keeps the original).
|
||||
for (const thumbnail_url of [undefined, null, '', thumb.replace('a%20%23%25_thumb', 'b_thumb'), thumb.replace('root=0', 'root=1'),
|
||||
thumb.replace(scopeA, scopeB), `https://other.test${thumb}`, thumb.replace('thumbnails/', ''), `${thumb}#x`]) {
|
||||
assert.equal(toRecord(envelope(entry), { ...entry, thumbnail_url }, extractGenerationFacts).preview, null, String(thumbnail_url));
|
||||
}
|
||||
const video = { root_index: 0, path: 'clip.MP4', filename: 'clip.MP4', url: `/prompt_manager/images/serve/clip.MP4?root=0&root_scope=${scopeA}`, root_scope: scopeA,
|
||||
thumbnail_url: `/prompt_manager/images/serve/thumbnails/clip_thumb.jpg?root=0&root_scope=${scopeA}` };
|
||||
const data = envelope(video); data.file = { ...data.file, media_type: 'video', extension: '.mp4', format: 'MP4' };
|
||||
assert.deepEqual(toRecord(data, video, extractGenerationFacts).preview, { url: video.thumbnail_url });
|
||||
});
|
||||
|
||||
test("names the server quotes (parentheses, '!*') still match their original and preview (#249 M11)", () => {
|
||||
const toRecord = contract('toInspectionRecord');
|
||||
const entry = { root_index: 0, path: "folder/a (1)!'*.png", filename: "a (1)!'*.png", root_scope: scopeA,
|
||||
url: `/prompt_manager/images/serve/folder/a%20%281%29%21%27%2A.png?root=0&root_scope=${scopeA}`,
|
||||
thumbnail_url: `/prompt_manager/images/serve/thumbnails/folder/a%20%281%29%21%27%2A_thumb.png?root=0&root_scope=${scopeA}` };
|
||||
const record = toRecord(envelope(entry), entry, extractGenerationFacts);
|
||||
assert.equal(record.download.url, entry.url);
|
||||
assert.deepEqual(record.preview, { url: entry.thumbnail_url });
|
||||
});
|
||||
|
||||
@@ -253,3 +253,11 @@ test('waiting stops when its signal aborts', async () => {
|
||||
gone.abort(new Error('already'));
|
||||
await assert.rejects(waitForJob(fakeActivity(), JOB, { signal: gone.signal }), { message: 'already' });
|
||||
});
|
||||
|
||||
test('the thumbnail row counts previews that could not be made (#249)', async () => {
|
||||
const { thumbnailTitle } = await overviewModule;
|
||||
assert.equal(thumbnailTitle({ pending: 12, failed: 2 }), '12 missing thumbnails · 2 failed');
|
||||
assert.equal(thumbnailTitle({ pending: 1, failed: 1 }), '1 missing thumbnail · 1 failed');
|
||||
assert.equal(thumbnailTitle({ pending: 0, failed: 0 }), 'No missing thumbnails');
|
||||
assert.equal(thumbnailTitle({ pending: 3, failed: 'x' }), '3 missing thumbnails');
|
||||
});
|
||||
|
||||
@@ -79,6 +79,8 @@ class AdminAPITestCase(AioHTTPTestCase):
|
||||
|
||||
async def tearDownAsync(self):
|
||||
await super().tearDownAsync() # closes the aiohttp test client
|
||||
# The automatic preview worker writes into the temp tree: stop it first.
|
||||
self.api.close_auto_thumbnails()
|
||||
GalleryConfig.MONITORING_DIRECTORIES = self._orig_gallery_dirs
|
||||
PromptManagerConfig.update_config(self._orig_pm_config)
|
||||
for key, value in self._orig_env.items():
|
||||
|
||||
@@ -131,6 +131,8 @@ class AutotagAPITestCase(AioHTTPTestCase):
|
||||
|
||||
async def tearDownAsync(self):
|
||||
await super().tearDownAsync() # closes the aiohttp test client
|
||||
# The automatic preview worker writes into the temp tree: stop it first.
|
||||
self.api.close_auto_thumbnails()
|
||||
self._service_patch.stop()
|
||||
if model_gate is not None:
|
||||
leaked = model_gate.holder()
|
||||
|
||||
@@ -62,6 +62,8 @@ class BackupRestoreTestCase(AioHTTPTestCase):
|
||||
|
||||
async def tearDownAsync(self):
|
||||
await super().tearDownAsync() # closes the aiohttp test client
|
||||
# The automatic preview worker writes into the temp tree: stop it first.
|
||||
self.api.close_auto_thumbnails()
|
||||
PromptModel.close_all_instances() # restore may have swapped api.db
|
||||
assert_closed(self, self.tmpdir)
|
||||
for name in os.listdir(self.tmpdir):
|
||||
|
||||
@@ -384,6 +384,8 @@ class CoreRoutesTestCase(AioHTTPTestCase):
|
||||
|
||||
async def tearDownAsync(self):
|
||||
await super().tearDownAsync() # closes the aiohttp test client
|
||||
# The automatic preview worker writes into the temp tree: stop it first.
|
||||
self.api.close_auto_thumbnails()
|
||||
self.api.db.close_all()
|
||||
assert_closed(self, self.tmpdir)
|
||||
shutil.rmtree(self.tmpdir, ignore_errors=True)
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -401,7 +401,7 @@ class HistoryOutputAPITests(ImageAPITestCase):
|
||||
self.assertEqual(response.status, 400)
|
||||
await response.read()
|
||||
|
||||
async def test_media_ranges_head_and_existing_thumbnail_without_generation(self):
|
||||
async def test_media_ranges_head_and_a_preview_made_on_first_view(self):
|
||||
path = self.output_dir / "clip.mp4"
|
||||
payload = bytes(range(256)) * 600
|
||||
path.write_bytes(payload)
|
||||
@@ -430,9 +430,14 @@ class HistoryOutputAPITests(ImageAPITestCase):
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(int(response.headers["Content-Length"]), len(payload))
|
||||
self.assertEqual(await response.read(), b"")
|
||||
# A missing preview is made on first view (#249): an undecodable video
|
||||
# gets the placeholder frame, never the whole file.
|
||||
preview = self.output_dir / "thumbnails" / "clip_thumb.jpg"
|
||||
response = await self.current(descriptor["thumbnail_url"])
|
||||
self.assertEqual(await response.read(), payload)
|
||||
preview = make_png(self.output_dir / "thumbnails" / "clip_thumb.jpg")
|
||||
self.assertEqual(response.headers["Content-Type"], "image/jpeg")
|
||||
self.assertEqual(await response.read(), preview.read_bytes())
|
||||
self.assertLess(preview.stat().st_size, len(payload))
|
||||
make_png(preview)
|
||||
response = await self.current(descriptor["thumbnail_url"])
|
||||
self.assertEqual(await response.read(), preview.read_bytes())
|
||||
|
||||
@@ -1111,6 +1116,73 @@ class HistoryOutputAPITests(ImageAPITestCase):
|
||||
response = await self.current(descriptor["thumbnail_url"])
|
||||
self.assert_error(response, await response.json(), 403, "not-authorized")
|
||||
|
||||
async def test_missing_preview_is_never_made_through_a_linked_out_folder(self):
|
||||
original = make_png(self.output_dir / "image.png", size=(64, 64))
|
||||
elsewhere = self.temp_dir / "elsewhere"
|
||||
elsewhere.mkdir()
|
||||
(self.output_dir / "thumbnails").symlink_to(elsewhere, target_is_directory=True)
|
||||
start = self.capture([self.ref()])
|
||||
descriptor = await self.media(start)
|
||||
response = await self.current(descriptor["thumbnail_url"])
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(await response.read(), original.read_bytes())
|
||||
self.assertEqual(os.listdir(elsewhere), [])
|
||||
|
||||
async def test_missing_preview_is_never_made_through_a_linked_sub_folder(self):
|
||||
# M6: every link segment of the preview's folder is checked.
|
||||
original = make_png(self.output_dir / "sub" / "image.png", size=(64, 64))
|
||||
elsewhere = self.temp_dir / "elsewhere"
|
||||
elsewhere.mkdir()
|
||||
(self.output_dir / "thumbnails").mkdir()
|
||||
(self.output_dir / "thumbnails" / "sub").symlink_to(
|
||||
elsewhere, target_is_directory=True
|
||||
)
|
||||
start = self.capture([self.ref(subfolder="sub")])
|
||||
descriptor = await self.media(start)
|
||||
response = await self.current(descriptor["thumbnail_url"])
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(await response.read(), original.read_bytes())
|
||||
self.assertEqual(os.listdir(elsewhere), [])
|
||||
|
||||
async def test_a_video_preview_that_cannot_be_made_is_the_placeholder(self):
|
||||
# I1: never the video bytes inside an <img>.
|
||||
from py.api.jobs import RESTORE_JOB_KIND
|
||||
|
||||
(self.output_dir / "clip.webm").write_bytes(b"AV1 stream no decoder here")
|
||||
start = self.capture([self.ref(filename="clip.webm")])
|
||||
descriptor = await self.media(start)
|
||||
gate = asyncio.Event()
|
||||
|
||||
async def runner(job):
|
||||
await gate.wait()
|
||||
|
||||
job = self.api._jobs.start(
|
||||
RESTORE_JOB_KIND,
|
||||
"restore",
|
||||
admission="records",
|
||||
cancellable=False,
|
||||
runner=runner,
|
||||
)
|
||||
try:
|
||||
response = await self.current(descriptor["thumbnail_url"])
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.headers["Content-Type"], "image/jpeg")
|
||||
self.assertTrue((await response.read()).startswith(b"\xff\xd8"))
|
||||
finally:
|
||||
gate.set()
|
||||
await asyncio.wait_for(job.task, 5)
|
||||
self.assertFalse((self.output_dir / "thumbnails").exists())
|
||||
|
||||
async def test_a_preview_that_cannot_be_made_serves_the_original(self):
|
||||
broken = self.output_dir / "image.png"
|
||||
broken.write_bytes(b"\x89PNG\r\n\x1a\n not an image")
|
||||
start = self.capture([self.ref()])
|
||||
descriptor = await self.media(start)
|
||||
response = await self.current(descriptor["thumbnail_url"])
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(await response.read(), broken.read_bytes())
|
||||
self.assertFalse((self.output_dir / "thumbnails" / "image_thumb.png").exists())
|
||||
|
||||
async def test_disappearance_before_open_is_missing_and_closes_handle(
|
||||
self,
|
||||
):
|
||||
|
||||
+69
-16
@@ -80,6 +80,8 @@ class ImageAPITestCase(AioHTTPTestCase):
|
||||
|
||||
async def tearDownAsync(self):
|
||||
await super().tearDownAsync() # closes the aiohttp test client
|
||||
# The automatic preview worker writes into the temp tree: stop it first.
|
||||
self.api.close_auto_thumbnails()
|
||||
self._anchor_patch.stop()
|
||||
# Windows refuses to unlink a database that still has open handles.
|
||||
self.api.db.close_all()
|
||||
@@ -1565,39 +1567,65 @@ class TestVideoThumbnailFallbacks(ImageRouteCoverageCase):
|
||||
with Image.open(self.thumb) as img:
|
||||
self.assertEqual(img.format, "JPEG")
|
||||
|
||||
def test_cv2_cannot_open_or_read(self):
|
||||
with patch.dict(sys.modules, {"cv2": self._fake_cv2(opened=False)}):
|
||||
closed = self.api._generate_video_thumbnail(
|
||||
self.video, self.thumb, self.SIZE
|
||||
)
|
||||
with patch.dict(sys.modules, {"cv2": self._fake_cv2(read_ok=False)}):
|
||||
unread = self.api._generate_video_thumbnail(
|
||||
self.video, self.thumb, self.SIZE
|
||||
)
|
||||
self.assertEqual((closed, unread), (False, False))
|
||||
def test_cv2_cannot_open_or_read_falls_through_to_ffmpeg_then_placeholder(self):
|
||||
# #249 I1: an OpenCV that cannot decode the clip (AV1) no longer ends
|
||||
# the run; ffmpeg and then the placeholder are tried.
|
||||
for fake in (self._fake_cv2(opened=False), self._fake_cv2(read_ok=False)):
|
||||
self.thumb.unlink(missing_ok=True)
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": fake}),
|
||||
patch("subprocess.Popen", side_effect=FileNotFoundError("ffmpeg")),
|
||||
):
|
||||
ok = self.api._generate_video_thumbnail(
|
||||
self.video, self.thumb, self.SIZE
|
||||
)
|
||||
self.assertTrue(ok)
|
||||
with Image.open(self.thumb) as img:
|
||||
self.assertEqual((img.format, img.size), ("JPEG", self.SIZE))
|
||||
|
||||
def test_cv2_crash_is_contained(self):
|
||||
cv2 = MagicMock()
|
||||
cv2.VideoCapture.side_effect = RuntimeError("driver crash")
|
||||
with patch.dict(sys.modules, {"cv2": cv2}):
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": cv2}),
|
||||
patch("subprocess.Popen", side_effect=FileNotFoundError("ffmpeg")),
|
||||
):
|
||||
ok = self.api._generate_video_thumbnail(self.video, self.thumb, self.SIZE)
|
||||
self.assertTrue(ok, "the placeholder is written")
|
||||
|
||||
def test_cv2_failure_keeps_an_existing_preview_on_rebuild(self):
|
||||
make_png(self.thumb)
|
||||
before = self.thumb.read_bytes()
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": self._fake_cv2(opened=False)}),
|
||||
patch("subprocess.Popen", side_effect=FileNotFoundError("ffmpeg")),
|
||||
):
|
||||
ok = self.api._generate_video_thumbnail(
|
||||
self.video, self.thumb, self.SIZE, keep_existing=True
|
||||
)
|
||||
self.assertFalse(ok)
|
||||
self.assertEqual(self.thumb.read_bytes(), before)
|
||||
|
||||
def test_ffmpeg_absent_yields_placeholder(self):
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": None}),
|
||||
patch("subprocess.run", side_effect=FileNotFoundError("ffmpeg")),
|
||||
patch("subprocess.Popen", side_effect=FileNotFoundError("ffmpeg")),
|
||||
):
|
||||
ok = self.api._generate_video_thumbnail(self.video, self.thumb, self.SIZE)
|
||||
self.assertTrue(ok)
|
||||
with Image.open(self.thumb) as img:
|
||||
self.assertEqual(img.size, self.SIZE)
|
||||
|
||||
@staticmethod
|
||||
def _process(returncode, stderr=""):
|
||||
process = MagicMock(returncode=returncode)
|
||||
process.communicate.return_value = ("", stderr)
|
||||
return process
|
||||
|
||||
def test_ffmpeg_failure_yields_placeholder(self):
|
||||
failed = MagicMock(returncode=1, stderr="no stream")
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": None}),
|
||||
patch("subprocess.run", return_value=failed),
|
||||
patch("subprocess.Popen", return_value=self._process(1, "no stream")),
|
||||
):
|
||||
ok = self.api._generate_video_thumbnail(self.video, self.thumb, self.SIZE)
|
||||
self.assertTrue(ok)
|
||||
@@ -1606,16 +1634,41 @@ class TestVideoThumbnailFallbacks(ImageRouteCoverageCase):
|
||||
def test_ffmpeg_success(self):
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": None}),
|
||||
patch("subprocess.run", return_value=MagicMock(returncode=0)) as run,
|
||||
patch("subprocess.Popen", return_value=self._process(0)) as run,
|
||||
):
|
||||
ok = self.api._generate_video_thumbnail(self.video, self.thumb, self.SIZE)
|
||||
self.assertTrue(ok)
|
||||
self.assertEqual(run.call_args.args[0][0], "ffmpeg")
|
||||
|
||||
def test_ffmpeg_is_tracked_and_a_hold_before_the_rename_writes_nothing(self):
|
||||
tracked = []
|
||||
|
||||
def writing(cmd, **kwargs):
|
||||
Path(cmd[-1]).write_bytes(b"\xff\xd8 frame")
|
||||
return self._process(0)
|
||||
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": None}),
|
||||
patch("subprocess.Popen", side_effect=writing),
|
||||
):
|
||||
with self.assertRaises(images_module.PreviewWriteHeld):
|
||||
self.api._generate_video_thumbnail(
|
||||
self.video,
|
||||
self.thumb,
|
||||
self.SIZE,
|
||||
guard=lambda: False,
|
||||
track=tracked.append,
|
||||
)
|
||||
self.assertEqual(len(tracked), 1)
|
||||
self.assertFalse(self.thumb.exists())
|
||||
self.assertEqual(
|
||||
[p.name for p in self.thumb.parent.iterdir() if "part" in p.name], []
|
||||
)
|
||||
|
||||
def test_placeholder_failure_returns_false(self):
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": None}),
|
||||
patch("subprocess.run", side_effect=FileNotFoundError),
|
||||
patch("subprocess.Popen", side_effect=FileNotFoundError),
|
||||
patch.object(
|
||||
images_module.Image, "new", side_effect=RuntimeError("no PIL")
|
||||
),
|
||||
|
||||
@@ -397,15 +397,20 @@ class TestQueueReads(ReviewRoutesTestCase):
|
||||
self._no_paths()
|
||||
|
||||
async def test_queue_cards_use_the_images_listing_thumbnail(self):
|
||||
# M10: a known thumbnail is the card's media, as on the Images page;
|
||||
# without one the card falls back to the original.
|
||||
# M10: the card's media is the Images listing's thumbnail. Without a
|
||||
# preview yet it is still the thumbnail URL (made on first view, #249).
|
||||
self._require_routes()
|
||||
from py.api.images import _output_image_entry
|
||||
|
||||
path = self._ambiguous("sub/f0.png")
|
||||
self._record(path)
|
||||
item = (await self._list())["items"][0]
|
||||
self.assertIsNone(item["thumbnail_url"])
|
||||
scope = item["target"]["root_scope"]
|
||||
self.assertEqual(
|
||||
item["thumbnail_url"],
|
||||
"/prompt_manager/images/serve/thumbnails/sub/f0_thumb.png"
|
||||
f"?root=0&root_scope={scope}",
|
||||
)
|
||||
thumb = self.output_dir / "thumbnails" / "sub" / "f0_thumb.png"
|
||||
thumb.parent.mkdir(parents=True)
|
||||
write_png(thumb, prompt=ambiguous_graph("cat", "dog"), color=(1, 2, 3))
|
||||
|
||||
@@ -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"
|
||||
)
|
||||
]
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -324,7 +324,7 @@ class TestThumbnailJob(ThumbnailJobCase):
|
||||
old = preview.read_bytes()
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": None}),
|
||||
patch("subprocess.run", side_effect=FileNotFoundError("no ffmpeg")),
|
||||
patch("subprocess.Popen", side_effect=FileNotFoundError("no ffmpeg")),
|
||||
):
|
||||
_, body = await self._start(mode="rebuild")
|
||||
job = await self._finished(body["job"]["id"])
|
||||
@@ -566,7 +566,7 @@ class TestThumbnailJob(ThumbnailJobCase):
|
||||
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": None}),
|
||||
patch("subprocess.run", side_effect=FileNotFoundError("no ffmpeg")),
|
||||
patch("subprocess.Popen", side_effect=FileNotFoundError("no ffmpeg")),
|
||||
patch.object(images_module, "_atomic_save", refuse),
|
||||
):
|
||||
_, body = await self._start()
|
||||
@@ -594,12 +594,14 @@ class TestThumbnailJob(ThumbnailJobCase):
|
||||
def fake_ffmpeg(cmd, **_kwargs):
|
||||
commands.append(cmd)
|
||||
Path(cmd[-1]).write_bytes(b"jpeg")
|
||||
return type("Done", (), {"returncode": 0, "stderr": ""})()
|
||||
done = type("Done", (), {"returncode": 0})()
|
||||
done.communicate = lambda timeout=None: ("", "")
|
||||
return done
|
||||
|
||||
for quality, qscale in (("high", "2"), ("medium", "4"), ("low", "8")):
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": None}),
|
||||
patch("subprocess.run", side_effect=fake_ffmpeg),
|
||||
patch("subprocess.Popen", side_effect=fake_ffmpeg),
|
||||
):
|
||||
_, body = await self._start(mode="rebuild", size=640, quality=quality)
|
||||
job = await self._finished(body["job"]["id"])
|
||||
@@ -618,12 +620,13 @@ class TestThumbnailJob(ThumbnailJobCase):
|
||||
real_save = images_module._atomic_save
|
||||
|
||||
def recording(img, dst, fmt=None, **kwargs):
|
||||
saved.append(kwargs)
|
||||
# The encoder settings (guard is the automatic previews' hold check).
|
||||
saved.append({k: v for k, v in kwargs.items() if k != "guard"})
|
||||
return real_save(img, dst, fmt, **kwargs)
|
||||
|
||||
with (
|
||||
patch.dict(sys.modules, {"cv2": None}),
|
||||
patch("subprocess.run", side_effect=FileNotFoundError("no ffmpeg")),
|
||||
patch("subprocess.Popen", side_effect=FileNotFoundError("no ffmpeg")),
|
||||
patch.object(images_module, "_atomic_save", recording),
|
||||
):
|
||||
_, body = await self._start(quality="low")
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -79,6 +79,10 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
# ``record_review_file``, wired at start-up): a linked output is
|
||||
# never queued and its pending items leave the queue. None: no pass.
|
||||
self.review_recorder: Optional[Callable[[str], str]] = None
|
||||
# Queues a settled output's preview (the API's
|
||||
# ``queue_capture_thumbnail``, wired at start-up; #249). It only
|
||||
# queues: the preview is written on its own worker. None: no queue.
|
||||
self.thumbnail_queue: Optional[Callable[[str], str]] = None
|
||||
|
||||
# Read from GalleryConfig if available, otherwise use defaults
|
||||
try:
|
||||
@@ -220,10 +224,30 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
)
|
||||
return
|
||||
self.process_new_image(image_path, prompt_snapshot=prompt_snapshot)
|
||||
self._queue_thumbnail(image_path)
|
||||
finally:
|
||||
with self._pending_lock:
|
||||
self._pending.pop(image_path, None)
|
||||
|
||||
def _queue_thumbnail(self, image_path: str) -> None:
|
||||
"""Queue a settled output's preview, best effort (linked or not).
|
||||
|
||||
A failure is logged, never raised: the preview is then made on its
|
||||
first view.
|
||||
"""
|
||||
queue_preview = getattr(self, "thumbnail_queue", None)
|
||||
if queue_preview is None:
|
||||
return
|
||||
try:
|
||||
outcome = queue_preview(image_path)
|
||||
except Exception:
|
||||
self.logger.warning(
|
||||
f"Preview of {os.path.basename(image_path)} could not be queued",
|
||||
exc_info=True,
|
||||
)
|
||||
return
|
||||
self.logger.debug(f"Preview {outcome}: {os.path.basename(image_path)}")
|
||||
|
||||
def wait_until_settled(self, image_path: str, stop=None) -> bool:
|
||||
"""True once the file's size is non-zero and unchanged across two reads.
|
||||
|
||||
@@ -636,6 +660,16 @@ class ImageMonitor:
|
||||
self.monitored_directories = []
|
||||
self.logger = get_logger("prompt_manager.image_monitor")
|
||||
self.review_recorder: Optional[Callable[[str], str]] = None
|
||||
self.thumbnail_queue: Optional[Callable[[str], str]] = None
|
||||
|
||||
def set_thumbnail_queue(
|
||||
self, queue_preview: Optional[Callable[[str], str]]
|
||||
) -> None:
|
||||
"""Queue every settled output's preview through ``queue_preview``
|
||||
(None stops it), on the running handler and every later one."""
|
||||
self.thumbnail_queue = queue_preview
|
||||
if self.handler is not None:
|
||||
self.handler.thumbnail_queue = queue_preview
|
||||
|
||||
def set_review_recorder(self, recorder: Optional[Callable[[str], str]]) -> None:
|
||||
"""Record ambiguous outputs through ``recorder`` (None stops recording).
|
||||
@@ -703,6 +737,7 @@ class ImageMonitor:
|
||||
# leaves the monitor stuck "already running" with nothing to stop.
|
||||
handler = ImageGenerationHandler(self.db_manager, self.prompt_tracker)
|
||||
handler.review_recorder = self.review_recorder
|
||||
handler.thumbnail_queue = self.thumbnail_queue
|
||||
observer = Observer()
|
||||
scheduled = []
|
||||
for output_dir in valid_directories:
|
||||
|
||||
+4
-2
@@ -488,7 +488,8 @@
|
||||
const index = from + offset;
|
||||
// Use thumbnail for display if available, fallback to original
|
||||
const displayUrl = image.thumbnail_url || image.url;
|
||||
const hasThumb = !!image.thumbnail_url;
|
||||
// The listing names a preview URL for every file (made on view, #249); the dot means the file exists.
|
||||
const hasThumb = typeof image.thumbnail_ready === 'boolean' ? image.thumbnail_ready : !!image.thumbnail_url;
|
||||
const isVideo = image.is_video || false;
|
||||
const mediaType = image.media_type || 'image';
|
||||
|
||||
@@ -551,7 +552,8 @@
|
||||
const index = from + offset;
|
||||
// Use thumbnail for display if available, fallback to original
|
||||
const displayUrl = image.thumbnail_url || image.url;
|
||||
const hasThumb = !!image.thumbnail_url;
|
||||
// The listing names a preview URL for every file (made on view, #249); the dot means the file exists.
|
||||
const hasThumb = typeof image.thumbnail_ready === 'boolean' ? image.thumbnail_ready : !!image.thumbnail_url;
|
||||
const isVideo = image.is_video || false;
|
||||
const mediaType = image.media_type || 'image';
|
||||
|
||||
|
||||
@@ -318,18 +318,25 @@ export function mountFileInspector(container, {api, sender, parentScope, variant
|
||||
});
|
||||
return button;
|
||||
}
|
||||
/** The concept's plain cover; the card image opens the viewer (D19). Video keeps its own element. */
|
||||
/**
|
||||
* The concept's plain cover; the card image opens the viewer (D19). #249: it is the file's preview (a video's
|
||||
* is a JPEG frame) whenever the listing gave one; without one the original, and a video keeps its own element.
|
||||
*/
|
||||
function cover(record) {
|
||||
const {file} = record, video = file?.media_type === 'video';
|
||||
const node = el(video ? 'video' : 'img', undefined, 'detail-cover'); node.src = record.download.url;
|
||||
if (video) {
|
||||
node.muted = true; node.playsInline = true; node.preload = 'metadata';
|
||||
contentScope.cleanup(() => {node.pause(); node.removeAttribute('src'); node.load();});
|
||||
} else {
|
||||
node.alt = ''; node.decoding = 'async';
|
||||
if (file?.width && file?.height) {node.width = file.width; node.height = file.height;}
|
||||
contentScope.cleanup(() => node.removeAttribute('src'));
|
||||
}
|
||||
const {file} = record, preview = record.preview?.url, isVideo = file?.media_type === 'video';
|
||||
if (isVideo && !preview) return videoCover(record);
|
||||
const node = el('img', undefined, 'detail-cover'); node.src = preview || record.download.url;
|
||||
node.alt = ''; node.decoding = 'async';
|
||||
if (file?.width && file?.height) {node.width = file.width; node.height = file.height;}
|
||||
contentScope.cleanup(() => node.removeAttribute('src'));
|
||||
// #249: a video whose preview cannot be loaded swaps back to its own element.
|
||||
if (isVideo) contentScope.listen(node, 'error', () => { if (node.isConnected) node.replaceWith(videoCover(record)); }, {once:true});
|
||||
return node;
|
||||
}
|
||||
function videoCover(record) {
|
||||
const node = el('video', undefined, 'detail-cover'); node.src = record.download.url;
|
||||
node.muted = true; node.playsInline = true; node.preload = 'metadata';
|
||||
contentScope.cleanup(() => {node.pause(); node.removeAttribute('src'); node.load();});
|
||||
return node;
|
||||
}
|
||||
/** "File & provenance": the file facts, the D21 prompt evidence, the occurrence and the saved association (D66 for a local file). */
|
||||
|
||||
@@ -64,7 +64,10 @@ export function toHistoryInspectionRecord(result, target, extractFacts) {
|
||||
const path = [target.reference.subfolder, target.reference.filename].filter(Boolean).join('/');
|
||||
const identity = {kind:'reported-location', key:target.location_key, path:target.reference.path_style === 'windows' ? path.replaceAll('\\', '/') : path,
|
||||
occurrence:{...target.occurrence}, history_epoch:target.history_epoch, current_scope:target.current_scope, reference:{...target.reference}};
|
||||
return projectInspectionPayload(result, identity, {url:urlFor(target, 'url', 'media'), filename:target.filename}, extractFacts);
|
||||
const record = projectInspectionPayload(result, identity, {url:urlFor(target, 'url', 'media'), filename:target.filename}, extractFacts);
|
||||
// #249: the inspector's cover is the recorded output's preview (made by the server on first view).
|
||||
record.preview = {url:urlFor(target, 'thumbnail_url', 'thumbnail')};
|
||||
return record;
|
||||
}
|
||||
export const reportedLocationAdapter = Object.freeze({
|
||||
validate,
|
||||
|
||||
@@ -17,19 +17,29 @@ function el(tag, text, className) {
|
||||
return node;
|
||||
}
|
||||
|
||||
/** The card's media: a video's first frame when it has no thumbnail, else the thumbnail or the image. */
|
||||
/** A video's own element (its first frame), for a video without a usable preview. */
|
||||
function videoMedia(image) {
|
||||
const node = el('video');
|
||||
node.src = window.ApiClient.resolveUrl(image.url);
|
||||
node.muted = true; node.playsInline = true; node.preload = 'metadata';
|
||||
return node;
|
||||
}
|
||||
|
||||
/**
|
||||
* The card's media: the thumbnail (a video's is its preview frame), else the image, else a video's own element.
|
||||
* #249: a video whose preview cannot be loaded swaps back to its own element.
|
||||
*/
|
||||
export function cardMedia(image, lazy = false) {
|
||||
const { resolveUrl } = window.ApiClient;
|
||||
const video = image.is_video && !image.thumbnail_url;
|
||||
const node = el(video ? 'video' : 'img');
|
||||
if (image.is_video && !image.thumbnail_url) return videoMedia(image);
|
||||
const node = el('img');
|
||||
node.src = resolveUrl(image.thumbnail_url || image.url);
|
||||
if (video) { node.muted = true; node.playsInline = true; node.preload = 'metadata'; }
|
||||
else {
|
||||
// The concept card image: alt is the card title, with the listing's intrinsic size when known.
|
||||
node.alt = cardTitle(image); node.decoding = 'async'; if (lazy) node.loading = 'lazy';
|
||||
if ([image.width, image.height].every(value => Number.isFinite(value) && value > 0)) {
|
||||
node.width = image.width; node.height = image.height;
|
||||
}
|
||||
// replaceWith does nothing for a card no longer in the page.
|
||||
if (image.is_video) node.addEventListener('error', () => node.replaceWith(videoMedia(image)), { once: true });
|
||||
// The concept card image: alt is the card title, with the listing's intrinsic size when known.
|
||||
node.alt = cardTitle(image); node.decoding = 'async'; if (lazy) node.loading = 'lazy';
|
||||
if ([image.width, image.height].every(value => Number.isFinite(value) && value > 0)) {
|
||||
node.width = image.width; node.height = image.height;
|
||||
}
|
||||
return node;
|
||||
}
|
||||
|
||||
@@ -44,6 +44,9 @@ export function inspectPath(image) {
|
||||
return `/prompt_manager/images/inspect/${image.path.split('/').map(encodeURIComponent).join('/')}?${query}`;
|
||||
}
|
||||
|
||||
/** One path segment as the server quotes it (Python `quote(..., safe="/")` also escapes `!'()*`, #249 M11). */
|
||||
const quoteSegment = segment => encodeURIComponent(segment).replace(/[!'()*]/g, c => `%${c.charCodeAt(0).toString(16).toUpperCase()}`);
|
||||
|
||||
function downloadFor(image, required) {
|
||||
try {
|
||||
if (typeof image.url !== 'string' || !image.url) throw new TypeError('Missing original media URL');
|
||||
@@ -55,7 +58,7 @@ function downloadFor(image, required) {
|
||||
const sameOrigin = 'Origin' in globalThis ?
|
||||
globalThis.Origin.from(url).isSameOrigin(globalThis.Origin.from(reference)) :
|
||||
url.origin !== 'null' && url.origin === reference.origin;
|
||||
const expected = new URL(resolveUrl(`/prompt_manager/images/serve/${image.path.split('/').map(encodeURIComponent).join('/')}`), reference);
|
||||
const expected = new URL(resolveUrl(`/prompt_manager/images/serve/${image.path.split('/').map(quoteSegment).join('/')}`), reference);
|
||||
const scope = descriptorScope(image);
|
||||
if (!sameOrigin || !['http:', 'https:'].includes(url.protocol) || url.username || url.password || url.hash ||
|
||||
url.pathname !== expected.pathname || url.searchParams.getAll('root').length !== 1 || url.searchParams.get('root') !== String(image.root_index) ||
|
||||
@@ -68,6 +71,39 @@ function downloadFor(image, required) {
|
||||
}
|
||||
}
|
||||
|
||||
// The server's video extensions (py/api/inspection_service.py VIDEO_EXTENSIONS): their preview is a JPEG.
|
||||
const VIDEO_EXTENSIONS = new Set(['.mp4', '.webm', '.avi', '.mov', '.mkv', '.m4v', '.wmv']);
|
||||
|
||||
/** `thumbnails/<path>_thumb<ext>` for a listed file, as the server names its preview (images.media_thumbnail_rel). */
|
||||
function previewPath(path) {
|
||||
const slash = path.lastIndexOf('/'), dot = path.lastIndexOf('.');
|
||||
const cut = dot > slash ? dot : path.length, extension = path.slice(cut).toLowerCase();
|
||||
return `thumbnails/${path.slice(0, cut)}_thumb${VIDEO_EXTENSIONS.has(extension) ? '.jpg' : extension}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* #249: the listing's thumbnail of `image` as `{url}` for the inspector's cover, or null. Only the same-origin
|
||||
* serve URL of this file's own preview under the same root and source scope is used; anything else keeps the
|
||||
* original (a missing preview is made by the server on first view).
|
||||
*/
|
||||
function previewFor(image) {
|
||||
try {
|
||||
if (typeof image.thumbnail_url !== 'string' || !image.thumbnail_url) return null;
|
||||
const location = globalThis.location || globalThis.window?.location;
|
||||
const reference = new URL(location?.href || 'http://localhost/');
|
||||
const resolveUrl = globalThis.window?.ApiClient?.resolveUrl || globalThis.ApiClient?.resolveUrl || (url => url);
|
||||
const value = resolveUrl(image.thumbnail_url), url = new URL(value, reference);
|
||||
const expected = new URL(resolveUrl(`/prompt_manager/images/serve/${previewPath(image.path).split('/').map(quoteSegment).join('/')}`), reference);
|
||||
const scope = descriptorScope(image);
|
||||
if (url.origin !== reference.origin || url.username || url.password || url.hash || url.pathname !== expected.pathname ||
|
||||
url.searchParams.getAll('root').length !== 1 || url.searchParams.get('root') !== String(image.root_index) ||
|
||||
url.searchParams.getAll('root_scope').length > 1 || (scope !== null && url.searchParams.get('root_scope') !== scope)) return null;
|
||||
return { url: value };
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
export function originalMediaDownload(image) {
|
||||
validateIdentity(image);
|
||||
return downloadFor(image, true);
|
||||
@@ -284,16 +320,18 @@ export function toInspectionRecord(result, image, extractFacts = globalThis.Comf
|
||||
if (!legacy && result.inspection_version !== 1) throw new TypeError('Unsupported inspection version');
|
||||
if (!legacy && result.success !== true) throw new TypeError('Invalid inspection success status');
|
||||
if ((!legacy || result.image !== undefined) && (!object(result.image) || result.image.root_index !== image.root_index || result.image.path !== image.path || result.image.filename !== image.path.split('/').at(-1))) throw new TypeError('Mismatched inspection identity');
|
||||
return projectInspectionPayload(result,
|
||||
const record = projectInspectionPayload(result,
|
||||
{ kind: 'rooted-file', key: imageKey(image), root_index: image.root_index, path: image.path, root_scope: scope },
|
||||
downloadFor(image, !legacy || scope !== null), extractFacts, legacy, parser);
|
||||
record.preview = record.download ? previewFor(image) : null;
|
||||
return record;
|
||||
}
|
||||
|
||||
/** Shared bounded payload projection; adapters validate identity before entering here. */
|
||||
export function projectInspectionPayload(result, identity, download, extractFacts = globalThis.ComfyMetadata?.extractGenerationFacts || globalThis.window?.ComfyMetadata?.extractGenerationFacts, legacy = false, parser = comfyMetadata()) {
|
||||
const record = {
|
||||
identity, file: null, prompts: promptsFor(result.prompts), library: null, metadata: null,
|
||||
generation: unavailableFacts(), historyTarget: null, download, imageTags: null,
|
||||
generation: unavailableFacts(), historyTarget: null, download, imageTags: null, preview: null,
|
||||
};
|
||||
if (legacy) return record;
|
||||
record.imageTags = imageTagsFor(result.image_tags);
|
||||
|
||||
@@ -88,7 +88,9 @@ export function thumbnailTitle(status, sources) {
|
||||
const missing = status?.pending;
|
||||
if (!Number.isSafeInteger(missing) || sources?.effective === 0) return 'Thumbnail status unavailable';
|
||||
if (missing === 0) return 'No missing thumbnails';
|
||||
return `${number(missing)} missing ${noun(missing, 'thumbnail')}`;
|
||||
// #249: files whose automatic preview could not be made this session (they show the original).
|
||||
const failed = Number.isSafeInteger(status.failed) && status.failed > 0 ? ` · ${number(status.failed)} failed` : '';
|
||||
return `${number(missing)} missing ${noun(missing, 'thumbnail')}${failed}`;
|
||||
}
|
||||
|
||||
const SOURCES_SHARED = 'Monitoring and enrichment share one source list.';
|
||||
|
||||
Reference in New Issue
Block a user