PythonでAthenaにクエリする

はじめに

今回、PythonでAthenaを使ってデータ取得する機会がありましたので、
その時に利用したawswranglerを使ったAthenaへのクエリ実行方法を紹介したいと思います。
※この記事ではAWS接続用のプロファイルの設定やAthenaへのデータ登録は既に済んでいることを前提にしています

インストール

pip install awswrangler

使い方

import awswrangler as wr

# クエリリクエスト
df = wr.athena.read_sql_query(
    ql='SELECT * FROM test_table',
    database='db_name'
)
# 結果を出力
print(df)

ドキュメント

接続情報

awswranglerがプロファイルの情報を取得して接続するので、スクリプト側で特に指定する必要はありません。
プロファイルが複数ある場合は環境変数に以下を追加すれば、指定したプロファイルを利用して接続を行います

AWS_DEFAULT_PROFILE=profile_name

リトライ処理

athenaに接続を試みて失敗した場合の最大リトライ回数を指定できます

import awswrangler as wr
import botocore

wr.config.botocore_config = botocore.config.Config(
    retries={"max_attempts": 5},
)

ドキュメント

ただし、待ち時間が以下のように短いため長めに設定したいのであればリトライ用のライブラリを利用して対応した方が良さそうです

2021-06-01 11:34:43 sleeping for 0.6901632031039504 seconds
2021-06-01 11:34:44 sleeping for 0.7869570857297759 seconds
2021-06-01 11:34:45 sleeping for 1.955833391490407 seconds
2021-06-01 11:34:47 sleeping for 2.6080917647958914 seconds

まとめ

boto3のみでも今回の処理は実現できますが、athenaへのクエリ実行を実装する場合は自前で実行状況を確認する処理を実装して状況に応じてデータを取得するようにしないといけなくなる為、すごい手間を要します。
しかし、今回紹介したawswranglerを使えば全てやってくれるので、おすすめです。

また、Athenaへのクエリ実行だけでなくRedshiftやS3など他のAWSサービスに関する機能も提供しているので、これがあればいろいろな箇所で実装の手間が省けそうです。