Files
ComfyAssets-ComfyUI_PromptM…/database/models.py
T
Vito Sansevero 23b0162a06 feat: sort param on /recent and /search; microsecond last_used_at
- /prompt_manager/recent and /search accept ?sort= (whitelisted in the DB
  layer; unknown values fall back to newest first).
- last_used_at is written with microseconds from Python so a re-run in the
  same millisecond as another save still sorts first; ordering stays correct
  against millisecond values from the backfill.

Refs #90
2026-10-07 11:39:49 -07:00

587 lines
23 KiB
Python

"""
Database schema and models for KikoTextEncode prompt storage.
"""
import datetime
import sqlite3
import os
import threading
from typing import Optional
# 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"
def utc_now_iso() -> str:
"""Current UTC time as ISO text with microseconds, e.g. 2026-10-07T18:39:06.758123+00:00."""
return datetime.datetime.now(datetime.timezone.utc).isoformat(
timespec="microseconds"
)
class PromptModel:
"""Database model for prompt storage and schema management."""
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
self._conn: Optional[sqlite3.Connection] = None
self._conn_lock = threading.Lock()
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
"""
try:
with sqlite3.connect(self.db_path) as conn:
conn.execute("PRAGMA journal_mode = WAL")
conn.execute("PRAGMA busy_timeout = 5000")
conn.execute("PRAGMA foreign_keys = ON")
self._create_tables(conn)
self._create_indexes(conn)
conn.commit()
except Exception as e:
self.logger.error(f"Error creating database: {e}")
raise
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,
FOREIGN KEY (prompt_id) REFERENCES prompts(id) ON DELETE CASCADE,
UNIQUE(prompt_id, filename)
)
""")
# 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
)
""")
# Add unique constraint to existing databases (migration)
self._migrate_add_unique_constraint(conn)
# 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 (last, after any migration that rebuilds prompts)
self._migrate_add_usage_columns(conn)
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 = [
"CREATE INDEX IF NOT EXISTS idx_prompts_text ON prompts(text)",
"CREATE INDEX IF NOT EXISTS idx_prompts_category ON prompts(category)",
"CREATE INDEX IF NOT EXISTS idx_prompts_created_at ON prompts(created_at)",
"CREATE INDEX IF NOT EXISTS idx_prompts_hash ON prompts(hash)",
"CREATE INDEX IF NOT EXISTS idx_prompts_rating ON prompts(rating)",
"CREATE INDEX IF NOT EXISTS idx_prompts_last_used ON prompts(last_used_at)",
"CREATE INDEX IF NOT EXISTS idx_prompts_run_count ON prompts(run_count)",
"CREATE INDEX IF NOT EXISTS idx_prompt_images ON generated_images(prompt_id)",
"CREATE INDEX IF NOT EXISTS idx_image_path ON generated_images(image_path)",
"CREATE INDEX IF NOT EXISTS idx_generation_time ON generated_images(generation_time)",
"CREATE INDEX IF NOT EXISTS idx_prompt_tags_tag ON prompt_tags(tag_id)",
"CREATE INDEX IF NOT EXISTS idx_tags_name ON tags(name)",
]
for index_sql in indexes:
conn.execute(index_sql)
def get_connection(self) -> sqlite3.Connection:
"""
Get the persistent database connection.
Returns a thread-safe, reusable connection protected by a lock.
The connection is created on first call and reused thereafter.
Callers use it as a context manager; the lock is held for the
duration of the ``with`` block.
Returns:
sqlite3.Connection: Configured database connection
"""
with self._conn_lock:
if self._conn is None:
self._conn = sqlite3.connect(self.db_path, check_same_thread=False)
self._conn.row_factory = sqlite3.Row
self._conn.execute("PRAGMA journal_mode = WAL")
self._conn.execute("PRAGMA foreign_keys = ON")
self._conn.execute("PRAGMA busy_timeout = 5000")
return self._conn
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.
"""
try:
# Check if workflow_name column exists
cursor = conn.execute("PRAGMA table_info(prompts)")
columns = [column[1] for column in cursor.fetchall()]
if "workflow_name" in columns:
self.logger.info("Migrating database: removing workflow_name column")
# Create new table without workflow_name
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
)
""")
# Copy data from old table to new table
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
""")
# Drop old table and rename new one
conn.execute("DROP TABLE prompts")
conn.execute("ALTER TABLE prompts_new RENAME TO prompts")
self.logger.info("Database migration completed")
except Exception as e:
self.logger.error(f"Migration error: {e}")
# If migration fails, the table creation will handle it
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.
"""
try:
# Check if generated_images table exists and has TEXT prompt_id
cursor = conn.execute("PRAGMA table_info(generated_images)")
columns = {column[1]: column[2] for column in cursor.fetchall()}
if "prompt_id" in columns and columns["prompt_id"] == "TEXT":
self.logger.info(
"Migrating foreign key types: prompt_id TEXT -> INTEGER"
)
# Create new table with correct types
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
)
""")
# Copy data, converting prompt_id from TEXT to INTEGER
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)
""")
# Drop old table and rename new one
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 Exception as e:
self.logger.error(f"Foreign key migration error: {e}")
# If migration fails, continue with existing schema
def _migrate_add_unique_constraint(self, conn: sqlite3.Connection) -> None:
"""
Add UNIQUE constraint on (prompt_id, filename) to prevent duplicate image entries.
This migration:
1. Checks if the constraint already exists
2. Removes duplicate entries (keeping the most recent)
3. Recreates the table with the UNIQUE constraint
Args:
conn: Active database connection
"""
try:
# Check if the unique constraint already exists by looking at table info
cursor = conn.execute("PRAGMA index_list(generated_images)")
indexes = cursor.fetchall()
# Check if we have a unique index on prompt_id, filename
has_unique_constraint = False
for idx in indexes:
if idx[2] == 1: # unique flag
cursor = conn.execute(f"PRAGMA index_info({idx[1]})")
columns = [col[2] for col in cursor.fetchall()]
if "prompt_id" in columns and "filename" in columns:
has_unique_constraint = True
break
if has_unique_constraint:
return # Already migrated
self.logger.info(
"Migrating database: adding UNIQUE constraint on (prompt_id, filename)"
)
# First, remove duplicates keeping only the most recent (highest id)
conn.execute("""
DELETE FROM generated_images
WHERE id NOT IN (
SELECT MAX(id) FROM generated_images
GROUP BY prompt_id, filename
)
""")
duplicates_removed = conn.total_changes
if duplicates_removed > 0:
self.logger.info(
f"Removed {duplicates_removed} duplicate image entries"
)
# Create new table with UNIQUE constraint
conn.execute("""
CREATE TABLE IF NOT EXISTS 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,
UNIQUE(prompt_id, filename)
)
""")
# Copy data
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, prompt_id, image_path, filename, generation_time, file_size,
width, height, format, workflow_data, prompt_metadata, parameters
FROM generated_images
""")
# Drop old table and rename new one
conn.execute("DROP TABLE generated_images")
conn.execute("ALTER TABLE generated_images_new RENAME TO generated_images")
# Recreate indexes
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_prompt_images ON generated_images(prompt_id)"
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_image_path ON generated_images(image_path)"
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_generation_time ON generated_images(generation_time)"
)
self.logger.info("UNIQUE constraint migration completed successfully")
except Exception as e:
self.logger.error(f"UNIQUE constraint migration error: {e}")
# Continue with existing schema if migration fails
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.
"""
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 Exception as e:
self.logger.error(f"Tag junction migration error: {e}")
def _migrate_add_usage_columns(self, conn: sqlite3.Connection) -> None:
"""
Add last_used_at and run_count (3.2.4) and fill in any prompt without usage data.
The backfill runs on every start but only touches rows whose last_used_at
is NULL, so it is a no-op once healed. That also repairs rows written by an
older version after a downgrade, or left unfilled by an interrupted upgrade.
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)")}
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")
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"
)
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 Exception as e:
self.logger.error(f"Usage tracking migration error: {e}")
def migrate_database(self) -> None:
"""
Apply any pending database migrations.
This method serves as an entry point for future schema migrations.
Add new migration logic here as the database evolves.
"""
# Future migrations can be added here
pass
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 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:
"""
Create a backup of the database.
Args:
backup_path: Path where the backup should be saved
Returns:
bool: True if backup was successful, False otherwise
"""
try:
import shutil
shutil.copy2(self.db_path, backup_path)
return True
except Exception as e:
self.logger.error(f"Error creating database backup: {e}")
return False