Spaces:
Running
Running
Download notification_store.py from ZHIWEI666/ComfyUI-Ranking-API: direct link, hf CLI and curl.
- Browser
- Download file 6.61 kB
-
https://huggingface.co/spaces/ZHIWEI666/ComfyUI-Ranking-API/resolve/main/notification_store.py
- Command line
-
hf download hf://spaces/ZHIWEI666/ComfyUI-Ranking-API/notification_store.py
-
curl -L -o notification_store.py https://huggingface.co/spaces/ZHIWEI666/ComfyUI-Ranking-API/resolve/main/notification_store.py
6.61 kB
| """通知与公告的严格读取、白名单输出和锁内更新。""" | |
| import copy | |
| import logging | |
| import math | |
| import time | |
| import uuid | |
| from fastapi import HTTPException | |
| import 数据库连接 as db | |
| logger = logging.getLogger(__name__) | |
| TTL = 7 * 24 * 3600 | |
| FIELDS = frozenset({"id", "type", "from_user", "from_name", "from_avatar", "target_item_id", | |
| "target_item_title", "content", "is_read", "created_at", "receiver"}) | |
| def _unavailable(): | |
| return HTTPException(status_code=503, detail="通知数据暂不可用") | |
| def _timestamp(message): | |
| value = message["created_at"] | |
| return value / 1000 if value >= 1e12 else value | |
| def _validate(message): | |
| if not isinstance(message, dict): | |
| raise _unavailable() | |
| if any(not isinstance(message.get(key), str) or not message[key] for key in ("id", "type")): | |
| raise _unavailable() | |
| timestamp = message.get("created_at") | |
| try: | |
| valid_timestamp = type(timestamp) in (int, float) and math.isfinite(timestamp) and timestamp >= 0 | |
| except (OverflowError, ValueError): | |
| valid_timestamp = False | |
| if not valid_timestamp: | |
| raise _unavailable() | |
| if not isinstance(message.get("content", ""), str) or type(message.get("is_read", False)) is not bool: | |
| raise _unavailable() | |
| return message | |
| def _messages(data, account): | |
| if not isinstance(data, dict) or not isinstance(data.get(account, []), list): | |
| raise _unavailable() | |
| messages = data.get(account, []) | |
| for message in messages: | |
| _validate(message) | |
| if "receiver" in message and message["receiver"] != account: | |
| raise _unavailable() | |
| return messages | |
| def _announcements(data): | |
| if not isinstance(data, list): | |
| raise _unavailable() | |
| for message in data: | |
| if _validate(message)["type"] != "system": | |
| raise _unavailable() | |
| return data | |
| def _read(name, default): | |
| try: | |
| return db.load_data(name, default_data=default, skip_cache=True, strict=True) | |
| except Exception as error: | |
| logger.warning("通知读取失败: %s", type(error).__name__) | |
| raise _unavailable() from None | |
| def _mutate(name, updater, default): | |
| try: | |
| return db.atomic_update(name, updater, default_data=default) | |
| except HTTPException: | |
| raise | |
| except Exception as error: | |
| logger.warning("通知写入失败: %s", type(error).__name__) | |
| raise _unavailable() from None | |
| def _profile(account): | |
| users = _read("users.json", {}) | |
| if not isinstance(users, dict) or not isinstance(users.get(account, {}), dict): | |
| raise _unavailable() | |
| return users.get(account, {}) | |
| def publish_announcement(account, content): | |
| if not isinstance(content, str) or not content.strip(): | |
| raise HTTPException(status_code=400, detail="公告内容不能为空") | |
| user = _profile(account) | |
| announcement = {"id": f"sys_{int(time.time())}_{uuid.uuid4().hex}", "type": "system", "from_user": account, | |
| "from_name": user.get("name", account), "from_avatar": user.get("avatarDataUrl", ""), | |
| "content": content, "created_at": int(time.time())} | |
| def update(data): | |
| _announcements(data).append(announcement) | |
| _mutate("announcements.json", update, []) | |
| def append_notification(account, payload): | |
| source = payload.get("from_user") | |
| if not account or source == account: | |
| return | |
| user = _profile(source) if source and source != "system" else {} | |
| message = {"id": f"msg_{int(time.time())}_{uuid.uuid4().hex}", "type": payload.get("type"), | |
| "from_user": source, "from_name": user.get("name", source) if source and source != "system" else "系统通知", | |
| "from_avatar": user.get("avatarDataUrl", ""), "target_item_id": payload.get("target_item_id", ""), | |
| "target_item_title": payload.get("target_item_title", ""), "content": payload.get("content", ""), | |
| "is_read": False, "created_at": int(time.time())} | |
| _validate(message) | |
| def update(data): | |
| messages = _messages(data, account) | |
| data[account] = [message, *messages] | |
| _mutate("messages.json", update, {}) | |
| def get_messages(account, count_only=False): | |
| now = time.time() | |
| announcements = _announcements(_read("announcements.json", [])) | |
| # 过期公告不重复注入已被清理的旧通知,也不在角标中重新出现。 | |
| active = [message for message in announcements if now - _timestamp(message) < TTL] | |
| messages = _messages(_read("messages.json", {}), account) | |
| def combine(messages): | |
| result = copy.deepcopy(messages) | |
| ids = {message["id"] for message in result} | |
| for message in active: | |
| if message["id"] not in ids: | |
| result.append({**copy.deepcopy(message), "receiver": account, "is_read": False}) | |
| ids.add(message["id"]) | |
| return result | |
| def unread(messages): | |
| return sum(not message.get("is_read", False) and message["type"] not in ("private", "plugin_update") | |
| and now - _timestamp(message) < TTL for message in messages) | |
| def project(messages): | |
| return [copy.deepcopy({key: value for key, value in message.items() if key in FIELDS}) for message in messages] | |
| combined = combine(messages) | |
| if count_only: | |
| return {"status": "success", "unread_count": unread(combined)} | |
| if len(combined) == len(messages) and all(message.get("is_read", False) and now - _timestamp(message) < TTL for message in messages): | |
| return {"status": "success", "data": project(messages), "unread_count": 0} | |
| def update(data): | |
| latest = combine(_messages(data, account)) | |
| count = unread(latest) | |
| kept = [message for message in latest if not message.get("is_read", False) or now - _timestamp(message) < TTL] | |
| for message in kept: | |
| message["is_read"] = True | |
| kept.sort(key=_timestamp, reverse=True) | |
| data[account] = kept | |
| return {"status": "success", "data": project(kept), "unread_count": count} | |
| return _mutate("messages.json", update, {}) | |
| def mark_messages_read(account): | |
| messages = _messages(_read("messages.json", {}), account) | |
| if all(message.get("is_read", False) for message in messages): | |
| return | |
| def update(data): | |
| for message in _messages(data, account): | |
| message["is_read"] = True | |
| _mutate("messages.json", update, {}) | |