Add workers.celery.securityContexts & workers.kubernetes.securityCont…

…exts

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

@@ -21,10 +21,10 @@
 {{- $affinity := or .Values.workers.affinity .Values.affinity }}
 {{- $tolerations := or .Values.workers.tolerations .Values.tolerations }}
 {{- $topologySpreadConstraints := or .Values.workers.topologySpreadConstraints .Values.topologySpreadConstraints }}
-{{- $securityContext := include "airflowPodSecurityContext" (list .Values.workers .Values) }}
+{{- $securityContext := include "airflowPodSecurityContext" (list .Values.workers.kubernetes .Values.workers .Values) }}
 {{- $containerSecurityContextKerberosSidecar := include "containerSecurityContext" (list .Values.workers.kerberosSidecar .Values) }}
 {{- $containerLifecycleHooksKerberosSidecar := or .Values.workers.kerberosSidecar.containerLifecycleHooks .Values.containerLifecycleHooks }}
-{{- $containerSecurityContext := include "containerSecurityContext" (list .Values.workers .Values) }}
+{{- $containerSecurityContext := include "containerSecurityContext" (list .Values.workers.kubernetes .Values.workers .Values) }}
 {{- $containerLifecycleHooks := or .Values.workers.containerLifecycleHooks .Values.containerLifecycleHooks }}
 {{- $safeToEvict := dict "cluster-autoscaler.kubernetes.io/safe-to-evict" (.Values.workers.safeToEvict | toString) }}
 {{- $podAnnotations := mergeOverwrite (deepCopy .Values.airflowPodAnnotations) $safeToEvict .Values.workers.podAnnotations }}

chart/templates/workers/worker-deployment.yaml

@@ -52,8 +52,8 @@
 {{- $tolerations := or .Values.workers.tolerations .Values.tolerations }}
 {{- $topologySpreadConstraints := or .Values.workers.topologySpreadConstraints .Values.topologySpreadConstraints }}
 {{- $revisionHistoryLimit := include "airflow.revisionHistoryLimit" (list .Values.workers.celery.revisionHistoryLimit .Values.workers.revisionHistoryLimit .Values.revisionHistoryLimit) }}
-{{- $securityContext := include "airflowPodSecurityContext" (list .Values.workers .Values) }}
-{{- $containerSecurityContext := include "containerSecurityContext" (list .Values.workers .Values) }}
+{{- $securityContext := include "airflowPodSecurityContext" (list .Values.workers.celery .Values.workers .Values) }}
+{{- $containerSecurityContext := include "containerSecurityContext" (list .Values.workers.celery .Values.workers .Values) }}
 {{- $containerSecurityContextPersistence := include "containerSecurityContext" (list .Values.workers.celery.persistence .Values.workers.persistence .Values) }}
 {{- $containerSecurityContextWaitForMigrations := include "containerSecurityContext" (list .Values.workers.waitForMigrations .Values) }}
 {{- $containerSecurityContextLogGroomerSidecar := include "containerSecurityContext" (list .Values.workers.logGroomerSidecar .Values) }}

chart/values.schema.json

@@ -2713,6 +2713,41 @@
                                 "additionalProperties": {}
                             }
                         },
+                        "securityContexts": {
+                            "description": "Security context definition for the Airflow Celery workers. If not set, the values from global `securityContexts` will be used.",
+                            "type": "object",
+                            "properties": {
+                                "pod": {
+                                    "description": "Pod security context definition.",
+                                    "type": "object",
+                                    "$ref": "#/definitions/io.k8s.api.core.v1.PodSecurityContext",
+                                    "default": {},
+                                    "examples": [
+                                        {
+                                            "runAsUser": 50000,
+                                            "runAsGroup": 0,
+                                            "fsGroup": 0
+                                        }
+                                    ]
+                                },
+                                "container": {
+                                    "description": "Container security context definition.",
+                                    "type": "object",
+                                    "$ref": "#/definitions/io.k8s.api.core.v1.SecurityContext",
+                                    "default": {},
+                                    "examples": [
+                                        {
+                                            "allowPrivilegeEscalation": false,
+                                            "capabilities": {
+                                                "drop": [
+                                                    "ALL"
+                                                ]
+                                            }
+                                        }
+                                    ]
+                                }
+                            }
+                        },
                         "serviceAccount": {
                             "description": "Create ServiceAccount.",
                             "type": "object",
@@ -2829,6 +2864,41 @@
                             },
                             "default": null
                         },
