feat(database): migrate old workflow rows in batches, then reclaim space
The workflow storage migration rewrites old-form image rows (text inline) to document ids, batch by batch: - Each batch is one BEGIN IMMEDIATE transaction of at most 200 rows or 16 MB of text, found by a rowid range search and read one row at a time. - Each text is stored or reused, then read back through the one accessor before its row changes; a text that would not come back identical fails and rolls back the whole batch. - A value that is not text stays inline and is counted as protected. - A crash or a stop leaves committed batches new form and the rest old form; both read the same, and a rerun continues. No media file is opened. - Unreferenced documents are swept at the end. The maintenance preview gains a read-only workflow_storage plan (records to rewrite, inline bytes where SQLite has octet_length, the Reclaim space and older-version notes). reclaim_space sweeps, runs VACUUM and truncates the WAL so the freed pages go back to the disk.
This commit is contained in:
Vito Sansevero
@@ -49,6 +49,8 @@ MAINTENANCE_OPERATIONS = (
|
||||
"optimize",
|
||||
"unused_components",
|
||||
"missing_links",
|
||||
# Applied as its own batched job (database/workflow_storage.py).
|
||||
"workflow_storage",
|
||||
)
|
||||
MAX_DETAILS = 50 # detail rows a plan lists
|
||||
MAX_RESULT_ROWS = 200 # changed records an apply describes (the job lists 200)
|
||||
@@ -420,6 +422,8 @@ class MaintenanceOperations:
|
||||
roots = [str(root) for root in roots]
|
||||
if operation == "optimize":
|
||||
return self._plan_optimize()
|
||||
if operation == "workflow_storage":
|
||||
return self._plan_workflow_storage()
|
||||
if operation == "missing_links":
|
||||
rows = _read_snapshot(self.model.db_path, self._link_rows)
|
||||
return self._plan_missing_links(rows, roots)
|
||||
@@ -668,7 +672,7 @@ class MaintenanceOperations:
|
||||
records changed and up to ``MAX_RESULT_ROWS`` descriptions of them
|
||||
(``{prompt_id[, kept_id]}`` or ``{image_path}``).
|
||||
"""
|
||||
if operation not in MAINTENANCE_OPERATIONS:
|
||||
if operation not in MAINTENANCE_OPERATIONS or operation == "workflow_storage":
|
||||
raise ValueError(f"Unknown maintenance operation {operation!r}")
|
||||
if operation == "optimize":
|
||||
# VACUUM cannot run in a transaction; the plan binds the database.
|
||||
|
||||
@@ -25,6 +25,7 @@ from .execution_outputs import ExecutionOutputOperations
|
||||
from .history_reads import HistoryReadOperations
|
||||
from .health import HealthOperations, iter_consistency_issues
|
||||
from .duplicate_copies import DuplicateCopyOperations
|
||||
from .workflow_storage import WorkflowStorageOperations
|
||||
from .workflow_documents import (
|
||||
DocumentReader,
|
||||
documents_of_rows,
|
||||
@@ -346,6 +347,7 @@ class PromptDatabase(
|
||||
HealthOperations,
|
||||
MaintenanceOperations,
|
||||
DuplicateCopyOperations,
|
||||
WorkflowStorageOperations,
|
||||
):
|
||||
"""Database operations class for managing prompts."""
|
||||
|
||||
|
||||
@@ -0,0 +1,321 @@
|
||||
"""The workflow storage migration and Reclaim space (#249 S6 carry-over).
|
||||
|
||||
``migrate_workflow_storage`` rewrites old-form image rows (workflow and
|
||||
prompt graph text inline) to the new form (a ``workflow_documents`` id,
|
||||
``database/workflow_documents.py``) in batches:
|
||||
|
||||
- each batch is one ``BEGIN IMMEDIATE`` transaction over at most
|
||||
``BATCH_ROWS`` rows or about ``BATCH_BYTES`` of text, found by a rowid
|
||||
range search (``NEXT_IDS_SQL``) and read one row at a time, so memory is
|
||||
bounded by one row plus the batch's distinct documents;
|
||||
- each text is stored (or its document reused) and read back through the
|
||||
one accessor before its row changes; a text that would not come back
|
||||
identical fails the batch (``DocumentMismatch``) and rolls it back whole;
|
||||
- a value that is not text (or is over ``MAX_DOCUMENT_BYTES``) stays inline
|
||||
and is counted as protected.
|
||||
|
||||
A crash or a stop between batches leaves committed batches in the new form
|
||||
and the rest in the old form; both read the same, and a rerun continues
|
||||
from the first old-form row. Only database rows change: no image or video
|
||||
file is opened.
|
||||
|
||||
``reclaim_space`` is the VACUUM afterwards: it sweeps unreferenced
|
||||
documents, rebuilds the file and checkpoints the WAL with TRUNCATE, so the
|
||||
freed pages go back to the disk.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
from contextlib import closing
|
||||
|
||||
from .workflow_documents import (
|
||||
TEXT_COLUMNS,
|
||||
DocumentMismatch,
|
||||
read_stored_text,
|
||||
release_documents,
|
||||
storable,
|
||||
store_text,
|
||||
sweep_documents,
|
||||
)
|
||||
|
||||
OPERATION = "workflow_storage"
|
||||
BATCH_ROWS = 200
|
||||
BATCH_BYTES = 16 * 2**20 # characters of text rewritten per transaction
|
||||
|
||||
# typeof() reads a column's type from the record header: no overflow page.
|
||||
NEXT_IDS_SQL = (
|
||||
"SELECT id FROM generated_images WHERE id>? AND "
|
||||
"(typeof(workflow_data)!='null' OR typeof(prompt_metadata)!='null') "
|
||||
"ORDER BY id LIMIT ?"
|
||||
)
|
||||
ROW_SQL = (
|
||||
"SELECT workflow_data, prompt_metadata, workflow_document_id, "
|
||||
"prompt_document_id FROM generated_images WHERE id=?"
|
||||
)
|
||||
UPDATE_SQL = (
|
||||
"UPDATE generated_images SET {assignments} WHERE id=?" # columns are constants
|
||||
)
|
||||
# One pass over the record headers: rows with inline text, and inline
|
||||
# values that are not text (they stay inline).
|
||||
INLINE_COUNTS_SQL = (
|
||||
"SELECT "
|
||||
"SUM(typeof(workflow_data)='text' OR typeof(prompt_metadata)='text'), "
|
||||
"SUM(typeof(workflow_data) NOT IN ('null','text')) "
|
||||
"+ SUM(typeof(prompt_metadata) NOT IN ('null','text')) "
|
||||
"FROM generated_images"
|
||||
)
|
||||
# octet_length() (SQLite 3.43+) reads no overflow page either.
|
||||
INLINE_BYTES_SQL = (
|
||||
"SELECT SUM(octet_length(workflow_data)), SUM(octet_length(prompt_metadata)) "
|
||||
"FROM generated_images"
|
||||
)
|
||||
|
||||
QUERY_PLAN_STATEMENTS = {
|
||||
"next old-form rows": (NEXT_IDS_SQL, (0, BATCH_ROWS)),
|
||||
"one row": (ROW_SQL, (1,)),
|
||||
}
|
||||
|
||||
|
||||
def _rewrite_row(conn, image_id, row, verified):
|
||||
"""Store the row's inline texts as documents; returns (changed, protected, size).
|
||||
|
||||
``verified`` maps document ids already read back in this batch to their
|
||||
text, so a document shared by many rows is decompressed once.
|
||||
"""
|
||||
changes, replaced = {}, []
|
||||
protected = size = 0
|
||||
for index, (column, reference) in enumerate(TEXT_COLUMNS):
|
||||
value, current = row[index], row[index + len(TEXT_COLUMNS)]
|
||||
if value is None:
|
||||
continue
|
||||
if not storable(value):
|
||||
protected += 1
|
||||
continue
|
||||
document_id = store_text(conn, value)
|
||||
if document_id not in verified:
|
||||
verified[document_id] = read_stored_text(conn, None, document_id)
|
||||
if verified[document_id] != value:
|
||||
raise DocumentMismatch(f"Image {image_id}: {column} would change")
|
||||
changes[column], changes[reference] = None, document_id
|
||||
if current is not None and current != document_id:
|
||||
replaced.append(current)
|
||||
size += len(value)
|
||||
if changes:
|
||||
assignments = ",".join(f"{key}=?" for key in changes)
|
||||
conn.execute(
|
||||
UPDATE_SQL.format(assignments=assignments),
|
||||
[*changes.values(), image_id],
|
||||
)
|
||||
release_documents(conn, replaced)
|
||||
return bool(changes), protected, size
|
||||
|
||||
|
||||
def migrate_batch(conn, after_id, max_rows=BATCH_ROWS, max_bytes=BATCH_BYTES):
|
||||
"""Rewrite the next batch after ``after_id`` in one transaction.
|
||||
|
||||
Returns ``{last_id, rows, protected, done}``: the last row looked at,
|
||||
rows rewritten, inline values kept (not text) and whether there was
|
||||
nothing left to look at.
|
||||
"""
|
||||
if conn.in_transaction:
|
||||
raise RuntimeError("The migration needs its own transaction")
|
||||
conn.execute("BEGIN IMMEDIATE")
|
||||
try:
|
||||
ids = [row[0] for row in conn.execute(NEXT_IDS_SQL, (after_id, max_rows))]
|
||||
result = {"last_id": after_id, "rows": 0, "protected": 0, "done": not ids}
|
||||
verified, size = {}, 0
|
||||
for image_id in ids:
|
||||
row = conn.execute(ROW_SQL, (image_id,)).fetchone()
|
||||
changed, protected, row_size = _rewrite_row(conn, image_id, row, verified)
|
||||
result["rows"] += int(changed)
|
||||
result["protected"] += protected
|
||||
result["last_id"] = image_id
|
||||
size += row_size
|
||||
if size >= max_bytes:
|
||||
break
|
||||
conn.commit()
|
||||
return result
|
||||
except BaseException:
|
||||
conn.rollback()
|
||||
raise
|
||||
|
||||
|
||||
def _file_bytes(db_path):
|
||||
return sum(
|
||||
os.path.getsize(path)
|
||||
for path in (db_path, db_path + "-wal")
|
||||
if os.path.exists(path)
|
||||
)
|
||||
|
||||
|
||||
class WorkflowStorageOperations:
|
||||
"""PromptDatabase mixin: the workflow storage migration and Reclaim space."""
|
||||
|
||||
def migrate_workflow_storage(
|
||||
self,
|
||||
*,
|
||||
on_batch=None,
|
||||
should_stop=None,
|
||||
max_rows=BATCH_ROWS,
|
||||
max_bytes=BATCH_BYTES,
|
||||
) -> dict:
|
||||
"""Rewrite every old-form row, batch by batch (blocking).
|
||||
|
||||
``on_batch(progress)`` runs after each committed batch with the
|
||||
running totals ``{rows, protected, batches, last_id}``;
|
||||
``should_stop()`` is asked before each batch. Returns ``{rows,
|
||||
protected, batches, stopped, documents_removed}``. Unreferenced
|
||||
documents are swept once every row was looked at.
|
||||
"""
|
||||
conn = self.model.get_connection()
|
||||
totals = {"rows": 0, "protected": 0, "batches": 0, "last_id": 0}
|
||||
stopped = False
|
||||
while True:
|
||||
if should_stop is not None and should_stop():
|
||||
stopped = True
|
||||
break
|
||||
batch = migrate_batch(conn, totals["last_id"], max_rows, max_bytes)
|
||||
if batch["done"]:
|
||||
break
|
||||
totals["rows"] += batch["rows"]
|
||||
totals["protected"] += batch["protected"]
|
||||
totals["batches"] += 1
|
||||
totals["last_id"] = batch["last_id"]
|
||||
if on_batch is not None:
|
||||
on_batch(dict(totals))
|
||||
removed = 0
|
||||
if not stopped:
|
||||
with conn:
|
||||
removed = sweep_documents(conn)
|
||||
return {
|
||||
"rows": totals["rows"],
|
||||
"protected": totals["protected"],
|
||||
"batches": totals["batches"],
|
||||
"stopped": stopped,
|
||||
"documents_removed": removed,
|
||||
}
|
||||
|
||||
def workflow_storage_counts(self, conn) -> dict:
|
||||
"""``{rows, protected, documents, inline_bytes}`` on ``conn`` (reads only).
|
||||
|
||||
``inline_bytes`` is None where SQLite has no ``octet_length``.
|
||||
"""
|
||||
rows, protected = conn.execute(INLINE_COUNTS_SQL).fetchone()
|
||||
documents = conn.execute("SELECT COUNT(*) FROM workflow_documents").fetchone()
|
||||
inline_bytes = None
|
||||
if sqlite3.sqlite_version_info >= (3, 43, 0):
|
||||
row = conn.execute(INLINE_BYTES_SQL).fetchone()
|
||||
inline_bytes = sum(value or 0 for value in row)
|
||||
return {
|
||||
"rows": rows or 0,
|
||||
"protected": protected or 0,
|
||||
"documents": documents[0],
|
||||
"inline_bytes": inline_bytes,
|
||||
}
|
||||
|
||||
def reclaim_space(self) -> dict:
|
||||
"""Sweep, VACUUM and truncate the WAL (blocking); file bytes before/after.
|
||||
|
||||
The caller holds the registry's ``exclusive`` slot. Raises whatever
|
||||
stopped VACUUM; the database is then unchanged.
|
||||
"""
|
||||
db_path = self.model.db_path
|
||||
conn = self.model.get_connection()
|
||||
with conn:
|
||||
removed = sweep_documents(conn)
|
||||
before = _file_bytes(db_path)
|
||||
self.model.vacuum_database()
|
||||
with closing(sqlite3.connect(db_path)) as checkpoint:
|
||||
busy = checkpoint.execute("PRAGMA wal_checkpoint(TRUNCATE)").fetchone()[0]
|
||||
return {
|
||||
"before_bytes": before,
|
||||
"after_bytes": _file_bytes(db_path),
|
||||
"documents_removed": removed,
|
||||
"checkpoint_busy": bool(busy),
|
||||
}
|
||||
|
||||
# ── preview ───────────────────────────────────────────────────────
|
||||
|
||||
def _plan_workflow_storage(self):
|
||||
"""The reviewed plan of the migration (one read-only snapshot)."""
|
||||
from .maintenance import (
|
||||
MaintenanceDraft,
|
||||
_Fingerprint,
|
||||
_megabytes,
|
||||
_plural,
|
||||
_read_snapshot,
|
||||
)
|
||||
|
||||
def read(conn):
|
||||
sizes = [
|
||||
conn.execute(f"PRAGMA {name}").fetchone()[0]
|
||||
for name in ("page_size", "page_count", "freelist_count")
|
||||
]
|
||||
return sizes, self.workflow_storage_counts(conn)
|
||||
|
||||
(page_size, page_count, free_pages), counts = _read_snapshot(
|
||||
self.model.db_path, read
|
||||
)
|
||||
size = page_size * page_count
|
||||
details = [
|
||||
{
|
||||
"label": "Database file",
|
||||
"detail": f"{_megabytes(size)}, of which "
|
||||
f"{_megabytes(free_pages * page_size)} is free pages",
|
||||
},
|
||||
{
|
||||
"label": "Image records",
|
||||
"detail": f"{_plural(counts['rows'], 'record')} keep their workflow "
|
||||
"or prompt graph inline; each distinct text is stored once, "
|
||||
"compressed",
|
||||
},
|
||||
]
|
||||
if counts["inline_bytes"] is not None:
|
||||
details.append(
|
||||
{
|
||||
"label": "Inline text",
|
||||
"detail": f"{_megabytes(counts['inline_bytes'])} now",
|
||||
}
|
||||
)
|
||||
if counts["documents"]:
|
||||
details.append(
|
||||
{
|
||||
"label": "Already stored once",
|
||||
"detail": _plural(counts["documents"], "workflow document"),
|
||||
}
|
||||
)
|
||||
details += [
|
||||
{"label": "Reclaim space", "detail": RECLAIM_NOTE},
|
||||
{"label": "Older versions", "detail": OLDER_VERSIONS_NOTE},
|
||||
]
|
||||
protected = counts["protected"]
|
||||
return MaintenanceDraft(
|
||||
operation=OPERATION,
|
||||
fingerprint=_Fingerprint(OPERATION).hexdigest(),
|
||||
targets=[],
|
||||
scope={
|
||||
"count": counts["rows"],
|
||||
"label": _plural(counts["rows"], "image record") + " to rewrite",
|
||||
},
|
||||
protected={
|
||||
"count": protected,
|
||||
"reasons": (
|
||||
[f"{_plural(protected, 'stored value')} not text; kept as is"]
|
||||
if protected
|
||||
else []
|
||||
),
|
||||
},
|
||||
media={"deletes_files": False, "label": MEDIA_NOTE},
|
||||
details=details,
|
||||
)
|
||||
|
||||
|
||||
MEDIA_NOTE = "No image or video file is opened or changed; only the database."
|
||||
RECLAIM_NOTE = (
|
||||
"When every record is rewritten, Reclaim space runs alone to give the freed "
|
||||
"space back to the disk (VACUUM). The database is locked while it runs."
|
||||
)
|
||||
OLDER_VERSIONS_NOTE = (
|
||||
"An older Prompt Manager version shows rewritten images without their "
|
||||
"workflow. Restoring the safety backup brings back the old form."
|
||||
)
|
||||
Reference in New Issue
Block a user