nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
f9835aa4d66f746c40c30b2f00668b82d6e2b549
ComfyAssets-ComfyUI_PromptM…/database/models.py
T
Vito 28ad3a1a1f feat(system): backup and restore backend (System S8) (#335) 
* refactor(database): stage restores before publishing

restore_from_file is now stage_restore + publish_staged (database/
restore_stage.py). The stage lives in a .pm-restore-* folder in the
library's own folder, is migrated, has damaged workflow documents repaired
from an inline copy or cleared, is verified with every document read and is
bound by its SHA-256. Publishing keeps a unique, standalone, 0600 safety
backup (microsecond names), copies the pages through the backup API as
before and then truncates the WAL, which stayed database-sized after a
restore. The safety backup writer moves to database/safety_backup.py so
the database layer and maintenance share it.

* fix(database): read the verify streams one row at a time

The provenance and review-evidence checks iterated their cursors; on a
managed connection iteration prefetches 128 rows, so up to 128 large values
were held at once while the docstrings promised one. Both now use fetchone.

* feat(backup): review a restore, run it as an Activity job, verify backups

- POST /restore/preview stages an upload in the library's folder and
  answers what it holds beside the live library, writing nothing live;
  DELETE discards it; at most two stages are held (429 too-large).
- POST /restore/apply rechecks the live fingerprint (409 plan-changed) and
  starts "Database restore": exclusive, not cancellable, indeterminate. It
  waits for request leases, publishes the exact staged bytes, swaps the
  database, prunes safety backups and logs counts, repairs and the backup.
- The legacy /restore runs the same stage and publish as an exclusive hold
  with its old response, so a restore and the health check, Move library
  and every other job refuse each other.
- JobRegistry.lease: Needs review decisions and maintenance and restore
  previews are refused while an exclusive job runs, and a restore waits
  for them before it switches the file.
- GET /backup holds a guarded slot, verifies a standalone snapshot before
  streaming it and removes the closed temp file; GET /backup/status lists
  the safety backups on disk and this session's verified download.

* docs(testing): backup and restore contract

* test(system): a restore refuses a running health check

The health route test relied on a legacy restore interleaving with a running
check; the restore is now exclusive and refuses it, naming the check. The
diagnostics backup tests follow the backup handler into
py/api/backup_restore.py.

* fix(backup): drop stale restore previews off the event loop

A restore drops the previews reviewed against the old library; their stage
files are now removed in the executor, with the old database's connections.

* test(backup): wait for the backup hold to end before reading its state

CI (py3.11/3.12 on Linux, every macOS) read the registry right after the
client had the last byte, before the handler left its hold and removed its
temp file: the job was still running. The tests now wait on that state.

* test(images): the scan dialog's Back up now leaves a finished backup hold

A backup download now holds a guarded registry slot while it streams
(#249 S8), so the jobs list holds that finished hold after the dialog's
Back up now; the spec still proves Cancel starts no scan.

* fix(backup): final review round for the restore and backup routes

- I1: once the stage's copy replaced the library the restore succeeded:
  pruning, reopening self.db, dropping old previews and closing the old
  object only log, so the job finishes and the legacy /restore answers 200
  instead of "the library was not changed". The generation is always bumped.
- M1: a backup download holds its guarded slot only while the snapshot is
  written; the stream runs outside it, so a stalled download blocks nothing.
- M2: a restore preview leases only its read of the live library, not the
  upload or the staging; a restore whose drain still times out keeps the
  reviewed stage under its token.
- M3: the reviewed-preview store changes only on the event loop; stage
  folders are removed in the executor.
- M4: a document whose body inflates completely while only its recorded
  hash or only its length disagrees is re-keyed (or merged into the
  document holding that text) instead of dropped.
- M5: a backup download is refused only when the snapshot is not SQLite or
  fails integrity_check; a semantic check failure still downloads, logged,
  sent as X-PromptManager-Verify and recorded as not verified.
- M6: a refused, failed or cancelled preview removes its stage, also when
  the cancel lands while a worker still stages; a timer removes expired
  stages without another request.
- M7, N1: repair reads rows one at a time and counts distinct images.
- N3: the live fingerprint never commits a transaction it did not open.

* test(backup): close the old library the close-failure test leaves open

CI (py3.12 Linux and the coverage gate) found the old object's
executor-thread connections still open at teardown: the test patches its
close_all to fail, and once the object is collected close_all_instances no
longer reaches it. The test now closes it itself.

* fix(backup): re-review round for restore repair, apply and backup cleanup

- R1: when re-keying merges an intact body into a holder whose own body was
  damaged, the rewritten holder leaves the damaged list, so the clear pass
  no longer deletes it and unlinks the images just moved to it.
- R2: a losing duplicate apply removes a stage only when the store still
  held it, never one another apply's job has taken.
- R3: a restore that times out waiting for requests keeps its review only
  while it is unexpired and its slot is free; otherwise the stage goes and
  the job says to try again.
- R4: a backup download cancelled while its snapshot is written removes
  the snapshot once the worker finishes.
2026-10-25 10:17:07 -07:00

1320 lines
54 KiB
Python
Raw Blame History

"""
Database schema and models for KikoTextEncode prompt storage.
"""
import contextlib
import datetime
import json
import math
import sqlite3
import os
import threading
import weakref
from contextlib import closing
from pathlib import Path
from typing import Optional, Set, Tuple
from .managed_connection import (
ManagedConnection,
await_closed,
retire_connections,
)
from . import restore_stage
from .library_schema import library_schema_error, migrate_library_state
from .image_favorites import image_favorites_schema_error, migrate_image_favorites
from .annotation_schema import annotation_schema_error, migrate_annotation_state
from .review_schema import (
migrate_review_state,
review_schema_error,
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,
)
# Import logging system
try:
from ..utils.logging_config import get_logger
except ImportError:
import sys
current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, current_dir)
from utils.logging_config import get_logger
# last_used_at is UTC ISO text so ORDER BY can compare it as text. The backfill
# writes milliseconds (SQLite's %f); live writes use microseconds so runs in the
# same millisecond still order correctly. Both forms sort correctly together.
ISO_TIMESTAMP_FORMAT = "%Y-%m-%dT%H:%M:%f+00:00"
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):
return "Pair provenance must be JSON text"
if len(encoded) > MAX_PAIR_PROVENANCE_LENGTH:
return "Pair provenance exceeds the size limit"
try:
value = json.loads(encoded)
except (ValueError, TypeError, RecursionError):
return "Pair provenance is not readable JSON"
if not isinstance(value, dict):
return "Pair provenance must be a JSON object"
pending = [(iter((value,)), 0)]
while pending:
children, depth = pending[-1]
try:
item = next(children)
except StopIteration:
pending.pop()
continue
if depth > MAX_PAIR_PROVENANCE_DEPTH:
return "Pair provenance exceeds the nesting limit"
if isinstance(item, float) and not math.isfinite(item):
return "Pair provenance contains nonfinite numbers"
if isinstance(item, dict):
pending.append((iter(item.values()), depth + 1))
elif isinstance(item, list):
pending.append((iter(item), depth + 1))
return None
def utc_now_iso() -> str:
"""Current UTC time as ISO text with microseconds.
Example: 2026-10-07T18:39:06.758123+00:00
"""
return datetime.datetime.now(datetime.timezone.utc).isoformat(
timespec="microseconds"
)
def normalize_image_path(path: str) -> str:
"""Canonical form of an image path for generated_images.file_path.
Absolute, with '..' and '.' collapsed and the platform's case folding
applied, so two spellings of one file (or the same basename in two date
folders) compare the way the filesystem would.
"""
return os.path.normcase(os.path.normpath(os.path.abspath(path)))
# How long a connection waits for a locked database before giving up. Thirty
# seconds covers a scan or a thumbnail run writing while ComfyUI saves images,
# and the slow disks of CI runners.
BUSY_TIMEOUT_SECONDS = 30
SQLITE_HEADER = b"SQLite format 3\x00"
REQUIRED_PROMPT_COLUMNS = ("id", "text", "created_at")
def configure_connection(conn: sqlite3.Connection) -> None:
"""Apply the locking and durability settings every managed connection shares.
WAL lets readers run beside the single writer. In WAL mode
``synchronous=NORMAL`` keeps the database consistent after an application
or ComfyUI crash and flushes only at checkpoints, so a commit no longer
holds SQLite's write lock across a disk flush. Under ``FULL`` a slow flush
(common on Windows) let other writers' polling busy handlers starve past
BUSY_TIMEOUT_SECONDS and fail with "database is locked" (#242). Without
WAL (for example an in-memory database) SQLite's FULL default is kept.
"""
journal_mode = conn.execute("PRAGMA journal_mode = WAL").fetchone()[0]
conn.execute("PRAGMA foreign_keys = ON")
conn.execute(f"PRAGMA busy_timeout = {BUSY_TIMEOUT_SECONDS * 1000}")
if str(journal_mode).lower() == "wal":
conn.execute("PRAGMA synchronous = NORMAL")
class PromptModel:
"""Database model for prompt storage and schema management."""
# Paths (normalised, see _schema_key) whose schema has been created and
# migrated in this process. Many PromptDatabase instances point at the
# same file; only the first one pays for the schema check.
_initialized_paths: Set[str] = set()
_init_lock = threading.Lock()
# Every live model, so a shutdown or a test teardown can close them all
# before the database file is deleted (Windows refuses otherwise).
_instances: "weakref.WeakSet[PromptModel]" = weakref.WeakSet()
_instances_lock = threading.Lock()
def __init__(self, db_path: str):
"""
Initialize the database model.
Args:
db_path: Absolute path to the SQLite database file
"""
self.logger = get_logger("prompt_manager.database.models")
self.logger.debug(f"Initializing database model with path: {db_path}")
self.db_path = db_path
# One connection per thread (see get_connection); every open connection
# is also tracked here so close_all can clean up across threads.
self._local = threading.local()
self._connections: Set[sqlite3.Connection] = set()
self._conn_lock = threading.Lock()
# Bumped by _close_all_connections so threads holding a closed
# connection reopen on their next get_connection().
self._generation = 0
with PromptModel._instances_lock:
PromptModel._instances.add(self)
self._ensure_database_exists()
def _ensure_database_exists(self) -> None:
"""
Create database and tables if they don't exist.
Sets up the database schema including tables for prompts and generated images,
creates necessary indexes, and applies any pending migrations.
Raises:
Exception: If database creation fails
"""
key = self._schema_key()
with PromptModel._init_lock:
if key in PromptModel._initialized_paths and os.path.exists(self.db_path):
return
try:
conn = sqlite3.connect(self.db_path, timeout=BUSY_TIMEOUT_SECONDS)
try:
configure_connection(conn)
self._create_tables(conn)
self._create_indexes(conn)
conn.commit()
finally:
conn.close()
except Exception as e:
self.logger.error(f"Error creating database: {e}")
raise
PromptModel._initialized_paths.add(key)
def _schema_key(self) -> str:
"""Normalised identity of db_path for the once-per-process schema cache."""
return os.path.normcase(os.path.realpath(self.db_path))
@classmethod
def reset_schema_cache(cls) -> None:
"""Forget which databases were initialised (tests and restores)."""
with cls._init_lock:
cls._initialized_paths.clear()
@contextlib.contextmanager
def _table_rebuild(self, conn: sqlite3.Connection):
"""Run a copy-and-swap table rebuild with foreign keys disabled.
DROP TABLE on a parent table with foreign_keys ON performs an implicit
DELETE that cascades into child rows (generated_images, prompt_tags),
so the rebuild runs with the pragma off. The pragma is a no-op inside
a transaction, hence the commits around it. Errors roll back and
propagate after the pragma is restored.
"""
conn.commit()
conn.execute("PRAGMA foreign_keys = OFF")
# Explicit BEGIN: in sqlite3's legacy transaction mode DDL autocommits,
# so without it a failed copy would leave the *_new table behind.
conn.execute("BEGIN")
try:
yield
conn.commit()
except BaseException:
conn.rollback()
raise
finally:
conn.execute("PRAGMA foreign_keys = ON")
def _create_tables(self, conn: sqlite3.Connection) -> None:
"""
Create the prompts and generated_images tables with all required columns.
Args:
conn: Active database connection
Creates:
- prompts table: Stores prompt text and metadata
- generated_images table: Links generated images to their source prompts
"""
conn.execute("""
CREATE TABLE IF NOT EXISTS prompts (
id INTEGER PRIMARY KEY AUTOINCREMENT,
text TEXT NOT NULL,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
category TEXT,
tags TEXT,
rating INTEGER CHECK(rating >= 1 AND rating <= 5),
notes TEXT,
hash TEXT UNIQUE,
last_used_at TIMESTAMP,
run_count INTEGER NOT NULL DEFAULT 0
)
""")
# Create images table for gallery functionality
conn.execute("""
CREATE TABLE IF NOT EXISTS generated_images (
id INTEGER PRIMARY KEY AUTOINCREMENT,
prompt_id INTEGER NOT NULL,
image_path TEXT NOT NULL,
filename TEXT NOT NULL,
generation_time TIMESTAMP DEFAULT CURRENT_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)
)
""")
# Create normalized tag tables (junction table pattern)
conn.execute("""
CREATE TABLE IF NOT EXISTS tags (
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT NOT NULL UNIQUE
)
""")
conn.execute("""
CREATE TABLE IF NOT EXISTS prompt_tags (
prompt_id INTEGER NOT NULL,
tag_id INTEGER NOT NULL,
PRIMARY KEY (prompt_id, tag_id),
FOREIGN KEY (prompt_id) REFERENCES prompts(id) ON DELETE CASCADE,
FOREIGN KEY (tag_id) REFERENCES tags(id) ON DELETE CASCADE
)
""")
# Check if we need to migrate from old schema with workflow_name
self._migrate_workflow_name_removal(conn)
# Fix foreign key data type mismatch
self._migrate_foreign_key_types(conn)
# Migrate JSON tags to normalized junction tables
self._migrate_json_tags_to_junction(conn)
# Usage tracking columns (after any migration that rebuilds prompts)
self._migrate_add_usage_columns(conn)
# Image uniqueness by full path instead of basename (3.2.4)
self._migrate_image_uniqueness_by_path(conn)
self._migrate_prompt_pairs(conn)
migrate_execution_history(conn)
migrate_library_state(conn)
migrate_image_favorites(conn)
migrate_saved_views(conn)
# A review group without its annotation group is refused before the
# annotation install could make it look valid (catalog only; the rows
# are validated once, in migrate_review_state).
review_error = review_without_annotations_error(conn)
if review_error:
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")
try:
conn.execute("""
CREATE TABLE IF NOT EXISTS prompt_pairs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
positive_prompt_id INTEGER NOT NULL,
negative_prompt_id INTEGER,
title TEXT,
category TEXT,
rating INTEGER CHECK(rating BETWEEN 1 AND 5),
notes TEXT,
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
last_used_at TEXT,
run_count INTEGER NOT NULL DEFAULT 0,
archived INTEGER NOT NULL DEFAULT 0 CHECK(archived IN (0, 1)),
FOREIGN KEY(positive_prompt_id)
REFERENCES prompts(id) ON DELETE RESTRICT,
FOREIGN KEY(negative_prompt_id)
REFERENCES prompts(id) ON DELETE RESTRICT,
UNIQUE(positive_prompt_id, negative_prompt_id)
)
""")
conn.execute("""
CREATE TABLE IF NOT EXISTS prompt_pair_tags (
pair_id INTEGER NOT NULL,
tag_id INTEGER NOT NULL,
PRIMARY KEY(pair_id, tag_id),
FOREIGN KEY(pair_id) REFERENCES prompt_pairs(id) ON DELETE CASCADE,
FOREIGN KEY(tag_id) REFERENCES tags(id) ON DELETE CASCADE
)
""")
conn.execute("""
CREATE TABLE IF NOT EXISTS image_prompt_pairs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
image_id INTEGER NOT NULL,
pair_id INTEGER NOT NULL,
provenance TEXT NOT NULL,
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY(image_id)
REFERENCES generated_images(id) ON DELETE CASCADE,
FOREIGN KEY(pair_id) REFERENCES prompt_pairs(id) ON DELETE RESTRICT,
UNIQUE(pair_id, image_id, provenance)
)
""")
self._create_pair_indexes(conn)
error = self._pair_schema_error(conn)
if error:
raise sqlite3.DatabaseError(error)
conn.execute("RELEASE prompt_pair_schema")
except BaseException:
conn.execute("ROLLBACK TO prompt_pair_schema")
conn.execute("RELEASE prompt_pair_schema")
raise
@staticmethod
def _create_pair_indexes(conn: sqlite3.Connection) -> None:
conn.execute(
"CREATE UNIQUE INDEX IF NOT EXISTS idx_prompt_pairs_positive_only "
"ON prompt_pairs(positive_prompt_id) WHERE negative_prompt_id IS NULL"
)
for name, target in (
("idx_prompt_pairs_negative", "prompt_pairs(negative_prompt_id)"),
("idx_prompt_pair_tags_tag", "prompt_pair_tags(tag_id)"),
("idx_image_prompt_pairs_image", "image_prompt_pairs(image_id)"),
):
conn.execute(f"CREATE INDEX IF NOT EXISTS {name} ON {target}")
@staticmethod
def _pair_schema_error(conn: sqlite3.Connection) -> Optional[str]:
required = {
"prompt_pairs": {
"id",
"positive_prompt_id",
"negative_prompt_id",
"title",
"category",
"rating",
"notes",
"created_at",
"updated_at",
"last_used_at",
"run_count",
"archived",
},
"prompt_pair_tags": {"pair_id", "tag_id"},
"image_prompt_pairs": {
"id",
"image_id",
"pair_id",
"provenance",
"created_at",
},
}
present = {
row[0]
for row in conn.execute("SELECT name FROM sqlite_master WHERE type='table'")
}
if not present.intersection(required):
return None
if not set(required).issubset(present):
return "Incomplete prompt pair schema"
foreign_keys = {
"prompt_pairs": {
("positive_prompt_id", "prompts", "id", "RESTRICT"),
("negative_prompt_id", "prompts", "id", "RESTRICT"),
},
"prompt_pair_tags": {
("pair_id", "prompt_pairs", "id", "CASCADE"),
("tag_id", "tags", "id", "CASCADE"),
},
"image_prompt_pairs": {
("pair_id", "prompt_pairs", "id", "RESTRICT"),
("image_id", "generated_images", "id", "CASCADE"),
},
}
unique = {
"prompt_pairs": ("positive_prompt_id", "negative_prompt_id"),
"prompt_pair_tags": ("pair_id", "tag_id"),
"image_prompt_pairs": ("pair_id", "image_id", "provenance"),
}
non_null = {
"prompt_pairs": {
"positive_prompt_id",
"created_at",
"updated_at",
"run_count",
"archived",
},
"prompt_pair_tags": {"pair_id", "tag_id"},
"image_prompt_pairs": {"pair_id", "image_id", "provenance", "created_at"},
}
integer_columns = {
"id",
"positive_prompt_id",
"negative_prompt_id",
"pair_id",
"tag_id",
"image_id",
"rating",
"run_count",
"archived",
}
for table, columns in required.items():
info = {row[1]: row for row in conn.execute(f"PRAGMA table_info({table})")}
if not columns.issubset(info):
return f"{table} is missing required pair columns"
if any(not info[name][3] for name in non_null[table]):
return f"{table} has nullable required pair columns"
if any(
info[name][2].upper()
!= ("INTEGER" if name in integer_columns else "TEXT")
for name in columns
):
return f"{table} has invalid pair column types"
primary = tuple(
name
for name in sorted(info, key=lambda name: info[name][5])
if info[name][5]
)
if primary != (
("pair_id", "tag_id") if table == "prompt_pair_tags" else ("id",)
):
return f"{table} has an invalid pair primary key"
actual_fks = {
(row[3], row[2], row[4], row[6])
for row in conn.execute(f"PRAGMA foreign_key_list({table})")
}
if actual_fks != foreign_keys[table]:
return f"{table} has invalid pair foreign key definitions"
if conn.execute(f"PRAGMA foreign_key_check({table})").fetchone():
return f"{table} contains foreign key violations"
indexes = list(conn.execute(f"PRAGMA index_list({table})"))
index_columns = {}
for row in indexes:
name = row[1].replace('"', '""')
index_columns[row[1]] = tuple(
item[2] for item in conn.execute(f'PRAGMA index_info("{name}")')
)
if not any(
row[2] and not row[4] and index_columns[row[1]] == unique[table]
for row in indexes
):
return f"{table} is missing pair uniqueness constraints"
if table == "prompt_pairs":
null_unique = False
for row in indexes:
if (
row[2]
and row[4]
and index_columns[row[1]] == ("positive_prompt_id",)
):
sql = conn.execute(
"SELECT sql FROM sqlite_master WHERE type='index' AND name"
"=?",
(row[1],),
).fetchone()[0]
if " ".join(sql.lower().split()).endswith(
"where negative_prompt_id is null"
):
null_unique = True
if not null_unique:
return "prompt_pairs is missing empty-negative pair uniqueness"
if conn.execute(
"SELECT 1 FROM image_prompt_pairs ip "
"JOIN prompt_pairs pp ON pp.id=ip.pair_id "
"JOIN generated_images gi ON gi.id=ip.image_id "
"WHERE pp.positive_prompt_id IS NOT gi.prompt_id LIMIT 1"
).fetchone():
return "Image pair association has a different positive prompt owner"
# fetchone, not iteration: a managed connection's cursor prefetches
# ITER_CHUNK rows when iterated, and each provenance may be ~4 MB.
cursor = conn.execute("SELECT provenance FROM image_prompt_pairs")
while (row := cursor.fetchone()) is not None:
error = pair_provenance_error(row[0])
if error:
return error
return None
def _create_indexes(self, conn: sqlite3.Connection) -> None:
"""
Create indexes for better query performance.
Args:
conn: Active database connection
Creates indexes on:
- Text content for search operations
- Categories and tags for filtering
- Timestamps for sorting
- Hash values for duplicate detection
"""
indexes = [
("idx_prompts_text", "prompts(text)"),
("idx_prompts_category", "prompts(category)"),
("idx_prompts_created_at", "prompts(created_at)"),
("idx_prompts_hash", "prompts(hash)"),
("idx_prompts_rating", "prompts(rating)"),
("idx_prompts_last_used", "prompts(last_used_at)"),
("idx_prompts_run_count", "prompts(run_count)"),
("idx_prompt_images", "generated_images(prompt_id)"),
("idx_image_path", "generated_images(image_path)"),
("idx_image_file_path", "generated_images(file_path)"),
("idx_generation_time", "generated_images(generation_time)"),
("idx_prompt_tags_tag", "prompt_tags(tag_id)"),
("idx_tags_name", "tags(name)"),
]
for name, target in indexes:
conn.execute(f"CREATE INDEX IF NOT EXISTS {name} ON {target}")
def get_connection(self) -> sqlite3.Connection:
"""
Get the calling thread's database connection.
Each thread gets its own ``sqlite3.Connection`` (kept in
``threading.local``), created on first use and reused by that thread
thereafter. Every connection gets ``row_factory = sqlite3.Row`` and the
shared settings of ``configure_connection``: WAL journaling with
``synchronous=NORMAL``, ``PRAGMA foreign_keys = ON`` and a
``busy_timeout`` of BUSY_TIMEOUT_SECONDS so concurrent writers wait
instead of failing. Callers use it as a context manager
(``with conn:``), which commits or rolls back that thread's transaction
and leaves the connection open. No lock is held while a connection is in use.
Returns:
sqlite3.Connection: The current thread's configured connection
"""
conn = getattr(self._local, "conn", None)
if conn is not None and self._local.generation == self._generation:
return conn
while True:
# The path and generation are read together, and the new
# connection is kept only if no Move library repointed the model
# meanwhile: otherwise it would write to the old file.
with self._conn_lock:
path, generation = self.db_path, self._generation
# check_same_thread=False only so close_all and Move library can
# retire other threads' connections; each connection is still
# used by one thread, and ManagedConnection never closes it under
# a running statement (see managed_connection).
conn = sqlite3.connect(
path,
check_same_thread=False,
timeout=BUSY_TIMEOUT_SECONDS,
factory=ManagedConnection,
)
try:
conn.row_factory = sqlite3.Row
configure_connection(conn)
except BaseException:
conn.close()
raise
with self._conn_lock:
if generation == self._generation:
self._connections.add(conn)
break
conn.close()
self._local.conn = conn
self._local.generation = generation
return conn
@property
def generation(self) -> int:
"""Bumped whenever every connection is closed or the model repointed."""
return self._generation
@classmethod
def close_all_instances(cls) -> None:
"""Close every connection of every live model in this process."""
with cls._instances_lock:
models = list(cls._instances)
for model in models:
model.close_all()
@classmethod
def repoint_instances(cls, old_path: str, new_path: str) -> list:
"""Point every model of ``old_path`` at ``new_path`` (Move library).
Every thread opens the new file on its next ``get_connection``. The
old connections are detached, not closed, and returned: the caller
retires them (``retire_detached``) while it still holds the old
file's write lock. A write waiting on that lock is then interrupted
and can never commit, so it fails (and, through
``PromptDatabase.survive_library_switch``, runs again on the new file).
"""
old_key = os.path.normcase(os.path.realpath(old_path))
with cls._instances_lock:
models = list(cls._instances)
with cls._init_lock:
cls._initialized_paths.add(os.path.normcase(os.path.realpath(new_path)))
detached = []
for model in models:
if model._schema_key() == old_key:
detached.extend(model._detach_connections(new_path))
return detached
@staticmethod
def retire_detached(connections, logger=None) -> list:
"""Retire detached connections without waiting for their owners.
Idle ones close now. One whose owner thread is inside a statement is
never closed from here (that frees memory the owner is using: an
access violation on Windows); it is interrupted, refuses to commit and
is closed by its owner as the statement ends. Returns those still
open, for ``await_detached`` once no lock the owners wait on is held.
"""
return retire_connections(connections, logger)
@staticmethod
def await_detached(connections, logger=None) -> bool:
"""Wait until retired connections are closed; False if some are not."""
return await_closed(connections, BUSY_TIMEOUT_SECONDS, logger)
@staticmethod
def close_detached(connections, logger=None) -> bool:
"""Retire connections ``repoint_instances`` detached and await them.
Do not call it while holding a lock their owners may wait on; retire
them there (``retire_detached``) and await them after releasing it.
"""
pending = retire_connections(connections, logger)
return await_closed(pending, BUSY_TIMEOUT_SECONDS, logger)
def close_all(self) -> None:
"""Close every thread's connection; they reopen lazily on next use.
Needed before deleting the database file (Windows keeps open files
locked) and by tests that ran queries through executor threads.
"""
self.close_detached(self._detach_connections(), self.logger)
def _detach_connections(self, new_path: Optional[str] = None) -> list:
"""Forget every connection (optionally pointing the model elsewhere).
Returns them for the caller to close; threads reopen on next use.
"""
with self._conn_lock:
if new_path is not None:
self.db_path = new_path
connections = list(self._connections)
self._connections.clear()
self._generation += 1
self._local.conn = None
return connections
def close(self) -> None:
"""Close the calling thread's connection, if it has one.
Other threads keep their connections. Call this before deleting or
replacing the database file: Windows refuses to delete an open file.
"""
conn = getattr(self._local, "conn", None)
if conn is None:
return
self._local.conn = None
with self._conn_lock:
self._connections.discard(conn)
self.close_detached([conn], self.logger)
def _migrate_workflow_name_removal(self, conn: sqlite3.Connection) -> None:
"""
Remove workflow_name column if it exists in existing database.
Args:
conn: Active database connection
This migration handles legacy schema updates by removing the deprecated
workflow_name column while preserving all other data.
Raises:
sqlite3.Error: If the rebuild fails (logged, nothing half-applied)
"""
try:
cursor = conn.execute("PRAGMA table_info(prompts)")
columns = [column[1] for column in cursor.fetchall()]
if "workflow_name" not in columns:
return
self.logger.info("Migrating database: removing workflow_name column")
with self._table_rebuild(conn):
# A previous attempt may have died between CREATE and RENAME
conn.execute("DROP TABLE IF EXISTS prompts_new")
conn.execute("""
CREATE TABLE prompts_new (
id INTEGER PRIMARY KEY AUTOINCREMENT,
text TEXT NOT NULL,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
category TEXT,
tags TEXT,
rating INTEGER CHECK(rating >= 1 AND rating <= 5),
notes TEXT,
hash TEXT UNIQUE
)
""")
conn.execute("""
INSERT INTO prompts_new (id, text, created_at, updated_at,
category, tags, rating, notes, hash)
SELECT id, text, created_at, updated_at, category, tags, rating,
notes, hash
FROM prompts
""")
conn.execute("DROP TABLE prompts")
conn.execute("ALTER TABLE prompts_new RENAME TO prompts")
self.logger.info("Database migration completed")
except sqlite3.Error as e:
self.logger.error(f"Migration error (workflow_name removal): {e}")
raise
def _migrate_foreign_key_types(self, conn: sqlite3.Connection) -> None:
"""
Fix foreign key data type mismatch in generated_images table.
Args:
conn: Active database connection
Converts prompt_id from TEXT to INTEGER type to match the prompts table's
primary key type, ensuring referential integrity.
Raises:
sqlite3.Error: If the rebuild fails (logged, nothing half-applied)
"""
try:
cursor = conn.execute("PRAGMA table_info(generated_images)")
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):
conn.execute("DROP TABLE IF EXISTS generated_images_new")
conn.execute("""
CREATE TABLE generated_images_new (
id INTEGER PRIMARY KEY AUTOINCREMENT,
prompt_id INTEGER NOT NULL,
image_path TEXT NOT NULL,
filename TEXT NOT NULL,
generation_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
file_size INTEGER,
width INTEGER,
height INTEGER,
format TEXT,
workflow_data TEXT,
prompt_metadata TEXT,
parameters TEXT,
FOREIGN KEY (prompt_id) REFERENCES prompts(id) ON DELETE CASCADE
)
""")
conn.execute("""
INSERT INTO generated_images_new
(id, prompt_id, image_path, filename, generation_time, file_size,
width, height, format, workflow_data, prompt_metadata, parameters)
SELECT id, CAST(prompt_id AS INTEGER), image_path, filename,
generation_time,
file_size, width, height, format, workflow_data,
prompt_metadata, parameters
FROM generated_images
WHERE prompt_id != '' AND prompt_id IS NOT NULL
AND CAST(prompt_id AS INTEGER) IN (SELECT id FROM prompts)
""")
conn.execute("DROP TABLE generated_images")
conn.execute(
"ALTER TABLE generated_images_new RENAME TO generated_images"
)
self.logger.info("Foreign key migration completed")
except sqlite3.Error as e:
self.logger.error(f"Migration error (foreign key types): {e}")
raise
def _migrate_json_tags_to_junction(self, conn: sqlite3.Connection) -> None:
"""
Populate tags and prompt_tags tables from legacy JSON tags column.
Runs once: skips if the tags table already has data. Uses json_each()
to extract tag names from the JSON arrays stored in prompts.tags.
Raises:
sqlite3.Error: If the migration fails (logged)
"""
try:
cursor = conn.execute("SELECT COUNT(*) FROM tags")
if cursor.fetchone()[0] > 0:
return # Already migrated
cursor = conn.execute(
"SELECT COUNT(*) FROM prompts "
"WHERE tags IS NOT NULL AND tags != '' AND tags != '[]'"
)
if cursor.fetchone()[0] == 0:
return # No tags to migrate
self.logger.info("Migrating JSON tags to normalized junction tables")
# Insert all unique tag names
conn.execute(
"INSERT OR IGNORE INTO tags (name) "
"SELECT DISTINCT je.value FROM prompts, json_each(prompts.tags) AS je "
"WHERE prompts.tags IS NOT NULL AND prompts.tags != '' AND "
"prompts.tags != '[]'"
)
# Populate junction table
conn.execute(
"INSERT OR IGNORE INTO prompt_tags (prompt_id, tag_id) "
"SELECT p.id, t.id "
"FROM prompts p, json_each(p.tags) AS je "
"JOIN tags t ON t.name = je.value "
"WHERE p.tags IS NOT NULL AND p.tags != '' AND p.tags != '[]'"
)
tag_count = conn.execute("SELECT COUNT(*) FROM tags").fetchone()[0]
link_count = conn.execute("SELECT COUNT(*) FROM prompt_tags").fetchone()[0]
self.logger.info(
f"Tag migration complete: {tag_count} unique tags, {link_count} "
"prompt-tag links"
)
except sqlite3.Error as e:
self.logger.error(f"Migration error (tag junction): {e}")
raise
def _migrate_add_usage_columns(self, conn: sqlite3.Connection) -> None:
"""
Add last_used_at and run_count (3.2.4) and estimate usage for existing rows.
The backfill runs only when the columns are added, i.e. on the upgrade
from a version without usage tracking. Afterwards a NULL last_used_at
means "never run": saving a prompt does not stamp it, only
record_prompt_use does, and "Recently Used" relies on that.
The estimate comes from linked images: run_count is the image count (at
least 1) and last_used_at the newest image time, falling back to created_at.
Timestamps are compared with julianday() because prompts use ISO 'T'
timestamps while generated_images uses SQLite's CURRENT_TIMESTAMP format.
"""
try:
columns = {row[1] for row in conn.execute("PRAGMA table_info(prompts)")}
added = False
if "last_used_at" not in columns:
self.logger.info("Adding prompt usage column last_used_at")
conn.execute("ALTER TABLE prompts ADD COLUMN last_used_at TIMESTAMP")
added = True
if "run_count" not in columns:
self.logger.info("Adding prompt usage column run_count")
conn.execute(
"ALTER TABLE prompts ADD COLUMN run_count INTEGER NOT NULL "
"DEFAULT 0"
)
added = True
if not added:
return
cursor = conn.execute(f"""
UPDATE prompts SET
run_count = MAX(run_count, 1, (
SELECT COUNT(*) FROM generated_images gi
WHERE gi.prompt_id = prompts.id
)),
last_used_at = strftime('{ISO_TIMESTAMP_FORMAT}', MAX(
COALESCE(julianday(created_at), julianday('now')),
COALESCE((
SELECT MAX(julianday(gi.generation_time))
FROM generated_images gi WHERE gi.prompt_id = prompts.id
), 0)
))
WHERE last_used_at IS NULL
""")
if cursor.rowcount:
self.logger.info(f"Estimated usage for {cursor.rowcount} prompts")
except sqlite3.Error as e:
self.logger.error(f"Migration error (usage columns): {e}")
raise
def _migrate_image_uniqueness_by_path(self, conn: sqlite3.Connection) -> None:
"""
Rebuild generated_images so an image is unique per (prompt_id, file_path).
Older schemas were unique per (prompt_id, filename), i.e. the basename,
so ComfyUI's same-named files in different date folders were silently
dropped. file_path holds normalize_image_path(image_path); rows whose
paths normalise to the same file collapse to the newest one.
Raises:
sqlite3.Error: If the rebuild fails (logged, nothing half-applied)
"""
try:
columns = {
row[1] for row in conn.execute("PRAGMA table_info(generated_images)")
}
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):
rows = conn.execute(
"SELECT id, prompt_id, image_path FROM generated_images ORDER BY id"
).fetchall()
keep = {}
for row_id, prompt_id, image_path in rows:
keep[(prompt_id, normalize_image_path(image_path))] = row_id
conn.execute("DROP TABLE IF EXISTS generated_images_new")
conn.execute("""
CREATE TABLE generated_images_new (
id INTEGER PRIMARY KEY AUTOINCREMENT,
prompt_id INTEGER NOT NULL,
image_path TEXT NOT NULL,
filename TEXT NOT NULL,
generation_time TIMESTAMP DEFAULT CURRENT_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.executemany(
"""
INSERT INTO generated_images_new
(id, prompt_id, image_path, filename, generation_time, file_size,
width, height, format, workflow_data, prompt_metadata, parameters,
file_path)
SELECT id, prompt_id, image_path, filename, generation_time,
file_size,
width, height, format, workflow_data, prompt_metadata,
parameters,
?
FROM generated_images WHERE id = ?
""",
[(file_path, row_id) for (_, file_path), row_id in keep.items()],
)
conn.execute("DROP TABLE generated_images")
conn.execute(
"ALTER TABLE generated_images_new RENAME TO generated_images"
)
dropped = len(rows) - len(keep)
if dropped:
self.logger.info(f"Collapsed {dropped} duplicate image rows")
self.logger.info("Image path uniqueness migration completed")
except sqlite3.Error as e:
self.logger.error(f"Migration error (image path uniqueness): {e}")
raise
def vacuum_database(self) -> None:
"""
Optimize database by running VACUUM command.
Reclaims unused space and defragments the database file,
improving query performance and reducing file size.
Raises:
Exception: whatever stopped VACUUM (logged first); the database is
then unchanged.
"""
try:
with closing(sqlite3.connect(self.db_path)) as conn:
conn.execute("VACUUM")
conn.commit()
except Exception as e:
self.logger.error(f"Error vacuuming database: {e}")
raise
def get_database_info(self) -> dict:
"""
Get information about the database.
Returns:
dict: Database statistics and information
"""
try:
with self.get_connection() as conn:
cursor = conn.execute("SELECT COUNT(*) as total_prompts FROM prompts")
total_prompts = cursor.fetchone()["total_prompts"]
cursor = conn.execute(
"SELECT COUNT(DISTINCT category) as unique_categories FROM "
"prompts WHERE category IS NOT NULL"
)
unique_categories = cursor.fetchone()["unique_categories"]
cursor = conn.execute(
"SELECT AVG(rating) as avg_rating FROM prompts WHERE rating IS "
"NOT NULL"
)
avg_rating = cursor.fetchone()["avg_rating"]
# Get database file size
db_size = (
os.path.getsize(self.db_path) if os.path.exists(self.db_path) else 0
)
return {
"total_prompts": total_prompts,
"unique_categories": unique_categories,
"average_rating": round(avg_rating, 2) if avg_rating else None,
"database_size_bytes": db_size,
"database_path": os.path.abspath(self.db_path),
}
except Exception as e:
self.logger.error(f"Error getting database info: {e}")
return {}
def backup_database(self, backup_path: str) -> bool:
"""
Write a consistent copy of the database to backup_path.
Uses the SQLite online backup API through this thread's connection,
so rows that are committed but still sit in the -wal file are
included and other threads may keep reading and writing meanwhile.
Args:
backup_path: Path where the backup should be saved (overwritten)
Returns:
bool: True if backup was successful, False otherwise
"""
try:
if os.path.exists(backup_path) and os.path.samefile(
backup_path, self.db_path
):
self.logger.error("Refusing to back up the database onto itself")
return False
if os.path.exists(backup_path):
os.remove(backup_path)
source = self.get_connection()
destination = sqlite3.connect(backup_path)
try:
source.backup(destination)
finally:
destination.close()
return True
except (OSError, sqlite3.Error) as e:
self.logger.error(f"Error creating database backup: {e}")
return False
@staticmethod
def verify_database_file(
path: str, *, require_library_state: bool = False, documents: bool = True
) -> Tuple[bool, str]:
"""
Check that a file is a healthy prompts database before restoring it.
Opens the file read-only, runs PRAGMA integrity_check and confirms the
prompts table has the columns every version of this extension needs.
Args:
path: File to inspect
require_library_state: Also require the tables every current
database has (Library state, image favorites, saved views,
image annotations and the review queue), as a migrated
restore stage must; old backups may lack them.
documents: Read every workflow document and refuse a damaged
one. A restore checks its upload without it and repairs
damaged documents on the stage (``database/restore_stage.py``).
Returns:
(ok, reason): reason is "ok" when the file is acceptable, otherwise
a human-readable explanation. Never raises for bad input.
"""
if not os.path.isfile(path):
return False, "File not found"
try:
with open(path, "rb") as fh:
header = fh.read(len(SQLITE_HEADER))
except OSError as e:
return False, f"Cannot read file: {e}"
if header != SQLITE_HEADER:
return False, "Not a SQLite database file"
uri = Path(os.path.abspath(path)).as_uri() + "?mode=ro"
try:
conn = sqlite3.connect(uri, uri=True)
except sqlite3.Error as e:
return False, f"Invalid SQLite database: {e}"
try:
row = conn.execute("PRAGMA integrity_check").fetchone()
verdict = row[0] if row else "no result"
if verdict != "ok":
return False, f"Integrity check failed: {verdict}"
table = conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' AND name='prompts'"
).fetchone()
if not table:
return False, "Database does not contain a 'prompts' table"
columns = {r[1] for r in conn.execute("PRAGMA table_info(prompts)")}
missing = [c for c in REQUIRED_PROMPT_COLUMNS if c not in columns]
if missing:
return (
False,
"prompts table is missing required columns: " + ", ".join(missing),
)
forbidden = conn.execute(
"SELECT type FROM sqlite_master WHERE type IN ('trigger', 'view') "
"LIMIT 1"
).fetchone()
if forbidden:
return False, f"Database contains an unsupported {forbidden[0]}"
pair_error = PromptModel._pair_schema_error(conn)
if pair_error:
return False, pair_error
history_error = history_schema_error(conn)
if history_error:
return False, history_error
library_error = library_schema_error(
conn, allow_missing=not require_library_state
)
if library_error:
return False, library_error
favorites_error = image_favorites_schema_error(
conn, allow_missing=not require_library_state
)
if favorites_error:
return False, favorites_error
views_error = saved_views_schema_error(
conn, allow_missing=not require_library_state
)
if views_error:
return False, views_error
annotation_error = annotation_schema_error(
conn, allow_missing=not require_library_state
)
if annotation_error:
return False, annotation_error
review_error = review_schema_error(
conn, allow_missing=not require_library_state
)
if review_error:
return False, review_error
documents_error = workflow_documents_schema_error(
conn, allow_missing=not require_library_state, documents=documents
)
if documents_error:
return False, documents_error
return True, "ok"
except sqlite3.Error as e:
return False, f"Invalid SQLite database: {e}"
finally:
conn.close()
def _replace_contents_from(self, src_path: str) -> None:
"""Copy every page of *src_path* into the live database.
Goes through SQLite's online backup API into this thread's normal
connection, so SQLite takes the locks itself and the other
PromptModel instances that share the file (nodes, image monitor,
queue hook, API) simply see the new pages on their next statement.
Nothing is swapped underneath open handles and no -wal/-shm file has
to be deleted, which also keeps Windows happy. If SQLite refuses
replacement (for example, incompatible page sizes in WAL mode),
propagate the error rather than copy files under other open owners.
"""
source = sqlite3.connect(
Path(src_path).resolve().as_uri() + "?mode=ro", uri=True
)
try:
target = self.get_connection()
# The target is written under backup's step too: mark it in use
# so a concurrent close_all retires it instead of closing it.
with target.using():
source.backup(target)
finally:
source.close()
def stage_restore(self, src_path: str):
"""Stage ``src_path`` beside this database for a restore (blocking).
Returns a ``StagedRestore`` (``database/restore_stage.py``): a
migrated, repaired and verified copy in the library's folder, its
digest and what it holds. Neither the live database nor the source
is touched. Raises ValueError when either does not verify.
"""
return restore_stage.stage_restore(PromptModel, self.db_path, src_path)
def publish_staged(self, staged) -> str:
"""Replace the live database with ``staged``; return the safety backup.
See ``restore_stage.publish_staged``: the staged bytes are rechecked,
the current database is kept as ``<db_path>.backup_<stamp>`` and the
WAL is truncated afterwards.
"""
backup_path = restore_stage.publish_staged(self, staged)
self.logger.info("Database restored from a staged copy")
from .execution_history import notify_history_restore
notify_history_restore(self.db_path)
return backup_path
def restore_from_file(self, src_path: str) -> str:
"""
Replace the live database with the file at src_path.
``stage_restore`` then ``publish_staged``: a disposable copy of the
source is migrated, repaired and verified in the library's folder
first. Only then is the current database backed up next to itself as
``<db_path>.backup_<YYYYmmdd_HHMMSS_ffffff>`` (WAL content included),
and the staged pages replace the live pages through SQLite's backup
API. Staging failures never replace live content or modify the
source. An online replacement failure leaves the live database intact
and propagates; the completed safety backup remains available. The
stage is removed either way.
Args:
src_path: SQLite file to restore
Returns:
Path of the backup taken of the previous database, or "" when
there was no database to back up.
Raises:
ValueError: If the source or its migrated copy does not verify
SafetyBackupFailed: If the previous database could not be backed up
sqlite3.OperationalError: If SQLite refuses online replacement
"""
staged = self.stage_restore(src_path)
try:
backup_path = self.publish_staged(staged)
finally:
staged.discard()
self.logger.info(f"Database restored from {src_path}")
return backup_path
Reference in New Issue View Git Blame Copy Permalink