feat: 保存消息到数据库，3天定时删除

Author: HenryXiaoYang

database/messsagDB.py

@@ -1,11 +1,11 @@
 import asyncio
 import logging
 import tomllib
-from datetime import datetime
+from datetime import datetime, timedelta
 from typing import Optional, List
 
 from pydantic import validate_arguments
-from sqlalchemy import Column, String, Integer, DateTime, Text, Boolean
+from sqlalchemy import Column, String, Integer, DateTime, Text, Boolean, delete
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_scoped_session
 from sqlalchemy.orm import declarative_base, sessionmaker
@@ -20,7 +20,7 @@ class Message(DeclarativeBase):
     __tablename__ = 'messages'
 
     id = Column(Integer, primary_key=True, autoincrement=True)
-    msg_id = Column(Integer, unique=True, index=True, comment='消息唯一ID（整型）')
+    msg_id = Column(Integer, index=True, comment='消息唯一ID（整型）')
     sender_wxid = Column(String(40), index=True, comment='消息发送人wxid')
     from_wxid = Column(String(40), index=True, comment='消息来源wxid')
     msg_type = Column(Integer, comment='消息类型（整型编码）')
@@ -123,7 +123,26 @@ async def close(self):
         """关闭数据库连接"""
         await self.engine.dispose()
 
+    async def cleanup_messages(self):
+        """每三天清理旧消息"""
+        while True:
+            async with self._async_session_factory() as session:
+                try:
+                    # 计算三天前的时间
+                    three_days_ago = datetime.now() - timedelta(days=3)
+                    # 删除三天前的消息
+                    await session.execute(
+                        delete(Message).where(Message.timestamp < three_days_ago)
+                    )
+                    await session.commit()
+                except Exception as e:
+                    logging.error(f"清理消息失败: {str(e)}")
+                    await session.rollback()
+            await asyncio.sleep(259200)  # 每三天（259200秒）执行一次
+
     async def __aenter__(self):
+        # 启动清理消息的定时任务
+        asyncio.create_task(self.cleanup_messages())
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):

requirements.txt

@@ -18,4 +18,5 @@ pymediainfo~=7.0.1
 matplotlib~=3.10.0
 numpy~=1.26.4
 requests~=2.32.3
-pillow~=10.4.0
+pillow~=10.4.0
+pydantic~=2.10.5

utils/xybot.py

@@ -6,6 +6,7 @@
 
 from WechatAPI import WechatAPIClient
 from WechatAPI.Client.protect import protector
+from database.messsagDB import MessageDB
 from utils.event_manager import EventManager
 
 
@@ -26,6 +27,8 @@ def __init__(self, bot_client: WechatAPIClient):
         self.whitelist = main_config.get("XYBot", {}).get("whitelist", [])
         self.blacklist = main_config.get("XYBot", {}).get("blacklist", [])
 
+        self.msg_db = MessageDB()
+
 
     def update_profile(self, wxid: str, nickname: str, alias: str, phone: str):
         """更新机器人信息"""
@@ -117,6 +120,16 @@ async def process_text_message(self, message: Dict[str, Any]):
             ats = []
         message["Ats"] = ats if ats and ats[0] != "" else []
 
+        # 保存消息到数据库
+        await self.msg_db.save_message(
+            msg_id=int(message["MsgId"]),
+            sender_wxid=message["SenderWxid"],
+            from_wxid=message["FromWxid"],
+            msg_type=int(message["MsgType"]),
+            content=message["Content"],
+            is_group=message["IsGroup"]
+        )
+
         if self.wxid in ats:
             logger.info("收到被@消息: 消息ID:{} 来自:{} 发送人:{} @:{} 内容:{}",
                         message["MsgId"],
@@ -171,6 +184,15 @@ async def process_image_message(self, message: Dict[str, Any]):
                     message["SenderWxid"],
                     message["Content"])
 
+        await self.msg_db.save_message(
+            msg_id=int(message["MsgId"]),
+            sender_wxid=message["SenderWxid"],
+            from_wxid=message["FromWxid"],
+            msg_type=int(message["MsgType"]),
+            content=message["MsgSource"],
+            is_group=message["IsGroup"]
+        )
+
         # 解析图片消息
         aeskey, cdnmidimgurl = None, None
         try:
@@ -219,6 +241,15 @@ async def process_voice_message(self, message: Dict[str, Any]):
                     message["SenderWxid"],
                     message["Content"])
 
+        await self.msg_db.save_message(
+            msg_id=int(message["MsgId"]),
+            sender_wxid=message["SenderWxid"],
+            from_wxid=message["FromWxid"],
+            msg_type=int(message["MsgType"]),
+            content=message["Content"],
+            is_group=message["IsGroup"]
+        )
+
         if message["IsGroup"] or not message.get("ImgBuf", {}).get("buffer", ""):
             # 解析语音消息
             voiceurl, length = None, None
@@ -265,6 +296,15 @@ async def process_xml_message(self, message: Dict[str, Any]):
                 message["FromWxid"] = message["ToWxid"]
             message["IsGroup"] = False
 
+        await self.msg_db.save_message(
+            msg_id=int(message["MsgId"]),
+            sender_wxid=message["SenderWxid"],
+            from_wxid=message["FromWxid"],
+            msg_type=int(message["MsgType"]),
+            content=message["Content"],
+            is_group=message["IsGroup"]
+        )
+
         try:
             root = ET.fromstring(message["Content"])
             type = int(root.find("appmsg").find("type").text)
@@ -413,6 +453,15 @@ async def process_video_message(self, message):
                     message["SenderWxid"],
                     str(message["Content"]).replace("\n", ""))
 
+        await self.msg_db.save_message(
+            msg_id=int(message["MsgId"]),
+            sender_wxid=message["SenderWxid"],
+            from_wxid=message["FromWxid"],
+            msg_type=int(message["MsgType"]),
+            content=message["Content"],
+            is_group=message["IsGroup"]
+        )
+
         message["Video"] = await self.bot.download_video(message["MsgId"])
 
         if self.ignore_check(message["FromWxid"], message["SenderWxid"]):
@@ -441,6 +490,15 @@ async def process_file_message(self, message: Dict[str, Any]):
                     message["SenderWxid"],
                     message["Content"])
 
+        await self.msg_db.save_message(
+            msg_id=int(message["MsgId"]),
+            sender_wxid=message["SenderWxid"],
+            from_wxid=message["FromWxid"],
+            msg_type=int(message["MsgType"]),
+            content=message["Content"],
+            is_group=message["IsGroup"]
+        )
+
         message["File"] = await self.bot.download_attach(attach_id)
 
         if self.ignore_check(message["FromWxid"], message["SenderWxid"]):
@@ -507,6 +565,15 @@ async def process_pat_message(self, message: Dict[str, Any]):
                     message["Patted"],
                     message["PatSuffix"])
 
+        await self.msg_db.save_message(
+            msg_id=int(message["MsgId"]),
+            sender_wxid=message["SenderWxid"],
+            from_wxid=message["FromWxid"],
+            msg_type=int(message["MsgType"]),
+            content=f"{message['Patter']} 拍了拍 {message['Patted']} {message['PatSuffix']}",
+            is_group=message["IsGroup"]
+        )
+
         if self.ignore_check(message["FromWxid"], message["SenderWxid"]):
             if self.ignore_protection or not protector.check(14400):
                 await EventManager.emit("pat_message", self.bot, message)