Spaces:
Running
Running
Download db_utils.py from ZHIWEI666/ComfyUI-Ranking-API: direct link, hf CLI and curl.
- Browser
- Download file 24.3 kB
-
https://huggingface.co/spaces/ZHIWEI666/ComfyUI-Ranking-API/resolve/main/db_utils.py
- Command line
-
hf download hf://spaces/ZHIWEI666/ComfyUI-Ranking-API/db_utils.py
-
curl -L -o db_utils.py https://huggingface.co/spaces/ZHIWEI666/ComfyUI-Ranking-API/resolve/main/db_utils.py
24.3 kB
| # 云端Space代码/db_utils.py | |
| # ========================================== | |
| # 🔧 P2代码质量优化:数据库工具函数 | |
| # ========================================== | |
| # 作用:封装 JSON 数据库常用操作,减少重复代码 | |
| # 关联文件: | |
| # - 数据库连接.py (基础读写) | |
| # - router_tasks.py (任务操作) | |
| # - router_posts.py (帖子操作) | |
| # - router_items.py (商品操作) | |
| # ========================================== | |
| from typing import Any, Dict, List, Optional, Callable, Union | |
| import copy | |
| from collections import OrderedDict | |
| import threading | |
| import 数据库连接 as db | |
| # ========================================== | |
| # 📖 查询工具函数 | |
| # ========================================== | |
| def read_public_source(store, file_name, default_data): | |
| """公开读取拒绝故障默认值及异常集合;保留已验证文件缓存。""" | |
| try: | |
| data = store.load_data(file_name, default_data=default_data, strict=True) | |
| if type(data) is not type(default_data): | |
| raise ValueError("invalid public source root") | |
| if isinstance(data, list): | |
| if any(not isinstance(row, dict) or not row for row in data): | |
| raise ValueError("invalid public source records") | |
| elif file_name == "users.json": | |
| if any(not isinstance(key, str) or not key or not isinstance(row, dict) | |
| for key, row in data.items()): | |
| raise ValueError("invalid public user records") | |
| elif file_name == "comments.json": | |
| if any(not isinstance(rows, list) or | |
| any(not isinstance(row, dict) or not row for row in rows) | |
| for rows in data.values()): | |
| raise ValueError("invalid public comment records") | |
| elif file_name == "versions.json": | |
| if any(row is not None and not isinstance(row, (str, dict)) or | |
| isinstance(row, dict) and row.get("hash") is not None and | |
| not isinstance(row.get("hash"), str) | |
| for row in data.values()): | |
| raise ValueError("invalid public version records") | |
| return data | |
| except Exception: | |
| from fastapi import HTTPException | |
| raise HTTPException(status_code=503, detail="ranking_data_unavailable") from None | |
| def get_by_id(file_name: str, item_id: str, id_field: str = "id") -> Optional[Dict]: | |
| """ | |
| 根据 ID 获取单个记录 | |
| 参数: | |
| file_name: JSON 文件名(如 tasks.json) | |
| item_id: 要查找的 ID | |
| id_field: ID 字段名(默认 "id") | |
| 返回: | |
| 找到的记录,或 None | |
| 示例: | |
| task = get_by_id("tasks.json", "task_123") | |
| user = get_by_id("users.json", "user@example.com", id_field="account") | |
| """ | |
| data = db.load_data(file_name, default_data=[]) | |
| if isinstance(data, dict): | |
| return data.get(item_id) | |
| return next((item for item in data if item.get(id_field) == item_id), None) | |
| def get_by_field(file_name: str, field: str, value: Any) -> Optional[Dict]: | |
| """ | |
| 根据指定字段获取单个记录 | |
| 参数: | |
| file_name: JSON 文件名 | |
| field: 字段名 | |
| value: 字段值 | |
| 返回: | |
| 找到的记录,或 None | |
| """ | |
| data = db.load_data(file_name, default_data=[]) | |
| if isinstance(data, dict): | |
| for item in data.values(): | |
| if isinstance(item, dict) and item.get(field) == value: | |
| return item | |
| return None | |
| return next((item for item in data if item.get(field) == value), None) | |
| def filter_by(file_name: str, **conditions) -> List[Dict]: | |
| """ | |
| 根据条件筛选记录 | |
| 参数: | |
| file_name: JSON 文件名 | |
| **conditions: 筛选条件(键值对) | |
| 返回: | |
| 符合条件的记录列表 | |
| 示例: | |
| open_tasks = filter_by("tasks.json", status="open") | |
| user_posts = filter_by("posts.json", author="user123", deleted=False) | |
| """ | |
| data = db.load_data(file_name, default_data=[]) | |
| if isinstance(data, dict): | |
| data = list(data.values()) | |
| result = [] | |
| for item in data: | |
| if all(item.get(key) == value for key, value in conditions.items()): | |
| result.append(item) | |
| return result | |
| def count_by(file_name: str, **conditions) -> int: | |
| """ | |
| 统计符合条件的记录数量 | |
| 参数: | |
| file_name: JSON 文件名 | |
| **conditions: 筛选条件 | |
| 返回: | |
| 符合条件的记录数量 | |
| """ | |
| return len(filter_by(file_name, **conditions)) | |
| # ========================================== | |
| # ✏️ 更新工具函数 | |
| # ========================================== | |
| class _Unchanged(Exception): | |
| def __init__(self, result): | |
| self.result = result | |
| def _mutate(file_name, callback): | |
| def update(data): | |
| if not isinstance(data, (dict, list)) or not all(isinstance(item, dict) for item in (data.values() if isinstance(data, dict) else data)): | |
| raise ValueError('invalid mutable record source') | |
| result = callback(data) | |
| if result is False or (type(result) is int and result == 0): | |
| raise _Unchanged(result) | |
| return result | |
| try: | |
| return db.atomic_update(file_name, update, default_data=[]) | |
| except _Unchanged as unchanged: | |
| return unchanged.result | |
| def _find_record(data, item_id, id_field): | |
| if isinstance(data, dict): | |
| return data.get(item_id) | |
| matches = [item for item in data if item.get(id_field) == item_id] | |
| if len(matches) > 1: | |
| raise ValueError('ambiguous mutable record identity') | |
| return matches[0] if matches else None | |
| def update_by_id(file_name: str, item_id: str, updates: Dict, id_field: str = "id") -> bool: | |
| """ | |
| 根据 ID 更新记录 | |
| 参数: | |
| file_name: JSON 文件名 | |
| item_id: 要更新的 ID | |
| updates: 要更新的字段(键值对) | |
| id_field: ID 字段名 | |
| 返回: | |
| True 更新成功 / False 记录不存在 | |
| 示例: | |
| update_by_id("tasks.json", "task_123", {"status": "completed"}) | |
| """ | |
| saved_updates = copy.deepcopy(updates) | |
| def update(data): | |
| item = _find_record(data, item_id, id_field) | |
| if item is None: return False | |
| item.update(copy.deepcopy(saved_updates)) | |
| return True | |
| return _mutate(file_name, update) | |
| def update_with_fn(file_name: str, item_id: str, update_fn: Callable[[Dict], None], id_field: str = "id") -> bool: | |
| """ | |
| 使用函数更新记录(支持复杂更新逻辑) | |
| 参数: | |
| file_name: JSON 文件名 | |
| item_id: 要更新的 ID | |
| update_fn: 更新函数,接收记录 dict,直接修改 | |
| id_field: ID 字段名 | |
| 返回: | |
| True 更新成功 / False 记录不存在 | |
| 示例: | |
| def increment_views(item): | |
| item["views"] = item.get("views", 0) + 1 | |
| update_with_fn("items.json", "item_123", increment_views) | |
| """ | |
| def update(data): | |
| item = _find_record(data, item_id, id_field) | |
| if item is None: return False | |
| update_fn(item) | |
| return True | |
| return _mutate(file_name, update) | |
| # ========================================== | |
| # ➕ 添加工具函数 | |
| # ========================================== | |
| def insert(file_name: str, item: Dict, prepend: bool = True) -> bool: | |
| """ | |
| 插入新记录 | |
| 参数: | |
| file_name: JSON 文件名 | |
| item: 要插入的记录 | |
| prepend: True 插入到开头 / False 插入到末尾 | |
| 返回: | |
| True 插入成功 | |
| 示例: | |
| insert("tasks.json", {"id": "task_123", "title": "新任务"}) | |
| """ | |
| saved = copy.deepcopy(item) | |
| if not isinstance(saved, dict): raise ValueError('invalid inserted record') | |
| def append(data): | |
| if isinstance(data, dict): | |
| item_id = saved.get('id') or saved.get('account') | |
| if not item_id: return False | |
| data[item_id] = copy.deepcopy(saved) | |
| elif prepend: | |
| data.insert(0, copy.deepcopy(saved)) | |
| else: | |
| data.append(copy.deepcopy(saved)) | |
| return True | |
| return _mutate(file_name, append) | |
| def insert_if_not_exists(file_name: str, item: Dict, id_field: str = "id") -> bool: | |
| """ | |
| 如果不存在则插入 | |
| 参数: | |
| file_name: JSON 文件名 | |
| item: 要插入的记录 | |
| id_field: ID 字段名 | |
| 返回: | |
| True 插入成功 / False 已存在 | |
| """ | |
| saved = copy.deepcopy(item) | |
| if not isinstance(saved, dict): raise ValueError('invalid inserted record') | |
| item_id = saved.get(id_field) | |
| if not item_id: return False | |
| def append(data): | |
| if isinstance(data, dict): | |
| if item_id in data: return False | |
| data[item_id] = copy.deepcopy(saved) | |
| else: | |
| if _find_record(data, item_id, id_field) is not None: return False | |
| data.insert(0, copy.deepcopy(saved)) | |
| return True | |
| return _mutate(file_name, append) | |
| # ========================================== | |
| # ❌ 删除工具函数 | |
| # ========================================== | |
| def delete_by_id(file_name: str, item_id: str, id_field: str = "id") -> bool: | |
| """ | |
| 根据 ID 删除记录 | |
| 参数: | |
| file_name: JSON 文件名 | |
| item_id: 要删除的 ID | |
| id_field: ID 字段名 | |
| 返回: | |
| True 删除成功 / False 记录不存在 | |
| """ | |
| return bool(batch_delete(file_name, [item_id], id_field)) | |
| def soft_delete_by_id(file_name: str, item_id: str, id_field: str = "id") -> bool: | |
| """ | |
| 软删除(标记为已删除,不物理删除) | |
| 参数: | |
| file_name: JSON 文件名 | |
| item_id: 要删除的 ID | |
| id_field: ID 字段名 | |
| 返回: | |
| True 成功 / False 记录不存在 | |
| """ | |
| import time | |
| return update_by_id(file_name, item_id, { | |
| "deleted": True, | |
| "deleted_at": int(time.time()) | |
| }, id_field) | |
| # ========================================== | |
| # 🔍 分页工具函数 | |
| # ========================================== | |
| def paginate( | |
| file_name: str, | |
| page: int = 1, | |
| limit: int = 20, | |
| sort_by: str = None, | |
| sort_desc: bool = True, | |
| **filters | |
| ) -> Dict: | |
| """ | |
| 分页查询 | |
| 参数: | |
| file_name: JSON 文件名 | |
| page: 页码(从 1 开始) | |
| limit: 每页数量 | |
| sort_by: 排序字段 | |
| sort_desc: 是否降序 | |
| **filters: 筛选条件 | |
| 返回: | |
| { | |
| "data": [...], # 当前页数据 | |
| "total": 100, # 总数 | |
| "page": 1, # 当前页 | |
| "limit": 20, # 每页数量 | |
| "pages": 5 # 总页数 | |
| } | |
| """ | |
| # 获取并筛选数据 | |
| if filters: | |
| data = filter_by(file_name, **filters) | |
| else: | |
| data = db.load_data(file_name, default_data=[]) | |
| if isinstance(data, dict): | |
| data = list(data.values()) | |
| # 排序 | |
| if sort_by: | |
| try: | |
| data = sorted(data, key=lambda x: x.get(sort_by, 0), reverse=sort_desc) | |
| except TypeError: | |
| pass # 排序失败时忽略 | |
| # 计算分页 | |
| total = len(data) | |
| pages = (total + limit - 1) // limit # 向上取整 | |
| start = (page - 1) * limit | |
| end = start + limit | |
| return { | |
| "data": data[start:end], | |
| "total": total, | |
| "page": page, | |
| "limit": limit, | |
| "pages": pages | |
| } | |
| # ========================================== | |
| # 👁️ 访问量记录工具函数 | |
| # ========================================== | |
| def record_view(data_file: str, item_id: str, user_account: str) -> Optional[Dict]: | |
| """ | |
| 记录访问量(原子操作,并发安全) | |
| 参数: | |
| data_file: JSON 文件名(如 Items.json) | |
| item_id: 要记录访问的 item/task ID | |
| user_account: 访问用户账号 | |
| 返回: | |
| {"views": N, "daily_views": N} 成功 | |
| None 记录不存在 | |
| 逻辑: | |
| - 初始化字段: views=0, viewed_by=[], daily_views=0, daily_views_date="" | |
| - daily_views 每次调用都增加 | |
| - views 只在用户首次访问时增加(user_account 不在 viewed_by 中) | |
| - 如果 daily_views_date 不是今天,重置 daily_views=0 并更新日期 | |
| 并发安全: | |
| - 使用 atomic_update 确保读-改-写在同一把锁内完成 | |
| - 高并发下不会丢失访问量 | |
| """ | |
| from datetime import datetime, timedelta | |
| # 用于在闭包中存储结果 | |
| result_container = [None] | |
| def updater(data): | |
| # 查找目标记录 | |
| target_item = None | |
| if isinstance(data, dict): | |
| if item_id in data: | |
| target_item = data[item_id] | |
| else: | |
| for item in data: | |
| if item.get("id") == item_id: | |
| target_item = item | |
| break | |
| if target_item is None: | |
| result_container[0] = None | |
| return | |
| # 初始化字段(如果不存在) | |
| if "views" not in target_item: | |
| target_item["views"] = 0 | |
| if "viewed_by" not in target_item: | |
| target_item["viewed_by"] = [] | |
| if "daily_views" not in target_item: | |
| target_item["daily_views"] = 0 | |
| if "daily_views_date" not in target_item: | |
| target_item["daily_views_date"] = "" | |
| # 获取今天的日期字符串(使用 UTC+8 时区) | |
| now_utc8 = datetime.utcnow() + timedelta(hours=8) | |
| today_str = now_utc8.date().isoformat() # "2026-04-02" | |
| # 检查是否需要重置日访问量 | |
| if target_item["daily_views_date"] != today_str: | |
| target_item["daily_views"] = 0 | |
| target_item["daily_views_date"] = today_str | |
| # 增加日访问量(每次调用都增加) | |
| target_item["daily_views"] += 1 | |
| # 检查用户是否已访问过 | |
| if user_account not in target_item["viewed_by"]: | |
| target_item["viewed_by"].append(user_account) | |
| target_item["views"] += 1 | |
| # 保存结果到闭包容器 | |
| result_container[0] = { | |
| "views": target_item["views"], | |
| "daily_views": target_item["daily_views"] | |
| } | |
| # 使用原子更新,整个读-改-写过程在同一把锁内完成 | |
| db.atomic_update(data_file, updater, default_data=[]) | |
| return result_container[0] | |
| # ========================================== | |
| # 🗂️ 排序结果缓存工具类 | |
| # ========================================== | |
| import time | |
| from typing import Callable, List, Dict, Any | |
| class SortCache: | |
| """ | |
| 排序结果缓存类 - 用于缓存列表排序结果,减少重复排序开销 | |
| 设计原则: | |
| 1. 缓存排序后的 ID 顺序,而非完整数据,避免内存浪费 | |
| 2. 使用 TTL (默认5分钟) 自动过期,确保数据新鲜度 | |
| 3. 写操作时清除相关缓存,确保数据一致性 | |
| 4. 在 load_data 返回的最新数据之上应用缓存的顺序 | |
| 使用示例: | |
| # 在列表接口中 | |
| def get_items(sort="time"): | |
| items = db.load_data("items.json", default_data=[]) | |
| cache_key = f"items:{sort}" | |
| def sort_fn(data): | |
| if sort == "likes": | |
| data.sort(key=lambda x: x.get("likes", 0), reverse=True) | |
| else: | |
| data.sort(key=lambda x: x.get("created_at", 0), reverse=True) | |
| return sort_cache.get_sorted(cache_key, items, sort_fn) | |
| # 在写操作后 | |
| def create_item(...): | |
| ... | |
| sort_cache.invalidate("items") | |
| """ | |
| def __init__(self, ttl: int = 300, max_entries: int = 128, max_ids: int = 100000): | |
| """ | |
| 初始化排序缓存 | |
| 参数: | |
| ttl: 缓存过期时间(秒),默认 300 秒(5分钟) | |
| max_entries: 最多保留的缓存条目,0 表示不保留 | |
| max_ids: 所有缓存合计保留的 ID 数量,0 表示不保留 | |
| """ | |
| if any(type(value) is not int or value < 0 for value in (max_entries, max_ids)): | |
| raise ValueError("sort cache limits must be nonnegative integers") | |
| self._cache = OrderedDict() # {cache_key: (sorted_ids, timestamp)} | |
| self._max_entries = max_entries | |
| self._max_ids = max_ids | |
| self._cached_ids = 0 | |
| self._ttl = ttl | |
| self._lock = threading.Lock() | |
| self._pending_sorts: Dict[str, object] = {} | |
| self._read_epoch = object() | |
| def capture_read_epoch(self): | |
| """在读取排序源数据前捕获;只标识本进程失效代次,不是文件版本。""" | |
| with self._lock: | |
| return self._read_epoch | |
| def _remove_cached(self, key: str): | |
| # 调用者持有 _lock;失效、过期和容量淘汰共用计数维护。 | |
| cached = self._cache.pop(key, None) | |
| if cached is not None: | |
| self._cached_ids -= len(cached[0]) | |
| def get_sorted(self, cache_key: str, items: List[Dict], sort_fn: Callable[[List[Dict]], None], *, read_epoch=None, fresh: bool = False) -> List[Dict]: | |
| """ | |
| 获取排序后的数据(带缓存) | |
| 参数: | |
| cache_key: 缓存键,应包含数据文件和排序参数 | |
| items: 原始数据列表(来自 load_data 的最新数据) | |
| sort_fn: 排序函数,接收 items 列表并原地排序 | |
| read_epoch: 读取前捕获的代次;已失效输入只排序返回,不读写缓存 | |
| fresh: 按本次输入直接排序,不读写排序缓存或发布标记 | |
| 返回: | |
| 排序后的 items 列表 | |
| """ | |
| if fresh: | |
| # 公开榜单源数据可能由其它进程更新,旧ID顺序不可作为当前排序依据。 | |
| # 直接排序比逐条序列化/摘要校验更轻;不触碰其它调用的缓存与代次。 | |
| sort_fn(items) | |
| return items | |
| now = time.time() | |
| with self._lock: | |
| stale_read = read_epoch is not None and read_epoch is not self._read_epoch | |
| hit = False | |
| if not stale_read: | |
| cached = self._cache.get(cache_key) | |
| hit = cached is not None and now - cached[1] < self._ttl | |
| if hit: | |
| self._cache.move_to_end(cache_key) | |
| else: | |
| # 未命中时机会性回收过期ID列表,不新增后台线程。 | |
| for key, (_, cached_time) in list(self._cache.items()): | |
| if now - cached_time >= self._ttl: | |
| self._remove_cached(key) | |
| token = object() | |
| # 同键后开始的排序取得发布权;旧计算只能返回自己的结果。 | |
| self._pending_sorts[cache_key] = token | |
| if stale_read: | |
| # 失效前读取的输入不能命中新代缓存,也不能撤销新代发布标记。 | |
| sort_fn(items) | |
| return items | |
| if hit: | |
| sorted_ids, cached_time = cached | |
| id_order = {id_: idx for idx, id_ in enumerate(sorted_ids)} | |
| return sorted(items, key=lambda x: id_order.get(x.get("id"), float('inf'))) | |
| completed = False | |
| try: | |
| # 耗时排序及业务回调均在锁外,失效/其它键不会被阻塞。 | |
| sort_fn(items) | |
| # 超预算仍返回原排序结果,不构造/保留额外ID列表,也不淘汰有效旧缓存。 | |
| sorted_ids = ([item.get("id") for item in items] | |
| if self._max_entries and len(items) <= self._max_ids else []) | |
| completed = True | |
| return items | |
| finally: | |
| with self._lock: | |
| if self._pending_sorts.get(cache_key) is token: | |
| del self._pending_sorts[cache_key] | |
| if completed and sorted_ids: | |
| self._remove_cached(cache_key) | |
| while (len(self._cache) >= self._max_entries or | |
| self._cached_ids + len(sorted_ids) > self._max_ids): | |
| self._remove_cached(next(iter(self._cache))) | |
| self._cache[cache_key] = (sorted_ids, now) | |
| self._cached_ids += len(sorted_ids) | |
| def invalidate(self, prefix: str = ""): | |
| """ | |
| 清除缓存 | |
| 参数: | |
| prefix: 缓存键前缀,如果指定则只清除匹配的缓存,否则清除所有缓存 | |
| """ | |
| with self._lock: | |
| self._read_epoch = object() | |
| if prefix: | |
| for key in list(self._cache): | |
| if key.startswith(prefix): | |
| self._remove_cached(key) | |
| for key in list(self._pending_sorts): | |
| if key.startswith(prefix): | |
| del self._pending_sorts[key] | |
| else: | |
| self._cache.clear() | |
| self._cached_ids = 0 | |
| self._pending_sorts.clear() | |
| def get_stats(self) -> Dict[str, Any]: | |
| """ | |
| 获取缓存统计信息(用于调试) | |
| """ | |
| now = time.time() | |
| with self._lock: | |
| snapshot = tuple(self._cache.values()) | |
| valid_count = sum(1 for _, ts in snapshot if now - ts < self._ttl) | |
| return { | |
| "total_cached": len(snapshot), | |
| "valid": valid_count, | |
| "expired": len(snapshot) - valid_count, | |
| "ttl": self._ttl | |
| } | |
| # 全局排序缓存实例(TTL 5分钟) | |
| sort_cache = SortCache(ttl=300) | |
| # ========================================== | |
| # 🔄 批量操作工具函数 | |
| # ========================================== | |
| def batch_update(file_name: str, item_ids: List[str], updates: Dict, id_field: str = "id") -> int: | |
| """ | |
| 批量更新记录 | |
| 参数: | |
| file_name: JSON 文件名 | |
| item_ids: 要更新的 ID 列表 | |
| updates: 要更新的字段 | |
| id_field: ID 字段名 | |
| 返回: | |
| 更新的记录数量 | |
| """ | |
| identities, saved_updates = set(item_ids), copy.deepcopy(updates) | |
| def update(data): | |
| selected = [item for key, item in data.items() if key in identities] if isinstance(data, dict) else [item for item in data if item.get(id_field) in identities] | |
| for item in selected: item.update(copy.deepcopy(saved_updates)) | |
| return len(selected) | |
| return _mutate(file_name, update) | |
| def batch_delete(file_name: str, item_ids: List[str], id_field: str = "id") -> int: | |
| """ | |
| 批量删除记录 | |
| 参数: | |
| file_name: JSON 文件名 | |
| item_ids: 要删除的 ID 列表 | |
| id_field: ID 字段名 | |
| 返回: | |
| 删除的记录数量 | |
| """ | |
| identities = set(item_ids) | |
| def delete(data): | |
| if isinstance(data, dict): | |
| selected = [key for key in data if key in identities] | |
| for key in selected: del data[key] | |
| return len(selected) | |
| old_count = len(data) | |
| data[:] = [item for item in data if item.get(id_field) not in identities] | |
| return old_count - len(data) | |
| return _mutate(file_name, delete) | |