Add statsd overrideMappings in helm chart values

chart/templates/configmaps/statsd-configmap.yaml

@@ -18,7 +18,7 @@
 ################################
 ## Airflow StatsD ConfigMap
 #################################
-{{- if and .Values.statsd.enabled .Values.statsd.extraMappings }}
+{{- if and .Values.statsd.enabled }}
 apiVersion: v1
 kind: ConfigMap
 metadata:
@@ -34,6 +34,13 @@ metadata:
 {{- end }}
 data:
   mappings.yml: |-
-{{ .Files.Get "dockerfiles/statsd-exporter/mappings.yml" | indent 4 }}
+{{- if .Values.statsd.overrideMappings }}
+    mappings:
+{{ toYaml .Values.statsd.overrideMappings | indent 6 }}
+{{- else }}
+{{ .Files.Get "files/statsd-mappings.yml" | indent 4 }}
+{{- if .Values.statsd.extraMappings }}
 {{ toYaml .Values.statsd.extraMappings | indent 6 }}
 {{- end }}
+{{- end }}
+{{- end }}

chart/templates/statsd/statsd-deployment.yaml

@@ -59,9 +59,7 @@ spec:
 {{- end }}
 {{- if or .Values.statsd.extraMappings .Values.statsd.podAnnotations }}
       annotations:
-        {{- if .Values.statsd.extraMappings }}
         checksum/statsd-config: {{ include (print $.Template.BasePath "/configmaps/statsd-configmap.yaml") . | sha256sum }}
-        {{- end }}
         {{- if .Values.statsd.podAnnotations }}
         {{- toYaml .Values.statsd.podAnnotations | nindent 8 }}
         {{- end }}
@@ -113,7 +111,6 @@ spec:
             initialDelaySeconds: 10
             periodSeconds: 10
             timeoutSeconds: 5
-{{- if .Values.statsd.extraMappings }}
           volumeMounts:
             - name: config
               mountPath: /etc/statsd-exporter/mappings.yml
@@ -123,4 +120,3 @@ spec:
           configMap:
             name: {{ .Release.Name }}-statsd
 {{- end }}
-{{- end }}

chart/values.schema.json

@@ -571,12 +571,12 @@
                         "repository": {
                             "description": "The StatsD image repository.",
                             "type": "string",
-                            "default": "apache/airflow"
+                            "default": "quay.io/prometheus/statsd-exporter"
                         },
                         "tag": {
                             "description": "The StatsD image tag.",
                             "type": "string",
-                            "default": "airflow-statsd-exporter-2021.04.28-v0.17.0"
+                            "default": "v0.22.8"
                         },
                         "pullPolicy": {
                             "description": "The StatsD image pull policy.",
@@ -4117,7 +4117,12 @@
                     }
                 },
                 "extraMappings": {
-                    "description": "Additional mappings for StatsD exporter.",
+                    "description": "Additional mappings for StatsD exporter.If set, will merge default mapping and extra mappings, default mapping has higher priority. So, if you want to change some default mapping, please use `overrideMappings`",
+                    "type": "array",
+                    "default": []
+                },
+                "overrideMappings": {
+                    "description": "Override mappings for StatsD exporter.If set, will ignore setting item in default and `extraMappings`. So, If you use it, ensure all mapping item contains in it.",
                     "type": "array",
                     "default": []
                 },

chart/values.yaml

@@ -76,8 +76,8 @@ images:
     tag: ~
     pullPolicy: IfNotPresent
   statsd:
-    repository: apache/airflow
-    tag: airflow-statsd-exporter-2021.04.28-v0.17.0
+    repository: quay.io/prometheus/statsd-exporter
+    tag: v0.22.8
     pullPolicy: IfNotPresent
   redis:
     repository: redis
@@ -1366,8 +1366,15 @@ statsd:
   priorityClassName: ~
 
   # Additional mappings for StatsD exporter.
+  # If set, will merge default mapping and extra mappings, default mapping has higher priority.
+  # So, if you want to change some default mapping, please use `overrideMappings`
   extraMappings: []
 
+  # Override mappings for StatsD exporter.
+  # If set, will ignore setting item in default and `extraMappings`.
+  # So, If you use it, ensure all mapping item contains in it.
+  overrideMappings: []
+
   podAnnotations: {}
 
 # PgBouncer settings

