重构消息存储逻辑,使用Peewee模型存储消息和撤回消息,添加时间戳处理
This commit is contained in:
@@ -1,9 +1,10 @@
|
||||
import re
|
||||
from typing import Union
|
||||
|
||||
from ...common.database.database import db
|
||||
# from ...common.database.database import db # db is now Peewee's SqliteDatabase instance
|
||||
from .message import MessageSending, MessageRecv
|
||||
from .chat_stream import ChatStream
|
||||
from ...common.database.database_model import Messages, RecalledMessages # Import Peewee models
|
||||
from src.common.logger import get_module_logger
|
||||
|
||||
logger = get_module_logger("message_storage")
|
||||
@@ -29,42 +30,65 @@ class MessageStorage:
|
||||
else:
|
||||
filtered_detailed_plain_text = ""
|
||||
|
||||
message_data = {
|
||||
"message_id": message.message_info.message_id,
|
||||
"time": message.message_info.time,
|
||||
"chat_id": chat_stream.stream_id,
|
||||
"chat_info": chat_stream.to_dict(),
|
||||
"user_info": message.message_info.user_info.to_dict(),
|
||||
# 使用过滤后的文本
|
||||
"processed_plain_text": filtered_processed_plain_text,
|
||||
"detailed_plain_text": filtered_detailed_plain_text,
|
||||
"memorized_times": message.memorized_times,
|
||||
}
|
||||
db.messages.insert_one(message_data)
|
||||
chat_info_dict = chat_stream.to_dict()
|
||||
user_info_dict = message.message_info.user_info.to_dict()
|
||||
|
||||
# Ensure message_id is an int if the model field is IntegerField
|
||||
try:
|
||||
msg_id = int(message.message_info.message_id)
|
||||
except ValueError:
|
||||
logger.error(f"Message ID {message.message_info.message_id} is not a valid integer. Storing as 0 or consider changing model field type.")
|
||||
msg_id = 0 # Or handle as appropriate, e.g. skip storing, or change model field to TextField
|
||||
|
||||
Messages.create(
|
||||
message_id=msg_id,
|
||||
time=float(message.message_info.time),
|
||||
chat_id=chat_stream.stream_id,
|
||||
# Flattened chat_info
|
||||
chat_info_stream_id=chat_info_dict.get("stream_id"),
|
||||
chat_info_platform=chat_info_dict.get("platform"),
|
||||
chat_info_user_platform=chat_info_dict.get("user_info", {}).get("platform"),
|
||||
chat_info_user_id=chat_info_dict.get("user_info", {}).get("user_id"),
|
||||
chat_info_user_nickname=chat_info_dict.get("user_info", {}).get("user_nickname"),
|
||||
chat_info_user_cardname=chat_info_dict.get("user_info", {}).get("user_cardname"),
|
||||
chat_info_group_platform=chat_info_dict.get("group_info", {}).get("platform"),
|
||||
chat_info_group_id=chat_info_dict.get("group_info", {}).get("group_id"),
|
||||
chat_info_group_name=chat_info_dict.get("group_info", {}).get("group_name"),
|
||||
chat_info_create_time=float(chat_info_dict.get("create_time", 0.0)),
|
||||
chat_info_last_active_time=float(chat_info_dict.get("last_active_time", 0.0)),
|
||||
# Flattened user_info (message sender)
|
||||
user_platform=user_info_dict.get("platform"),
|
||||
user_id=user_info_dict.get("user_id"),
|
||||
user_nickname=user_info_dict.get("user_nickname"),
|
||||
user_cardname=user_info_dict.get("user_cardname"),
|
||||
# Text content
|
||||
processed_plain_text=filtered_processed_plain_text,
|
||||
detailed_plain_text=filtered_detailed_plain_text,
|
||||
memorized_times=message.memorized_times,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("存储消息失败")
|
||||
|
||||
@staticmethod
|
||||
async def store_recalled_message(message_id: str, time: str, chat_stream: ChatStream) -> None:
|
||||
"""存储撤回消息到数据库"""
|
||||
if "recalled_messages" not in db.list_collection_names():
|
||||
db.create_collection("recalled_messages")
|
||||
else:
|
||||
try:
|
||||
message_data = {
|
||||
"message_id": message_id,
|
||||
"time": time,
|
||||
"stream_id": chat_stream.stream_id,
|
||||
}
|
||||
db.recalled_messages.insert_one(message_data)
|
||||
except Exception:
|
||||
logger.exception("存储撤回消息失败")
|
||||
# Table creation is handled by initialize_database in database_model.py
|
||||
try:
|
||||
RecalledMessages.create(
|
||||
message_id=message_id,
|
||||
time=float(time), # Assuming time is a string representing a float timestamp
|
||||
stream_id=chat_stream.stream_id,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("存储撤回消息失败")
|
||||
|
||||
@staticmethod
|
||||
async def remove_recalled_message(time: str) -> None:
|
||||
"""删除撤回消息"""
|
||||
try:
|
||||
db.recalled_messages.delete_many({"time": {"$lt": time - 300}})
|
||||
# Assuming input 'time' is a string timestamp that can be converted to float
|
||||
current_time_float = float(time)
|
||||
RecalledMessages.delete().where(RecalledMessages.time < (current_time_float - 300)).execute()
|
||||
except Exception:
|
||||
logger.exception("删除撤回消息失败")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user