From 31f1924986fa7655725047bf295a08354192aeb9 Mon Sep 17 00:00:00 2001
From: Ash Berlin-Taylor <ash@apache.org>
Date: Thu, 18 Nov 2021 11:23:27 +0000
Subject: [PATCH 1/2] Use built-in check-migrations command for Airflow>=2 in
 helm chart

---
 chart/templates/_helpers.yaml | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/chart/templates/_helpers.yaml b/chart/templates/_helpers.yaml
index 20cef7880e357..e2981a66cbb09 100644
--- a/chart/templates/_helpers.yaml
+++ b/chart/templates/_helpers.yaml
@@ -542,7 +542,11 @@ Create the name of the cleanup service account to use
 {{- end -}}
 
 {{ define "wait-for-migrations-command" }}
-  {{/* From Airflow 2.0.0 this can become [airflow, db, check-migrations] */}}
+  {{- if semverCompare ">=2.0.0" .Values.airflowVersion }}
+  - airflow
+  - db
+  - check-migrations
+  {{- else }}
   - python
   - -c
   - |
@@ -581,6 +585,7 @@ Create the name of the cleanup service account to use
                 ticker += 1
                 time.sleep(1)
                 logging.info('Waiting for migrations... %s second(s)', ticker)
+  {{- end }}
 {{- end }}
 
 {{ define "registry_docker_config" -}}

From 7481671b095c89ccbcdc711782446269a985f72e Mon Sep 17 00:00:00 2001
From: Ash Berlin-Taylor <ash@apache.org>
Date: Thu, 18 Nov 2021 13:11:25 +0000
Subject: [PATCH 2/2] fixup! Use built-in check-migrations command for
 Airflow>=2 in helm chart

---
 chart/tests/test_webserver.py | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/chart/tests/test_webserver.py b/chart/tests/test_webserver.py
index 329e6117dbec9..821d3f2236ee8 100644
--- a/chart/tests/test_webserver.py
+++ b/chart/tests/test_webserver.py
@@ -137,6 +137,24 @@ def test_should_add_extra_containers(self):
             "image": "test-registry/test-repo:test-tag",
         } == jmespath.search("spec.template.spec.containers[-1]", docs[0])
 
+    @parameterized.expand(
+        [
+            ("2.0.0", ["airflow", "db", "check-migrations"]),
+            ("2.1.0", ["airflow", "db", "check-migrations"]),
+            ("1.10.2", ["python", "-c"]),
+        ],
+    )
+    def test_wait_for_migration_airflow_version(self, airflow_version, expected_arg):
+        docs = render_chart(
+            values={
+                "airflowVersion": airflow_version,
+            },
+            show_only=["templates/webserver/webserver-deployment.yaml"],
+        )
+        # Don't test the full string, just the length of the expect matches
+        actual = jmespath.search("spec.template.spec.initContainers[0].args", docs[0])
+        assert expected_arg == actual[: len(expected_arg)]
+
     def test_should_add_extra_init_containers(self):
         docs = render_chart(
             values={