yili-docker-images/charts/test/airflow/templates/_helpers.yaml

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

{{/*
Create a default fully qualified app name.
We truncate at 63 chars because some Kubernetes name fields are limited to this (by the DNS naming spec).
If release name contains chart name it will be used as a full name.
*/}}
{{- define "airflow.fullname" -}}
{{- if .Values.fullnameOverride }}
{{- .Values.fullnameOverride | trunc 63 | trimSuffix "-" }}
{{- else }}
{{- $name := default .Chart.Name .Values.nameOverride }}
{{- if contains $name .Release.Name }}
{{- .Release.Name | trunc 63 | trimSuffix "-" }}
{{- else }}
{{- printf "%s-%s" .Release.Name $name | trunc 63 | trimSuffix "-" }}
{{- end }}
{{- end }}
{{- end }}


{{/* Standard Airflow environment variables */}}
{{- define "standard_airflow_environment" }}
  # Hard Coded Airflow Envs
  {{- if .Values.enableBuiltInSecretEnvVars.AIRFLOW__CORE__FERNET_KEY }}
  - name: AIRFLOW__CORE__FERNET_KEY
    valueFrom:
      secretKeyRef:
        name: {{ template "fernet_key_secret" . }}
        key: fernet-key
  {{- end }}
  # For Airflow <2.3, backward compatibility; moved to [database] in 2.3
  {{- if .Values.enableBuiltInSecretEnvVars.AIRFLOW__CORE__SQL_ALCHEMY_CONN }}
  - name: AIRFLOW__CORE__SQL_ALCHEMY_CONN
    valueFrom:
      secretKeyRef:
        name: {{ template "airflow_metadata_secret" . }}
        key: connection
  {{- end }}
  {{- if .Values.enableBuiltInSecretEnvVars.AIRFLOW__DATABASE__SQL_ALCHEMY_CONN }}
  - name: AIRFLOW__DATABASE__SQL_ALCHEMY_CONN
    valueFrom:
      secretKeyRef:
        name: {{ template "airflow_metadata_secret" . }}
        key: connection
  {{- end }}
  {{- if .Values.enableBuiltInSecretEnvVars.AIRFLOW_CONN_AIRFLOW_DB }}
  - name: AIRFLOW_CONN_AIRFLOW_DB
    valueFrom:
      secretKeyRef:
        name: {{ template "airflow_metadata_secret" . }}
        key: connection
  {{- end }}
  {{- if .Values.enableBuiltInSecretEnvVars.AIRFLOW__WEBSERVER__SECRET_KEY }}
  - name: AIRFLOW__WEBSERVER__SECRET_KEY
    valueFrom:
      secretKeyRef:
        name: {{ template "webserver_secret_key_secret" . }}
        key: webserver-secret-key
  {{- end }}
  {{- if or (eq .Values.executor "CeleryExecutor") (eq .Values.executor "CeleryKubernetesExecutor") }}
    {{- if .Values.enableBuiltInSecretEnvVars.AIRFLOW__CELERY__CELERY_RESULT_BACKEND }}
  # (Airflow 1.10.* variant)
  - name: AIRFLOW__CELERY__CELERY_RESULT_BACKEND
    valueFrom:
      secretKeyRef:
        name: {{ template "airflow_result_backend_secret" . }}
        key: connection
    {{- end }}
    {{- if .Values.enableBuiltInSecretEnvVars.AIRFLOW__CELERY__RESULT_BACKEND }}
  - name: AIRFLOW__CELERY__RESULT_BACKEND
    valueFrom:
      secretKeyRef:
        name: {{ template "airflow_result_backend_secret" . }}
        key: connection
    {{- end }}
    {{- if .Values.enableBuiltInSecretEnvVars.AIRFLOW__CELERY__BROKER_URL }}
  - name: AIRFLOW__CELERY__BROKER_URL
    valueFrom:
      secretKeyRef:
        name: {{ default (printf "%s-broker-url" .Release.Name) .Values.data.brokerUrlSecretName }}
        key: connection
    {{- end }}
  {{- end }}
  {{- if .Values.elasticsearch.enabled }}
  # The elasticsearch variables were updated to the shorter names in v1.10.4
    {{- if .Values.enableBuiltInSecretEnvVars.AIRFLOW__ELASTICSEARCH__HOST }}
  - name: AIRFLOW__ELASTICSEARCH__HOST
    valueFrom:
      secretKeyRef:
        name: {{ template "elasticsearch_secret" . }}
        key: connection
    {{- end }}
    {{- if .Values.enableBuiltInSecretEnvVars.AIRFLOW__ELASTICSEARCH__ELASTICSEARCH_HOST }}
  # This is the older format for these variable names, kept here for backward compatibility
  - name: AIRFLOW__ELASTICSEARCH__ELASTICSEARCH_HOST
    valueFrom:
      secretKeyRef:
        name: {{ template "elasticsearch_secret" . }}
        key: connection
    {{- end }}
  {{- end }}
{{- end }}

