whirl

#!/usr/bin/env bash
AIRFLOW_UI_PORT=5000
# https://stackoverflow.com/a/4774063
SCRIPT_DIR="$( cd -- "$(dirname "$0")" >/dev/null 2>&1 ; pwd -P )"

set -e
# load and export environment variables
# priority order:
#     .whirl.env in home directory
#     .whirl.env in whirl repository directory
#     .whirl.env in environment directory
#     .whirl.env in dag directory
function export_environment_vars() {
  set -a
  DOCKER_CONTEXT_FOLDER=${SCRIPT_DIR}/docker
  DAG_FOLDER=$(pwd)
  PROJECTNAME=$(basename "${DAG_FOLDER}")
  WHIRL_INITIATOR=$(whoami)
  WHIRL_SETUP_FOLDER=/etc/airflow/whirl.setup.d

  if [ -f ~/.whirl.env ]; then
    echo "Reading configuration from ~/.whirl.env"
    . ~/.whirl.env
  fi
  if [ -f "${SCRIPT_DIR}/.whirl.env" ]; then
    echo "Reading configuration from ${SCRIPT_DIR}/.whirl.env"
    . "${SCRIPT_DIR}/.whirl.env"
  fi

  # determine whether to use the environment set at the commandline or
  # in the DAG FOLDER .whirl.env
  if [ -z "${WHIRL_ENVIRONMENT_ARG}" ]; then
    if [ -f "${DAG_FOLDER}/.whirl.env" ]; then
      if grep -qRE "^(\s+)?WHIRL_ENVIRONMENT=(.*)" "${DAG_FOLDER}/.whirl.env"; then
          echo "Found WHIRL_ENVIRONMENT in ${DAG_FOLDER}/.whirl.env";
          WHIRL_ENVIRONMENT=$(grep -oRE "^(\s+)?WHIRL_ENVIRONMENT=(.*)" "${DAG_FOLDER}/.whirl.env" | \
                              sed -e 's/.*=\(.*\)/\1/g')
          echo "environment is ${WHIRL_ENVIRONMENT}"
      fi
    fi
  else
    WHIRL_ENVIRONMENT=${WHIRL_ENVIRONMENT_ARG}
  fi

  # determine whether to use the environment dir set at the commandline or
  # in the DAG FOLDER .whirl.env
  echo ""
  if [ -z "${WHIRL_ENVIRONMENT_DIR_ARG}" ]; then
    if [ -f "${DAG_FOLDER}/.whirl.env" ]; then
      if grep -qRE "^(\s+)?WHIRL_ENVIRONMENT_DIR=(.*)" "${DAG_FOLDER}/.whirl.env"; then
          echo "Found WHIRL_ENVIRONMENT_DIR in ${DAG_FOLDER}/.whirl.env";

          grep -oRE "^(\s+)?WHIRL_ENVIRONMENT_DIR=(.*)" "${DAG_FOLDER}/.whirl.env";
          grep -oRE "^(\s+)?WHIRL_ENVIRONMENT_DIR=(.*)" "${DAG_FOLDER}/.whirl.env" | sed -e 's/.*=\(.*\)/\1/g';
          WHIRL_ENVIRONMENT_DIR=$(grep -oRE "^(\s+)?WHIRL_ENVIRONMENT_DIR=(.*)" "${DAG_FOLDER}/.whirl.env" | \
                              sed -e 's/.*=\(.*\)/\1/g')
          echo "environment directory is ${WHIRL_ENVIRONMENT_DIR}"
      fi
    fi
  else
    WHIRL_ENVIRONMENT_DIR=${WHIRL_ENVIRONMENT_DIR_ARG}
  fi

  if [[ -z "${WHIRL_ENVIRONMENT_DIR}" ]]; then
    WHIRL_ENVIRONMENT_DIR=${SCRIPT_DIR}/envs
  fi
  ENVIRONMENT_FOLDER=${WHIRL_ENVIRONMENT_DIR}/${WHIRL_ENVIRONMENT};

  if [[ -z "${WHIRL_ENVIRONMENT}" || ! -d ${ENVIRONMENT_FOLDER} ]]; then
    echo "No valid environment '${WHIRL_ENVIRONMENT}' specified"
    exit 2;
  fi

  if [ -f "${ENVIRONMENT_FOLDER}/.whirl.env" ]; then
    echo "Reading configuration from ${ENVIRONMENT_FOLDER}/.whirl.env"
    . "${ENVIRONMENT_FOLDER}/.whirl.env"
  fi

  if [ -f "${DAG_FOLDER}/.whirl.env" ]; then
    echo "Reading configuration from ${DAG_FOLDER}/.whirl.env"
    . "${DAG_FOLDER}/.whirl.env"
  fi

  # in case DAG_FOLDER specifies WHIRL_ENVIRONMENT, commandline needs to overrule if set
  if [[ -n "${WHIRL_ENVIRONMENT_ARG}" ]]; then
    WHIRL_ENVIRONMENT=${WHIRL_ENVIRONMENT_ARG}
  fi
  set +a
}

