Agent/examples/auto_put_ad_mini/roi_control/odps_client.py

#!/usr/bin/env python
# coding=utf-8

import os

from odps import ODPS, options


class ODPSClient(object):
    def __init__(self, project="loghubods"):
        self.accessId = os.environ.get("ODPS_ACCESS_ID")
        self.accessSecret = os.environ.get("ODPS_ACCESS_SECRET")
        if not self.accessId or not self.accessSecret:
            raise RuntimeError(
                "请设置 ODPS_ACCESS_ID 和 ODPS_ACCESS_SECRET 环境变量"
            )

        self.endpoint = os.getenv(
            "ODPS_ENDPOINT", "http://service.odps.aliyun.com/api"
        )

        # 调长超时，应对复杂 SQL 长时间运行（20 分钟）
        options.connect_timeout = 60
        options.read_timeout = 1200
        options.retry_times = 3

        self.odps = ODPS(
            self.accessId,
            self.accessSecret,
            project,
            self.endpoint,
        )

    def execute_sql(self, sql: str):
        hints = {
            'odps.sql.submit.mode': 'script'
        }
        # 异步提交以便先打印 instanceId / logview，方便用户在 ODPS 控制台跟踪进度
        inst = self.odps.run_sql(sql, hints=hints)
        try:
            logview = inst.get_logview_address()
        except Exception:
            logview = '(获取 logview 失败)'
        print(f'[ODPS] InstanceId: {inst.id}', flush=True)
        print(f'[ODPS] LogView   : {logview}', flush=True)
        inst.wait_for_success()
        with inst.open_reader(tunnel=True) as reader:
            pd_df = reader.to_pandas()
        return pd_df

    def execute_sql_result_save_file(self, sql: str, output_file: str):
        data_df = self.execute_sql(sql)
        data_df.to_csv(output_file, index=False)