feat(chat): 聊天功能

core/routers/chat/handlers.py

@@ -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
+        )

core/routers/chat/managers.py

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

core/routers/chat/schema.py

@@ -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

core/routers/chat/services.py

@@ -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