From 1bf847914e5a97b817059214ffe54bbae7ad2051 Mon Sep 17 00:00:00 2001 From: mewa Date: Mon, 3 Jun 2024 10:28:15 +0200 Subject: [PATCH] Enable templating in extraContainers and extraInitContainers (#38507) * Add test cases for templating extraContainers and extraInitContainers * Use templating for extra airflow containers --- .../pod-template-file.kubernetes-helm-yaml | 4 +-- .../dag-processor-deployment.yaml | 2 +- .../scheduler/scheduler-deployment.yaml | 2 +- .../triggerer/triggerer-deployment.yaml | 2 +- .../webserver/webserver-deployment.yaml | 2 +- .../templates/workers/worker-deployment.yaml | 2 +- .../airflow_aux/test_create_user_job.py | 14 +++++++++ .../airflow_aux/test_migrate_database_job.py | 14 +++++++++ .../airflow_aux/test_pod_template_file.py | 30 +++++++++++++++++++ helm_tests/airflow_core/test_dag_processor.py | 30 +++++++++++++++++++ helm_tests/airflow_core/test_scheduler.py | 29 ++++++++++++++++++ helm_tests/airflow_core/test_triggerer.py | 28 +++++++++++++++++ helm_tests/airflow_core/test_worker.py | 29 ++++++++++++++++++ helm_tests/webserver/test_webserver.py | 29 ++++++++++++++++++ 14 files changed, 210 insertions(+), 7 deletions(-) diff --git a/chart/files/pod-template-file.kubernetes-helm-yaml b/chart/files/pod-template-file.kubernetes-helm-yaml index bdfd9d2b186bf..421922561608c 100644 --- a/chart/files/pod-template-file.kubernetes-helm-yaml +++ b/chart/files/pod-template-file.kubernetes-helm-yaml @@ -50,7 +50,7 @@ spec: {{- include "git_sync_container" (dict "Values" .Values "is_init" "true" "Template" .Template) | nindent 4 }} {{- end }} {{- if .Values.workers.extraInitContainers }} - {{- toYaml .Values.workers.extraInitContainers | nindent 4 }} + {{- tpl (toYaml .Values.workers.extraInitContainers) . | nindent 4 }} {{- end }} {{- if and (semverCompare ">=2.8.0" .Values.airflowVersion) .Values.workers.kerberosInitContainer.enabled }} - name: kerberos-init @@ -180,7 +180,7 @@ spec: {{- include "standard_airflow_environment" . | indent 6 }} {{- end }} {{- if .Values.workers.extraContainers }} - {{- toYaml .Values.workers.extraContainers | nindent 4 }} + {{- tpl (toYaml .Values.workers.extraContainers) . | nindent 4 }} {{- end }} {{- if .Values.workers.priorityClassName }} priorityClassName: {{ .Values.workers.priorityClassName }} diff --git a/chart/templates/dag-processor/dag-processor-deployment.yaml b/chart/templates/dag-processor/dag-processor-deployment.yaml index ed9b92a213868..90162d5cc37de 100644 --- a/chart/templates/dag-processor/dag-processor-deployment.yaml +++ b/chart/templates/dag-processor/dag-processor-deployment.yaml @@ -144,7 +144,7 @@ spec: {{- include "git_sync_container" (dict "Values" .Values "is_init" "true" "Template" .Template) | nindent 8 }} {{- end }} {{- if .Values.dagProcessor.extraInitContainers }} - {{- toYaml .Values.dagProcessor.extraInitContainers | nindent 8 }} + {{- tpl (toYaml .Values.dagProcessor.extraInitContainers) . | nindent 8 }} {{- end }} containers: - name: dag-processor diff --git a/chart/templates/scheduler/scheduler-deployment.yaml b/chart/templates/scheduler/scheduler-deployment.yaml index 92bd327587a12..d21f60202bdc5 100644 --- a/chart/templates/scheduler/scheduler-deployment.yaml +++ b/chart/templates/scheduler/scheduler-deployment.yaml @@ -169,7 +169,7 @@ spec: {{- include "git_sync_container" (dict "Values" .Values "is_init" "true" "Template" .Template) | nindent 8 }} {{- end }} {{- if .Values.scheduler.extraInitContainers }} - {{- toYaml .Values.scheduler.extraInitContainers | nindent 8 }} + {{- tpl (toYaml .Values.scheduler.extraInitContainers) . | nindent 8 }} {{- end }} containers: # Always run the main scheduler container. diff --git a/chart/templates/triggerer/triggerer-deployment.yaml b/chart/templates/triggerer/triggerer-deployment.yaml index a80637e843773..828fd0e516d9e 100644 --- a/chart/templates/triggerer/triggerer-deployment.yaml +++ b/chart/templates/triggerer/triggerer-deployment.yaml @@ -160,7 +160,7 @@ spec: {{- include "git_sync_container" (dict "Values" .Values "is_init" "true" "Template" .Template) | nindent 8 }} {{- end }} {{- if .Values.triggerer.extraInitContainers }} - {{- toYaml .Values.triggerer.extraInitContainers | nindent 8 }} + {{- tpl (toYaml .Values.triggerer.extraInitContainers) . | nindent 8 }} {{- end }} containers: - name: triggerer diff --git a/chart/templates/webserver/webserver-deployment.yaml b/chart/templates/webserver/webserver-deployment.yaml index f44ce6d7b26b8..cd40a7ba47ad9 100644 --- a/chart/templates/webserver/webserver-deployment.yaml +++ b/chart/templates/webserver/webserver-deployment.yaml @@ -163,7 +163,7 @@ spec: {{- include "git_sync_container" (dict "Values" .Values "is_init" "true" "Template" .Template) | nindent 8 }} {{- end }} {{- if .Values.webserver.extraInitContainers }} - {{- toYaml .Values.webserver.extraInitContainers | nindent 8 }} + {{- tpl (toYaml .Values.webserver.extraInitContainers) . | nindent 8 }} {{- end }} containers: - name: webserver diff --git a/chart/templates/workers/worker-deployment.yaml b/chart/templates/workers/worker-deployment.yaml index 970a05c52245a..0a2f7c6720ef6 100644 --- a/chart/templates/workers/worker-deployment.yaml +++ b/chart/templates/workers/worker-deployment.yaml @@ -220,7 +220,7 @@ spec: {{- include "git_sync_container" (dict "Values" .Values "is_init" "true" "Template" .Template) | nindent 8 }} {{- end }} {{- if .Values.workers.extraInitContainers }} - {{- toYaml .Values.workers.extraInitContainers | nindent 8 }} + {{- tpl (toYaml .Values.workers.extraInitContainers) . | nindent 8 }} {{- end }} containers: - name: worker diff --git a/helm_tests/airflow_aux/test_create_user_job.py b/helm_tests/airflow_aux/test_create_user_job.py index 73b6b9d92b723..6ee458bcbbce2 100644 --- a/helm_tests/airflow_aux/test_create_user_job.py +++ b/helm_tests/airflow_aux/test_create_user_job.py @@ -166,6 +166,20 @@ def test_should_add_extra_containers(self): "image": "test-registry/test-repo:test-tag", } == jmespath.search("spec.template.spec.containers[-1]", docs[0]) + def test_should_template_extra_containers(self): + docs = render_chart( + values={ + "createUserJob": { + "extraContainers": [{"name": "{{ .Release.Name }}-test-container"}], + }, + }, + show_only=["templates/jobs/create-user-job.yaml"], + ) + + assert {"name": "release-name-test-container"} == jmespath.search( + "spec.template.spec.containers[-1]", docs[0] + ) + def test_should_add_extra_volumes(self): docs = render_chart( values={ diff --git a/helm_tests/airflow_aux/test_migrate_database_job.py b/helm_tests/airflow_aux/test_migrate_database_job.py index fd26b61193027..89fc0ababfcba 100644 --- a/helm_tests/airflow_aux/test_migrate_database_job.py +++ b/helm_tests/airflow_aux/test_migrate_database_job.py @@ -190,6 +190,20 @@ def test_should_add_extra_containers(self): "image": "test-registry/test-repo:test-tag", } == jmespath.search("spec.template.spec.containers[-1]", docs[0]) + def test_should_template_extra_containers(self): + docs = render_chart( + values={ + "migrateDatabaseJob": { + "extraContainers": [{"name": "{{ .Release.Name }}-test-container"}], + }, + }, + show_only=["templates/jobs/migrate-database-job.yaml"], + ) + + assert {"name": "release-name-test-container"} == jmespath.search( + "spec.template.spec.containers[-1]", docs[0] + ) + def test_set_resources(self): docs = render_chart( values={ diff --git a/helm_tests/airflow_aux/test_pod_template_file.py b/helm_tests/airflow_aux/test_pod_template_file.py index 41a4a2a427510..821634751219f 100644 --- a/helm_tests/airflow_aux/test_pod_template_file.py +++ b/helm_tests/airflow_aux/test_pod_template_file.py @@ -702,6 +702,21 @@ def test_should_add_extra_init_containers(self): "image": "test-registry/test-repo:test-tag", } == jmespath.search("spec.initContainers[-1]", docs[0]) + def test_should_template_extra_init_containers(self): + docs = render_chart( + values={ + "workers": { + "extraInitContainers": [{"name": "{{ .Release.Name }}-test-init-container"}], + }, + }, + show_only=["templates/pod-template-file.yaml"], + chart_dir=self.temp_chart_dir, + ) + + assert { + "name": "release-name-test-init-container", + } == jmespath.search("spec.initContainers[-1]", docs[0]) + def test_should_add_extra_containers(self): docs = render_chart( values={ @@ -720,6 +735,21 @@ def test_should_add_extra_containers(self): "image": "test-registry/test-repo:test-tag", } == jmespath.search("spec.containers[-1]", docs[0]) + def test_should_template_extra_containers(self): + docs = render_chart( + values={ + "workers": { + "extraContainers": [{"name": "{{ .Release.Name }}-test-container"}], + }, + }, + show_only=["templates/pod-template-file.yaml"], + chart_dir=self.temp_chart_dir, + ) + + assert { + "name": "release-name-test-container", + } == jmespath.search("spec.containers[-1]", docs[0]) + def test_should_add_pod_labels(self): docs = render_chart( values={"labels": {"label1": "value1", "label2": "value2"}}, diff --git a/helm_tests/airflow_core/test_dag_processor.py b/helm_tests/airflow_core/test_dag_processor.py index 17d29ca1b3ede..31256b4095cf9 100644 --- a/helm_tests/airflow_core/test_dag_processor.py +++ b/helm_tests/airflow_core/test_dag_processor.py @@ -110,6 +110,21 @@ def test_should_add_extra_containers(self): "image": "test-registry/test-repo:test-tag", } == jmespath.search("spec.template.spec.containers[-1]", docs[0]) + def test_should_template_extra_containers(self): + docs = render_chart( + values={ + "dagProcessor": { + "enabled": True, + "extraContainers": [{"name": "{{ .Release.Name }}-test-container"}], + }, + }, + show_only=["templates/dag-processor/dag-processor-deployment.yaml"], + ) + + assert {"name": "release-name-test-container"} == jmespath.search( + "spec.template.spec.containers[-1]", docs[0] + ) + def test_should_add_extra_init_containers(self): docs = render_chart( values={ @@ -128,6 +143,21 @@ def test_should_add_extra_init_containers(self): "image": "test-registry/test-repo:test-tag", } == jmespath.search("spec.template.spec.initContainers[-1]", docs[0]) + def test_should_template_extra_init_containers(self): + docs = render_chart( + values={ + "dagProcessor": { + "enabled": True, + "extraInitContainers": [{"name": "{{ .Release.Name }}-test-init-container"}], + }, + }, + show_only=["templates/dag-processor/dag-processor-deployment.yaml"], + ) + + assert {"name": "release-name-test-init-container"} == jmespath.search( + "spec.template.spec.initContainers[-1]", docs[0] + ) + def test_should_add_extra_volume_and_extra_volume_mount(self): docs = render_chart( values={ diff --git a/helm_tests/airflow_core/test_scheduler.py b/helm_tests/airflow_core/test_scheduler.py index acf7cbeca935e..eef2f8756426d 100644 --- a/helm_tests/airflow_core/test_scheduler.py +++ b/helm_tests/airflow_core/test_scheduler.py @@ -69,6 +69,21 @@ def test_should_add_extra_containers(self): "image": "test-registry/test-repo:test-tag", } == jmespath.search("spec.template.spec.containers[-1]", docs[0]) + def test_should_template_extra_containers(self): + docs = render_chart( + values={ + "executor": "CeleryExecutor", + "scheduler": { + "extraContainers": [{"name": "{{ .Release.Name }}-test-container"}], + }, + }, + show_only=["templates/scheduler/scheduler-deployment.yaml"], + ) + + assert {"name": "release-name-test-container"} == jmespath.search( + "spec.template.spec.containers[-1]", docs[0] + ) + def test_disable_wait_for_migration(self): docs = render_chart( values={ @@ -100,6 +115,20 @@ def test_should_add_extra_init_containers(self): "image": "test-registry/test-repo:test-tag", } == jmespath.search("spec.template.spec.initContainers[-1]", docs[0]) + def test_should_template_extra_init_containers(self): + docs = render_chart( + values={ + "scheduler": { + "extraInitContainers": [{"name": "{{ .Release.Name }}-test-init-container"}], + }, + }, + show_only=["templates/scheduler/scheduler-deployment.yaml"], + ) + + assert {"name": "release-name-test-init-container"} == jmespath.search( + "spec.template.spec.initContainers[-1]", docs[0] + ) + def test_should_add_extra_volume_and_extra_volume_mount(self): docs = render_chart( values={ diff --git a/helm_tests/airflow_core/test_triggerer.py b/helm_tests/airflow_core/test_triggerer.py index 69a948a7aa614..c92683e80dc06 100644 --- a/helm_tests/airflow_core/test_triggerer.py +++ b/helm_tests/airflow_core/test_triggerer.py @@ -107,6 +107,20 @@ def test_should_add_extra_containers(self): "image": "test-registry/test-repo:test-tag", } == jmespath.search("spec.template.spec.containers[-1]", docs[0]) + def test_should_template_extra_containers(self): + docs = render_chart( + values={ + "triggerer": { + "extraContainers": [{"name": "{{ .Release.Name }}-test-container"}], + }, + }, + show_only=["templates/triggerer/triggerer-deployment.yaml"], + ) + + assert {"name": "release-name-test-container"} == jmespath.search( + "spec.template.spec.containers[-1]", docs[0] + ) + def test_should_add_extra_init_containers(self): docs = render_chart( values={ @@ -124,6 +138,20 @@ def test_should_add_extra_init_containers(self): "image": "test-registry/test-repo:test-tag", } == jmespath.search("spec.template.spec.initContainers[-1]", docs[0]) + def test_should_template_extra_init_containers(self): + docs = render_chart( + values={ + "triggerer": { + "extraInitContainers": [{"name": "{{ .Release.Name }}-test-init-container"}], + }, + }, + show_only=["templates/triggerer/triggerer-deployment.yaml"], + ) + + assert {"name": "release-name-test-init-container"} == jmespath.search( + "spec.template.spec.initContainers[-1]", docs[0] + ) + def test_should_add_extra_volume_and_extra_volume_mount(self): docs = render_chart( values={ diff --git a/helm_tests/airflow_core/test_worker.py b/helm_tests/airflow_core/test_worker.py index 76891d405882e..100d53c28c4a1 100644 --- a/helm_tests/airflow_core/test_worker.py +++ b/helm_tests/airflow_core/test_worker.py @@ -82,6 +82,21 @@ def test_should_add_extra_containers(self): "image": "test-registry/test-repo:test-tag", } == jmespath.search("spec.template.spec.containers[-1]", docs[0]) + def test_should_template_extra_containers(self): + docs = render_chart( + values={ + "executor": "CeleryExecutor", + "workers": { + "extraContainers": [{"name": "{{ .Release.Name }}-test-container"}], + }, + }, + show_only=["templates/workers/worker-deployment.yaml"], + ) + + assert {"name": "release-name-test-container"} == jmespath.search( + "spec.template.spec.containers[-1]", docs[0] + ) + def test_disable_wait_for_migration(self): docs = render_chart( values={ @@ -113,6 +128,20 @@ def test_should_add_extra_init_containers(self): "image": "test-registry/test-repo:test-tag", } == jmespath.search("spec.template.spec.initContainers[-1]", docs[0]) + def test_should_template_extra_init_containers(self): + docs = render_chart( + values={ + "workers": { + "extraInitContainers": [{"name": "{{ .Release.Name }}-test-init-container"}], + }, + }, + show_only=["templates/workers/worker-deployment.yaml"], + ) + + assert {"name": "release-name-test-init-container"} == jmespath.search( + "spec.template.spec.initContainers[-1]", docs[0] + ) + def test_should_add_extra_volume_and_extra_volume_mount(self): docs = render_chart( values={ diff --git a/helm_tests/webserver/test_webserver.py b/helm_tests/webserver/test_webserver.py index 3b597af215428..364a7e4b1607f 100644 --- a/helm_tests/webserver/test_webserver.py +++ b/helm_tests/webserver/test_webserver.py @@ -200,6 +200,21 @@ def test_should_add_extra_containers(self): "image": "test-registry/test-repo:test-tag", } == jmespath.search("spec.template.spec.containers[-1]", docs[0]) + def test_should_template_extra_containers(self): + docs = render_chart( + values={ + "executor": "CeleryExecutor", + "webserver": { + "extraContainers": [{"name": "{{ .Release.Name }}-test-container"}], + }, + }, + show_only=["templates/webserver/webserver-deployment.yaml"], + ) + + assert { + "name": "release-name-test-container", + } == jmespath.search("spec.template.spec.containers[-1]", docs[0]) + def test_should_add_extraEnvs(self): docs = render_chart( values={ @@ -315,6 +330,20 @@ def test_should_add_extra_init_containers(self): "image": "test-registry/test-repo:test-tag", } == jmespath.search("spec.template.spec.initContainers[-1]", docs[0]) + def test_should_template_extra_init_containers(self): + docs = render_chart( + values={ + "webserver": { + "extraInitContainers": [{"name": "{{ .Release.Name }}-init-container"}], + }, + }, + show_only=["templates/webserver/webserver-deployment.yaml"], + ) + + assert { + "name": "release-name-init-container", + } == jmespath.search("spec.template.spec.initContainers[-1]", docs[0]) + def test_should_add_component_specific_labels(self): docs = render_chart( values={