{{/* User defined Airflow environment variables */}}
{{- define "custom_airflow_environment" }}
  # Dynamically created environment variables
  {{- range $i, $config := .Values.env }}
  - name: {{ $config.name }}
    value: {{ $config.value | quote }}
    {{- if or (eq $.Values.executor "KubernetesExecutor") (eq $.Values.executor "LocalKubernetesExecutor") (eq $.Values.executor "CeleryKubernetesExecutor") }}
  - name: AIRFLOW__KUBERNETES_ENVIRONMENT_VARIABLES__{{ $config.name }}
    value: {{ $config.value | quote }}
    {{- end }}
  {{- end }}
  # Dynamically created secret envs
  {{- range $i, $config := .Values.secret }}
  - name: {{ $config.envName }}
    valueFrom:
      secretKeyRef:
        name: {{ $config.secretName }}
        key: {{ default "value" $config.secretKey }}
  {{- end }}
  {{- if or (eq $.Values.executor "LocalKubernetesExecutor") (eq $.Values.executor "KubernetesExecutor") (eq $.Values.executor "CeleryKubernetesExecutor") }}
    {{- range $i, $config := .Values.secret }}
  - name: AIRFLOW__KUBERNETES_SECRETS__{{ $config.envName }}
    value: {{ printf "%s=%s" $config.secretName $config.secretKey }}
    {{- end }}
  {{ end }}
  # Extra env
  {{- $Global := . }}
  {{- with .Values.extraEnv }}
  {{- tpl . $Global | nindent 2 }}
  {{- end }}
{{- end }}

{{/* User defined Airflow environment from */}}
{{- define "custom_airflow_environment_from" }}
  {{- $Global := . }}
  {{- with .Values.extraEnvFrom }}
  {{- tpl . $Global | nindent 2 }}
  {{- end }}
{{- end }}

{{/*  Git ssh key volume */}}
{{- define "git_sync_ssh_key_volume"}}
- name: git-sync-ssh-key
  secret:
    secretName: {{ .Values.dags.gitSync.sshKeySecret }}
    defaultMode: 288
{{- end }}

