はじめに
この記事は@ktatさんのaws-cli のラッパースクリプトを書いて作業を超効率化するを見て、AthenaとLambdaのラッパースクリプトを作ってみたものになります。
現状
業務上Amazon AthenaやAWS Lambdaをよく使うのですが、「クエリを書いて結果を見たい」「関数を実行させたい」場合に逐一コンソールへ入ってそれを行うのが面倒でした。
また、aws cliも
# 一々execution idを取得した後、そのステータスを確認し、終了していたらget-query-resultsを実行して結果を得る
$ aws athena start-query-execution --query-string --そのほかの設定
$ aws athena get-query-execution --query-execution-id "start-query-executionで発行されたID"
$ aws athena get-query-results --query-execution-id "start-query-executionで発行されたID" #処理がsucceededで終わっていたら結果が返ってくる
# --payloadに入れる値はjson形式の文字列で入れることがほとんどなので、--cli-binary-format raw-in-base64-outがほぼ必須。cliでresponseを受け取りたい場合--cli-read-timeout 0(無制限)がほぼ必須
$ aws lambda invoke --function-name "function-name" --payload '{"key1":"value1"}' --cli-binary-format raw-in-base64-out --cli-read-timeout 0
と、どちらもあまり効率が良くないですし、コマンドも忘れやすいです。
スクリプト
ということで以下を作りました
- aws-athena -> aws athenaのラッパー
- aws-lambda -> aws lambdaのラッパー
aws-athena
「コマンドに打ち込んだクエリ文字列」 or 「.sqlファイルのクエリ」を読み取って結果を返してくれます
#!/bin/bash
COMMAND=$1;
QUERY=$2;
help() {
echo
echo $0 ... aws athena wrapper command
echo
echo "$0 query [query string] ... execution and get result the query"
echo "$0 file [.sql file] ... execution and get result from the .sql file"
echo
exit 1
}
get_query_results() {
local sql_query=$1
# クエリ実行の試行とエラーハンドリング
execution_response=$(aws athena start-query-execution --query-string "$sql_query" --output json 2>&1)
if echo "$execution_response" | grep -q "InvalidRequestException"; then
echo "Error starting query execution: $execution_response"
exit 1
fi
# クエリ実行IDの抽出
execution_id=$(echo "$execution_response" | jq -r '.QueryExecutionId')
echo "Query Execution ID: $execution_id"
# クエリ実行結果がSUCCEEDEDになったら結果を表示
while true; do
execution_result=$(aws athena get-query-execution --query-execution-id "$execution_id" --output json)
status=$(echo $execution_result | jq -r '.QueryExecution.Status.State')
if [ "$status" = "SUCCEEDED" ]; then
echo "Query succeeded. Fetching results..."
result=$(aws athena get-query-results --query-execution-id "$execution_id" --output json)
header=$(echo "$result" | jq -r '.ResultSet.ResultSetMetadata.ColumnInfo | map(.Label) | @tsv' | paste -sd '\t')
data=$(echo "$result" | jq -r '.ResultSet.Rows[1:][] | .Data | map(.VarCharValue) | @tsv')
output="$header\n$data"
echo -e "$output" | column -s $'\t' -t
break
elif [ "$status" = "FAILED" ]; then
echo "Query failed."
echo "(echo $execution_result | jq '.QueryExecution.Status.StateChangeReason')"
break
elif [ "$status" = "CANCELLED" ]; then
echo "Query was cancelled."
break
else
echo "Query is still running. Retrying in 1 second..."
sleep 1
fi
done
}
# コマンドがサポートしている文字列で打たれているか
if [ "$COMMAND" != "query" ] && [ "$COMMAND" != "file" ]; then
echo "COMMAND is required as 1st arg: query/file";
help;
fi
# queryコマンドの場合、次の引数にクエリがあるか
if [ "$COMMAND" = "query" ]; then
if [ "$QUERY" = "" ]; then
echo "query requires second arg: query sentence";
help;
else
get_query_results "$QUERY"
fi
fi
# fileコマンドの場合、次の引数に.sqlファイルが指定されているか
if [ "$COMMAND" = "file" ]; then
if [[ "$QUERY" != *.sql ]]; then
echo "file requires second arg: .sql file";
help;
else
# SQLファイルからクエリを読み取る
sql_query=$(cat "$QUERY")
get_query_results "$sql_query"
fi
fi
実際に使ってみます
- queryオプションでコマンドに直接クエリを書く
$ bash ./aws-athena.sh query "select * from awsdatacatalog.sampledb.elb_logs limit 1"
Query Execution ID: d6358ffc-4875-4274-b56a-2a5183bb1e74
Query is still running. Retrying in 1 second...
Query succeeded. Fetching results...
request_timestamp elb_name request_ip request_port backend_ip backend_port request_processing_time backend_processing_time client_response_time elb_response_code backend_response_code received_bytes sent_bytes request_verb url protocol user_agent ssl_cipher ssl_protocol
2015-01-04T04:00:00.516940Z elb_demo_005 244.139.233.91 13307 172.30.133.232 443 0.001177 9.59E-4 0.001703 200 200 0 2914 GET https://www.example.com/articles/856 HTTP/1.1 "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11_6) AppleWebKit/602.1.50 (KHTML, like Gecko) Version/10.0 Safari/602.1.50" DHE-RSA-AES128-SHA TLSv1.2
- fileオプションで.sqlファイルのクエリを読む
select
request_timestamp as "りくえすとたいむすたんぷ"
from
"awsdatacatalog"."sampledb"."elb_logs"
limit 10
$ bash ./aws-athena.sh file test_query.sql
Query Execution ID: 8f1ebbcd-f8c3-465d-bec4-dea9117b2515
Query is still running. Retrying in 1 second...
Query succeeded. Fetching results...
りくえすとたいむすたんぷ 2015-01-04T16:00:01.206255Z
2015-01-04T16:00:01.612598Z
2015-01-04T16:00:02.793335Z
2015-01-04T16:00:03.068897Z
2015-01-04T16:00:03.470121Z
2015-01-04T16:00:04.159502Z
2015-01-04T16:00:04.778187Z
2015-01-04T16:00:06.178798Z
2015-01-04T16:00:06.607063Z
2015-01-04T16:00:06.625672Z
めちゃくちゃ便利ですね!
statusに応じて処理を分岐しているので、構文エラーがあった場合も分かります
- 構文エラー
$ bash ./aws-athena.sh query "select aaaa from awsdatacatalog.sampledb.elb_logs limit 1"
Query Execution ID: fdd58ad3-3756-47b7-b1db-1f55844c749b
Query failed.
(echo {
"QueryExecution": {
"QueryExecutionId": "fdd58ad3-3756-47b7-b1db-1f55844c749b",
"Query": "select aaaa from awsdatacatalog.sampledb.elb_logs limit 1",
"StatementType": "DML",
"ResultConfiguration": {
"OutputLocation": "s3://xxxxxxxx/primary_folder/fdd58ad3-3756-47b7-b1db-1f55844c749b.csv"
},
"QueryExecutionContext": {},
"Status": {
"State": "FAILED",
"StateChangeReason": "COLUMN_NOT_FOUND: line 1:8: Column 'aaaa' cannot be resolved or requester is not authorized to access requested resources",
"SubmissionDateTime": "2024-09-21T13:22:54.254000+09:00",
"CompletionDateTime": "2024-09-21T13:22:54.626000+09:00",
"AthenaError": {
"ErrorCategory": 2,
"ErrorType": 1006,
"Retryable": false,
"ErrorMessage": "COLUMN_NOT_FOUND: line 1:8: Column 'aaaa' cannot be resolved or requester is not authorized to access requested resources"
}
},
"Statistics": {
"EngineExecutionTimeInMillis": 177,
"DataScannedInBytes": 0,
"TotalExecutionTimeInMillis": 372,
"QueryQueueTimeInMillis": 70,
"ServiceProcessingTimeInMillis": 60
},
"WorkGroup": "primary",
"EngineVersion": {
"SelectedEngineVersion": "Athena engine version 3",
"EffectiveEngineVersion": "Athena engine version 3"
}
}
} | jq '.QueryExecution.Status.StateChangeReason')
これで気軽にテーブルの中身を見たい時や、.sql
ファイルのクエリを実行したいときに時間を掛けることが無くなりました。
aws-lambda
「関数名一覧を表示する」「特定の関数を実行する」ことができます。
#!/bin/bash
COMMAND=$1;
FUNCTION_NAME=$2;
PAYLOAD=$3;
help() {
echo
echo $0 ... aws lambda wrapper command
echo
echo "$0 list ... list lambda functions"
echo "$0 invoke [lambda function name] [json format payload] ... invoke lambda function"
echo
exit 1
}
# コマンドがサポートしている文字列で打たれているか
if [ "$COMMAND" != "list" ] && [ "$COMMAND" != "invoke" ]; then
echo "COMMAND is required as 1st arg: list/invoke";
help;
fi
# invokeコマンドの場合、次の引数にLambda関数名とPayloadがあるか
if [ "$COMMAND" = "invoke" ]; then
if [ "$FUNCTION_NAME" = "" ]; then
echo "invoke requires second arg: lambda function name"
help;
elif [ "$PAYLOAD" = "" ]; then
echo "invoke requires third arg: lambda payload json"
help;
elif ! echo "$PAYLOAD" | jq empty >/dev/null 2>&1; then
echo "third arg json format invalid"
exit 1
fi
fi
# listコマンド
if [ "$COMMAND" = "list" ]; then
list_result=$(aws lambda list-functions | jq -r '.Functions[] | .FunctionName')
echo "$list_result"
# invokeコマンド
else
timestamp=$(date +%Y%m%d_%H%M%S)
response_file_name="response_${timestamp}.json"
payload_json=$(echo "$PAYLOAD" | jq -c .)
invoke_result=$(aws lambda invoke --function-name "$FUNCTION_NAME" --payload "$payload_json" --cli-binary-format raw-in-base64-out --cli-read-timeout 0 "$response_file_name")
echo "AWS CLI Output:"
echo "$invoke_result" | jq .
echo
echo "Lambda Response:"
jq '.' "$response_file_name"
rm "$response_file_name"
fi
- listコマンドで関数一覧を表示
$ bash ./aws-lambda.sh list
athena_query_scanning
athena_data_insert_test
slack_notification
graph_plot_to_slack
delete_s3_objects
athena_droptable_detection
create_parquet_datasource
slack_webhook_test
yuzu_data_get
get_secret_value
container_test
sync_github_codecommit
function_log_test
s3_data_transfer
plotly_html_graph_plot
plot_test
athena_query_scan
athena_query_execute
graph_notification
yuzu_data_upload
- invokeコマンドで関数を実行
# 関数は単純にkey1の値を返す関数
$ bash aws-lambda.sh invoke function_log_test '{"key1":"aaa"}'
AWS CLI Output:
{
"StatusCode": 200,
"ExecutedVersion": "$LATEST"
}
Lambda Response:
{
"statusCode": 200,
"body": "aaa"
}
コチラも良いですね!!
Lambda関数の実行による詳細はaws-logsのラッパースクリプトを使うと良いと思います。
$ bash aws-log tail /aws/lambda/function_log_test 10
tail since: 2024-09-21T04:30:01 (UTC)
2024-09-21T04:32:26.790000+00:00 2024/09/21/[$LATEST]xxxxxxxxxxxxxxxxxxxxxxxx INIT_START Runtime Version: python:3.11.v39 Runtime Version ARN: arn:aws:lambda:ap-northeast-1::runtime:yyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyy
2024-09-21T04:32:26.866000+00:00 2024/09/21/[$LATEST]xxxxxxxxxxxxxxxxxxxxxxxx START RequestId: zzzzzzzzzzzzzzzzzzzzzzzzzzzz Version: $LATEST
2024-09-21T04:32:26.866000+00:00 2024/09/21/[$LATEST]xxxxxxxxxxxxxxxxxxxxxxxx [INFO] 2024-09-21T04:32:26.866Z zzzzzzzzzzzzzzzzzzzzzzzzzzzz Lambda function has started
2024-09-21T04:32:26.866000+00:00 2024/09/21/[$LATEST]xxxxxxxxxxxxxxxxxxxxxxxx [INFO] 2024-09-21T04:32:26.866Z zzzzzzzzzzzzzzzzzzzzzzzzzzzz Received event: {"key1": "aaa"}
2024-09-21T04:32:26.866000+00:00 2024/09/21/[$LATEST]xxxxxxxxxxxxxxxxxxxxxxxx [INFO] 2024-09-21T04:32:26.866Z zzzzzzzzzzzzzzzzzzzzzzzzzzzz Value for 'key1': aaa
2024-09-21T04:32:26.866000+00:00 2024/09/21/[$LATEST]xxxxxxxxxxxxxxxxxxxxxxxx [INFO] 2024-09-21T04:32:26.866Z zzzzzzzzzzzzzzzzzzzzzzzzzzzz Processing succeeded
2024-09-21T04:32:26.868000+00:00 2024/09/21/[$LATEST]xxxxxxxxxxxxxxxxxxxxxxxx END RequestId: zzzzzzzzzzzzzzzzzzzzzzzzzzzz
2024-09-21T04:32:26.868000+00:00 2024/09/21/[$LATEST]xxxxxxxxxxxxxxxxxxxxxxxx REPORT RequestId: zzzzzzzzzzzzzzzzzzzzzzzzzzzz Duration: 2.48 ms Billed Duration: 3 ms Memory Size: 128 MB Max Memory Used: 33 MB Init Duration: 74.28 ms
う~ん便利
さいごに
今回紹介したコードは以下で管理しています。
今後はGlue JobとBatchのSubmit Jobを作りたいな~と思ってます。
これら2つもLambda同様結果はaws-logsで見れるので、実行命令だけ飛ばすようなものを作ろうかと
また、見よう見まねで行ったので、おかしな部分があると思います。
ご指摘などあればコメントいただけますと嬉しいです。
ここまで見ていただきありがとうございました。