nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

style: apply Black formatter (line-length=88) to all Python files

Browse Source 
Automated formatting pass across 30 files to establish consistent code
style enforced by CI. No logic changes.
This commit is contained in:
Vito Sansevero
2026-02-07 08:30:09 -08:00
parent 224fcc7803
commit 2737a4b1ef
30 changed files with 3278 additions and 2531 deletions
Download Patch File Download Diff File Expand all files Collapse all files
database/models.py
+84 -68
View File
@@ -12,6 +12,7 @@ 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
@@ -19,28 +20,28 @@ except ImportError:
class PromptModel:
"""Database model for prompt storage and schema management."""
def __init__(self, db_path: str = "prompts.db"):
"""
Initialize the database model.
Args:
db_path: Path to the SQLite database file
"""
self.logger = get_logger('prompt_manager.database.models')
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
"""
@@ -55,14 +56,14 @@ class PromptModel:
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
@@ -80,7 +81,7 @@ class PromptModel:
hash TEXT UNIQUE
)
""")
# Create images table for gallery functionality
conn.execute("""
CREATE TABLE IF NOT EXISTS generated_images (
@@ -129,14 +130,14 @@ class PromptModel:
# Migrate JSON tags to normalized junction tables
self._migrate_json_tags_to_junction(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
@@ -155,10 +156,10 @@ class PromptModel:
"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.
@@ -173,22 +174,20 @@ class PromptModel:
"""
with self._conn_lock:
if self._conn is None:
self._conn = sqlite3.connect(
self.db_path, check_same_thread=False
)
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.
"""
@@ -196,10 +195,10 @@ class PromptModel:
# 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:
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 (
@@ -214,31 +213,31 @@ class PromptModel:
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.
"""
@@ -246,10 +245,12 @@ class PromptModel:
# 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")
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 (
@@ -268,7 +269,7 @@ class PromptModel:
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
@@ -280,17 +281,19 @@ class PromptModel:
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")
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.
@@ -314,14 +317,16 @@ class PromptModel:
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:
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)")
self.logger.info(
"Migrating database: adding UNIQUE constraint on (prompt_id, filename)"
)
# First, remove duplicates keeping only the most recent (highest id)
conn.execute("""
@@ -334,7 +339,9 @@ class PromptModel:
duplicates_removed = conn.total_changes
if duplicates_removed > 0:
self.logger.info(f"Removed {duplicates_removed} duplicate image entries")
self.logger.info(
f"Removed {duplicates_removed} duplicate image entries"
)
# Create new table with UNIQUE constraint
conn.execute("""
@@ -371,9 +378,15 @@ class PromptModel:
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)")
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")
@@ -436,11 +449,11 @@ class PromptModel:
"""
# 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.
"""
@@ -450,57 +463,60 @@ class PromptModel:
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']
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']
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']
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
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)
"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
return False