From a762bd94f5d1bc2e0d4a3e0ca75d287c2d88587e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebasti=C3=A0=20Galm=C3=A9s?= <sebastia.galmes@teya.com>
Date: Fri, 25 Apr 2025 11:40:15 +0200
Subject: [PATCH] feat: support prometheus as backend for keda trigger

---
 chart/templates/_helpers.yaml                 |  2 +-
 .../secrets/metadata-connection-secret.yaml   |  6 ++-
 .../triggerer/triggerer-kedaautoscaler.yaml   |  9 +++-
 .../workers/worker-kedaautoscaler.yaml        |  9 +++-
 chart/values.schema.json                      | 34 ++++++++++++++
 chart/values.yaml                             | 14 ++++++
 .../tests/helm_tests/other/test_keda.py       | 44 +++++++++++++++++++
 7 files changed, 113 insertions(+), 5 deletions(-)

diff --git a/chart/templates/_helpers.yaml b/chart/templates/_helpers.yaml
index 4e9b9c787b5aa..ab5f4987bbef8 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 (or (eq .Values.data.metadataConnection.protocol "mysql") (and .Values.pgbouncer.enabled (not .Values.workers.keda.usePgbouncer))) }}
+  {{- if and .Values.workers.keda.enabled (not .Values.workers.keda.prometheus.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
index d304414cc18b0..e0380caf8e147 100644
--- a/chart/templates/secrets/metadata-connection-secret.yaml
+++ b/chart/templates/secrets/metadata-connection-secret.yaml
@@ -47,13 +47,15 @@ data:
   {{- with .Values.data.metadataConnection }}
   connection: {{ urlJoin (dict "scheme" .protocol "userinfo" (printf "%s:%s" (.user | urlquery) (.pass | urlquery) ) "host" (printf "%s:%s" $host $port) "path" (printf "/%s" $database) "query" $query) | b64enc | quote }}
   {{- end }}
-  {{- if and .Values.workers.keda.enabled .Values.pgbouncer.enabled (not .Values.workers.keda.usePgbouncer) }}
+  {{- if or (and .Values.workers.keda.enabled (not .Values.workers.keda.prometheus.enabled)) (and .Values.triggerer.keda.enabled (not .Values.triggerer.keda.prometheus.enabled)) }}
+  {{- if and .Values.pgbouncer.enabled (not .Values.workers.keda.usePgbouncer) }}
   {{- 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 (or .Values.workers.keda.enabled .Values.triggerer.keda.enabled) (eq .Values.data.metadataConnection.protocol "mysql") }}
+  {{- else if (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) | trimPrefix "//" | b64enc | quote }}
   {{- end }}
   {{- end }}
+  {{- end }}
 {{- end }}
diff --git a/chart/templates/triggerer/triggerer-kedaautoscaler.yaml b/chart/templates/triggerer/triggerer-kedaautoscaler.yaml
index 4eda2858f0272..0330f482bfcd3 100644
--- a/chart/templates/triggerer/triggerer-kedaautoscaler.yaml
+++ b/chart/templates/triggerer/triggerer-kedaautoscaler.yaml
@@ -49,7 +49,14 @@ spec:
   advanced: {{- toYaml .Values.triggerer.keda.advanced | nindent 4 }}
   {{- end }}
   triggers:
-    {{- if eq .Values.data.metadataConnection.protocol "mysql" }}
+    {{- if .Values.triggerer.keda.prometheus.enabled }}
+    - type: "prometheus"
+      metadata:
+        serverAddress: {{ .Values.triggerer.keda.prometheus.serverAddress | quote }}
+        metricName: "airflow-triggerer"
+        query: {{ tpl .Values.triggerer.keda.query . | quote }}
+        threshold: "1"
+    {{- else if eq .Values.data.metadataConnection.protocol "mysql" }}
     - type: "mysql"
       metadata:
         queryValue: "1"
diff --git a/chart/templates/workers/worker-kedaautoscaler.yaml b/chart/templates/workers/worker-kedaautoscaler.yaml
index ab794c899636b..2b350802ca121 100644
--- a/chart/templates/workers/worker-kedaautoscaler.yaml
+++ b/chart/templates/workers/worker-kedaautoscaler.yaml
@@ -48,7 +48,14 @@ spec:
   advanced: {{- toYaml .Values.workers.keda.advanced | nindent 4 }}
   {{- end }}
   triggers:
-    {{- if eq .Values.data.metadataConnection.protocol "mysql" }}
+    {{- if .Values.workers.keda.prometheus.enabled }}
+    - type: "prometheus"
+      metadata:
+        serverAddress: {{ .Values.workers.keda.prometheus.serverAddress | quote }}
+        metricName: "airflow-workers"
+        query: {{ tpl .Values.workers.keda.query . | quote }}
+        threshold: "1"
+    {{- else if eq .Values.data.metadataConnection.protocol "mysql" }}
     - type: "mysql"
       metadata:
         queryValue: "1"
diff --git a/chart/values.schema.json b/chart/values.schema.json
index 93036cac30e42..54d4c82d91b44 100644
--- a/chart/values.schema.json
+++ b/chart/values.schema.json
@@ -1720,6 +1720,23 @@
                             "description": "Weather to use PGBouncer to connect to the database or not when it is enabled. This configuration will be ignored if PGBouncer is not enabled.",
                             "type": "boolean",
                             "default": true
+                        },
+                        "prometheus": {
+                            "description": "Prometheus configuration",
+                            "type": "object",
+                            "additionalProperties": false,
+                            "properties": {
+                                "enabled": {
+                                    "description": "Use prometheus for Keda autoscaling",
+                                    "type": "boolean",
+                                    "default": false
+                                },
+                                "serverAddress": {
+                                    "description": "Address of Prometheus server",
+                                    "type": "string",
+                                    "default": ""
+                                }
+                            }
                         }
                     }
                 },
