# -*- 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, ChatParticipant from core.routers.auth.schema import LoginUser from core.utils.helper import format_date 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 return None async def get_chat_unread_messages_count(db_session: AsyncSession, login_user: LoginUser, chats): """查询未读消息的预览""" # 查询消息情况 read_status_alias = aliased(ReadStatus) query = ( select(Message.chat_id, 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_(chats), ) ) .group_by(Message.chat_id) ) result = await db_session.execute(query) unread_messages_info = result.fetchall() return unread_messages_info async def get_chat_messages_preview(db_session: AsyncSession, chats): """查询未读消息的预览""" latest_subquery = ( select( Message.chat_id, func.max(Message.created_at).label("max_created") ) .group_by(Message.chat_id) .subquery() ) LatestMessage = aliased(Message) query = ( select(LatestMessage) .join(latest_subquery, (LatestMessage.chat_id == latest_subquery.c.chat_id) & (LatestMessage.created_at == latest_subquery.c.max_created)) ) result = await db_session.execute(query) unread_messages_info = result.scalars().all() return unread_messages_info async def get_unread_messages_per_chat( db_session: AsyncSession, chats: list[Chat], login_user: LoginUser ) -> list[dict]: """获取聊天未读消息数和消息预览""" messages_info_per_chat: dict[int, dict] = {chat.id: {'message_count': 0} for chat in chats} # 未读消息数和预览 unread_messages_count = await get_chat_unread_messages_count( db_session, login_user, messages_info_per_chat ) for messages_info in unread_messages_count: messages_info_per_chat[messages_info[0]] = { 'message_count': messages_info[1], } # 已读消息预览 no_unread_message_chats = [ chat_id for chat_id, message_info in messages_info_per_chat.items() if message_info.get('message') is None ] messages_info = await get_chat_messages_preview(db_session, no_unread_message_chats) for message_info in messages_info: messages_info_per_chat[message_info.chat_id].update({ 'message': message_info }) 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 = 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'), is_deleted=chat.is_deleted, unread_messages_count=messages_info.get('message_count'), )) return results async def get_chat(chat_guid, login_user: User, db_session: AsyncSession, is_query_messages: bool = False) -> Chat | None: subquery = ( select(ChatParticipant.chat_id) .where( and_( ChatParticipant.user_id == login_user.user_id, ChatParticipant.is_deleted == False ) ) ).scalar_subquery() conditions = [ Chat.id.in_(subquery), Chat.guid == chat_guid ] if is_query_messages is False: conditions.append( Chat.is_deleted == False ) query = ( select(Chat) .where(and_( *conditions )) ) result = await db_session.execute(query) await db_session.commit() return result.scalar_one_or_none()