Chart: Add support for securityContext (#18249)

This adds the ability to set both a global `securityContext` and `securityContext` by deployment,
allowing for greater flexibility in configuring how Airflow is run.

chart/files/pod-template-file.kubernetes-helm-yaml

@@ -18,6 +18,7 @@
 {{- $nodeSelector := or .Values.nodeSelector .Values.workers.nodeSelector }}
 {{- $affinity := or .Values.affinity .Values.workers.affinity }}
 {{- $tolerations := or .Values.tolerations .Values.workers.tolerations }}
+{{- $securityContext := include "airflowSecurityContext" (list . .Values.workers) }}
 apiVersion: v1
 kind: Pod
 metadata:
@@ -83,9 +84,7 @@ spec:
     - name: {{ template "registry_secret" . }}
   {{- end }}
   restartPolicy: Never
-  securityContext:
-    runAsUser: {{ .Values.uid }}
-    fsGroup: {{ .Values.gid }}
+  securityContext: {{ $securityContext | nindent 4 }}
   nodeSelector: {{ toYaml $nodeSelector | nindent 4 }}
   affinity: {{ toYaml $affinity | nindent 4 }}
   tolerations: {{ toYaml $tolerations | nindent 4 }}

chart/templates/_helpers.yaml

@@ -143,8 +143,7 @@ If release name contains chart name it will be used as a full name.
 - name: {{ .Values.dags.gitSync.containerName }}{{ if .is_init }}-init{{ end }}
   image: {{ template "git_sync_image" . }}
   imagePullPolicy: {{ .Values.images.gitSync.pullPolicy }}
-  securityContext:
-    runAsUser: {{ .Values.dags.gitSync.uid }}
+  securityContext: {{ include "localSecurityContext" .Values.dags.gitSync | nindent 4 }}
   env:
     {{- if .Values.dags.gitSync.sshKeySecret }}
     - name: GIT_SSH_KEY_FILE
@@ -616,3 +615,84 @@ Create the name of the cleanup service account to use
   {{- 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 -}}

chart/templates/cleanup/cleanup-cronjob.yaml

@@ -22,6 +22,7 @@
 {{- $nodeSelector := or .Values.cleanup.nodeSelector .Values.nodeSelector }}
 {{- $affinity := or .Values.cleanup.affinity .Values.affinity }}
 {{- $tolerations := or .Values.cleanup.tolerations .Values.tolerations }}
+{{- $securityContext := include "airflowSecurityContext" (list . .Values.cleanup) }}
 {{- if semverCompare ">= 1.21.x" (include "kubeVersion" .) }}
 apiVersion: batch/v1
 {{- else }}
@@ -70,6 +71,7 @@ spec:
           imagePullSecrets:
             - name: {{ template "registry_secret" . }}
           {{- end }}
+          securityContext: {{ $securityContext | nindent 12 }}
           containers:
             - name: airflow-cleanup-pods
               image: {{ template "airflow_image" . }}

chart/templates/flower/flower-deployment.yaml

@@ -23,6 +23,7 @@
 {{- $nodeSelector := or .Values.flower.nodeSelector .Values.nodeSelector }}
 {{- $affinity := or .Values.flower.affinity .Values.affinity }}
 {{- $tolerations := or .Values.flower.tolerations .Values.tolerations }}
+{{- $securityContext := include "airflowSecurityContext" (list . .Values.flower) }}
 kind: Deployment
 apiVersion: apps/v1
 metadata:
@@ -67,8 +68,7 @@ spec:
 {{ toYaml $tolerations | indent 8 }}
       serviceAccountName: {{ include "flower.serviceAccountName" . }}
       restartPolicy: Always
-      securityContext:
-        runAsUser: {{ .Values.uid }}
+      securityContext: {{ $securityContext | nindent 8 }}
       {{- if or .Values.registry.secretName .Values.registry.connection }}
       imagePullSecrets:
         - name: {{ template "registry_secret" . }}

chart/templates/jobs/create-user-job.yaml

@@ -22,6 +22,7 @@
 {{- $nodeSelector := or .Values.createUserJob.nodeSelector .Values.nodeSelector }}
 {{- $affinity := or .Values.createUserJob.affinity .Values.affinity }}
 {{- $tolerations := or .Values.createUserJob.tolerations .Values.tolerations }}
+{{- $securityContext := include "airflowSecurityContext" (list . .Values.createUserJob) }}
 apiVersion: batch/v1
 kind: Job
 metadata:
@@ -65,8 +66,7 @@ spec:
         {{- end }}
       {{- end }}
     spec:
-      securityContext:
-          runAsUser: {{ .Values.uid }}
+      securityContext: {{ $securityContext | nindent 8 }}
       restartPolicy: OnFailure
       nodeSelector:
 {{ toYaml $nodeSelector | indent 8 }}

chart/templates/jobs/migrate-database-job.yaml

@@ -21,6 +21,7 @@
 {{- $nodeSelector := or .Values.migrateDatabaseJob.nodeSelector .Values.nodeSelector }}
 {{- $affinity := or .Values.migrateDatabaseJob.affinity .Values.affinity }}
 {{- $tolerations := or .Values.migrateDatabaseJob.tolerations .Values.tolerations }}
+{{- $securityContext := include "airflowSecurityContext" (list . .Values.migrateDatabaseJob) }}
 apiVersion: batch/v1
 kind: Job
 metadata:
@@ -64,8 +65,7 @@ spec:
         {{- end }}
       {{- end }}
     spec:
-      securityContext:
-          runAsUser: {{ .Values.uid }}
+      securityContext: {{ $securityContext | nindent 8 }}
       restartPolicy: OnFailure
       nodeSelector:
 {{ toYaml $nodeSelector | indent 8 }}

chart/templates/scheduler/scheduler-deployment.yaml

@@ -31,6 +31,7 @@
 {{- $nodeSelector := or .Values.scheduler.nodeSelector .Values.nodeSelector }}
 {{- $affinity := or .Values.scheduler.affinity .Values.affinity }}
 {{- $tolerations := or .Values.scheduler.tolerations .Values.tolerations }}
+{{- $securityContext := include "airflowSecurityContext" (list . .Values.scheduler) }}
 
 kind: {{ if $stateful }}StatefulSet{{ else }}Deployment{{ end }}
 apiVersion: apps/v1
@@ -98,9 +99,7 @@ spec:
       restartPolicy: Always
       terminationGracePeriodSeconds: 10
       serviceAccountName: {{ include "scheduler.serviceAccountName" . }}
-      securityContext:
-        runAsUser: {{ .Values.uid }}
-        fsGroup: {{ .Values.gid }}
+      securityContext: {{ $securityContext | nindent 8 }}
       {{- if or .Values.registry.secretName .Values.registry.connection }}
       imagePullSecrets:
         - name: {{ template "registry_secret" . }}

chart/templates/statsd/statsd-deployment.yaml

@@ -22,6 +22,7 @@
 {{- $nodeSelector := or .Values.statsd.nodeSelector .Values.nodeSelector }}
 {{- $affinity := or .Values.statsd.affinity .Values.affinity }}
 {{- $tolerations := or .Values.statsd.tolerations .Values.tolerations }}
+{{- $securityContext := include "localSecurityContext" .Values.statsd }}
 kind: Deployment
 apiVersion: apps/v1
 metadata:
@@ -63,8 +64,7 @@ spec:
       tolerations:
 {{ toYaml $tolerations | indent 8 }}
       serviceAccountName: {{ include "statsd.serviceAccountName" . }}
-      securityContext:
-        runAsUser: {{ .Values.statsd.uid }}
+      securityContext: {{ $securityContext | nindent 8 }}
       restartPolicy: Always
       {{- if or .Values.registry.secretName .Values.registry.connection }}
       imagePullSecrets:

chart/templates/triggerer/triggerer-deployment.yaml

@@ -23,6 +23,7 @@
 {{- $nodeSelector := or .Values.nodeSelector .Values.triggerer.nodeSelector }}
 {{- $affinity := or .Values.affinity .Values.triggerer.affinity }}
 {{- $tolerations := or .Values.tolerations .Values.triggerer.tolerations }}
+{{- $securityContext := include "airflowSecurityContext" (list . .Values.triggerer) }}
 kind: Deployment
 apiVersion: apps/v1
 metadata:
@@ -81,9 +82,7 @@ spec:
       terminationGracePeriodSeconds: {{ .Values.triggerer.terminationGracePeriodSeconds }}
       restartPolicy: Always
       serviceAccountName: {{ include "triggerer.serviceAccountName" . }}
-      securityContext:
-        runAsUser: {{ .Values.uid }}
-        fsGroup: {{ .Values.gid }}
+      securityContext: {{ $securityContext | nindent 8 }}
       {{- if or .Values.registry.secretName .Values.registry.connection }}
       imagePullSecrets:
         - name: {{ template "registry_secret" . }}

chart/templates/webserver/webserver-deployment.yaml

@@ -21,6 +21,7 @@
 {{- $nodeSelector := or .Values.webserver.nodeSelector .Values.nodeSelector }}
 {{- $affinity := or .Values.webserver.affinity .Values.affinity }}
 {{- $tolerations := or .Values.webserver.tolerations .Values.tolerations }}
+{{- $securityContext := include "airflowSecurityContext" (list . .Values.webserver) }}
 kind: Deployment
 apiVersion: apps/v1
 metadata:
@@ -93,9 +94,7 @@ spec:
       tolerations:
 {{ toYaml $tolerations | indent 8 }}
       restartPolicy: Always
-      securityContext:
-        runAsUser: {{ .Values.uid }}
-        fsGroup: {{ .Values.gid }}
+      securityContext: {{ $securityContext | nindent 8 }}
       {{- if or .Values.registry.secretName .Values.registry.connection }}
       imagePullSecrets:
         - name: {{ template "registry_secret" . }}

chart/templates/workers/worker-deployment.yaml

@@ -23,6 +23,7 @@
 {{- $nodeSelector := or .Values.nodeSelector .Values.workers.nodeSelector }}
 {{- $affinity := or .Values.affinity .Values.workers.affinity }}
 {{- $tolerations := or .Values.tolerations .Values.workers.tolerations }}
+{{- $securityContext := include "airflowSecurityContext" (list . .Values.workers) }}
 kind: {{ if $persistence }}StatefulSet{{ else }}Deployment{{ end }}
 apiVersion: apps/v1
 metadata:
@@ -95,9 +96,7 @@ spec:
       terminationGracePeriodSeconds: {{ .Values.workers.terminationGracePeriodSeconds }}
       restartPolicy: Always
       serviceAccountName: {{ include "worker.serviceAccountName" . }}
-      securityContext:
-        runAsUser: {{ .Values.uid }}
-        fsGroup: {{ .Values.gid }}
+      securityContext: {{ $securityContext | nindent 8 }}
       {{- if or .Values.registry.secretName .Values.registry.connection }}
       imagePullSecrets:
         - name: {{ template "registry_secret" . }}
@@ -112,7 +111,7 @@ spec:
           command:
             - chown
             - -R
-            - "{{ .Values.uid }}:{{ .Values.gid }}"
+            - "{{ include "airflowSecurityContextIds" (list . .Values.workers) }}"
             - {{ template "airflow_logs" . }}
           securityContext:
             runAsUser: 0