Zenn
🐇

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

2020/11/10に公開
AWS
Bash
S3
athena
tech

以下をシェルスクリプトで実装します。

  • Athenaにクエリを投げます。
  • 出力結果のCSVをS3からダウンロードします。

尚、前提としてAthenaのworkspace及びdatabase,table等のセットアップは済んでいるものとします。

ディレクトリ構造

run.shをDockerコンテナ内で実行することを想定以下のようなディレクトリ構成となっています

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

queriesにはAthenaに投げるクエリのSQLファイルを配置します。csvはCSVファイルの格納フォルダです。

準備

Athenaに投げるSQL

./queriesにSQLファイルを配置し、Athenaに投げるクエリを記載します。Athenaでは日付でパーティションが切られていることが多いと思うのでWHERE句を変数にして置換できるようにしておきます。

./queries/my_tables.sql
SELECT * FROM my_database.my_table WHERE date = :DATE_CONDITION

シェルスクリプト

シェルスクリプトを実装します。以下は完成したものです。

run.sh
#!/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

以下、処理概要となります。

  • ./queries/*.sqlからクエリを読み込みます。このときsedコマンドで変数を置換します。
  • 読み込んだクエリをAthenaに投げ、処理が完了するまで待機します。
  • 処理が完了したらS3から実行結果のCSVファイルをダウンロードします。

Dockerfile

Dockerfileです。用意したクエリとシェルスクリプト、必要なコマンドラインツールをインストールしておきます。また、AWSのサービスに依存したスクリプトなのでベースイメージはamazonlinuxにしました。

Dockerfile
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"]

docker-compose.yml

docker-compose.ymlです。volumesにダウンロードしたCSVファイルの格納フォルダを指定します。また、環境変数には以下を指定します。

  • S3_BUCKET: CSV出力先のS3バケット
  • UNIX_TIMESTAMP: Athenaにクエリを投げる時に指定する日付パーティションの条件
  • AWS...: AWSのクレデンシャル
docker-compose.yml
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

NOTE

クエリがシンプルで処理が軽いうちはいいですが、ちょっと凝ったことをしたくなったり並列実行などしたくなった場合は、素直にSDKを使って別の言語で対応したほうがいいと思います。

おわり

Discussion