Allow migration jobs and init containers to be optional (#22195)

apache · Mar 22, 2022 · d78e086 · d78e086
1 parent d7dbfb7
commit d78e086
Show file tree

Hide file tree

Showing 10 changed files with 115 additions and 0 deletions.
diff --git a/chart/templates/jobs/migrate-database-job.yaml b/chart/templates/jobs/migrate-database-job.yaml
@@ -18,6 +18,7 @@
 ################################
 ## Airflow Run Migrations
 #################################
+{{- if .Values.migrateDatabaseJob.enabled }}
 {{- $nodeSelector := or .Values.migrateDatabaseJob.nodeSelector .Values.nodeSelector }}
 {{- $affinity := or .Values.migrateDatabaseJob.affinity .Values.affinity }}
 {{- $tolerations := or .Values.migrateDatabaseJob.tolerations .Values.tolerations }}
@@ -113,3 +114,4 @@ spec:
 {{- if .Values.migrateDatabaseJob.extraVolumes }}
 {{ toYaml .Values.migrateDatabaseJob.extraVolumes | nindent 8 }}
 {{- end }}
+{{- end }}
diff --git a/chart/templates/scheduler/scheduler-deployment.yaml b/chart/templates/scheduler/scheduler-deployment.yaml
@@ -119,6 +119,7 @@ spec:
         - name: {{ template "registry_secret" . }}
       {{- end }}
       initContainers:
+        {{- if .Values.scheduler.waitForMigrations.enabled }}
         - name: wait-for-airflow-migrations
           resources:
 {{ toYaml .Values.scheduler.resources | indent 12 }}
@@ -136,6 +137,7 @@ spec:
           env:
           {{- include "custom_airflow_environment" . | indent 10 }}
           {{- include "standard_airflow_environment" . | indent 10 }}
+        {{- end }}
         {{- if .Values.dags.gitSync.enabled }}
         {{- include "git_sync_container" (dict "Values" .Values "is_init" "true") | nindent 8 }}
         {{- end }}

diff --git a/chart/templates/triggerer/triggerer-deployment.yaml b/chart/templates/triggerer/triggerer-deployment.yaml
@@ -102,6 +102,7 @@ spec:
         - name: {{ template "registry_secret" . }}
       {{- end }}
       initContainers:
+        {{- if .Values.triggerer.waitForMigrations.enabled }}
         - name: wait-for-airflow-migrations
           resources:
             {{- toYaml .Values.triggerer.resources | nindent 12 }}
@@ -119,6 +120,7 @@ spec:
           env:
           {{- include "custom_airflow_environment" . | nindent 10 }}
           {{- include "standard_airflow_environment" . | nindent 10 }}
+        {{- 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/templates/webserver/webserver-deployment.yaml b/chart/templates/webserver/webserver-deployment.yaml
@@ -114,6 +114,7 @@ spec:
         - name: {{ template "registry_secret" . }}
       {{- end }}
       initContainers:
+        {{- if .Values.webserver.waitForMigrations.enabled }}
         - name: wait-for-airflow-migrations
           resources:
 {{ toYaml .Values.webserver.resources | indent 12 }}
@@ -131,6 +132,7 @@ spec:
           env:
           {{- include "custom_airflow_environment" . | indent 10 }}
           {{- include "standard_airflow_environment" . | indent 10 }}
+        {{- end }}
         {{- if and (.Values.dags.gitSync.enabled) (not .Values.dags.persistence.enabled) (semverCompare "<2.0.0" .Values.airflowVersion) }}
         {{- 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
@@ -1786,6 +1786,18 @@
                             "fsGroup": 0
                         }
                     ]
+                },
+                "waitForMigrations": {
+                    "description": "wait-for-airflow-migrations init container.",
+                    "type": "object",
+                    "additionalProperties": false,
+                    "properties": {
+                        "enabled": {
+                            "description": "Enable wait-for-airflow-migrations init container.",
+                            "type": "boolean",
+                            "default": true
+                        }
+                    }
                 }
             }
         },
@@ -2023,6 +2035,18 @@
                             "fsGroup": 0
                         }
                     ]
+                },
+                "waitForMigrations": {
+                    "description": "wait-for-airflow-migrations init container.",
+                    "type": "object",
+                    "additionalProperties": false,
+                    "properties": {
+                        "enabled": {
+                            "description": "Enable wait-for-airflow-migrations init container.",
+                            "type": "boolean",
+                            "default": true
+                        }
+                    }
                 }
             }
         },
