diff --git a/.vscode/settings.json b/.vscode/settings.json index 58bfe96..b2d8246 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -13,6 +13,7 @@ "htmlrender", "localstore", "parameterless", + "postprocessor", "pyecharts", "pygal", "sqlalchemy", diff --git a/README.md b/README.md index e0148e9..35d76bd 100644 --- a/README.md +++ b/README.md @@ -158,12 +158,6 @@ __！！注意！！__ - [x] 私聊的查询（超级用户可以任意查询群聊的信息）一半完成 -- [ ] 提供多样化的渲染器配置（html 渲染，pillow 渲染，统计绘图软件渲染） - -- [ ] 使用管理员权限直接获取 QQ 官方统计的今日消息量以优化代码运行速度 - -- [ ] 为 pillow 渲染方式提供插件的加载方式（什么？插件里的插件？？？） - - [ ] 特殊的储存方案优化消息统计 - [ ] 查询带某关键词的消息量 diff --git a/nonebot_plugin_dialectlist/__init__.py b/nonebot_plugin_dialectlist/__init__.py index 1e2566c..af36671 100644 --- a/nonebot_plugin_dialectlist/__init__.py +++ b/nonebot_plugin_dialectlist/__init__.py @@ -16,6 +16,7 @@ from datetime import datetime, timedelta from arclet.alconna import ArparmaBehavior from arclet.alconna.arparma import Arparma +from nonebot import on_command from nonebot.log import logger from nonebot.typing import T_State from nonebot.params import Arg, Depends @@ -30,8 +31,7 @@ from nonebot_plugin_alconna import ( from nonebot_plugin_chatrecorder import get_message_records from nonebot_plugin_session import Session, SessionIdType, extract_session -# from . import migrations #抄词云的部分代码，还不知道这有什么用 -# from .function import * +from .storage import get_cache,build_cache from .config import Config, plugin_config from .usage import __usage__ from .time import ( @@ -45,7 +45,6 @@ from .utils import ( get_rank_image, persist_id2user_id, get_user_infos, - # get_user_info2, ) __plugin_meta__ = PluginMetadata( @@ -58,7 +57,6 @@ __plugin_meta__ = PluginMetadata( "nonebot_plugin_chatrecorder", "nonebot_plugin_saa", "nonebot_plugin_alconna" ), config=Config, - # extra={"orm_version_location": migrations}, ) @@ -76,6 +74,13 @@ def wrapper(slot: Union[int, str], content: Optional[str], context) -> str: return content return "" # pragma: no cover +build_cache_cmd = on_command("build_cache", aliases={"重建缓存"}, block=True) + +@build_cache_cmd.handle() +async def _build_cache(bot: Bot, event: Event): + await saa.Text("正在重建缓存，请稍等。").send(reply=True) + await build_cache() + await saa.Text("重建缓存完成。").send(reply=True) rank_cmd = on_alconna( Alconna( @@ -187,8 +192,6 @@ async def _group_message( except ValueError: await rank_cmd.finish("请输入正确的日期，不然我没法理解呢！") - logger.debug(f"命令解析花费时间:{t.time() - t1}") - @rank_cmd.got( "start", @@ -221,8 +224,9 @@ async def handle_rank( await saa.Text("没有指定群哦").finish() if plugin_config.counting_cache: - raise Exception("我草缓存功能还没端上来呢，你怎么就先用上了") + raw_rank = await get_cache(start, stop, id) else: + t1 = t.time() messages = await get_message_records( id2s=[id], id_type=SessionIdType.GROUP, @@ -233,12 +237,12 @@ async def handle_rank( time_stop=stop, exclude_id1s=plugin_config.excluded_people, ) - - if not messages: - await saa.Text("明明这个时间段都没有人说话怎么会有话痨榜呢？").finish() - + logger.debug(f"获取计数消息花费时间:{t.time() - t1}") raw_rank = msg_counter(messages) + if not raw_rank: + await saa.Text("明明这个时间段都没有人说话怎么会有话痨榜呢？").finish() + rank = got_rank(raw_rank) logger.debug(rank) ids = await persist_id2user_id([int(i[0]) for i in rank]) @@ -276,10 +280,3 @@ async def handle_rank( logger.debug(f"群聊消息渲染图片花费时间:{t.time() - t1}") await msg.finish(reply=True) - - -# @scheduler.scheduled_job( -# "dialectlist", day="*/2", id="xxx", args=[1], kwargs={"arg2": 2} -# ) -# async def __(): -# pass diff --git a/nonebot_plugin_dialectlist/config.py b/nonebot_plugin_dialectlist/config.py index 90c2d1d..69703c4 100644 --- a/nonebot_plugin_dialectlist/config.py +++ b/nonebot_plugin_dialectlist/config.py @@ -7,11 +7,11 @@ class ScopedConfig(BaseModel): get_num: int = 5 # 获取人数数量 font: str = "SimHei" # 字体格式 suffix: bool = False # 是否显示后缀 - excluded_self: bool = True + excluded_self: bool = True # 是否排除自己 visualization: bool = True # 是否可视化 - counting_cache: bool = False # 计数缓存（没有完工） + counting_cache: bool = True # 计数缓存(能够提高回复速度) excluded_people: List[str] = [] # 排除的人的QQ号 - timezone: Optional[str] = "Asia/Shanghai" + timezone: Optional[str] = "Asia/Shanghai" # 时区，影响统计时间 string_suffix: str = "统计花费时间:{timecost}秒" # 消息格式后缀 template_path: str = "./template/rank_template.j2" # 模板路径 string_format: str = "第{index}名：\n{nickname},{chatdatanum}条消息\n" # 消息格式 diff --git a/nonebot_plugin_dialectlist/model.py b/nonebot_plugin_dialectlist/model.py index 2848ed5..17e69e3 100644 --- a/nonebot_plugin_dialectlist/model.py +++ b/nonebot_plugin_dialectlist/model.py @@ -14,14 +14,9 @@ class UserRankInfo(UserInfo): user_avatar_bytes: bytes -# class MsgCountDayData(BaseModel): -# session_id: str -# session_bnum: int - - class MessageCountCache(Model): __table_args__ = {"extend_existing": True} - id: Mapped[int] = mapped_column(primary_key=True) + id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) time: Mapped[datetime] - session_id: Mapped[int] = mapped_column(Integer) + session_id: Mapped[int] = mapped_column(Integer, index=True) session_bnum: Mapped[int] = mapped_column(Integer) diff --git a/nonebot_plugin_dialectlist/storage.py b/nonebot_plugin_dialectlist/storage.py index 020215f..62225dc 100644 --- a/nonebot_plugin_dialectlist/storage.py +++ b/nonebot_plugin_dialectlist/storage.py @@ -1,2 +1,129 @@ -# TODO 使用计数缓存进行数据库查询优化，避免一次性查询过多消息导致内存爆炸。 -from nonebot_plugin_orm import Model +import os +import json +from datetime import datetime +from sqlalchemy import delete, or_, select + +from nonebot import get_driver +from nonebot.log import logger +from nonebot.params import Depends +from nonebot.adapters import Event,Bot +from nonebot.message import event_postprocessor + +from .model import MessageCountCache +from .config import plugin_config + +from nonebot_plugin_localstore import get_data_file +from nonebot_plugin_chatrecorder import get_message_records +from nonebot_plugin_chatrecorder.utils import remove_timezone +from nonebot_plugin_session import extract_session, Session +from nonebot_plugin_session_orm import SessionModel, get_session_persist_id +from nonebot_plugin_orm import get_session + + +async def get_cache(time_start: datetime, time_stop: datetime, group_id: str): + async with get_session() as db_session: + where = [or_(SessionModel.id2 == group_id)] + statement = select(SessionModel).where(*where) + + sessions = (await db_session.scalars(statement)).all() + + where = [ + or_(MessageCountCache.session_id == session.id) for session in sessions + ] + where.append(or_(MessageCountCache.time >= remove_timezone(time_start))) + where.append(or_(MessageCountCache.time <= remove_timezone(time_stop))) + statement = select(MessageCountCache).where(*where) + + user_caches = (await db_session.scalars(statement)).all() + raw_rank = {} + for i in user_caches: + raw_rank[i.session_id] = i.session_bnum + return raw_rank + + +async def build_cache(): + async with get_session() as db_session: + await db_session.execute(delete(MessageCountCache)) + await db_session.commit() + logger.info("先前可能存在的缓存已清空") + messages = await get_message_records(types=["message"]) + async with get_session() as db_session: + for msg in messages: + msg_session_id = msg.session_persist_id + + where = [or_(MessageCountCache.session_id == msg_session_id)] + where.append( + or_( + MessageCountCache.time + == remove_timezone(msg.time.replace(hour=1, minute=0, second=0, microsecond=0)) + ) + ) + statement = select(MessageCountCache).where(*where) + + user_cache = (await db_session.scalars(statement)).all() + + if user_cache: + user_cache[0].session_bnum += 1 + else: + user_cache = MessageCountCache( + session_id=msg.session_persist_id, + time=remove_timezone(msg.time.replace(hour=1, minute=0, second=0, microsecond=0)), + session_bnum=1, + ) + db_session.add(user_cache) + await db_session.commit() + + logger.info("缓存构建完成") + + +driver = get_driver() + + +@driver.on_startup +async def _(): + if not plugin_config.counting_cache: + return + f_name = get_data_file("nonebot-plugin-dialectlist", "is-pre-cached.json") + if not os.path.exists(f_name): + with open(f_name, "w", encoding="utf-8") as f: + s = json.dumps({"is-pre-cached": False}, ensure_ascii=False, indent=4) + f.write(s) + + with open(f_name, "r", encoding="utf-8") as f: + if json.load(f)["is-pre-cached"]: + return + logger.info("未检查到缓存，开始构建缓存") + with open(f_name, "w", encoding="utf-8") as f: + await build_cache() + json.dump({"is-pre-cached": True}, f, ensure_ascii=False, indent=4) + + +@event_postprocessor +async def _(bot: Bot, event: Event,session: Session = Depends(extract_session)): + if not plugin_config.counting_cache: + return + if not session.id2: + return + if event.get_type() != "message": + return + now = datetime.now() + now = now.replace(hour=1, minute=0, second=0, microsecond=0) + + async with get_session() as db_session: + session_id = await get_session_persist_id(session) + where = [or_(MessageCountCache.session_id == session_id)] + where = [or_(MessageCountCache.time == remove_timezone(now))] + statement = select(MessageCountCache).where(*where) + user_cache = (await db_session.scalars(statement)).all() + + if user_cache: + user_cache[0].session_bnum += 1 + else: + user_cache = MessageCountCache( + session_id=session_id, + time=remove_timezone(now), + session_bnum=1, + ) + db_session.add(user_cache) + await db_session.commit() + logger.debug("已计入缓存") diff --git a/nonebot_plugin_dialectlist/utils.py b/nonebot_plugin_dialectlist/utils.py index e182951..453b8dc 100644 --- a/nonebot_plugin_dialectlist/utils.py +++ b/nonebot_plugin_dialectlist/utils.py @@ -26,16 +26,6 @@ from .config import plugin_config cache_path = get_cache_dir("nonebot_plugin_dialectlist") -# 暂时不做考虑 -# def admin_permission(): -# permission = SUPERUSER -# with contextlib.suppress(ImportError): -# from nonebot.adapters.onebot.v11.permission import GROUP_ADMIN, GROUP_OWNER - -# permission = permission | GROUP_ADMIN | GROUP_OWNER - -# return permission - async def ensure_group(matcher: Matcher, session: Session = Depends(extract_session)): """确保在群组中使用"""