{{/*  Git sync container */}}
{{- define "git_sync_container"}}
- name: {{ .Values.dags.gitSync.containerName }}{{ if .is_init }}-init{{ end }}
  image: {{ template "git_sync_image" . }}
  imagePullPolicy: {{ .Values.images.gitSync.pullPolicy }}
  securityContext: {{ include "localSecurityContext" .Values.dags.gitSync | nindent 4 }}
  env:
    {{- if .Values.dags.gitSync.sshKeySecret }}
    - name: GIT_SSH_KEY_FILE
      value: "/etc/git-secret/ssh"
    - name: GIT_SYNC_SSH
      value: "true"
    {{- if .Values.dags.gitSync.knownHosts }}
    - name: GIT_KNOWN_HOSTS
      value: "true"
    - name: GIT_SSH_KNOWN_HOSTS_FILE
      value: "/etc/git-secret/known_hosts"
    {{- else }}
    - name: GIT_KNOWN_HOSTS
      value: "false"
    {{- end }}
    {{ else if .Values.dags.gitSync.credentialsSecret }}
    - name: GIT_SYNC_USERNAME
      valueFrom:
        secretKeyRef:
          name: {{ .Values.dags.gitSync.credentialsSecret | quote }}
          key: GIT_SYNC_USERNAME
    - name: GIT_SYNC_PASSWORD
      valueFrom:
        secretKeyRef:
          name: {{ .Values.dags.gitSync.credentialsSecret | quote }}
          key: GIT_SYNC_PASSWORD
    {{- end }}
    - name: GIT_SYNC_REV
      value: {{ .Values.dags.gitSync.rev | quote }}
    - name: GIT_SYNC_BRANCH
      value: {{ .Values.dags.gitSync.branch | quote }}
    - name: GIT_SYNC_REPO
      value: {{ .Values.dags.gitSync.repo | quote }}
    - name: GIT_SYNC_DEPTH
      value: {{ .Values.dags.gitSync.depth | quote }}
    - name: GIT_SYNC_ROOT
      value: "/git"
    - name: GIT_SYNC_DEST
      value: "repo"
    - name: GIT_SYNC_ADD_USER
      value: "true"
    - name: GIT_SYNC_WAIT
      value: {{ .Values.dags.gitSync.wait | quote }}
    - name: GIT_SYNC_MAX_SYNC_FAILURES
      value: {{ .Values.dags.gitSync.maxFailures | quote }}
    {{- if .is_init }}
    - name: GIT_SYNC_ONE_TIME
      value: "true"
    {{- end }}
    {{- range $i, $config := .Values.dags.gitSync.env }}
    - name: {{ $config.name }}
      value: {{ $config.value | quote }}
    {{- end }}
  resources: {{ toYaml .Values.dags.gitSync.resources | nindent 6 }}
  volumeMounts:
  - name: dags
    mountPath: /git
  {{- if .Values.dags.gitSync.sshKeySecret }}
  - name: git-sync-ssh-key
    mountPath: /etc/git-secret/ssh
    readOnly: true
    subPath: gitSshKey
  {{- if .Values.dags.gitSync.knownHosts }}
  - name: config
    mountPath: /etc/git-secret/known_hosts
    readOnly: true
    subPath: known_hosts
  {{- end }}
  {{- end }}
{{- if .Values.dags.gitSync.extraVolumeMounts }}
{{ toYaml .Values.dags.gitSync.extraVolumeMounts | indent 2 }}
{{- end }}
{{- end }}

# This helper will change when customers deploy a new image.
{{ define "airflow_image" -}}
{{ printf "%s:%s" (.Values.images.airflow.repository | default .Values.defaultAirflowRepository) (.Values.images.airflow.tag | default .Values.defaultAirflowTag) }}
{{- end }}

{{ define "pod_template_image" -}}
{{ printf "%s:%s" (.Values.images.pod_template.repository | default .Values.defaultAirflowRepository) (.Values.images.pod_template.tag | default .Values.defaultAirflowTag) }}
{{- end }}

# This helper is used for airflow containers that do not need the users code.
{{ define "default_airflow_image" -}}
{{ printf "%s:%s" .Values.defaultAirflowRepository .Values.defaultAirflowTag }}
{{- end }}

{{ define "airflow_image_for_migrations" -}}
{{- if .Values.images.useDefaultImageForMigration -}}
{{ template "default_airflow_image" . }}
{{- else -}}
{{ template "airflow_image" . }}
{{- end -}}
{{- end }}

{{ define "flower_image" -}}
{{ printf "%s:%s" (.Values.images.flower.repository | default .Values.defaultAirflowRepository) (.Values.images.flower.tag | default .Values.defaultAirflowTag) }}
{{- end }}

{{ define "statsd_image" -}}
{{ printf "%s:%s" .Values.images.statsd.repository .Values.images.statsd.tag }}
{{- end }}

{{ define "redis_image" -}}
{{ printf "%s:%s" .Values.images.redis.repository .Values.images.redis.tag }}
{{- end }}

{{ define "pgbouncer_image" -}}
{{ printf "%s:%s" .Values.images.pgbouncer.repository .Values.images.pgbouncer.tag }}
{{- end }}

