feat: add support for .extraEnv (defaulting to []) for the various po…

…d types As per #629, add support for .extraEnv for the various pod types, including some of the validation and warnings. This is generally better where we either use an environment variable that's something new, or that we are overriding something in airflow.config (eg. setting the log level at DEBUG on workers but INFO everywhere else) Signed-off-by: Matthew Byng-Maddick <[email protected]>
airflow-helm · Aug 30, 2023 · 75ba4bd · 75ba4bd
1 parent c4ade4e
commit 75ba4bd
Show file tree

Hide file tree

Showing 10 changed files with 73 additions and 23 deletions.
diff --git a/charts/airflow/README.md b/charts/airflow/README.md
@@ -252,6 +252,7 @@ Parameter | Description | Default
 `scheduler.podDisruptionBudget.*` | configs for the PodDisruptionBudget of the scheduler | `<see values.yaml>`
 `scheduler.logCleanup.*` | configs for the log-cleanup sidecar of the scheduler | `<see values.yaml>`
 `scheduler.numRuns` | the value of the `airflow --num_runs` parameter used to run the airflow scheduler | `-1`
+`scheduler.extraEnv` | extra environment to make available in the scheduler Pods | `[]`
 `scheduler.extraPipPackages` | extra pip packages to install in the scheduler Pods | `[]`
 `scheduler.extraVolumeMounts` | extra VolumeMounts for the scheduler Pods | `[]`
 `scheduler.extraVolumes` | extra Volumes for the scheduler Pods | `[]`
@@ -281,6 +282,7 @@ Parameter | Description | Default
 `web.service.*` | configs for the Service of the web pods | `<see values.yaml>`
 `web.readinessProbe.*` | configs for the web Pods' readiness probe | `<see values.yaml>`
 `web.livenessProbe.*` | configs for the web Pods' liveness probe | `<see values.yaml>`
+`web.extraEnv` | extra environment to make available in the web Pods | `[]`
 `web.extraPipPackages` | extra pip packages to install in the web Pods | `[]`
 `web.extraVolumeMounts` | extra VolumeMounts for the web Pods | `[]`
 `web.extraVolumes` | extra Volumes for the web Pods | `[]`
@@ -309,6 +311,7 @@ Parameter | Description | Default
 `workers.celery.*` | configs for the celery worker Pods | `<see values.yaml>`
 `workers.terminationPeriod` | how many seconds to wait after SIGTERM before SIGKILL of the celery worker | `60`
 `workers.logCleanup.*` | configs for the log-cleanup sidecar of the worker Pods | `<see values.yaml>`
+`workers.extraEnv` | extra environment to make available in the worker Pods | `[]`
 `workers.livenessProbe.*` | configs for the worker Pods' liveness probe | `<see values.yaml>`
 `workers.extraPipPackages` | extra pip packages to install in the worker Pods | `[]`
 `workers.extraVolumeMounts` | extra VolumeMounts for the worker Pods | `[]`
@@ -336,6 +339,7 @@ Parameter | Description | Default
 `triggerer.podDisruptionBudget.*` | configs for the PodDisruptionBudget of the triggerer Deployment | `<see values.yaml>`
 `triggerer.capacity` | maximum number of triggers each triggerer will run at once (sets `AIRFLOW__TRIGGERER__DEFAULT_CAPACITY`) | `1000`
 `triggerer.livenessProbe.*` | configs for the triggerer Pods' liveness probe | `<see values.yaml>`
+`triggerer.extraEnv` | extra environment to make available in the triggerer Pods | `[]`
 `triggerer.extraPipPackages` | extra pip packages to install in the triggerer Pods | `[]`
 `triggerer.extraVolumeMounts` | extra VolumeMounts for the triggerer Pods | `[]`
 `triggerer.extraVolumes` | extra Volumes for the triggerer Pods | `[]`
@@ -362,6 +366,7 @@ Parameter | Description | Default
 `flower.basicAuthSecret` | the name of a pre-created secret containing the basic authentication value for flower | `""`
 `flower.basicAuthSecretKey` | the key within `flower.basicAuthSecret` containing the basic authentication string | `""`
 `flower.service.*` | configs for the Service of the flower Pods | `<see values.yaml>`
