import boto3
import json
import time

def query_athena_and_return_json(query, database, output_location):
    """
    执行Athena查询并返回JSON格式结果
    
    :param query: SQL查询语句
    :param database: Athena数据库名称
    :param output_location: S3输出路径
    :return: JSON格式的查询结果
    """
    # 创建Athena客户端
    athena_client = boto3.client('athena')
    
    # 启动查询执行
    response = athena_client.start_query_execution(
        QueryString=query,
        QueryExecutionContext={
            'Database': database
        },
        ResultConfiguration={
            'OutputLocation': output_location
        }
    )
    
    query_execution_id = response['QueryExecutionId']
    
    # 等待查询完成
    while True:
        status = athena_client.get_query_execution(
            QueryExecutionId=query_execution_id
        )
        state = status['QueryExecution']['Status']['State']
        
        if state in ['SUCCEEDED', 'FAILED', 'CANCELLED']:
            break
        
        time.sleep(1)  # 避免频繁轮询
    
    if state != 'SUCCEEDED':
        raise Exception(f"Query failed with state: {state}")
    
    # 获取查询结果
    results = athena_client.get_query_results(
        QueryExecutionId=query_execution_id
    )
    
    # 转换为JSON格式
    columns = [col['Name'] for col in results['ResultSet']['ResultSetMetadata']['ColumnInfo']]
    rows = results['ResultSet']['Rows'][1:]  # 跳过标题行
    
    json_result = []
    for row in rows:
        json_row = {}
        for i, data in enumerate(row['Data']):
            json_row[columns[i]] = data.get('VarCharValue', None)
        json_result.append(json_row)
    
    return json.dumps(json_result, indent=2)

# 启动查询执行
QUERY_ID=$(aws athena start-query-execution \
    --query-string "SELECT * FROM your_table LIMIT 10" \
    --query-execution-context Database=your_database \
    --result-configuration OutputLocation=s3://your-bucket/path/ \
    --output text)

# 等待查询完成
while true; do
    STATUS=$(aws athena get-query-execution --query-execution-id $QUERY_ID \
             --query 'QueryExecution.Status.State' --output text)
    if [[ $STATUS == "SUCCEEDED" ]]; then
        break
    elif [[ $STATUS == "FAILED" || $STATUS == "CANCELLED" ]]; then
        echo "Query failed or was cancelled"
        exit 1
    fi
    sleep 1
done

# 获取结果并转换为JSON
aws athena get-query-results --query-execution-id $QUERY_ID \
    --output json | jq '.ResultSet.Rows[1:] | map(.Data | map_values(.[]?))'