""" Менеджер базы данных SQLite: папки, файлы, метаданные, теги, избранное. """ import sqlite3 import logging from pathlib import Path from typing import Optional, List, Dict, Any from contextlib import contextmanager from comfy_studio.db.schema import PRAGMA_FOREIGN_KEYS, CREATE_TABLES_SQL, CREATE_INDEXES_SQL logger = logging.getLogger(__name__) def _normalize(path: str) -> str: return str(Path(path).resolve().as_posix()) class DBManager: def __init__(self, db_path: str = "comfystudio.db"): self.db_path = str(db_path) self._initialize_db() @contextmanager def _get_connection(self): conn = None try: conn = sqlite3.connect(self.db_path) conn.execute(PRAGMA_FOREIGN_KEYS) conn.row_factory = sqlite3.Row yield conn except sqlite3.Error as e: logger.error(f"Ошибка соединения с БД {self.db_path}: {e}") if conn: conn.rollback() raise finally: if conn: conn.close() def _initialize_db(self) -> bool: try: with self._get_connection() as conn: cursor = conn.cursor() cursor.executescript(CREATE_TABLES_SQL) cursor.executescript(CREATE_INDEXES_SQL) conn.commit() return True except sqlite3.Error as e: logger.critical(f"Критическая ошибка инициализации БД: {e}") return False # ──────────────────────────────── # Папки # ──────────────────────────────── def add_folder(self, folder_path: str, parent_id: Optional[int] = None) -> Optional[int]: if not folder_path: return None normalized_path = _normalize(folder_path) try: with self._get_connection() as conn: cursor = conn.cursor() cursor.execute( "INSERT INTO folders (path, parent_id) VALUES (?, ?) " "ON CONFLICT(path) DO UPDATE SET parent_id=excluded.parent_id RETURNING id", (normalized_path, parent_id) ) result = cursor.fetchone() if not result: cursor.execute("SELECT id FROM folders WHERE path = ?", (normalized_path,)) result = cursor.fetchone() conn.commit() return result[0] if result else None except sqlite3.Error as e: logger.error(f"Ошибка при добавлении папки {normalized_path}: {e}") return None def remove_folder_by_path(self, folder_path: str) -> bool: normalized_path = _normalize(folder_path) try: with self._get_connection() as conn: cursor = conn.cursor() cursor.execute("DELETE FROM folders WHERE path = ?", (normalized_path,)) conn.commit() return cursor.rowcount > 0 except sqlite3.Error as e: logger.error(f"Ошибка при удалении папки {normalized_path}: {e}") return False # ──────────────────────────────── # Файлы # ──────────────────────────────── def register_file(self, folder_id: int, filename: str, filepath: str, size: int, mtime: float) -> Optional[int]: if not filename or not filepath: return None normalized_filepath = _normalize(filepath) try: with self._get_connection() as conn: cursor = conn.cursor() cursor.execute( """ INSERT INTO files (folder_id, filename, filepath, size, mtime) VALUES (?, ?, ?, ?, ?) ON CONFLICT(filepath) DO UPDATE SET size = excluded.size, mtime = excluded.mtime, folder_id = excluded.folder_id RETURNING id """, (folder_id, filename, normalized_filepath, size, mtime) ) result = cursor.fetchone() conn.commit() return result[0] if result else None except sqlite3.Error as e: logger.error(f"Ошибка регистрации файла {normalized_filepath}: {e}") return None def remove_file_by_path(self, filepath: str) -> bool: normalized_filepath = _normalize(filepath) try: with self._get_connection() as conn: cursor = conn.cursor() cursor.execute("DELETE FROM files WHERE filepath = ?", (normalized_filepath,)) conn.commit() return cursor.rowcount > 0 except sqlite3.Error as e: logger.error(f"Ошибка удаления файла {normalized_filepath}: {e}") return False # ──────────────────────────────── # Метаданные # ──────────────────────────────── def save_metadata(self, file_id: int, meta_payload: Dict[str, Any]) -> bool: try: with self._get_connection() as conn: cursor = conn.cursor() cursor.execute( """ INSERT INTO metadata ( file_id, prompt_json, workflow_json, positive_prompt, negative_prompt, seed, model_name, sampler, steps, cfg ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(file_id) DO UPDATE SET prompt_json = excluded.prompt_json, workflow_json = excluded.workflow_json, positive_prompt = excluded.positive_prompt, negative_prompt = excluded.negative_prompt, seed = excluded.seed, model_name = excluded.model_name, sampler = excluded.sampler, steps = excluded.steps, cfg = excluded.cfg """, ( file_id, meta_payload.get("prompt_json"), meta_payload.get("workflow_json"), meta_payload.get("positive_prompt"), meta_payload.get("negative_prompt"), meta_payload.get("seed"), meta_payload.get("model_name"), meta_payload.get("sampler"), meta_payload.get("steps"), meta_payload.get("cfg") ) ) conn.commit() return True except sqlite3.Error as e: logger.error(f"Ошибка сохранения метаданных файла ID {file_id}: {e}") return False # ──────────────────────────────── # Выборки # ──────────────────────────────── def get_files_in_folder(self, folder_path: str, search_query: str = "", favorites_only: bool = False, tag: Optional[str] = None) -> List[Dict[str, Any]]: """Файлы папки с фильтрами: поиск, избранное, тег.""" normalized_path = _normalize(folder_path) sql = """ SELECT f.id, f.filename, f.filepath, f.rating, f.favorite, CASE WHEN m.workflow_json IS NOT NULL AND m.workflow_json != '' THEN 1 ELSE 0 END as has_workflow FROM files f JOIN folders fo ON f.folder_id = fo.id LEFT JOIN metadata m ON f.id = m.file_id WHERE fo.path = ? """ params: List[Any] = [normalized_path] if search_query: like = f"%{search_query}%" sql += " AND (f.filename LIKE ? OR m.positive_prompt LIKE ? OR m.negative_prompt LIKE ?)" params.extend([like, like, like]) if favorites_only: sql += " AND f.favorite = 1" if tag: sql += """ AND f.id IN ( SELECT ft.file_id FROM file_tags ft JOIN tags t ON ft.tag_id = t.id WHERE t.name = ? )""" params.append(tag) sql += " ORDER BY f.filename ASC" try: with self._get_connection() as conn: cursor = conn.cursor() cursor.execute(sql, params) return [dict(row) for row in cursor.fetchall()] except sqlite3.Error as e: logger.error(f"Ошибка получения файлов в папке {normalized_path}: {e}") return [] def get_file_details(self, file_id: int) -> Optional[Dict[str, Any]]: try: with self._get_connection() as conn: cursor = conn.cursor() cursor.execute( """ SELECT f.id, f.filename, f.filepath, f.rating, f.favorite, m.prompt_json, m.workflow_json, m.positive_prompt, m.negative_prompt, m.seed, m.model_name, m.sampler, m.steps, m.cfg FROM files f LEFT JOIN metadata m ON f.id = m.file_id WHERE f.id = ? """, (file_id,) ) row = cursor.fetchone() if not row: return None details = dict(row) details["tags"] = self.get_file_tags(file_id) return details except sqlite3.Error as e: logger.error(f"Ошибка получения деталей файла ID {file_id}: {e}") return None def update_file_details(self, file_id: int, positive: str, negative: str, rating: int) -> bool: try: with self._get_connection() as conn: cursor = conn.cursor() cursor.execute("UPDATE files SET rating = ? WHERE id = ?", (rating, file_id)) cursor.execute( "UPDATE metadata SET positive_prompt = ?, negative_prompt = ? WHERE file_id = ?", (positive, negative, file_id) ) conn.commit() return True except sqlite3.Error as e: logger.error(f"Ошибка сохранения правок для файла ID {file_id}: {e}") return False def update_json_field(self, file_id: int, field: str, json_text: Optional[str]) -> bool: """Обновляет workflow_json или prompt_json файла. При изменении prompt_json перечитывает производные параметры генерации.""" if field not in ("workflow_json", "prompt_json"): logger.error(f"Недопустимое JSON-поле: {field}") return False try: with self._get_connection() as conn: cursor = conn.cursor() # Гарантируем существование строки metadata cursor.execute( "INSERT OR IGNORE INTO metadata (file_id) VALUES (?)", (file_id,)) cursor.execute( f"UPDATE metadata SET {field} = ? WHERE file_id = ?", (json_text or None, file_id) ) if field == "prompt_json": # Пересчитываем параметры из нового prompt-графа from comfy_studio.parsers.metadata_parser import MetadataParser params = MetadataParser.parse_comfy_parameters(json_text) cursor.execute( """ UPDATE metadata SET positive_prompt = ?, negative_prompt = ?, seed = ?, model_name = ?, sampler = ?, steps = ?, cfg = ? WHERE file_id = ? """, ( params["positive_prompt"], params["negative_prompt"], params["seed"], params["model_name"], params["sampler"], params["steps"], params["cfg"], file_id ) ) conn.commit() return True except sqlite3.Error as e: logger.error(f"Ошибка сохранения {field} для файла ID {file_id}: {e}") return False # ──────────────────────────────── # Избранное # ──────────────────────────────── def set_favorite(self, file_id: int, favorite: bool) -> bool: try: with self._get_connection() as conn: cursor = conn.cursor() cursor.execute("UPDATE files SET favorite = ? WHERE id = ?", (1 if favorite else 0, file_id)) conn.commit() return cursor.rowcount > 0 except sqlite3.Error as e: logger.error(f"Ошибка изменения избранного для ID {file_id}: {e}") return False # ──────────────────────────────── # Теги # ──────────────────────────────── def add_tag_to_file(self, file_id: int, tag_name: str) -> bool: tag_name = tag_name.strip().lower() if not tag_name: return False try: with self._get_connection() as conn: cursor = conn.cursor() cursor.execute("INSERT OR IGNORE INTO tags (name) VALUES (?)", (tag_name,)) cursor.execute("SELECT id FROM tags WHERE name = ?", (tag_name,)) tag_id = cursor.fetchone()[0] cursor.execute( "INSERT OR IGNORE INTO file_tags (file_id, tag_id) VALUES (?, ?)", (file_id, tag_id) ) conn.commit() return True except sqlite3.Error as e: logger.error(f"Ошибка добавления тега '{tag_name}' к файлу ID {file_id}: {e}") return False def remove_tag_from_file(self, file_id: int, tag_name: str) -> bool: try: with self._get_connection() as conn: cursor = conn.cursor() cursor.execute( """ DELETE FROM file_tags WHERE file_id = ? AND tag_id = (SELECT id FROM tags WHERE name = ?) """, (file_id, tag_name.strip().lower()) ) # Подчищаем теги-сироты cursor.execute( "DELETE FROM tags WHERE id NOT IN (SELECT DISTINCT tag_id FROM file_tags)" ) conn.commit() return True except sqlite3.Error as e: logger.error(f"Ошибка удаления тега '{tag_name}' у файла ID {file_id}: {e}") return False def get_file_tags(self, file_id: int) -> List[str]: try: with self._get_connection() as conn: cursor = conn.cursor() cursor.execute( """ SELECT t.name FROM tags t JOIN file_tags ft ON t.id = ft.tag_id WHERE ft.file_id = ? ORDER BY t.name """, (file_id,) ) return [row[0] for row in cursor.fetchall()] except sqlite3.Error as e: logger.error(f"Ошибка получения тегов файла ID {file_id}: {e}") return [] def get_all_tags(self) -> List[str]: try: with self._get_connection() as conn: cursor = conn.cursor() cursor.execute("SELECT name FROM tags ORDER BY name") return [row[0] for row in cursor.fetchall()] except sqlite3.Error as e: logger.error(f"Ошибка получения списка тегов: {e}") return []