#!/bin/sh

# Global defaults and back-compat
: "${AIRFLOW_HOME:="/usr/local/airflow"}"
: "${AIRFLOW__CORE__EXECUTOR:=${EXECUTOR:-Celery}Executor}"
: "${AIRFLOW__CELERY__RESULT_BACKEND:="db+postgresql://airflow:airflow@database:5432/airflow"}"


export \
  AIRFLOW_HOME \
  AIRFLOW__CORE__EXECUTOR \
  AIRFLOW__CELERY__RESULT_BACKEND \

# Waiting for database is ready
DATABASE=$(echo "$AIRFLOW__CORE__SQL_ALCHEMY_CONN" | cut -d '/' -f 3 | cut -d '@' -f 2)
while ! nc -z $DATABASE 5432; do
  sleep 3
done

# Check the database is created just now
EMPTY=$(python -c "\
import sqlalchemy as sa;\
engine = sa.create_engine('$AIRFLOW__CORE__SQL_ALCHEMY_CONN');\
tables = sa.inspect(engine).get_table_names();
if tables == []: print('empty')\
")

case "$1" in
  webserver)
    if [ $EMPTY ]; then
      echo "Init database"
      airflow db init
      airflow users create --username admin --firstname admin --lastname admin --role Admin --email [email protected] --password admin
    fi
    exec airflow webserver
    ;;
  worker|scheduler)
    # Give the webserver time to run initdb.
    sleep 10
    exec airflow "$@"
    ;;
  flower)
    sleep 10
    exec airflow "$@"
    ;;
  version)
    exec airflow "$@"
    ;;
  *)
    # The command is something like bash, not an airflow subcommand. Just run it in the right environment.
    exec "$@"
    ;;
esac

# Create Jupyter notebooks from plain python code for interactive debugging
python ./plugins/utils/converter.py

jupyter lab --ip=0.0.0.0 --no-browser --notebook-dir=/tmp/notebooks --ServerApp.base_url=/jupyter/