Skip to content

Commit c85246c

Browse files
authored
Add missing support for: securityContexts and containerLifecycleHooks (apache#60677)
* keep consistency * support kerberos init securityContext/lifecycle in workers
1 parent cdee423 commit c85246c

File tree

6 files changed

+318
-0
lines changed

6 files changed

+318
-0
lines changed

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@
2424
{{- $securityContext := include "airflowPodSecurityContext" (list .Values.workers.kubernetes .Values.workers .Values) }}
2525
{{- $containerSecurityContextKerberosSidecar := include "containerSecurityContext" (list .Values.workers.kerberosSidecar .Values) }}
2626
{{- $containerLifecycleHooksKerberosSidecar := or .Values.workers.kerberosSidecar.containerLifecycleHooks .Values.containerLifecycleHooks }}
27+
{{- $containerSecurityContextKerberosInitContainer := include "containerSecurityContext" (list .Values.workers.kubernetes.kerberosInitContainer .Values.workers.kerberosInitContainer .Values) }}
28+
{{- $containerLifecycleHooksKerberosInitContainer := or .Values.workers.kubernetes.kerberosInitContainer.containerLifecycleHooks .Values.workers.kerberosInitContainer.containerLifecycleHooks .Values.containerLifecycleHooks }}
2729
{{- $containerSecurityContext := include "containerSecurityContext" (list .Values.workers.kubernetes .Values.workers .Values) }}
2830
{{- $containerLifecycleHooks := or .Values.workers.containerLifecycleHooks .Values.containerLifecycleHooks }}
2931
{{- $safeToEvict := dict "cluster-autoscaler.kubernetes.io/safe-to-evict" (.Values.workers.safeToEvict | toString) }}
@@ -57,6 +59,10 @@ spec:
5759
- name: kerberos-init
5860
image: {{ template "airflow_image" . }}
5961
imagePullPolicy: {{ .Values.images.airflow.pullPolicy }}
62+
securityContext: {{ $containerSecurityContextKerberosInitContainer | nindent 8 }}
63+
{{- if $containerLifecycleHooksKerberosInitContainer }}
64+
lifecycle: {{- tpl (toYaml $containerLifecycleHooksKerberosInitContainer) . | nindent 8 }}
65+
{{- end }}
6066
args: ["kerberos", "-o"]
6167
resources: {{- toYaml (.Values.workers.kubernetes.kerberosInitContainer.resources | default .Values.workers.kerberosInitContainer.resources) | nindent 8 }}
6268
volumeMounts:

chart/templates/workers/worker-deployment.yaml

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,8 @@
4848
{{- $containerSecurityContextWaitForMigrations := include "containerSecurityContext" (list .Values.workers.waitForMigrations .Values) }}
4949
{{- $containerSecurityContextLogGroomerSidecar := include "containerSecurityContext" (list .Values.workers.logGroomerSidecar .Values) }}
5050
{{- $containerSecurityContextKerberosSidecar := include "containerSecurityContext" (list .Values.workers.kerberosSidecar .Values) }}
51+
{{- $containerSecurityContextKerberosInitContainer := include "containerSecurityContext" (list .Values.workers.kerberosInitContainer .Values) }}
52+
{{- $containerLifecycleHooksKerberosInitContainer := or .Values.workers.kerberosInitContainer.containerLifecycleHooks .Values.containerLifecycleHooks }}
5153
{{- $containerLifecycleHooks := or .Values.workers.containerLifecycleHooks .Values.containerLifecycleHooks }}
5254
{{- $containerLifecycleHooksLogGroomerSidecar := or .Values.workers.logGroomerSidecar.containerLifecycleHooks .Values.containerLifecycleHooks }}
5355
{{- $containerLifecycleHooksKerberosSidecar := or .Values.workers.kerberosSidecar.containerLifecycleHooks .Values.containerLifecycleHooks }}
@@ -180,6 +182,10 @@ spec:
180182
{{- if and (semverCompare ">=2.8.0" .Values.airflowVersion) .Values.workers.kerberosInitContainer.enabled }}
181183
- name: kerberos-init
182184
image: {{ template "airflow_image" . }}
185+
securityContext: {{ $containerSecurityContextKerberosInitContainer | nindent 12 }}
186+
{{- if $containerLifecycleHooksKerberosInitContainer }}
187+
lifecycle: {{- tpl (toYaml $containerLifecycleHooksKerberosInitContainer) . | nindent 12 }}
188+
{{- end }}
183189
imagePullPolicy: {{ .Values.images.airflow.pullPolicy }}
184190
args: ["kerberos", "-o"]
185191
resources: {{- toYaml .Values.workers.kerberosInitContainer.resources | nindent 12 }}

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

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1384,3 +1384,105 @@ def test_kerberos_init_container_resources(self, workers_values):
13841384
"memory": "4Mi",
13851385
},
13861386
}
1387+
1388+
@pytest.mark.parametrize(
1389+
"workers_values",
1390+
[
1391+
{
1392+
"kerberosInitContainer": {
1393+
"enabled": True,
1394+
"securityContexts": {"container": {"runAsUser": 2000}},
1395+
}
1396+
},
1397+
{
1398+
"kubernetes": {
1399+
"kerberosInitContainer": {
1400+
"enabled": True,
1401+
"securityContexts": {"container": {"runAsUser": 2000}},
1402+
}
1403+
}
1404+
},
1405+
{
1406+
"kerberosInitContainer": {
1407+
"enabled": True,
1408+
"securityContexts": {"container": {"runAsUser": 1000}},
1409+
},
1410+
"kubernetes": {
1411+
"kerberosInitContainer": {
1412+
"enabled": True,
1413+
"securityContexts": {"container": {"runAsUser": 2000}},
1414+
}
1415+
},
1416+
},
1417+
],
1418+
)
1419+
def test_kerberos_init_container_security_context(self, workers_values):
1420+
docs = render_chart(
1421+
values={
1422+
"workers": workers_values,
1423+
},
1424+
show_only=["templates/pod-template-file.yaml"],
1425+
chart_dir=self.temp_chart_dir,
1426+
)
1427+
1428+
assert jmespath.search(
1429+
"spec.initContainers[?name=='kerberos-init'] | [0].securityContext", docs[0]
1430+
) == {"runAsUser": 2000}
1431+
1432+
@pytest.mark.parametrize(
1433+
("workers_values", "expected"),
1434+
[
1435+
(
1436+
{
1437+
"kerberosInitContainer": {
1438+
"enabled": True,
1439+
"containerLifecycleHooks": {"postStart": {"exec": {"command": ["echo", "base"]}}},
1440+
}
1441+
},
1442+
{"postStart": {"exec": {"command": ["echo", "base"]}}},
1443+
),
1444+
(
1445+
{
1446+
"kubernetes": {
1447+
"kerberosInitContainer": {
1448+
"enabled": True,
1449+
"containerLifecycleHooks": {
1450+
"postStart": {"exec": {"command": ["echo", "kubernetes"]}}
1451+
},
1452+
}
1453+
}
1454+
},
1455+
{"postStart": {"exec": {"command": ["echo", "kubernetes"]}}},
1456+
),
1457+
(
1458+
{
1459+
"kerberosInitContainer": {
1460+
"enabled": True,
1461+
"containerLifecycleHooks": {"preStop": {"exec": {"command": ["echo", "base"]}}},
1462+
},
1463+
"kubernetes": {
1464+
"kerberosInitContainer": {
1465+
"enabled": True,
1466+
"containerLifecycleHooks": {
1467+
"postStart": {"exec": {"command": ["echo", "kubernetes"]}}
1468+
},
1469+
}
1470+
},
1471+
},
1472+
{"postStart": {"exec": {"command": ["echo", "kubernetes"]}}},
1473+
),
1474+
],
1475+
)
1476+
def test_kerberos_init_container_lifecycle_hooks(self, workers_values, expected):
1477+
docs = render_chart(
1478+
values={
1479+
"workers": workers_values,
1480+
},
1481+
show_only=["templates/pod-template-file.yaml"],
1482+
chart_dir=self.temp_chart_dir,
1483+
)
1484+
1485+
assert (
1486+
jmespath.search("spec.initContainers[?name=='kerberos-init'] | [0].lifecycle", docs[0])
1487+
== expected
1488+
)