+                        "securityContexts": {
+                            "description": "Security context definition for the pod-template-file. If not set, the values from global `securityContexts` will be used.",
+                            "type": "object",
+                            "properties": {
+                                "pod": {
+                                    "description": "Pod security context definition.",
+                                    "type": "object",
+                                    "$ref": "#/definitions/io.k8s.api.core.v1.PodSecurityContext",
+                                    "default": {},
+                                    "examples": [
+                                        {
+                                            "runAsUser": 50000,
+                                            "runAsGroup": 0,
+                                            "fsGroup": 0
+                                        }
+                                    ]
+                                },
+                                "container": {
+                                    "description": "Container security context definition.",
+                                    "type": "object",
+                                    "$ref": "#/definitions/io.k8s.api.core.v1.SecurityContext",
+                                    "default": {},
+                                    "examples": [
+                                        {
+                                            "allowPrivilegeEscalation": false,
+                                            "capabilities": {
+                                                "drop": [
+                                                    "ALL"
+                                                ]
+                                            }
+                                        }
+                                    ]
+                                }
+                            }
+                        },
                         "serviceAccount": {
                             "description": "Create ServiceAccount.",
                             "type": "object",

chart/values.yaml

@@ -1074,6 +1074,7 @@ workers:
       failureThreshold: 5
       periodSeconds: 60
       command: ~
+
     # List of worker sets. Each item can override values from the parent 'workers' section.
     sets: []
     # sets:
@@ -1110,6 +1111,11 @@ workers:
     # while preserving its uniqueness and identity
     # podManagementPolicy: Parallel
 
+    # Detailed default security context for Airflow Celery workers for container and pod level
+    securityContexts:
+      pod: {}
+      container: {}
+
     # Create ServiceAccount for Airflow Celery workers
     serviceAccount:
       # default value is true
@@ -1156,6 +1162,11 @@ workers:
     # Command to use in pod-template-file (templated)
     command: ~
 
+    # Detailed default security context for pod-template-file for container and pod level
+    securityContexts:
+      pod: {}
+      container: {}
+
     # Create ServiceAccount for pods created with pod-template-file
     serviceAccount:
       # Auto mount service account token into the pod. Default value is true.

helm-tests/tests/helm_tests/airflow_aux/test_pod_template_file.py

@@ -637,6 +637,7 @@ def test_container_security_context_default(self):
             {"securityContexts": {"pod": {"runAsUser": 10}}},
             {"workers": {"securityContext": {"runAsUser": 10}}},
             {"workers": {"securityContexts": {"pod": {"runAsUser": 10}}}},
+            {"workers": {"kubernetes": {"securityContexts": {"pod": {"runAsUser": 10}}}}},
         ],
     )
     def test_pod_security_context_set(self, values):
@@ -653,6 +654,11 @@ def test_pod_security_context_set(self, values):
         [
             {"securityContexts": {"containers": {"allowPrivilegeEscalation": False}}},
             {"workers": {"securityContexts": {"container": {"allowPrivilegeEscalation": False}}}},
+            {
+                "workers": {
+                    "kubernetes": {"securityContexts": {"container": {"allowPrivilegeEscalation": False}}}
+                }
+            },
         ],
     )
     def test_container_security_context_set(self, values):
@@ -674,13 +680,29 @@ def test_container_security_context_set(self, values):
                 "securityContexts": {"pod": {"runAsUser": 5}},
                 "workers": {"securityContexts": {"pod": {"runAsUser": 10}}},
             },
+            {
+                "securityContexts": {"pod": {"runAsUser": 5}},
+                "workers": {"kubernetes": {"securityContexts": {"pod": {"runAsUser": 10}}}},
+            },
+            {
+                "workers": {
+                    "securityContexts": {"pod": {"runAsUser": 5}},
+                    "kubernetes": {"securityContexts": {"pod": {"runAsUser": 10}}},
+                },
+            },
             {"securityContext": {"runAsUser": 5}, "securityContexts": {"pod": {"runAsUser": 10}}},
             {
                 "workers": {
                     "securityContext": {"runAsUser": 5},
                     "securityContexts": {"pod": {"runAsUser": 10}},
                 }
             },
+            {
+                "workers": {
+                    "securityContext": {"runAsUser": 5},
+                    "kubernetes": {"securityContexts": {"pod": {"runAsUser": 10}}},
+                }
+            },
         ],
     )
     def test_pod_security_context_overwrite(self, values):
@@ -692,12 +714,30 @@ def test_pod_security_context_overwrite(self, values):
 
         assert jmespath.search("spec.securityContext", docs[0]) == {"runAsUser": 10}
 
