data_query_skill/skills/odps-ad-risk-analysis/scripts/odps_module.py

#!/usr/bin/env python3
"""Minimal PyODPS client configured only through environment variables."""

from __future__ import annotations

import os
from pathlib import Path

from odps import ODPS, options


ENV_ACCESS_ID = "ODPS_ACCESS_ID"
ENV_ACCESS_SECRET = "ODPS_ACCESS_SECRET"
ENV_PROJECT = "ODPS_PROJECT"
ENV_ENDPOINT = "ODPS_ENDPOINT"


def require_env(name: str) -> str:
    value = os.getenv(name)
    if not value:
        raise RuntimeError(f"Missing required environment variable: {name}")
    return value


class ODPSClient:
    def __init__(self, project: str | None = None, endpoint: str | None = None):
        access_id = require_env(ENV_ACCESS_ID)
        access_secret = require_env(ENV_ACCESS_SECRET)
        project_name = project or require_env(ENV_PROJECT)
        endpoint_url = endpoint or require_env(ENV_ENDPOINT)

        options.connect_timeout = 60
        options.read_timeout = 1200
        options.retry_times = 3

        self.odps = ODPS(
            access_id,
            access_secret,
            project=project_name,
            endpoint=endpoint_url,
        )

    def execute_sql(self, sql: str):
        if not sql.strip():
            raise ValueError("SQL must not be empty")

        instance = self.odps.run_sql(
            sql,
            hints={"odps.sql.submit.mode": "script"},
        )
        print(f"[ODPS] InstanceId: {instance.id}", flush=True)
        instance.wait_for_success()
        with instance.open_reader(tunnel=True) as reader:
            return reader.to_pandas()

    def execute_sql_result_save_file(self, sql: str, output_file: str | Path):
        output_path = Path(output_file).expanduser()
        output_path.parent.mkdir(parents=True, exist_ok=True)
        data = self.execute_sql(sql)
        data.to_csv(output_path, index=False, encoding="utf-8-sig")
        return output_path