ComfyUI-Ranking-API / notification_store.py
ZHIWEI666's picture
3.0更新
c8e5eaa verified
Raw History Blame Contribute Delete
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, {})