{{ define "pgbouncer_exporter_image" -}}
{{ printf "%s:%s" .Values.images.pgbouncerExporter.repository .Values.images.pgbouncerExporter.tag }}
{{- end }}

{{ define "git_sync_image" -}}
{{ printf "%s:%s" .Values.images.gitSync.repository .Values.images.gitSync.tag }}
{{- end }}

{{ define "fernet_key_secret" -}}
{{ default (printf "%s-fernet-key" .Release.Name) .Values.fernetKeySecretName }}
{{- end }}

{{ define "webserver_secret_key_secret" -}}
{{ default (printf "%s-webserver-secret-key" .Release.Name) .Values.webserverSecretKeySecretName }}
{{- end }}

{{ define "redis_password_secret" -}}
{{ default (printf "%s-redis-password" .Release.Name) .Values.redis.passwordSecretName }}
{{- end }}

{{ define "airflow_metadata_secret" -}}
{{ default (printf "%s-airflow-metadata" .Release.Name) .Values.data.metadataSecretName }}
{{- end }}

{{ define "airflow_result_backend_secret" -}}
{{ default (printf "%s-airflow-result-backend" .Release.Name) .Values.data.resultBackendSecretName }}
{{- end }}

{{ define "airflow_pod_template_file" -}}
{{ (printf "%s/pod_templates" .Values.airflowHome) }}
{{- end }}

{{ define "pgbouncer_config_secret" -}}
{{ default (printf "%s-pgbouncer-config" .Release.Name) .Values.pgbouncer.configSecretName }}
{{- end }}

{{ define "pgbouncer_certificates_secret" -}}
{{ .Release.Name }}-pgbouncer-certificates
{{- end }}

{{ define "pgbouncer_stats_secret" -}}
{{ .Release.Name }}-pgbouncer-stats
{{- end }}

{{ define "registry_secret" -}}
{{ default (printf "%s-registry" .Release.Name) .Values.registry.secretName }}
{{- end }}

{{ define "elasticsearch_secret" -}}
{{ default (printf "%s-elasticsearch" .Release.Name) .Values.elasticsearch.secretName }}
{{- end }}

{{ define "flower_secret" -}}
{{ default (printf "%s-flower" .Release.Name) .Values.flower.secretName }}
{{- end }}

{{ define "kerberos_keytab_secret" -}}
{{ .Release.Name }}-kerberos-keytab
{{- end }}

{{ define "kerberos_ccache_path" -}}
{{ printf "%s/%s" .Values.kerberos.ccacheMountPath .Values.kerberos.ccacheFileName }}
{{- end }}

{{ define "pgbouncer_config" }}
{{- $resultBackendConnection := .Values.data.resultBackendConnection | default .Values.data.metadataConnection }}
{{- $pgMetadataHost := .Values.data.metadataConnection.host | default (printf "%s-%s.%s" .Release.Name "postgresql" .Release.Namespace) }}
{{- $pgResultBackendHost := $resultBackendConnection.host | default (printf "%s-%s.%s" .Release.Name "postgresql" .Release.Namespace) }}
[databases]
{{ .Release.Name }}-metadata = host={{ $pgMetadataHost }} dbname={{ .Values.data.metadataConnection.db }} port={{ .Values.data.metadataConnection.port }} pool_size={{ .Values.pgbouncer.metadataPoolSize }} {{ .Values.pgbouncer.extraIniMetadata | default "" }}
{{ .Release.Name }}-result-backend = host={{ $pgResultBackendHost }} dbname={{ $resultBackendConnection.db }} port={{ $resultBackendConnection.port }} pool_size={{ .Values.pgbouncer.resultBackendPoolSize }} {{ .Values.pgbouncer.extraIniResultBackend | default "" }}

