From b71e8af9746b47d022f2724e1ba0ca628b8d7be6 Mon Sep 17 00:00:00 2001
From: Romain Forlot <romain.forlot@dailymotion.com>
Date: Mon, 11 Dec 2023 14:19:03 +0100
Subject: [PATCH] feat: support mysql backend as keda trigger

#36166

Signed-off-by: Romain Forlot <romain.forlot@dailymotion.com>
---
 chart/templates/_helpers.yaml                         |  2 +-
 .../templates/secrets/metadata-connection-secret.yaml |  4 ++++
 chart/templates/workers/worker-kedaautoscaler.yaml    | 11 ++++++++++-
 3 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/chart/templates/_helpers.yaml b/chart/templates/_helpers.yaml
index 719b6647b585f4..761527b8545843 100644
--- 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 (or (and .Values.workers.keda.enabled (eq .Values.data.metadataConnection.protocol "mysql")) (and .Values.workers.keda.enabled .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
index c25156d118e789..220a516682eceb 100644
--- 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
index f93484efd02922..95adb3e1267634 100644
--- 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 }}