Add support for custom command and args in jobs

chart/templates/jobs/create-user-job.yaml

@@ -83,27 +83,12 @@ spec:
         - name: create-user
           image: {{ template "airflow_image" . }}
           imagePullPolicy: {{ .Values.images.airflow.pullPolicy }}
-          args:
-            - "bash"
-            - "-c"
-            {{- if semverCompare ">=2.0.0" .Values.airflowVersion }}
-            - 'airflow users create "$@"'
-            {{- else }}
-            - 'airflow create_user "$@"'
-            {{- end }}
-            - --
-            - "-r"
-            - {{ .Values.webserver.defaultUser.role }}
-            - "-u"
-            - {{ .Values.webserver.defaultUser.username }}
-            - "-e"
-            - {{ .Values.webserver.defaultUser.email }}
-            - "-f"
-            - {{ .Values.webserver.defaultUser.firstName }}
-            - "-l"
-            - {{ .Values.webserver.defaultUser.lastName }}
-            - "-p"
-            - "{{ .Values.webserver.defaultUser.password }}"
+          {{- if .Values.createUserJob.command }}
+          command: {{ tpl (toYaml .Values.createUserJob.command) . | nindent 12 }}
+          {{- end }}
+          {{- if .Values.createUserJob.args }}
+          args: {{ tpl (toYaml .Values.createUserJob.args) . | nindent 12 }}
+          {{- end }}
           envFrom:
           {{- include "custom_airflow_environment_from" . | default "\n  []" | indent 10 }}
           env:

chart/templates/jobs/migrate-database-job.yaml

@@ -82,11 +82,11 @@ spec:
         - name: run-airflow-migrations
           image: {{ template "airflow_image_for_migrations" . }}
           imagePullPolicy: {{ .Values.images.airflow.pullPolicy }}
-          # Support running against 1.10.x and 2.x
-          {{- if semverCompare ">=2.0.0" .Values.airflowVersion }}
-          args: ["bash", "-c", "airflow db upgrade"]
-          {{- else }}
-          args: ["bash", "-c", "airflow upgradedb"]
+          {{- if .Values.migrateDatabaseJob.command }}
+          command: {{ tpl (toYaml .Values.migrateDatabaseJob.command) . | nindent 12 }}
+          {{- end }}
+          {{- if .Values.migrateDatabaseJob.args }}
+          args: {{ tpl (toYaml .Values.migrateDatabaseJob.args) . | nindent 12 }}
           {{- end }}
           envFrom:
           {{- include "custom_airflow_environment_from" . | default "\n  []" | indent 10 }}

chart/tests/test_create_user_job.py

@@ -18,6 +18,7 @@
 import unittest
 
 import jmespath
+from parameterized import parameterized
 
 from chart.tests.helm_template_generator import render_chart
 
@@ -168,3 +169,102 @@ def test_should_add_extra_volume_mounts(self):
         assert {"name": "foobar", "mountPath": "foo/bar"} == jmespath.search(
             "spec.template.spec.containers[0].volumeMounts[-1]", docs[0]
         )
