Vito Sansevero
0ad791a6f4
The reviewed maintenance plan gains workflow_storage. Applying it starts "Store each workflow once": kind maintenance, records admission, cancellable, groups changed/protected/errors. - The promised verified safety backup comes first; a failed backup ends the job with nothing rewritten. The apply's up-front disk check covers it. - Old-form rows are rewritten batch by batch on the event loop's executor, with progress after each committed batch. Cancel stops between batches; a failed batch rolls back whole and fails the job with a coded message. Either way every record reads as before and the next apply continues. - When every row was looked at, "Reclaim space" starts: exclusive (runs alone), not cancellable. It checks the disk for VACUUM's copies of the live pages, then sweeps, runs VACUUM and truncates the WAL. If another job holds the registry it does not start, and the finished migration says so and names Optimize database. - No image or video file is opened; tests track every open under the output folder and compare the media bytes and mtimes.
334 lines
12 KiB
Python
334 lines
12 KiB
Python
"""The workflow storage migration and Reclaim space (#249 S6 carry-over).
|
|
|
|
``migrate_workflow_storage`` rewrites old-form image rows (workflow and
|
|
prompt graph text inline) to the new form (a ``workflow_documents`` id,
|
|
``database/workflow_documents.py``) in batches:
|
|
|
|
- each batch is one ``BEGIN IMMEDIATE`` transaction over at most
|
|
``BATCH_ROWS`` rows or about ``BATCH_BYTES`` of text, found by a rowid
|
|
range search (``NEXT_IDS_SQL``) and read one row at a time, so memory is
|
|
bounded by one row plus the batch's distinct documents;
|
|
- each text is stored (or its document reused) and read back through the
|
|
one accessor before its row changes; a text that would not come back
|
|
identical fails the batch (``DocumentMismatch``) and rolls it back whole;
|
|
- a value that is not text (or is over ``MAX_DOCUMENT_BYTES``) stays inline
|
|
and is counted as protected.
|
|
|
|
A crash or a stop between batches leaves committed batches in the new form
|
|
and the rest in the old form; both read the same, and a rerun continues
|
|
from the first old-form row. Only database rows change: no image or video
|
|
file is opened.
|
|
|
|
``reclaim_space`` is the VACUUM afterwards: it sweeps unreferenced
|
|
documents, rebuilds the file and checkpoints the WAL with TRUNCATE, so the
|
|
freed pages go back to the disk.
|
|
"""
|
|
|
|
import os
|
|
import sqlite3
|
|
from contextlib import closing
|
|
|
|
from .workflow_documents import (
|
|
TEXT_COLUMNS,
|
|
DocumentMismatch,
|
|
read_stored_text,
|
|
release_documents,
|
|
storable,
|
|
store_text,
|
|
sweep_documents,
|
|
)
|
|
|
|
OPERATION = "workflow_storage"
|
|
BATCH_ROWS = 200
|
|
BATCH_BYTES = 16 * 2**20 # characters of text rewritten per transaction
|
|
|
|
# typeof() reads a column's type from the record header: no overflow page.
|
|
NEXT_IDS_SQL = (
|
|
"SELECT id FROM generated_images WHERE id>? AND "
|
|
"(typeof(workflow_data)!='null' OR typeof(prompt_metadata)!='null') "
|
|
"ORDER BY id LIMIT ?"
|
|
)
|
|
ROW_SQL = (
|
|
"SELECT workflow_data, prompt_metadata, workflow_document_id, "
|
|
"prompt_document_id FROM generated_images WHERE id=?"
|
|
)
|
|
UPDATE_SQL = (
|
|
"UPDATE generated_images SET {assignments} WHERE id=?" # columns are constants
|
|
)
|
|
# One pass over the record headers: rows with inline text, and inline
|
|
# values that are not text (they stay inline).
|
|
INLINE_COUNTS_SQL = (
|
|
"SELECT "
|
|
"SUM(typeof(workflow_data)='text' OR typeof(prompt_metadata)='text'), "
|
|
"SUM(typeof(workflow_data) NOT IN ('null','text')) "
|
|
"+ SUM(typeof(prompt_metadata) NOT IN ('null','text')) "
|
|
"FROM generated_images"
|
|
)
|
|
# octet_length() (SQLite 3.43+) reads no overflow page either.
|
|
INLINE_BYTES_SQL = (
|
|
"SELECT SUM(octet_length(workflow_data)), SUM(octet_length(prompt_metadata)) "
|
|
"FROM generated_images"
|
|
)
|
|
|
|
QUERY_PLAN_STATEMENTS = {
|
|
"next old-form rows": (NEXT_IDS_SQL, (0, BATCH_ROWS)),
|
|
"one row": (ROW_SQL, (1,)),
|
|
}
|
|
|
|
|
|
def _rewrite_row(conn, image_id, row, verified):
|
|
"""Store the row's inline texts as documents; returns (changed, protected, size).
|
|
|
|
``verified`` maps document ids already read back in this batch to their
|
|
text, so a document shared by many rows is decompressed once.
|
|
"""
|
|
changes, replaced = {}, []
|
|
protected = size = 0
|
|
for index, (column, reference) in enumerate(TEXT_COLUMNS):
|
|
value, current = row[index], row[index + len(TEXT_COLUMNS)]
|
|
if value is None:
|
|
continue
|
|
if not storable(value):
|
|
protected += 1
|
|
continue
|
|
document_id = store_text(conn, value)
|
|
if document_id not in verified:
|
|
verified[document_id] = read_stored_text(conn, None, document_id)
|
|
if verified[document_id] != value:
|
|
raise DocumentMismatch(f"Image {image_id}: {column} would change")
|
|
changes[column], changes[reference] = None, document_id
|
|
if current is not None and current != document_id:
|
|
replaced.append(current)
|
|
size += len(value)
|
|
if changes:
|
|
assignments = ",".join(f"{key}=?" for key in changes)
|
|
conn.execute(
|
|
UPDATE_SQL.format(assignments=assignments),
|
|
[*changes.values(), image_id],
|
|
)
|
|
release_documents(conn, replaced)
|
|
return bool(changes), protected, size
|
|
|
|
|
|
def migrate_batch(conn, after_id, max_rows=BATCH_ROWS, max_bytes=BATCH_BYTES):
|
|
"""Rewrite the next batch after ``after_id`` in one transaction.
|
|
|
|
Returns ``{last_id, rows, protected, done}``: the last row looked at,
|
|
rows rewritten, inline values kept (not text) and whether there was
|
|
nothing left to look at.
|
|
"""
|
|
if conn.in_transaction:
|
|
raise RuntimeError("The migration needs its own transaction")
|
|
conn.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
ids = [row[0] for row in conn.execute(NEXT_IDS_SQL, (after_id, max_rows))]
|
|
result = {"last_id": after_id, "rows": 0, "protected": 0, "done": not ids}
|
|
verified, size = {}, 0
|
|
for image_id in ids:
|
|
row = conn.execute(ROW_SQL, (image_id,)).fetchone()
|
|
changed, protected, row_size = _rewrite_row(conn, image_id, row, verified)
|
|
result["rows"] += int(changed)
|
|
result["protected"] += protected
|
|
result["last_id"] = image_id
|
|
size += row_size
|
|
if size >= max_bytes:
|
|
break
|
|
conn.commit()
|
|
return result
|
|
except BaseException:
|
|
conn.rollback()
|
|
raise
|
|
|
|
|
|
def _file_bytes(db_path):
|
|
return sum(
|
|
os.path.getsize(path)
|
|
for path in (db_path, db_path + "-wal")
|
|
if os.path.exists(path)
|
|
)
|
|
|
|
|
|
class WorkflowStorageOperations:
|
|
"""PromptDatabase mixin: the workflow storage migration and Reclaim space."""
|
|
|
|
def migrate_workflow_storage(
|
|
self,
|
|
*,
|
|
on_batch=None,
|
|
should_stop=None,
|
|
max_rows=BATCH_ROWS,
|
|
max_bytes=BATCH_BYTES,
|
|
) -> dict:
|
|
"""Rewrite every old-form row, batch by batch (blocking).
|
|
|
|
``on_batch(progress)`` runs after each committed batch with the
|
|
running totals ``{rows, protected, batches, last_id}``;
|
|
``should_stop()`` is asked before each batch. Returns ``{rows,
|
|
protected, batches, stopped, documents_removed}``. Unreferenced
|
|
documents are swept once every row was looked at.
|
|
"""
|
|
totals = {"rows": 0, "protected": 0, "batches": 0, "last_id": 0}
|
|
stopped = False
|
|
while True:
|
|
if should_stop is not None and should_stop():
|
|
stopped = True
|
|
break
|
|
batch = self.migrate_workflow_batch(totals["last_id"], max_rows, max_bytes)
|
|
if batch["done"]:
|
|
break
|
|
totals["rows"] += batch["rows"]
|
|
totals["protected"] += batch["protected"]
|
|
totals["batches"] += 1
|
|
totals["last_id"] = batch["last_id"]
|
|
if on_batch is not None:
|
|
on_batch(dict(totals))
|
|
removed = 0 if stopped else self.sweep_workflow_documents()
|
|
return {
|
|
"rows": totals["rows"],
|
|
"protected": totals["protected"],
|
|
"batches": totals["batches"],
|
|
"stopped": stopped,
|
|
"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 counts["inline_bytes"] is not None:
|
|
details.append(
|
|
{
|
|
"label": "Inline text",
|
|
"detail": f"{_megabytes(counts['inline_bytes'])} now",
|
|
}
|
|
)
|
|
if counts["documents"]:
|
|
details.append(
|
|
{
|
|
"label": "Already stored once",
|
|
"detail": _plural(counts["documents"], "workflow document"),
|
|
}
|
|
)
|
|
details += [
|
|
{"label": "Reclaim space", "detail": RECLAIM_NOTE},
|
|
{"label": "Older versions", "detail": OLDER_VERSIONS_NOTE},
|
|
]
|
|
protected = counts["protected"]
|
|
return MaintenanceDraft(
|
|
operation=OPERATION,
|
|
fingerprint=_Fingerprint(OPERATION).hexdigest(),
|
|
targets=[],
|
|
scope={
|
|
"count": counts["rows"],
|
|
"label": _plural(counts["rows"], "image record") + " to rewrite",
|
|
},
|
|
protected={
|
|
"count": protected,
|
|
"reasons": (
|
|
[f"{_plural(protected, 'stored value')} not text; kept as is"]
|
|
if protected
|
|
else []
|
|
),
|
|
},
|
|
media={"deletes_files": False, "label": MEDIA_NOTE},
|
|
details=details,
|
|
)
|
|
|
|
|
|
MEDIA_NOTE = "No image or video file is opened or changed; only the database."
|
|
RECLAIM_NOTE = (
|
|
"When every record is rewritten, Reclaim space runs alone to give the freed "
|
|
"space back to the disk (VACUUM). The database is locked while it runs."
|
|
)
|
|
OLDER_VERSIONS_NOTE = (
|
|
"An older Prompt Manager version shows rewritten images without their "
|
|
"workflow. Restoring the safety backup brings back the old form."
|
|
)
|