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()