Vito
b7cf210f0f
* feat(annotations): add image annotation schema with restore validation
Install the image annotation group beside image_favorites and saved_views:
media_annotation_metadata (one row, schema version 1), media_annotation_files
(one private locator per physical file), media_annotation_bytes (one record
per observed SHA-256 with the verified size and mtime_ns) and media_image_tags
(the ordered tag set of a byte record, joined into the global tags registry).
The group migrates in a SAVEPOINT that validates before and after creating
it, with no backfill from prompt or pair tags and no filesystem probe.
annotation_schema_error reads only: it compares the catalog objects after
whitespace compaction, the columns, foreign keys and indexes, and the
persisted values (UUIDs, hex digests, bounds, the 4096-character locator cap
measured past embedded NULs, exactly one metadata row). verify_database_file
checks it with allow_missing=not require_library_state, so old backups
restore and a migrated restore stage must carry it.
Tests pin the exact DDL, a fresh install and idempotent reopen, a current
database without the group gaining it with legacy rows untouched, every
partial or malformed shape rejected without repair, install failure rolling
back and retrying, a WAL backup/restore round trip, staged-failure
immutability of source, live and safety files with every handle closed, and
the require_library_state gate.
* feat(annotations): persist per-image tag sets in the shared vocabulary
Add database/annotations.py (AnnotationOperations, mixed into
PromptDatabase): read_image_tags, replace_image_tags and add_image_tags
run in one _pair_transaction fenced by the storage epoch, own the byte
record by exact locator and digest, always supply UUID ids, rewrite the
ordered junction and bump tag_revision plus image_index_revision.
Revision mismatches raise AnnotationCommandError("tag-conflict") with
the current set; the identical set at the current or previous revision
is replayed without writes; a union over 100 names is "tag-limit".
normalize_image_tags validates each name with validate_vocabulary_tag_name
and folds case-insensitive repeats keeping the first spelling.
refresh_image_fingerprint, current_image_tags (chunked, fingerprint
matched, newest verification wins) and iter_image_tag_records serve the
S1 cheap reads.
Vocabulary participation: _move_tag_links also carries media_image_tags,
legacy merge_tags moves image-only sources, tag rows gain image_count
(files with any byte record carrying the tag), impact gains images and
rename/delete receipts gain images_changed. vocabulary-state.mjs accepts
the optional counts and names images in the delete text.
Browser fixtures clear the three annotation data tables on reset (never
the metadata row); TESTING.md documents the schema beside favorites.
* feat(annotations): add bounded private media capture store
Add utils/media_capture.py: CaptureStore copies one rooted media file
into an owner-only prompt-manager-captures-* temp directory, streaming
it in 1 MiB chunks while computing the whole-file SHA-256, and compares
the opened descriptor with the path (dev, ino where meaningful, size,
mtime_ns, ctime_ns) before and after the copy; any disagreement raises
CaptureChanged and removes the copy. Symlinked files or directories,
root escapes, non-media names and directories raise CaptureDenied; the
locator is normcase(join(realpath(root), relative)), the listing's
identity, so a symlinked root matches it.
Limits are the plan's: 256 MiB per file on stat and stream
(CaptureTooLarge), 8 completed captures, 512 MiB reserved and 2
capture/verify workers with nonblocking admission (CaptureCapacity),
900 s idle and 1800 s hard lifetimes on an injected clock. acquire and
release_use keep an in-use capture past expiry, release and
invalidate_all until its last use ends; close removes only owned files
and the owned directory. A header-only probe of the copy marks Generate
unsupported above 4096 px per side or 16,777,216 px, for unreadable
headers and for video; verify re-hashes the current file without a
copy. Every handle closes before unlink.
tests/test_media_capture.py covers exact bytes and owner-only modes,
alias and escape denial, the symlinked-root locator, change detection
during the copy (size, same-size rewrite with restored mtime through
the POSIX change time, inode replacement), both size caps, count, byte
and worker admission, idle and hard expiry with in-use retention,
cleanup scope, handle release at unlink, dimension and video
readiness, and the allowlist equality with the inspection service.
TESTING.md documents the store and the Windows change-time limit.
* test(annotations): wait on the POSIX change time only where it exists
The same-size rewrite case waits for the filesystem clock to pass the
source's ctime so the detection is deterministic on Linux's coarse
timestamps; Windows reports the creation time as st_ctime, where the
probe could never advance, so the wait runs only on POSIX.
* fix(annotations): compare the change time only between fstat reads
Python 3.12+ on Windows reports the change time from stat but the
creation time from fstat, so a freshly written file whose creation and
write straddle a clock tick failed the pre-read identity check with
CaptureChanged. The path-versus-descriptor checks now compare size,
mtime_ns and dev/ino; ctime_ns is compared between the fstat before
and after the copy, where it still detects a same-size rewrite with a
restored mtime on POSIX.
The symlinked-root test asserted the locator against the unresolved
temp path, which macOS (/var -> /private/var) and Windows (short names)
resolve differently; it now compares against the real path.
* test(browser): give the inspector evidence matrix the concept comparisons' budget
The touch run of 'full inspector attaches matched theme and mobile
evidence with reachable content' captures 32 screenshots across four
viewports and took 23.5 s, then 31.6 s, in the sharded coverage job,
hitting the 30 s default. Like the concept comparison specs it now sets
an explicit 120 s timeout; the CI-like pinned run (4 CPUs, 3 workers,
repeat-each 3) passes 198/198 with the matrix at 18-20 s on touch.
* fix(annotations): compare capture identity like with like and keep paths out of errors
Native Windows CI (Python 3.12/3.13) rejected untouched files with
CaptureChanged because the path stat and the descriptor fstat disagree
there (change time versus creation time, possibly dev/ino encodings).
The descriptor is now checked against the path only on the API-stable
facts (size, mtime_ns); the full identity (size, mtime_ns, ctime_ns,
dev/ino where meaningful) is compared between the fstat before and
after the copy and between the path stats before and after. A test
injects shifted fstat results to simulate the Windows discrepancy, and
a fresh file is captured twice with no platform guard.
Error messages carry strerror only, never the path (the path goes to
the debug log); a test asserts no CaptureError names the root. Files
are unlinked outside the store lock; close() during an in-flight
capture lets that worker remove its temp file and then the directory.
get() documents that readers of capture.path must acquire/release_use.
The test fixture resolves its temp root (macOS /private/var, Windows
8.3 names). TESTING.md states the real identity contract.
* fix(annotations): move image revisions when a fingerprint or the vocabulary changes
A replayed save that observes a new verified (size, mtime_ns) now bumps
image_index_revision, like refresh_image_fingerprint, so caches keyed
on it see the record become current again; an identical replay still
moves nothing.
Vocabulary rename, merge and delete (native and legacy paths) bump the
tag_revision of every image set carrying the source tag, in the same
transaction and before the links move or cascade, plus one
image_index_revision bump when any set was affected. A Tag dialog
opened before the change therefore gets tag-conflict with the renamed
set as current instead of recreating the old name.
* test(annotations): accept change detection wherever fstat exposes a change time
Native Windows CI on Python 3.13 detected the same-size, mtime-restored
rewrite through the descriptor's change time (3.10-3.12 report the
creation time there and complete the torn copy). The test asserted the
platform, not the signal: it now accepts CaptureChanged plus temp
removal wherever it happens, requires detection on POSIX and on
Windows 3.13+, and otherwise requires the digest to describe the
copied bytes. A new test injects a drifting fstat change time so the
detection path runs on every platform; TESTING.md records the limit as
"Windows before Python 3.13".
* test(annotations): assert the rewrite outcome, not the platform
Native Windows CI detected the same-size, mtime-restored rewrite on
3.12 and missed it on 3.13 in the same run: the path stat's change time
(Python 3.12+ fast path) only moves when the clock ticks between the
pre-stat and the rewrite, and the tick wait was skipped on Windows. The
wait now runs everywhere with fresh probe names (Windows tunnels the
creation time of a reused name), and the test asserts the outcome: a
detected change leaves no temp file; an undetected one is accepted on
Windows only and must produce a digest of the copied bytes; POSIX must
detect it. The injected-fstat test keeps the detection path covered on
every platform.
* fix(images): keep the Filters label node across a listing reload
Every listing reload reports the active refinement count, and
setActiveFilters rewrote the button's textContent even when the label
was unchanged, replacing its text node. WebKit drops a click whose
mousedown hit a node that was detached before mouseup, so a Safari
user who types a search and clicks Filters inside the 250 ms debounce
lost the click, and the webkit page workflows timed out in the
source-change recovery tests (1 in ~30 runs of the pinned image; an
instrumented probe showed pointerdown/mousedown, a DOM mutation,
pointerup/mouseup and no click event).
The label is rewritten only when it changes. The toolbar unit test
counts textContent writes (1 for an unchanged count before, 0 after),
and a workspace-images test pins the text node identity across a
search reload: 0/6 before the fix on chromium and webkit, 6/6 after,
with the recovery and Refine tests 36/36 in the pinned CI-like run.
1196 lines
49 KiB
Python
1196 lines
49 KiB
Python
"""
|
|
Database schema and models for KikoTextEncode prompt storage.
|
|
"""
|
|
|
|
import contextlib
|
|
import datetime
|
|
import json
|
|
import math
|
|
import sqlite3
|
|
import os
|
|
import threading
|
|
import tempfile
|
|
import weakref
|
|
from contextlib import closing
|
|
from pathlib import Path
|
|
from typing import Optional, Set, Tuple
|
|
|
|
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 .saved_views import migrate_saved_views, saved_views_schema_error
|
|
from .execution_history_schema import (
|
|
history_schema_error,
|
|
migrate_execution_history,
|
|
rotate_history_epoch,
|
|
)
|
|
|
|
# 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 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)
|
|
migrate_annotation_state(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"
|
|
for row in conn.execute("SELECT provenance FROM image_prompt_pairs"):
|
|
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
|
|
|
|
# check_same_thread=False only so close_all can clean up other
|
|
# threads' connections; each connection is still used by one thread.
|
|
conn = sqlite3.connect(
|
|
self.db_path, check_same_thread=False, timeout=BUSY_TIMEOUT_SECONDS
|
|
)
|
|
conn.row_factory = sqlite3.Row
|
|
configure_connection(conn)
|
|
self._local.conn = conn
|
|
self._local.generation = self._generation
|
|
with self._conn_lock:
|
|
self._connections.add(conn)
|
|
return conn
|
|
|
|
@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()
|
|
|
|
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.
|
|
"""
|
|
with self._conn_lock:
|
|
connections = list(self._connections)
|
|
self._connections.clear()
|
|
self._generation += 1
|
|
self._local.conn = None
|
|
for conn in connections:
|
|
try:
|
|
conn.close()
|
|
except sqlite3.Error as e:
|
|
self.logger.warning(f"Error closing database connection: {e}")
|
|
|
|
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)
|
|
try:
|
|
conn.close()
|
|
except sqlite3.Error as e:
|
|
self.logger.warning(f"Error closing database connection: {e}")
|
|
|
|
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
|
|
|
|
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
|
|
|
|
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.
|
|
"""
|
|
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}")
|
|
|
|
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
|
|
) -> 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
|
|
and image annotations), as a migrated restore stage must;
|
|
old backups may lack them.
|
|
|
|
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
|
|
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:
|
|
source.backup(self.get_connection())
|
|
finally:
|
|
source.close()
|
|
|
|
def restore_from_file(self, src_path: str) -> str:
|
|
"""
|
|
Replace the live database with the file at src_path.
|
|
|
|
A disposable SQLite backup of the source is migrated and verified first.
|
|
Only then is the current database backed up next to itself as
|
|
``<db_path>.backup_<YYYYmmdd_HHMMSS>`` (WAL content included), and the
|
|
staged pages replace the live pages through SQLite's backup API.
|
|
Migration 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.
|
|
|
|
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:
|
|
RuntimeError: If the previous database could not be backed up
|
|
sqlite3.OperationalError: If SQLite refuses online replacement
|
|
"""
|
|
ok, reason = self.verify_database_file(src_path)
|
|
if not ok:
|
|
raise ValueError(reason)
|
|
with tempfile.TemporaryDirectory(prefix="prompt-manager-restore-") as directory:
|
|
staged_path = str(Path(directory) / "staged.db")
|
|
source = sqlite3.connect(
|
|
Path(src_path).resolve().as_uri() + "?mode=ro", uri=True
|
|
)
|
|
try:
|
|
with closing(sqlite3.connect(staged_path)) as destination:
|
|
source.backup(destination)
|
|
finally:
|
|
source.close()
|
|
staged = None
|
|
try:
|
|
staged = PromptModel(staged_path)
|
|
rotate_history_epoch(staged.get_connection())
|
|
ok, reason = self.verify_database_file(
|
|
staged_path, require_library_state=True
|
|
)
|
|
if not ok:
|
|
raise ValueError(reason)
|
|
staged.close_all()
|
|
backup_path = ""
|
|
if os.path.exists(self.db_path):
|
|
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
backup_path = f"{self.db_path}.backup_{timestamp}"
|
|
if not self.backup_database(backup_path):
|
|
raise RuntimeError("Could not back up the current database")
|
|
self.logger.info(f"Current database backed up to: {backup_path}")
|
|
try:
|
|
self._replace_contents_from(staged_path)
|
|
except sqlite3.OperationalError as exc:
|
|
exc._prompt_manager_restore_backup_created = bool(backup_path)
|
|
raise
|
|
with PromptModel._init_lock:
|
|
PromptModel._initialized_paths.add(self._schema_key())
|
|
self.logger.info(f"Database restored from {src_path}")
|
|
from .execution_history import notify_history_restore
|
|
|
|
notify_history_restore(self.db_path)
|
|
return backup_path
|
|
finally:
|
|
if staged is not None:
|
|
staged.close_all()
|
|
with PromptModel._init_lock:
|
|
PromptModel._initialized_paths.discard(
|
|
os.path.normcase(os.path.realpath(staged_path))
|
|
)
|