277 lines
9.9 KiB
Python
277 lines
9.9 KiB
Python
# -*- coding = utf-8 -*-
|
|
# @Time : 2025/2/13 下午4:53
|
|
# @File : routes.py
|
|
# @Software : PyCharm
|
|
# @Author : xingc
|
|
# @Desc :
|
|
import logging
|
|
from typing import List, Dict
|
|
from json.decoder import JSONDecodeError
|
|
|
|
from fastapi import APIRouter, Query, Depends, WebSocket, WebSocketDisconnect, HTTPException, status
|
|
from fastapi_limiter.depends import WebSocketRateLimiter
|
|
import redis.asyncio as aioredis
|
|
from sqlalchemy import select, func, and_, exists
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from sqlalchemy.orm import selectinload
|
|
|
|
from core.base import schema as base_schema
|
|
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.base.exceptions import WebsocketTooManyRequests
|
|
|
|
from . import schema
|
|
from .handlers import socket_manager
|
|
from .services import websocket_callback, get_user_active_direct_chats, get_unread_messages_per_chat
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
chat_router = APIRouter(prefix='/chat', tags=['聊天模块'])
|
|
|
|
|
|
@chat_router.get(
|
|
'/all/', summary='获取所有聊天',
|
|
response_model=base_schema.BaseResponse[
|
|
base_schema.Paginated[
|
|
List[schema.ChatUnreadOut]
|
|
]
|
|
]
|
|
)
|
|
async def get_all_chats(
|
|
login_user: LoginUser = Depends(get_current_user),
|
|
db_session: AsyncSession = Depends(get_async_session),
|
|
page: int = Query(default=1, gt=0),
|
|
page_size: int = Query(default=15, gt=0, le=200),
|
|
):
|
|
filter_conditions = (
|
|
exists().where(
|
|
(chat_participant.c.chat_id == Chat.id) &
|
|
(chat_participant.c.user_id == login_user.user_id)
|
|
) &
|
|
Chat.is_deleted.is_(False)
|
|
)
|
|
query = (
|
|
select(Chat)
|
|
.where(filter_conditions).order_by(Chat.updated_at.desc())
|
|
)
|
|
|
|
total_query = select(func.count()).select_from(Chat).filter(filter_conditions)
|
|
total_result = await db_session.execute(total_query)
|
|
total = total_result.scalar()
|
|
|
|
query = query.limit(page_size).offset((page - 1) * page_size)
|
|
result = await db_session.execute(query)
|
|
chats = result.scalars().all()
|
|
|
|
unread_messages_chats_count = await get_unread_messages_per_chat(db_session, chats, login_user)
|
|
|
|
return {
|
|
'data': {
|
|
'list': unread_messages_chats_count,
|
|
'total': total,
|
|
'page': page,
|
|
'page_size': page_size,
|
|
}
|
|
}
|
|
|
|
|
|
@chat_router.get(
|
|
'/{chat_guid}/messages/', summary='获取单个聊天的历史消息',
|
|
response_model=base_schema.BaseResponse[
|
|
base_schema.Paginated[
|
|
List[Dict]
|
|
]
|
|
]
|
|
)
|
|
async def get_single_chat_messages(
|
|
chat_guid: str,
|
|
login_user: LoginUser = Depends(get_current_user),
|
|
db_session: AsyncSession = Depends(get_async_session),
|
|
page: int = Query(default=1, gt=0),
|
|
page_size: int = Query(default=15, gt=0, le=200),
|
|
last_read_message_id: int = Query(default=0, ge=0),
|
|
):
|
|
db_crud = CRUD(db_session, Chat)
|
|
chat: Chat | None = await db_crud.get(filters={'guid': chat_guid, 'is_deleted': False})
|
|
if not chat:
|
|
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail='您访问的聊天不存在')
|
|
|
|
chat_user_ids = [user.user_id for user in await chat.awaitable_attrs.users]
|
|
if login_user.user_id not in chat_user_ids:
|
|
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail='您无权访问此聊天')
|
|
|
|
# 组装查询条件
|
|
conditions = and_(
|
|
Message.chat_id == chat.id,
|
|
Message.id > last_read_message_id
|
|
)
|
|
|
|
# 记录总数查询
|
|
total_stmt = (
|
|
select(func.count()).where(
|
|
conditions
|
|
)
|
|
)
|
|
total = await db_session.execute(total_stmt)
|
|
# 分页记录查询
|
|
stmt = (
|
|
select(Message)
|
|
.where(conditions)
|
|
.order_by(Message.created_at.desc())
|
|
.limit(page_size).offset((page - 1) * page_size)
|
|
.options(selectinload(Message.user), selectinload(Message.chat))
|
|
)
|
|
records = await db_session.execute(stmt)
|
|
|
|
return {
|
|
'data': {
|
|
'list': [
|
|
{
|
|
**card.dict(),
|
|
**user_collection.dict()
|
|
} for card, user_collection in records
|
|
],
|
|
'total': total.scalar(),
|
|
'page': page,
|
|
'page_size': page_size,
|
|
},
|
|
}
|
|
|
|
|
|
@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),
|
|
db_session: AsyncSession = Depends(get_async_session),
|
|
):
|
|
if payload.recipient_user_id == login_user.user_id:
|
|
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail='操作禁止')
|
|
|
|
recipient_user: User | None = await CRUD(db_session, User).get(filters={'user_id': payload.recipient_user_id})
|
|
if recipient_user is None:
|
|
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail='收件人不存在')
|
|
|
|
initiator_user: User = await CRUD(db_session, User).get(filters={'user_id': login_user.user_id})
|
|
# 查询是否已存在 对应且未关闭的聊天
|
|
query = (
|
|
select(Chat)
|
|
.where(
|
|
exists().where(
|
|
(chat_participant.c.chat_id == Chat.id) &
|
|
(chat_participant.c.user_id == login_user.user_id)
|
|
) &
|
|
exists().where(
|
|
(chat_participant.c.chat_id == Chat.id) &
|
|
(chat_participant.c.user_id == recipient_user.user_id)
|
|
) &
|
|
Chat.is_deleted == False
|
|
).order_by(Chat.created_at.desc())
|
|
)
|
|
result = await db_session.execute(query)
|
|
chat = result.scalar_one_or_none()
|
|
|
|
if chat is None:
|
|
try:
|
|
chat = Chat()
|
|
chat.users.append(initiator_user)
|
|
chat.users.append(recipient_user)
|
|
db_session.add(chat)
|
|
await db_session.flush()
|
|
|
|
initiator_read_status = ReadStatus(chat_id=chat.id, user_id=initiator_user.user_id, last_read_message_id=0)
|
|
recipient_read_status = ReadStatus(chat_id=chat.id, user_id=recipient_user.user_id, last_read_message_id=0)
|
|
db_session.add_all([initiator_read_status, recipient_read_status])
|
|
await db_session.commit()
|
|
|
|
except Exception as exc_info:
|
|
await db_session.rollback()
|
|
raise exc_info
|
|
|
|
return {
|
|
'data': {
|
|
'chat_guid': chat.guid,
|
|
'friend_user': recipient_user,
|
|
}
|
|
}
|
|
|
|
|
|
@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),
|
|
db_session: AsyncSession = Depends(get_async_session),
|
|
):
|
|
chat = await CRUD(db_session, Chat).update(
|
|
data={'is_deleted': True},
|
|
filters={'chat_guid': chat_guid, 'user_id': login_user.user_id}
|
|
)
|
|
if chat is None:
|
|
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail='聊天不存在')
|
|
|
|
return {'message': '聊天已删除'}
|
|
|
|
|
|
@chat_router.websocket('/ws')
|
|
async def get_chat_websocket(
|
|
websocket: WebSocket,
|
|
login_user: LoginUser = Depends(get_current_user),
|
|
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)
|
|
user_id = str(login_user.user_id)
|
|
await socket_manager.add_user_socket_connection(user_id, websocket)
|
|
|
|
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)
|
|
else:
|
|
chats = dict()
|
|
|
|
try:
|
|
while True:
|
|
try:
|
|
incoming_message = await websocket.receive_json()
|
|
await ratelimit(websocket)
|
|
|
|
try:
|
|
message = schema.Message(**incoming_message)
|
|
except ValueError:
|
|
await socket_manager.send_error(websocket, 'The message type or format you sent is invalid')
|
|
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')
|
|
continue
|
|
|
|
await handler(
|
|
websocket=websocket,
|
|
db_session=db_session,
|
|
cache=redis_conn,
|
|
current_user=current_user,
|
|
chats=chats,
|
|
message=message,
|
|
)
|
|
|
|
except (JSONDecodeError, AttributeError) as exc:
|
|
logger.exception(f'Websocket error, detail: {exc}')
|
|
await websocket.send_json({'status': 'error', 'message': 'Wrong message format'})
|
|
|
|
except ValueError as exc:
|
|
logger.exception(f'Websocket error, detail: {exc}')
|
|
await websocket.send_json({'status': 'error', 'message': 'Could not validate incoming message'})
|
|
|
|
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'})
|
|
except WebSocketDisconnect:
|
|
logging.info('Websocket is disconnected')
|