Prometheus: Add interpolation of groupByKeys query parameters into expr (#86360)

---------

Co-authored-by: ismail simsek <ismailsimsek09@gmail.com>

packages/grafana-prometheus/src/dataquery.ts

@@ -45,4 +45,5 @@ export interface Prometheus extends common.DataQuery {
   range?: boolean;
   scopes?: ScopeSpec[];
   adhocFilters?: ScopeSpecFilter[];
+  groupByKeys?: string[];
 }

packages/grafana-prometheus/src/datasource.ts

@@ -377,6 +377,10 @@ export class PrometheusDatasource
       processedTarget.scopes = (request.scopes ?? []).map((scope) => scope.spec);
     }
 
+    if (config.featureToggles.groupByVariable) {
+      processedTarget.groupByKeys = request.groupByKeys;
+    }
+
     if (target.instant && target.range) {
       // We have query type "Both" selected
       // We should send separate queries with different refId

pkg/promlib/models/query.go

@@ -70,6 +70,9 @@ type PrometheusQueryProperties struct {
 
 	// Additional Ad-hoc filters that take precedence over Scope on conflict.
 	AdhocFilters []ScopeFilter `json:"adhocFilters,omitempty"`
+
+	// Group By parameters to apply to aggregate expressions in the query
+	GroupByKeys []string `json:"groupByKeys,omitempty"`
 }
 
 // ScopeSpec is a hand copy of the ScopeSpec struct from pkg/apis/scope/v0alpha1/types.go
@@ -234,7 +237,7 @@ func Parse(span trace.Span, query backend.DataQuery, dsScrapeInterval string, in
 			}()))
 		}
 
-		expr, err = ApplyQueryFilters(expr, scopeFilters, model.AdhocFilters)
+		expr, err = ApplyFiltersAndGroupBy(expr, scopeFilters, model.AdhocFilters, model.GroupByKeys)
 		if err != nil {
 			return nil, err
 		}

pkg/promlib/models/query.panel.schema.json

@@ -85,6 +85,13 @@
             ],
             "x-enum-description": {}
           },
+          "groupByKeys": {
+            "description": "Group By parameters to apply to aggregate expressions in the query",
+            "type": "array",
+            "items": {
+              "type": "string"
+            }
+          },
           "hide": {
             "description": "true if query is disabled (ie should not be returned to the dashboard)\nNOTE: this does not always imply that the query should not be executed since\nthe results from a hidden query may be used as the input to other queries (SSE etc)",
             "type": "boolean"

pkg/promlib/models/query.request.schema.json

@@ -95,6 +95,13 @@
             ],
             "x-enum-description": {}
           },
