如何检查 Athena 查询在 Python 中是否成功运行？

Question

我想使用boto3 athena 客户端函数start_query_execution从 Python 运行查询。此处的函数文档：https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/athena.html#Athena.Client.start_query_execution

用法如下：

query = "SELECT * FROM TABLE"
athena_client = boto3.client("athena")
start_response = athena_client.start_query_execution(
    QueryString=query,
    QueryExecutionContext={"Database": ATHENA.database_name},
    ResultConfiguration={
        "OutputLocation": s3_output,
    },
)

我正在寻找一个函数/包装器，以确保此查询成功运行，并且仅在运行完成后才返回。从我的搜索中找不到 aws 包装器。

Answer 1

我实现了一个通用函数，该函数执行特定查询，并通过按时间间隔轮询查询 ID 来确保它成功运行：

import time
import logging
import boto3

def run_query(query: str, s3_output: str) -> None:
    """Generic function to run athena query and ensures it is successfully completed

    Parameters
    ----------
    query : str
        formatted string containing athena sql query
    s3_output : str
        query output path
    """
    athena_client = boto3.client("athena")
    start_response = athena_client.start_query_execution(
        QueryString=query,
        QueryExecutionContext={"Database": ATHENA.database_name},
        ResultConfiguration={
            "OutputLocation": s3_output,
        },
    )
    query_id = start_response["QueryExecutionId"]

    while True:
        finish_state = athena_client.get_query_execution(QueryExecutionId=query_id)[
            "QueryExecution"
        ]["Status"]["State"]
        if finish_state == "RUNNING" or finish_state == "QUEUED":
            time.sleep(10)
        else:
            break

    assert finish_state == "SUCCEEDED", f"query state is {finish_state}"
    logging.info(f"Query {query_id} complete")