+`flower.extraEnv` | extra environment to make available in the flower Pod | `[]`
 `flower.extraPipPackages` | extra pip packages to install in the flower Pod | `[]`
 `flower.extraVolumeMounts` | extra VolumeMounts for the flower Pods | `[]`
 `flower.extraVolumes` | extra Volumes for the flower Pods | `[]`

diff --git a/charts/airflow/templates/NOTES.txt b/charts/airflow/templates/NOTES.txt
@@ -11,6 +11,11 @@
   {{- $remote_logging_enabled = true }}
   {{- end }}
 {{- end }}
+{{- range $env := .Values.workers.extraEnv }}
+  {{- if has $env.name $remote_logging_envvars }}
+  {{- $remote_logging_enabled = true }}
+  {{- end }}
+{{- end }}
 
 {{- /* if an extra volume has been mounted for worker logs */ -}}
 {{- $extra_volumes_worker_logs := false }}
@@ -44,6 +49,7 @@
   {{- $fernet_key_warning = false }}
   {{- end }}
 {{- end }}
+{{/* we still warn even if we set the fernet key in sub-extraEnvs as they might be different, and that would be bad */}}
 
 {{- /* if we show the webserver secret_key warning */ -}}
 {{- $web_secret_warning := true }}
@@ -61,6 +67,11 @@
   {{- $web_secret_warning = false }}
   {{- end }}
 {{- end }}
+{{- range $env := .Values.web.extraEnv }}
+  {{- if has $env.name $web_secret_envvars }}
+  {{- $web_secret_warning = false }}
+  {{- end }}
+{{- end }}
 
 {{- /* if we show the external database password warning */ -}}
 {{- $external_database_password_warning := false }}

