feat: support prometheus as backend for keda trigger

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:

chart/templates/secrets/metadata-connection-secret.yaml

@@ -51,13 +51,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 }}

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"

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"

chart/values.schema.json

@@ -1832,6 +1832,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": ""
+                                }
+                            }
                         }
                     }
                 },
@@ -3781,6 +3798,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": ""
+                                }
+                            }
                         }
                     }
                 }

chart/values.yaml

@@ -743,6 +743,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 for Airflow Celery workers (KEDA must be disabled)
   hpa:
     enabled: false
@@ -2021,6 +2028,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: ~

helm-tests/tests/helm_tests/other/test_keda.py

@@ -337,3 +337,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()