[pgbouncer]
pool_mode = transaction
listen_port = {{ .Values.ports.pgbouncer }}
listen_addr = *
auth_type = {{ .Values.pgbouncer.auth_type }}
auth_file = {{ .Values.pgbouncer.auth_file }}
stats_users = {{ .Values.data.metadataConnection.user }}
ignore_startup_parameters = extra_float_digits
max_client_conn = {{ .Values.pgbouncer.maxClientConn }}
verbose = {{ .Values.pgbouncer.verbose }}
log_disconnections = {{ .Values.pgbouncer.logDisconnections }}
log_connections = {{ .Values.pgbouncer.logConnections }}

server_tls_sslmode = {{ .Values.pgbouncer.sslmode }}
server_tls_ciphers = {{ .Values.pgbouncer.ciphers }}

{{- if .Values.pgbouncer.ssl.ca }}
server_tls_ca_file = /etc/pgbouncer/root.crt
{{- end }}
{{- if .Values.pgbouncer.ssl.cert }}
server_tls_cert_file = /etc/pgbouncer/server.crt
{{- end }}
{{- if .Values.pgbouncer.ssl.key }}
server_tls_key_file = /etc/pgbouncer/server.key
{{- end }}

{{- if .Values.pgbouncer.extraIni }}
{{ .Values.pgbouncer.extraIni }}
{{- end }}
{{- end }}

{{ define "pgbouncer_users" }}
{{- $resultBackendConnection := .Values.data.resultBackendConnection | default .Values.data.metadataConnection }}
{{ .Values.data.metadataConnection.user | quote }} {{ .Values.data.metadataConnection.pass | quote }}
{{ $resultBackendConnection.user | quote }} {{ $resultBackendConnection.pass | quote }}
{{- end }}

{{ define "airflow_logs" -}}
{{ (printf "%s/logs" .Values.airflowHome) | quote }}
{{- end }}

{{ define "airflow_logs_no_quote" -}}
{{ (printf "%s/logs" .Values.airflowHome) }}
{{- end }}

{{ define "airflow_logs_volume_claim" -}}
{{- if .Values.logs.persistence.existingClaim -}}
{{ .Values.logs.persistence.existingClaim }}
{{- else -}}
{{ .Release.Name }}-logs
{{- end -}}
{{- end -}}

{{ define "airflow_dags" -}}
{{- if .Values.dags.gitSync.enabled -}}
{{ (printf "%s/dags/repo/%s" .Values.airflowHome .Values.dags.gitSync.subPath) }}
{{- else -}}
{{ (printf "%s/dags" .Values.airflowHome) }}
{{- end -}}
{{- end -}}

{{ define "airflow_dags_volume_claim" -}}
{{- if .Values.dags.persistence.existingClaim -}}
{{ .Values.dags.persistence.existingClaim }}
{{- else -}}
{{ .Release.Name }}-dags
{{- end -}}
{{- end -}}

{{ define "airflow_dags_mount" -}}
- name: dags
  mountPath: {{ (printf "%s/dags" .Values.airflowHome) }}
  {{ if .Values.dags.persistence.subPath -}}
  subPath: {{ .Values.dags.persistence.subPath }}
  {{- end }}
  readOnly: {{ .Values.dags.gitSync.enabled | ternary "True" "False" }}
{{- end -}}

{{ define "airflow_config_path" -}}
{{ (printf "%s/airflow.cfg" .Values.airflowHome) | quote }}
{{- end }}

{{ define "airflow_webserver_config_path" -}}
{{ (printf "%s/webserver_config.py" .Values.airflowHome) | quote }}
{{- end }}

{{ define "airflow_local_setting_path" -}}
{{ (printf "%s/config/airflow_local_settings.py" .Values.airflowHome) | quote }}
{{- end }}

{{ define "airflow_config" -}}
{{ (printf "%s-airflow-config" .Release.Name) }}
{{- end }}

{{/*
Create the name of the webserver service account to use
*/}}
{{- define "webserver.serviceAccountName" -}}
{{- if .Values.webserver.serviceAccount.create -}}
    {{ default (printf "%s-webserver" (include "airflow.fullname" .)) .Values.webserver.serviceAccount.name }}
{{- else -}}
    {{ default "default" .Values.webserver.serviceAccount.name }}
{{- end -}}
{{- end -}}

