410 lines
17 KiB
Python
410 lines
17 KiB
Python
"""
|
|
SQLite database implementation for the ComfyUI Queue Manager.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import sqlite3
|
|
import threading
|
|
from contextlib import contextmanager
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any, Generator
|
|
|
|
from error_handler import ErrorHandler, with_database_error_handling
|
|
from exceptions import ConnectionError, DatabaseError, SchemaError, ValidationError
|
|
from interfaces import DatabaseInterface
|
|
from logging_config import get_logger
|
|
from models import QueueConfig, QueueFilter, QueueItem, QueueStatus
|
|
from performance_monitor import get_performance_monitor, time_function
|
|
|
|
logger = get_logger("database")
|
|
|
|
|
|
class SQLiteDatabase(DatabaseInterface):
|
|
"""SQLite implementation of the database interface."""
|
|
|
|
# Database schema version for migrations
|
|
SCHEMA_VERSION = 1
|
|
|
|
# SQL statements for table creation
|
|
CREATE_TABLES_SQL = """
|
|
-- Queue Items Table
|
|
CREATE TABLE IF NOT EXISTS queue_items (
|
|
id TEXT PRIMARY KEY,
|
|
workflow_name TEXT NOT NULL,
|
|
workflow_data TEXT NOT NULL,
|
|
status TEXT NOT NULL,
|
|
created_at TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL,
|
|
started_at TEXT,
|
|
completed_at TEXT,
|
|
error_message TEXT,
|
|
result_data TEXT
|
|
);
|
|
|
|
-- Queue Configuration Table
|
|
CREATE TABLE IF NOT EXISTS queue_config (
|
|
key TEXT PRIMARY KEY,
|
|
value TEXT NOT NULL
|
|
);
|
|
|
|
-- Schema Version Table
|
|
CREATE TABLE IF NOT EXISTS schema_version (
|
|
version INTEGER PRIMARY KEY
|
|
);
|
|
|
|
-- Indexes for better query performance
|
|
CREATE INDEX IF NOT EXISTS idx_queue_items_status ON queue_items(status);
|
|
CREATE INDEX IF NOT EXISTS idx_queue_items_created_at ON queue_items(created_at);
|
|
CREATE INDEX IF NOT EXISTS idx_queue_items_workflow_name ON queue_items(workflow_name);
|
|
"""
|
|
|
|
def __init__(self, db_path: str | Path = "queue_manager.db"):
|
|
"""Initialize the SQLite database.
|
|
|
|
Args:
|
|
db_path: Path to the SQLite database file
|
|
"""
|
|
self.db_path = Path(db_path)
|
|
self._connection_lock = threading.RLock()
|
|
self._local = threading.local()
|
|
self._initialized = False
|
|
|
|
# Ensure the directory exists
|
|
self.db_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
def _get_connection(self) -> sqlite3.Connection:
|
|
"""Get a thread-local database connection."""
|
|
if not hasattr(self._local, 'connection') or self._local.connection is None:
|
|
try:
|
|
conn = sqlite3.connect(
|
|
str(self.db_path),
|
|
check_same_thread=False,
|
|
timeout=30.0,
|
|
isolation_level=None # Autocommit mode
|
|
)
|
|
conn.row_factory = sqlite3.Row
|
|
# Enable WAL mode for better concurrency
|
|
conn.execute("PRAGMA journal_mode=WAL")
|
|
# Enable foreign key constraints
|
|
conn.execute("PRAGMA foreign_keys=ON")
|
|
self._local.connection = conn
|
|
logger.debug(f"Created new database connection for thread {threading.current_thread().name}")
|
|
except sqlite3.Error as e:
|
|
logger.error(f"Failed to connect to database: {e}")
|
|
raise ErrorHandler.handle_database_error(e, "connect", context={"db_path": str(self.db_path)})
|
|
|
|
return self._local.connection
|
|
|
|
@contextmanager
|
|
def _get_cursor(self) -> Generator[sqlite3.Cursor, None, None]:
|
|
"""Get a database cursor with automatic error handling."""
|
|
conn = self._get_connection()
|
|
cursor = conn.cursor()
|
|
try:
|
|
yield cursor
|
|
except sqlite3.Error as e:
|
|
logger.error(f"Database operation failed: {e}")
|
|
try:
|
|
conn.rollback()
|
|
except sqlite3.Error as rollback_error:
|
|
logger.error(f"Failed to rollback transaction: {rollback_error}")
|
|
raise ErrorHandler.handle_database_error(e, "cursor_operation")
|
|
finally:
|
|
try:
|
|
cursor.close()
|
|
except sqlite3.Error as close_error:
|
|
logger.warning(f"Failed to close cursor: {close_error}")
|
|
|
|
@with_database_error_handling(operation="initialize")
|
|
def initialize(self) -> bool:
|
|
"""Initialize the database and create tables if needed."""
|
|
if self._initialized:
|
|
return True
|
|
|
|
with self._connection_lock:
|
|
with self._get_cursor() as cursor:
|
|
# Create all tables
|
|
cursor.executescript(self.CREATE_TABLES_SQL)
|
|
|
|
# Check and update schema version
|
|
cursor.execute("SELECT version FROM schema_version LIMIT 1")
|
|
result = cursor.fetchone()
|
|
|
|
if result is None:
|
|
# First time setup
|
|
cursor.execute("INSERT INTO schema_version (version) VALUES (?)", (self.SCHEMA_VERSION,))
|
|
logger.info(f"Database initialized with schema version {self.SCHEMA_VERSION}")
|
|
else:
|
|
current_version = result[0]
|
|
if current_version < self.SCHEMA_VERSION:
|
|
self._migrate_schema(cursor, current_version, self.SCHEMA_VERSION)
|
|
elif current_version > self.SCHEMA_VERSION:
|
|
logger.warning(f"Database schema version {current_version} is newer than expected {self.SCHEMA_VERSION}")
|
|
raise SchemaError(
|
|
f"Database schema version {current_version} is newer than expected {self.SCHEMA_VERSION}",
|
|
schema_version=current_version
|
|
)
|
|
|
|
# Set default configuration if not exists
|
|
self._initialize_default_config(cursor)
|
|
|
|
self._initialized = True
|
|
logger.info("Database initialization completed successfully")
|
|
return True
|
|
|
|
def _migrate_schema(self, cursor: sqlite3.Cursor, from_version: int, to_version: int) -> None:
|
|
"""Migrate database schema from one version to another."""
|
|
logger.info(f"Migrating database schema from version {from_version} to {to_version}")
|
|
|
|
# Add migration logic here as schema evolves
|
|
# For now, we only have version 1, so no migrations needed
|
|
|
|
cursor.execute("UPDATE schema_version SET version = ?", (to_version,))
|
|
logger.info(f"Schema migration completed to version {to_version}")
|
|
|
|
def _initialize_default_config(self, cursor: sqlite3.Cursor) -> None:
|
|
"""Initialize default configuration values."""
|
|
default_config = QueueConfig()
|
|
config_dict = default_config.to_dict()
|
|
|
|
for key, value in config_dict.items():
|
|
cursor.execute(
|
|
"INSERT OR IGNORE INTO queue_config (key, value) VALUES (?, ?)",
|
|
(key, json.dumps(value))
|
|
)
|
|
|
|
@with_database_error_handling(operation="INSERT", table="queue_items")
|
|
@time_function("database.create_queue_item")
|
|
def create_queue_item(self, item: QueueItem) -> bool:
|
|
"""Create a new queue item in the database."""
|
|
# Validate item data
|
|
if not item.id:
|
|
raise ValidationError("Queue item ID cannot be empty", field="id")
|
|
if not item.workflow_name:
|
|
raise ValidationError("Workflow name cannot be empty", field="workflow_name")
|
|
|
|
with self._get_cursor() as cursor:
|
|
cursor.execute("""
|
|
INSERT INTO queue_items (
|
|
id, workflow_name, workflow_data, status, created_at, updated_at,
|
|
started_at, completed_at, error_message, result_data
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
""", (
|
|
item.id,
|
|
item.workflow_name,
|
|
json.dumps(item.workflow_data),
|
|
item.status.value,
|
|
item.created_at.isoformat(),
|
|
item.updated_at.isoformat(),
|
|
item.started_at.isoformat() if item.started_at else None,
|
|
item.completed_at.isoformat() if item.completed_at else None,
|
|
item.error_message,
|
|
json.dumps(item.result_data) if item.result_data else None
|
|
))
|
|
logger.debug(f"Created queue item {item.id}")
|
|
return True
|
|
|
|
@with_database_error_handling(operation="SELECT", table="queue_items")
|
|
@time_function("database.get_queue_item")
|
|
def get_queue_item(self, item_id: str) -> QueueItem | None:
|
|
"""Retrieve a queue item by ID."""
|
|
if not item_id:
|
|
raise ValidationError("Item ID cannot be empty", field="item_id")
|
|
|
|
with self._get_cursor() as cursor:
|
|
cursor.execute("SELECT * FROM queue_items WHERE id = ?", (item_id,))
|
|
row = cursor.fetchone()
|
|
if row:
|
|
return self._row_to_queue_item(row)
|
|
return None
|
|
|
|
@with_database_error_handling(operation="SELECT", table="queue_items")
|
|
@time_function("database.get_all_queue_items")
|
|
def get_all_queue_items(self) -> list[QueueItem]:
|
|
"""Retrieve all queue items."""
|
|
with self._get_cursor() as cursor:
|
|
cursor.execute("SELECT * FROM queue_items ORDER BY created_at DESC")
|
|
rows = cursor.fetchall()
|
|
return [self._row_to_queue_item(row) for row in rows]
|
|
|
|
@with_database_error_handling(operation="UPDATE", table="queue_items")
|
|
def update_queue_item(self, item: QueueItem) -> bool:
|
|
"""Update an existing queue item."""
|
|
if not item.id:
|
|
raise ValidationError("Queue item ID cannot be empty", field="id")
|
|
|
|
# Update the updated_at timestamp
|
|
item.updated_at = datetime.now(timezone.utc)
|
|
|
|
with self._get_cursor() as cursor:
|
|
cursor.execute("""
|
|
UPDATE queue_items SET
|
|
workflow_name = ?, workflow_data = ?, status = ?, updated_at = ?,
|
|
started_at = ?, completed_at = ?, error_message = ?, result_data = ?
|
|
WHERE id = ?
|
|
""", (
|
|
item.workflow_name,
|
|
json.dumps(item.workflow_data),
|
|
item.status.value,
|
|
item.updated_at.isoformat(),
|
|
item.started_at.isoformat() if item.started_at else None,
|
|
item.completed_at.isoformat() if item.completed_at else None,
|
|
item.error_message,
|
|
json.dumps(item.result_data) if item.result_data else None,
|
|
item.id
|
|
))
|
|
|
|
if cursor.rowcount > 0:
|
|
logger.debug(f"Updated queue item {item.id}")
|
|
return True
|
|
else:
|
|
logger.warning(f"Queue item {item.id} not found for update")
|
|
return False
|
|
|
|
@with_database_error_handling(operation="DELETE", table="queue_items")
|
|
def delete_queue_item(self, item_id: str) -> bool:
|
|
"""Delete a queue item by ID."""
|
|
if not item_id:
|
|
raise ValidationError("Item ID cannot be empty", field="item_id")
|
|
|
|
with self._get_cursor() as cursor:
|
|
cursor.execute("DELETE FROM queue_items WHERE id = ?", (item_id,))
|
|
if cursor.rowcount > 0:
|
|
logger.debug(f"Deleted queue item {item_id}")
|
|
return True
|
|
else:
|
|
logger.warning(f"Queue item {item_id} not found for deletion")
|
|
return False
|
|
|
|
def get_items_by_status(self, status: QueueStatus) -> list[QueueItem]:
|
|
"""Retrieve queue items by status."""
|
|
try:
|
|
with self._get_cursor() as cursor:
|
|
cursor.execute(
|
|
"SELECT * FROM queue_items WHERE status = ? ORDER BY created_at DESC",
|
|
(status.value,)
|
|
)
|
|
rows = cursor.fetchall()
|
|
return [self._row_to_queue_item(row) for row in rows]
|
|
except Exception as e:
|
|
logger.error(f"Failed to get items by status {status}: {e}")
|
|
return []
|
|
|
|
def filter_queue_items(self, filter_criteria: QueueFilter) -> list[QueueItem]:
|
|
"""Filter queue items based on criteria."""
|
|
try:
|
|
query_parts = ["SELECT * FROM queue_items WHERE 1=1"]
|
|
params = []
|
|
|
|
if filter_criteria.status:
|
|
placeholders = ",".join("?" * len(filter_criteria.status))
|
|
query_parts.append(f"AND status IN ({placeholders})")
|
|
params.extend([s.value for s in filter_criteria.status])
|
|
|
|
if filter_criteria.workflow_name:
|
|
query_parts.append("AND workflow_name LIKE ?")
|
|
params.append(f"%{filter_criteria.workflow_name}%")
|
|
|
|
if filter_criteria.date_range:
|
|
start_date, end_date = filter_criteria.date_range
|
|
query_parts.append("AND created_at BETWEEN ? AND ?")
|
|
params.extend([start_date.isoformat(), end_date.isoformat()])
|
|
|
|
if filter_criteria.search_term:
|
|
query_parts.append("AND (workflow_name LIKE ? OR workflow_data LIKE ?)")
|
|
search_param = f"%{filter_criteria.search_term}%"
|
|
params.extend([search_param, search_param])
|
|
|
|
query_parts.append("ORDER BY created_at DESC")
|
|
query = " ".join(query_parts)
|
|
|
|
with self._get_cursor() as cursor:
|
|
cursor.execute(query, params)
|
|
rows = cursor.fetchall()
|
|
return [self._row_to_queue_item(row) for row in rows]
|
|
except Exception as e:
|
|
logger.error(f"Failed to filter queue items: {e}")
|
|
return []
|
|
|
|
def bulk_update_status(self, item_ids: list[str], status: QueueStatus) -> bool:
|
|
"""Update status for multiple items."""
|
|
if not item_ids:
|
|
return True
|
|
|
|
try:
|
|
placeholders = ",".join("?" * len(item_ids))
|
|
updated_at = datetime.now(timezone.utc).isoformat()
|
|
|
|
with self._get_cursor() as cursor:
|
|
cursor.execute(f"""
|
|
UPDATE queue_items
|
|
SET status = ?, updated_at = ?
|
|
WHERE id IN ({placeholders})
|
|
""", [status.value, updated_at] + item_ids)
|
|
|
|
logger.debug(f"Bulk updated {cursor.rowcount} items to status {status}")
|
|
return cursor.rowcount > 0
|
|
except Exception as e:
|
|
logger.error(f"Failed to bulk update status: {e}")
|
|
return False
|
|
|
|
def get_config(self, key: str) -> str | None:
|
|
"""Get a configuration value."""
|
|
try:
|
|
with self._get_cursor() as cursor:
|
|
cursor.execute("SELECT value FROM queue_config WHERE key = ?", (key,))
|
|
row = cursor.fetchone()
|
|
return row[0] if row else None
|
|
except Exception as e:
|
|
logger.error(f"Failed to get config {key}: {e}")
|
|
return None
|
|
|
|
def set_config(self, key: str, value: str) -> bool:
|
|
"""Set a configuration value."""
|
|
try:
|
|
with self._get_cursor() as cursor:
|
|
cursor.execute("""
|
|
INSERT OR REPLACE INTO queue_config (key, value) VALUES (?, ?)
|
|
""", (key, value))
|
|
logger.debug(f"Set config {key} = {value}")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to set config {key}: {e}")
|
|
return False
|
|
|
|
def close(self) -> None:
|
|
"""Close the database connection."""
|
|
if hasattr(self._local, 'connection') and self._local.connection:
|
|
try:
|
|
self._local.connection.close()
|
|
self._local.connection = None
|
|
logger.debug("Database connection closed")
|
|
except Exception as e:
|
|
logger.error(f"Error closing database connection: {e}")
|
|
|
|
def _row_to_queue_item(self, row: sqlite3.Row) -> QueueItem:
|
|
"""Convert a database row to a QueueItem object."""
|
|
return QueueItem(
|
|
id=row["id"],
|
|
workflow_name=row["workflow_name"],
|
|
workflow_data=json.loads(row["workflow_data"]) if row["workflow_data"] else {},
|
|
status=QueueStatus(row["status"]),
|
|
created_at=datetime.fromisoformat(row["created_at"]),
|
|
updated_at=datetime.fromisoformat(row["updated_at"]),
|
|
started_at=datetime.fromisoformat(row["started_at"]) if row["started_at"] else None,
|
|
completed_at=datetime.fromisoformat(row["completed_at"]) if row["completed_at"] else None,
|
|
error_message=row["error_message"],
|
|
result_data=json.loads(row["result_data"]) if row["result_data"] else None
|
|
)
|
|
|
|
def __enter__(self):
|
|
"""Context manager entry."""
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
"""Context manager exit."""
|
|
self.close() |