From 8d66d4416d354524a9e0cb66833cbd35d1ca2333 Mon Sep 17 00:00:00 2001
From: pgvishnuram <81585115+pgvishnuram@users.noreply.github.com>
Date: Mon, 4 Nov 2024 10:48:25 +0530
Subject: [PATCH] support pod securitycontext and container security context
 (#496)

* support pod securitycontext and container security context

* fix test cases

* rename securitycontext for dag service

* rename security context similar to airflow components

* extraEnv is a dict, not a list

* Revert "extraEnv is a dict, not a list"

This reverts commit 6c5f9dba73a84e49552df89a0cd2708278330e2e.

* Update templates/dag-deploy/dag-server-statefulset.yaml

Co-authored-by: Daniel Hoherd <danielh@astronomer.io>

* fixing test and securitycontexts.pod

* update test

* update tests

* rename tests

---------

Co-authored-by: Daniel Hoherd <danielh@astronomer.io>
Co-authored-by: Daniel Hoherd <daniel.hoherd@gmail.com>
Co-authored-by: Chioma Onyekpere <chioma.onyekpere@gmail.com>
Co-authored-by: Shubham-astro <shubham.singh@astronomer.io>
---
 .../dag-deploy/dag-server-statefulset.yaml    |  3 ++-
 tests/chart/test_dag_server_statefulset.py    | 24 +++++++++++++++----
 values.yaml                                   |  6 ++---
 3 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/templates/dag-deploy/dag-server-statefulset.yaml b/templates/dag-deploy/dag-server-statefulset.yaml
index d425e3b0..83d7fc9d 100644
--- a/templates/dag-deploy/dag-server-statefulset.yaml
+++ b/templates/dag-deploy/dag-server-statefulset.yaml
@@ -43,7 +43,7 @@ spec:
         - name: {{ template "astro.registry_secret" . }}
       {{- end }}
       serviceAccountName: {{ template "astro.dagDeploy.serviceAccountName" . }}
-      securityContext: {{ toYaml .Values.dagDeploy.securityContext | nindent 8 }}
+      securityContext: {{ toYaml .Values.dagDeploy.securityContexts.pod | nindent 8 }}
       containers:
         - name: dag-server
           image: "{{ .Values.dagDeploy.images.dagServer.repository }}:{{ .Values.dagDeploy.images.dagServer.tag }}"
@@ -59,6 +59,7 @@ spec:
           ports:
           - name: server
             containerPort: {{ .Values.dagDeploy.ports.dagServerHttp }}
+          securityContext: {{ toYaml .Values.dagDeploy.securityContexts.container | nindent 12 }}
           resources:
             {{- toYaml .Values.dagDeploy.resources | nindent 12 }}
           env:
diff --git a/tests/chart/test_dag_server_statefulset.py b/tests/chart/test_dag_server_statefulset.py
index e1c44e2e..90a0f567 100644
--- a/tests/chart/test_dag_server_statefulset.py
+++ b/tests/chart/test_dag_server_statefulset.py
@@ -106,8 +106,23 @@ def test_dag_server_statefulset_with_resource_overrides(self, kube_version):
 
     def test_dag_server_statefulset_with_securitycontext_overrides(self, kube_version):
         """Test that dag-server statefulset are configurable with custom securitycontext."""
-        dag_serversecuritycontext = {"runAsUser": 12345, "privileged": True}
-        values = {"dagDeploy": {"enabled": True, "securityContext": dag_serversecuritycontext}}
+        dag_server_pod_securitycontext = {
+            "runAsUser": 12345,
+            "allowPrivilegeEscalation": True,
+            "runAsGroup": 1000,
+            "fsGroup": 2000,
+            "readOnlyRootFilesystem": True,
+        }
+        dag_server_container_securitycontext = {"allowPrivilegeEscalation": False}
+        values = {
+            "dagDeploy": {
+                "enabled": True,
+                "securityContexts": {
+                    "pod": dag_server_pod_securitycontext,
+                    "container": dag_server_container_securitycontext,
+                },
+            }
+        }
 
         docs = render_chart(
             kube_version=kube_version,
@@ -118,8 +133,9 @@ def test_dag_server_statefulset_with_securitycontext_overrides(self, kube_versio
         doc = docs[0]
 
         common_default_tests(doc)
-
-        assert dag_serversecuritycontext == doc["spec"]["template"]["spec"]["securityContext"]
+        spec = doc["spec"]["template"]["spec"]
+        assert dag_server_pod_securitycontext == spec["securityContext"]
+        assert dag_server_container_securitycontext == spec["containers"][0]["securityContext"]
 
     def test_dag_server_statefulset_with_custom_registry_secret(self, kube_version):
         """Test dag-server statefulset with custom registry secret."""
diff --git a/values.yaml b/values.yaml
index a12a60ee..77aaf060 100644
--- a/values.yaml
+++ b/values.yaml
@@ -594,9 +594,9 @@ dagDeploy:
   # when it wants to scale a node down.
   safeToEvict: true
 
-  securityContext: {}
-  #  runAsUser: 999
-  #  runAsGroup: 0
+  securityContexts:
+    pod: {}
+    container: {}
 
 gitSyncRelay:
   enabled: ~