SupplyAgent/supply_infra/odps/client.py

from __future__ import annotations

import logging
from functools import lru_cache
from typing import Any

from supply_infra.config import get_infra_settings

logger = logging.getLogger(__name__)


class ODPSClient:
    """ODPS / MaxCompute 查询客户端封装。"""

    def __init__(
        self,
        access_id: str,
        access_key: str,
        project: str,
        endpoint: str,
    ) -> None:
        self.access_id = access_id
        self.access_key = access_key
        self.project = project
        self.endpoint = endpoint
        self._client: Any = None

    def _get_client(self) -> Any:
        if self._client is None:
            try:
                from odps import ODPS
            except ImportError as e:
                raise ImportError(
                    "pyodps 未安装，请执行: pip install pyodps"
                ) from e
            self._client = ODPS(
                self.access_id,
                self.access_key,
                project=self.project,
                endpoint=self.endpoint,
            )
        return self._client

    def execute_sql(self, sql: str) -> list[dict[str, Any]]:
        """执行 SQL 并返回字典列表。"""
        client = self._get_client()
        logger.info("ODPS executing SQL: %s", sql[:200])
        instance = client.execute_sql(sql)
        with instance.open_reader() as reader:
            columns = [col.name for col in reader._schema.columns]  # type: ignore[attr-defined]
            return [dict(zip(columns, row.values)) for row in reader]
        return []

    def fetch_global_categories(self, bizdate: str) -> list[dict[str, Any]]:
        """拉取 global_category 中仍有效的“实质”分类树。"""
        sql = f"""
        SELECT  stable_id,name,description,level,parent_stable_id
        FROM    loghubods.global_category
        WHERE   dt = '{bizdate}'
        AND     retired_at_execution_id IS NULL
        AND     source_type = '实质'
        ORDER BY level
        """
        return self.execute_sql(sql)

    def fetch_multi_demand_pool(self, bizdate: str) -> list[dict[str, Any]]:
        """拉取 dwd_multi_demand_pool_di 策略需求天级数据（video_list 仅取前 10 个）。"""
        sql = f"""
        SELECT  strategy
                ,demand_id
                ,demand_name
                ,weight
                ,`type`
                ,video_count
                ,SLICE(video_list, 1, 10) AS video_list
                ,extend
        FROM    loghubods.dwd_multi_demand_pool_di
        WHERE   dt = '{bizdate}'
        """
        return self.execute_sql(sql)

    def count_multi_demand_pool(self, bizdate: str) -> int:
        """统计 dwd_multi_demand_pool_di 分区去重后行数（strategy + demand_id）。"""
        sql = f"""
        SELECT  COUNT(1) AS cnt
        FROM    (
                    SELECT  strategy
                            ,demand_id
                    FROM    loghubods.dwd_multi_demand_pool_di
                    WHERE   dt = '{bizdate}'
                    GROUP BY strategy
                             ,demand_id
                ) t
        """
        rows = self.execute_sql(sql)
        if not rows:
            return 0
        return int(rows[0].get("cnt") or 0)

    def fetch_topic_decode_results(
        self,
        dt: str,
        vids: list[str],
        batch_size: int = 100,
    ) -> list[dict[str, Any]]:
        """按 vid 批量拉取 dwd_topic_decode_result_di 的 vid、url1、url2、decode_result。"""
        # 保序去重
        unique_vids = list(
            dict.fromkeys(
                str(v).strip() for v in vids if v is not None and str(v).strip()
            )
        )
        if not unique_vids:
            return []

        results: list[dict[str, Any]] = []
        for i in range(0, len(unique_vids), batch_size):
            batch = unique_vids[i : i + batch_size]
            in_list = ",".join("'" + v.replace("'", "''") + "'" for v in batch)
            sql = f"""
            SELECT  vid
                    ,decode_result
                    ,url1
                    ,url2
            FROM    loghubods.dwd_topic_decode_result_di
            WHERE   dt = '{dt}'
            AND     vid IN ({in_list})
            """
            results.extend(self.execute_sql(sql))
        return results

    def fetch_real_rov_vov_7d(
        self,
        dt_left: str,
        dt_right: str,
        limit: int = 1000,
    ) -> list[dict[str, Any]]:
        """拉取近 N 日人工/自动 AGC 的 rov_diff / vov_diff（相对全局基线）。"""
        sql = f"""
        WITH base AS (
            SELECT
                特征维度,
                特征值,
                特征值_寻找词,
                供给类型,
                当日分发曝光pv,
                当日分发回流uv,
                当日分发拉回曝光pv
            FROM loghubods.dwd_video_produce_plan_stat_hour
            WHERE dt BETWEEN '{dt_left}' AND '{dt_right}'
                AND 供给类型 IN ('人工AGC', '自动AGC')
                AND 特征值 NOT IN ('-', '', 'null')
        ),
        grouped AS (
            SELECT
                1 AS row_order,
                特征维度,
                特征值,
                特征值_寻找词,
                供给类型,
                SUM(当日分发曝光pv) AS exp,
                SUM(当日分发回流uv) AS return_uv,
                SUM(当日分发拉回曝光pv) AS new_exp
            FROM base
            GROUP BY
                特征维度,
                特征值,
                特征值_寻找词,
                供给类型
        ),
        result_rows AS (
            SELECT
                0 AS row_order,
                '全局SUM' AS 特征维度,
                '全局SUM' AS 特征值,
                '全局SUM' AS 特征值_寻找词,
                '全局SUM' AS 供给类型,
                SUM(exp) AS exp,
                SUM(return_uv) AS return_uv,
                SUM(new_exp) AS new_exp,
                COALESCE(ROUND(SUM(return_uv) / NULLIF(SUM(exp), 0), 4), 0) AS rov,
                COALESCE(ROUND(SUM(new_exp) / NULLIF(SUM(exp), 0), 4), 0) AS vov,
                0 AS rov_diff,
                0 AS vov_diff
            FROM grouped
            UNION ALL
            SELECT
                row_order,
                特征维度,
                特征值,
                特征值_寻找词,
                供给类型,
                exp,
                return_uv,
                new_exp,
                COALESCE(ROUND(return_uv / NULLIF(exp, 0), 4), 0) AS rov,
                COALESCE(ROUND(new_exp / NULLIF(exp, 0), 4), 0) AS vov,
                COALESCE(
                    ROUND(
                        (return_uv / NULLIF(exp, 0))
                        / NULLIF(SUM(return_uv) OVER () / NULLIF(SUM(exp) OVER (), 0), 0)
                        - 1,
                        4
                    ),
                    0
                ) AS rov_diff,
                COALESCE(
                    ROUND(
                        (new_exp / NULLIF(exp, 0))
                        / NULLIF(SUM(new_exp) OVER () / NULLIF(SUM(exp) OVER (), 0), 0)
                        - 1,
                        4
                    ),
                    0
                ) AS vov_diff
            FROM grouped
        )
        SELECT
            特征维度,
            特征值,
            特征值_寻找词,
            供给类型,
            exp,
            return_uv AS `return`,
            new_exp,
            rov,
            vov,
            rov_diff,
            vov_diff
        FROM result_rows
        ORDER BY
            row_order,
            exp DESC
        LIMIT {int(limit)}
        """
        return self.execute_sql(sql)


@lru_cache
def get_odps_client() -> ODPSClient:
    settings = get_infra_settings()
    return ODPSClient(
        access_id=settings.odps_access_id,
        access_key=settings.odps_access_key,
        project=settings.odps_project,
        endpoint=settings.odps_endpoint,
    )