From e3f28ad67c1c9ce7f3e19d30fc4ee10db101f14f Mon Sep 17 00:00:00 2001 From: xingc Date: Wed, 16 Apr 2025 14:44:52 +0800 Subject: [PATCH] =?UTF-8?q?feat(chat):=20=E5=AE=8C=E5=96=84=E8=81=8A?= =?UTF-8?q?=E5=A4=A9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- core/routers/chat/handlers.py | 59 ++++++++++++++++-------------- core/routers/chat/managers.py | 8 +++-- core/routers/chat/routes.py | 45 +++++++++++------------ core/routers/chat/schema.py | 68 +++++++++++++++++++++-------------- core/routers/chat/services.py | 9 +++-- 5 files changed, 110 insertions(+), 79 deletions(-) diff --git a/core/routers/chat/handlers.py b/core/routers/chat/handlers.py index de318e5..c07e37d 100644 --- a/core/routers/chat/handlers.py +++ b/core/routers/chat/handlers.py @@ -10,6 +10,7 @@ from typing import List, Dict import redis.asyncio as aioredis from fastapi import WebSocket +from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from core.db.crud import CRUD @@ -26,19 +27,20 @@ socket_manager = SocketManager() async def send_new_message( websocket: WebSocket, db_session: AsyncSession, - cache: aioredis.Redis, chats: dict, current_user: User, - incoming_message: schema.Message, + incoming_message: schema.WSMessage, ): chat_guid: str = str(incoming_message.chat_guid) - chat: Chat | None = await CRUD(db_session, Chat).get(filters={'chat_guid': chat_guid}) + chat: Chat | None = await CRUD(db_session, Chat).get(filters={'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: + logger.warning(f'chat {chat_guid} not found') await socket_manager.send_error('聊天尚未添加', websocket) return @@ -49,64 +51,59 @@ async def send_new_message( user_id=current_user.user_id, chat_id=chat_id, message_type=incoming_message.message_type, - content=incoming_message.body, + content=incoming_message.content, ) db_session.add(message) - await db_session.flush() chat.updated_at = datetime.now() db_session.add(chat) + await db_session.flush() await db_session.commit() - await db_session.refresh(message, attribute_names=["user", "chat"]) - await db_session.refresh(chat, attribute_names=["users"]) + # 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( + send_message_schema = schema.WSMessageOut( 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, + friend_user=current_user.dict(), + ) + send_status_replay = schema.WSMessageOut( + type='chat.send_status_reply', + message_id=message.id, + message_guid=message.guid, + message_type=message.message_type, + chat_guid=chat.guid, ) - 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 + await socket_manager.broadcast_to_chat(chat_guid, send_message_schema.model_dump_json()) + await websocket.send_json(send_status_replay.model_dump_json()) @socket_manager.handler('chat.message_read') async def read_message( websocket: WebSocket, db_session: AsyncSession, - cache: aioredis.Redis, current_user: User, - message: schema.Message, + incoming_message: schema.WSMessage, ): - last_read_message_id: str | None = message.body + last_read_message_id: str | None = incoming_message.content 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, + 'chat_guid': incoming_message.chat_guid, } result = await db_crud.update(filters=condition, data={'last_read_message_id': last_read_message_id}) if result is None: @@ -117,3 +114,13 @@ async def read_message( }, filters=condition ) + + +@socket_manager.handler('chat.ping') +async def ping( + websocket: WebSocket, + db_session: AsyncSession, + current_user: User, + incoming_message: schema.WSMessage, +): + await websocket.send_json({'type': 'chat.send_status_reply', 'status': 'ok'}) diff --git a/core/routers/chat/managers.py b/core/routers/chat/managers.py index f5b6743..0227bff 100644 --- a/core/routers/chat/managers.py +++ b/core/routers/chat/managers.py @@ -10,6 +10,7 @@ import logging import redis.asyncio as aioredis import ujson from fastapi import WebSocket +from fastapi.websockets import WebSocketState from core.db.engine import get_cache @@ -72,7 +73,7 @@ class SocketManager: 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)) + 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): @@ -100,6 +101,8 @@ class SocketManager: sockets = self.chats.get(chat_guid) if sockets: for socket in sockets: + if socket.client_state != WebSocketState.CONNECTED: + continue data = message['data'] await socket.send_text(data) else: @@ -109,7 +112,8 @@ class SocketManager: logger.exception(f'Exception occurred: {exc}') async def send_error(self, message: str, websocket: WebSocket): - await websocket.send_json({'status': 'error', 'message': message}) + if websocket.client_state == WebSocketState.CONNECTED: + await websocket.send_json({'status': 'error', 'content': message}) async def quit(self): self._task_cancel = False diff --git a/core/routers/chat/routes.py b/core/routers/chat/routes.py index 3cf7a18..036cc37 100644 --- a/core/routers/chat/routes.py +++ b/core/routers/chat/routes.py @@ -20,7 +20,7 @@ from core.db.crud import CRUD from core.db.engine import get_async_session, get_cache from core.db.models import Message, Chat, User, ReadStatus, chat_participant from core.routers.auth.schema import LoginUser -from core.routers.auth.services import get_current_user +from core.routers.auth.services import get_current_user_form_http, get_current_user_form_ws from core.base.exceptions import WebsocketTooManyRequests from . import schema @@ -41,7 +41,7 @@ chat_router = APIRouter(prefix='/chat', tags=['聊天模块']) ] ) async def get_all_chats( - login_user: LoginUser = Depends(get_current_user), + login_user: LoginUser = Depends(get_current_user_form_http), db_session: AsyncSession = Depends(get_async_session), page: int = Query(default=1, gt=0), page_size: int = Query(default=15, gt=0, le=200), @@ -82,13 +82,13 @@ async def get_all_chats( '/{chat_guid}/messages/', summary='获取单个聊天的历史消息', response_model=base_schema.BaseResponse[ base_schema.Paginated[ - List[Dict] + List[schema.MessageOut] ] ] ) async def get_single_chat_messages( chat_guid: str, - login_user: LoginUser = Depends(get_current_user), + login_user: LoginUser = Depends(get_current_user_form_http), db_session: AsyncSession = Depends(get_async_session), page: int = Query(default=1, gt=0), page_size: int = Query(default=15, gt=0, le=200), @@ -125,15 +125,18 @@ async def get_single_chat_messages( .options(selectinload(Message.user), selectinload(Message.chat)) ) records = await db_session.execute(stmt) + results = [ + { + 'user': record[0].user.dict(), + 'chat_guid': record[0].chat.guid, + 'message_guid': record[0].guid, + **record[0].dict(exclude={'guid'}) + } for record in records + ] return { 'data': { - 'list': [ - { - **card.dict(), - **user_collection.dict() - } for card, user_collection in records - ], + 'list': results, 'total': total.scalar(), 'page': page, 'page_size': page_size, @@ -144,7 +147,7 @@ async def get_single_chat_messages( @chat_router.post('/create', summary='创建新聊天', response_model=base_schema.BaseResponse[schema.Chat]) async def create_new_chat( payload: schema.ChatCreateIn, - login_user: LoginUser = Depends(get_current_user), + login_user: LoginUser = Depends(get_current_user_form_http), db_session: AsyncSession = Depends(get_async_session), ): if payload.recipient_user_id == login_user.user_id: @@ -201,7 +204,7 @@ async def create_new_chat( @chat_router.delete('/{chat_guid}/', summary='删除聊天', response_model=base_schema.BaseResponse) async def delete_single_chat( chat_guid: str, - login_user: LoginUser = Depends(get_current_user), + login_user: LoginUser = Depends(get_current_user_form_http), db_session: AsyncSession = Depends(get_async_session), ): chat = await CRUD(db_session, Chat).update( @@ -217,17 +220,16 @@ async def delete_single_chat( @chat_router.websocket('/ws') async def get_chat_websocket( websocket: WebSocket, - login_user: LoginUser = Depends(get_current_user), + login_user: LoginUser = Depends(get_current_user_form_ws), db_session: AsyncSession = Depends(get_async_session), - redis_conn: aioredis.Redis = Depends(get_cache) ): - current_user = await CRUD(db_session, User).get(filters={'user_id': login_user.user_id}) await socket_manager.connect_socket(websocket=websocket) - logger.info('Websocket connection is established') - ratelimit = WebSocketRateLimiter(times=50, seconds=10, callback=websocket_callback) + logger.info(f'Websocket connection is established: {login_user.user_id}') user_id = str(login_user.user_id) await socket_manager.add_user_socket_connection(user_id, websocket) + ratelimit = WebSocketRateLimiter(times=50, seconds=10, callback=websocket_callback) + current_user = await CRUD(db_session, User).get(filters={'user_id': login_user.user_id}) if chats := await get_user_active_direct_chats(db_session, current_user=current_user): for chat_guid in chats: await socket_manager.add_user_to_chat(chat_guid, websocket) @@ -238,10 +240,10 @@ async def get_chat_websocket( while True: try: incoming_message = await websocket.receive_json() + logger.info(f'新消息：{incoming_message}') await ratelimit(websocket) - try: - message = schema.Message(**incoming_message) + message = schema.WSMessage(**incoming_message) except ValueError: await socket_manager.send_error(websocket, 'The message type or format you sent is invalid') continue @@ -255,10 +257,9 @@ async def get_chat_websocket( await handler( websocket=websocket, db_session=db_session, - cache=redis_conn, current_user=current_user, chats=chats, - message=message, + incoming_message=message, ) except (JSONDecodeError, AttributeError) as exc: @@ -273,4 +274,4 @@ async def get_chat_websocket( logger.exception(f'User: {login_user} sent too many ws requests') await websocket.send_json({'status': 'error', 'message': 'You have sent too many requests'}) except WebSocketDisconnect: - logging.info('Websocket is disconnected') + logging.info(f'Websocket is disconnected: {login_user.user_id}') diff --git a/core/routers/chat/schema.py b/core/routers/chat/schema.py index 52c3dcb..9f9f3f8 100644 --- a/core/routers/chat/schema.py +++ b/core/routers/chat/schema.py @@ -17,31 +17,6 @@ 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 @@ -51,10 +26,51 @@ class ChatCreateIn(BaseModel): recipient_user_id: int +class PreviewMessage(BaseModel): + message_type: MessageType + content: str + + class ChatUnreadOut(BaseModel): chat_guid: UUID4 created_at: str updated_at: str recipient_user: LiteUser - last_message: str | None + last_message: PreviewMessage | None unread_messages_count: int + + +class MessageOut(BaseModel): + chat_id: int + chat_guid: UUID4 + id: int + message_guid: UUID4 + message_type: MessageType + content: str + user: LiteUser + created_at: str + updated_at: str + + +class WSMessage(BaseModel): + type: Literal[ + 'chat.new_message', 'chat.message_read', 'chat.ping' + ] + message_guid: UUID4 + message_type: MessageType + chat_guid: UUID4 + receive_user: Optional[int] = None + content: str + + +class WSMessageOut(BaseModel): + type: Literal[ + 'chat.new_message', 'chat.send_status_reply', + ] + status: Optional[Literal['ok', 'error']] = 'ok' + message_id: int + message_guid: UUID4 + message_type: MessageType + chat_guid: UUID4 + content: Optional[str] = None + friend_user: Optional[LiteUser] = None diff --git a/core/routers/chat/services.py b/core/routers/chat/services.py index b388fbe..38196fc 100644 --- a/core/routers/chat/services.py +++ b/core/routers/chat/services.py @@ -50,7 +50,7 @@ async def get_unread_messages_per_chat( # 查询消息情况 read_status_alias = aliased(ReadStatus) query = ( - select(Message.chat_id, Message.content, func.count().label("message_count")) + select(Message.chat_id, Message.content, Message.message_type, 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), @@ -70,8 +70,11 @@ async def get_unread_messages_per_chat( for messages_info in unread_messages_info: unread_messages_info_per_chat[messages_info[0]] = { - 'message': messages_info[1], - 'message_count': messages_info[2], + 'message': { + 'message_type': messages_info[2], + 'content': messages_info[1], + }, + 'message_count': messages_info[3], } results = []