apiVersion: v1
kind: ConfigMap
metadata:
  name: airflow-webserver-config
data:
  webserver_config.py: |-
    import os
    from flask_appbuilder.security.manager import AUTH_REMOTE_USER
    from airflow.configuration import conf
    from flask import g
    from flask import get_flashed_messages, request, redirect, flash
    from flask_appbuilder import expose
    from flask_appbuilder._compat import as_unicode
    from flask_appbuilder.security.views import AuthView
    from flask_login import login_user, logout_user

    from airflow.www_rbac.security import AirflowSecurityManager

    class CustomAuthRemoteUserView(AuthView):
      login_template = ""

      @expose("/login/")
      def login(self):
          if g.user is not None and g.user.is_authenticated:
              return redirect(self.appbuilder.get_url_for_index)
          username = request.environ.get('HTTP_REMOTE_USER')
          if username:
              # https://github.com/dpgaspar/Flask-AppBuilder/blob/55b0976e1450295d5a26a06d28c5b992fb0b561e/flask_appbuilder/security/manager.py#L1201
              user = self.appbuilder.sm.auth_user_remote_user(username)
              if user is None:
                  flash(as_unicode(self.invalid_login_message), "warning")
              else:
                  login_user(user)
          else:
              flash(as_unicode(self.invalid_login_message), "warning")

          # Flush "Access is Denied" flash messaage
          get_flashed_messages()
          return redirect(self.appbuilder.get_url_for_index)

      @expose("/logout/")
      def logout(self):
          logout_user()
          return redirect("/oauth/logout")

    class CustomAirflowSecurityManager(AirflowSecurityManager):
      authremoteuserview = CustomAuthRemoteUserView

    SECURITY_MANAGER_CLASS = CustomAirflowSecurityManager

    basedir = os.path.abspath(os.path.dirname(__file__))
    # The SQLAlchemy connection string.
    SQLALCHEMY_DATABASE_URI = conf.get('core', 'SQL_ALCHEMY_CONN')
    # Flask-WTF flag for CSRF
    WTF_CSRF_ENABLED = True
    AUTH_TYPE = AUTH_REMOTE_USER
  users_create.sh: |-
    #! /bin/sh
    echo "Creating User $ADMIN_USER_NAME"
    res=$(airflow users create --username $ADMIN_USER_NAME --password $ADMIN_PASSWORD --firstname admin --lastname admin --email admin@example.com --role 'Admin')
    exit 0
  postgres_hang.sh: |-
    #!/bin/sh
    {{- if .Values.airflow.postgresql.enabled }}
      DATABASE_HOST="{{ .Release.Name }}-postgresql"
      DATABASE_PORT="5432"
      {{- else }}
      DATABASE_HOST="{{ .Values.airflow.externalDatabase.host }}"
      DATABASE_PORT="{{ .Values.airflow.externalDatabase.port }}"
    {{- end }}
    pg_uri=$PG_URL
    pg_host=$PG_HOST
    pg_user=postgres
    # make sure pg is ready to accept connections
    until pg_isready -h $DATABASE_HOST -p $DATABASE_PORT
    do
      echo "Waiting for postgres at: $DATABASE_HOST"
      sleep 2;
    done
    echo "postgres is live"
