feat: support mysql backend as keda trigger (apache#36167)

flacode · Jan 22, 2024 · 5f5c99b · 5f5c99b
1 parent 4c1477f
commit 5f5c99b
Show file tree

Hide file tree

Showing 4 changed files with 51 additions and 2 deletions.
diff --git a/chart/templates/_helpers.yaml b/chart/templates/_helpers.yaml
@@ -84,7 +84,7 @@ If release name contains chart name it will be used as a full name.
         name: {{ template "airflow_metadata_secret" . }}
         key: connection
   {{- end }}
-  {{- if and .Values.workers.keda.enabled .Values.pgbouncer.enabled (not .Values.workers.keda.usePgbouncer) }}
+  {{- if and .Values.workers.keda.enabled (or (eq .Values.data.metadataConnection.protocol "mysql") (and .Values.pgbouncer.enabled (not .Values.workers.keda.usePgbouncer))) }}
   - name: KEDA_DB_CONN
     valueFrom:
       secretKeyRef:

diff --git a/chart/templates/secrets/metadata-connection-secret.yaml b/chart/templates/secrets/metadata-connection-secret.yaml
@@ -51,5 +51,9 @@ data:
   {{- with .Values.data.metadataConnection }}
   kedaConnection: {{ urlJoin (dict "scheme" .protocol "userinfo" (printf "%s:%s" (.user | urlquery) (.pass | urlquery) ) "host" (printf "%s:%s" $metadataHost $metadataPort) "path" (printf "/%s" $metadataDatabase) "query" $query) | b64enc | quote }}
   {{- end }}
+  {{- else if and .Values.workers.keda.enabled (eq .Values.data.metadataConnection.protocol "mysql") }}
+  {{- with .Values.data.metadataConnection }}
+  kedaConnection: {{ urlJoin (dict "userinfo" (printf "%s:%s" (.user | urlquery) (.pass | urlquery) ) "host" (printf "tcp(%s:%s)" $metadataHost $metadataPort) "path" (printf "/%s" $metadataDatabase) "query" $query) | b64enc | quote }}
+  {{- end }}
   {{- end }}
 {{- end }}
diff --git a/chart/templates/workers/worker-kedaautoscaler.yaml b/chart/templates/workers/worker-kedaautoscaler.yaml
@@ -47,7 +47,14 @@ spec:
   advanced: {{- toYaml .Values.workers.keda.advanced | nindent 4 }}
   {{- end }}
   triggers:
-    - type: postgresql
+    {{- if eq .Values.data.metadataConnection.protocol "mysql" }}
+    - type: "mysql"
+      metadata:
+        queryValue: "1"
+        connectionStringFromEnv: KEDA_DB_CONN
+        query: {{ tpl .Values.workers.keda.query . | quote }}
+    {{- else }}
+    - type: "postgresql"
       metadata:
         targetQueryValue: "1"
         {{- if and .Values.pgbouncer.enabled (not .Values.workers.keda.usePgbouncer) }}
@@ -56,4 +63,6 @@ spec:
         connectionFromEnv: AIRFLOW_CONN_AIRFLOW_DB
         {{- end }}
         query: {{ tpl .Values.workers.keda.query . | quote }}
+    {{- end }}
+
 {{- end }}
diff --git a/helm_tests/other/test_keda.py b/helm_tests/other/test_keda.py
@@ -275,3 +275,39 @@ def test_default_keda_db_connection_pgbouncer_enabled_usePgbouncer_false(self):
             "spec.triggers[0].metadata.connectionFromEnv", keda_autoscaler
         )
         assert autoscaler_connection_env_var == "KEDA_DB_CONN"
+
+    def test_mysql_keda_db_connection(self):
+        """Verify keda db connection when pgbouncer is enabled."""
+        docs = render_chart(
+            values={
+                "data": {"metadataConnection": {"protocol": "mysql", "port": 3306}},
+                "workers": {"keda": {"enabled": True}},
+                "executor": "CeleryExecutor",
+            },
+            show_only=[
+                "templates/workers/worker-deployment.yaml",
+                "templates/workers/worker-kedaautoscaler.yaml",
+                "templates/secrets/metadata-connection-secret.yaml",
+            ],
+        )
+        worker_deployment = docs[0]
+        keda_autoscaler = docs[1]
+        metadata_connection_secret = docs[2]
+
+        worker_container_env_vars = jmespath.search(
+            "spec.template.spec.containers[?name=='worker'].env[].name", worker_deployment
+        )
+        assert "AIRFLOW_CONN_AIRFLOW_DB" in worker_container_env_vars
+        assert "KEDA_DB_CONN" in worker_container_env_vars
+
+        keda_autoscaler_metadata = jmespath.search("spec.triggers[0].metadata", keda_autoscaler)
+        assert "queryValue" in keda_autoscaler_metadata
+
+        secret_data = jmespath.search("data", metadata_connection_secret)
+        assert "connection" in secret_data.keys()
+        assert "kedaConnection" in secret_data.keys()
+
+        autoscaler_connection_env_var = jmespath.search(
+            "spec.triggers[0].metadata.connectionStringFromEnv", keda_autoscaler
+        )
+        assert autoscaler_connection_env_var == "KEDA_DB_CONN"