{{/*
Create the name of the redis service account to use
*/}}
{{- define "redis.serviceAccountName" -}}
{{- if .Values.redis.serviceAccount.create -}}
    {{ default (printf "%s-redis" (include "airflow.fullname" .)) .Values.redis.serviceAccount.name }}
{{- else -}}
    {{ default "default" .Values.redis.serviceAccount.name }}
{{- end -}}
{{- end -}}

{{/*
Create the name of the flower service account to use
*/}}
{{- define "flower.serviceAccountName" -}}
{{- if .Values.flower.serviceAccount.create -}}
    {{ default (printf "%s-flower" (include "airflow.fullname" .)) .Values.flower.serviceAccount.name }}
{{- else -}}
    {{ default "default" .Values.flower.serviceAccount.name }}
{{- end -}}
{{- end -}}

{{/*
Create the name of the scheduler service account to use
*/}}
{{- define "scheduler.serviceAccountName" -}}
{{- if .Values.scheduler.serviceAccount.create -}}
    {{ default (printf "%s-scheduler" (include "airflow.fullname" .)) .Values.scheduler.serviceAccount.name }}
{{- else -}}
    {{ default "default" .Values.scheduler.serviceAccount.name }}
{{- end -}}
{{- end -}}

{{/*
Create the name of the StatsD service account to use
*/}}
{{- define "statsd.serviceAccountName" -}}
{{- if .Values.statsd.serviceAccount.create -}}
    {{ default (printf "%s-statsd" (include "airflow.fullname" .)) .Values.statsd.serviceAccount.name }}
{{- else -}}
    {{ default "default" .Values.statsd.serviceAccount.name }}
{{- end -}}
{{- end -}}

{{/*
Create the name of the create user job service account to use
*/}}
{{- define "createUserJob.serviceAccountName" -}}
{{- if .Values.createUserJob.serviceAccount.create -}}
    {{ default (printf "%s-create-user-job" (include "airflow.fullname" .)) .Values.createUserJob.serviceAccount.name }}
{{- else -}}
    {{ default "default" .Values.createUserJob.serviceAccount.name }}
{{- end -}}
{{- end -}}

{{/*
Create the name of the migrate database job service account to use
*/}}
{{- define "migrateDatabaseJob.serviceAccountName" -}}
{{- if .Values.migrateDatabaseJob.serviceAccount.create -}}
    {{ default (printf "%s-migrate-database-job" (include "airflow.fullname" .)) .Values.migrateDatabaseJob.serviceAccount.name }}
{{- else -}}
    {{ default "default" .Values.migrateDatabaseJob.serviceAccount.name }}
{{- end -}}
{{- end -}}

{{/*
Create the name of the worker service account to use
*/}}
{{- define "worker.serviceAccountName" -}}
{{- if .Values.workers.serviceAccount.create -}}
    {{ default (printf "%s-worker" (include "airflow.fullname" .)) .Values.workers.serviceAccount.name }}
{{- else -}}
    {{ default "default" .Values.workers.serviceAccount.name }}
{{- end -}}
{{- end -}}

{{/*
Create the name of the triggerer service account to use
*/}}
{{- define "triggerer.serviceAccountName" -}}
{{- if .Values.triggerer.serviceAccount.create -}}
    {{ default (printf "%s-triggerer" (include "airflow.fullname" .)) .Values.triggerer.serviceAccount.name }}
{{- else -}}
    {{ default "default" .Values.triggerer.serviceAccount.name }}
{{- end -}}
{{- end -}}

{{/*
Create the name of the pgbouncer service account to use
*/}}
{{- define "pgbouncer.serviceAccountName" -}}
{{- if .Values.pgbouncer.serviceAccount.create -}}
    {{ default (printf "%s-pgbouncer" (include "airflow.fullname" .)) .Values.pgbouncer.serviceAccount.name }}
{{- else -}}
    {{ default "default" .Values.pgbouncer.serviceAccount.name }}
{{- end -}}
{{- end -}}

