import json import os import sqlite3 from datetime import datetime from system_module.gelbooru import gelbooru_date_parse class DB: def __init__(self, debug: bool): if debug: db_path = "data_test" else: db_path = "data" if not os.path.exists(db_path): os.makedirs(db_path) self.conn = sqlite3.connect(os.path.join(db_path, "db.sqlite"), check_same_thread=False) self.conn.row_factory = sqlite3.Row self._create_table() def _create_table(self): with self.conn: self.conn.execute(""" CREATE TABLE IF NOT EXISTS files ( md5 TEXT PRIMARY KEY, local_filename TEXT, is_video INTEGER, is_gif INTEGER, tags TEXT, rating TEXT, score INTEGER, timestamp TEXT, parsed_time INTEGER, display_date TEXT, source TEXT, width INTEGER, height INTEGER ) """) def _row_to_dict(self, row: sqlite3.Row) -> dict: if not row: return None d = dict(row) d["is_video"] = bool(d["is_video"]) d["is_gif"] = bool(d["is_gif"]) # Парсинг тегов из JSON-строки обратно в массив d["tags"] = json.loads(d["tags"]) if d["tags"] else [] # Восстанавливаем другие поля (например, display_date и source), # если в них были переданы кортежи или списки for field in ["display_date", "source"]: val = d.get(field) if isinstance(val, str) and (val.startswith('[') or val.startswith('{')): try: d[field] = json.loads(val) except ValueError: pass return d def add(self, data_add: dict): parsed_time = 0 if "timestamp" in data_add: try: parsed_time = int(datetime.strptime(data_add["timestamp"], "%a %b %d %H:%M:%S %z %Y").timestamp()) except ValueError: pass row_data = { "md5": data_add.get("md5"), "local_filename": data_add.get("local_filename", data_add.get("image")), "is_video": int(data_add.get("is_video", False)), "is_gif": int(data_add.get("is_gif", False)), "tags": data_add.get("tags", []), "rating": data_add.get("rating"), "score": data_add.get("score"), "timestamp": data_add.get("timestamp", data_add.get("created_at")), "parsed_time": parsed_time, "display_date": data_add.get("display_date"), "source": data_add.get("source"), "width": data_add.get("width"), "height": data_add.get("height") } # БЕЗОПАСНАЯ СЕРИАЛИЗАЦИЯ ДЛЯ SQLITE: # SQLite принимает только числа, строки и байты. # Любые списки (tags) или кортежи (display_date) мы перегоняем в JSON. for key, value in row_data.items(): if isinstance(value, (list, tuple, dict)): row_data[key] = json.dumps(value) columns = ", ".join(row_data.keys()) placeholders = ", ".join(["?"] * len(row_data)) with self.conn: self.conn.execute( f"INSERT OR REPLACE INTO files ({columns}) VALUES ({placeholders})", tuple(row_data.values()) ) def add_raw(self, data_add: dict): is_video = data_add["image"].split(".")[-1] in ["webm", "mp4", "avi"] is_gif = data_add["image"].split(".")[-1] == "gif" dt_obj = datetime.strptime(data_add["created_at"], "%a %b %d %H:%M:%S %z %Y") parsed_time = int(dt_obj.timestamp()) add_data = { "md5": data_add["md5"], "local_filename": data_add["image"], "is_video": is_video, "is_gif": is_gif, "tags": data_add["tags"].split(" "), "rating": data_add["rating"], "score": data_add["score"], "timestamp": data_add["created_at"], "parsed_time": parsed_time, "display_date": gelbooru_date_parse(data_add["created_at"]), "source": data_add["source"], "width": data_add["width"], "height": data_add["height"] } self.add(add_data) def get_all(self): cursor = self.conn.execute("SELECT * FROM files ORDER BY parsed_time DESC") return [self._row_to_dict(row) for row in cursor.fetchall()] def get_search(self, rating_tmp: list[str], tags: list[str]): mapping_rating = { "e": "explicit", "g": "general", "s": "sensitive", "q": "questionable" } query = "SELECT * FROM files WHERE 1=1" params = [] ratings = [mapping_rating.get(x, x) for x in rating_tmp if x] if ratings: placeholders = ",".join("?" * len(ratings)) query += f" AND rating IN ({placeholders})" params.extend(ratings) tags = [t for t in tags if t] if tags: for tag in tags: query += " AND EXISTS (SELECT 1 FROM json_each(files.tags) WHERE value = ?)" params.append(tag) query += " ORDER BY parsed_time DESC" cursor = self.conn.execute(query, tuple(params)) return [self._row_to_dict(row) for row in cursor.fetchall()] def get_unknow(self): cursor = self.conn.execute("SELECT * FROM files WHERE tags = '[]' OR tags IS NULL") return [self._row_to_dict(row) for row in cursor.fetchall()] def check_uni(self, md5: str): query = "SELECT 1 FROM files WHERE md5 = ? LIMIT 1" cursor = self.conn.execute(query, (md5,)) return cursor.fetchone() is not None def get_top_tag(self, tag_slise: str, count_tag_get: int = 9999999): query = """ SELECT value AS tag, COUNT(*) AS count FROM files, json_each(files.tags) WHERE value LIKE ? GROUP BY value ORDER BY count DESC LIMIT {count_tag_get} """.format(count_tag_get=count_tag_get) cursor = self.conn.execute(query, (f"{tag_slise.lower()}%",)) return [{"tag": row["tag"], "count": row["count"]} for row in cursor.fetchall()] def edit_data(self, id_file: str, edit_data: dict): if not edit_data: return set_clauses = [] params = [] for key, value in edit_data.items(): # Если передаем списки/кортежи, пакуем в JSON if isinstance(value, (list, tuple, dict)): value = json.dumps(value) set_clauses.append(f"{key} = ?") if key == "timestamp": params.append(value) try: set_clauses.append("parsed_time = ?") params.append(int(datetime.strptime(value, "%a %b %d %H:%M:%S %z %Y").timestamp())) except ValueError: pass else: params.append(value) params.append(id_file) query = f"UPDATE files SET {', '.join(set_clauses)} WHERE md5 = ?" with self.conn: self.conn.execute(query, tuple(params)) def get_from_id(self, id_file: str): cursor = self.conn.execute("SELECT * FROM files WHERE md5 = ?", (id_file,)) return self._row_to_dict(cursor.fetchone()) def close(self): self.conn.close()