MWAA の制御のスクリプト化

前回のテック記事から期間が空いてしまいましたが、
継続は力なりということで久方振りに記事を書きます。

概要

MWAA（Managed Workflows for Apache Airflow）は、AWSが提供するAirflowのマネージドサービスです。
Airflowは、ワークフローの定義とスケジューリングを行うためのプラットフォームです。

MWAAの設定変更やスケール変更のタイミングなどで
MWAA上で稼働しているDAGを一時停止したり、再開したいことがあります。

MWAAはWEBコンソールからの操作が可能ですが、
多数のDAGを操作したい場合、一つ一つ変更する手間があるのと、
対象を誤る懸念があるため、スクリプト化して一括で操作できるようにしました。

スクリプト

MWAAをCLIで制御するためには一時的なCLIトークンが必要ですが、
有効期限が60秒と短いため、変更リクエスト毎にトークンを発行するようにしています。

• mwaa_control.sh

#!/bin/sh

# このシェルは、MWAAのDAGを操作するためのシェルです。
# 第一引数にMWAA名　
# 第二引数に unpause(起動)、pause(停止)　を指定してください。

export ENVIRONMENT_NAME=$1

# コマンド指定
Command=$2

# スクリプトのディレクトリを取得
SCRIPT_DIR=$(dirname "$0")
# 対象リストファイル
export Targetfile="${SCRIPT_DIR}/$1.list"

while read line ;do
    export WEB_SERVER_HOSTNAME=$(aws mwaa create-cli-token --name ${ENVIRONMENT_NAME} | jq -r .WebServerHostname)
    export CLI_TOKEN=$(aws mwaa create-cli-token --name ${ENVIRONMENT_NAME} | jq -r .CliToken)
    echo "[INFO] ==== ${line} ${Command} ==== "
    # curlの出力を変数に保存
    response=$(curl -s --request POST "https://${WEB_SERVER_HOSTNAME}/aws_mwaa/cli" \
         --header "Authorization: Bearer ${CLI_TOKEN}" \
         --header "Content-Type: text/plain" \
         --data-raw "dags ${Command} ${line}")
    
    # 変数をjqでパース
    echo $response | jq -r .stdout | base64 -d
    sleep 1
done < ./${Targetfile}

• <MWAA名>.list

制御対象のDAG名を記載します。

sample_dag1
sample_dag2

使い方

./mwaa_control.sh <MWAA名> <コマンド>

実行例

• 一時停止

./mwaa_control.sh sample_airflow pause

• 再開

./mwaa_control.sh sample_airflow unpause

まとめ

GUI操作と比較し操作時間としては差異がないかもしれません。
ただ、スクリプト化することで毎回同じ操作で結果に差異が出ないことや、
誤操作を防ぐことができるため、運用の安定性向上に繋がると考えています。