+          "groupByKeys": {
+            "description": "Group By parameters to apply to aggregate expressions in the query",
+            "type": "array",
+            "items": {
+              "type": "string"
+            }
+          },
           "hide": {
             "description": "true if query is disabled (ie should not be returned to the dashboard)\nNOTE: this does not always imply that the query should not be executed since\nthe results from a hidden query may be used as the input to other queries (SSE etc)",
             "type": "boolean"

pkg/promlib/models/query.types.json

@@ -8,7 +8,7 @@
     {
       "metadata": {
         "name": "default",
-        "resourceVersion": "1715777575561",
+        "resourceVersion": "1715781995240",
         "creationTimestamp": "2024-03-25T13:19:04Z"
       },
       "spec": {
@@ -69,6 +69,13 @@
               "type": "string",
               "x-enum-description": {}
             },
+            "groupByKeys": {
+              "description": "Group By parameters to apply to aggregate expressions in the query",
+              "items": {
+                "type": "string"
+              },
+              "type": "array"
+            },
             "instant": {
               "description": "Returns only the latest value that Prometheus has scraped for the requested time series",
               "type": "boolean"

pkg/promlib/models/scope.go

@@ -7,7 +7,7 @@ import (
 	"github.com/prometheus/prometheus/promql/parser"
 )
 
-func ApplyQueryFilters(rawExpr string, scopeFilters, adHocFilters []ScopeFilter) (string, error) {
+func ApplyFiltersAndGroupBy(rawExpr string, scopeFilters, adHocFilters []ScopeFilter, groupBy []string) (string, error) {
 	expr, err := parser.ParseExpr(rawExpr)
 	if err != nil {
 		return "", err
@@ -50,7 +50,17 @@ func ApplyQueryFilters(rawExpr string, scopeFilters, adHocFilters []ScopeFilter)
 			}
 
 			return nil
-
+		case *parser.AggregateExpr:
+			found := make(map[string]bool)
+			for _, lName := range v.Grouping {
+				found[lName] = true
+			}
+			for _, k := range groupBy {
+				if !found[k] {
+					v.Grouping = append(v.Grouping, k)
+				}
+			}
+			return nil
 		default:
 			return nil
 		}

pkg/promlib/models/scope_test.go

@@ -6,7 +6,7 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestApplyQueryFilters(t *testing.T) {
+func TestApplyQueryFiltersAndGroupBy_Filters(t *testing.T) {
 	tests := []struct {
 		name         string
 		query        string
@@ -92,7 +92,105 @@ func TestApplyQueryFilters(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			expr, err := ApplyQueryFilters(tt.query, tt.scopeFilters, tt.adhocFilters)
+			expr, err := ApplyFiltersAndGroupBy(tt.query, tt.scopeFilters, tt.adhocFilters, nil)
+
+			if tt.expectErr {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, expr, tt.expected)
+			}
+		})
+	}
+}
+
+func TestApplyQueryFiltersAndGroupBy_GroupBy(t *testing.T) {
+	tests := []struct {
+		name      string
+		query     string
+		groupBy   []string
+		expected  string
+		expectErr bool
+	}{
+		{
+			name:      "GroupBy with no aggregate expression",
+			groupBy:   []string{"job"},
+			query:     `http_requests_total`,
+			expected:  `http_requests_total`,
+			expectErr: false,
+		},
+		{
+			name:      "No GroupBy with aggregate expression",
+			query:     `sum by () (http_requests_total)`,
+			expected:  `sum(http_requests_total)`,
+			expectErr: false,
+		},
+		{
+			name:      "GroupBy with aggregate expression with no existing group by",
+			groupBy:   []string{"job"},
+			query:     `sum(http_requests_total)`,
+			expected:  `sum by (job) (http_requests_total)`,
+			expectErr: false,
+		},
+		{
+			name:      "GroupBy with aggregate expression with existing group by",
+			groupBy:   []string{"status"},
+			query:     `sum by (job) (http_requests_total)`,
+			expected:  `sum by (job, status) (http_requests_total)`,
+			expectErr: false,
+		},
+		{
+			name:      "GroupBy with aggregate expression with existing group by (already exists)",
+			groupBy:   []string{"job"},
+			query:     `sum by (job) (http_requests_total)`,
+			expected:  `sum by (job) (http_requests_total)`,
+			expectErr: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			expr, err := ApplyFiltersAndGroupBy(tt.query, nil, nil, tt.groupBy)
+
+			if tt.expectErr {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, expr, tt.expected)
+			}
+		})
+	}
+}
+
+func TestApplyQueryFiltersAndGroupBy(t *testing.T) {
+	tests := []struct {
+		name         string
+		query        string
+		adhocFilters []ScopeFilter
+		scopeFilters []ScopeFilter
+		groupby      []string
+		expected     string
+		expectErr    bool
+	}{
+
+		{
+			name:  "Adhoc filters with more complex expression",
+			query: `sum(capacity_bytes{job="prometheus"} + available_bytes{job="grafana"}) / 1024`,
+			adhocFilters: []ScopeFilter{
+				{Key: "job", Value: "alloy", Operator: FilterOperatorEquals},
+			},
+			scopeFilters: []ScopeFilter{
+				{Key: "vol", Value: "/", Operator: FilterOperatorEquals},
+			},
+			groupby:   []string{"job"},
+			expected:  `sum by (job) (capacity_bytes{job="alloy",vol="/"} + available_bytes{job="alloy",vol="/"}) / 1024`,
+			expectErr: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			expr, err := ApplyFiltersAndGroupBy(tt.query, tt.scopeFilters, tt.adhocFilters, tt.groupby)
 
 			if tt.expectErr {
 				require.Error(t, err)