Skip to content

Commit f22dbc0

Browse files
MiretplMiretpl
authored
Add workers.celery.containerLifecycleHooks & workers.kubernetes.containerLifecycleHooks (#61369)
* Add workers.celery.containerLifecycleHooks & workers.kubernetes.containerLifecycleHooks * Refactor some new tests & simplify them
1 parent 00c69f5 commit f22dbc0

File tree

8 files changed

+282
-130
lines changed

8 files changed

+282
-130
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
{{- $containerSecurityContextKerberosInitContainer := include "containerSecurityContext" (list .Values.workers.kubernetes.kerberosInitContainer .Values.workers.kerberosInitContainer .Values) }}
2828
{{- $containerLifecycleHooksKerberosInitContainer := or .Values.workers.kubernetes.kerberosInitContainer.containerLifecycleHooks .Values.workers.kerberosInitContainer.containerLifecycleHooks .Values.containerLifecycleHooks }}
2929
{{- $containerSecurityContext := include "containerSecurityContext" (list .Values.workers.kubernetes .Values.workers .Values) }}
30-
{{- $containerLifecycleHooks := or .Values.workers.containerLifecycleHooks .Values.containerLifecycleHooks }}
30+
{{- $containerLifecycleHooks := or .Values.workers.kubernetes.containerLifecycleHooks .Values.workers.containerLifecycleHooks .Values.containerLifecycleHooks }}
3131
{{- $safeToEvict := dict "cluster-autoscaler.kubernetes.io/safe-to-evict" (.Values.workers.safeToEvict | toString) }}
3232
{{- $podAnnotations := mergeOverwrite (deepCopy .Values.airflowPodAnnotations) $safeToEvict .Values.workers.podAnnotations }}
3333
{{- $schedulerName := or .Values.workers.schedulerName .Values.schedulerName }}

‎chart/values.schema.json‎

Lines changed: 59 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2452,7 +2452,7 @@
24522452
]
24532453
},
24542454
"containerLifecycleHooks": {
2455-
"description": "Container Lifecycle Hooks definition for Airflow Celery workers and pods created with pod-template-file. If not set, the values from global `containerLifecycleHooks` will be used.",
2455+
"description": "Container Lifecycle Hooks definition for Airflow Celery workers and pods created with pod-template-file. If not set, the values from global `containerLifecycleHooks` will be used. Use `workers.celery.containerLifecycleHooks` and/or `workers.kubernetes.containerLifecycleHooks` to separate value between Celery workers and pod-template-file.",
24562456
"type": "object",
24572457
"$ref": "#/definitions/io.k8s.api.core.v1.Lifecycle",
24582458
"default": {},
@@ -2850,6 +2850,35 @@
28502850
}
28512851
}
28522852
},
2853+
"containerLifecycleHooks": {
2854+
"description": "Container Lifecycle Hooks definition for Airflow Celery workers. If not set, the values from `workers.containerLifecycleHooks` will be used.",
2855+
"type": "object",
2856+
"$ref": "#/definitions/io.k8s.api.core.v1.Lifecycle",
2857+
"default": {},
2858+
"x-docsSection": "Kubernetes",
2859+
"examples": [
2860+
{
2861+
"postStart": {
2862+
"exec": {
2863+
"command": [
2864+
"/bin/sh",
2865+
"-c",
2866+
"echo postStart handler > /usr/share/message"
2867+
]
2868+
}
2869+
},
2870+
"preStop": {
2871+
"exec": {
2872+
"command": [
2873+
"/bin/sh",
2874+
"-c",
2875+
"echo preStop handler > /usr/share/message"
2876+
]
2877+
}
2878+
}
2879+
}
2880+
]
2881+
},
28532882
"persistence": {
28542883
"description": "Persistence configuration for Airflow Celery workers.",
28552884
"type": "object",
@@ -3065,6 +3094,35 @@
30653094
}
30663095
}
30673096
},
3097+
"containerLifecycleHooks": {
3098+
"description": "Container Lifecycle Hooks definition for pods created with pod-template-file. If not set, the values from `workers.containerLifecycleHooks` will be used.",
3099+
"type": "object",
3100+
"$ref": "#/definitions/io.k8s.api.core.v1.Lifecycle",
3101+
"default": {},
3102+
"x-docsSection": "Kubernetes",
3103+
"examples": [
3104+
{
3105+
"postStart": {
3106+
"exec": {
3107+
"command": [
3108+
"/bin/sh",
3109+
"-c",
3110+
"echo postStart handler > /usr/share/message"
3111+
]
3112+
}
3113+
},
3114+
"preStop": {
3115+
"exec": {
3116+
"command": [
3117+
"/bin/sh",
3118+
"-c",
3119+
"echo preStop handler > /usr/share/message"
3120+
]
3121+
}
3122+
}
3123+
}
3124+
]
3125+
},
30683126
"kerberosInitContainer": {
30693127
"description": "Kerberos init container for pods created with pod-template-file.",
30703128
"type": "object",

‎chart/values.yaml‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -699,6 +699,8 @@ workers:
699699

700700
# Container level Lifecycle Hooks definition for
701701
# Airflow Celery workers and pods created with pod-template-file
702+
# Use workers.celery.containerLifecycleHooks and/or workers.kubernetes.containerLifecycleHooks
703+
# to separate value between Celery workers and pod-template-file.
702704
containerLifecycleHooks: {}
703705

704706
# Worker pod disruption budget
@@ -1105,6 +1107,9 @@ workers:
11051107
pod: {}
11061108
container: {}
11071109

1110+
# Container level Lifecycle Hooks definition for Airflow Celery workers
1111+
containerLifecycleHooks: {}
1112+
11081113
# Persistence volume configuration for Airflow Celery workers
11091114
persistence:
11101115
# Enable persistent volumes
@@ -1166,6 +1171,9 @@ workers:
11661171
pod: {}
11671172
container: {}
11681173

1174+
# Container level Lifecycle Hooks definition for pods created with pod-template-file
1175+
containerLifecycleHooks: {}
1176+
11691177
# Kerberos init container configuration for pods created with pod-template-file
11701178
# If not set, the values from `workers.kubernetesInitContainer` section will be used.
11711179
kerberosInitContainer:

‎helm-tests/tests/helm_tests/airflow_aux/test_container_lifecycle.py‎

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -107,16 +107,38 @@ def test_global_setting_external(self, hook_type):
107107
assert jmespath.search("spec.template.spec.containers[0].lifecycle", doc) != LIFECYCLE_PARSED
108108

109109
# <local>.containerLifecycleWebhooks > containerLifecycleWebhooks
110-
@pytest.mark.parametrize("hook_type", ["preStop", "postStart"])
111-
def test_check_main_container_setting(self, hook_type):
110+
@pytest.mark.parametrize(
111+
("hook_type", "workers_values"),
112+
[
113+
("preStop", {"containerLifecycleHooks": {"preStop": LIFECYCLE_TEMPLATE}}),
114+
("preStop", {"celery": {"containerLifecycleHooks": {"preStop": LIFECYCLE_TEMPLATE}}}),
115+
(
116+
"preStop",
117+
{
118+
"containerLifecycleHooks": {"postStart": LIFECYCLE_TEMPLATE},
119+
"celery": {"containerLifecycleHooks": {"preStop": LIFECYCLE_TEMPLATE}},
120+
},
121+
),
122+
("postStart", {"containerLifecycleHooks": {"postStart": LIFECYCLE_TEMPLATE}}),
123+
("postStart", {"celery": {"containerLifecycleHooks": {"postStart": LIFECYCLE_TEMPLATE}}}),
124+
(
125+
"postStart",
126+
{
127+
"containerLifecycleHooks": {"preStop": LIFECYCLE_TEMPLATE},
128+
"celery": {"containerLifecycleHooks": {"postStart": LIFECYCLE_TEMPLATE}},
129+
},
130+
),
131+
],
132+
)
133+
def test_check_main_container_setting(self, hook_type, workers_values):
112134
docs = render_chart(
113135
name=RELEASE_NAME,
114136
values={
115137
"containerLifecycleHooks": {hook_type: LIFECYCLE_TEMPLATE},
116138
"flower": {"containerLifecycleHooks": {hook_type: LIFECYCLE_TEMPLATE}},
117139
"scheduler": {"containerLifecycleHooks": {hook_type: LIFECYCLE_TEMPLATE}},
118140
"webserver": {"containerLifecycleHooks": {hook_type: LIFECYCLE_TEMPLATE}},
119-
"workers": {"containerLifecycleHooks": {hook_type: LIFECYCLE_TEMPLATE}},
141+
"workers": workers_values,
120142
"migrateDatabaseJob": {"containerLifecycleHooks": {hook_type: LIFECYCLE_TEMPLATE}},
121143
"triggerer": {"containerLifecycleHooks": {hook_type: LIFECYCLE_TEMPLATE}},
122144
"redis": {"containerLifecycleHooks": {hook_type: LIFECYCLE_TEMPLATE}},

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

Lines changed: 61 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -1075,22 +1075,43 @@ def test_workers_priority_class_name(self):
10751075

10761076
assert jmespath.search("spec.priorityClassName", docs[0]) == "test-priority"
10771077

1078-
def test_workers_container_lifecycle_webhooks_are_configurable(self):
1079-
docs = render_chart(
1080-
name="test-release",
1081-
values={
1082-
"workers": {
1078+
@pytest.mark.parametrize(
1079+
"workers_values",
1080+
[
1081+
{
1082+
"containerLifecycleHooks": {
1083+
"preStop": {"exec": {"command": ["echo", "preStop", "{{ .Release.Name }}"]}}
1084+
}
1085+
},
1086+
{
1087+
"kubernetes": {
1088+
"containerLifecycleHooks": {
1089+
"preStop": {"exec": {"command": ["echo", "preStop", "{{ .Release.Name }}"]}}
1090+
}
1091+
}
1092+
},
1093+
{
1094+
"containerLifecycleHooks": {
1095+
"postStart": {"exec": {"command": ["echo", "postStart", "{{ .Release.Name }}"]}}
1096+
},
1097+
"kubernetes": {
10831098
"containerLifecycleHooks": {
10841099
"preStop": {"exec": {"command": ["echo", "preStop", "{{ .Release.Name }}"]}}
10851100
}
10861101
},
10871102
},
1103+
],
1104+
)
1105+
def test_workers_container_lifecycle_webhooks_are_configurable(self, workers_values):
1106+
docs = render_chart(
1107+
name="test-release",
1108+
values={"workers": workers_values},
10881109
show_only=["templates/pod-template-file.yaml"],
10891110
chart_dir=self.temp_chart_dir,
10901111
)
10911112

1092-
assert jmespath.search("spec.containers[0].lifecycle.preStop", docs[0]) == {
1093-
"exec": {"command": ["echo", "preStop", "test-release"]}
1113+
assert jmespath.search("spec.containers[0].lifecycle", docs[0]) == {
1114+
"preStop": {"exec": {"command": ["echo", "preStop", "test-release"]}}
10941115
}
10951116

10961117
def test_termination_grace_period_seconds(self):
@@ -1385,7 +1406,7 @@ def test_kerberos_init_container_resources(self, workers_values):
13851406
{
13861407
"kerberosInitContainer": {
13871408
"enabled": True,
1388-
"securityContexts": {"container": {"runAsUser": 1000}},
1409+
"securityContexts": {"container": {"allowPrivilegeEscalation": False}},
13891410
},
13901411
"kubernetes": {
13911412
"kerberosInitContainer": {
@@ -1398,9 +1419,7 @@ def test_kerberos_init_container_resources(self, workers_values):
13981419
)
13991420
def test_kerberos_init_container_security_context(self, workers_values):
14001421
docs = render_chart(
1401-
values={
1402-
"workers": workers_values,
1403-
},
1422+
values={"workers": workers_values},
14041423
show_only=["templates/pod-template-file.yaml"],
14051424
chart_dir=self.temp_chart_dir,
14061425
)
@@ -1410,59 +1429,52 @@ def test_kerberos_init_container_security_context(self, workers_values):
14101429
) == {"runAsUser": 2000}
14111430

14121431
@pytest.mark.parametrize(
1413-
("workers_values", "expected"),
1432+
"workers_values",
14141433
[
1415-
(
1416-
{
1434+
{
1435+
"kerberosInitContainer": {
1436+
"enabled": True,
1437+
"containerLifecycleHooks": {
1438+
"postStart": {"exec": {"command": ["echo", "{{ .Release.Name }}"]}}
1439+
},
1440+
}
1441+
},
1442+
{
1443+
"kubernetes": {
14171444
"kerberosInitContainer": {
14181445
"enabled": True,
1419-
"containerLifecycleHooks": {"postStart": {"exec": {"command": ["echo", "base"]}}},
1420-
}
1421-
},
1422-
{"postStart": {"exec": {"command": ["echo", "base"]}}},
1423-
),
1424-
(
1425-
{
1426-
"kubernetes": {
1427-
"kerberosInitContainer": {
1428-
"enabled": True,
1429-
"containerLifecycleHooks": {
1430-
"postStart": {"exec": {"command": ["echo", "kubernetes"]}}
1431-
},
1432-
}
1446+
"containerLifecycleHooks": {
1447+
"postStart": {"exec": {"command": ["echo", "{{ .Release.Name }}"]}}
1448+
},
14331449
}
1450+
}
1451+
},
1452+
{
1453+
"kerberosInitContainer": {
1454+
"enabled": True,
1455+
"containerLifecycleHooks": {"preStop": {"exec": {"command": ["echo", "base"]}}},
14341456
},
1435-
{"postStart": {"exec": {"command": ["echo", "kubernetes"]}}},
1436-
),
1437-
(
1438-
{
1457+
"kubernetes": {
14391458
"kerberosInitContainer": {
14401459
"enabled": True,
1441-
"containerLifecycleHooks": {"preStop": {"exec": {"command": ["echo", "base"]}}},
1442-
},
1443-
"kubernetes": {
1444-
"kerberosInitContainer": {
1445-
"enabled": True,
1446-
"containerLifecycleHooks": {
1447-
"postStart": {"exec": {"command": ["echo", "kubernetes"]}}
1448-
},
1449-
}
1450-
},
1460+
"containerLifecycleHooks": {
1461+
"postStart": {"exec": {"command": ["echo", "{{ .Release.Name }}"]}}
1462+
},
1463+
}
14511464
},
1452-
{"postStart": {"exec": {"command": ["echo", "kubernetes"]}}},
1453-
),
1465+
},
14541466
],
14551467
)
1456-
def test_kerberos_init_container_lifecycle_hooks(self, workers_values, expected):
1468+
def test_kerberos_init_container_lifecycle_hooks(self, workers_values):
14571469
docs = render_chart(
1470+
name="test-release",
14581471
values={
14591472
"workers": workers_values,
14601473
},
14611474
show_only=["templates/pod-template-file.yaml"],
14621475
chart_dir=self.temp_chart_dir,
14631476
)
14641477

1465-
assert (
1466-
jmespath.search("spec.initContainers[?name=='kerberos-init'] | [0].lifecycle", docs[0])
1467-
== expected
1468-
)
1478+
assert jmespath.search("spec.initContainers[?name=='kerberos-init'] | [0].lifecycle", docs[0]) == {
1479+
"postStart": {"exec": {"command": ["echo", "test-release"]}}
1480+
}

0 commit comments

Comments
 (0)