-    def test_container_security_context_overwrite(self):
-        docs = render_chart(
-            values={
+    @pytest.mark.parametrize(
+        "values",
+        [
+            {
                 "securityContexts": {"containers": {"allowPrivilegeEscalation": True}},
                 "workers": {"securityContexts": {"container": {"allowPrivilegeEscalation": False}}},
             },
+            {
+                "securityContexts": {"containers": {"allowPrivilegeEscalation": True}},
+                "workers": {
+                    "kubernetes": {"securityContexts": {"container": {"allowPrivilegeEscalation": False}}}
+                },
+            },
+            {
+                "workers": {
+                    "securityContexts": {"container": {"allowPrivilegeEscalation": True}},
+                    "kubernetes": {"securityContexts": {"container": {"allowPrivilegeEscalation": False}}},
+                },
+            },
+        ],
+    )
+    def test_container_security_context_overwrite(self, values):
+        docs = render_chart(
+            values=values,
             show_only=["templates/pod-template-file.yaml"],
             chart_dir=self.temp_chart_dir,
         )

helm-tests/tests/helm_tests/security/test_security_context.py

@@ -333,15 +333,22 @@ def test_check_local_setting(self):
 
     # Test priority:
     # <local>.securityContexts > securityContexts > uid + gid
-    def test_check_local_setting_default_version(self):
+    @pytest.mark.parametrize(
+        "workers_values",
+        [
+            {"securityContexts": {"pod": {"runAsUser": 9000, "fsGroup": 90}}},
+            {"celery": {"securityContexts": {"pod": {"runAsUser": 9000, "fsGroup": 90}}}},
+        ],
+    )
+    def test_check_local_setting_default_version(self, workers_values):
         component_contexts = {"securityContexts": {"pod": {"runAsUser": 9000, "fsGroup": 90}}}
         docs = render_chart(
             values={
                 "uid": 3000,
                 "gid": 30,
                 "securityContexts": {"pod": {"runAsUser": 6000, "fsGroup": 60}},
                 "apiServer": component_contexts,
-                "workers": component_contexts,
+                "workers": workers_values,
                 "dagProcessor": component_contexts,
                 "flower": {"enabled": True, **component_contexts},
                 "scheduler": component_contexts,
@@ -446,7 +453,14 @@ def test_global_security_context(self):
         assert default_ctx_value_pod_redis == jmespath.search("spec.template.spec.securityContext", docs[-1])
 
     # Test securityContexts for main containers
-    def test_main_container_setting_airflow_2(self):
+    @pytest.mark.parametrize(
+        "workers_values",
+        [
+            {"securityContexts": {"container": {"allowPrivilegeEscalation": False}}},
+            {"celery": {"securityContexts": {"container": {"allowPrivilegeEscalation": False}}}},
+        ],
+    )
+    def test_main_container_setting_airflow_2(self, workers_values):
         ctx_value = {"allowPrivilegeEscalation": False}
         security_context = {"securityContexts": {"container": ctx_value}}
         docs = render_chart(
@@ -455,7 +469,7 @@ def test_main_container_setting_airflow_2(self):
                 "cleanup": {"enabled": True, **security_context},
                 "scheduler": {**security_context},
                 "webserver": {**security_context},
-                "workers": {**security_context},
+                "workers": workers_values,
                 "flower": {"enabled": True, **security_context},
                 "statsd": {**security_context},
                 "createUserJob": {**security_context},
@@ -488,7 +502,14 @@ def test_main_container_setting_airflow_2(self):
             assert ctx_value == jmespath.search("spec.template.spec.containers[0].securityContext", doc)
 
     # Test securityContexts for main containers
-    def test_main_container_setting(self):
+    @pytest.mark.parametrize(
+        "workers_values",
+        [
+            {"securityContexts": {"container": {"allowPrivilegeEscalation": False}}},
+            {"celery": {"securityContexts": {"container": {"allowPrivilegeEscalation": False}}}},
+        ],
+    )
+    def test_main_container_setting(self, workers_values):
         ctx_value = {"allowPrivilegeEscalation": False}
         security_context = {"securityContexts": {"container": ctx_value}}
         docs = render_chart(
@@ -498,7 +519,7 @@ def test_main_container_setting(self):
                 "scheduler": {**security_context},
                 "dagProcessor": {**security_context},
                 "apiServer": {**security_context},
-                "workers": {**security_context},
+                "workers": workers_values,
                 "flower": {"enabled": True, **security_context},
                 "statsd": {**security_context},
                 "createUserJob": {**security_context},
@@ -677,7 +698,14 @@ def test_volume_permissions_init_container_setting(self, workers_values):
         }
 
     # Test securityContexts for main pods
-    def test_main_pod_setting_airflow_2(self):
+    @pytest.mark.parametrize(
+        "workers_values",
+        [
+            {"securityContexts": {"pod": {"runAsUser": 7000}}},
+            {"celery": {"securityContexts": {"pod": {"runAsUser": 7000}}}},
+        ],
+    )
+    def test_main_pod_setting_airflow_2(self, workers_values):
         ctx_value = {"runAsUser": 7000}
         security_context = {"securityContexts": {"pod": ctx_value}}
         docs = render_chart(
@@ -686,7 +714,7 @@ def test_main_pod_setting_airflow_2(self):
                 "cleanup": {"enabled": True, **security_context},
                 "scheduler": {**security_context},
                 "webserver": {**security_context},
-                "workers": {**security_context},
+                "workers": workers_values,
                 "flower": {"enabled": True, **security_context},
                 "statsd": {**security_context},
                 "createUserJob": {**security_context},
@@ -716,7 +744,14 @@ def test_main_pod_setting_airflow_2(self):
         for doc in docs[1:]:
             assert ctx_value == jmespath.search("spec.template.spec.securityContext", doc)
 
-    def test_main_pod_setting(self):
+    @pytest.mark.parametrize(
+        "workers_values",
+        [
+            {"securityContexts": {"pod": {"runAsUser": 7000}}},
+            {"celery": {"securityContexts": {"pod": {"runAsUser": 7000}}}},
+        ],
+    )
+    def test_main_pod_setting(self, workers_values):
         ctx_value = {"runAsUser": 7000}
         security_context = {"securityContexts": {"pod": ctx_value}}
         docs = render_chart(
@@ -725,7 +760,7 @@ def test_main_pod_setting(self):
                 "cleanup": {"enabled": True, **security_context},
                 "scheduler": {**security_context},
                 "apiServer": {**security_context},
-                "workers": {**security_context},
+                "workers": workers_values,
                 "flower": {"enabled": True, **security_context},
                 "statsd": {**security_context},
                 "createUserJob": {**security_context},
@@ -852,7 +887,20 @@ def test_deprecated_overwrite_global(self):
             assert jmespath.search("spec.template.spec.securityContext.runAsUser", doc) == 9000
             assert jmespath.search("spec.template.spec.securityContext.fsGroup", doc) == 90
 
-    def test_deprecated_overwrite_local(self):
+    @pytest.mark.parametrize(
+        "workers_values",
+        [
+            {
+                "securityContext": {"runAsUser": 6000, "fsGroup": 60},
+                "securityContexts": {"pod": {"runAsUser": 9000, "fsGroup": 90}},
+            },
+            {
+                "securityContext": {"runAsUser": 6000, "fsGroup": 60},
+                "celery": {"securityContexts": {"pod": {"runAsUser": 9000, "fsGroup": 90}}},
+            },
+        ],
+    )
+    def test_deprecated_overwrite_local(self, workers_values):
         context = {
             "securityContext": {"runAsUser": 6000, "fsGroup": 60},
             "securityContexts": {"pod": {"runAsUser": 9000, "fsGroup": 90}},
@@ -863,7 +911,7 @@ def test_deprecated_overwrite_local(self):
                 "flower": context,
                 "scheduler": context,
                 "triggerer": context,
-                "workers": context,
+                "workers": workers_values,
                 "dagProcessor": context,
             },
             show_only=[
@@ -878,3 +926,17 @@ def test_deprecated_overwrite_local(self):
         for doc in docs:
             assert jmespath.search("spec.template.spec.securityContext.runAsUser", doc) == 9000
             assert jmespath.search("spec.template.spec.securityContext.fsGroup", doc) == 90
+
+    def test_workers_overwrite_local(self):
+        docs = render_chart(
+            values={
+                "workers": {
+                    "securityContexts": {"pod": {"runAsUser": 6000, "fsGroup": 60}},
+                    "celery": {"securityContexts": {"pod": {"runAsUser": 9000, "fsGroup": 90}}},
+                },
+            },
+            show_only=["templates/workers/worker-deployment.yaml"],
+        )
+
+        assert jmespath.search("spec.template.spec.securityContext.runAsUser", docs[0]) == 9000
+        assert jmespath.search("spec.template.spec.securityContext.fsGroup", docs[0]) == 90