aws-cli Athena/Lambdaのラッパースクリプトを書いて作業を効率化する

はじめに

この記事は@ktatさんのaws-cli のラッパースクリプトを書いて作業を超効率化するを見て、AthenaとLambdaのラッパースクリプトを作ってみたものになります。

現状

業務上Amazon AthenaやAWS Lambdaをよく使うのですが、「クエリを書いて結果を見たい」「関数を実行させたい」場合に逐一コンソールへ入ってそれを行うのが面倒でした。

また、aws cliも

Athenaの場合

# 一々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で終わっていたら結果が返ってくる

Lambdaの場合

# --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ファイルのクエリ」を読み取って結果を返してくれます

aws-athena.sh

#!/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ファイルのクエリを読む

test_query.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

「関数名一覧を表示する」「特定の関数を実行する」ことができます。

aws-lambda.sh

#!/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で見れるので、実行命令だけ飛ばすようなものを作ろうかと

また、見よう見まねで行ったので、おかしな部分があると思います。
ご指摘などあればコメントいただけますと嬉しいです。

ここまで見ていただきありがとうございました。