+
+    @parameterized.expand(
+        [
+            ("1.10.14", "airflow create_user"),
+            ("2.0.2", "airflow users create"),
+        ],
+    )
+    def test_default_command_and_args_airflow_version(self, airflow_version, expected_arg):
+        docs = render_chart(
+            values={
+                "airflowVersion": airflow_version,
+            },
+            show_only=["templates/jobs/create-user-job.yaml"],
+        )
+
+        assert jmespath.search("spec.template.spec.containers[0].command", docs[0]) is None
+        assert [
+            "bash",
+            "-c",
+            f"exec \\\n{expected_arg} \"$@\"",
+            "--",
+            "-r",
+            "Admin",
+            "-u",
+            "admin",
+            "-e",
+            "[email protected]",
+            "-f",
+            "admin",
+            "-l",
+            "user",
+            "-p",
+            "admin",
+        ] == jmespath.search("spec.template.spec.containers[0].args", docs[0])
+
+    @parameterized.expand(
+        [
+            (None, None),
+            (None, ["custom", "args"]),
+            (["custom", "command"], None),
+            (["custom", "command"], ["custom", "args"]),
+        ]
+    )
+    def test_command_and_args_overrides(self, command, args):
+        docs = render_chart(
+            values={"createUserJob": {"command": command, "args": args}},
+            show_only=["templates/jobs/create-user-job.yaml"],
+        )
+
+        assert command == jmespath.search("spec.template.spec.containers[0].command", docs[0])
+        assert args == jmespath.search("spec.template.spec.containers[0].args", docs[0])
+
+    def test_command_and_args_overrides_are_templated(self):
+        docs = render_chart(
+            values={
+                "createUserJob": {"command": ["{{ .Release.Name }}"], "args": ["{{ .Release.Service }}"]}
+            },
+            show_only=["templates/jobs/create-user-job.yaml"],
+        )
+
+        assert ["RELEASE-NAME"] == jmespath.search("spec.template.spec.containers[0].command", docs[0])
+        assert ["Helm"] == jmespath.search("spec.template.spec.containers[0].args", docs[0])
+
+    def test_default_user_overrides(self):
+        docs = render_chart(
+            values={
+                "webserver": {
+                    "defaultUser": {
+                        "role": "SomeRole",
+                        "username": "jdoe",
+                        "email": "[email protected]",
+                        "firstName": "John",
+                        "lastName": "Doe",
+                        "password": "whereisjane?",
+                    }
+                }
+            },
+            show_only=["templates/jobs/create-user-job.yaml"],
+        )
+
+        assert jmespath.search("spec.template.spec.containers[0].command", docs[0]) is None
+        assert [
+            "bash",
+            "-c",
+            "exec \\\nairflow users create \"$@\"",
+            "--",
+            "-r",
+            "SomeRole",
+            "-u",
+            "jdoe",
+            "-e",
+            "[email protected]",
+            "-f",
+            "John",
+            "-l",
+            "Doe",
+            "-p",
+            "whereisjane?",
+        ] == jmespath.search("spec.template.spec.containers[0].args", docs[0])

chart/tests/test_migrate_database_job.py

@@ -17,6 +17,7 @@
 
 import jmespath
 import pytest
+from parameterized import parameterized
 
 from tests.helm_template_generator import render_chart
 
@@ -200,3 +201,52 @@ def test_should_add_extra_volume_mounts(self):
         assert {"name": "foobar", "mountPath": "foo/bar"} == jmespath.search(
             "spec.template.spec.containers[0].volumeMounts[-1]", docs[0]
         )
+
+    @parameterized.expand(
+        [
+            ("1.10.14", "airflow upgradedb"),
+            ("2.0.2", "airflow db upgrade"),
+        ],
+    )
+    def test_default_command_and_args_airflow_version(self, airflow_version, expected_arg):
+        docs = render_chart(
+            values={
+                "airflowVersion": airflow_version,
+            },
+            show_only=["templates/jobs/migrate-database-job.yaml"],
+        )
+
+        assert jmespath.search("spec.template.spec.containers[0].command", docs[0]) is None
+        assert [
+            "bash",
+            "-c",
+            f"exec \\\n{expected_arg}",
+        ] == jmespath.search("spec.template.spec.containers[0].args", docs[0])
+
+    @parameterized.expand(
+        [
+            (None, None),
+            (None, ["custom", "args"]),
+            (["custom", "command"], None),
+            (["custom", "command"], ["custom", "args"]),
+        ]
+    )
+    def test_command_and_args_overrides(self, command, args):
+        docs = render_chart(
+            values={"migrateDatabaseJob": {"command": command, "args": args}},
+            show_only=["templates/jobs/migrate-database-job.yaml"],
+        )
+
+        assert command == jmespath.search("spec.template.spec.containers[0].command", docs[0])
+        assert args == jmespath.search("spec.template.spec.containers[0].args", docs[0])
+
+    def test_command_and_args_overrides_are_templated(self):
+        docs = render_chart(
+            values={
+                "migrateDatabaseJob": {"command": ["{{ .Release.Name }}"], "args": ["{{ .Release.Service }}"]}
+            },
+            show_only=["templates/jobs/migrate-database-job.yaml"],
+        )
+
+        assert ["RELEASE-NAME"] == jmespath.search("spec.template.spec.containers[0].command", docs[0])
+        assert ["Helm"] == jmespath.search("spec.template.spec.containers[0].args", docs[0])

