external_demand/app/gap_script_demand/odps_demand_pool.py

"""从 ODPS 需求池拉取当下供需 gap 数据。"""

from __future__ import annotations

import re

from app.aliyun_odps.client import get_odps_client
from app.gap_script_demand.exceptions import GapScriptDemandError

IDENTIFIER_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)?$")


def _safe_identifier(name: str) -> str:
    value = name.strip()
    if not IDENTIFIER_RE.match(value):
        raise GapScriptDemandError(f"invalid sql identifier: {name}")
    return value


def _escape_sql_string(value: str) -> str:
    return value.replace("'", "''")


def _normalize_demand_key(value: str) -> str:
    return "".join(value.split())


def _dedupe_demand_names(demand_names: list[str]) -> list[str]:
    deduped: list[str] = []
    seen: set[str] = set()
    for raw_name in demand_names:
        demand_name = str(raw_name).strip()
        if not demand_name:
            continue
        keys = {demand_name, _normalize_demand_key(demand_name)}
        if keys & seen:
            continue
        seen.update(keys)
        deduped.append(demand_name)
    return deduped


def fetch_demand_names(
    *,
    partition_dt: str,
    strategy: str,
    source_table: str,
) -> list[str]:
    """按分区日期与 strategy 从 ODPS 需求池拉取 demand_name。"""
    table = _safe_identifier(source_table)
    dt = _escape_sql_string(partition_dt.strip())
    strategy_value = _escape_sql_string(strategy.strip())
    if not dt:
        raise GapScriptDemandError("partition_dt is required")
    if not strategy_value:
        raise GapScriptDemandError("strategy is required")

    sql = f"""
    SELECT TRIM(demand_name) AS demand_name
    FROM {table}
    WHERE dt = '{dt}'
      AND strategy = '{strategy_value}'
      AND demand_name IS NOT NULL
      AND TRIM(demand_name) <> ''
    """

    odps_client = get_odps_client()
    instance = odps_client.execute_sql(sql)
    raw_demand_names: list[str] = []
    with instance.open_reader(tunnel=True) as reader:
        for record in reader:
            demand_name = str(record["demand_name"] or "").strip()
            if demand_name:
                raw_demand_names.append(demand_name)
    return _dedupe_demand_names(raw_demand_names)