feat(chat): 聊天功能
This commit is contained in:
xingc
119
core/routers/chat/handlers.py
Normal file
119
core/routers/chat/handlers.py
Normal file
@@ -0,0 +1,119 @@
|
||||
# -*- coding = utf-8 -*-
|
||||
# @Time : 2025/3/1 下午4:40
|
||||
# @File : handlers.py
|
||||
# @Software : PyCharm
|
||||
# @Author : xingc
|
||||
# @Desc :
|
||||
import logging
|
||||
from datetime import datetime
|
||||
from typing import List, Dict
|
||||
|
||||
import redis.asyncio as aioredis
|
||||
from fastapi import WebSocket
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from core.db.crud import CRUD
|
||||
from core.db.models import ReadStatus, User, Chat, Message
|
||||
|
||||
from .managers import SocketManager
|
||||
from . import schema
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
socket_manager = SocketManager()
|
||||
|
||||
|
||||
@socket_manager.handler('chat.new_message')
|
||||
async def send_new_message(
|
||||
websocket: WebSocket,
|
||||
db_session: AsyncSession,
|
||||
cache: aioredis.Redis,
|
||||
chats: dict,
|
||||
current_user: User,
|
||||
incoming_message: schema.Message,
|
||||
):
|
||||
chat_guid: str = str(incoming_message.chat_guid)
|
||||
chat: Chat | None = await CRUD(db_session, Chat).get(filters={'chat_guid': chat_guid})
|
||||
# 判断是否为新聊天
|
||||
if not chats or chat_guid not in chats:
|
||||
if chat:
|
||||
chats[chat_guid] = chat.id
|
||||
await socket_manager.add_user_to_chat(chat_guid, websocket)
|
||||
else:
|
||||
await socket_manager.send_error('聊天尚未添加', websocket)
|
||||
return
|
||||
|
||||
chat_id = chats.get(chat_guid)
|
||||
try:
|
||||
message = Message(
|
||||
guid=incoming_message.message_guid,
|
||||
user_id=current_user.user_id,
|
||||
chat_id=chat_id,
|
||||
message_type=incoming_message.message_type,
|
||||
content=incoming_message.body,
|
||||
)
|
||||
db_session.add(message)
|
||||
await db_session.flush()
|
||||
|
||||
chat.updated_at = datetime.now()
|
||||
db_session.add(chat)
|
||||
|
||||
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()
|
||||
logger.exception(f"[new_message] Exception, rolling back session, detail: {exc_info}")
|
||||
raise exc_info
|
||||
|
||||
send_message_schema = schema.MessageOut(
|
||||
type='chat.new_message',
|
||||
message_id=message.id,
|
||||
message_guid=message.guid,
|
||||
message_type=message.message_type,
|
||||
chat_guid=chat.guid,
|
||||
user_id=current_user.user_id,
|
||||
content=message.content,
|
||||
created_at=message.created_at,
|
||||
friend_user=current_user,
|
||||
)
|
||||
await socket_manager.broadcast_to_chat(chat_guid, send_message_schema.dict())
|
||||
|
||||
|
||||
@socket_manager.handler('chat.refresh_unread_messages')
|
||||
async def refresh_unread_messages(
|
||||
websocket: WebSocket,
|
||||
db_session: AsyncSession,
|
||||
cache: aioredis.Redis,
|
||||
current_user: User,
|
||||
message: schema.Message,
|
||||
):
|
||||
pass
|
||||
|
||||
|
||||
@socket_manager.handler('chat.message_read')
|
||||
async def read_message(
|
||||
websocket: WebSocket,
|
||||
db_session: AsyncSession,
|
||||
cache: aioredis.Redis,
|
||||
current_user: User,
|
||||
message: schema.Message,
|
||||
):
|
||||
last_read_message_id: str | None = message.body
|
||||
if not last_read_message_id:
|
||||
return
|
||||
|
||||
db_crud = CRUD(db_session, ReadStatus)
|
||||
condition = {
|
||||
'user_id': current_user.user_id,
|
||||
'chat_guid': message.chat_guid,
|
||||
}
|
||||
result = await db_crud.update(filters=condition, data={'last_read_message_id': last_read_message_id})
|
||||
if result is None:
|
||||
await db_crud.create(
|
||||
data={
|
||||
**condition,
|
||||
'last_read_message_id': last_read_message_id
|
||||
},
|
||||
filters=condition
|
||||
)
|
||||
116
core/routers/chat/managers.py
Normal file
116
core/routers/chat/managers.py
Normal file
@@ -0,0 +1,116 @@
|
||||
# -*- coding = utf-8 -*-
|
||||
# @Time : 2025/3/1 下午8:47
|
||||
# @File : managers.py
|
||||
# @Software : PyCharm
|
||||
# @Author : xingc
|
||||
# @Desc :
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
import redis.asyncio as aioredis
|
||||
import ujson
|
||||
from fastapi import WebSocket
|
||||
|
||||
from core.db.engine import get_cache
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RedisPubSubManager:
|
||||
def __init__(self):
|
||||
self.redis_connection = None
|
||||
self.pubsub = None
|
||||
|
||||
async def _get_redis_connection(self) -> aioredis:
|
||||
return await get_cache()
|
||||
|
||||
async def connect(self):
|
||||
self.redis_connection = await self._get_redis_connection()
|
||||
self.pubsub = self.redis_connection.pubsub()
|
||||
|
||||
async def subscribe(self, chat_guid: str) -> aioredis:
|
||||
await self.pubsub.subscribe(chat_guid)
|
||||
return self.pubsub
|
||||
|
||||
async def unsubscribe(self, chat_guid: str):
|
||||
await self.pubsub.unsubscribe(chat_guid)
|
||||
|
||||
async def publish(self, chat_guid: str, message: str):
|
||||
await self.redis_connection.publish(chat_guid, message)
|
||||
|
||||
async def disconnect(self):
|
||||
if self.redis_connection is not None:
|
||||
await self.redis_connection.close()
|
||||
|
||||
|
||||
class SocketManager:
|
||||
|
||||
def __init__(self):
|
||||
self.handlers: dict = {}
|
||||
self.chats: dict = {}
|
||||
self.pubsub_client = RedisPubSubManager()
|
||||
self.user_guid_to_websocket: dict = {}
|
||||
self._task_cancel = True
|
||||
|
||||
def handler(self, message_type):
|
||||
def decorator(func):
|
||||
self.handlers[message_type] = func
|
||||
return func
|
||||
|
||||
return decorator
|
||||
|
||||
async def connect_socket(self, websocket: WebSocket):
|
||||
await websocket.accept()
|
||||
|
||||
async def add_user_socket_connection(self, user_guid: str, websocket: WebSocket):
|
||||
self.user_guid_to_websocket.setdefault(user_guid, set()).add(websocket)
|
||||
|
||||
async def add_user_to_chat(self, chat_guid: str, websocket: WebSocket):
|
||||
if chat_guid in self.chats:
|
||||
self.chats[chat_guid].add(websocket)
|
||||
else:
|
||||
self.chats[chat_guid] = {websocket}
|
||||
await self.pubsub_client.connect()
|
||||
pubsub_subscriber = await self.pubsub_client.subscribe(chat_guid)
|
||||
await asyncio.create_task(self._pubsub_data_reader(pubsub_subscriber))
|
||||
|
||||
async def broadcast_to_chat(self, chat_guid: str, message: str | dict) -> None:
|
||||
if isinstance(message, dict):
|
||||
message = ujson.dumps(message)
|
||||
await self.pubsub_client.publish(chat_guid, message)
|
||||
|
||||
async def remove_user_from_chat(self, chat_guid: str, websocket: WebSocket) -> None:
|
||||
self.chats[chat_guid].remove(websocket)
|
||||
if len(self.chats[chat_guid]) == 0:
|
||||
del self.chats[chat_guid]
|
||||
logger.info(f'Removing user from PubSub channel {chat_guid}')
|
||||
await self.pubsub_client.unsubscribe(chat_guid)
|
||||
|
||||
async def remove_user_guid_to_websocket(self, user_guid: str, websocket: WebSocket):
|
||||
if user_guid in self.user_guid_to_websocket:
|
||||
self.user_guid_to_websocket.get(user_guid).remove(websocket)
|
||||
|
||||
# https://github.com/redis/redis-py/issues/2523
|
||||
async def _pubsub_data_reader(self, pubsub_subscriber):
|
||||
try:
|
||||
while True:
|
||||
message = await pubsub_subscriber.get_message(ignore_subscribe_messages=True)
|
||||
if message is not None:
|
||||
chat_guid = message['channel']
|
||||
sockets = self.chats.get(chat_guid)
|
||||
if sockets:
|
||||
for socket in sockets:
|
||||
data = message['data']
|
||||
await socket.send_text(data)
|
||||
else:
|
||||
if self._task_cancel is False:
|
||||
break
|
||||
except Exception as exc:
|
||||
logger.exception(f'Exception occurred: {exc}')
|
||||
|
||||
async def send_error(self, message: str, websocket: WebSocket):
|
||||
await websocket.send_json({'status': 'error', 'message': message})
|
||||
|
||||
async def quit(self):
|
||||
self._task_cancel = False
|
||||
await self.pubsub_client.disconnect()
|
||||
60
core/routers/chat/schema.py
Normal file
60
core/routers/chat/schema.py
Normal file
@@ -0,0 +1,60 @@
|
||||
# -*- coding = utf-8 -*-
|
||||
# @Time : 2025/3/8 下午10:54
|
||||
# @File : schema.py
|
||||
# @Software : PyCharm
|
||||
# @Author : xingc
|
||||
# @Desc :
|
||||
from datetime import datetime
|
||||
from typing import Literal, Optional, List
|
||||
|
||||
from pydantic import BaseModel, conint, UUID4
|
||||
|
||||
from core.base.schema import MessageType
|
||||
from core.routers.user.schema import LiteUser
|
||||
|
||||
|
||||
class ChatUser(LiteUser):
|
||||
pass
|
||||
|
||||
|
||||
class Message(BaseModel):
|
||||
type: Literal[
|
||||
'chat.new_message', 'chat.message_read'
|
||||
]
|
||||
message_guid: UUID4
|
||||
message_type: MessageType
|
||||
chat_guid: UUID4
|
||||
receive_user: Optional[int] = None
|
||||
body: str
|
||||
|
||||
|
||||
class MessageOut(BaseModel):
|
||||
type: Literal[
|
||||
'chat.new_message', 'chat.send_status_reply'
|
||||
]
|
||||
code: Optional[int] = 200
|
||||
status: Optional[Literal['ok', 'error']] = 'ok'
|
||||
message_id: int
|
||||
message_guid: UUID4
|
||||
message_type: MessageType
|
||||
chat_guid: UUID4
|
||||
body: Optional[str] = None
|
||||
friend_user: LiteUser
|
||||
|
||||
|
||||
class Chat(BaseModel):
|
||||
chat_guid: UUID4
|
||||
friend_user: LiteUser
|
||||
|
||||
|
||||
class ChatCreateIn(BaseModel):
|
||||
recipient_user_id: int
|
||||
|
||||
|
||||
class ChatUnreadOut(BaseModel):
|
||||
chat_guid: UUID4
|
||||
created_at: str
|
||||
updated_at: str
|
||||
recipient_user: LiteUser
|
||||
last_message: str | None
|
||||
unread_messages_count: int
|
||||
90
core/routers/chat/services.py
Normal file
90
core/routers/chat/services.py
Normal file
@@ -0,0 +1,90 @@
|
||||
# -*- coding = utf-8 -*-
|
||||
# @Time : 2025/3/1 下午4:44
|
||||
# @File : services.py
|
||||
# @Software : PyCharm
|
||||
# @Author : xingc
|
||||
# @Desc :
|
||||
from typing import Dict
|
||||
|
||||
from sqlalchemy import select, and_, func
|
||||
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.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")
|
||||
|
||||
|
||||
async def get_user_active_direct_chats(
|
||||
db_session: AsyncSession, current_user: User
|
||||
) -> Dict[str, int] | None:
|
||||
"""获取活跃聊天"""
|
||||
direct_chats_dict = dict()
|
||||
query = (
|
||||
select(Chat)
|
||||
.where(and_(Chat.is_deleted.is_(False), Chat.users.contains(current_user)))
|
||||
.options(selectinload(Chat.users))
|
||||
)
|
||||
result = await db_session.execute(query)
|
||||
direct_chats: list[Chat] = result.scalars().all()
|
||||
|
||||
if direct_chats:
|
||||
for direct_chat in direct_chats:
|
||||
direct_chats_dict[str(direct_chat.guid)] = direct_chat.id
|
||||
return direct_chats_dict
|
||||
|
||||
|
||||
async def get_unread_messages_per_chat(
|
||||
db_session: AsyncSession, chats: list[Chat], login_user: LoginUser
|
||||
) -> list[dict]:
|
||||
"""获取聊天未读消息数"""
|
||||
unread_messages_info_per_chat = {chat.id: {'message_count': 0} for chat in chats}
|
||||
|
||||
# 查询消息情况
|
||||
read_status_alias = aliased(ReadStatus)
|
||||
query = (
|
||||
select(Message.chat_id, Message.content, func.count().label("message_count"))
|
||||
.join(
|
||||
read_status_alias,
|
||||
and_(read_status_alias.user_id == login_user.user_id, read_status_alias.chat_id == Message.chat_id),
|
||||
)
|
||||
.where(
|
||||
and_(
|
||||
Message.user_id != login_user.user_id,
|
||||
Message.id > func.coalesce(read_status_alias.last_read_message_id, 0),
|
||||
Message.chat_id.in_(unread_messages_info_per_chat),
|
||||
)
|
||||
)
|
||||
.order_by(Message.created_at.desc())
|
||||
.group_by(Message.chat_id)
|
||||
)
|
||||
result = await db_session.execute(query)
|
||||
unread_messages_info = result.fetchall()
|
||||
|
||||
for messages_info in unread_messages_info:
|
||||
unread_messages_info_per_chat[messages_info[0]] = {
|
||||
'message': messages_info[1],
|
||||
'message_count': messages_info[2],
|
||||
}
|
||||
|
||||
results = []
|
||||
for chat in chats:
|
||||
users = [user for user in await chat.awaitable_attrs.users if user.user_id != login_user.user_id]
|
||||
messages_info = unread_messages_info_per_chat[chat.id]
|
||||
results.append(dict(
|
||||
chat_guid=chat.guid,
|
||||
created_at=format_date(chat.created_at),
|
||||
updated_at=format_date(chat.updated_at),
|
||||
recipient_user=users[0],
|
||||
last_message=messages_info.get('message'),
|
||||
unread_messages_count=messages_info.get('message_count'),
|
||||
))
|
||||
|
||||
return results
|
||||
Reference in New Issue
Block a user