377 lines
13 KiB
Python
377 lines
13 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
|
|
from json.decoder import JSONDecodeError
|
|
|
|
from fastapi import APIRouter, Query, Depends, WebSocket, WebSocketDisconnect, HTTPException, status
|
|
from fastapi_limiter.depends import WebSocketRateLimiter
|
|
from sqlalchemy import select, func, and_, exists, update
|
|
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
|
|
from core.db.models import Message, Chat, User, ChatParticipant
|
|
from core.routers.auth.schema import LoginUser
|
|
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
|
|
from .handlers import socket_manager
|
|
from .services import websocket_callback, get_user_active_direct_chats, get_unread_messages_per_chat, get_chat, \
|
|
get_total_unread_messages_count
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
chat_router = APIRouter(prefix='/chat', tags=['聊天模块'])
|
|
|
|
|
|
@chat_router.get(
|
|
'/all/', summary='获取所有聊天',
|
|
response_model=base_schema.BaseResponse[
|
|
schema.ChatPaginated[
|
|
List[schema.ChatUnreadOut]
|
|
]
|
|
]
|
|
)
|
|
async def get_all_chats(
|
|
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),
|
|
):
|
|
filter_conditions = (
|
|
exists().where(
|
|
and_(
|
|
ChatParticipant.chat_id == Chat.id,
|
|
ChatParticipant.user_id == login_user.user_id,
|
|
ChatParticipant.is_deleted == 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 = await get_unread_messages_per_chat(db_session, chats, login_user)
|
|
unread_messages_total = await get_total_unread_messages_count(db_session, login_user)
|
|
|
|
return {
|
|
'data': {
|
|
'list': unread_messages_chats,
|
|
'unread_messages_total': unread_messages_total,
|
|
'next_page': (page * page_size) < total,
|
|
'total': total,
|
|
'page': page,
|
|
'page_size': page_size,
|
|
}
|
|
}
|
|
|
|
|
|
@chat_router.get(
|
|
'/{chat_guid}/messages/', summary='获取单个聊天的历史消息',
|
|
response_model=base_schema.BaseResponse[
|
|
base_schema.Paginated[
|
|
List[schema.MessageOut]
|
|
]
|
|
]
|
|
)
|
|
async def get_single_chat_messages(
|
|
chat_guid: str,
|
|
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),
|
|
last_read_message_id: int = Query(ge=0)
|
|
):
|
|
chat: Chat | None = await get_chat(chat_guid, login_user, db_session, is_query_messages=True)
|
|
if not chat:
|
|
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail='您访问的聊天不存在')
|
|
|
|
read_status = await CRUD(db_session, ChatParticipant).get(
|
|
filters={'chat_id': chat.id, 'user_id': login_user.user_id})
|
|
|
|
# 组装查询条件
|
|
conditions = and_(
|
|
Message.chat_id == chat.id,
|
|
Message.id > (read_status.last_del_message_id or 0),
|
|
Message.id <= last_read_message_id,
|
|
)
|
|
|
|
# 记录总数查询
|
|
total_stmt = (
|
|
select(func.count()).where(
|
|
conditions
|
|
)
|
|
)
|
|
total_result = await db_session.execute(total_stmt)
|
|
total = total_result.scalar()
|
|
# 分页记录查询
|
|
stmt = (
|
|
select(Message)
|
|
.where(conditions)
|
|
.options(selectinload(Message.user), selectinload(Message.chat))
|
|
.order_by(Message.id.desc()).
|
|
limit(page_size).offset((page - 1) * page_size)
|
|
)
|
|
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': results,
|
|
'next_page': (page * page_size) < total,
|
|
'total': total,
|
|
'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_form_http),
|
|
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='操作禁止')
|
|
|
|
db_crud = CRUD(db_session, User)
|
|
|
|
recipient_user: User | None = await db_crud.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 db_crud.get(filters={'user_id': login_user.user_id})
|
|
# 查询是否已存在 对应且未关闭的聊天
|
|
query = (
|
|
select(Chat)
|
|
.where(
|
|
and_(
|
|
exists().where(
|
|
(ChatParticipant.chat_id == Chat.id) &
|
|
(ChatParticipant.user_id == login_user.user_id)
|
|
),
|
|
exists().where(
|
|
(ChatParticipant.chat_id == Chat.id) &
|
|
(ChatParticipant.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()
|
|
await db_session.commit()
|
|
|
|
except Exception as exc_info:
|
|
await db_session.rollback()
|
|
raise exc_info
|
|
|
|
if chat.is_deleted is True:
|
|
await db_crud.updates(model=Chat, data={'is_deleted': False}, filters={'id': chat.id})
|
|
await db_crud.updates(model=ChatParticipant, data={'is_deleted': False}, filters={'chat_id': chat.id})
|
|
|
|
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_form_http),
|
|
db_session: AsyncSession = Depends(get_async_session),
|
|
):
|
|
subquery = (
|
|
select(ChatParticipant.chat_id)
|
|
.where(ChatParticipant.user_id == login_user.user_id)
|
|
).scalar_subquery()
|
|
|
|
query = (
|
|
select(Chat)
|
|
.where(and_(
|
|
Chat.id.in_(subquery),
|
|
Chat.guid == chat_guid
|
|
))
|
|
.options(selectinload(Chat.chat_participants))
|
|
)
|
|
result = await db_session.execute(query)
|
|
chat = result.scalar_one_or_none()
|
|
|
|
if chat is None:
|
|
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail='聊天不存在')
|
|
|
|
deleted_status = [
|
|
chat_participant.is_deleted
|
|
for chat_participant in chat.chat_participants if chat_participant.user_id != login_user.user_id
|
|
]
|
|
if all(deleted_status):
|
|
await db_session.delete(chat)
|
|
await db_session.commit()
|
|
else:
|
|
db_curd = CRUD(db_session, Chat)
|
|
await db_curd.update(
|
|
data={'is_deleted': True},
|
|
filters={'id': chat.id},
|
|
)
|
|
|
|
message_query = (
|
|
select(Message).where(Message.chat_id == chat.id).order_by(Message.id.desc()).limit(1)
|
|
)
|
|
message_result = await db_session.execute(message_query)
|
|
message: Message | None = message_result.scalar()
|
|
|
|
condition = {'user_id': login_user.user_id, 'chat_id': chat.id}
|
|
data = {'is_deleted': True}
|
|
if message:
|
|
data['last_del_message_id'] = message.id
|
|
|
|
await db_curd.update(
|
|
model=ChatParticipant,
|
|
data=data,
|
|
filters=condition,
|
|
)
|
|
|
|
return {'message': '聊天已删除'}
|
|
|
|
|
|
@chat_router.post(
|
|
'/messages/make-all-read',
|
|
summary='聊天信息标记为已读',
|
|
description='支持批量操作多个聊天', response_model=base_schema.BaseResponse
|
|
)
|
|
async def make_messages_read(
|
|
payload: schema.MakeMessageReadIn,
|
|
login_user: LoginUser = Depends(get_current_user_form_http),
|
|
db_session: AsyncSession = Depends(get_async_session),
|
|
):
|
|
db_curd = CRUD(db_session, Chat)
|
|
chat_mapping = {str(chat.chat_guid): chat.last_read_message_id for chat in payload.chats}
|
|
chats = await db_curd.get_all(filters={'guid__in': chat_mapping.keys()})
|
|
|
|
for chat in chats:
|
|
last_read_message_id = chat_mapping.get(chat.guid)
|
|
if last_read_message_id and await db_curd.exists(
|
|
model=Message,
|
|
filters={
|
|
'id': last_read_message_id,
|
|
'chat_id': chat.id,
|
|
}
|
|
):
|
|
await db_curd.update(
|
|
model=ChatParticipant,
|
|
data={'last_read_message_id': last_read_message_id},
|
|
filters={
|
|
'chat_id': chat.id,
|
|
'user_id': login_user.user_id,
|
|
'last_read_message_id__lt': last_read_message_id,
|
|
},
|
|
)
|
|
|
|
return {'message': 'ok'}
|
|
|
|
|
|
@chat_router.websocket('/ws')
|
|
async def get_chat_websocket(
|
|
websocket: WebSocket,
|
|
login_user: LoginUser = Depends(get_current_user_form_ws),
|
|
db_session: AsyncSession = Depends(get_async_session),
|
|
):
|
|
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=30, 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)
|
|
else:
|
|
chats = dict()
|
|
|
|
try:
|
|
while True:
|
|
try:
|
|
incoming_message = await websocket.receive_json()
|
|
logger.info(f'新消息：{incoming_message}')
|
|
await ratelimit(websocket)
|
|
try:
|
|
message = schema.WSMessage(**incoming_message)
|
|
except ValueError:
|
|
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(f'Type: {message.type} was not found', websocket)
|
|
continue
|
|
|
|
await handler(
|
|
websocket=websocket,
|
|
db_session=db_session,
|
|
current_user=current_user,
|
|
chats=chats,
|
|
incoming_message=message,
|
|
)
|
|
|
|
send_status_replay = schema.WSMessageOut(
|
|
type='chat.send_status_reply',
|
|
message_id=-1,
|
|
message_guid=message.message_guid,
|
|
message_type=message.message_type,
|
|
chat_guid=message.chat_guid,
|
|
)
|
|
await websocket.send_text(send_status_replay.model_dump_json(exclude_none=True))
|
|
|
|
except (JSONDecodeError, AttributeError) as exc:
|
|
logger.exception(f'Websocket error, detail: {exc}')
|
|
await socket_manager.send_error('Wrong message format', websocket)
|
|
|
|
except ValueError as exc:
|
|
logger.exception(f'Websocket error, detail: {exc}')
|
|
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 socket_manager.send_error('You have sent too many requests', websocket)
|
|
except WebSocketDisconnect:
|
|
logging.info(f'Websocket is disconnected: {login_user.user_id}')
|