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

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

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

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

#!/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