add env support for migratedatabase job (apache#42345)

* add env support for migratedatabase job * add test case for env config * update schema json for migrateDatabaseJob * fix ci failures * fix pre-commit ci for json schema
astronomer · Sep 25, 2024 · b9b7bfc · b9b7bfc
1 parent 4b08a0f
commit b9b7bfc
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 0 deletions.
diff --git a/chart/templates/jobs/migrate-database-job.yaml b/chart/templates/jobs/migrate-database-job.yaml
@@ -117,6 +117,9 @@ spec:
             - name: PYTHONUNBUFFERED
               value: "1"
           {{- include "standard_airflow_environment" . | indent 10 }}
+          {{- if .Values.migrateDatabaseJob.env }}
+              {{- tpl (toYaml .Values.migrateDatabaseJob.env) $ | nindent 12 }}
+          {{- end }}
           resources: {{- toYaml .Values.migrateDatabaseJob.resources | nindent 12 }}
           volumeMounts:
             {{- include "airflow_config_mount" . | nindent 12 }}

diff --git a/chart/values.schema.json b/chart/values.schema.json
@@ -4651,6 +4651,16 @@
                         "null"
                     ],
                     "default": 300
+                },
+                "env": {
+                    "description": "Add additional env vars to migrate database job.",
+                    "items": {
+                        "$ref": "#/definitions/io.k8s.api.core.v1.EnvVar"
+                    },
+                    "type": "array",
+                    "default": [],
+                    "x-kubernetes-patch-merge-key": "name",
+                    "x-kubernetes-patch-strategy": "merge"
                 }
             }
         },

diff --git a/chart/values.yaml b/chart/values.yaml
@@ -1236,6 +1236,7 @@ migrateDatabaseJob:
   # Disable this if you are using ArgoCD for example
   useHelmHooks: true
   applyCustomEnv: true
+  env: []
 
 # rpcServer support is experimental / dev purpose only and will later be renamed
 _rpcServer:

diff --git a/helm_tests/airflow_aux/test_migrate_database_job.py b/helm_tests/airflow_aux/test_migrate_database_job.py
@@ -455,3 +455,15 @@ def test_overridden_automount_service_account_token(self):
             show_only=["templates/jobs/migrate-database-job-serviceaccount.yaml"],
         )
         assert jmespath.search("automountServiceAccountToken", docs[0]) is False
+
+    def test_should_add_component_specific_env(self):
+        env = {"name": "test_env_key", "value": "test_env_value"}
+        docs = render_chart(
+            values={
+                "migrateDatabaseJob": {
+                    "env": [env],
+                },
+            },
+            show_only=["templates/jobs/migrate-database-job.yaml"],
+        )
+        assert env in jmespath.search("spec.template.spec.containers[0].env", docs[0])