SupplyAgent/supply_infra/db/repositories/demand_video_expansion_repo.py

from __future__ import annotations

from sqlalchemy import func, select
from sqlalchemy.dialects.mysql import insert

from supply_infra.db.models.demand_video_expansion import (
    DemandVideoExpansion,
    DemandVideoExpansionRun,
)
from supply_infra.db.repositories.base import BaseRepository

_BATCH_SIZE = 500


class DemandVideoExpansionRepository(BaseRepository[DemandVideoExpansion]):
    """需求视频点位拓展结果 repository。"""

    model = DemandVideoExpansion

    def get_active_by_id(self, expansion_id: int) -> DemandVideoExpansion | None:
        stmt = select(DemandVideoExpansion).where(
            DemandVideoExpansion.id == int(expansion_id),
            DemandVideoExpansion.is_delete == 0,
        )
        return self.session.scalars(stmt).first()

    def has_video(
        self,
        *,
        biz_dt: str,
        source_demand_grade_id: int,
        video_id: str,
    ) -> bool:
        stmt = select(DemandVideoExpansion.id).where(
            DemandVideoExpansion.biz_dt == biz_dt,
            DemandVideoExpansion.source_demand_grade_id == int(source_demand_grade_id),
            DemandVideoExpansion.video_id == video_id,
            DemandVideoExpansion.is_delete == 0,
        )
        return self.session.scalar(stmt.limit(1)) is not None

    def list_by_demand_grade(
        self, biz_dt: str, source_demand_grade_id: int
    ) -> list[DemandVideoExpansion]:
        """按业务日与来源需求 id 查询拓展结果（未删除）。"""
        stmt = (
            select(DemandVideoExpansion)
            .where(
                DemandVideoExpansion.biz_dt == biz_dt,
                DemandVideoExpansion.source_demand_grade_id == int(source_demand_grade_id),
                DemandVideoExpansion.is_delete == 0,
            )
            .order_by(
                DemandVideoExpansion.video_id,
                DemandVideoExpansion.point_type,
                DemandVideoExpansion.id,
            )
        )
        return list(self.session.scalars(stmt).all())

    def count_distinct_videos_by_demand_grade(self, biz_dt: str) -> dict[int, int]:
        """按业务日统计每条需求的有效命中视频数。"""
        stmt = (
            select(
                DemandVideoExpansion.source_demand_grade_id,
                func.count(func.distinct(DemandVideoExpansion.video_id)),
            )
            .where(
                DemandVideoExpansion.biz_dt == biz_dt,
                DemandVideoExpansion.is_delete == 0,
                DemandVideoExpansion.video_id != "",
                DemandVideoExpansion.point_type.in_(("inspiration", "purpose", "key")),
            )
            .group_by(DemandVideoExpansion.source_demand_grade_id)
        )
        return {
            int(demand_grade_id): int(video_count or 0)
            for demand_grade_id, video_count in self.session.execute(stmt).all()
        }

    def bulk_upsert(self, rows: list[dict]) -> int:
        """按唯一键批量 upsert，冲突时更新 reason / point_desc。"""
        if not rows:
            return 0

        affected = 0
        for i in range(0, len(rows), _BATCH_SIZE):
            batch = rows[i : i + _BATCH_SIZE]
            stmt = insert(DemandVideoExpansion).values(batch)
            stmt = stmt.on_duplicate_key_update(
                reason=stmt.inserted.reason,
                point_desc=stmt.inserted.point_desc,
                run_id=stmt.inserted.run_id,
                is_delete=0,
            )
            result = self.session.execute(stmt)
            affected += result.rowcount or 0
        return affected


class DemandVideoExpansionRunRepository(BaseRepository[DemandVideoExpansionRun]):
    """拓展任务执行记录 repository。"""

    model = DemandVideoExpansionRun

    def list_finished_grade_ids(self, biz_dt: str) -> set[int]:
        """返回指定业务日已成功完成拓展判断的 demand_grade.id 集合。"""
        stmt = select(DemandVideoExpansionRun.source_demand_grade_id).where(
            DemandVideoExpansionRun.biz_dt == biz_dt,
            DemandVideoExpansionRun.status == "finished",
        )
        return {int(v) for v in self.session.scalars(stmt).all() if v is not None}

    def upsert_run(
        self,
        *,
        biz_dt: str,
        run_id: str,
        source_demand_grade_id: int,
        saved_count: int,
        status: str = "finished",
        error_message: str | None = None,
    ) -> None:
        """记录单条需求的拓展判断完成状态。"""
        row = {
            "biz_dt": biz_dt,
            "run_id": run_id,
            "source_demand_grade_id": int(source_demand_grade_id),
            "saved_count": int(saved_count),
            "status": status,
            "error_message": error_message,
        }
        stmt = insert(DemandVideoExpansionRun).values(row)
        stmt = stmt.on_duplicate_key_update(
            run_id=stmt.inserted.run_id,
            saved_count=stmt.inserted.saved_count,
            status=stmt.inserted.status,
            error_message=stmt.inserted.error_message,
        )
        self.session.execute(stmt)

    def list_by_biz_dt(self, biz_dt: str) -> list[DemandVideoExpansionRun]:
        stmt = (
            select(DemandVideoExpansionRun)
            .where(DemandVideoExpansionRun.biz_dt == biz_dt)
            .order_by(DemandVideoExpansionRun.source_demand_grade_id)
        )
        return list(self.session.scalars(stmt).all())

    def get_by_demand_grade(
        self, biz_dt: str, source_demand_grade_id: int
    ) -> DemandVideoExpansionRun | None:
        stmt = select(DemandVideoExpansionRun).where(
            DemandVideoExpansionRun.biz_dt == biz_dt,
            DemandVideoExpansionRun.source_demand_grade_id == int(source_demand_grade_id),
        )
        return self.session.scalars(stmt).first()