chart/values.schema.json

@@ -2003,6 +2003,45 @@
             "x-docsSection": "Jobs",
             "additionalProperties": false,
             "properties": {
+                "command": {
+                    "description": "Command to use when running create user job (templated).",
+                    "type": [
+                        "array",
+                        "null"
+                    ],
+                    "items": {
+                        "type": "string"
+                    },
+                    "default": null
+                },
+                "args": {
+                    "description": "Args to use when running create user job (templated).",
+                    "type": [
+                        "array",
+                        "null"
+                    ],
+                    "items": {
+                        "type": "string"
+                    },
+                    "default": [
+                        "bash",
+                        "-c",
+                        "exec \\\nairflow {{ semverCompare \">=2.0.0\" .Values.airflowVersion | ternary \"users create\" \"create_user\" }} \"$@\"",
+                        "--",
+                        "-r",
+                        "{{ .Values.webserver.defaultUser.role }}",
+                        "-u",
+                        "{{ .Values.webserver.defaultUser.username }}",
+                        "-e",
+                        "{{ .Values.webserver.defaultUser.email }}",
+                        "-f",
+                        "{{ .Values.webserver.defaultUser.firstName }}",
+                        "-l",
+                        "{{ .Values.webserver.defaultUser.lastName }}",
+                        "-p",
+                        "{{ .Values.webserver.defaultUser.password }}"
+                    ]
+                },
                 "annotations": {
                     "description": "Annotations to add to the create user job pod.",
                     "type": "object",
@@ -2138,6 +2177,32 @@
             "x-docsSection": "Jobs",
             "additionalProperties": false,
             "properties": {
+                "command": {
+                    "description": "Command to use when running migrate database job (templated).",
+                    "type": [
+                        "array",
+                        "null"
+                    ],
+                    "items": {
+                        "type": "string"
+                    },
+                    "default": null
+                },
+                "args": {
+                    "description": "Args to use when running migrate database job (templated).",
+                    "type": [
+                        "array",
+                        "null"
+                    ],
+                    "items": {
+                        "type": "string"
+                    },
+                    "default": [
+                        "bash",
+                        "-c",
+                        "exec \\\nairflow {{ semverCompare \">=2.0.0\" .Values.airflowVersion | ternary \"db upgrade\" \"upgradedb\" }}"
+                    ]
+                },
                 "annotations": {
                     "description": "Annotations to add to the migrate database job pod.",
                     "type": "object",

chart/values.yaml

@@ -662,6 +662,30 @@ scheduler:
 
 # Airflow create user job settings
 createUserJob:
+  # Command to use when running the create user job (templated).
+  command: ~
+  # Args to use when running the create user job (templated).
+  args:
+    - "bash"
+    - "-c"
+    # The format below is necessary to get `helm lint` happy
+    - |-
+      exec \
+      airflow {{ semverCompare ">=2.0.0" .Values.airflowVersion | ternary "users create" "create_user" }} "$@"
+    - --
+    - "-r"
+    - "{{ .Values.webserver.defaultUser.role }}"
+    - "-u"
+    - "{{ .Values.webserver.defaultUser.username }}"
+    - "-e"
+    - "{{ .Values.webserver.defaultUser.email }}"
+    - "-f"
+    - "{{ .Values.webserver.defaultUser.firstName }}"
+    - "-l"
+    - "{{ .Values.webserver.defaultUser.lastName }}"
+    - "-p"
+    - "{{ .Values.webserver.defaultUser.password }}"
+
   # Annotations on the create user job pod
   annotations: {}
   # jobAnnotations are annotations on the create user job
@@ -708,6 +732,17 @@ createUserJob:
 
 # Airflow database migration job settings
 migrateDatabaseJob:
+  # Command to use when running the migrate database job (templated).
+  command: ~
+  # Args to use when running the migrate database job (templated).
+  args:
+    - "bash"
+    - "-c"
+    # The format below is necessary to get `helm lint` happy
+    - |-
+      exec \
+      airflow {{ semverCompare ">=2.0.0" .Values.airflowVersion | ternary "db upgrade" "upgradedb" }}
+
   # Annotations on the database migration pod
   annotations: {}
   # jobAnnotations are annotations on the database migration job