external_demand/app/gap_script_demand/odps_writer.py

"""脚本主驱需求写入 ODPS 需求池表。"""

from __future__ import annotations

import re
from typing import Any

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 build_odps_row(
    *,
    strategy: str,
    demand_id: str,
    demand_name: str,
    weight: float = 1.0,
    demand_type: str = "特征点",
) -> dict[str, Any]:
    return {
        "strategy": strategy,
        "demand_id": demand_id,
        "demand_name": demand_name,
        "weight": weight,
        "type": demand_type,
        "video_count": 0,
        "extend": "",
    }


def _build_row_select(row: dict[str, Any]) -> str:
    strategy = _escape_sql_string(str(row["strategy"]))
    demand_id = _escape_sql_string(str(row["demand_id"]))
    demand_name = _escape_sql_string(str(row["demand_name"]))
    weight = float(row.get("weight") or 1.0)
    demand_type = _escape_sql_string(str(row.get("type") or "特征点"))
    extend = _escape_sql_string(str(row.get("extend") or ""))
    return f"""
        SELECT
            '{strategy}' AS strategy,
            '{demand_id}' AS demand_id,
            '{demand_name}' AS demand_name,
            {weight} AS weight,
            '{demand_type}' AS type,
            CAST(0 AS BIGINT) AS video_count,
            array() AS video_list,
            '{extend}' AS extend
    """


def sync_gap_script_demands_to_odps(
    *,
    rows: list[dict[str, Any]],
    partition_dt: str,
    target_table: str,
) -> dict[str, Any]:
    """将脚本主驱需求追加写入 ODPS 分区（不覆盖分区已有数据）。"""
    partition_value = partition_dt.strip()
    if not partition_value:
        raise GapScriptDemandError("partition_dt is required")
    if not rows:
        return {
            "target_table": target_table,
            "partition_dt": partition_value,
            "written_count": 0,
        }

    table_name = _safe_identifier(target_table)
    odps_client = get_odps_client()
    select_sql = " UNION ALL ".join(_build_row_select(row) for row in rows)
    sql = f"""
    INSERT INTO TABLE {table_name} PARTITION (dt='{_escape_sql_string(partition_value)}')
    {select_sql}
    """
    instance = odps_client.execute_sql(sql)
    instance.wait_for_success()
    return {
        "target_table": table_name,
        "partition_dt": partition_value,
        "written_count": len(rows),
    }