From 99f15cff84a4be4c1abb91fae2a67360e58d3a2c Mon Sep 17 00:00:00 2001 From: xingc Date: Wed, 14 May 2025 21:19:46 +0800 Subject: [PATCH] =?UTF-8?q?fix(chat):=20=E4=BF=AE=E5=A4=8D=E5=A4=9A?= =?UTF-8?q?=E4=B8=AA=E8=81=8A=E5=A4=A9=E9=94=99=E8=AF=AF=20=E6=B6=88?= =?UTF-8?q?=E6=81=AF=E6=A0=BC=E5=BC=8F=E9=94=99=E8=AF=AF=E5=AF=BC=E8=87=B4?= =?UTF-8?q?=E6=96=AD=E8=BF=9E=E3=80=81=E6=96=B0=E6=B6=88=E6=81=AF=E4=BC=9A?= =?UTF-8?q?=E9=80=9A=E7=9F=A5=E5=88=B0=E5=8F=91=E4=BB=B6=E4=BA=BA=E3=80=81?= =?UTF-8?q?=E5=8A=A0=E5=BC=BA=E6=B6=88=E6=81=AF=E5=8F=91=E9=80=81=E7=9A=84?= =?UTF-8?q?=E6=88=90=E5=91=98=E9=AA=8C=E8=AF=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- core/routers/chat/handlers.py | 6 +++--- core/routers/chat/managers.py | 5 ++++- core/routers/chat/routes.py | 11 ++++++----- core/routers/chat/services.py | 23 ++++++++++++++++++++--- 4 files changed, 33 insertions(+), 12 deletions(-) diff --git a/core/routers/chat/handlers.py b/core/routers/chat/handlers.py index 5d94ae3..19884a1 100644 --- a/core/routers/chat/handlers.py +++ b/core/routers/chat/handlers.py @@ -14,6 +14,7 @@ from core.db.models import ReadStatus, User, Chat, Message from .managers import SocketManager from . import schema +from .services import get_chat logger = logging.getLogger(__name__) socket_manager = SocketManager() @@ -28,7 +29,8 @@ async def send_new_message( incoming_message: schema.WSMessage, ): chat_guid: str = str(incoming_message.chat_guid) - chat: Chat | None = await CRUD(db_session, Chat).get(filters={'guid': chat_guid}) + # chat: Chat | None = await CRUD(db_session, Chat).get(filters={'guid': chat_guid, 'is_deleted': False} + chat: Chat | None = await get_chat(chat_guid, current_user, db_session) # 判断是否为新聊天 if not chats or chat_guid not in chats: @@ -56,8 +58,6 @@ async def send_new_message( 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"]) except Exception as exc_info: await db_session.rollback() diff --git a/core/routers/chat/managers.py b/core/routers/chat/managers.py index 0227bff..0d31495 100644 --- a/core/routers/chat/managers.py +++ b/core/routers/chat/managers.py @@ -11,6 +11,7 @@ import redis.asyncio as aioredis import ujson from fastapi import WebSocket from fastapi.websockets import WebSocketState +from simple_spider_tool import jsonpath from core.db.engine import get_cache @@ -104,7 +105,9 @@ class SocketManager: if socket.client_state != WebSocketState.CONNECTED: continue data = message['data'] - await socket.send_text(data) + if f',"user_id":{socket.user_id}' in data: + continue + await socket.send_json(data) else: if self._task_cancel is False: break diff --git a/core/routers/chat/routes.py b/core/routers/chat/routes.py index 8091635..aa72fbe 100644 --- a/core/routers/chat/routes.py +++ b/core/routers/chat/routes.py @@ -239,6 +239,7 @@ async def get_chat_websocket( await socket_manager.connect_socket(websocket=websocket) logger.info(f'Websocket connection is established: {login_user.user_id}') user_id = str(login_user.user_id) + websocket.user_id = login_user.user_id await socket_manager.add_user_socket_connection(user_id, websocket) ratelimit = WebSocketRateLimiter(times=50, seconds=10, callback=websocket_callback) @@ -258,13 +259,13 @@ async def get_chat_websocket( try: message = schema.WSMessage(**incoming_message) except ValueError: - await socket_manager.send_error(websocket, 'The message type or format you sent is invalid') + await socket_manager.send_error('The message type or format you sent is invalid', websocket) continue handler = socket_manager.handlers.get(message.type) if not handler: logger.error(f'No handler [{message.type}] exists') - await socket_manager.send_error(websocket, f'Type: {message.type} was not found') + await socket_manager.send_error(f'Type: {message.type} was not found', websocket) continue await handler( @@ -277,14 +278,14 @@ async def get_chat_websocket( except (JSONDecodeError, AttributeError) as exc: logger.exception(f'Websocket error, detail: {exc}') - await websocket.send_json({'status': 'error', 'message': 'Wrong message format'}) + await socket_manager.send_error('Wrong message format', websocket) except ValueError as exc: logger.exception(f'Websocket error, detail: {exc}') - await websocket.send_json({'status': 'error', 'message': 'Could not validate incoming message'}) + await socket_manager.send_error('Could not validate incoming message', websocket) except WebsocketTooManyRequests: logger.exception(f'User: {login_user} sent too many ws requests') - await websocket.send_json({'status': 'error', 'message': 'You have sent too many requests'}) + await socket_manager.send_error('You have sent too many requests', websocket) except WebSocketDisconnect: logging.info(f'Websocket is disconnected: {login_user.user_id}') diff --git a/core/routers/chat/services.py b/core/routers/chat/services.py index 38196fc..2927f6b 100644 --- a/core/routers/chat/services.py +++ b/core/routers/chat/services.py @@ -11,12 +11,10 @@ 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.db.models import User, Chat, ReadStatus, Message, chat_participant 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") @@ -91,3 +89,22 @@ async def get_unread_messages_per_chat( )) return results + + +async def get_chat(chat_guid, login_user: User, db_session: AsyncSession) -> Chat | None: + subquery = ( + select(chat_participant.c.chat_id) + .where(chat_participant.c.user_id == login_user.user_id) + ).scalar_subquery() + + query = ( + select(Chat) + .where(and_( + Chat.id.in_(subquery), + Chat.guid == chat_guid, + Chat.is_deleted == False + )) + ) + result = await db_session.execute(query) + await db_session.commit() + return result.scalar_one_or_none()