如何为awswrangler.athena.read_sql_query配置重试与超时机制？

AWS Wrangler基于boto3构建，你可以通过以下几种方式实现针对读超时、字节数错误等场景的自动重试：

1. 配置boto3全局重试策略 #

通过自定义boto3的客户端配置，指定重试次数、重试覆盖的错误类型，同时延长超时时间减少触发概率：

import boto3
from botocore.config import Config
import awswrangler as wr

# 构建重试配置：覆盖读超时、连接错误等可重试场景
retry_config = Config(
    retries={
        'max_attempts': 5,
        'mode': 'standard'  # 自动识别并重试AWS定义的可重试错误
    },
    read_timeout=300,  # 延长读超时至5分钟
    connect_timeout=60
)

# 创建带重试配置的会话
session = boto3.Session(client_config=retry_config)

# 使用该会话执行Athena查询
df = wr.athena.read_sql_query(
    sql="SELECT * FROM your_target_table",
    database="your_database",
    boto3_session=session
)

2. 自定义特定错误的重试逻辑 #

如果需要更精准控制重试范围（比如只针对ReadTimeoutError和字节数错误），可以用tenacity库实现装饰器式重试：

import awswrangler as wr
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
from botocore.exceptions import ReadTimeoutError, ClientError

# 定义需要重试的错误判断逻辑
def is_retryable_error(exc):
    if isinstance(exc, ReadTimeoutError):
        return True
    if isinstance(exc, ClientError):
        # 根据实际字节数错误对应的ErrorCode调整，示例为假设值
        return exc.response['Error']['Code'] in ['InvalidContentLength', 'ConnectionError']
    return False

# 应用重试规则
@retry(
    stop=stop_after_attempt(5),
    wait=wait_exponential(multiplier=1, min=2, max=10),  # 指数退避等待，避免频繁请求
    retry=retry_if_exception_type((ReadTimeoutError, ClientError)) & is_retryable_error
)
def execute_athena_query(sql, db):
    return wr.athena.read_sql_query(sql=sql, database=db)

# 调用查询
df = execute_athena_query("SELECT * FROM your_target_table", "your_database")

3. 结合AWS Wrangler内置的查询状态轮询重试 #

AWS Wrangler的Athena方法自带max_tries参数，用于重试查询状态的轮询（比如查询未完成时的重试），可配合前两种方法一起使用：

df = wr.athena.read_sql_query(
    sql="SELECT * FROM your_target_table",
    database="your_database",
    max_tries=5  # 轮询查询执行状态的重试次数
)

注意：该参数仅针对查询执行状态的轮询，不处理HTTP层面的超时或字节数错误，需和boto3配置/tenacity重试配合使用。

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