From 3564b98fc2e19d4f339b87f30b894c0afdb2e7a1 Mon Sep 17 00:00:00 2001
From: Gregory Jones <greg@wequilt.com>
Date: Tue, 30 May 2023 19:11:49 -0400
Subject: [PATCH 1/4] Add missing waitForMigrations for workers

---
 chart/templates/workers/worker-deployment.yaml | 2 ++
 chart/values.schema.json                       | 5 +++++
 chart/values.yaml                              | 2 ++
 3 files changed, 9 insertions(+)

diff --git a/chart/templates/workers/worker-deployment.yaml b/chart/templates/workers/worker-deployment.yaml
index 211949ca12cee..b222ab26156c0 100644
--- a/chart/templates/workers/worker-deployment.yaml
+++ b/chart/templates/workers/worker-deployment.yaml
@@ -141,6 +141,7 @@ spec:
             - name: logs
               mountPath: {{ template "airflow_logs" . }}
         {{- end }}
+        {{- if .Values.workers.waitForMigrations.enabled }}
         - name: wait-for-airflow-migrations
           resources: {{- toYaml .Values.workers.resources | nindent 12 }}
           image: {{ template "airflow_image_for_migrations" . }}
@@ -164,6 +165,7 @@ spec:
             {{- if .Values.workers.waitForMigrations.env }}
               {{- tpl (toYaml .Values.workers.waitForMigrations.env) $ | nindent 12 }}
             {{- end }}
+        {{- end }}
         {{- if and (.Values.dags.gitSync.enabled) (not .Values.dags.persistence.enabled) }}
           {{- include "git_sync_container" (dict "Values" .Values "is_init" "true") | nindent 8 }}
         {{- end }}
diff --git a/chart/values.schema.json b/chart/values.schema.json
index 59194b8ac4b0b..005948802346a 100644
--- a/chart/values.schema.json
+++ b/chart/values.schema.json
@@ -1645,6 +1645,11 @@
                     "type": "object",
                     "additionalProperties": false,
                     "properties": {
+                        "enabled": {
+                            "description": "Enable wait-for-airflow-migrations init container.",
+                            "type": "boolean",
+                            "default": true
+                        },
                         "env": {
                             "description": "Add additional env vars to wait-for-airflow-migrations init container.",
                             "type": "array",
diff --git a/chart/values.yaml b/chart/values.yaml
index e888c3520c263..1bbcb9fc7d476 100644
--- a/chart/values.yaml
+++ b/chart/values.yaml
@@ -646,6 +646,8 @@ workers:
     #   memory: 128Mi
 
   waitForMigrations:
+    # Whether to create init container to wait for db migrations
+    enabled: true
     env: []
 
   env: []

From 128c4fa61293b9e209b7b7e910ae0353d9624267 Mon Sep 17 00:00:00 2001
From: Gregory Jones <grjones@gmail.com>
Date: Tue, 30 May 2023 21:13:34 -0400
Subject: [PATCH 2/4] Add worker test waitForMigrations.enabled

---
 tests/charts/airflow_core/test_worker.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/tests/charts/airflow_core/test_worker.py b/tests/charts/airflow_core/test_worker.py
index 33870fa459876..da0af97b7779a 100644
--- a/tests/charts/airflow_core/test_worker.py
+++ b/tests/charts/airflow_core/test_worker.py
@@ -82,6 +82,18 @@ def test_should_add_extra_containers(self):
             "image": "test-registry/test-repo:test-tag",
         } == jmespath.search("spec.template.spec.containers[-1]", docs[0])
 
+    def test_disable_wait_for_migration(self):
+        docs = render_chart(
+            values={
+                "workers": {
+                    "waitForMigrations": {"enabled": False},
+                },
+            },
+            show_only=["templates/workers/worker-deployment.yaml"],
+        )
+        actual = jmespath.search("spec.template.spec.initContainers", docs[0])
+        assert actual is None
+
     def test_should_add_extra_init_containers(self):
         docs = render_chart(
             values={

From 3ef0e1de75008022a225a8aa0ed22ece4828d1ce Mon Sep 17 00:00:00 2001
From: Gregory Jones <greg@wequilt.com>
Date: Wed, 31 May 2023 22:09:06 -0400
Subject: [PATCH 3/4] Narrow test scope for waitForMigrations container only

---
 tests/charts/airflow_core/test_worker.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/charts/airflow_core/test_worker.py b/tests/charts/airflow_core/test_worker.py
index da0af97b7779a..bf3ed2881720c 100644
--- a/tests/charts/airflow_core/test_worker.py
+++ b/tests/charts/airflow_core/test_worker.py
@@ -91,7 +91,7 @@ def test_disable_wait_for_migration(self):
             },
             show_only=["templates/workers/worker-deployment.yaml"],
         )
-        actual = jmespath.search("spec.template.spec.initContainers", docs[0])
+        actual = jmespath.search("spec.template.spec.initContainers[?name=='wait-for-airflow-migrations']", docs[0])
         assert actual is None
 
     def test_should_add_extra_init_containers(self):

From 1a4ddd3fd40965172c9db4d7493a6a10fa1566dc Mon Sep 17 00:00:00 2001
From: Hussein Awala <hussein@awala.fr>
Date: Thu, 1 Jun 2023 13:13:32 +0200
Subject: [PATCH 4/4] Update tests/charts/airflow_core/test_worker.py

---
 tests/charts/airflow_core/test_worker.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/charts/airflow_core/test_worker.py b/tests/charts/airflow_core/test_worker.py
index bf3ed2881720c..23075f45f683b 100644
--- a/tests/charts/airflow_core/test_worker.py
+++ b/tests/charts/airflow_core/test_worker.py
@@ -91,7 +91,9 @@ def test_disable_wait_for_migration(self):
             },
             show_only=["templates/workers/worker-deployment.yaml"],
         )
-        actual = jmespath.search("spec.template.spec.initContainers[?name=='wait-for-airflow-migrations']", docs[0])
+        actual = jmespath.search(
+            "spec.template.spec.initContainers[?name=='wait-for-airflow-migrations']", docs[0]
+        )
         assert actual is None
 
     def test_should_add_extra_init_containers(self):