diff --git a/charts/airflow/templates/_helpers/pods.tpl b/charts/airflow/templates/_helpers/pods.tpl
@@ -582,6 +582,9 @@ EXAMPLE USAGE: {{ include "airflow.env" (dict "Release" .Release "Values" .Value
 {{- end }}
 
 {{- /* user-defined environment variables */ -}}
+{{- if .extraEnv }}
+{{ toYaml .extraEnv }}
+{{- end }}
 {{- if .Values.airflow.extraEnv }}
 {{ toYaml .Values.airflow.extraEnv }}
 {{- end }}

diff --git a/charts/airflow/templates/_helpers/validate-values.tpl b/charts/airflow/templates/_helpers/validate-values.tpl
@@ -51,24 +51,35 @@
   {{- end }}
 {{- end }}
 
-{{/* Checks for `airflow.config` */}}
-{{- if .Values.airflow.config.AIRFLOW__CORE__EXECUTOR }}
-  {{ required "Don't define `airflow.config.AIRFLOW__CORE__EXECUTOR`, it will be automatically set from `airflow.executor`!" nil }}
-{{- end }}
-{{- if or .Values.airflow.config.AIRFLOW__CORE__DAGS_FOLDER }}
-  {{ required "Don't define `airflow.config.AIRFLOW__CORE__DAGS_FOLDER`, it will be automatically set from `dags.path`!" nil }}
-{{- end }}
-{{- if or (.Values.airflow.config.AIRFLOW__CELERY__BROKER_URL) (.Values.airflow.config.AIRFLOW__CELERY__BROKER_URL_CMD) }}
-  {{ required "Don't define `airflow.config.AIRFLOW__CELERY__BROKER_URL`, it will be automatically set by the chart!" nil }}
-{{- end }}
-{{- if or (.Values.airflow.config.AIRFLOW__CELERY__RESULT_BACKEND) (.Values.airflow.config.AIRFLOW__CELERY__RESULT_BACKEND_CMD) }}
-  {{ required "Don't define `airflow.config.AIRFLOW__CELERY__RESULT_BACKEND`, it will be automatically set by the chart!" nil }}
-{{- end }}
-{{- if or (.Values.airflow.config.AIRFLOW__CORE__SQL_ALCHEMY_CONN) (.Values.airflow.config.AIRFLOW__CORE__SQL_ALCHEMY_CONN_CMD) }}
-  {{ required "Don't define `airflow.config.AIRFLOW__CORE__SQL_ALCHEMY_CONN`, it will be automatically set by the chart!" nil }}
+{{/* defines checks for various environment overrides */}}
+{{- define "airflow.environment_check" }}
+  {{- if .Values.AIRFLOW__CORE__EXECUTOR }}
+    {{ required (printf "Don't define `AIRFLOW__CORE__EXECUTOR` under %s, it will be automatically set from `airflow.executor`!" .TreeName) nil }}
+  {{- end }}
+  {{- if .Values.AIRFLOW__CORE__DAGS_FOLDER }}
+    {{ required (printf "Don't define `AIRFLOW__CORE__DAGS_FOLDER` under %s, it will be automatically set from `dags.path`!" .TreeName) nil }}
+  {{- end }}
+  {{- if or (.Values.AIRFLOW__CELERY__BROKER_URL) (.Values.AIRFLOW__CELERY__BROKER_URL_CMD) }}
+    {{ required (printf "Don't define `AIRFLOW__CELERY__BROKER_URL` under %s, it will be automatically set by the chart!") nil }}
+  {{- end }}
+  {{- if or (.Values.AIRFLOW__CELERY__RESULT_BACKEND) (.Values.AIRFLOW__CELERY__RESULT_BACKEND_CMD) }}
+    {{ required (printf "Don't define `AIRFLOW__CELERY__RESULT_BACKEND` under %s, it will be automatically set by the chart!") nil }}
+  {{- end }}
+  {{- if or (.Values.AIRFLOW__CORE__SQL_ALCHEMY_CONN) (.Values.AIRFLOW__CORE__SQL_ALCHEMY_CONN_CMD) }}
+    {{ required (printf "Don't define `AIRFLOW__CORE__SQL_ALCHEMY_CONN` under %s, it will be automatically set by the chart!") nil }}
+  {{- end }}
+  {{- if or (.Values.AIRFLOW__DATABASE__SQL_ALCHEMY_CONN) (.Values.AIRFLOW__DATABASE__SQL_ALCHEMY_CONN_CMD) }}
+    {{ required (printf "Don't define `AIRFLOW__DATABASE__SQL_ALCHEMY_CONN` under %s, it will be automatically set by the chart!") nil }}
+  {{- end }}
 {{- end }}
-{{- if or (.Values.airflow.config.AIRFLOW__DATABASE__SQL_ALCHEMY_CONN) (.Values.airflow.config.AIRFLOW__DATABASE__SQL_ALCHEMY_CONN_CMD) }}
-  {{ required "Don't define `airflow.config.AIRFLOW__DATABASE__SQL_ALCHEMY_CONN`, it will be automatically set by the chart!" nil }}
+{{/* Checks for `airflow.config` and the extraEnvs */}}
+{{- include "airflow.environment_check" (dict "Values" .Values.airflow.config "TreeName" "`airflow.config`") }}
+{{- range $tree := (list "airflow" "flower" "scheduler" "triggerer" "web" "workers") }}
+  {{- $envDict := (dict) }}
+  {{- range $env := (index $.Values $tree "extraEnv") }}
+    {{- $_ := (set $envDict $env.name "1") }}
+  {{- end }}
+  {{- include "airflow.environment_check" (dict "Values" $envDict "TreeName" (printf "`%s.extraEnv`" $tree)) }}
 {{- end }}
 
 {{/* Checks for `scheduler.logCleanup` */}}
@@ -220,4 +231,4 @@
     {{ required "If `externalRedis.host` is set, then `redis.enabled` should be `false`!" nil }}
     {{- end }}
   {{- end }}
-{{- end }}
+{{- end }}
diff --git a/charts/airflow/templates/flower/flower-deployment.yaml b/charts/airflow/templates/flower/flower-deployment.yaml
@@ -3,6 +3,7 @@
 {{- $podAffinity := include "airflow.podAffinity" (dict "Release" .Release "Values" .Values "affinity" .Values.flower.affinity) }}
 {{- $podTolerations := include "airflow.podTolerations" (dict "Release" .Release "Values" .Values "tolerations" .Values.flower.tolerations) }}
 {{- $podSecurityContext := include "airflow.podSecurityContext" (dict "Release" .Release "Values" .Values "securityContext" .Values.flower.securityContext) }}
+{{- $env := include "airflow.env" (dict "Release" .Release "Values" .Values "extraEnv" .Values.flower.extraEnv) }}
 {{- $extraPipPackages := concat .Values.airflow.extraPipPackages .Values.flower.extraPipPackages }}
 {{- $extraVolumeMounts := .Values.flower.extraVolumeMounts }}
 {{- $volumeMounts := include "airflow.volumeMounts" (dict "Release" .Release "Values" .Values "extraPipPackages" $extraPipPackages "extraVolumeMounts" $extraVolumeMounts) }}
@@ -100,7 +101,7 @@ spec:
           envFrom:
             {{- include "airflow.envFrom" . | indent 12 }}
           env:
-            {{- include "airflow.env" . | indent 12 }}
+            {{- $env | indent 12 }}
           ports:
             - name: flower
               containerPort: 5555

diff --git a/charts/airflow/templates/scheduler/scheduler-deployment.yaml b/charts/airflow/templates/scheduler/scheduler-deployment.yaml
@@ -2,6 +2,7 @@
 {{- $podAffinity := include "airflow.podAffinity" (dict "Release" .Release "Values" .Values "affinity" .Values.scheduler.affinity) }}
 {{- $podTolerations := include "airflow.podTolerations" (dict "Release" .Release "Values" .Values "tolerations" .Values.scheduler.tolerations) }}
 {{- $podSecurityContext := include "airflow.podSecurityContext" (dict "Release" .Release "Values" .Values "securityContext" .Values.scheduler.securityContext) }}
+{{- $env := include "airflow.env" (dict "Release" .Release "Values" .Values "extraEnv" .Values.scheduler.extraEnv) }}
 {{- $extraPipPackages := concat .Values.airflow.extraPipPackages .Values.scheduler.extraPipPackages }}
 {{- $extraVolumeMounts := .Values.scheduler.extraVolumeMounts }}
 {{- $volumeMounts := include "airflow.volumeMounts" (dict "Release" .Release "Values" .Values "extraPipPackages" $extraPipPackages "extraVolumeMounts" $extraVolumeMounts) }}
@@ -110,7 +111,7 @@ spec:
           envFrom:
             {{- include "airflow.envFrom" . | indent 12 }}
           env:
-            {{- include "airflow.env" . | indent 12 }}
+            {{- $env | indent 12 }}
           command:
             {{- include "airflow.command" . | indent 12 }}
           args:

diff --git a/charts/airflow/templates/triggerer/triggerer-deployment.yaml b/charts/airflow/templates/triggerer/triggerer-deployment.yaml
@@ -3,6 +3,7 @@
 {{- $podAffinity := include "airflow.podAffinity" (dict "Release" .Release "Values" .Values "affinity" .Values.triggerer.affinity) }}
 {{- $podTolerations := include "airflow.podTolerations" (dict "Release" .Release "Values" .Values "tolerations" .Values.triggerer.tolerations) }}
 {{- $podSecurityContext := include "airflow.podSecurityContext" (dict "Release" .Release "Values" .Values "securityContext" .Values.triggerer.securityContext) }}
+{{- $env := include "airflow.env" (dict "Release" .Release "Values" .Values "extraEnv" .Values.triggerer.extraEnv) }}
 {{- $extraPipPackages := concat .Values.airflow.extraPipPackages .Values.triggerer.extraPipPackages }}
 {{- $extraVolumeMounts := .Values.triggerer.extraVolumeMounts }}
 {{- $volumeMounts := include "airflow.volumeMounts" (dict "Release" .Release "Values" .Values "extraPipPackages" $extraPipPackages "extraVolumeMounts" $extraVolumeMounts) }}
@@ -99,7 +100,7 @@ spec:
           envFrom:
             {{- include "airflow.envFrom" . | indent 12 }}
           env:
-            {{- include "airflow.env" . | indent 12 }}
+            {{- $env | indent 12 }}
           command:
             {{- include "airflow.command" . | indent 12 }}
           args:

diff --git a/charts/airflow/templates/webserver/webserver-deployment.yaml b/charts/airflow/templates/webserver/webserver-deployment.yaml
@@ -2,6 +2,7 @@
 {{- $podAffinity := include "airflow.podAffinity" (dict "Release" .Release "Values" .Values "affinity" .Values.web.affinity) }}
 {{- $podTolerations := include "airflow.podTolerations" (dict "Release" .Release "Values" .Values "tolerations" .Values.web.tolerations) }}
 {{- $podSecurityContext := include "airflow.podSecurityContext" (dict "Release" .Release "Values" .Values "securityContext" .Values.web.securityContext) }}
+{{- $env := include "airflow.env" (dict "Release" .Release "Values" .Values "extraEnv" .Values.web.extraEnv) }}
 {{- $extraPipPackages := concat .Values.airflow.extraPipPackages .Values.web.extraPipPackages }}
 {{- $extraVolumeMounts := .Values.web.extraVolumeMounts }}
 {{- $volumeMounts := include "airflow.volumeMounts" (dict "Release" .Release "Values" .Values "extraPipPackages" $extraPipPackages "extraVolumeMounts" $extraVolumeMounts) }}
@@ -105,7 +106,7 @@ spec:
           envFrom:
             {{- include "airflow.envFrom" . | indent 12 }}
           env:
-            {{- include "airflow.env" . | indent 12 }}
+            {{- $env | indent 12 }}
           command:
             {{- include "airflow.command" . | indent 12 }}
           args:

diff --git a/charts/airflow/templates/worker/worker-statefulset.yaml b/charts/airflow/templates/worker/worker-statefulset.yaml
@@ -3,6 +3,7 @@
 {{- $podAffinity := include "airflow.podAffinity" (dict "Release" .Release "Values" .Values "affinity" .Values.workers.affinity) }}
 {{- $podTolerations := include "airflow.podTolerations" (dict "Release" .Release "Values" .Values "tolerations" .Values.workers.tolerations) }}
 {{- $podSecurityContext := include "airflow.podSecurityContext" (dict "Release" .Release "Values" .Values "securityContext" .Values.workers.securityContext) }}
+{{- $env := include "airflow.env" (dict "Release" .Release "Values" .Values "extraEnv" .Values.workers.extraEnv) }}
 {{- $extraPipPackages := concat .Values.airflow.extraPipPackages .Values.workers.extraPipPackages }}
 {{- $extraVolumeMounts := .Values.workers.extraVolumeMounts }}
 {{- $volumeMounts := include "airflow.volumeMounts" (dict "Release" .Release "Values" .Values "extraPipPackages" $extraPipPackages "extraVolumeMounts" $extraVolumeMounts) }}
@@ -104,7 +105,7 @@ spec:
           envFrom:
             {{- include "airflow.envFrom" . | indent 12 }}
           env:
-            {{- include "airflow.env" . | indent 12 }}
+            {{- $env | indent 12 }}
             # have dumb-init only send signals to direct child process (needed for celery workers to warm shutdown)
             - name: DUMB_INIT_SETSID
               value: "0"

diff --git a/charts/airflow/values.yaml b/charts/airflow/values.yaml
@@ -645,6 +645,9 @@ scheduler:
       ##
       schedulerAgeBeforeCheck: 180
 
+  ## extra environment for just the scheduler Pods
+  extraEnv: []
+
   ## extra pip packages to install in the scheduler Pods
   ##
   ## ____ EXAMPLE _______________
@@ -813,6 +816,9 @@ web:
     timeoutSeconds: 5
     failureThreshold: 6
 
+  ## extra environment for just the web Pods
+  extraEnv: []
+
   ## extra pip packages to install in the web Pods
   ##
   ## ____ EXAMPLE _______________
@@ -994,6 +1000,9 @@ workers:
     ##
     intervalSeconds: 900
 
+  ## extra environment for just the worker Pods
+  extraEnv: []
+
   ## configs for the worker Pods' liveness probe
   ##
   livenessProbe:
@@ -1122,6 +1131,9 @@ triggerer:
     timeoutSeconds: 60
     failureThreshold: 5
 
+  ## extra environment for just the triggerer Pods
+  extraEnv: []
+
   ## extra pip packages to install in the triggerer Pod
   ##
   ## ____ EXAMPLE _______________
@@ -1264,6 +1276,9 @@ flower:
     timeoutSeconds: 5
     failureThreshold: 6
 
+  ## extra environment for just the flower Pods
+  extraEnv: []
+
   ## extra pip packages to install in the flower Pod
   ##
   ## ____ EXAMPLE _______________