diff --git a/chart/files/pod-template-file.kubernetes-helm-yaml b/chart/files/pod-template-file.kubernetes-helm-yaml index e62d0ea8bb7d5..aa42c1b2542c3 100644 --- a/chart/files/pod-template-file.kubernetes-helm-yaml +++ b/chart/files/pod-template-file.kubernetes-helm-yaml @@ -211,7 +211,11 @@ spec: terminationGracePeriodSeconds: {{ .Values.workers.terminationGracePeriodSeconds }} tolerations: {{- toYaml $tolerations | nindent 4 }} topologySpreadConstraints: {{- toYaml $topologySpreadConstraints | nindent 4 }} + {{- if .Values.workers.useWorkerDedicatedServiceAccounts }} + serviceAccountName: {{ include "worker.kubernetes.serviceAccountName" . }} + {{- else }} serviceAccountName: {{ include "worker.serviceAccountName" . }} + {{- end }} volumes: {{- if .Values.dags.persistence.enabled }} - name: dags diff --git a/chart/templates/_helpers.yaml b/chart/templates/_helpers.yaml index 8838b837dca0f..58eebf7226785 100644 --- a/chart/templates/_helpers.yaml +++ b/chart/templates/_helpers.yaml @@ -636,13 +636,23 @@ server_tls_key_file = /etc/pgbouncer/server.key {{- end }} {{- end }} -{{/* Helper to generate service account name respecting .Values.$section.serviceAccount flags */}} +{{/* Helper for service account name generation */}} +{{- define "_serviceAccountNameGen" -}} + {{- if .sa.create }} + {{- default (printf "%s-%s" (include "airflow.serviceAccountName" .) (default .key .nameSuffix )) .sa.name | quote }} + {{- else }} + {{- default "default" .sa.name | quote }} + {{- end }} +{{- end }} + +{{/* Helper to generate service account name respecting .Values.$section.serviceAccount or .Values.$section.$subSection.serviceAccount flags */}} {{- define "_serviceAccountName" -}} - {{- $sa := get (get .Values .key) "serviceAccount" }} - {{- if $sa.create }} - {{- default (printf "%s-%s" (include "airflow.serviceAccountName" .) (default .key .nameSuffix )) $sa.name | quote }} + {{- if .subKey }} + {{- $sa := get (get (get .Values .key) .subKey) "serviceAccount" -}} + {{- include "_serviceAccountNameGen" (merge (dict "sa" $sa "key" .key "nameSuffix" .nameSuffix) .) }} {{- else }} - {{- default "default" $sa.name | quote }} + {{- $sa := get (get .Values .key) "serviceAccount" }} + {{- include "_serviceAccountNameGen" (merge (dict "sa" $sa "key" .key "nameSuffix" .nameSuffix) .) }} {{- end }} {{- end }} @@ -692,6 +702,16 @@ server_tls_key_file = /etc/pgbouncer/server.key {{- include "_serviceAccountName" (merge (dict "key" "workers" "nameSuffix" "worker") .) -}} {{- end }} +{{/* Create the name of the worker celery service account to use */}} +{{- define "worker.celery.serviceAccountName" -}} + {{- include "_serviceAccountName" (merge (dict "key" "workers" "subKey" "celery" "nameSuffix" "worker-celery") .) -}} +{{- end }} + +{{/* Create the name of the worker kubernetes service account to use */}} +{{- define "worker.kubernetes.serviceAccountName" -}} + {{- include "_serviceAccountName" (merge (dict "key" "workers" "subKey" "kubernetes" "nameSuffix" "worker-kubernetes") .) -}} +{{- end }} + {{/* Create the name of the triggerer service account to use */}} {{- define "triggerer.serviceAccountName" -}} {{- include "_serviceAccountName" (merge (dict "key" "triggerer") .) -}} diff --git a/chart/templates/rbac/pod-launcher-rolebinding.yaml b/chart/templates/rbac/pod-launcher-rolebinding.yaml index adbee311d730d..b70cc03b413fb 100644 --- a/chart/templates/rbac/pod-launcher-rolebinding.yaml +++ b/chart/templates/rbac/pod-launcher-rolebinding.yaml @@ -68,12 +68,22 @@ subjects: {{- end }} {{- end }} {{- $workerAdded := false }} + {{- $workersDedicatedSA := .Values.workers.useWorkerDedicatedServiceAccounts -}} {{- range $executor := $executors }} {{- if and (has $executor $workerLaunchExecutors) (not $workerAdded) }} {{- $workerAdded = true }} + {{- if $workersDedicatedSA }} + - kind: ServiceAccount + name: {{ include "worker.celery.serviceAccountName" $ }} + namespace: "{{ $.Release.Namespace }}" + - kind: ServiceAccount + name: {{ include "worker.kubernetes.serviceAccountName" $ }} + namespace: "{{ $.Release.Namespace }}" + {{- else }} - kind: ServiceAccount name: {{ include "worker.serviceAccountName" $ }} namespace: "{{ $.Release.Namespace }}" {{- end }} {{- end }} + {{- end }} {{- end }} diff --git a/chart/templates/rbac/security-context-constraint-rolebinding.yaml b/chart/templates/rbac/security-context-constraint-rolebinding.yaml index aa4cf05a73e27..3a070a38b479c 100644 --- a/chart/templates/rbac/security-context-constraint-rolebinding.yaml +++ b/chart/templates/rbac/security-context-constraint-rolebinding.yaml @@ -55,10 +55,19 @@ subjects: name: {{ include "webserver.serviceAccountName" . }} namespace: "{{ .Release.Namespace }}" {{- if $hasWorkers }} + {{- if .Values.workers.useWorkerDedicatedServiceAccounts }} + - kind: ServiceAccount + name: {{ include "worker.celery.serviceAccountName" . }} + namespace: "{{ .Release.Namespace }}" + - kind: ServiceAccount + name: {{ include "worker.kubernetes.serviceAccountName" . }} + namespace: "{{ .Release.Namespace }}" + {{- else }} - kind: ServiceAccount name: {{ include "worker.serviceAccountName" . }} namespace: "{{ .Release.Namespace }}" {{- end }} + {{- end }} - kind: ServiceAccount name: {{ include "scheduler.serviceAccountName" . }} namespace: "{{ .Release.Namespace }}" diff --git a/chart/templates/workers/worker-celery-serviceaccount.yaml b/chart/templates/workers/worker-celery-serviceaccount.yaml new file mode 100644 index 0000000000000..9341e9fd861c4 --- /dev/null +++ b/chart/templates/workers/worker-celery-serviceaccount.yaml @@ -0,0 +1,41 @@ +{{/* + Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. +*/}} + +####################################### +## Airflow Worker Celery ServiceAccount +####################################### +{{- if and .Values.workers.celery.serviceAccount.create .Values.workers.useWorkerDedicatedServiceAccounts (or (contains "CeleryExecutor" .Values.executor) (contains "CeleryKubernetesExecutor" .Values.executor)) }} +apiVersion: v1 +kind: ServiceAccount +automountServiceAccountToken: {{ .Values.workers.celery.serviceAccount.automountServiceAccountToken }} +metadata: + name: {{ include "worker.celery.serviceAccountName" . }} + labels: + tier: airflow + component: worker + release: {{ .Release.Name }} + chart: "{{ .Chart.Name }}-{{ .Chart.Version }}" + heritage: {{ .Release.Service }} + {{- if or .Values.labels .Values.workers.labels .Values.workers.celery.labels }} + {{- mustMerge .Values.workers.celery.labels .Values.workers.labels .Values.labels | toYaml | nindent 4 }} + {{- end }} + {{- with .Values.workers.celery.serviceAccount.annotations }} + annotations: {{- toYaml . | nindent 4 }} + {{- end }} +{{- end }} diff --git a/chart/templates/workers/worker-deployment.yaml b/chart/templates/workers/worker-deployment.yaml index 998e059d1747d..9078560988a35 100644 --- a/chart/templates/workers/worker-deployment.yaml +++ b/chart/templates/workers/worker-deployment.yaml @@ -135,7 +135,11 @@ spec: {{- end }} terminationGracePeriodSeconds: {{ .Values.workers.terminationGracePeriodSeconds }} restartPolicy: Always + {{- if .Values.workers.useWorkerDedicatedServiceAccounts }} + serviceAccountName: {{ include "worker.celery.serviceAccountName" . }} + {{- else }} serviceAccountName: {{ include "worker.serviceAccountName" . }} + {{- end }} securityContext: {{ $securityContext | nindent 8 }} {{- if or .Values.registry.secretName .Values.registry.connection }} imagePullSecrets: diff --git a/chart/templates/workers/worker-kubernetes-serviceaccount.yaml b/chart/templates/workers/worker-kubernetes-serviceaccount.yaml new file mode 100644 index 0000000000000..296f1e4506caa --- /dev/null +++ b/chart/templates/workers/worker-kubernetes-serviceaccount.yaml @@ -0,0 +1,41 @@ +{{/* + Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. +*/}} + +########################################### +## Airflow Worker Kubernetes ServiceAccount +########################################### +{{- if and .Values.workers.kubernetes.serviceAccount.create .Values.workers.useWorkerDedicatedServiceAccounts (or (contains "CeleryKubernetesExecutor" .Values.executor) (contains "KubernetesExecutor" .Values.executor) (contains "LocalKubernetesExecutor" .Values.executor)) }} +apiVersion: v1 +kind: ServiceAccount +automountServiceAccountToken: {{ .Values.workers.kubernetes.serviceAccount.automountServiceAccountToken }} +metadata: + name: {{ include "worker.kubernetes.serviceAccountName" . }} + labels: + tier: airflow + component: worker + release: {{ .Release.Name }} + chart: "{{ .Chart.Name }}-{{ .Chart.Version }}" + heritage: {{ .Release.Service }} + {{- if or .Values.labels .Values.workers.labels .Values.workers.kubernetes.labels }} + {{- mustMerge .Values.workers.kubernetes.labels .Values.workers.labels .Values.labels | toYaml | nindent 4 }} + {{- end }} + {{- with .Values.workers.kubernetes.serviceAccount.annotations }} + annotations: {{- toYaml . | nindent 4 }} + {{- end }} +{{- end }} diff --git a/chart/templates/workers/worker-serviceaccount.yaml b/chart/templates/workers/worker-serviceaccount.yaml index 0feec8de3d99f..7fb46342ce1ab 100644 --- a/chart/templates/workers/worker-serviceaccount.yaml +++ b/chart/templates/workers/worker-serviceaccount.yaml @@ -20,7 +20,7 @@ ################################ ## Airflow Worker ServiceAccount ################################# -{{- if and .Values.workers.serviceAccount.create (or (contains "CeleryExecutor" .Values.executor) (contains "CeleryKubernetesExecutor" .Values.executor) (contains "KubernetesExecutor" .Values.executor) (contains "LocalKubernetesExecutor" .Values.executor)) }} +{{- if and .Values.workers.serviceAccount.create (not .Values.workers.useWorkerDedicatedServiceAccounts) (or (contains "CeleryExecutor" .Values.executor) (contains "CeleryKubernetesExecutor" .Values.executor) (contains "KubernetesExecutor" .Values.executor) (contains "LocalKubernetesExecutor" .Values.executor)) }} apiVersion: v1 kind: ServiceAccount automountServiceAccountToken: {{ .Values.workers.serviceAccount.automountServiceAccountToken }} diff --git a/chart/values.schema.json b/chart/values.schema.json index 01654f9042f9f..b8fb47fe557d9 100644 --- a/chart/values.schema.json +++ b/chart/values.schema.json @@ -2530,6 +2530,89 @@ } } ] + }, + "useWorkerDedicatedServiceAccounts": { + "description": "One common Service Account for all workers will be created if flag is set to false. If true, dedicated Service Accounts for every worker type will be created.", + "type": "boolean", + "default": false + }, + "celery": { + "description": "Airflow Celery Workers configuration.", + "type": "object", + "x-docsSection": "Workers", + "properties": { + "serviceAccount": { + "description": "Create ServiceAccount.", + "type": "object", + "properties": { + "automountServiceAccountToken": { + "description": "Specifies if ServiceAccount's API credentials should be mounted onto Pods.", + "type": "boolean", + "default": true + }, + "create": { + "description": "Specifies whether a ServiceAccount should be created.", + "type": "boolean", + "default": true + }, + "name": { + "description": "The name of the ServiceAccount to use. If not set and create is true, a name is generated using the release name.", + "type": [ + "string", + "null" + ], + "default": null + }, + "annotations": { + "description": "Annotations to add to the Airflow Celery worker Kubernetes ServiceAccount.", + "type": "object", + "default": {}, + "additionalProperties": { + "type": "string" + } + } + } + } + } + }, + "kubernetes": { + "description": "Airflow pod-template-file configuration.", + "type": "object", + "x-docsSection": "Workers", + "properties": { + "serviceAccount": { + "description": "Create ServiceAccount.", + "type": "object", + "properties": { + "automountServiceAccountToken": { + "description": "Specifies if ServiceAccount's API credentials should be mounted onto Pods.", + "type": "boolean", + "default": true + }, + "create": { + "description": "Specifies whether a ServiceAccount should be created.", + "type": "boolean", + "default": true + }, + "name": { + "description": "The name of the ServiceAccount to use. If not set and create is true, a name is generated using the release name.", + "type": [ + "string", + "null" + ], + "default": null + }, + "annotations": { + "description": "Annotations to add to the worker Kubernetes ServiceAccount.", + "type": "object", + "default": {}, + "additionalProperties": { + "type": "string" + } + } + } + } + } } } }, diff --git a/chart/values.yaml b/chart/values.yaml index dd57e4778226f..fe09ac49ee57a 100644 --- a/chart/values.yaml +++ b/chart/values.yaml @@ -985,6 +985,38 @@ workers: # requests: # storage: "20Gi" + # One common Service Account for all workers will be created if flag is set to false. + # If true, dedicated Service Accounts for every worker type will be created. + useWorkerDedicatedServiceAccounts: false + + celery: + # Create ServiceAccount for Airflow Celery workers + serviceAccount: + # default value is true + # ref: https://kubernetes.io/docs/tasks/configure-pod-container/configure-service-account/ + automountServiceAccountToken: true + # Specifies whether a ServiceAccount should be created + create: true + # The name of the ServiceAccount to use. + # If not set and create is true, a name is generated using the release name + name: ~ + # Annotations to add to worker kubernetes service account. + annotations: {} + + kubernetes: + # Create ServiceAccount for pods created with pod-template-file + serviceAccount: + # Auto mount service account token into the pod. Default value is true. + # ref: https://kubernetes.io/docs/tasks/configure-pod-container/configure-service-account/ + automountServiceAccountToken: true + # Specifies whether a ServiceAccount should be created + create: true + # The name of the ServiceAccount to use. + # If not set and create is true, a name is generated using the release name. + name: ~ + # Annotations to add to worker kubernetes service account + annotations: {} + # Airflow scheduler settings scheduler: enabled: true diff --git a/helm-tests/tests/helm_tests/airflow_aux/test_annotations.py b/helm-tests/tests/helm_tests/airflow_aux/test_annotations.py index 6c0d4788da36e..bf007e9bfd9bd 100644 --- a/helm-tests/tests/helm_tests/airflow_aux/test_annotations.py +++ b/helm-tests/tests/helm_tests/airflow_aux/test_annotations.py @@ -101,6 +101,90 @@ class TestServiceAccountAnnotations: "example": "worker", }, ), + ( + { + "workers": { + "useWorkerDedicatedServiceAccounts": True, + "celery": { + "serviceAccount": { + "annotations": { + "example": "worker", + }, + }, + }, + }, + }, + "templates/workers/worker-celery-serviceaccount.yaml", + { + "example": "worker", + }, + ), + ( + { + "workers": { + "useWorkerDedicatedServiceAccounts": True, + "serviceAccount": { + "annotations": { + "example": "missing", + }, + }, + "celery": { + "serviceAccount": { + "annotations": { + "example": "worker", + }, + }, + }, + }, + }, + "templates/workers/worker-celery-serviceaccount.yaml", + { + "example": "worker", + }, + ), + ( + { + "executor": "KubernetesExecutor", + "workers": { + "useWorkerDedicatedServiceAccounts": True, + "kubernetes": { + "serviceAccount": { + "annotations": { + "example": "worker", + }, + }, + }, + }, + }, + "templates/workers/worker-kubernetes-serviceaccount.yaml", + { + "example": "worker", + }, + ), + ( + { + "executor": "KubernetesExecutor", + "workers": { + "useWorkerDedicatedServiceAccounts": True, + "serviceAccount": { + "annotations": { + "example": "missing", + }, + }, + "kubernetes": { + "serviceAccount": { + "annotations": { + "example": "worker", + }, + }, + }, + }, + }, + "templates/workers/worker-kubernetes-serviceaccount.yaml", + { + "example": "worker", + }, + ), ( { "flower": { 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 d497f1df71c41..727540fbd51d6 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 @@ -1092,3 +1092,20 @@ def test_base_contains_kerberos_env(self, workers_values): scheduler_env = jmespath.search("spec.containers[0].env[*].name", docs[0]) assert set(["KRB5_CONFIG", "KRB5CCNAME"]).issubset(scheduler_env) + + def test_workers_kubernetes_service_account_custom_names_in_objects(self): + k8s_objects = render_chart( + "test-rbac", + values={ + "workers": { + "useWorkerDedicatedServiceAccounts": True, + "kubernetes": {"serviceAccount": {"name": "TestWorkerKubernetes"}}, + }, + }, + show_only=[ + "templates/pod-template-file.yaml", + ], + chart_dir=self.temp_chart_dir, + ) + + assert jmespath.search("spec.serviceAccountName", k8s_objects[0]) == "TestWorkerKubernetes" 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 bdb386194468f..570c0ccdc8abf 100644 --- a/helm-tests/tests/helm_tests/airflow_core/test_worker.py +++ b/helm-tests/tests/helm_tests/airflow_core/test_worker.py @@ -1256,19 +1256,100 @@ def test_should_add_component_specific_labels(self): assert jmespath.search("metadata.labels", docs[0])["test_label"] == "test_label_value" @pytest.mark.parametrize( - "executor, creates_service_account", + "executor", [ - ("LocalExecutor", False), - ("CeleryExecutor", True), - ("CeleryKubernetesExecutor", True), - ("CeleryExecutor,KubernetesExecutor", True), - ("KubernetesExecutor", True), - ("LocalKubernetesExecutor", True), + "CeleryKubernetesExecutor", + "CeleryExecutor,KubernetesExecutor", + "KubernetesExecutor", + "LocalKubernetesExecutor", ], ) - def test_should_create_worker_service_account_for_specific_executors( - self, executor, creates_service_account + @pytest.mark.parametrize( + "workers_values, obj", + [ + ({"serviceAccount": {"create": True}}, "worker"), + ( + { + "useWorkerDedicatedServiceAccounts": True, + "kubernetes": {"serviceAccount": {"create": True}}, + }, + "worker-kubernetes", + ), + ], + ) + def test_should_create_worker_service_account_for_specific_kubernetes_executors( + self, executor, workers_values, obj ): + docs = render_chart( + values={ + "executor": executor, + "workers": workers_values, + }, + show_only=[f"templates/workers/{obj}-serviceaccount.yaml"], + ) + + assert len(docs) == 1 + assert jmespath.search("kind", docs[0]) == "ServiceAccount" + + @pytest.mark.parametrize( + "executor", + [ + "CeleryExecutor", + "CeleryKubernetesExecutor", + "CeleryExecutor,KubernetesExecutor", + ], + ) + @pytest.mark.parametrize( + "workers_values, obj", + [ + ({"serviceAccount": {"create": True}}, "worker"), + ( + { + "useWorkerDedicatedServiceAccounts": True, + "celery": {"serviceAccount": {"create": True}}, + }, + "worker-celery", + ), + ], + ) + def test_should_create_worker_service_account_for_specific_celery_executors( + self, executor, workers_values, obj + ): + docs = render_chart( + values={ + "executor": executor, + "workers": workers_values, + }, + show_only=[f"templates/workers/{obj}-serviceaccount.yaml"], + ) + + assert len(docs) == 1 + assert jmespath.search("kind", docs[0]) == "ServiceAccount" + + def test_worker_service_account_creation_for_local_executor(self): + docs = render_chart( + values={ + "executor": "LocalExecutor", + "workers": { + "serviceAccount": {"create": True}, + }, + }, + show_only=["templates/workers/worker-serviceaccount.yaml"], + ) + + assert len(docs) == 0 + + @pytest.mark.parametrize( + "executor", + [ + "CeleryExecutor", + "CeleryKubernetesExecutor", + "CeleryExecutor,KubernetesExecutor", + "KubernetesExecutor", + "LocalKubernetesExecutor", + ], + ) + def test_worker_service_account_labels_per_executor(self, executor): docs = render_chart( values={ "executor": executor, @@ -1279,12 +1360,11 @@ def test_should_create_worker_service_account_for_specific_executors( }, show_only=["templates/workers/worker-serviceaccount.yaml"], ) - if creates_service_account: - assert jmespath.search("kind", docs[0]) == "ServiceAccount" - assert "test_label" in jmespath.search("metadata.labels", docs[0]) - assert jmespath.search("metadata.labels", docs[0])["test_label"] == "test_label_value" - else: - assert docs == [] + + assert len(docs) == 1 + assert jmespath.search("kind", docs[0]) == "ServiceAccount" + assert "test_label" in jmespath.search("metadata.labels", docs[0]) + assert jmespath.search("metadata.labels", docs[0])["test_label"] == "test_label_value" def test_default_automount_service_account_token(self): docs = render_chart( @@ -1297,13 +1377,75 @@ def test_default_automount_service_account_token(self): ) assert jmespath.search("automountServiceAccountToken", docs[0]) is True - def test_overridden_automount_service_account_token(self): + @pytest.mark.parametrize( + "workers_values, obj", + [ + ( + {"useWorkerDedicatedServiceAccounts": True, "celery": {"serviceAccount": {"create": True}}}, + "worker-celery", + ), + ( + { + "useWorkerDedicatedServiceAccounts": True, + "kubernetes": {"serviceAccount": {"create": True}}, + }, + "worker-kubernetes", + ), + ], + ) + def test_dedicated_service_account_token_automount(self, workers_values, obj): docs = render_chart( values={ - "workers": { - "serviceAccount": {"create": True, "automountServiceAccountToken": False}, + "executor": "CeleryExecutor,KubernetesExecutor", + "workers": workers_values, + }, + show_only=[f"templates/workers/{obj}-serviceaccount.yaml"], + ) + assert len(docs) == 1 + assert jmespath.search("automountServiceAccountToken", docs[0]) is True + + @pytest.mark.parametrize( + "workers_values, obj", + [ + ({"serviceAccount": {"create": True, "automountServiceAccountToken": False}}, "worker"), + ( + { + "useWorkerDedicatedServiceAccounts": True, + "celery": {"serviceAccount": {"create": True, "automountServiceAccountToken": False}}, + }, + "worker-celery", + ), + ( + { + "serviceAccount": {"create": True, "automountServiceAccountToken": True}, + "useWorkerDedicatedServiceAccounts": True, + "celery": {"serviceAccount": {"create": True, "automountServiceAccountToken": False}}, }, + "worker-celery", + ), + ( + { + "useWorkerDedicatedServiceAccounts": True, + "kubernetes": {"serviceAccount": {"create": True, "automountServiceAccountToken": False}}, + }, + "worker-kubernetes", + ), + ( + { + "serviceAccount": {"create": True, "automountServiceAccountToken": True}, + "useWorkerDedicatedServiceAccounts": True, + "kubernetes": {"serviceAccount": {"create": True, "automountServiceAccountToken": False}}, + }, + "worker-kubernetes", + ), + ], + ) + def test_overridden_automount_service_account_token(self, workers_values, obj): + docs = render_chart( + values={ + "executor": "CeleryExecutor,KubernetesExecutor", + "workers": workers_values, }, - show_only=["templates/workers/worker-serviceaccount.yaml"], + show_only=[f"templates/workers/{obj}-serviceaccount.yaml"], ) assert jmespath.search("automountServiceAccountToken", docs[0]) is False diff --git a/helm-tests/tests/helm_tests/security/test_rbac.py b/helm-tests/tests/helm_tests/security/test_rbac.py index aa93a91b691fe..819483459b3f2 100644 --- a/helm-tests/tests/helm_tests/security/test_rbac.py +++ b/helm-tests/tests/helm_tests/security/test_rbac.py @@ -64,7 +64,6 @@ SERVICE_ACCOUNT_NAME_TUPLES = [ ("ServiceAccount", "test-rbac-cleanup"), ("ServiceAccount", "test-rbac-scheduler"), - ("ServiceAccount", "test-rbac-worker"), ("ServiceAccount", "test-rbac-triggerer"), ("ServiceAccount", "test-rbac-pgbouncer"), ("ServiceAccount", "test-rbac-flower"), @@ -90,6 +89,8 @@ (CUSTOM_POSTGRESQL_NAME := "TestPostgresql"), ) CUSTOM_WEBSERVER_NAME = "TestWebserver" +CUSTOM_WORKER_CELERY_NAME = "TestWorkerCelery" +CUSTOM_WORKER_KUBERNETES_NAME = "TestWorkerKubernetes" parametrize_version = pytest.mark.parametrize("version", ["2.3.2", "2.4.0", "3.0.0", "default"]) @@ -105,7 +106,7 @@ def _get_values_with_version(self, values, version): def _is_airflow_3_or_above(self, version): return version == "default" or (parse_version(version) >= parse_version("3.0.0")) - def _get_object_tuples(self, version, sa: bool = True): + def _get_object_tuples(self, version, sa: bool = True, dedicated_workers_sa: None | bool = None): tuples = copy(DEPLOYMENT_NO_RBAC_NO_SA_KIND_NAME_TUPLES) if version in {"default", "3.0.0"}: tuples.append(("Service", "test-rbac-triggerer")) @@ -138,15 +139,34 @@ def _get_object_tuples(self, version, sa: bool = True): if sa: tuples.append(("ServiceAccount", "test-rbac-webserver")) + if dedicated_workers_sa is not None: + if dedicated_workers_sa: + tuples.append(("ServiceAccount", "test-rbac-worker-celery")) + tuples.append(("ServiceAccount", "test-rbac-worker-kubernetes")) + else: + tuples.append(("ServiceAccount", "test-rbac-worker")) + return tuples @parametrize_version - def test_deployments_no_rbac_no_sa(self, version): + @pytest.mark.parametrize( + "workers_values", + [ + {"serviceAccount": {"create": False}}, + { + "useWorkerDedicatedServiceAccounts": True, + "celery": {"serviceAccount": {"create": False}}, + "kubernetes": {"serviceAccount": {"create": False}}, + }, + ], + ) + def test_deployments_no_rbac_no_sa(self, version, workers_values): k8s_objects = render_chart( "test-rbac", values=self._get_values_with_version( values={ "fullnameOverride": "test-rbac", + "executor": "CeleryExecutor,KubernetesExecutor", "rbac": {"create": False}, "cleanup": { "enabled": True, @@ -165,7 +185,7 @@ def test_deployments_no_rbac_no_sa(self, version): "dagProcessor": {"serviceAccount": {"create": False}}, "webserver": {"serviceAccount": {"create": False}}, "apiServer": {"serviceAccount": {"create": False}}, - "workers": {"serviceAccount": {"create": False}}, + "workers": workers_values, "triggerer": {"serviceAccount": {"create": False}}, "statsd": {"serviceAccount": {"create": False}}, "createUserJob": {"serviceAccount": {"create": False}}, @@ -181,16 +201,19 @@ def test_deployments_no_rbac_no_sa(self, version): assert sorted(list_of_kind_names_tuples) == sorted(self._get_object_tuples(version, sa=False)) @parametrize_version - def test_deployments_no_rbac_with_sa(self, version): + @pytest.mark.parametrize("dedicated_workers_sa", [False, True]) + def test_deployments_no_rbac_with_sa(self, version, dedicated_workers_sa): k8s_objects = render_chart( "test-rbac", values=self._get_values_with_version( values={ "fullnameOverride": "test-rbac", + "executor": "CeleryExecutor,KubernetesExecutor", "rbac": {"create": False}, "cleanup": {"enabled": True}, "flower": {"enabled": True}, "pgbouncer": {"enabled": True}, + "workers": {"useWorkerDedicatedServiceAccounts": dedicated_workers_sa}, }, version=version, ), @@ -198,16 +221,31 @@ def test_deployments_no_rbac_with_sa(self, version): list_of_kind_names_tuples = [ (k8s_object["kind"], k8s_object["metadata"]["name"]) for k8s_object in k8s_objects ] - real_list_of_kind_names = self._get_object_tuples(version) + SERVICE_ACCOUNT_NAME_TUPLES + real_list_of_kind_names = ( + self._get_object_tuples(version, dedicated_workers_sa=dedicated_workers_sa) + + SERVICE_ACCOUNT_NAME_TUPLES + ) assert sorted(list_of_kind_names_tuples) == sorted(real_list_of_kind_names) @parametrize_version - def test_deployments_with_rbac_no_sa(self, version): + @pytest.mark.parametrize( + "workers_values", + [ + {"serviceAccount": {"create": False}}, + { + "useWorkerDedicatedServiceAccounts": True, + "celery": {"serviceAccount": {"create": False}}, + "kubernetes": {"serviceAccount": {"create": False}}, + }, + ], + ) + def test_deployments_with_rbac_no_sa(self, version, workers_values): k8s_objects = render_chart( "test-rbac", values=self._get_values_with_version( values={ "fullnameOverride": "test-rbac", + "executor": "CeleryExecutor,KubernetesExecutor", "cleanup": { "enabled": True, "serviceAccount": { @@ -218,7 +256,7 @@ def test_deployments_with_rbac_no_sa(self, version): "dagProcessor": {"serviceAccount": {"create": False}}, "webserver": {"serviceAccount": {"create": False}}, "apiServer": {"serviceAccount": {"create": False}}, - "workers": {"serviceAccount": {"create": False}}, + "workers": workers_values, "triggerer": {"serviceAccount": {"create": False}}, "flower": {"enabled": True, "serviceAccount": {"create": False}}, "statsd": {"serviceAccount": {"create": False}}, @@ -242,15 +280,18 @@ def test_deployments_with_rbac_no_sa(self, version): assert sorted(list_of_kind_names_tuples) == sorted(real_list_of_kind_names) @parametrize_version - def test_deployments_with_rbac_with_sa(self, version): + @pytest.mark.parametrize("dedicated_workers_sa", [False, True]) + def test_deployments_with_rbac_with_sa(self, version, dedicated_workers_sa): k8s_objects = render_chart( "test-rbac", values=self._get_values_with_version( values={ "fullnameOverride": "test-rbac", + "executor": "CeleryExecutor,KubernetesExecutor", "cleanup": {"enabled": True}, "flower": {"enabled": True}, "pgbouncer": {"enabled": True}, + "workers": {"useWorkerDedicatedServiceAccounts": dedicated_workers_sa}, }, version=version, ), @@ -259,7 +300,9 @@ def test_deployments_with_rbac_with_sa(self, version): (k8s_object["kind"], k8s_object["metadata"]["name"]) for k8s_object in k8s_objects ] real_list_of_kind_names = ( - self._get_object_tuples(version) + SERVICE_ACCOUNT_NAME_TUPLES + RBAC_ENABLED_KIND_NAME_TUPLES + self._get_object_tuples(version, dedicated_workers_sa=dedicated_workers_sa) + + SERVICE_ACCOUNT_NAME_TUPLES + + RBAC_ENABLED_KIND_NAME_TUPLES ) assert sorted(list_of_kind_names_tuples) == sorted(real_list_of_kind_names) @@ -301,6 +344,33 @@ def test_service_account_custom_names(self): ] assert sorted(list_of_sa_names) == sorted(CUSTOM_SERVICE_ACCOUNT_NAMES) + def test_workers_service_account_custom_name(self): + k8s_objects = render_chart( + "test-rbac", + values={ + "airflowVersion": "3.0.0", + "fullnameOverride": "test-rbac", + "executor": "CeleryExecutor,KubernetesExecutor", + "workers": { + "useWorkerDedicatedServiceAccounts": True, + "celery": {"serviceAccount": {"name": CUSTOM_WORKER_CELERY_NAME}}, + "kubernetes": {"serviceAccount": {"name": CUSTOM_WORKER_KUBERNETES_NAME}}, + }, + }, + show_only=[ + "templates/workers/worker-celery-serviceaccount.yaml", + "templates/workers/worker-kubernetes-serviceaccount.yaml", + ], + ) + + list_of_sa_names = [ + k8s_object["metadata"]["name"] + for k8s_object in k8s_objects + if k8s_object["kind"] == "ServiceAccount" + ] + assert len(k8s_objects) == 2 + assert sorted(list_of_sa_names) == [CUSTOM_WORKER_CELERY_NAME, CUSTOM_WORKER_KUBERNETES_NAME] + def test_webserver_service_account_name_airflow_2(self): k8s_objects = render_chart( "test-rbac", @@ -360,6 +430,27 @@ def test_service_account_custom_names_in_objects(self): assert sorted(list_of_sa_names_in_objects) == sorted(CUSTOM_SERVICE_ACCOUNT_NAMES) + def test_workers_celery_service_account_custom_names_in_objects(self): + k8s_objects = render_chart( + "test-rbac", + values={ + "airflowVersion": "3.0.0", + "fullnameOverride": "test-rbac", + "workers": { + "useWorkerDedicatedServiceAccounts": True, + "celery": {"serviceAccount": {"name": CUSTOM_WORKER_CELERY_NAME}}, + }, + }, + show_only=[ + "templates/workers/worker-deployment.yaml", + ], + ) + + assert ( + jmespath.search("spec.template.spec.serviceAccountName", k8s_objects[0]) + == CUSTOM_WORKER_CELERY_NAME + ) + def test_service_account_without_resource(self): k8s_objects = render_chart( "test-rbac",