From 7382948a312426e7889583ad66d864251514202d Mon Sep 17 00:00:00 2001
From: DerekHeldtWerle <derek.heldt-werle@viasat.com>
Date: Wed, 12 May 2021 16:34:32 -0700
Subject: [PATCH] Chart: Allow setting annotations on Airflow pods &
 `Configmap` (#15238)

This PR adds a new field (`airflowConfigAnnotations`) that allows users to add `annotations` to the main `configmap.yaml` file.

I ended up setting up a new testing file as I didn't find a file where this specifically fit, but if it should be moved elsewhere let me know.

closes https://github.com/apache/airflow/issues/13643

GitOrigin-RevId: 6d64cc54a6b7d1b22d0de89b5815035e21bfaf8c
---
 chart/templates/configmaps/configmap.yaml |  4 ++
 chart/tests/test_configmap.py             | 47 +++++++++++++++++++++++
 chart/values.schema.json                  |  4 ++
 chart/values.yaml                         |  4 ++
 docs/helm-chart/parameters-ref.rst        |  6 +++
 5 files changed, 65 insertions(+)
 create mode 100644 chart/tests/test_configmap.py

diff --git a/chart/templates/configmaps/configmap.yaml b/chart/templates/configmaps/configmap.yaml
index 120ca85c94a..9a26d181d48 100644
--- a/chart/templates/configmaps/configmap.yaml
+++ b/chart/templates/configmaps/configmap.yaml
@@ -31,6 +31,10 @@ metadata:
 {{- with .Values.labels }}
 {{ toYaml . | indent 4 }}
 {{- end -}}
+{{- if .Values.airflowConfigAnnotations }}
+  annotations:
+{{- toYaml .Values.airflowConfigAnnotations | nindent 4 }}
+{{- end }}
 {{- $Global := . }}
 data:
   # These are system-specified config overrides.
diff --git a/chart/tests/test_configmap.py b/chart/tests/test_configmap.py
new file mode 100644
index 00000000000..6a9e2868703
--- /dev/null
+++ b/chart/tests/test_configmap.py
@@ -0,0 +1,47 @@
+# 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.
+
+import unittest
+
+import jmespath
+
+from tests.helm_template_generator import render_chart
+
+
+class ConfigmapTest(unittest.TestCase):
+    def test_single_annotation(self):
+        docs = render_chart(
+            values={
+                "airflowConfigAnnotations": {"key": "value"},
+            },
+            show_only=["templates/configmaps/configmap.yaml"],
+        )
+
+        annotations = jmespath.search("metadata.annotations", docs[0])
+        assert "value" == annotations.get("key")
+
+    def test_multiple_annotations(self):
+        docs = render_chart(
+            values={
+                "airflowConfigAnnotations": {"key": "value", "key-two": "value-two"},
+            },
+            show_only=["templates/configmaps/configmap.yaml"],
+        )
+
+        annotations = jmespath.search("metadata.annotations", docs[0])
+        assert "value" == annotations.get("key")
+        assert "value-two" == annotations.get("key-two")
diff --git a/chart/values.schema.json b/chart/values.schema.json
index e6a29ba6412..ab85567a9ea 100644
--- a/chart/values.schema.json
+++ b/chart/values.schema.json
@@ -166,6 +166,10 @@
             "description": "Extra annotations to apply to all Airflow pods.",
             "type": "object"
         },
+        "airflowConfigAnnotations": {
+            "description": "Extra annotations to apply to the main Airflow configmap.",
+            "type": "object"
+        },
         "rbac": {
             "description": "Enable RBAC (default on most clusters these days).",
             "type": "object",
diff --git a/chart/values.yaml b/chart/values.yaml
index c2b53e78187..2171392a834 100644
--- a/chart/values.yaml
+++ b/chart/values.yaml
@@ -148,6 +148,10 @@ networkPolicies:
 # Airflow pods
 airflowPodAnnotations: {}
 
+# Extra annotations to apply to
+# main Airflow configmap
+airflowConfigAnnotations: {}
+
 # Enable RBAC (default on most clusters these days)
 rbac:
   # Specifies whether RBAC resources should be created
diff --git a/docs/helm-chart/parameters-ref.rst b/docs/helm-chart/parameters-ref.rst
index 9370bebe328..889104901a5 100644
--- a/docs/helm-chart/parameters-ref.rst
+++ b/docs/helm-chart/parameters-ref.rst
@@ -75,6 +75,12 @@ The following tables lists the configurable parameters of the Airflow chart and
    * - ``rbac.create``
      - Deploy pods with Kubernetes RBAC enabled
      - ``true``
+   * - ``airflowPodAnnotations``
+     - Extra annotations to apply to all Airflow pods.
+     - ``{}``
+   * - ``airflowConfigAnnotations``
+     - Extra annotations to apply to the main Airflow configmap.
+     - ``{}``
    * - ``executor``
      - Airflow executor (eg SequentialExecutor, LocalExecutor, CeleryExecutor, KubernetesExecutor)
      - ``1``