refactor: 移除MongoDBMessageStorage类中的self.db引用,直接使用db

This commit is contained in:
春河晴
2025-04-16 15:04:00 +09:00
parent 9fc74cb066
commit 1aad7f4f6e

View File

@@ -50,21 +50,18 @@ class MessageStorage(ABC):
class MongoDBMessageStorage(MessageStorage): class MongoDBMessageStorage(MessageStorage):
"""MongoDB消息存储实现""" """MongoDB消息存储实现"""
def __init__(self):
self.db = db
async def get_messages_after(self, chat_id: str, message_time: float) -> List[Dict[str, Any]]: async def get_messages_after(self, chat_id: str, message_time: float) -> List[Dict[str, Any]]:
query = {"chat_id": chat_id} query = {"chat_id": chat_id}
# print(f"storage_check_message: {message_time}") # print(f"storage_check_message: {message_time}")
query["time"] = {"$gt": message_time} query["time"] = {"$gt": message_time}
return list(self.db.messages.find(query).sort("time", 1)) return list(db.messages.find(query).sort("time", 1))
async def get_messages_before(self, chat_id: str, time_point: float, limit: int = 5) -> List[Dict[str, Any]]: async def get_messages_before(self, chat_id: str, time_point: float, limit: int = 5) -> List[Dict[str, Any]]:
query = {"chat_id": chat_id, "time": {"$lt": time_point}} query = {"chat_id": chat_id, "time": {"$lt": time_point}}
messages = list(self.db.messages.find(query).sort("time", -1).limit(limit)) messages = list(db.messages.find(query).sort("time", -1).limit(limit))
# 将消息按时间正序排列 # 将消息按时间正序排列
messages.reverse() messages.reverse()
@@ -73,7 +70,7 @@ class MongoDBMessageStorage(MessageStorage):
async def has_new_messages(self, chat_id: str, after_time: float) -> bool: async def has_new_messages(self, chat_id: str, after_time: float) -> bool:
query = {"chat_id": chat_id, "time": {"$gt": after_time}} query = {"chat_id": chat_id, "time": {"$gt": after_time}}
return self.db.messages.find_one(query) is not None return db.messages.find_one(query) is not None
# # 创建一个内存消息存储实现,用于测试 # # 创建一个内存消息存储实现,用于测试