trauma/api/message/db_requests.py

import asyncio

from fastapi import HTTPException

from trauma.api.account.model import AccountModel
from trauma.api.chat.model import ChatModel
from trauma.api.message.dto import Author
from trauma.api.message.model import MessageModel
from trauma.api.message.schemas import CreateMessageRequest
from trauma.core.config import settings


async def get_all_chat_messages_obj(
        chat_id: str, account: AccountModel
) -> list[MessageModel]:
    messages, chat = await asyncio.gather(
        settings.DB_CLIENT.messages.find({"chatId": chat_id}).to_list(length=None),
        settings.DB_CLIENT.chats.find_one({"id": chat_id})
    )
    messages = [MessageModel.from_mongo(message) for message in messages]

    if not chat:
        raise HTTPException(status_code=404, detail="Chat not found")

    chat = ChatModel.from_mongo(chat)
    if account and chat.account != account:
        raise HTTPException(status_code=403, detail="Chat account not match")

    return messages

async def create_message_obj(
        chat_id: str, message_data: CreateMessageRequest, account: AccountModel
) -> tuple[MessageModel, ChatModel]:
    chat = await settings.DB_CLIENT.chats.find_one({"id": chat_id})
    if not chat:
        raise HTTPException(status_code=404, detail="Chat not found")

    chat = ChatModel.from_mongo(chat)
    if account and chat.account != account:
        raise HTTPException(status_code=403, detail="Chat account not match")

    message = MessageModel(**message_data.model_dump(), chatId=chat_id, author=Author.User)
    await settings.DB_CLIENT.messages.insert_one(message.to_mongo())

    return message, chat