如何通过work_group使用PyAthena和SQLAlchemy连接Athena数据库

不需要在连接串里硬编码AWS访问密钥，PyAthena底层依赖boto3的默认凭证链自动完成鉴权，会按顺序读取环境变量、本地~/.aws/credentials配置、云资源绑定的IAM角色凭证，只要当前运行身份有对应Work Group的操作权限即可。

具体修改步骤 #

• 移除原连接串中{aws_access_key_id}:{aws_secret_access_key}@的鉴权字段段，不需要手动传入任何密钥参数
• 在连接串的查询参数末尾添加work_group参数，传入你要使用的Athena工作组名称，传参时记得用quote_plus做URL编码
• 如果你使用的Work Group已经在AWS控制台配置了默认查询结果S3路径，可以直接移除s3_staging_dir参数，复用工作组的默认配置

修改后的可运行代码 #

from urllib.parse import quote_plus
from sqlalchemy.engine import create_engine
from sqlalchemy.sql.expression import select
from sqlalchemy.sql.functions import func
from sqlalchemy.sql.schema import Table, MetaData

# 无硬编码密钥、指定work_group的连接串
conn_str = "awsathena+rest://athena.{region_name}.amazonaws.com:443/"\
           "{schema_name}?s3_staging_dir={s3_staging_dir}&work_group={work_group}"

engine = create_engine(conn_str.format(
    region_name="us-west-2",
    schema_name="default",
    s3_staging_dir=quote_plus("s3://YOUR_S3_BUCKET/path/to/"), # 若Work Group已配置默认结果路径可删除此行
    work_group=quote_plus("YOUR_WORK_GROUP_NAME") # 替换为实际使用的Athena工作组名称
))

with engine.connect() as connection:
    many_rows = Table("many_rows", MetaData(), autoload_with=connection)
    result = connection.execute(select(func.count("*"), from_obj=many_rows))
    print(result.scalar())

权限注意事项 #

确保运行代码的身份（本地配置的AWS用户、云资源绑定的IAM角色）至少拥有以下权限：指定Work Group的Athena查询权限、查询结果S3暂存路径的读写权限、待查询数据表对应的S3源数据读权限，否则会触发权限拒绝错误。

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