diff --git a/chart/templates/flower/flower-deployment.yaml b/chart/templates/flower/flower-deployment.yaml index cb162414174f5..a5e145525377d 100644 --- a/chart/templates/flower/flower-deployment.yaml +++ b/chart/templates/flower/flower-deployment.yaml @@ -54,9 +54,9 @@ spec: {{- with .Values.labels }} {{ toYaml . | indent 8 }} {{- end }} - {{- if .Values.airflowPodAnnotations }} + {{- if or (.Values.airflowPodAnnotations) (.Values.flower.podAnnotations) }} annotations: - {{- toYaml .Values.airflowPodAnnotations | nindent 8 }} + {{- mustMerge .Values.flower.podAnnotations .Values.airflowPodAnnotations | toYaml | nindent 8 }} {{- end }} spec: nodeSelector: diff --git a/chart/templates/scheduler/scheduler-deployment.yaml b/chart/templates/scheduler/scheduler-deployment.yaml index 0ac86f4aa0b57..a73c6b90e4ced 100644 --- a/chart/templates/scheduler/scheduler-deployment.yaml +++ b/chart/templates/scheduler/scheduler-deployment.yaml @@ -85,6 +85,9 @@ spec: {{- if .Values.airflowPodAnnotations }} {{- toYaml .Values.airflowPodAnnotations | nindent 8 }} {{- end }} + {{- if .Values.scheduler.podAnnotations }} + {{- toYaml .Values.scheduler.podAnnotations | nindent 8 }} + {{- end }} spec: nodeSelector: {{ toYaml $nodeSelector | indent 8 }} diff --git a/chart/templates/triggerer/triggerer-deployment.yaml b/chart/templates/triggerer/triggerer-deployment.yaml index fbfbcd2795583..e97f45a822a09 100644 --- a/chart/templates/triggerer/triggerer-deployment.yaml +++ b/chart/templates/triggerer/triggerer-deployment.yaml @@ -67,6 +67,9 @@ spec: {{- if .Values.airflowPodAnnotations }} {{- toYaml .Values.airflowPodAnnotations | nindent 8 }} {{- end }} + {{- if .Values.triggerer.podAnnotations }} + {{- toYaml .Values.triggerer.podAnnotations | nindent 8 }} + {{- end }} spec: nodeSelector: {{- toYaml $nodeSelector | nindent 8 }} diff --git a/chart/templates/webserver/webserver-deployment.yaml b/chart/templates/webserver/webserver-deployment.yaml index da89222d84105..57e564d421c91 100644 --- a/chart/templates/webserver/webserver-deployment.yaml +++ b/chart/templates/webserver/webserver-deployment.yaml @@ -81,6 +81,9 @@ spec: {{- if .Values.airflowPodAnnotations }} {{- toYaml .Values.airflowPodAnnotations | nindent 8 }} {{- end }} + {{- if .Values.webserver.podAnnotations }} + {{- toYaml .Values.webserver.podAnnotations | nindent 8 }} + {{- end }} spec: serviceAccountName: {{ include "webserver.serviceAccountName" . }} nodeSelector: diff --git a/chart/templates/workers/worker-deployment.yaml b/chart/templates/workers/worker-deployment.yaml index 5af41dc1e0f40..de6b4c64622a3 100644 --- a/chart/templates/workers/worker-deployment.yaml +++ b/chart/templates/workers/worker-deployment.yaml @@ -77,6 +77,9 @@ spec: {{- if .Values.airflowPodAnnotations }} {{- toYaml .Values.airflowPodAnnotations | nindent 8 }} {{- end }} + {{- if .Values.workers.podAnnotations }} + {{- toYaml .Values.workers.podAnnotations | nindent 8 }} + {{- end }} spec: nodeSelector: {{ toYaml $nodeSelector | indent 8 }} diff --git a/chart/tests/test_airflow_common.py b/chart/tests/test_airflow_common.py index 9122d35d2d87c..66244be5fa5e3 100644 --- a/chart/tests/test_airflow_common.py +++ b/chart/tests/test_airflow_common.py @@ -74,16 +74,20 @@ def test_annotations(self): release_name = "TEST-BASIC" k8s_objects = render_chart( name=release_name, - values={"airflowPodAnnotations": {"test-annotation/safe-to-evict": "true"}}, + values={ + "airflowVersion": "2.2.0", # Needed for triggerer to be enabled. + "airflowPodAnnotations": {"test-annotation/safe-to-evict": "true"}, + }, show_only=[ "templates/scheduler/scheduler-deployment.yaml", "templates/workers/worker-deployment.yaml", "templates/webserver/webserver-deployment.yaml", "templates/flower/flower-deployment.yaml", + "templates/triggerer/triggerer-deployment.yaml", ], ) - assert 4 == len(k8s_objects) + assert 5 == len(k8s_objects) for k8s_object in k8s_objects: annotations = k8s_object["spec"]["template"]["metadata"]["annotations"] diff --git a/chart/tests/test_annotations.py b/chart/tests/test_annotations.py index 899d697131a53..48dd28f688f1e 100644 --- a/chart/tests/test_annotations.py +++ b/chart/tests/test_annotations.py @@ -15,124 +15,309 @@ # specific language governing permissions and limitations # under the License. -import unittest +import pytest from tests.helm_template_generator import render_chart -# Values for each service mapped to the 'example' -# key annotation -CUSTOM_ANNOTATION_VALUES = ( - CUSTOM_SCHEDULER_ANNOTATION, - CUSTOM_WEBSERVER_ANNOTATION, - CUSTOM_WORKER_ANNOTATION, - CUSTOM_CLEANUP_ANNOTATION, - CUSTOM_FLOWER_ANNOTATION, - CUSTOM_PGBOUNCER_ANNOTATION, - CUSTOM_STATSD_ANNOTATION, - CUSTOM_CREATE_USER_JOB_ANNOTATION, - CUSTOM_MIGRATE_DATABASE_JOB_ANNOTATION, - CUSTOM_REDIS_ANNOTATION, -) = ( - "scheduler", - "webserver", - "worker", - "cleanup", - "flower", - "pgbouncer", - "statsd", - "createuser", - "migratedb", - "redis", -) - -class AnnotationsTest(unittest.TestCase): - def test_service_account_annotations(self): - k8s_objects = render_chart( - values={ - "cleanup": { - "enabled": True, - "serviceAccount": { - "annotations": { - "example": CUSTOM_CLEANUP_ANNOTATION, +class TestServiceAccountAnnotations: + @pytest.mark.parametrize( + "values,show_only,expected_annotations", + [ + ( + { + "cleanup": { + "enabled": True, + "serviceAccount": { + "annotations": { + "example": "cleanup", + }, }, }, }, - "scheduler": { - "serviceAccount": { - "annotations": { - "example": CUSTOM_SCHEDULER_ANNOTATION, + "templates/cleanup/cleanup-serviceaccount.yaml", + { + "example": "cleanup", + }, + ), + ( + { + "scheduler": { + "serviceAccount": { + "annotations": { + "example": "scheduler", + }, }, }, }, - "webserver": { - "serviceAccount": { - "annotations": { - "example": CUSTOM_WEBSERVER_ANNOTATION, + "templates/scheduler/scheduler-serviceaccount.yaml", + { + "example": "scheduler", + }, + ), + ( + { + "webserver": { + "serviceAccount": { + "annotations": { + "example": "webserver", + }, }, }, }, - "workers": { - "serviceAccount": { - "annotations": { - "example": CUSTOM_WORKER_ANNOTATION, + "templates/webserver/webserver-serviceaccount.yaml", + { + "example": "webserver", + }, + ), + ( + { + "workers": { + "serviceAccount": { + "annotations": { + "example": "worker", + }, }, }, }, - "flower": { - "serviceAccount": { - "annotations": { - "example": CUSTOM_FLOWER_ANNOTATION, + "templates/workers/worker-serviceaccount.yaml", + { + "example": "worker", + }, + ), + ( + { + "flower": { + "serviceAccount": { + "annotations": { + "example": "flower", + }, }, }, }, - "statsd": { - "serviceAccount": { - "annotations": { - "example": CUSTOM_STATSD_ANNOTATION, + "templates/flower/flower-serviceaccount.yaml", + { + "example": "flower", + }, + ), + ( + { + "statsd": { + "serviceAccount": { + "annotations": { + "example": "statsd", + }, }, }, }, - "redis": { - "serviceAccount": { - "annotations": { - "example": CUSTOM_REDIS_ANNOTATION, + "templates/statsd/statsd-serviceaccount.yaml", + { + "example": "statsd", + }, + ), + ( + { + "redis": { + "serviceAccount": { + "annotations": { + "example": "redis", + }, }, }, }, - "pgbouncer": { - "enabled": True, - "serviceAccount": { - "annotations": { - "example": CUSTOM_PGBOUNCER_ANNOTATION, + "templates/redis/redis-serviceaccount.yaml", + { + "example": "redis", + }, + ), + ( + { + "pgbouncer": { + "enabled": True, + "serviceAccount": { + "annotations": { + "example": "pgbouncer", + }, }, }, }, - "createUserJob": { - "serviceAccount": { - "annotations": { - "example": CUSTOM_CREATE_USER_JOB_ANNOTATION, + "templates/pgbouncer/pgbouncer-serviceaccount.yaml", + { + "example": "pgbouncer", + }, + ), + ( + { + "createUserJob": { + "serviceAccount": { + "annotations": { + "example": "createuser", + }, }, }, }, - "migrateDatabaseJob": { - "serviceAccount": { - "annotations": { - "example": CUSTOM_MIGRATE_DATABASE_JOB_ANNOTATION, + "templates/jobs/create-user-job-serviceaccount.yaml", + { + "example": "createuser", + }, + ), + ( + { + "migrateDatabaseJob": { + "serviceAccount": { + "annotations": { + "example": "migratedb", + }, }, }, }, - "executor": "CeleryExecutor", # create worker deployment + "templates/jobs/migrate-database-job-serviceaccount.yaml", + { + "example": "migratedb", + }, + ), + ( + { + "airflowVersion": "2.2.0", # Needed for triggerer to be enabled. + "triggerer": { + "serviceAccount": { + "annotations": { + "example": "triggerer", + }, + }, + }, + }, + "templates/triggerer/triggerer-serviceaccount.yaml", + { + "example": "triggerer", + }, + ), + ], + ) + def test_annotations_are_added(self, values, show_only, expected_annotations): + k8s_objects = render_chart( + values=values, + show_only=[show_only], + ) + + # This test relies on the convention that the helm chart puts a single + # ServiceAccount in its own .yaml file, so by specifying `show_only`, + # we should only get a single k8s_object here - the target object that + # we hope to test on. + assert len(k8s_objects) == 1 + obj = k8s_objects[0] + + for k, v in expected_annotations.items(): + assert k in obj["metadata"]["annotations"] + assert v == obj["metadata"]["annotations"][k] + + +@pytest.mark.parametrize( + "values,show_only,expected_annotations", + [ + ( + { + "scheduler": { + "podAnnotations": { + "example": "scheduler", + }, + }, + }, + "templates/scheduler/scheduler-deployment.yaml", + { + "example": "scheduler", + }, + ), + ( + { + "webserver": { + "podAnnotations": { + "example": "webserver", + }, + }, }, + "templates/webserver/webserver-deployment.yaml", + { + "example": "webserver", + }, + ), + ( + { + "workers": { + "podAnnotations": { + "example": "worker", + }, + }, + }, + "templates/workers/worker-deployment.yaml", + { + "example": "worker", + }, + ), + ( + { + "flower": { + "podAnnotations": { + "example": "flower", + }, + }, + }, + "templates/flower/flower-deployment.yaml", + { + "example": "flower", + }, + ), + ( + { + "airflowVersion": "2.2.0", # Needed for triggerer to be enabled. + "triggerer": { + "podAnnotations": { + "example": "triggerer", + }, + }, + }, + "templates/triggerer/triggerer-deployment.yaml", + { + "example": "triggerer", + }, + ), + ], +) +class TestPerComponentPodAnnotations: + def test_annotations_are_added(self, values, show_only, expected_annotations): + k8s_objects = render_chart( + values=values, + show_only=[show_only], ) - list_of_annotation_values_in_objects = [ - k8s_object['metadata']['annotations']['example'] - for k8s_object in k8s_objects - if k8s_object['kind'] == "ServiceAccount" - ] + # This test relies on the convention that the helm chart puts a single + # Deployment in its own .yaml file, so by specifying `show_only`, + # we should only get a single k8s_object here - the target object that + # we hope to test on. + assert len(k8s_objects) == 1 + obj = k8s_objects[0] + + for k, v in expected_annotations.items(): + assert k in obj["spec"]["template"]["metadata"]["annotations"] + assert v == obj["spec"]["template"]["metadata"]["annotations"][k] + + def test_precedence(self, values, show_only, expected_annotations): + values_global_annotations = {"airflowPodAnnotations": {k: "GLOBAL" for k in expected_annotations}} + + values_merged = {**values, **values_global_annotations} - self.assertCountEqual( - list_of_annotation_values_in_objects, - CUSTOM_ANNOTATION_VALUES, + k8s_objects = render_chart( + values=values_merged, + show_only=[show_only], ) + + # This test relies on the convention that the helm chart puts a single + # Deployment in its own .yaml file, so by specifying `show_only`, + # we should only get a single k8s_object here - the target object that + # we hope to test on. + assert len(k8s_objects) == 1 + obj = k8s_objects[0] + + for k, v in expected_annotations.items(): + assert k in obj["spec"]["template"]["metadata"]["annotations"] + assert v == obj["spec"]["template"]["metadata"]["annotations"][k] diff --git a/chart/values.schema.json b/chart/values.schema.json index 4c01287ca3028..012fc2d4d0478 100644 --- a/chart/values.schema.json +++ b/chart/values.schema.json @@ -1141,6 +1141,11 @@ } ] }, + "podAnnotations": { + "description": "Annotations to add to the worker pods.", + "type": "object", + "default": {} + }, "logGroomerSidecar": { "description": "Configuration for worker log groomer sidecar", "type": "object", @@ -1384,6 +1389,11 @@ "type": "array", "default": [] }, + "podAnnotations": { + "description": "Annotations to add to the scheduler pods.", + "type": "object", + "default": {} + }, "logGroomerSidecar": { "description": "Configuration for the schedulers log groomer sidecar.", "type": "object", @@ -1609,6 +1619,11 @@ "description": "Grace period for tasks to finish after SIGTERM is sent from Kubernetes.", "type": "integer", "default": 60 + }, + "podAnnotations": { + "description": "Annotations to add to the triggerer pods.", + "type": "object", + "default": {} } } }, @@ -2067,6 +2082,11 @@ "description": "Specify Tolerations for webserver pods.", "type": "array", "default": [] + }, + "podAnnotations": { + "description": "Annotations to add to the webserver pods.", + "type": "object", + "default": {} } } }, @@ -2287,6 +2307,11 @@ "description": "Specify Tolerations for Flower pods.", "type": "array", "default": [] + }, + "podAnnotations": { + "description": "Annotations to add to the Flower pods.", + "type": "object", + "default": {} } } }, diff --git a/chart/values.yaml b/chart/values.yaml index 0bc344c8098ba..b535edc48330d 100644 --- a/chart/values.yaml +++ b/chart/values.yaml @@ -485,6 +485,8 @@ workers: # hostnames: # - "test.hostname.two" + podAnnotations: {} + logGroomerSidecar: # Command to use when running the Airflow worker log groomer sidecar (templated). command: ~ @@ -578,6 +580,8 @@ scheduler: topologyKey: "kubernetes.io/hostname" tolerations: [] + podAnnotations: {} + logGroomerSidecar: # Whether to deploy the Airflow scheduler log groomer sidecar. enabled: true @@ -594,6 +598,7 @@ scheduler: # requests: # cpu: 100m # memory: 128Mi + # Airflow create user job settings createUserJob: # Annotations on the create user job pod @@ -762,6 +767,8 @@ webserver: topologyKey: "kubernetes.io/hostname" tolerations: [] + podAnnotations: {} + # Airflow Triggerer Config triggerer: # Number of airflow triggerers in the deployment @@ -834,6 +841,8 @@ triggerer: topologyKey: "kubernetes.io/hostname" tolerations: [] + podAnnotations: {} + # Flower settings flower: # Enable flower. @@ -915,6 +924,8 @@ flower: affinity: {} tolerations: [] + podAnnotations: {} + # Statsd settings statsd: enabled: true