Hell13Cat
217 lines
7.9 KiB
Python
217 lines
7.9 KiB
Python
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() |