From b71f32857af9b6e468f4a8a96a70069f5cc346ba Mon Sep 17 00:00:00 2001 From: 0nikod Date: Thu, 14 May 2026 08:40:55 +0800 Subject: [PATCH] refactor tag data update service --- simple_prompt/api/handlers.py | 235 +--------------------------- simple_prompt/core/data_update.py | 251 ++++++++++++++++++++++++++++++ 2 files changed, 256 insertions(+), 230 deletions(-) create mode 100644 simple_prompt/core/data_update.py diff --git a/simple_prompt/api/handlers.py b/simple_prompt/api/handlers.py index 9f25334..ec9460c 100644 --- a/simple_prompt/api/handlers.py +++ b/simple_prompt/api/handlers.py @@ -4,27 +4,19 @@ API 处理器层 这些函数接收和返回纯 Python 数据结构，不依赖任何 Web 框架 """ -import hashlib import logging import os -from urllib.parse import urlparse from typing import Any, Dict, List, Optional from simple_prompt.core import config -from simple_prompt.core.database import DB_WRITE_LOCK, get_db_connection, reinit_duckdb +from simple_prompt.core.database import get_db_connection from simple_prompt.core import tags as tags_module from simple_prompt.core import categories as categories_module from simple_prompt.core import presets as presets_module +from simple_prompt.core import data_update as data_update_module logger = logging.getLogger("SimplePrompt") -DEFAULT_RELEASE_URL = "https://api.github.com/repos/0nikod/danbooru_tag_process/releases/latest" -ALLOWED_RELEASE_HOSTS = {"api.github.com"} -ALLOWED_ASSET_HOSTS = {"github.com", "objects.githubusercontent.com"} -MAX_TAG_DOWNLOAD_BYTES = 512 * 1024 * 1024 -DOWNLOAD_CHUNK_SIZE = 1024 * 1024 -TAG_ASSET_NAMES = ["tags_processed.parquet", "tags.parquet"] - # -------------------------------------------------------------------------------- # Tags Handlers @@ -289,45 +281,7 @@ async def handle_check_update() -> Dict[str, Any]: Returns: 更新信息 """ - import aiohttp - - release_url = DEFAULT_RELEASE_URL - - local_sha256 = "" - if os.path.exists(config.TAGS_PARQUET_PATH): - try: - sha256_hash = hashlib.sha256() - with open(config.TAGS_PARQUET_PATH, "rb") as f: - for byte_block in iter(lambda: f.read(4096), b""): - sha256_hash.update(byte_block) - local_sha256 = sha256_hash.hexdigest() - except Exception as e: - logger.error(f"Error calculating local SHA256: {e}") - - timeout = aiohttp.ClientTimeout(total=30) - async with aiohttp.ClientSession(timeout=timeout) as session: - headers = {"User-Agent": "ComfyUI-Simple-Prompt"} - async with session.get(release_url, headers=headers) as resp: - if resp.status != 200: - raise RuntimeError(f"Failed to fetch release info: {resp.status}") - release_info = await resp.json() - - target_asset = _find_tag_asset(release_info) - - if not target_asset: - raise RuntimeError("Tag data file not found") - - remote_sha256 = _extract_remote_sha256(target_asset) - - update_available = local_sha256 != remote_sha256 - - return { - "update_available": update_available, - "local_sha256": local_sha256, - "remote_sha256": remote_sha256, - "version": release_info.get("tag_name", "unknown"), - "name": release_info.get("name", ""), - } + return await data_update_module.check_update() async def handle_update_tags(url: Optional[str] = None) -> Dict[str, str]: @@ -340,120 +294,7 @@ async def handle_update_tags(url: Optional[str] = None) -> Dict[str, str]: Returns: 操作结果 """ - import aiohttp - - release_url = _validate_release_url(url or DEFAULT_RELEASE_URL) - timeout = aiohttp.ClientTimeout(total=300) - - async with aiohttp.ClientSession(timeout=timeout) as session: - headers = {"User-Agent": "ComfyUI-Simple-Prompt"} - async with session.get(release_url, headers=headers) as resp: - if resp.status != 200: - raise RuntimeError(f"Failed to fetch release info: {resp.status}") - release_info = await resp.json() - - target_asset = _find_tag_asset(release_info) - download_url = target_asset.get("browser_download_url") if target_asset else None - - if not download_url: - raise RuntimeError("Tag data file not found") - - remote_sha256 = _extract_remote_sha256(target_asset) - temp_path = config.TAGS_PARQUET_PATH + ".tmp" - download_url = _validate_asset_url(download_url) - - try: - actual_sha256 = await _download_tag_asset(session, download_url, temp_path) - - if remote_sha256 and actual_sha256 != remote_sha256: - raise RuntimeError("Downloaded tag data checksum mismatch") - - _validate_parquet_file(temp_path) - - with DB_WRITE_LOCK: - os.replace(temp_path, config.TAGS_PARQUET_PATH) - reinit_duckdb() - finally: - if os.path.exists(temp_path): - os.remove(temp_path) - - return {"status": "success", "message": "Tags updated successfully"} - - -def _find_tag_asset(release_info: Dict[str, Any]) -> Optional[Dict[str, Any]]: - for name in TAG_ASSET_NAMES: - for asset in release_info.get("assets", []): - if asset.get("name") == name: - return asset - return None - - -def _extract_remote_sha256(asset: Dict[str, Any]) -> str: - remote_digest = asset.get("digest", "") - if remote_digest.startswith("sha256:"): - return remote_digest[7:] - return remote_digest - - -def _validate_release_url(url: str) -> str: - parsed = urlparse(url) - if parsed.scheme != "https" or parsed.netloc not in ALLOWED_RELEASE_HOSTS: - raise ValueError("Unsupported release URL") - if not parsed.path.startswith("/repos/0nikod/danbooru_tag_process/releases/"): - raise ValueError("Unsupported release URL") - return url - - -def _validate_asset_url(url: str) -> str: - parsed = urlparse(url) - if parsed.scheme != "https" or parsed.netloc not in ALLOWED_ASSET_HOSTS: - raise ValueError("Unsupported asset URL") - return url - - -async def _download_tag_asset(session: Any, download_url: str, temp_path: str) -> str: - os.makedirs(config.DATA_DIR, exist_ok=True) - sha256_hash = hashlib.sha256() - total_size = 0 - headers = {"User-Agent": "ComfyUI-Simple-Prompt"} - - async with session.get(download_url, headers=headers) as resp: - if resp.status != 200: - raise RuntimeError(f"Failed to download: {resp.status}") - - final_host = urlparse(str(resp.url)).netloc - if final_host not in ALLOWED_ASSET_HOSTS: - raise ValueError("Unsupported asset redirect URL") - - with open(temp_path, "wb") as f: - async for chunk in resp.content.iter_chunked(DOWNLOAD_CHUNK_SIZE): - total_size += len(chunk) - if total_size > MAX_TAG_DOWNLOAD_BYTES: - raise RuntimeError("Downloaded tag data is too large") - sha256_hash.update(chunk) - f.write(chunk) - - return sha256_hash.hexdigest() - - -def _validate_parquet_file(path: str) -> None: - if os.path.getsize(path) < 8: - raise RuntimeError("Downloaded tag data is not a valid parquet file") - - with open(path, "rb") as f: - header = f.read(4) - f.seek(-4, os.SEEK_END) - footer = f.read(4) - - if header != b"PAR1" or footer != b"PAR1": - raise RuntimeError("Downloaded tag data is not a valid parquet file") - - conn = get_db_connection() - if not conn: - raise RuntimeError("Database unavailable") - - path_sql = path.replace("\\", "/") - conn.execute(f"SELECT name, category, post_count, alias FROM read_parquet('{path_sql}') LIMIT 0") + return await data_update_module.update_tags(url) async def handle_update_data(action: str) -> Dict[str, str]: @@ -466,70 +307,4 @@ async def handle_update_data(action: str) -> Dict[str, str]: Returns: 操作结果 """ - conn = get_db_connection() - if not conn: - raise RuntimeError("Database unavailable") - - if action == "update_liked": - target_path = config.LIKED_TAGS_PATH - elif action == "update_user": - target_path = config.USER_TAGS_PATH - else: - raise ValueError("Unknown action") - - from simple_prompt.core.database import ensure_parquet_exists - - ensure_parquet_exists(target_path, config.TAGS_SCHEMA) - target_path_sql = target_path.replace("\\", "/") - - with DB_WRITE_LOCK: - conn.execute("DROP TABLE IF EXISTS temp_target") - conn.execute(f"CREATE TABLE temp_target AS SELECT * FROM read_parquet('{target_path_sql}')") - - # Build source views - other_sources = [] - if action == "update_liked": - if os.path.exists(config.USER_TAGS_PATH): - other_sources.append( - f"SELECT name, category, post_count, alias, 2 as p FROM read_parquet('{config.USER_TAGS_PATH.replace(os.sep, '/')}')" - ) - if os.path.exists(config.DEFAULT_TAGS_PATH): - other_sources.append( - f"SELECT name, category, post_count, alias, 3 as p FROM read_parquet('{config.DEFAULT_TAGS_PATH.replace(os.sep, '/')}')" - ) - if os.path.exists(config.TAGS_PARQUET_PATH): - other_sources.append( - f"SELECT name, category, post_count, alias, 4 as p FROM read_parquet('{config.TAGS_PARQUET_PATH.replace(os.sep, '/')}')" - ) - - if other_sources: - union_q = " UNION ALL ".join(other_sources) - conn.execute(f"CREATE OR REPLACE VIEW others_raw AS {union_q}") - conn.execute(""" - CREATE OR REPLACE VIEW others_best AS - SELECT * EXCLUDE (rn) FROM ( - SELECT *, ROW_NUMBER() OVER (PARTITION BY name ORDER BY p ASC) as rn FROM others_raw - ) WHERE rn = 1 - """) - - conn.execute(""" - UPDATE temp_target - SET category = others_best.category, - post_count = others_best.post_count, - alias = others_best.alias - FROM others_best - WHERE temp_target.name = others_best.name - """) - - conn.execute("DROP VIEW others_best") - conn.execute("DROP VIEW others_raw") - - conn.execute(f"COPY temp_target TO '{target_path_sql}' (FORMAT PARQUET)") - conn.execute("DROP TABLE temp_target") - - reinit_duckdb() - - if action == "update_liked": - return {"status": "success", "message": "Liked tags data updated from main DB."} - else: - return {"status": "success", "message": "User tags data updated."} + return await data_update_module.update_data(action) diff --git a/simple_prompt/core/data_update.py b/simple_prompt/core/data_update.py new file mode 100644 index 0000000..e02cfb9 --- /dev/null +++ b/simple_prompt/core/data_update.py @@ -0,0 +1,251 @@ +""" +Tag data update service. + +This module owns downloading, validating, and refreshing tag data files. +""" + +import hashlib +import logging +import os +from typing import Any, Dict, Optional +from urllib.parse import urlparse + +from . import config +from .database import DB_WRITE_LOCK, ensure_parquet_exists, get_db_connection, reinit_duckdb + +logger = logging.getLogger("SimplePrompt") + +DEFAULT_RELEASE_URL = "https://api.github.com/repos/0nikod/danbooru_tag_process/releases/latest" +ALLOWED_RELEASE_HOSTS = {"api.github.com"} +ALLOWED_ASSET_HOSTS = {"github.com", "objects.githubusercontent.com"} +MAX_TAG_DOWNLOAD_BYTES = 512 * 1024 * 1024 +DOWNLOAD_CHUNK_SIZE = 1024 * 1024 +TAG_ASSET_NAMES = ["tags_processed.parquet", "tags.parquet"] + + +async def check_update() -> Dict[str, Any]: + import aiohttp + + local_sha256 = _calculate_local_sha256(config.TAGS_PARQUET_PATH) + + timeout = aiohttp.ClientTimeout(total=30) + async with aiohttp.ClientSession(timeout=timeout) as session: + release_info = await _fetch_release_info(session, DEFAULT_RELEASE_URL) + target_asset = _find_tag_asset(release_info) + + if not target_asset: + raise RuntimeError("Tag data file not found") + + remote_sha256 = _extract_remote_sha256(target_asset) + + return { + "update_available": local_sha256 != remote_sha256, + "local_sha256": local_sha256, + "remote_sha256": remote_sha256, + "version": release_info.get("tag_name", "unknown"), + "name": release_info.get("name", ""), + } + + +async def update_tags(url: Optional[str] = None) -> Dict[str, str]: + import aiohttp + + release_url = _validate_release_url(url or DEFAULT_RELEASE_URL) + timeout = aiohttp.ClientTimeout(total=300) + + async with aiohttp.ClientSession(timeout=timeout) as session: + release_info = await _fetch_release_info(session, release_url) + target_asset = _find_tag_asset(release_info) + download_url = target_asset.get("browser_download_url") if target_asset else None + + if not download_url: + raise RuntimeError("Tag data file not found") + + remote_sha256 = _extract_remote_sha256(target_asset) + temp_path = config.TAGS_PARQUET_PATH + ".tmp" + download_url = _validate_asset_url(download_url) + + try: + actual_sha256 = await _download_tag_asset(session, download_url, temp_path) + + if remote_sha256 and actual_sha256 != remote_sha256: + raise RuntimeError("Downloaded tag data checksum mismatch") + + _validate_parquet_file(temp_path) + + with DB_WRITE_LOCK: + os.replace(temp_path, config.TAGS_PARQUET_PATH) + reinit_duckdb() + finally: + if os.path.exists(temp_path): + os.remove(temp_path) + + return {"status": "success", "message": "Tags updated successfully"} + + +async def update_data(action: str) -> Dict[str, str]: + conn = get_db_connection() + if not conn: + raise RuntimeError("Database unavailable") + + if action == "update_liked": + target_path = config.LIKED_TAGS_PATH + elif action == "update_user": + target_path = config.USER_TAGS_PATH + else: + raise ValueError("Unknown action") + + ensure_parquet_exists(target_path, config.TAGS_SCHEMA) + target_path_sql = target_path.replace("\\", "/") + + with DB_WRITE_LOCK: + conn.execute("DROP TABLE IF EXISTS temp_target") + conn.execute(f"CREATE TABLE temp_target AS SELECT * FROM read_parquet('{target_path_sql}')") + + other_sources = _build_update_sources(action) + if other_sources: + union_q = " UNION ALL ".join(other_sources) + conn.execute(f"CREATE OR REPLACE VIEW others_raw AS {union_q}") + conn.execute(""" + CREATE OR REPLACE VIEW others_best AS + SELECT * EXCLUDE (rn) FROM ( + SELECT *, ROW_NUMBER() OVER (PARTITION BY name ORDER BY p ASC) as rn FROM others_raw + ) WHERE rn = 1 + """) + + conn.execute(""" + UPDATE temp_target + SET category = others_best.category, + post_count = others_best.post_count, + alias = others_best.alias + FROM others_best + WHERE temp_target.name = others_best.name + """) + + conn.execute("DROP VIEW others_best") + conn.execute("DROP VIEW others_raw") + + conn.execute(f"COPY temp_target TO '{target_path_sql}' (FORMAT PARQUET)") + conn.execute("DROP TABLE temp_target") + + reinit_duckdb() + + if action == "update_liked": + return {"status": "success", "message": "Liked tags data updated from main DB."} + return {"status": "success", "message": "User tags data updated."} + + +def _calculate_local_sha256(path: str) -> str: + if not os.path.exists(path): + return "" + + try: + sha256_hash = hashlib.sha256() + with open(path, "rb") as f: + for byte_block in iter(lambda: f.read(4096), b""): + sha256_hash.update(byte_block) + return sha256_hash.hexdigest() + except Exception as e: + logger.error(f"Error calculating local SHA256: {e}") + return "" + + +async def _fetch_release_info(session: Any, release_url: str) -> Dict[str, Any]: + headers = {"User-Agent": "ComfyUI-Simple-Prompt"} + async with session.get(release_url, headers=headers) as resp: + if resp.status != 200: + raise RuntimeError(f"Failed to fetch release info: {resp.status}") + return await resp.json() + + +def _find_tag_asset(release_info: Dict[str, Any]) -> Optional[Dict[str, Any]]: + for name in TAG_ASSET_NAMES: + for asset in release_info.get("assets", []): + if asset.get("name") == name: + return asset + return None + + +def _extract_remote_sha256(asset: Dict[str, Any]) -> str: + remote_digest = asset.get("digest", "") + if remote_digest.startswith("sha256:"): + return remote_digest[7:] + return remote_digest + + +def _validate_release_url(url: str) -> str: + parsed = urlparse(url) + if parsed.scheme != "https" or parsed.netloc not in ALLOWED_RELEASE_HOSTS: + raise ValueError("Unsupported release URL") + if not parsed.path.startswith("/repos/0nikod/danbooru_tag_process/releases/"): + raise ValueError("Unsupported release URL") + return url + + +def _validate_asset_url(url: str) -> str: + parsed = urlparse(url) + if parsed.scheme != "https" or parsed.netloc not in ALLOWED_ASSET_HOSTS: + raise ValueError("Unsupported asset URL") + return url + + +async def _download_tag_asset(session: Any, download_url: str, temp_path: str) -> str: + os.makedirs(config.DATA_DIR, exist_ok=True) + sha256_hash = hashlib.sha256() + total_size = 0 + headers = {"User-Agent": "ComfyUI-Simple-Prompt"} + + async with session.get(download_url, headers=headers) as resp: + if resp.status != 200: + raise RuntimeError(f"Failed to download: {resp.status}") + + final_host = urlparse(str(resp.url)).netloc + if final_host not in ALLOWED_ASSET_HOSTS: + raise ValueError("Unsupported asset redirect URL") + + with open(temp_path, "wb") as f: + async for chunk in resp.content.iter_chunked(DOWNLOAD_CHUNK_SIZE): + total_size += len(chunk) + if total_size > MAX_TAG_DOWNLOAD_BYTES: + raise RuntimeError("Downloaded tag data is too large") + sha256_hash.update(chunk) + f.write(chunk) + + return sha256_hash.hexdigest() + + +def _validate_parquet_file(path: str) -> None: + if os.path.getsize(path) < 8: + raise RuntimeError("Downloaded tag data is not a valid parquet file") + + with open(path, "rb") as f: + header = f.read(4) + f.seek(-4, os.SEEK_END) + footer = f.read(4) + + if header != b"PAR1" or footer != b"PAR1": + raise RuntimeError("Downloaded tag data is not a valid parquet file") + + conn = get_db_connection() + if not conn: + raise RuntimeError("Database unavailable") + + path_sql = path.replace("\\", "/") + conn.execute(f"SELECT name, category, post_count, alias FROM read_parquet('{path_sql}') LIMIT 0") + + +def _build_update_sources(action: str) -> list[str]: + other_sources = [] + if action == "update_liked" and os.path.exists(config.USER_TAGS_PATH): + other_sources.append( + f"SELECT name, category, post_count, alias, 2 as p FROM read_parquet('{config.USER_TAGS_PATH.replace(os.sep, '/')}')" + ) + if os.path.exists(config.DEFAULT_TAGS_PATH): + other_sources.append( + f"SELECT name, category, post_count, alias, 3 as p FROM read_parquet('{config.DEFAULT_TAGS_PATH.replace(os.sep, '/')}')" + ) + if os.path.exists(config.TAGS_PARQUET_PATH): + other_sources.append( + f"SELECT name, category, post_count, alias, 4 as p FROM read_parquet('{config.TAGS_PARQUET_PATH.replace(os.sep, '/')}')" + ) + return other_sources