feat: Prometheus Metrics Provider Should Evaluate Prometheus Query Before Sending to Prometheus

docs/analysis/prometheus.md

@@ -39,6 +39,41 @@ you validate your [PromQL expression](https://prometheus.io/docs/prometheus/late
 
 See the [Analysis Overview page](../../features/analysis) for more details on the available options.
 
+## Query expression evaluation
+
+There are some cases in which you will want to be able to handle multiple different cases in the same template. For those cases, first the string in the `Query` field is evaluated as expression before the result then runs in prometheus.
+
+For example:
+
+```yaml
+apiVersion: argoproj.io/v1alpha1
+kind: AnalysisTemplate
+metadata:
+  name: success-rate
+spec:
+  args:
+  - name: service-name
+  - name: some-arg
+  metrics:
+  - name: success-rate
+    interval: 5m
+    # NOTE: prometheus queries return results in the form of a vector.
+    # So it is common to access the index 0 of the returned array to obtain the value
+    successCondition: result[0] >= 0.95
+    failureLimit: 3
+    provider:
+      prometheus:
+        address: http://prometheus.example.com:9090
+        # timeout is expressed in seconds
+        timeout: 40
+        headers:
+        - key: X-Scope-OrgID
+          value: tenant_a
+        # in this case the resulting query is query2
+        query: |
+          "{{ args.some-arg }}" == "not_some_arg_value" ? "query1" : "query2"
+```
+
 ## Range queries
 
 ```yaml

metricproviders/prometheus/prometheus.go

@@ -51,12 +51,23 @@ func (p *Provider) Type() string {
 func (p *Provider) GetMetadata(metric v1alpha1.Metric) map[string]string {
 	metricsMetadata := make(map[string]string)
 	if metric.Provider.Prometheus.Query != "" {
-		metricsMetadata[ResolvedPrometheusQuery] = metric.Provider.Prometheus.Query
+		evaluatedQuery, err := evaluate.EvalQuery(metric.Provider.Prometheus.Query)
+		if err != nil {
+			evaluatedQuery = metric.Provider.Prometheus.Query
+		}
+		metricsMetadata[ResolvedPrometheusQuery] = evaluatedQuery
 	}
 	return metricsMetadata
 }
 
 func (p *Provider) executeQuery(ctx context.Context, metric v1alpha1.Metric) (model.Value, v1.Warnings, error) {
+
+	evalutedQuery, err := evaluate.EvalQuery(metric.Provider.Prometheus.Query)
+	if err != nil {
+		// If err != nil we want the old behavior to just pass in the query
+		evalutedQuery = metric.Provider.Prometheus.Query
+	}
+
 	if metric.Provider.Prometheus.RangeQuery != nil {
 		start, err := evaluate.EvalTime(metric.Provider.Prometheus.RangeQuery.Start)
 		if err != nil {
@@ -70,13 +81,13 @@ func (p *Provider) executeQuery(ctx context.Context, metric v1alpha1.Metric) (mo
 		if err != nil {
 			return nil, nil, fmt.Errorf("failed to parse rangeQuery.step as duration: %w", err)
 		}
-		return p.api.QueryRange(ctx, metric.Provider.Prometheus.Query, v1.Range{
+		return p.api.QueryRange(ctx, evalutedQuery, v1.Range{
 			Start: start,
 			End:   end,
 			Step:  stepDuration,
 		})
 	} else {
-		return p.api.Query(ctx, metric.Provider.Prometheus.Query, time.Now())
+		return p.api.Query(ctx, evalutedQuery, time.Now())
 	}
 }
 

metricproviders/prometheus/prometheus_test.go

@@ -395,6 +395,27 @@ func TestGetStatusReturnsResolvedQuery(t *testing.T) {
 	assert.Equal(t, "resolved-query", metricsMetadata["ResolvedPrometheusQuery"])
 }
 
+func TestGetStatusReturnsResolvedQueryFromEvalQuery(t *testing.T) {
+
+	expectedEvaluatedQuery := `sum(rate(some_metric{filter1="filter1_value",filter2="filter2_value",filter3=~"filter3_value",filter4=~"filter4_value",filter5!~"filter5_value"}[5m])) by(some_value))`
+	query := `"some_arg" == "not_some_arg" ? "query1" : ( "some_arg" == "some_arg" ? 'sum(rate(some_metric{filter1="filter1_value",filter2="filter2_value",filter3=~"filter3_value",filter4=~"filter4_value",filter5!~"filter5_value"}[5m])) by(some_value))' : "query2")`
+	e := log.Entry{}
+	mock := &mockAPI{}
+	metric := v1alpha1.Metric{
+		Name: "foo",
+		Provider: v1alpha1.MetricProvider{
+			Prometheus: &v1alpha1.PrometheusMetric{
+				Query: query,
+			},
+		},
+	}
+	p, err := NewPrometheusProvider(mock, e, metric)
+	metricsMetadata := p.GetMetadata(metric)
+	assert.NotNil(t, metricsMetadata)
+	assert.NoError(t, err)
+	assert.Equal(t, expectedEvaluatedQuery, metricsMetadata["ResolvedPrometheusQuery"])
+}
+
 func TestRunWithEvaluationError(t *testing.T) {
 	e := log.WithField("", "")
 	mock := &mockAPI{}

utils/evaluate/evaluate.go

@@ -90,6 +90,36 @@ func EvalTime(expression string) (time.Time, error) {
 	}
 }
 
+func EvalQuery(query string) (string, error) {
+	var err error
+
+	env := map[string]any{}
+
+	unwrapFileErr := func(e error) error {
+		if fileErr, ok := err.(*file.Error); ok {
+			e = errors.New(fileErr.Message)
+		}
+		return e
+	}
+
+	program, err := expr.Compile(query, expr.Env(env))
+	if err != nil {
+		return "", unwrapFileErr(err)
+	}
+
+	q, err := expr.Run(program, env)
+	if err != nil {
+		return "", unwrapFileErr(err)
+	}
+
+	switch val := q.(type) {
+	case string:
+		return val, nil
+	default:
+		return "", fmt.Errorf("expected string, but got %T", val)
+	}
+}
+
 // EvalCondition evaluates the condition with the resultValue as an input
 func EvalCondition(resultValue any, condition string) (bool, error) {
 	var err error

utils/evaluate/evaluate_test.go

@@ -348,3 +348,31 @@ func TestEvalTimeWithInvalidExpression(t *testing.T) {
 	assert.Equal(t, time.Time{}, status)
 	assert.Error(t, err)
 }
+
+func TestEvalQueryFromValidExpression(t *testing.T) {
+	expectedEvaluatedQuery := `sum(rate(some_metric{filter1="filter1_value",filter2="filter2_value",filter3=~"filter3_value",filter4=~"filter4_value",filter5!~"filter5_value"}[5m])) by(some_value))`
+	evaluatedQuery, err := EvalQuery(`"some_arg" == "some_arg" ? 'sum(rate(some_metric{filter1="filter1_value",filter2="filter2_value",filter3=~"filter3_value",filter4=~"filter4_value",filter5!~"filter5_value"}[5m])) by(some_value))' : "query2"`)
+	assert.Nil(t, err)
+	assert.Equal(t, expectedEvaluatedQuery, evaluatedQuery)
+}
+
+func TestEvalQueryFromValidButNestedExpression(t *testing.T) {
+	expectedEvaluatedQuery := `sum(rate(some_metric{filter1="filter1_value",filter2="filter2_value",filter3=~"filter3_value",filter4=~"filter4_value",filter5!~"filter5_value"}[5m])) by(some_value))`
+	evaluatedQuery, err := EvalQuery(`"some_arg" == "not_some_arg" ? "query1" : ( "some_arg" == "some_arg" ? 'sum(rate(some_metric{filter1="filter1_value",filter2="filter2_value",filter3=~"filter3_value",filter4=~"filter4_value",filter5!~"filter5_value"}[5m])) by(some_value))' : "query2")`)
+	assert.Nil(t, err)
+	assert.Equal(t, expectedEvaluatedQuery, evaluatedQuery)
+}
+
+func TestEvalQueryFromInvalidExpression(t *testing.T) {
+	expectedEvaluatedQuery := ""
+	evaluatedQuery, err := EvalQuery(`"some_arg" == "not_some_arg" ? "query1"`)
+	assert.Error(t, err)
+	assert.Equal(t, expectedEvaluatedQuery, evaluatedQuery)
+}
+
+func TestEvalQueryWithOrOperatorInEachOption(t *testing.T) {
+	expectedEvaluatedQuery := `old_query or sum(rate(some_metric{filter1="filter1_value",filter2="filter2_value",filter3=~"filter3_value",filter4=~"filter4_value",filter5!~"filter5_value"}[5m])) by(some_value))`
+	evaluatedQuery, err := EvalQuery(`"some_arg" == "not_some_arg" ? "old_query or query1" : ( "some_arg" == "some_arg" ? 'old_query or sum(rate(some_metric{filter1="filter1_value",filter2="filter2_value",filter3=~"filter3_value",filter4=~"filter4_value",filter5!~"filter5_value"}[5m])) by(some_value))' : "old_query or query2")`)
+	assert.Nil(t, err)
+	assert.Equal(t, expectedEvaluatedQuery, evaluatedQuery)
+}