From 4d06b97cf4e9878b580f646ac8ac3f5292a64114 Mon Sep 17 00:00:00 2001 From: xingc Date: Thu, 3 Apr 2025 17:27:05 +0800 Subject: [PATCH] =?UTF-8?q?feat(chat):=20=E8=81=8A=E5=A4=A9=E5=8A=9F?= =?UTF-8?q?=E8=83=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- core/routers/chat/handlers.py | 119 ++++++++++++++++++++++++++++++++++ core/routers/chat/managers.py | 116 +++++++++++++++++++++++++++++++++ core/routers/chat/schema.py | 60 +++++++++++++++++ core/routers/chat/services.py | 90 +++++++++++++++++++++++++ 4 files changed, 385 insertions(+) create mode 100644 core/routers/chat/handlers.py create mode 100644 core/routers/chat/managers.py create mode 100644 core/routers/chat/schema.py create mode 100644 core/routers/chat/services.py diff --git a/core/routers/chat/handlers.py b/core/routers/chat/handlers.py new file mode 100644 index 0000000..de318e5 --- /dev/null +++ b/core/routers/chat/handlers.py @@ -0,0 +1,119 @@ +# -*- coding = utf-8 -*- +# @Time : 2025/3/1 下午4:40 +# @File : handlers.py +# @Software : PyCharm +# @Author : xingc +# @Desc : +import logging +from datetime import datetime +from typing import List, Dict + +import redis.asyncio as aioredis +from fastapi import WebSocket +from sqlalchemy.ext.asyncio import AsyncSession + +from core.db.crud import CRUD +from core.db.models import ReadStatus, User, Chat, Message + +from .managers import SocketManager +from . import schema + +logger = logging.getLogger(__name__) +socket_manager = SocketManager() + + +@socket_manager.handler('chat.new_message') +async def send_new_message( + websocket: WebSocket, + db_session: AsyncSession, + cache: aioredis.Redis, + chats: dict, + current_user: User, + incoming_message: schema.Message, +): + chat_guid: str = str(incoming_message.chat_guid) + chat: Chat | None = await CRUD(db_session, Chat).get(filters={'chat_guid': chat_guid}) + # 判断是否为新聊天 + if not chats or chat_guid not in chats: + if chat: + chats[chat_guid] = chat.id + await socket_manager.add_user_to_chat(chat_guid, websocket) + else: + await socket_manager.send_error('聊天尚未添加', websocket) + return + + chat_id = chats.get(chat_guid) + try: + message = Message( + guid=incoming_message.message_guid, + user_id=current_user.user_id, + chat_id=chat_id, + message_type=incoming_message.message_type, + content=incoming_message.body, + ) + db_session.add(message) + await db_session.flush() + + chat.updated_at = datetime.now() + db_session.add(chat) + + await db_session.commit() + await db_session.refresh(message, attribute_names=["user", "chat"]) + await db_session.refresh(chat, attribute_names=["users"]) + + except Exception as exc_info: + await db_session.rollback() + logger.exception(f"[new_message] Exception, rolling back session, detail: {exc_info}") + raise exc_info + + send_message_schema = schema.MessageOut( + type='chat.new_message', + message_id=message.id, + message_guid=message.guid, + message_type=message.message_type, + chat_guid=chat.guid, + user_id=current_user.user_id, + content=message.content, + created_at=message.created_at, + friend_user=current_user, + ) + await socket_manager.broadcast_to_chat(chat_guid, send_message_schema.dict()) + + +@socket_manager.handler('chat.refresh_unread_messages') +async def refresh_unread_messages( + websocket: WebSocket, + db_session: AsyncSession, + cache: aioredis.Redis, + current_user: User, + message: schema.Message, +): + pass + + +@socket_manager.handler('chat.message_read') +async def read_message( + websocket: WebSocket, + db_session: AsyncSession, + cache: aioredis.Redis, + current_user: User, + message: schema.Message, +): + last_read_message_id: str | None = message.body + if not last_read_message_id: + return + + db_crud = CRUD(db_session, ReadStatus) + condition = { + 'user_id': current_user.user_id, + 'chat_guid': message.chat_guid, + } + result = await db_crud.update(filters=condition, data={'last_read_message_id': last_read_message_id}) + if result is None: + await db_crud.create( + data={ + **condition, + 'last_read_message_id': last_read_message_id + }, + filters=condition + ) diff --git a/core/routers/chat/managers.py b/core/routers/chat/managers.py new file mode 100644 index 0000000..f5b6743 --- /dev/null +++ b/core/routers/chat/managers.py @@ -0,0 +1,116 @@ +# -*- coding = utf-8 -*- +# @Time : 2025/3/1 下午8:47 +# @File : managers.py +# @Software : PyCharm +# @Author : xingc +# @Desc : +import asyncio +import logging + +import redis.asyncio as aioredis +import ujson +from fastapi import WebSocket + +from core.db.engine import get_cache + +logger = logging.getLogger(__name__) + + +class RedisPubSubManager: + def __init__(self): + self.redis_connection = None + self.pubsub = None + + async def _get_redis_connection(self) -> aioredis: + return await get_cache() + + async def connect(self): + self.redis_connection = await self._get_redis_connection() + self.pubsub = self.redis_connection.pubsub() + + async def subscribe(self, chat_guid: str) -> aioredis: + await self.pubsub.subscribe(chat_guid) + return self.pubsub + + async def unsubscribe(self, chat_guid: str): + await self.pubsub.unsubscribe(chat_guid) + + async def publish(self, chat_guid: str, message: str): + await self.redis_connection.publish(chat_guid, message) + + async def disconnect(self): + if self.redis_connection is not None: + await self.redis_connection.close() + + +class SocketManager: + + def __init__(self): + self.handlers: dict = {} + self.chats: dict = {} + self.pubsub_client = RedisPubSubManager() + self.user_guid_to_websocket: dict = {} + self._task_cancel = True + + def handler(self, message_type): + def decorator(func): + self.handlers[message_type] = func + return func + + return decorator + + async def connect_socket(self, websocket: WebSocket): + await websocket.accept() + + async def add_user_socket_connection(self, user_guid: str, websocket: WebSocket): + self.user_guid_to_websocket.setdefault(user_guid, set()).add(websocket) + + async def add_user_to_chat(self, chat_guid: str, websocket: WebSocket): + if chat_guid in self.chats: + self.chats[chat_guid].add(websocket) + else: + self.chats[chat_guid] = {websocket} + await self.pubsub_client.connect() + pubsub_subscriber = await self.pubsub_client.subscribe(chat_guid) + await asyncio.create_task(self._pubsub_data_reader(pubsub_subscriber)) + + async def broadcast_to_chat(self, chat_guid: str, message: str | dict) -> None: + if isinstance(message, dict): + message = ujson.dumps(message) + await self.pubsub_client.publish(chat_guid, message) + + async def remove_user_from_chat(self, chat_guid: str, websocket: WebSocket) -> None: + self.chats[chat_guid].remove(websocket) + if len(self.chats[chat_guid]) == 0: + del self.chats[chat_guid] + logger.info(f'Removing user from PubSub channel {chat_guid}') + await self.pubsub_client.unsubscribe(chat_guid) + + async def remove_user_guid_to_websocket(self, user_guid: str, websocket: WebSocket): + if user_guid in self.user_guid_to_websocket: + self.user_guid_to_websocket.get(user_guid).remove(websocket) + + # https://github.com/redis/redis-py/issues/2523 + async def _pubsub_data_reader(self, pubsub_subscriber): + try: + while True: + message = await pubsub_subscriber.get_message(ignore_subscribe_messages=True) + if message is not None: + chat_guid = message['channel'] + sockets = self.chats.get(chat_guid) + if sockets: + for socket in sockets: + data = message['data'] + await socket.send_text(data) + else: + if self._task_cancel is False: + break + except Exception as exc: + logger.exception(f'Exception occurred: {exc}') + + async def send_error(self, message: str, websocket: WebSocket): + await websocket.send_json({'status': 'error', 'message': message}) + + async def quit(self): + self._task_cancel = False + await self.pubsub_client.disconnect() diff --git a/core/routers/chat/schema.py b/core/routers/chat/schema.py new file mode 100644 index 0000000..52c3dcb --- /dev/null +++ b/core/routers/chat/schema.py @@ -0,0 +1,60 @@ +# -*- coding = utf-8 -*- +# @Time : 2025/3/8 下午10:54 +# @File : schema.py +# @Software : PyCharm +# @Author : xingc +# @Desc : +from datetime import datetime +from typing import Literal, Optional, List + +from pydantic import BaseModel, conint, UUID4 + +from core.base.schema import MessageType +from core.routers.user.schema import LiteUser + + +class ChatUser(LiteUser): + pass + + +class Message(BaseModel): + type: Literal[ + 'chat.new_message', 'chat.message_read' + ] + message_guid: UUID4 + message_type: MessageType + chat_guid: UUID4 + receive_user: Optional[int] = None + body: str + + +class MessageOut(BaseModel): + type: Literal[ + 'chat.new_message', 'chat.send_status_reply' + ] + code: Optional[int] = 200 + status: Optional[Literal['ok', 'error']] = 'ok' + message_id: int + message_guid: UUID4 + message_type: MessageType + chat_guid: UUID4 + body: Optional[str] = None + friend_user: LiteUser + + +class Chat(BaseModel): + chat_guid: UUID4 + friend_user: LiteUser + + +class ChatCreateIn(BaseModel): + recipient_user_id: int + + +class ChatUnreadOut(BaseModel): + chat_guid: UUID4 + created_at: str + updated_at: str + recipient_user: LiteUser + last_message: str | None + unread_messages_count: int diff --git a/core/routers/chat/services.py b/core/routers/chat/services.py new file mode 100644 index 0000000..b388fbe --- /dev/null +++ b/core/routers/chat/services.py @@ -0,0 +1,90 @@ +# -*- coding = utf-8 -*- +# @Time : 2025/3/1 下午4:44 +# @File : services.py +# @Software : PyCharm +# @Author : xingc +# @Desc : +from typing import Dict + +from sqlalchemy import select, and_, func +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.orm import aliased, selectinload + +from core.base.exceptions import WebsocketTooManyRequests +from core.db.models import User, Chat, ReadStatus, Message +from core.routers.auth.schema import LoginUser +from core.utils.helper import format_date + +from .schema import ChatUnreadOut + + +async def websocket_callback(ws, pexpire): + raise WebsocketTooManyRequests("Too many requests") + + +async def get_user_active_direct_chats( + db_session: AsyncSession, current_user: User +) -> Dict[str, int] | None: + """获取活跃聊天""" + direct_chats_dict = dict() + query = ( + select(Chat) + .where(and_(Chat.is_deleted.is_(False), Chat.users.contains(current_user))) + .options(selectinload(Chat.users)) + ) + result = await db_session.execute(query) + direct_chats: list[Chat] = result.scalars().all() + + if direct_chats: + for direct_chat in direct_chats: + direct_chats_dict[str(direct_chat.guid)] = direct_chat.id + return direct_chats_dict + + +async def get_unread_messages_per_chat( + db_session: AsyncSession, chats: list[Chat], login_user: LoginUser +) -> list[dict]: + """获取聊天未读消息数""" + unread_messages_info_per_chat = {chat.id: {'message_count': 0} for chat in chats} + + # 查询消息情况 + read_status_alias = aliased(ReadStatus) + query = ( + select(Message.chat_id, Message.content, func.count().label("message_count")) + .join( + read_status_alias, + and_(read_status_alias.user_id == login_user.user_id, read_status_alias.chat_id == Message.chat_id), + ) + .where( + and_( + Message.user_id != login_user.user_id, + Message.id > func.coalesce(read_status_alias.last_read_message_id, 0), + Message.chat_id.in_(unread_messages_info_per_chat), + ) + ) + .order_by(Message.created_at.desc()) + .group_by(Message.chat_id) + ) + result = await db_session.execute(query) + unread_messages_info = result.fetchall() + + for messages_info in unread_messages_info: + unread_messages_info_per_chat[messages_info[0]] = { + 'message': messages_info[1], + 'message_count': messages_info[2], + } + + results = [] + for chat in chats: + users = [user for user in await chat.awaitable_attrs.users if user.user_id != login_user.user_id] + messages_info = unread_messages_info_per_chat[chat.id] + results.append(dict( + chat_guid=chat.guid, + created_at=format_date(chat.created_at), + updated_at=format_date(chat.updated_at), + recipient_user=users[0], + last_message=messages_info.get('message'), + unread_messages_count=messages_info.get('message_count'), + )) + + return results