diff --git a/chart/files/pod-template-file.kubernetes-helm-yaml b/chart/files/pod-template-file.kubernetes-helm-yaml index afbf33f4fe4aa..fdd168bcb32ce 100644 --- a/chart/files/pod-template-file.kubernetes-helm-yaml +++ b/chart/files/pod-template-file.kubernetes-helm-yaml @@ -122,8 +122,8 @@ spec: lifecycle: {{- tpl (toYaml $containerLifecycleHooks) . | nindent 8 }} {{- end }} name: base - {{- if .Values.workers.command }} - command: {{ tpl (toYaml .Values.workers.command) . | nindent 8 }} + {{- if or .Values.workers.kubernetes.command .Values.workers.command }} + command: {{ tpl (toYaml (.Values.workers.kubernetes.command | default .Values.workers.command)) . | nindent 8 }} {{- end }} resources: {{- toYaml .Values.workers.resources | nindent 8 }} volumeMounts: diff --git a/chart/templates/workers/worker-deployment.yaml b/chart/templates/workers/worker-deployment.yaml index 5af194eb790e0..b438e8df23549 100644 --- a/chart/templates/workers/worker-deployment.yaml +++ b/chart/templates/workers/worker-deployment.yaml @@ -249,8 +249,8 @@ spec: {{- if $containerLifecycleHooks }} lifecycle: {{- tpl (toYaml $containerLifecycleHooks) . | nindent 12 }} {{- end }} - {{- if .Values.workers.command }} - command: {{ tpl (toYaml .Values.workers.command) . | nindent 12 }} + {{- if or .Values.workers.celery.command .Values.workers.command }} + command: {{ tpl (toYaml (.Values.workers.celery.command | default .Values.workers.command)) . | nindent 12 }} {{- end }} {{- if .Values.workers.args }} args: {{ tpl (toYaml .Values.workers.args) . | nindent 12 }} diff --git a/chart/values.schema.json b/chart/values.schema.json index f8c53da11cce9..452810075cb53 100644 --- a/chart/values.schema.json +++ b/chart/values.schema.json @@ -1649,7 +1649,7 @@ "x-docsSection": null }, "command": { - "description": "Command to use when running Airflow Celery workers and using pod-template-file (templated).", + "description": "Command to use when running Airflow Celery workers and using pod-template-file (templated). Use workers.celery.command and/or workers.kubernetes.command to separate value between Celery workers and pod-template-file.", "type": [ "array", "null" @@ -2622,6 +2622,17 @@ "default": null, "x-docsSection": null }, + "command": { + "description": "Command to use when running Airflow Celery workers (templated).", + "type": [ + "array", + "null" + ], + "items": { + "type": "string" + }, + "default": null + }, "serviceAccount": { "description": "Create ServiceAccount.", "type": "object", @@ -2661,6 +2672,17 @@ "type": "object", "x-docsSection": "Workers", "properties": { + "command": { + "description": "Command to use in pod-template-file (templated).", + "type": [ + "array", + "null" + ], + "items": { + "type": "string" + }, + "default": null + }, "serviceAccount": { "description": "Create ServiceAccount.", "type": "object", diff --git a/chart/values.yaml b/chart/values.yaml index a748e5fc1ba0b..66746b47ee2c7 100644 --- a/chart/values.yaml +++ b/chart/values.yaml @@ -645,7 +645,10 @@ workers: revisionHistoryLimit: ~ # Command to use when running Airflow Celery workers and using pod-template-file (templated) + # Use workers.celery.command and/or workers.kubernetes.command to separate value between + # Celery workers and pod-template-file command: ~ + # Args to use when running Airflow Celery workers (templated) args: - "bash" @@ -1018,6 +1021,9 @@ workers: # Max number of old Airflow Celery workers ReplicaSets to retain revisionHistoryLimit: ~ + # Command to use when running Airflow Celery workers (templated) + command: ~ + # Create ServiceAccount for Airflow Celery workers serviceAccount: # default value is true @@ -1032,6 +1038,9 @@ workers: annotations: {} kubernetes: + # Command to use in pod-template-file (templated) + command: ~ + # Create ServiceAccount for pods created with pod-template-file serviceAccount: # Auto mount service account token into the pod. Default value is true. diff --git a/helm-tests/tests/helm_tests/airflow_aux/test_pod_template_file.py b/helm-tests/tests/helm_tests/airflow_aux/test_pod_template_file.py index ae74e995e8a67..926face3f4f47 100644 --- a/helm-tests/tests/helm_tests/airflow_aux/test_pod_template_file.py +++ b/helm-tests/tests/helm_tests/airflow_aux/test_pod_template_file.py @@ -1035,27 +1035,43 @@ def test_airflow_kerberos_init_container( assert initContainers[0]["args"] == ["kerberos", "-o"] @pytest.mark.parametrize( - ("cmd", "expected"), + ("workers_values", "expected"), [ - (["test", "command", "to", "run"], ["test", "command", "to", "run"]), - (["cmd", "{{ .Release.Name }}"], ["cmd", "release-name"]), + ({"command": ["test", "command", "to", "run"]}, ["test", "command", "to", "run"]), + ({"command": ["cmd", "{{ .Release.Name }}"]}, ["cmd", "release-name"]), + ({"kubernetes": {"command": ["test", "command", "to", "run"]}}, ["test", "command", "to", "run"]), + ({"kubernetes": {"command": ["cmd", "{{ .Release.Name }}"]}}, ["cmd", "release-name"]), + ( + {"command": ["test"], "kubernetes": {"command": ["test", "command", "to", "run"]}}, + ["test", "command", "to", "run"], + ), + ( + {"command": ["test"], "kubernetes": {"command": ["cmd", "{{ .Release.Name }}"]}}, + ["cmd", "release-name"], + ), ], ) - def test_should_add_command(self, cmd, expected): + def test_should_add_command(self, workers_values, expected): docs = render_chart( - values={ - "workers": {"command": cmd}, - }, + values={"workers": workers_values}, show_only=["templates/pod-template-file.yaml"], chart_dir=self.temp_chart_dir, ) assert expected == jmespath.search("spec.containers[0].command", docs[0]) - @pytest.mark.parametrize("cmd", [None, []]) - def test_should_not_add_command(self, cmd): + @pytest.mark.parametrize( + "workers_values", + [ + {"command": None}, + {"command": []}, + {"kubernetes": {"command": None}}, + {"kubernetes": {"command": []}}, + ], + ) + def test_should_not_add_command(self, workers_values): docs = render_chart( - values={"workers": {"command": cmd}}, + values={"workers": workers_values}, show_only=["templates/pod-template-file.yaml"], chart_dir=self.temp_chart_dir, ) diff --git a/helm-tests/tests/helm_tests/airflow_core/test_worker.py b/helm-tests/tests/helm_tests/airflow_core/test_worker.py index 970b43eb32614..ed6ebe7fcdbc6 100644 --- a/helm-tests/tests/helm_tests/airflow_core/test_worker.py +++ b/helm-tests/tests/helm_tests/airflow_core/test_worker.py @@ -809,25 +809,68 @@ def test_default_command_and_args_airflow_version(self, airflow_version, expecte f"exec \\\n{expected_arg}", ] == jmespath.search("spec.template.spec.containers[0].args", docs[0]) - @pytest.mark.parametrize("command", [None, ["custom", "command"]]) - @pytest.mark.parametrize("args", [None, ["custom", "args"]]) - def test_command_and_args_overrides(self, command, args): + @pytest.mark.parametrize( + ("workers_values", "expected"), + [ + ({"command": ["custom", "command"]}, ["custom", "command"]), + ({"command": ["custom", "{{ .Release.Name }}"]}, ["custom", "release-name"]), + ({"celery": {"command": ["custom", "command"]}}, ["custom", "command"]), + ({"celery": {"command": ["custom", "{{ .Release.Name }}"]}}, ["custom", "release-name"]), + ({"command": ["test"], "celery": {"command": ["custom", "command"]}}, ["custom", "command"]), + ( + {"command": ["test"], "celery": {"command": ["custom", "{{ .Release.Name }}"]}}, + ["custom", "release-name"], + ), + ], + ) + def test_should_add_command(self, workers_values, expected): + docs = render_chart( + values={"workers": workers_values}, + show_only=["templates/workers/worker-deployment.yaml"], + ) + + assert expected == jmespath.search("spec.template.spec.containers[0].command", docs[0]) + + @pytest.mark.parametrize( + "workers_values", + [ + {"command": None}, + {"command": []}, + {"celery": {"command": None}}, + {"celery": {"command": []}}, + ], + ) + def test_should_not_add_command(self, workers_values): + docs = render_chart( + values={"workers": workers_values}, + show_only=["templates/workers/worker-deployment.yaml"], + ) + + assert jmespath.search("spec.template.spec.containers[0].command", docs[0]) is None + + @pytest.mark.parametrize( + ("args", "expected"), + [ + (["custom", "args"], ["custom", "args"]), + (["custom", "{{ .Release.Service }}"], ["custom", "Helm"]), + ], + ) + def test_should_add_args(self, args, expected): docs = render_chart( - values={"workers": {"command": command, "args": args}}, + values={"workers": {"args": args}}, show_only=["templates/workers/worker-deployment.yaml"], ) - assert command == jmespath.search("spec.template.spec.containers[0].command", docs[0]) - assert args == jmespath.search("spec.template.spec.containers[0].args", docs[0]) + assert expected == jmespath.search("spec.template.spec.containers[0].args", docs[0]) - def test_command_and_args_overrides_are_templated(self): + @pytest.mark.parametrize("args", [None, []]) + def test_should_not_add_args(self, args): docs = render_chart( - values={"workers": {"command": ["{{ .Release.Name }}"], "args": ["{{ .Release.Service }}"]}}, + values={"workers": {"args": args}}, show_only=["templates/workers/worker-deployment.yaml"], ) - assert jmespath.search("spec.template.spec.containers[0].command", docs[0]) == ["release-name"] - assert jmespath.search("spec.template.spec.containers[0].args", docs[0]) == ["Helm"] + assert jmespath.search("spec.template.spec.containers[0].args", docs[0]) is None def test_dags_gitsync_sidecar_and_init_container(self): docs = render_chart(