AWS Athena のクエリを同期的に実行するシェルスクリプト

.
├── docker-compose.yml
├── Dockerfile
├── queries
├── run.sh
└── csv

SELECT * FROM my_database.my_table WHERE date = :DATE_CONDITION

#!/usr/bin/env bash
set -eu

# Athenaのworkspace及びdatabase
WORK_GROUP=myapp
DATABASE=my_database

# UNIX_TIMESTAMPをフォーマットした文字列(yyyy-mm-dd)
CURRENT_DATE=$(date --date=@${UNIX_TIMESTAMP} +%Y-%m-%d)

# Athenaのクエリ結果(CSV)を出力するパス
QUERY_RESULT_PATH="s3://${S3_BUCKET}/my_database/${UNIX_TIMESTAMP}"

# Athenaのクエリが完了するまで待機する
wait_until_query_done() {
  QUERY_EXECUTION_ID=$1
  while :
  do
    QUERY_STATUS=$(aws athena get-query-execution \
      --query-execution-id ${QUERY_EXECUTION_ID} \
      | jq -r .QueryExecution.Status.State)
    echo "status of query which is executing on aws athena is $QUERY_STATUS"
    if [[ ${QUERY_STATUS} = "SUCCEEDED" ]] || \
        [[ ${QUERY_STATUS} = "FAILED" ]] || \
        [[ ${QUERY_STATUS} = "CANCELLED" ]]; then
      break;
    fi
    sleep 1
  done
}

# Athenaにクエリを投げる
# 返り値: AthenaのQUERY_EXECUTION_ID
start_query_on_athena() {
  TABLE=$1
  QUERY=$(sed -e "s/:DATE_CONDITION/\'${CURRENT_DATE}\'/g" < ./queries/${TABLE}.sql)

  QUERY_EXEC_ID=$(aws athena start-query-execution \
    --query-string "${QUERY}" \
    --query-execution-context Database=${DATABASE} \
    --result-configuration OutputLocation=${QUERY_RESULT_PATH}/${TABLE} \
    --work-group ${WORK_GROUP} \
    | jq -r '.QueryExecutionId')

  echo ${QUERY_EXEC_ID}
}

# Athenaにクエリを打って結果のCSVファイルをダウンロードする
download_files_from_s3() {
  TABLE=$1
  QUERY_EXEC_ID=$(start_query_on_athena $TABLE)
  echo "query_execution_id of query which is executing on aws athena is ${QUERY_EXEC_ID}"
  wait_until_query_done ${QUERY_EXEC_ID}
  aws s3 cp ${QUERY_RESULT_PATH}/${TABLE}/${QUERY_EXEC_ID}.csv /tmp/csv/${TABLE}.csv
}

download_files_from_s3 my_table

FROM amazonlinux:2

RUN yum update -y
RUN yum install -y sudo date jq
RUN curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip" && \
        unzip awscliv2.zip && \
        sudo ./aws/install

COPY ./queries ./queries
COPY ./run.sh ./run.sh

ENV AWS_DEFAULT_REGION ap-northeast-1
ENV TZ Asia/Tokyo

ENTRYPOINT ["sh", "run.sh"]

version: '3'
services:
  download-athena-result:
      build:
        context: ./
        dockerfile: ./Dockerfile
      image: download-athena-result
      environment:
        - S3_BUCKET
        - UNIX_TIMESTAMP
        - AWS_ACCESS_KEY_ID
        - AWS_SECRET_ACCESS_KEY
        - AWS_SESSION_TOKEN
      volumes: ./csv:/tmp/csv

export S3_BUCKET=my-bucket
export UNIX_TIMESTAMP=xxx
export AWS_ACCESS_KEY_ID=xxx
export AWS_SECRET_ACCESS_KEY=xxx
export AWS_SESSION_TOKEN=xxx

docker-compose up