ローカルPC上でAirflow環境を実行するためのスクリプト

airflowctl.sh

#!/bin/sh
# ローカルPC上でAirflow環境を実行するためのスクリプト
#
# 環境変数PYTHON_BINに、Python3.7系のpythonのパスを指定する必要があります.
# macosの場合:
#  export PYTHON_BIN=/Library/Frameworks/Python.framework/Versions/3.7/bin/python3
# pyenv(Python3.7.9)の場合:
#  export PYTHON_BIN=~/.pyenv/versions/3.7.9/bin/python3
#
# ディレクトリ構造
# - PROJECT_DIR
#    - dags
#       - foo_dag.py
#       - bar_dag.py
#    - tests
#       - test_foo.py
#       - test_bar.py
#    - airflowctl.sh
#    - requirements.txt
#    - venv

# Airflowの環境設定、必要に応じて追記して下さい
export AIRFLOW__CORE__LOAD_EXAMPLES=False
# Airflowの環境設定、ここまで

AIRFLOW_VERSION=2.0.2
PYTHON_BIN=${PYTHON_BIN:-python3}
PROJECT_DIR=$(cd $(dirname $0); pwd)
export AIRFLOW_HOME=$PROJECT_DIR

init() {
  # venvによる仮想環境作成
  if [ ! -e $PROJECT_DIR/venv ]; then
    # Pythonのversionが3.7系か否かをチェック
    $PYTHON_BIN --version | grep 3.7 > /dev/null
    if [ $? != 0 ]; then
      echo "Error: invlid python version, please use 3.7"
      exit 1
    fi
    $PYTHON_BIN -m venv $PROJECT_DIR/venv
  fi

  source $PROJECT_DIR/venv/bin/activate

  # requirementsファイルの作成
  if [ ! -e $PROJECT_DIR/requirements.txt ]; then
    PYTHON_VERSION="$(python --version | cut -d " " -f 2 | cut -d "." -f 1-2)"
    CONSTRAINT_URL="https://raw.githubusercontent.com/apache/airflow/constraints-${AIRFLOW_VERSION}/constraints-${PYTHON_VERSION}.txt"
    pip install "apache-airflow==${AIRFLOW_VERSION}" --constraint "${CONSTRAINT_URL}"
    pip freeze > $PROJECT_DIR/requirements.txt
  else
    pip install -r $PROJECT_DIR/requirements.txt
  fi

  pip install pytest

  # airflow.dbの初期化
  rm -f $AIRFLOW_HOME/airflow.db
  rm -f $AIRFLOW_HOME/airflow.cfg
  rm -f $AIRFLOW_HOME/webserver_config.py
  airflow db init
  airflow users create \
    --username admin \
    --password admin \
    --firstname Peter \
    --lastname Parker \
    --role Admin \
    --email spiderman@superhero.org
}

clean() {
  rm -rf $PROJECT_DIR/venv
}

unittest() {
  source $PROJECT_DIR/venv/bin/activate
  cd $PROJECT_DIR && pytest
}

run_airflow() {
  source $PROJECT_DIR/venv/bin/activate
  airflow scheduler &
  PID_SCHEDULER=$!
  airflow webserver &
  PID_WEBSERVER=$!

  sleep 1
  ps -p $PID_SCHEDULER > /dev/null || exit 1
  ps -p $PID_WEBSERVER > /dev/null || exit 1

  trap 'kill $PID_WEBSERVER && kill $PID_SCHEDULER && sleep 3 && exit 0 || exit 1' SIGINT
  while true; do
    sleep 1
  done
}

# main routine
case $1 in
  init)
    init
    ;;
  clean)
    clean
    ;;
  test)
    unittest
    ;;
  run)
    run_airflow
    ;;
  *)
    echo "usage: airflowctl.sh [init|clean|test|run]"
esac