AWS LambdaでAthenaのクエリ履歴をS3に定期保存する

Amazon Athena はAmazon S3内のデータをSQLで検索できるサービスです。サーバを用意する必要がなくクエリ課金なので、私のいる部署でも重宝しています。
Athenaの問題点はクエリ履歴が45日間しか表示されないことです。セキュリティ要件としてAWS RDSやRedshiftのように監査ログとして実行したクエリを長期保存したいのですが、Athenaではその機能が提供されていません。
そこで、AWS Lambdaを使ってAthenaのクエリ履歴をS3に定期保存するようにしたので紹介します。

スクリプトはGistにあげています。Python 2.7で動きます。

https://gist.github.com/ohsawa0515/06019739eb9cdc4d6d3813c43fbd1e29

実行に必要なIAM権限はこちらです。

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": [
        "athena:GetQueryExecution",
        "athena:ListQueryExecutions",
        "s3:PutObject"
      ],
      "Effect": "Allow",
      "Resource": "*"
    }
  ]
}

Lambdaでの実行を想定していますが、ローカル環境での実行も可能です。環境変数 S3_BUCKET に保存先のS3バケット名を指定してください。

$ S3_BUCKET='<S3 Bucket>' python logging_athena_history.py

スクリプトを実行したその日のAthenaのクエリ履歴が /tmp/athena_history_YYYYMMDD.csv として保存され、S3にアップロードされます。
ログのサンプルは以下の通りです。

QueryExecutionId,Query,OutputLocation,Database,State,StateChangeReason,SubmissionDateTime,CompletionDateTime,EngineExecutionTimeInMillis,DataScannedInBytes
12121212-3434-5656-7878-abababababab,"SELECT * FROM ""foo"".""bar"" limit 10",s3://aws-athena-query-results-xxxxxxx/Unsaved/2018/09/03/12121212-3434-5656-7878-abababababab.csv,foo,SUCCEEDED,,2018-09-03 18:13:37.872000+09:00,2018-09-03 18:13:38.433000+09:00,341,60435
23232323-4545-6767-8989-cdcdcdcdcdcd,"SELECT * FROM ""foo"".""baz"" limit 10",s3://aws-athena-query-results-xxxxxxx/Unsaved/2018/09/03/23232323-4545-6767-8989-cdcdcdcdcdcd.csv,foo,SUCCEEDED,,2018-09-03 18:13:36.250000+09:00,2018-09-03 18:13:36.781000+09:00,349,62145
abababab-cdcd-efef-ghgh-ijijijijijij,"select * from hogehoge by id asclimit 1,s3://aws-athena-query-results-xxxxxxx/Unsaved/2018/09/03/abababab-cdcd-efef-ghgh-ijijijijijij.csv,hogehoge,FAILED,SYNTAX_ERROR: line 5:10: Column 'xxx' cannot be resolved,2018-09-03 17:53:18.777000+09:00,2018-09-03 17:53:19.740000+09:00,607,0

定期実行する場合は、CloudWatch Events と連携して日次実行するように設定してください。