@@ -2206,6 +2230,11 @@
             "x-docsSection": "Jobs",
             "additionalProperties": false,
             "properties": {
+                "enabled": {
+                    "description": "Enable migrate database job.",
+                    "type": "boolean",
+                    "default": true
+                },
                 "command": {
                     "description": "Command to use when running migrate database job (templated).",
                     "type": [
@@ -2807,6 +2836,18 @@
                     "additionalProperties": {
                         "type": "string"
                     }
+                },
+                "waitForMigrations": {
+                    "description": "wait-for-airflow-migrations init container.",
+                    "type": "object",
+                    "additionalProperties": false,
+                    "properties": {
+                        "enabled": {
+                            "description": "Enable wait-for-airflow-migrations init container.",
+                            "type": "boolean",
+                            "default": true
+                        }
+                    }
                 }
             }
         },

diff --git a/chart/values.yaml b/chart/values.yaml
@@ -661,6 +661,10 @@ scheduler:
     #   cpu: 100m
     #   memory: 128Mi
 
+  waitForMigrations:
+    # Whether to create init container to wait for db migrations
+    enabled: true
+
 # Airflow create user job settings
 createUserJob:
   # Command to use when running the create user job (templated).
@@ -733,6 +737,7 @@ createUserJob:
 
 # Airflow database migration job settings
 migrateDatabaseJob:
+  enabled: true
   # Command to use when running the migrate database job (templated).
   command: ~
   # Args to use when running the migrate database job (templated).
@@ -928,6 +933,10 @@ webserver:
 
   podAnnotations: {}
 
+  waitForMigrations:
+    # Whether to create init container to wait for db migrations
+    enabled: true
+
 # Airflow Triggerer Config
 triggerer:
   enabled: true
@@ -1013,6 +1022,10 @@ triggerer:
 
   podAnnotations: {}
 
+  waitForMigrations:
+    # Whether to create init container to wait for db migrations
+    enabled: true
+
 # Flower settings
 flower:
   # Enable flower.

diff --git a/tests/charts/test_migrate_database_job.py b/tests/charts/test_migrate_database_job.py
@@ -29,6 +29,23 @@ def test_should_run_by_default(self):
         assert "run-airflow-migrations" == jmespath.search("spec.template.spec.containers[0].name", docs[0])
         assert 50000 == jmespath.search("spec.template.spec.securityContext.runAsUser", docs[0])
 
+    @pytest.mark.parametrize(
+        "migrate_database_job_enabled,created",
+        [
+            (False, False),
+            (True, True),
+        ],
+    )
+    def test_enable_migrate_database_job(self, migrate_database_job_enabled, created):
+        docs = render_chart(
+            values={
+                "migrateDatabaseJob": {"enabled": migrate_database_job_enabled},
+            },
+            show_only=["templates/jobs/migrate-database-job.yaml"],
+        )
+
+        assert bool(docs) is created
+
     def test_should_support_annotations(self):
         docs = render_chart(
             values={"migrateDatabaseJob": {"annotations": {"foo": "bar"}, "jobAnnotations": {"fiz": "fuz"}}},

diff --git a/tests/charts/test_scheduler.py b/tests/charts/test_scheduler.py
@@ -67,6 +67,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={
+                "scheduler": {
+                    "waitForMigrations": {"enabled": False},
+                },
+            },
+            show_only=["templates/scheduler/scheduler-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={

diff --git a/tests/charts/test_triggerer.py b/tests/charts/test_triggerer.py
@@ -51,6 +51,18 @@ def test_can_be_disabled(self):
 
         assert 0 == len(docs)
 
+    def test_disable_wait_for_migration(self):
+        docs = render_chart(
+            values={
+                "triggerer": {
+                    "waitForMigrations": {"enabled": False},
+                },
+            },
+            show_only=["templates/triggerer/triggerer-deployment.yaml"],
+        )
+        actual = jmespath.search("spec.template.spec.initContainers", docs[0])
+        assert actual is None
+
     def test_should_add_extra_containers(self):
         docs = render_chart(
             values={

diff --git a/tests/charts/test_webserver.py b/tests/charts/test_webserver.py
@@ -155,6 +155,18 @@ def test_wait_for_migration_airflow_version(self, airflow_version, expected_arg)
         actual = jmespath.search("spec.template.spec.initContainers[0].args", docs[0])
         assert expected_arg == actual[: len(expected_arg)]
 
+    def test_disable_wait_for_migration(self):
+        docs = render_chart(
+            values={
+                "webserver": {
+                    "waitForMigrations": {"enabled": False},
+                },
+            },
+            show_only=["templates/webserver/webserver-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={