如何通过Lambda与Boto3实现Athena查询的分页处理

1. 从请求参数中接收分页参数 #

让前端（如Postman）传递分页核心参数，推荐用page_size（每页条数）和page_number（当前页码），后端再转换为LIMIT和OFFSET，同时做合法性校验避免非法值：

import boto3
import time
import csv
import json
import re

DATABASE = "your-database-name"
output = "s3://your-bucket/path/"
output_bucket = "your-bucket"

def lambda_handler(event, context):
    athena = boto3.client('athena')
    s3 = boto3.client('s3')
    
    # 获取分页参数，设置默认值并做范围限制
    page_size = event.get('page_size', 10)
    page_number = event.get('page_number', 1)
    # 限制每页最大条数，防止一次性查询过多数据
    limit = max(1, min(int(page_size), 100))
    # 计算偏移量，页码从1开始时，offset = (页码-1)*每页条数
    offset = (max(1, int(page_number)) - 1) * limit
    
    base_query = event['query']

2. 安全注入分页参数到SQL #

绝对禁止直接字符串拼接参数（避免SQL注入），推荐两种安全方式：

方式一：使用Athena参数化查询（优先推荐） #

Athena支持Prepared Statements，通过ExecutionParameters传递参数，无需手动拼接：

# 改造原始查询，添加LIMIT和OFFSET的占位符
paginated_query = f"{base_query} LIMIT ? OFFSET ?"

# 执行参数化查询
query_id = athena.start_query_execution(
    QueryString=paginated_query,
    QueryExecutionContext={'Database': DATABASE},
    ResultConfiguration={'OutputLocation': output},
    # 参数必须以字符串形式传递
    ExecutionParameters=[str(limit), str(offset)]
)['QueryExecutionId']

方式二：严格类型校验后的字符串格式化 #

如果不使用参数化查询，必须确保limit和offset为整数后再拼接：

# 校验参数类型，非法则返回错误
if not isinstance(limit, int) or not isinstance(offset, int):
    return {
        'statusCode': 400,
        'body': json.dumps({'error': 'Invalid pagination parameters'})
    }

# 安全拼接分页语句
paginated_query = f"{base_query} LIMIT {limit} OFFSET {offset}"

query_id = athena.start_query_execution(
    QueryString=paginated_query,
    QueryExecutionContext={'Database': DATABASE},
    ResultConfiguration={'OutputLocation': output}
)['QueryExecutionId']

3. 等待查询完成并返回分页结果 #

Athena查询是异步执行的，需等待查询完成后从S3读取结果并转换为JSON：

# 轮询等待查询完成
while True:
    query_status = athena.get_query_execution(QueryExecutionId=query_id)
    status = query_status['QueryExecution']['Status']['State']
    if status in ['SUCCEEDED', 'FAILED', 'CANCELLED']:
        break
    time.sleep(1)

if status == 'SUCCEEDED':
    # 从S3读取查询结果（Athena默认输出为CSV）
    result_key = f"{query_id}.csv"
    s3_response = s3.get_object(Bucket=output_bucket, Key=result_key)
    csv_content = s3_response['Body'].read().decode('utf-8')
    
    # 转换CSV为JSON格式
    reader = csv.DictReader(csv_content.splitlines())
    paginated_data = list(reader)
    
    # 可选：返回分页元数据（总条数、当前页等）
    return {
        'statusCode': 200,
        'body': json.dumps({
            'data': paginated_data,
            'pagination': {
                'page_size': limit,
                'page_number': page_number,
                'total_count': get_total_records(base_query, athena)
            }
        })
    }
else:
    return {
        'statusCode': 500,
        'body': json.dumps({'error': f"Query failed with status: {status}"})
    }

4. 可选：获取总条数用于计算总页数 #

如果UI需要显示总页数，可通过额外执行COUNT(*)查询获取总条数：

def get_total_records(base_query, athena_client):
    # 替换原始查询的SELECT部分为COUNT(*)，并移除原有排序/分页语句
    count_query = re.sub(r'SELECT.*FROM', 'SELECT COUNT(*) AS total FROM', base_query, flags=re.IGNORECASE)
    count_query = re.sub(r'ORDER BY.*$', '', count_query, flags=re.IGNORECASE)
    count_query = re.sub(r'LIMIT.*$', '', count_query, flags=re.IGNORECASE)
    
    count_query_id = athena_client.start_query_execution(
        QueryString=count_query,
        QueryExecutionContext={'Database': DATABASE},
        ResultConfiguration={'OutputLocation': output}
    )['QueryExecutionId']
    
    # 等待COUNT查询完成
    while True:
        status = athena_client.get_query_execution(QueryExecutionId=count_query_id)['QueryExecution']['Status']['State']
        if status in ['SUCCEEDED', 'FAILED', 'CANCELLED']:
            break
        time.sleep(1)
    
    if status == 'SUCCEEDED':
        s3_response = s3.get_object(Bucket=output_bucket, Key=f"{count_query_id}.csv")
        csv_content = s3_response['Body'].read().decode('utf-8')
        return int(csv_content.splitlines()[1])
    return 0

关键注意事项 #

• 性能优化：OFFSET在大数据量场景下性能较差（需扫描前置数据），若数据有唯一排序键（如id），可改用WHERE id > last_id替代OFFSET，大幅提升查询效率。
• 超时控制：Lambda最长执行时间为15分钟，若查询数据量极大，需考虑异步处理或优化查询逻辑。
• 参数校验：始终对分页参数做范围限制，避免恶意请求导致资源耗尽。

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