{{/*
Create the name of the cleanup service account to use
*/}}
{{- define "cleanup.serviceAccountName" -}}
{{- if .Values.cleanup.serviceAccount.create -}}
    {{ default (printf "%s-cleanup" (include "airflow.fullname" .)) .Values.cleanup.serviceAccount.name }}
{{- else -}}
    {{ default "default" .Values.cleanup.serviceAccount.name }}
{{- end -}}
{{- end -}}

{{ define "wait-for-migrations-command" }}
  {{- if semverCompare ">=2.0.0" .Values.airflowVersion }}
  - airflow
  - db
  - check-migrations
  - --migration-wait-timeout={{ .Values.images.migrationsWaitTimeout }}
  {{- else }}
  - python
  - -c
  - |
        import airflow
        import logging
        import os
        import time

        from alembic.config import Config
        from alembic.runtime.migration import MigrationContext
        from alembic.script import ScriptDirectory

        from airflow import settings

        package_dir = os.path.abspath(os.path.dirname(airflow.__file__))
        directory = os.path.join(package_dir, 'migrations')
        config = Config(os.path.join(package_dir, 'alembic.ini'))
        config.set_main_option('script_location', directory)
        config.set_main_option('sqlalchemy.url', settings.SQL_ALCHEMY_CONN.replace('%', '%%'))
        script_ = ScriptDirectory.from_config(config)

        timeout=60

        with settings.engine.connect() as connection:
            context = MigrationContext.configure(connection)
            ticker = 0
            while True:
                source_heads = set(script_.get_heads())

                db_heads = set(context.get_current_heads())
                if source_heads == db_heads:
                    break

                if ticker >= timeout:
                    raise TimeoutError("There are still unapplied migrations after {} seconds.".format(ticker))
                ticker += 1
                time.sleep(1)
                logging.info('Waiting for migrations... %s second(s)', ticker)
  {{- end }}
{{- end }}

{{define "scheduler_liveness_check_command"}}

  {{- if semverCompare ">=2.0.0" .Values.airflowVersion }}
  - sh
  - -c
  - |
    CONNECTION_CHECK_MAX_COUNT=0 AIRFLOW__LOGGING__LOGGING_LEVEL=ERROR exec /entrypoint \
    airflow jobs check --job-type SchedulerJob --hostname $(hostname)
  {{- else }}
  - sh
  - -c
  - |
    CONNECTION_CHECK_MAX_COUNT=0 exec /entrypoint python -Wignore -c "
    import os
    os.environ['AIRFLOW__CORE__LOGGING_LEVEL'] = 'ERROR'
    os.environ['AIRFLOW__LOGGING__LOGGING_LEVEL'] = 'ERROR'
    from airflow.jobs.scheduler_job import SchedulerJob
    from airflow.utils.db import create_session
    from airflow.utils.net import get_hostname
    import sys
    with create_session() as session:
        job = session.query(SchedulerJob).filter_by(hostname=get_hostname()).order_by(
            SchedulerJob.latest_heartbeat.desc()).limit(1).first()
    sys.exit(0 if job.is_alive() else 1)"
  {{- end }}
{{- end }}

{{define "triggerer_liveness_check_command"}}
  - sh
  - -c
  - |
    CONNECTION_CHECK_MAX_COUNT=0 AIRFLOW__LOGGING__LOGGING_LEVEL=ERROR exec /entrypoint \
    airflow jobs check --job-type TriggererJob --hostname $(hostname)
{{- end }}

{{ define "registry_docker_config" -}}
  {{- $host := .Values.registry.connection.host }}
  {{- $email := .Values.registry.connection.email }}
  {{- $user := .Values.registry.connection.user -}}
  {{- $pass := .Values.registry.connection.pass -}}

  {{- $config := dict "auths" -}}
  {{- $auth := dict -}}
  {{- $data := dict -}}
  {{- $_ := set $data "username" $user -}}
  {{- $_ := set $data "password" $pass -}}
  {{- $_ := set $data "email" $email -}}
  {{- $_ := set $data "auth" (printf "%v:%v" $user $pass | b64enc) -}}
  {{- $_ := set $auth $host $data -}}
  {{- $_ := set $config "auths" $auth -}}
  {{ $config | toJson | print }}
{{- end }}