@@ -3586,6 +3603,23 @@
                             "description": "Whether to use PGBouncer to connect to the database or not when it is enabled. This configuration will be ignored if PGBouncer is not enabled.",
                             "type": "boolean",
                             "default": false
+                        },
+                        "prometheus": {
+                            "description": "Prometheus configuration",
+                            "type": "object",
+                            "additionalProperties": false,
+                            "properties": {
+                                "enabled": {
+                                    "description": "Use prometheus for Keda autoscaling",
+                                    "type": "boolean",
+                                    "default": false
+                                },
+                                "serverAddress": {
+                                    "description": "Address of Prometheus server",
+                                    "type": "string",
+                                    "default": ""
+                                }
+                            }
                         }
                     }
                 }
diff --git a/chart/values.yaml b/chart/values.yaml
index 9066cbbb7f2c7..951c78f716680 100644
--- a/chart/values.yaml
+++ b/chart/values.yaml
@@ -666,6 +666,13 @@ workers:
     # This configuration will be ignored if PGBouncer is not enabled
     usePgbouncer: true
 
+    # Prometheus configuration
+    prometheus:
+      # Use Prometheus for Keda autoscaling
+      enabled: false
+      # Address of Prometheus server
+      serverAddress: ""
+
   # Allow HPA (KEDA must be disabled).
   hpa:
     enabled: false
@@ -1813,6 +1820,13 @@ triggerer:
     # This configuration will be ignored if PGBouncer is not enabled
     usePgbouncer: false
 
+    # Prometheus configuration
+    prometheus:
+      # Use Prometheus for Keda autoscaling
+      enabled: false
+      # Address of Prometheus server
+      serverAddress: ""
+
 # Airflow Dag Processor Config
 dagProcessor:
   enabled: ~
diff --git a/helm-tests/tests/helm_tests/other/test_keda.py b/helm-tests/tests/helm_tests/other/test_keda.py
index 40b319a267848..a0a53de84ae8d 100644
--- a/helm-tests/tests/helm_tests/other/test_keda.py
+++ b/helm-tests/tests/helm_tests/other/test_keda.py
@@ -334,3 +334,47 @@ def test_mysql_keda_db_connection(self):
             "spec.triggers[0].metadata.connectionStringFromEnv", keda_autoscaler
         )
         assert autoscaler_connection_env_var == "KEDA_DB_CONN"
+
+    def test_prometheus_keda(self):
+        """Verify keda prometheus configuration."""
+
+        docs = render_chart(
+            values={
+                "workers": {
+                    "keda": {
+                        "enabled": True,
+                        "prometheus": {
+                            "enabled": True,
+                            "serverAddress": "http://prometheus:9090",
+                        },
+                    }
+                },
+                "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 "KEDA_DB_CONN" not in worker_container_env_vars
+
+        keda_autoscaler = jmespath.search("spec.triggers[0]", keda_autoscaler)
+        assert keda_autoscaler["type"] == "prometheus"
+
+        keda_autoscaler_metadata = jmespath.search("spec.triggers[0].metadata", keda_autoscaler)
+        assert "threshold" in keda_autoscaler_metadata
+        assert "serverAddress" in keda_autoscaler_metadata
+        assert "metricName" in keda_autoscaler_metadata
+        assert "query" in keda_autoscaler_metadata
+
+        secret_data = jmespath.search("data", metadata_connection_secret)
+        assert "connection" in secret_data.keys()
+        assert "kedaConnection" not in secret_data.keys()
