sales-assistant-py-new/app/repository/archive_messages_repository.py

from typing import List, Optional
from datetime import datetime, timedelta

from sqlalchemy import text, update
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.mysql import ArchiveMessages


class ArchiveMessagesRepository:
    def __init__(self, session: AsyncSession):
        self.session = session

    async def get_message_statistics_by_date(self, start_date: int = None, end_date: int = None) -> List[ArchiveMessages]:
        """
        根据日期范围获取消息记录

        Args:
            start_date: 开始日期 (格式: YYYY-MM-DD)
            end_date: 结束日期 (格式: YYYY-MM-DD)

        Returns:
            ArchiveMessages 对象列表
        """
        base_query = """
            SELECT DATE(FROM_UNIXTIME(msgtime / 1000)) AS created_at, from_user
            FROM archive_messages 
        """
        conditions = []
        params = {}

        if start_date:
            conditions.append(" and msgtime >= :start_date")
            params["start_date"] = start_date

        if end_date:
            conditions.append(" msgtime <= :end_date")
            params["end_date"] = end_date
        base_query += " WHERE 1=1 and roomid='' and msgtype in ('text','voice') "
        #base_query += " WHERE (from_user='BuShouShiJinBuGaiMing') and roomid='' and msgtype='text' "
        if conditions:
            base_query += "  AND ".join(conditions)+"  "
        base_query += " GROUP BY DATE(FROM_UNIXTIME(msgtime / 1000)), from_user ORDER BY created_at DESC, from_user;"

        query = text(base_query)
        result = await self.session.execute(query, params)
        return [ArchiveMessages(**dict(row)) for row in result.mappings().fetchall()]


    async def get_message_statistics_by_date_user(self, day: str = None, from_user: str = None, to_user: str = None) -> List[ArchiveMessages]:
        """
        根据日期+发送人+接收人查询记录
        Returns:
            ArchiveMessages 对象列表
        """
        if to_user:
            sql = """
                SELECT * FROM archive_messages 
                WHERE DATE(FROM_UNIXTIME(msgtime / 1000)) = :day 
                  AND from_user = :from_user 
                  AND to_user = :to_user and msgtype in ('text','voice')
                ORDER BY created_at DESC
            """
            result = await self.session.execute(text(sql), {"day": day, "from_user": from_user, "to_user": to_user})
        else :
            sql = """
                SELECT from_user,to_user FROM archive_messages 
                WHERE DATE(FROM_UNIXTIME(msgtime / 1000)) = :day 
                  AND from_user = :from_user and msgtype in ('text','voice')
                group by to_user
            """
            result = await self.session.execute(text(sql), {"day": day, "from_user": from_user})
        return [ArchiveMessages(**dict(row)) for row in result.mappings().fetchall()]

    async def update_content_by_id(self, message_id: int, content: str) -> bool:
        """
        根据ID更新消息的content字段

        Args:
            message_id: 消息ID
            content: 新的内容

        Returns:
            是否更新成功
        """
        try:
            stmt = (
                update(ArchiveMessages)
                .where(ArchiveMessages.id == message_id)
                .values(content=content)
            )
            result = await self.session.execute(stmt)
            await self.session.commit()
            return result.rowcount > 0
        except Exception as e:
            await self.session.rollback()
            raise e

    async def get_conversation_messages(
        self,
        from_user: str,
        to_user: str,
        start_time: Optional[int] = None,
        end_time: Optional[int] = None,
    ) -> List[ArchiveMessages]:

        toUser_m = "[\"" + to_user + "\"]"
        fromUser_m = "[\"" + from_user + "\"]"
        # 默认时间范围：当天开始(含) ~ 次日开始(不含)，单位：UTC 毫秒
        if start_time is None or end_time is None:
            today_start = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0)
            next_day_start = today_start + timedelta(days=1)
            if start_time is None:
                start_time = int(today_start.timestamp() * 1000)
            if end_time is None:
                end_time = int(next_day_start.timestamp() * 1000)

        sql = """
            (
                SELECT msgtime, created_at, from_role, content
                FROM archive_messages
                WHERE from_user = :from_user_ab
                  AND to_user = :to_user_ab
                  AND msgtime >= :start_time
                  AND msgtime < :end_time
                  AND msgtype IN ('text', 'voice')
                  AND content IS NOT NULL
                  AND content != ''
                LIMIT 100
            )
            UNION ALL
            (
                SELECT msgtime, created_at, from_role, content
                FROM archive_messages
                WHERE from_user = :from_user_ba
                  AND to_user = :to_user_ba
                  AND msgtime >= :start_time
                  AND msgtime < :end_time
                  AND msgtype IN ('text', 'voice')
                  AND content IS NOT NULL
                  AND content != ''
                LIMIT 100
            )
            ORDER BY msgtime ASC
        """
        params = {
            # 方向1: from_user -> to_user
            "from_user_ab": from_user,
            "to_user_ab": toUser_m,
            # 方向2: to_user -> from_user
            "from_user_ba": to_user,
            "to_user_ba": fromUser_m,
            # 共用时间范围
            "start_time": start_time,
            "end_time": end_time,
        }
        result = await self.session.execute(text(sql), params)
        return [ArchiveMessages(**dict(row)) for row in result.mappings().fetchall()]