{{/* Allow Kubernetes Version to be overridden. Credit to https://github.com/prometheus-community/helm-charts for Regex. */}}
{{- define "kubeVersion" -}}
  {{- $kubeVersion := default .Capabilities.KubeVersion.Version .Values.kubeVersionOverride -}}
  {{/* Special use case for Amazon EKS, Google GKE */}}
  {{- if and (regexMatch "\\d+\\.\\d+\\.\\d+-(?:eks|gke).+" $kubeVersion) (not .Values.kubeVersionOverride) -}}
    {{- $kubeVersion = regexFind "\\d+\\.\\d+\\.\\d+" $kubeVersion -}}
  {{- end -}}
  {{- $kubeVersion -}}
{{- end -}}

{{/*
Set the default value for securityContext
If no value is passed for securityContext or <node>.securityContext, defaults to global uid and gid.

    +------------------------+      +-----------------+      +-------------------------+
    | <node>.securityContext |  ->  | securityContext |  ->  | Values.uid + Values.gid |
    +------------------------+      +-----------------+      +-------------------------+

Values are not accumulated meaning that if runAsUser is set to 10 in <node>.securityContext,
any extra values set to securityContext or uid+gid will be ignored.

The template can be called like so:
   include "airflowSecurityContext" (list . .Values.webserver)

Where `.` is the global variables scope and `.Values.webserver` the local variables scope for the webserver template.
*/}}
{{- define "airflowSecurityContext" -}}
  {{- $ := index . 0 -}}
  {{- with index . 1 }}
    {{- if .securityContext -}}
{{ toYaml .securityContext | print }}
    {{- else if $.Values.securityContext -}}
{{ toYaml $.Values.securityContext | print }}
    {{- else -}}
runAsUser: {{ $.Values.uid }}
fsGroup: {{ $.Values.gid }}
    {{- end -}}
  {{- end -}}
{{- end -}}

{{/*
Set the default value for securityContext
If no value is passed for securityContext or <node>.securityContext, defaults to UID in the local node.

    +------------------------+     +-------------+
    | <node>.securityContext |  >  | <node>.uid  |
    +------------------------+     +-------------+

The template can be called like so:
  include "localSecurityContext" .Values.statsd

It is important to pass the local variables scope to this template as it is used to determine the local node value for uid.
*/}}
{{- define "localSecurityContext" -}}
  {{- if .securityContext -}}
{{ toYaml .securityContext | print }}
  {{- else -}}
runAsUser: {{ .uid }}
  {{- end -}}
{{- end -}}

{{/*
Set the default value for workers chown for persistent storage
If no value is passed for securityContext or <node>.securityContext, defaults to global uid and gid.
The template looks for `runAsUser` and `fsGroup` specifically, any other parameter will be ignored.

    +------------------------+      +-----------------+      +-------------------------+
    | <node>.securityContext |  ->  | securityContext |  ->  | Values.uid + Values.gid |
    +------------------------+      +-----------------+      +-------------------------+

Values are not accumulated meaning that if runAsUser is set to 10 in <node>.securityContext,
any extra values set to securityContext or uid+gid will be ignored.

The template can be called like so:
   include "airflowSecurityContextIds" (list . .Values.workers)

Where `.` is the global variables scope and `.Values.workers` the local variables scope for the workers template.
*/}}
{{- define "airflowSecurityContextIds" -}}
  {{- $ := index . 0 -}}
  {{- with index . 1 }}
    {{- if .securityContext -}}
{{ pluck "runAsUser" .securityContext | first | default $.Values.uid }}:{{ pluck "fsGroup" .securityContext | first | default $.Values.gid }}
    {{- else if $.Values.securityContext -}}
{{ pluck "runAsUser" $.Values.securityContext | first | default $.Values.uid }}:{{ pluck "fsGroup" $.Values.securityContext | first | default $.Values.gid }}
    {{- else -}}
{{ $.Values.uid }}:{{ $.Values.gid }}
    {{- end -}}
  {{- end -}}
{{- end -}}