AWS CLI で Athena のクエリ実行を同期的に行う #AWS

Athena に対するクエリを AWS CLI で実行する場合、一度のコマンドでクエリ結果を取得することは出来ません。

いくつかのコマンドを実行して結果を取得するような流れになってます。

クエリ結果取得までのコマンド実行の流れ

1. クエリ実行開始 ( start-query-execution )

start-query-execution \
  --query-string <value> \
  --result-configuration <value>

クエリは非同期的に処理されるため時間の掛かるクエリであってもコマンドの結果はすぐに返ってくる。
結果は QueryExecutionId (クエリ実行 ID) のみ。

2. クエリ実行状況の問い合わせ ( get-query-execution )

get-query-execution --query-execution-id <value>

start-query-execution で得られたクエリ実行 ID を渡してクエリの実行状況を問い合わせる。
クエリの実行ステータスを取得できる。
- SUBMITTED: 実行キューに追加された
- RUNNING: 実行中
- SUCCEEDED: 実行完了した
- CANCELLED: キャンセルされた
- FAILED: 失敗した
SUCCEEDED の場合は一緒に実行結果の出力先の S3 Path (.csv) が返ってくる。

3. クエリ結果取得 ( get-query-results )

get-query-results --query-execution-id <value>

start-query-execution で得られたクエリ実行 ID を引数に渡してクエリの実行結果を JSON format として取得する。
ページング取得に対応。
S3 へ出力された get-query-execution の実行結果 CSV を JSON 形式で取得するだけのコマンドのため、生の CSV 形式でもよければ S3 から直接取得すればよい。

つまりスクリプトでクエリを実行して結果を取得するには...

start-query-execution でクエリ実行を開始
定期的に get-query-execution を実行してクエリ実行完了を待つ
完了したら aws s3 cp もしくは get-query-results で実行結果を取得する

クエリ実行完了を待つコードを毎回書くのが地味に面倒なので関数化した

#!/usr/bin/env bash

# 同期的に Athena Query を実行します。
# params:
#   $1 string クエリ出力先パス
#   $2 string 実行するSQL
# stdout:
#   query_execution_result の結果 (JSON)
# return:
#   1 - Any error
#   2 - Query timed out
# see:
#   https://docs.aws.amazon.com/cli/latest/reference/athena/get-query-execution.html
function athena_query_execution_sync() {
  local output_path="$1"
  local sql="$2"

  # クエリ実行開始
  query_id=$(aws athena start-query-execution \
    --result-configuration OutputLocation="$output_path" \
    --query-string "$sql" | jq -r '.QueryExecutionId')

  local max_retry=10
  local fetch_interval_seconds=10

  for ((i = 0; i < $max_retry; i++)); do
    # クエリ実行完了を待つため一定時間待機
    sleep ${fetch_interval_seconds}

    # クエリの実行状況を問い合わせ
    query_execution_result=$(aws athena get-query-execution --query-execution-id ${query_id})
    query_state=$(echo ${query_execution_result} | jq -r '.QueryExecution.Status.State')

    # 想定されるクエリ実行ステータス:
    # QUEUED, RUNNING, SUCCEEDED, FAILED, CANCELLED
    case ${query_state} in
      CANCELLED | FAILED)
        return 1
        ;;
      SUCCEEDED)
        echo "${query_execution_result}"
        return 0
        ;;
    esac
  done

  # Timed out
  return 2
}

# エラーメッセージを表示してスクリプトを終了します。
# params:
#   $1 string エラーメッセージ
# return:
#   1 - Any time
function error_exit() {
  echo "[ERROR] $1" >&2
  exit 1
}

# Athena のクエリ結果出力先
readonly OUTPUT_LOCATION="s3://aws-athena-query-results/"
# 任意の実行したい SQL
readonly SQL="SELECT 1"

# クエリを実行
query_result="$(athena_query_execution_sync "${OUTPUT_LOCATION}" "${SQL}")"
case $? in
  1) error_exit "An error occurred...";;
  2) error_exit "Timed out...";;
esac

# クエリ実行結果 (CSV) をダウンロード
result_s3_path="$(echo "${query_result}" | jq -r '.QueryExecution.ResultConfiguration.OutputLocation')"
aws s3 cp "${result_s3_path}" results.csv