Python+Boto3调用AWS Athena：合并查询与结果获取脚本问题

我明白你的问题了——分开运行两个函数能拿到结果，但整合到一个脚本里就不行，核心原因是Athena的查询是异步执行的：你提交查询后立刻去调用get_query_results，此时查询可能还在排队或者运行中，自然拿不到有效数据。

下面是完整的整合脚本，我会附带关键逻辑的说明：

import boto3
import time

def run_athena_query(query, database, s3_output):
    # 初始化Athena客户端，替换成你的AWS区域
    athena_client = boto3.client('athena', region_name='us-east-1')

    # 提交查询请求
    submit_response = athena_client.start_query_execution(
        QueryString=query,
        QueryExecutionContext={'Database': database},
        ResultConfiguration={'OutputLocation': s3_output}
    )
    query_id = submit_response['QueryExecutionId']
    print(f"查询已提交，ID: {query_id}")

    # 轮询检查查询状态，直到进入终端状态
    while True:
        status_check = athena_client.get_query_execution(QueryExecutionId=query_id)
        current_status = status_check['QueryExecution']['Status']['State']
        
        # 终端状态：成功/失败/取消
        if current_status in ['SUCCEEDED', 'FAILED', 'CANCELLED']:
            break
        print(f"当前查询状态: {current_status}，5秒后再次检查...")
        time.sleep(5)

    # 根据状态处理结果
    if current_status == 'SUCCEEDED':
        print("查询执行完成，开始获取结果...")
        results = athena_client.get_query_results(QueryExecutionId=query_id)
        
        # 解析并打印结果（可根据需求自定义格式化逻辑）
        columns = [col['Label'] for col in results['ResultSet']['ResultSetMetadata']['ColumnInfo']]
        print("\n查询列名:", columns)
        
        # 跳过表头行，遍历数据行
        data_rows = results['ResultSet']['Rows'][1:]
        for row in data_rows:
            row_values = [data.get('VarCharValue', '') for data in row['Data']]
            print(row_values)
        return results
    else:
        error_reason = status_check['QueryExecution']['Status'].get('StateChangeReason', '无详细原因')
        print(f"查询失败/被取消，状态: {current_status}，原因: {error_reason}")
        return None

# 示例调用，替换成你的实际参数
if __name__ == "__main__":
    MY_QUERY = "SELECT * FROM your_target_table LIMIT 10;"
    MY_DATABASE = "your_athena_database"
    S3_OUTPUT_PATH = "s3://your-bucket/athena-query-results/"
    
    run_athena_query(MY_QUERY, MY_DATABASE, S3_OUTPUT_PATH)

关键逻辑说明 #

1. 异步等待机制：
   用get_query_execution轮询查询状态，必须等查询进入终端状态后再调用get_query_results，这是整合脚本的核心——解决了分开运行时手动等待查询完成的问题。

2. 状态处理：
   覆盖了所有可能的终端状态，避免无限循环，同时在查询失败时返回具体错误原因，方便排查问题。

3. 结果解析：
   Athena返回的结果结构有固定格式，表头存在ResultSetMetadata中，数据行从第二行开始（第一行是表头），示例里做了基础的解析和打印，你可以根据需求调整成CSV、JSON等格式。

额外注意事项 #

• 可以根据查询数据量调整time.sleep(5)的间隔，大数据查询可以设为10-15秒，减少不必要的API调用。
• 如果查询结果超过1000行，get_query_results会返回NextToken，需要循环调用直到获取所有结果（示例为简化逻辑未实现分页，若需要可自行添加循环处理NextToken）。
• 确保你的IAM角色拥有Athena执行权限、S3输出桶的读写权限，以及查询目标数据集的访问权限。

内容的提问来源于stack exchange，提问作者mindaJalaj