feat(chat): 完善聊天

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}')