long_articles_agentic_supply/src/domains/fetch_demand/_mapper.py

"""需求搜索队列 MySQL 读写 + 视频搜索词 PG 查询 —— 仅负责 DB I/O，不含业务逻辑"""

from typing import Dict, List

from src.infra.database.mysql.manager import MysqlManager
from src.infra.database.postgresql.manager import PgManager

from ._const import DemandConst


class DemandQueueMapper(DemandConst):
    """demand_search_queue 表的数据访问层

    职责: SQL 拼接 + 参数绑定 + 返回原生 dict。
    不负责: 业务校验、格式转换（由上层 service 负责）。
    """

    def __init__(self, pool: MysqlManager):
        self.pool = pool

    # ═══════════════════════════════════════════════════════════════
    # 写入
    # ═══════════════════════════════════════════════════════════════

    async def insert_batch(self, items: List[Dict]) -> int:
        """批量写入队列，返回受影响行数

        items 每条须包含:
          account, search_key, dt, video_id, channel_name,
          key_type, experiment_id, drive_type
        """
        if not items:
            return 0
        query = f"""
            INSERT IGNORE INTO {self.QUEUE_TABLE}
                (account, search_key, dt, video_id,
                 channel_name, category, key_type, experiment_id,
                 status, drive_type)
            VALUES
                (%s, %s, %s, %s,
                 %s, %s, %s, %s,
                 %s, %s)
        """
        params = [
            (
                it["account"],
                it["search_key"],
                it["dt"],
                it["video_id"],
                it["channel_name"],
                it.get("category", ""),
                it["key_type"],
                it["experiment_id"],
                self.QueueStatus.INIT,
                it["drive_type"],
            )
            for it in items
        ]
        return await self.pool.save(query=query, params=params, batch=True)

    # ═══════════════════════════════════════════════════════════════
    # 读取（供下游 consumer 使用）
    # ═══════════════════════════════════════════════════════════════

    async def fetch_pending(
        self,
        dt: str,
        *,
        limit: int = 500,
        offset: int = 0,
    ) -> List[Dict]:
        """查询待搜索的队列项（按 dt）"""
        query = f"""
            SELECT *
            FROM {self.QUEUE_TABLE}
            WHERE dt = %s AND status = %s
            ORDER BY id ASC
            LIMIT %s OFFSET %s
        """
        return await self.pool.fetch(
            query=query,
            params=(dt, self.QueueStatus.INIT, limit, offset),
        )

    async def count_by_status(self, dt: str, status: int) -> int:
        """统计指定状态的队列项数量"""
        query = f"""
            SELECT COUNT(*) AS cnt
            FROM {self.QUEUE_TABLE}
            WHERE dt = %s AND status = %s
        """
        row = await self.pool.fetch_one(query=query, params=(dt, status))
        return row["cnt"] if row else 0

    # ═══════════════════════════════════════════════════════════════
    # 状态更新
    # ═══════════════════════════════════════════════════════════════

    async def mark_processing(self, item_id: int) -> int:
        query = f"""
            UPDATE {self.QUEUE_TABLE}
            SET status = %s
            WHERE id = %s AND status = %s
        """
        return await self.pool.save(
            query=query,
            params=(self.QueueStatus.PROCESSING, item_id, self.QueueStatus.INIT),
        )

    async def mark_processing_batch(self, ids: List[int]) -> int:
        """批量标记为处理中"""
        if not ids:
            return 0
        placeholders = ",".join(["%s"] * len(ids))
        query = f"""
            UPDATE {self.QUEUE_TABLE}
            SET status = %s
            WHERE id IN ({placeholders})
        """
        return await self.pool.save(
            query=query,
            params=(self.QueueStatus.PROCESSING, *ids),
        )

    async def mark_done(self, item_id: int) -> int:
        query = f"""
            UPDATE {self.QUEUE_TABLE}
            SET status = %s
            WHERE id = %s
        """
        return await self.pool.save(
            query=query, params=(self.QueueStatus.SUCCESS, item_id),
        )

    async def mark_failed(self, item_id: int, reason: str = "") -> int:
        query = f"""
            UPDATE {self.QUEUE_TABLE}
            SET status = %s, fail_reason = %s
            WHERE id = %s
        """
        return await self.pool.save(
            query=query,
            params=(self.QueueStatus.FAIL, reason[:512], item_id),
        )


class VideoDeconstructMapper:
    """video_vectors 表的数据访问层（PostgreSQL）

    职责: 按 video_id 查询搜索词，每行 (config_code, text) 为一个搜索词。
    config_code 直接对应 DemandConst.KeyType。
    """

    COLUMNS = ("video_id", "config_code", "text")

    def __init__(self, pg: PgManager):
        self.pg = pg

    async def find_by_video_ids(self, video_ids: list[int]) -> dict[int, list[Dict]]:
        """批量查询 video_vectors，一次网络往返

        Returns:
            {video_id: [{"config_code": "VIDEO_KEYPOINT", "text": "xx"}, ...], ...}
            不存在的 video_id 不会出现在返回 key 中。
        """
        if not video_ids:
            return {}
        cols = ", ".join(self.COLUMNS)
        rows = await self.pg.fetch(
            f"SELECT {cols} FROM video_vectors WHERE video_id = ANY($1)",
            params=(video_ids,),
        )
        result: dict[int, list[Dict]] = {}
        for row in rows:
            vid = int(row["video_id"])
            config_code = (row.get("config_code") or "").strip()
            text_val = (row.get("text") or "").strip()
            if config_code and text_val:
                if vid not in result:
                    result[vid] = []
                result[vid].append({"config_code": config_code, "text": text_val})
        return result


__all__ = ["DemandQueueMapper", "VideoDeconstructMapper"]