detect_potential_dag() {
  test "$(find . -type f -name '*.py' -o -name '*.zip' | wc -l)" -gt 0
}

test_dag_state() {
  local DAG_ID=$1
  local STATE=$2
  local SHOW_RESULT=${3:-false}
  echo "Checking dag state(${STATE}) for ${DAG_ID}" >&2
  result=$(curl -s -u admin:admin "http://localhost:${AIRFLOW_UI_PORT}/api/v1/dags/${DAG_ID}/dagRuns" | jq ".dag_runs | map(select(.state == \"${STATE}\")) | length")
  echo "Result of checking dag state(${STATE}) for ${DAG_ID} is ${result}" >&2
  if [ "${SHOW_RESULT}" == true ]; then
    output=$(curl -s -u admin:admin "http://localhost:${AIRFLOW_UI_PORT}/api/v1/dags/${DAG_ID}/dagRuns" | jq ".dag_runs" )
    echo "Dagruns: $output" >&2
    output=$(curl -s -u admin:admin "http://localhost:${AIRFLOW_UI_PORT}/api/v1/dags/${DAG_ID}/dagRuns/~/taskInstances" | jq ".task_instances" )
    echo "Tasks: $output" >&2
  else
    echo "$result"
  fi
}

run_compose_setup_scripts() {
  if [[ -d ${ENVIRONMENT_FOLDER}/compose.setup.d ]]; then
    echo "============================================"
    echo "== Setup docker-compose specifics =========="
    echo "============================================"
    for filename in "${ENVIRONMENT_FOLDER}"/compose.setup.d/*.sh; do
      echo "Executing compose prepare script: $filename"
      . "$filename"
    done
  fi
}

start() {
    echo "Starting airflow local run for environment ${WHIRL_ENVIRONMENT}"

    # Possible cleanup before starting
    run_compose_setup_scripts

    docker build --build-arg PYTHON_VERSION="${PYTHON_VERSION}" --build-arg AIRFLOW_VERSION="${AIRFLOW_VERSION}" -t "docker-whirl-airflow:py-${PYTHON_VERSION}-local" "${SCRIPT_DIR}/docker/airflow-python"

    DAEMON=""
    if [ "${CI_MODE}" == true ]; then
      if [ -z "${DAG_ID}" ]; then
        DAG_ID=$(grep -Eo "dag_id=['\"](.*)['\"]" "${DAG_FOLDER}/dag.py" | sed -E "s/.*=['\"](.*)['\"]/\1/")
        if [ -z "${DAG_ID}" ]; then
          echo "Unable to determine dag id for CI_MODE and DAG_ID not provided"
          exit 4
        fi
      fi
      DAEMON="-d"
      export AIRFLOW__API__AUTH_BACKEND=airflow.api.auth.backend.basic_auth
    fi

    docker-compose -f "${ENVIRONMENT_FOLDER}/docker-compose.yml" up ${DAEMON} --build --force-recreate

    if [ "${CI_MODE}" == true ]; then
      while [[ "$(curl -s -o /dev/null -w %\{http_code\} http://localhost:${AIRFLOW_UI_PORT})" != "302" ]]; do
        echo "Waiting for Airflow UI to come up..."
        sleep 10;
      done
      if [[ "$(curl -s -o /dev/null -w %\{http_code\} -X PATCH  -H 'Content-Type: application/json' --user admin:admin "http://localhost:${AIRFLOW_UI_PORT}/api/v1/dags/${DAG_ID}?update_mask=is_paused" -d '{ "is_paused": false }')" != "200" ]]; then
        echo "Unable to unpause dag with id ${DAG_ID}."
        echo "Command: curl -X PATCH  -H 'Content-Type: application/json' --user admin:admin \"http://localhost:${AIRFLOW_UI_PORT}/api/v1/dags/${DAG_ID}?update_mask=is_paused\" -d '{ \"is_paused\": false }'"
        echo "Output:"
        curl -X PATCH  -H 'Content-Type: application/json' --user admin:admin "http://localhost:${AIRFLOW_UI_PORT}/api/v1/dags/${DAG_ID}?update_mask=is_paused" -d '{ "is_paused": false }'
        stop
        exit 3
      else
        echo "Unpaused dag with id ${DAG_ID}."
        sleep 5;
        while [ "$(test_dag_state "${DAG_ID}" "running")" -gt 0 ]; do
          echo "Dag '${DAG_ID}' is (still) running..."
          sleep 5;
        done
        SUCCES_STATE="success"
        FAILURE_STATE="failed"
        if [ "${WHIRL_CI_EXPECT_FAILURE}" == true ]; then
          SUCCES_STATE="failed"
          FAILURE_STATE="success"
        fi
        if [ "$(test_dag_state "${DAG_ID}" "${SUCCES_STATE}")" -ge 1 ]; then
          echo "Dag '${DAG_ID}' run(s) successfully finished"
          stop
          exit 0
        fi
        if [ "$(test_dag_state "${DAG_ID}" "${FAILURE_STATE}")" -ge 1 ]; then
          echo "Dag '${DAG_ID}' run(s) failed"
          # output dagrun result
          $(test_dag_state "${DAG_ID}" "${FAILURE_STATE}" true)
          if [ "$STOP_ON_FAILURE" = true ]; then
            stop
            exit 1
          fi
        fi
      fi
    fi

    echo "Closing down the environment"
    docker-compose -f "${ENVIRONMENT_FOLDER}/docker-compose.yml" down
}

stop() {
  echo "Stopping airflow-localrun containers..."
  docker-compose -f "${ENVIRONMENT_FOLDER}/docker-compose.yml" down --volumes --rmi local --remove-orphans
}

logs() {
  echo "Showing logs of airflow-localrun service ${WHIRL_SERVICE_NAME}"
  docker-compose -f "${ENVIRONMENT_FOLDER}/docker-compose.yml" logs -f "${WHIRL_SERVICE_NAME}"
}

usage() {
  echo "usage: ${BASH_SOURCE[0]} [-h|--help] [-e|--environment env] [start|stop|ci]"
  echo "  -h|--help                          display usage"
  echo "  -e|--environment environment       specify environment to use"
  echo "  -d|--directory environment_folder  specify the folder that contains the environments (defaults to SCRIPT_DIR)"
  echo "  -l|--logs servicename              tail the logs of the service"
  echo "  start|stop                         start or stop all"
  echo "  ci                                 runs in daemonized mode and awaits dag run completion"
  echo "  -i|--dag_id dag_id                 specify the dag_id to check (only in ci mode)"
  exit 1
}

function read_arguments() {
  while [[ $# -gt 0 ]]
  do
      key="${1}"
      case ${key} in
      -e|--environment)
          WHIRL_ENVIRONMENT_ARG="${2}"
          shift # past argument
          shift # past value
          ;;
      -d|--directory)
          WHIRL_ENVIRONMENT_DIR_ARG="${2}"
          shift # past argument
          shift # past value
          ;;
      -l|--logs)
          LOGS=true
          WHIRL_SERVICE_NAME="${2}"
          shift # past argument
          shift # past value
          ;;
      -i|--dag_id)
          DAG_ID="${2}"
          shift # past argument
          shift # past value
          ;;
      start)
          CI_MODE=false
          shift
          ;;
      ci)
          CI_MODE=true
          shift # past argument
          ;;
      stop)
          STOP=true
          shift # past argument
          ;;
      -h|--help)
          usage
          ;;
      *)  # unknown option
          echo "WARNING: Skipping unknown commandline argument: '${key}'"
          shift # past argument
          ;;
      esac
  done
}

function main() {
  read_arguments "$@"

  if detect_potential_dag; then
    export_environment_vars

    if [ -z "${LOGS}" ]; then
      if [ -z "${STOP}" ]; then
        start
      else
        stop
      fi
    else
      logs
    fi
  else
    echo "No .py or .zip files found that may contain an Apache Airflow DAG"
  fi
}

main "$@"