Agent/examples/tencent_realtime_control/odps_source.py

"""ODPS source for the overall hourly CPM signal."""

from __future__ import annotations

import os
from dataclasses import dataclass
from datetime import date

import pandas as pd
from odps import ODPS


@dataclass(frozen=True)
class HourlyCpm:
    partition: str
    hour_of_day: int
    impressions: int
    cpm: float


def build_odps_client() -> ODPS:
    missing = [
        key
        for key in ("ODPS_ACCESS_ID", "ODPS_ACCESS_SECRET")
        if not os.getenv(key)
    ]
    if missing:
        raise RuntimeError(f"Missing ODPS environment variables: {', '.join(missing)}")
    return ODPS(
        os.environ["ODPS_ACCESS_ID"],
        os.environ["ODPS_ACCESS_SECRET"],
        os.getenv("ODPS_PROJECT", "loghubods"),
        endpoint=os.getenv("ODPS_ENDPOINT", "http://service.odps.aliyun.com/api"),
    )


def fetch_hourly_cpm(client: ODPS, data_date: str) -> pd.DataFrame:
    sql = f"""
SELECT dt,
       CAST(SUBSTR(dt, 9, 2) AS BIGINT) AS hour_of_day,
       `曝光次数_总` AS impressions,
       `真实cpm_总` AS cpm
FROM loghubods.advertiser_data_da_hour
WHERE dt LIKE '{data_date}%'
  AND `名称` = 'SUM'
  AND advertisercode = 'SUM'
  AND company = 'SUM'
  AND `行业` = 'SUM'
  AND `客户` = 'SUM'
  AND `落地页类型` = 'SUM'
ORDER BY dt
"""
    instance = client.execute_sql(sql, hints={"odps.sql.submit.mode": "script"})
    with instance.open_reader(tunnel=True) as reader:
        frame = reader.to_pandas()
    if frame.empty:
        return frame
    frame["hour_of_day"] = pd.to_numeric(frame["hour_of_day"]).astype("int64")
    frame["impressions"] = pd.to_numeric(frame["impressions"]).astype("int64")
    frame["cpm"] = pd.to_numeric(frame["cpm"]).astype(float)
    return frame


def fetch_latest_cpm(client: ODPS, data_date: date) -> HourlyCpm | None:
    frame = fetch_hourly_cpm(client, data_date.strftime("%Y%m%d"))
    if frame.empty:
        return None
    row = frame.iloc[-1]
    return HourlyCpm(
        partition=str(row["dt"]),
        hour_of_day=int(row["hour_of_day"]),
        impressions=int(row["impressions"]),
        cpm=float(row["cpm"]),
    )