SupplyAgent/supply_infra/db/repositories/demand_popularity_stats_repo.py

from __future__ import annotations

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

from supply_infra.db.models.demand_popularity_stats import DemandPopularityStats
from supply_infra.db.repositories.base import BaseRepository

_BATCH_SIZE = 500

_UPSERT_COLUMNS = (
    "demand_word_name",
    "ext_pop_avg",
    "ext_pop_count",
    "plat_sust_pop_avg",
    "plat_sust_pop_count",
    "plat_ly_pop_avg",
    "plat_ly_pop_count",
    "recent_pop_avg",
    "recent_pop_count",
    "real_rov_7d_avg",
    "real_rov_7d_count",
    "real_vov_7d_avg",
    "real_vov_7d_count",
)


class DemandPopularityStatsRepository(BaseRepository[DemandPopularityStats]):
    """需求分类热度统计表 repository。"""

    model = DemandPopularityStats

    def get_latest_biz_dt(self) -> str | None:
        """返回热度统计表中最新业务日；无数据时返回 None。"""
        stmt = select(func.max(DemandPopularityStats.biz_dt))
        return self.session.scalar(stmt)

    def has_biz_dt(self, biz_dt: str) -> bool:
        """指定业务日是否存在热度统计数据。"""
        stmt = select(DemandPopularityStats.id).where(
            DemandPopularityStats.biz_dt == biz_dt
        ).limit(1)
        return self.session.scalars(stmt).first() is not None

    def list_by_biz_dt(self, biz_dt: str) -> list[DemandPopularityStats]:
        """返回指定业务日的全部热度统计行。"""
        stmt = select(DemandPopularityStats).where(DemandPopularityStats.biz_dt == biz_dt)
        return list(self.session.scalars(stmt).all())

    def search_by_word_name(
        self,
        keyword: str,
        biz_dt: str | None = None,
    ) -> list[DemandPopularityStats]:
        """按 demand_word_name 精确+模糊搜索；未指定 biz_dt 时不限日期，按 biz_dt 降序返回。"""
        keyword = (keyword or "").strip()
        if not keyword:
            return []

        stmt = select(DemandPopularityStats).where(
            DemandPopularityStats.demand_word_name.like(f"%{keyword}%")
        )
        if biz_dt:
            stmt = stmt.where(DemandPopularityStats.biz_dt == biz_dt)
        stmt = stmt.order_by(DemandPopularityStats.biz_dt.desc())
        return list(self.session.scalars(stmt).all())

    def list_by_biz_dt_and_belong_ids(
        self, biz_dt: str, belong_ids: list[int]
    ) -> list[DemandPopularityStats]:
        """按业务日 + demand_belong_category.id 列表查询热度行。"""
        if not belong_ids:
            return []
        rows: list[DemandPopularityStats] = []
        for i in range(0, len(belong_ids), _BATCH_SIZE):
            batch = belong_ids[i : i + _BATCH_SIZE]
            stmt = select(DemandPopularityStats).where(
                DemandPopularityStats.biz_dt == biz_dt,
                DemandPopularityStats.demand_category_id.in_(batch),
            )
            rows.extend(self.session.scalars(stmt).all())
        return rows

    def upsert_rows(self, rows: list[dict]) -> int:
        """按 (demand_category_id, biz_dt) 批量 upsert。"""
        if not rows:
            return 0

        affected = 0
        for i in range(0, len(rows), _BATCH_SIZE):
            batch = rows[i : i + _BATCH_SIZE]
            stmt = insert(DemandPopularityStats).values(batch)
            stmt = stmt.on_duplicate_key_update(
                **{col: stmt.inserted[col] for col in _UPSERT_COLUMNS}
            )
            result = self.session.execute(stmt)
            affected += result.rowcount or 0
        return affected