From a762bd94f5d1bc2e0d4a3e0ca75d287c2d88587e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sebasti=C3=A0=20Galm=C3=A9s?= 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()