o2o-negosium-original/backend/crud/chat_crud.py

from abc import ABC, abstractmethod
from datetime import datetime, timezone
from typing import Optional, Tuple

from sqlalchemy import asc, desc, func, select, update
from sqlalchemy.ext.asyncio import AsyncSession

from common.database.db_session_manager import DB_SESSION_MNG
from common.database.model.models import chats, items, sessions
from common.enums import ChatSender, ErrorType
from common.logger import LOG


# 협상 채팅 CRUD. 메시지 로그(negotiation.chats)와 종료 시 세션 입찰 확정(negotiation.sessions)을 다룬다.
# chats / sessions 모두 NEGOTIATION 논리 DB 라 한 트랜잭션(execute_lambda_run)으로 묶을 수 있다.
class IChatCRUD(ABC):
    @abstractmethod
    async def list_by_session(self, cdb: AsyncSession, session_id) -> Tuple[ErrorType, list]:
        pass

    @abstractmethod
    async def get_last(self, cdb: AsyncSession, session_id) -> Tuple[ErrorType, Tuple[int, Optional[int]]]:
        """마지막 메시지의 (seq, sender). 없으면 (0, None). 동시전송 가드 + seq 채번에 사용."""
        pass

    @abstractmethod
    async def insert_message(self, cdb: AsyncSession, message: chats) -> ErrorType:
        pass

    @abstractmethod
    async def soft_delete_message(self, cdb: AsyncSession, chat_id) -> ErrorType:
        pass

    @abstractmethod
    async def count_bot_messages(self, cdb: AsyncSession, session_id) -> Tuple[ErrorType, int]:
        pass

    @abstractmethod
    async def get_item_by_id(self, cdb: AsyncSession, item_id) -> Tuple[ErrorType, items]:
        pass

    @abstractmethod
    async def finalize_session(
        self, cdb: AsyncSession, session_id, status: int,
        bid_price: Optional[int] = None, reject_reason: Optional[str] = None, reject_price: Optional[int] = None,
    ) -> ErrorType:
        pass


class ChatCRUD(IChatCRUD):
    async def list_by_session(self, cdb: AsyncSession, session_id) -> Tuple[ErrorType, list]:
        try:
            # (session_id, seq) 유니크 인덱스가 정렬 스캔을 커버한다.
            query = (
                select(chats)
                .where(chats.session_id == session_id, chats.deleted == False)  # noqa: E712
                .order_by(asc(chats.seq))
            )
            err_type, rows = await DB_SESSION_MNG.execute(cdb, query, "list_by_session failed.")
            if err_type != ErrorType.SUCCESS:
                return err_type, []
            return ErrorType.SUCCESS, rows
        except Exception as ex:
            LOG.e_no_callstack(ex)
            return ErrorType.DB_RUN_FAILED, []

    async def get_last(self, cdb: AsyncSession, session_id) -> Tuple[ErrorType, Tuple[int, Optional[int]]]:
        try:
            query = (
                select(chats.seq, chats.sender)
                .where(chats.session_id == session_id, chats.deleted == False)  # noqa: E712
                .order_by(desc(chats.seq))
                .limit(1)
            )
            err_type, rows = await DB_SESSION_MNG.execute(cdb, query, "get_last failed.")
            if err_type != ErrorType.SUCCESS:
                return err_type, (0, None)
            if not rows:
                return ErrorType.SUCCESS, (0, None)
            return ErrorType.SUCCESS, (rows[0][0], rows[0][1])
        except Exception as ex:
            LOG.e_no_callstack(ex)
            return ErrorType.DB_RUN_FAILED, (0, None)

    async def insert_message(self, cdb: AsyncSession, message: chats) -> ErrorType:
        try:
            return await DB_SESSION_MNG.insert(cdb, message, raise_error=False)
        except Exception as ex:
            LOG.e_no_callstack(ex)
            return ErrorType.DB_RUN_FAILED

    async def soft_delete_message(self, cdb: AsyncSession, chat_id) -> ErrorType:
        # agent 실패 시 선점(pre-claim)한 유저 메시지를 되돌린다. 부분 유니크(WHERE deleted=FALSE)라 seq 가 다시 비워진다.
        try:
            query = update(chats).where(chats.chat_id == chat_id).values(deleted=True)
            return await DB_SESSION_MNG.add(cdb, query)
        except Exception as ex:
            LOG.e_no_callstack(ex)
            return ErrorType.DB_RUN_FAILED

    async def count_bot_messages(self, cdb: AsyncSession, session_id) -> Tuple[ErrorType, int]:
        # mock agent 진행(turn) 계산용. 실제 agent 는 자체 세션 상태로 진행하므로 무시한다.
        try:
            query = select(func.count()).select_from(chats).where(
                chats.session_id == session_id,
                chats.sender == ChatSender.BOT.value,
                chats.deleted == False,  # noqa: E712
            )
            err_type, rows = await DB_SESSION_MNG.execute(cdb, query, "count_bot_messages failed.")
            if err_type != ErrorType.SUCCESS:
                return err_type, 0
            return ErrorType.SUCCESS, (rows[0] if rows else 0)
        except Exception as ex:
            LOG.e_no_callstack(ex)
            return ErrorType.DB_RUN_FAILED, 0

    async def get_item_by_id(self, cdb: AsyncSession, item_id) -> Tuple[ErrorType, items]:
        try:
            query = select(items).where(items.item_id == item_id, items.deleted == False).limit(1)  # noqa: E712
            err_type, row_list = await DB_SESSION_MNG.execute(cdb, query, f"get_item_by_id({item_id}) failed.")
            if err_type != ErrorType.SUCCESS:
                return err_type, None
            if len(row_list) != 1:
                return ErrorType.DB_INVALID_KEY, None
            return ErrorType.SUCCESS, row_list[0]
        except Exception as ex:
            LOG.e_no_callstack(ex)
            return ErrorType.DB_RUN_FAILED, None

    async def finalize_session(
        self, cdb: AsyncSession, session_id, status: int,
        bid_price: Optional[int] = None, reject_reason: Optional[str] = None, reject_price: Optional[int] = None,
    ) -> ErrorType:
        try:
            values = {"status": status}
            if bid_price is not None:
                values["bid_price"] = bid_price
                values["bid_at"] = datetime.now(timezone.utc)
            if reject_reason is not None:
                values["reject_reason"] = reject_reason[:255]
            if reject_price is not None:
                values["reject_price"] = reject_price
            query = update(sessions).where(sessions.session_id == session_id).values(**values)
            return await DB_SESSION_MNG.add(cdb, query)
        except Exception as ex:
            LOG.e_no_callstack(ex)
            return ErrorType.DB_RUN_FAILED