helm-tests/tests/helm_tests/airflow_core/test_worker.py

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1050,6 +1050,65 @@ def test_kerberos_init_container_resources(self, workers_values):
10501050
},
10511051
}
10521052

1053+
@pytest.mark.parametrize(
1054+
("workers_values", "expected"),
1055+
[
1056+
(
1057+
{
1058+
"kerberosInitContainer": {
1059+
"enabled": True,
1060+
"containerLifecycleHooks": {"postStart": {"exec": {"command": ["echo", "base"]}}},
1061+
}
1062+
},
1063+
{"postStart": {"exec": {"command": ["echo", "base"]}}},
1064+
),
1065+
(
1066+
{
1067+
"celery": {
1068+
"kerberosInitContainer": {
1069+
"enabled": True,
1070+
"containerLifecycleHooks": {
1071+
"postStart": {"exec": {"command": ["echo", "celery"]}}
1072+
},
1073+
}
1074+
}
1075+
},
1076+
{"postStart": {"exec": {"command": ["echo", "celery"]}}},
1077+
),
1078+
(
1079+
{
1080+
"kerberosInitContainer": {
1081+
"enabled": True,
1082+
"containerLifecycleHooks": {"postStart": {"exec": {"command": ["echo", "base"]}}},
1083+
},
1084+
"celery": {
1085+
"kerberosInitContainer": {
1086+
"enabled": True,
1087+
"containerLifecycleHooks": {
1088+
"postStart": {"exec": {"command": ["echo", "celery"]}}
1089+
},
1090+
}
1091+
},
1092+
},
1093+
{"postStart": {"exec": {"command": ["echo", "celery"]}}},
1094+
),
1095+
],
1096+
)
1097+
def test_kerberos_init_container_lifecycle_hooks(self, workers_values, expected):
1098+
docs = render_chart(
1099+
values={
1100+
"workers": workers_values,
1101+
},
1102+
show_only=["templates/workers/worker-deployment.yaml"],
1103+
)
1104+
1105+
assert (
1106+
jmespath.search(
1107+
"spec.template.spec.initContainers[?name=='kerberos-init'] | [0].lifecycle", docs[0]
1108+
)
1109+
== expected
1110+
)
1111+
10531112
@pytest.mark.parametrize(
10541113
("airflow_version", "expected_arg"),
10551114
[

helm-tests/tests/helm_tests/airflow_core/test_worker_sets.py

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -841,6 +841,108 @@ def test_overwrite_kerberos_init_container_resources(self, values):
841841
"limits": {"cpu": "3m", "memory": "4Mi"},
842842
}
843843

844+
@pytest.mark.parametrize(
845+
"values",
846+
[
847+
{
848+
"celery": {
849+
"enableDefault": False,
850+
"sets": [
851+
{
852+
"name": "test",
853+
"kerberosInitContainer": {
854+
"enabled": True,
855+
"securityContexts": {
856+
"container": {"runAsUser": 10},
857+
},
858+
},
859+
}
860+
],
861+
}
862+
},
863+
{
864+
"kerberosInitContainer": {
865+
"securityContexts": {
866+
"container": {"allowPrivilegeEscalation": False},
867+
}
868+
},
869+
"celery": {
870+
"enableDefault": False,
871+
"sets": [
872+
{
873+
"name": "test",
874+
"kerberosInitContainer": {
875+
"enabled": True,
876+
"securityContexts": {
877+
"container": {"runAsUser": 10},
878+
},
879+
},
880+
}
881+
],
882+
},
883+
},
884+
],
885+
)
886+
def test_overwrite_kerberos_init_container_security_context(self, values):
887+
docs = render_chart(
888+
values={"workers": values},
889+
show_only=["templates/workers/worker-deployment.yaml"],
890+
)
891+
892+
assert jmespath.search(
893+
"spec.template.spec.initContainers[?name=='kerberos-init'] | [0].securityContext", docs[0]
894+
) == {"runAsUser": 10}
895+
896+
@pytest.mark.parametrize(
897+
"values",
898+
[
899+
{
900+
"celery": {
901+
"enableDefault": False,
902+
"sets": [
903+
{
904+
"name": "test",
905+
"kerberosInitContainer": {
906+
"enabled": True,
907+
"containerLifecycleHooks": {
908+
"postStart": {"exec": {"command": ["echo", "{{ .Release.Name }}"]}},
909+
},
910+
},
911+
}
912+
],
913+
}
914+
},
915+
{
916+
"kerberosInitContainer": {
917+
"containerLifecycleHooks": {"preStop": {"exec": {"command": ["echo", "test"]}}}
918+
},
919+
"celery": {
920+
"enableDefault": False,
921+
"sets": [
922+
{
923+
"name": "test",
924+
"kerberosInitContainer": {
925+
"enabled": True,
926+
"containerLifecycleHooks": {
927+
"postStart": {"exec": {"command": ["echo", "{{ .Release.Name }}"]}},
928+
},
929+
},
930+
}
931+
],
932+
},
933+
},
934+
],
935+
)
936+
def test_overwrite_kerberos_init_container_lifecycle_hooks(self, values):
937+
docs = render_chart(
938+
values={"workers": values},
939+
show_only=["templates/workers/worker-deployment.yaml"],
940+
)
941+
942+
assert jmespath.search(
943+
"spec.template.spec.initContainers[?name=='kerberos-init'] | [0].lifecycle", docs[0]
944+
) == {"postStart": {"exec": {"command": ["echo", "release-name"]}}}
945+
844946
def test_overwrite_container_lifecycle_hooks(self):
845947
docs = render_chart(
846948
values={

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

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -632,6 +632,49 @@ def test_worker_kerberos_container_setting(self):
632632

633633
assert ctx_value == jmespath.search("spec.template.spec.containers[2].securityContext", docs[0])
634634

635+
@pytest.mark.parametrize(
636+
"workers_values",
637+
[
638+
{
639+
"kerberosInitContainer": {
640+
"enabled": True,
641+
"securityContexts": {"container": {"runAsUser": 2000}},
642+
}
643+
},
644+
{
645+
"celery": {
646+
"kerberosInitContainer": {
647+
"enabled": True,
648+
"securityContexts": {"container": {"runAsUser": 2000}},
649+
}
650+
}
651+
},
652+
{
653+
"kerberosInitContainer": {
654+
"enabled": True,
655+
"securityContexts": {"container": {"runAsUser": 1000}},
656+
},
657+
"celery": {
658+
"kerberosInitContainer": {
659+
"enabled": True,
660+
"securityContexts": {"container": {"runAsUser": 2000}},
661+
}
662+
},
663+
},
664+
],
665+
)
666+
def test_worker_kerberos_init_container_security_context(self, workers_values):
667+
docs = render_chart(
668+
values={
669+
"workers": workers_values,
670+
},
671+
show_only=["templates/workers/worker-deployment.yaml"],
672+
)
673+
674+
assert jmespath.search(
675+
"spec.template.spec.initContainers[?name=='kerberos-init'] | [0].securityContext", docs[0]
676+
) == {"runAsUser": 2000}
677+
635678
# Test securityContexts for the wait-for-migrations init containers
636679
def test_wait_for_migrations_init_container_setting_airflow_2(self):
637680
ctx_value = {"allowPrivilegeEscalation": False}

0 commit comments

Comments
 (0)