scripts/ci/pre_commit/pre_commit_chart_schema.py

@@ -36,6 +36,7 @@
     "$['properties']['ingress']['properties']['web']['properties']['succeedingPaths']",
     # The value of this parameter is passed to statsd_exporter, which does not have a strict type definition.
     "$['properties']['statsd']['properties']['extraMappings']",
+    "$['properties']['statsd']['properties']['overrideMappings']",
 }
 VENDORED_PATHS = {
     # We don't want to check the upstream k8s definitions

tests/charts/test_basic_helm_chart.py

@@ -27,7 +27,7 @@
 
 from tests.charts.helm_template_generator import render_chart
 
-OBJECT_COUNT_IN_BASIC_DEPLOYMENT = 34
+OBJECT_COUNT_IN_BASIC_DEPLOYMENT = 35
 
 
 class TestBaseChartTest(unittest.TestCase):
@@ -81,6 +81,7 @@ def test_basic_deployments(self, version):
             ('Secret', 'test-basic-postgresql'),
             ('Secret', 'test-basic-redis-password'),
             ('ConfigMap', 'test-basic-airflow-config'),
+            ('ConfigMap', 'test-basic-statsd'),
             ('Role', 'test-basic-pod-launcher-role'),
             ('Role', 'test-basic-pod-log-reader-role'),
             ('RoleBinding', 'test-basic-pod-launcher-rolebinding'),
@@ -158,6 +159,7 @@ def test_basic_deployment_with_standalone_dag_processor(self, version):
             ('Secret', 'test-basic-postgresql'),
             ('Secret', 'test-basic-redis-password'),
             ('ConfigMap', 'test-basic-airflow-config'),
+            ('ConfigMap', 'test-basic-statsd'),
             ('Role', 'test-basic-pod-launcher-role'),
             ('Role', 'test-basic-pod-log-reader-role'),
             ('RoleBinding', 'test-basic-pod-launcher-rolebinding'),
@@ -208,6 +210,23 @@ def test_basic_deployment_without_default_users(self, version):
         assert ('Job', 'test-basic-create-user') not in list_of_kind_names_tuples
         assert expected_object_count_in_basic_deployment - 2 == len(k8s_objects)
 
+    @parameterized.expand(["2.3.2", "2.4.0", "default"])
+    def test_basic_deployment_without_statsd(self, version):
+        expected_object_count_in_basic_deployment = self._get_object_count(version)
+        k8s_objects = render_chart(
+            "test-basic",
+            values=self._get_values_with_version(values={"statsd": {'enabled': False}}, version=version),
+        )
+        list_of_kind_names_tuples = [
+            (k8s_object['kind'], k8s_object['metadata']['name']) for k8s_object in k8s_objects
+        ]
+        assert ('ServiceAccount', 'test-basic-statsd') not in list_of_kind_names_tuples
+        assert ('ConfigMap', 'test-basic-statsd') not in list_of_kind_names_tuples
+        assert ('Service', 'test-basic-statsd') not in list_of_kind_names_tuples
+        assert ('Deployment', 'test-basic-statsd') not in list_of_kind_names_tuples
+
+        assert expected_object_count_in_basic_deployment - 4 == len(k8s_objects)
+
     def test_network_policies_are_valid(self):
         k8s_objects = render_chart(
             "test-basic",

tests/charts/test_rbac.py

@@ -29,6 +29,7 @@
     ('Secret', 'test-rbac-pgbouncer-config'),
     ('Secret', 'test-rbac-pgbouncer-stats'),
     ('ConfigMap', 'test-rbac-airflow-config'),
+    ('ConfigMap', 'test-rbac-statsd'),
     ('Service', 'test-rbac-postgresql-headless'),
     ('Service', 'test-rbac-postgresql'),
     ('Service', 'test-rbac-statsd'),

tests/charts/test_statsd.py

@@ -19,6 +19,7 @@
 import unittest
 
 import jmespath
+import yaml
 from parameterized import parameterized
 
 from tests.charts.helm_template_generator import render_chart
@@ -32,6 +33,16 @@ def test_should_create_statsd_default(self):
 
         assert "statsd" == jmespath.search("spec.template.spec.containers[0].name", docs[0])
 
+        assert {"name": "config", "configMap": {"name": "release-name-statsd"}} in jmespath.search(
+            "spec.template.spec.volumes", docs[0]
+        )
+
+        assert {
+            "name": "config",
+            "mountPath": "/etc/statsd-exporter/mappings.yml",
+            "subPath": "mappings.yml",
+        } in jmespath.search("spec.template.spec.containers[0].volumeMounts", docs[0])
+
     def test_should_add_volume_and_volume_mount_when_exist_extra_mappings(self):
         extra_mapping = {
             "match": "airflow.pool.queued_slots.*",
@@ -53,6 +64,27 @@ def test_should_add_volume_and_volume_mount_when_exist_extra_mappings(self):
             "subPath": "mappings.yml",
         } in jmespath.search("spec.template.spec.containers[0].volumeMounts", docs[0])
 
+    def test_should_add_volume_and_volume_mount_when_exist_override_mappings(self):
+        override_mapping = {
+            "match": "airflow.pool.queued_slots.*",
+            "name": "airflow_pool_queued_slots",
+            "labels": {"pool": "$1"},
+        }
+        docs = render_chart(
+            values={"statsd": {"enabled": True, "overrideMappings": [override_mapping]}},
+            show_only=["templates/statsd/statsd-deployment.yaml"],
+        )
+
+        assert {"name": "config", "configMap": {"name": "release-name-statsd"}} in jmespath.search(
+            "spec.template.spec.volumes", docs[0]
+        )
+
+        assert {
+            "name": "config",
+            "mountPath": "/etc/statsd-exporter/mappings.yml",
+            "subPath": "mappings.yml",
+        } in jmespath.search("spec.template.spec.containers[0].volumeMounts", docs[0])
+
     @parameterized.expand([(8, 10), (10, 8), (8, None), (None, 10), (None, None)])
     def test_revision_history_limit(self, revision_history_limit, global_revision_history_limit):
         values = {"statsd": {"enabled": True}}
@@ -134,3 +166,46 @@ def test_statsd_resources_are_not_added_by_default(self):
             show_only=["templates/statsd/statsd-deployment.yaml"],
         )
         assert jmespath.search("spec.template.spec.containers[0].resources", docs[0]) == {}
+
+    def test_statsd_configmap_by_default(self):
+        docs = render_chart(show_only=["templates/configmaps/statsd-configmap.yaml"])
+
+        mappings_yml = jmespath.search('data."mappings.yml"', docs[0])
+        mappings_yml_obj = yaml.safe_load(mappings_yml)
+
+        assert "airflow_dagrun_dependency_check" == mappings_yml_obj["mappings"][0]["name"]
+        assert "airflow_pool_starving_tasks" == mappings_yml_obj["mappings"][-1]["name"]
+
+    def test_statsd_configmap_when_exist_extra_mappings(self):
+        extra_mapping = {
+            "match": "airflow.pool.queued_slots.*",
+            "name": "airflow_pool_queued_slots",
+            "labels": {"pool": "$1"},
+        }
+        docs = render_chart(
+            values={"statsd": {"enabled": True, "extraMappings": [extra_mapping]}},
+            show_only=["templates/configmaps/statsd-configmap.yaml"],
+        )
+
+        mappings_yml = jmespath.search('data."mappings.yml"', docs[0])
+        mappings_yml_obj = yaml.safe_load(mappings_yml)
+
+        assert "airflow_dagrun_dependency_check" == mappings_yml_obj["mappings"][0]["name"]
+        assert "airflow_pool_queued_slots" == mappings_yml_obj["mappings"][-1]["name"]
+
+    def test_statsd_configmap_when_exist_override_mappings(self):
+        override_mapping = {
+            "match": "airflow.pool.queued_slots.*",
+            "name": "airflow_pool_queued_slots",
+            "labels": {"pool": "$1"},
+        }
+        docs = render_chart(
+            values={"statsd": {"enabled": True, "overrideMappings": [override_mapping]}},
+            show_only=["templates/configmaps/statsd-configmap.yaml"],
+        )
+
+        mappings_yml = jmespath.search('data."mappings.yml"', docs[0])
+        mappings_yml_obj = yaml.safe_load(mappings_yml)
+
+        assert 1 == len(mappings_yml_obj["mappings"])
+        assert "airflow_pool_queued_slots" == mappings_yml_obj["mappings"][0]["name"]