Chart: enable mysql keda support for triggerer (#37365)

GitOrigin-RevId: 44a53005948a0e96b94714e71ce4aa2b5f1f9d6b
GoogleCloudPlatform · Nov 9, 2024 · b52657b · b52657b
1 parent 4324e94
commit b52657b
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 1 deletion.
diff --git a/chart/templates/triggerer/triggerer-kedaautoscaler.yaml b/chart/templates/triggerer/triggerer-kedaautoscaler.yaml
@@ -48,10 +48,22 @@ spec:
   advanced: {{- toYaml .Values.triggerer.keda.advanced | nindent 4 }}
   {{- end }}
   triggers:
+    {{- if eq .Values.data.metadataConnection.protocol "mysql" }}
+    - type: "mysql"
+      metadata:
+        queryValue: "1"
+        connectionStringFromEnv: KEDA_DB_CONN
+        query: {{ tpl .Values.triggerer.keda.query . | quote }}
+    {{- else }}
     - type: postgresql
       metadata:
         targetQueryValue: "1"
+        {{- if and .Values.pgbouncer.enabled (not .Values.triggerer.keda.usePgbouncer) }}
+        connectionFromEnv: KEDA_DB_CONN
+        {{- else }}
         connectionFromEnv: AIRFLOW_CONN_AIRFLOW_DB
+        {{- end }}
         query: {{ tpl .Values.triggerer.keda.query . | quote }}
+  {{- end }}
 {{- end }}
 {{- end }}
diff --git a/chart/templates/workers/worker-kedaautoscaler.yaml b/chart/templates/workers/worker-kedaautoscaler.yaml
@@ -64,5 +64,4 @@ spec:
         {{- end }}
         query: {{ tpl .Values.workers.keda.query . | quote }}
     {{- end }}
-
 {{- end }}
diff --git a/helm_tests/airflow_core/test_triggerer.py b/helm_tests/airflow_core/test_triggerer.py
@@ -738,3 +738,20 @@ def test_should_use_keda_query(self, query, expected_query):
             show_only=["templates/triggerer/triggerer-kedaautoscaler.yaml"],
         )
         assert expected_query == jmespath.search("spec.triggers[0].metadata.query", docs[0])
+
+    def test_mysql_db_backend_keda(self):
+        docs = render_chart(
+            values={
+                "data": {"metadataConnection": {"protocol": "mysql"}},
+                "triggerer": {
+                    "enabled": True,
+                    "keda": {"enabled": True},
+                },
+            },
+            show_only=["templates/triggerer/triggerer-kedaautoscaler.yaml"],
+        )
+        assert "1" == jmespath.search("spec.triggers[0].metadata.queryValue", docs[0])
+        assert jmespath.search("spec.triggers[0].metadata.targetQueryValue", docs[0]) is None
+
+        assert "KEDA_DB_CONN" == jmespath.search("spec.triggers[0].metadata.connectionStringFromEnv", docs[0])
+        assert jmespath.search("spec.triggers[0].metadata.connectionFromEnv", docs[0]) is None
diff --git a/helm_tests/airflow_core/test_worker.py b/helm_tests/airflow_core/test_worker.py
@@ -1030,6 +1030,22 @@ def test_should_use_keda_query(self, query, executor, expected_query):
         )
         assert expected_query == jmespath.search("spec.triggers[0].metadata.query", docs[0])
 
+    def test_mysql_db_backend_keda(self):
+        docs = render_chart(
+            values={
+                "data": {"metadataConnection": {"protocol": "mysql"}},
+                "workers": {
+                    "keda": {"enabled": True},
+                },
+            },
+            show_only=["templates/workers/worker-kedaautoscaler.yaml"],
+        )
+        assert "1" == jmespath.search("spec.triggers[0].metadata.queryValue", docs[0])
+        assert jmespath.search("spec.triggers[0].metadata.targetQueryValue", docs[0]) is None
+
+        assert "KEDA_DB_CONN" == jmespath.search("spec.triggers[0].metadata.connectionStringFromEnv", docs[0])
+        assert jmespath.search("spec.triggers[0].metadata.connectionFromEnv", docs[0]) is None
+
 
 class TestWorkerHPAAutoScaler:
     """Tests worker HPA auto scaler."""