query rejection - add API type

Signed-off-by: Erlan Zholdubai uulu <[email protected]>

Author: erlan-z

docs/configuration/config-file-reference.md

@@ -5367,6 +5367,10 @@ limits:
 ### `QueryAttribute`
 
 ```yaml
+# API type for the query. Should be one of the query, query_range, series,
+# labels, label_values. If not set, it won't be checked.
+[api_type: <string> | default = ""]
+
 # Regex that the query string (or at least one of the matchers in metadata
 # query) should match. If not set, it won't be checked.
 [regex: <string> | default = ""]

integration/query_frontend_test.go

@@ -631,7 +631,8 @@ func TestQueryFrontendQueryRejection(t *testing.T) {
     query_rejection:
       enabled: true
       query_attributes:
-        - regex: .*rate.*
+        - api_type: "query"
+          regex: .*rate.*
           query_step_limit:
            min: 6s
            max: 20m
@@ -659,13 +660,19 @@ func TestQueryFrontendQueryRejection(t *testing.T) {
 	require.NoError(t, err)
 
 	now := time.Now()
-	// We expect request to be rejected as it matches query_attribute of query_rejection (contains rate, contains dashboard header dash123). step limit is ignored for instant queries
+	// We expect request to be rejected, as it matches query_attribute of query_rejection (contains rate, contains dashboard header dash123). step limit is ignored for instant queries
 	// Query shouldn't be checked against attributes that is not provided in query_attribute config(time_window, time_range_limit, user_agent_regex, panel_id)
 	resp, body, err := c.QueryRaw(`min_over_time( rate(http_requests_total[5m])[30m:5s] )`, now, map[string]string{"X-Dashboard-Uid": "dash123", "User-Agent": "grafana-agent/v0.19.0"})
 	require.NoError(t, err)
 	require.Equal(t, http.StatusUnprocessableEntity, resp.StatusCode)
 	require.Contains(t, string(body), tripperware.QueryRejectErrorMessage)
 
+	// We expect request not to be rejected, as it doesn't match api_type
+	resp, body, err = c.QueryRangeRaw(`min_over_time( rate(http_requests_total[5m])[30m:5s] )`, now.Add(-11*time.Hour), now.Add(-8*time.Hour), 25*time.Minute, map[string]string{"X-Dashboard-Uid": "dash123", "User-Agent": "grafana-agent/v0.19.0"})
+	require.NoError(t, err)
+	require.Equal(t, http.StatusOK, resp.StatusCode)
+	require.NotContains(t, string(body), tripperware.QueryRejectErrorMessage)
+
 	// update runtime config
 	newRuntimeConfig = []byte(`overrides:
   user-1:

pkg/querier/tripperware/query_attribute_matcher.go

@@ -58,7 +58,7 @@ func rejectQueryOrSetPriority(r *http.Request, now time.Time, lookbackDelta time
 
 	if queryReject := limits.QueryRejection(userStr); queryReject.Enabled && (op == "series" || op == "labels" || op == "label_values") {
 		for _, attribute := range queryReject.QueryAttributes {
-			if matchAttributeForMetadataQuery(attribute, r, now) {
+			if matchAttributeForMetadataQuery(attribute, op, r, now) {
 				rejectedQueriesPerTenant.WithLabelValues(op, userStr).Inc()
 				return httpgrpc.Errorf(http.StatusUnprocessableEntity, QueryRejectErrorMessage)
 			}
@@ -86,6 +86,9 @@ func getOperation(r *http.Request) string {
 }
 
 func matchAttributeForExpressionQuery(attribute validation.QueryAttribute, op string, r *http.Request, query string, now time.Time, minTime, maxTime int64) bool {
+	if attribute.ApiType != "" && attribute.ApiType != op {
+		return false
+	}
 	if attribute.Regex != "" && attribute.Regex != ".*" && attribute.Regex != ".+" {
 		if attribute.CompiledRegex != nil && !attribute.CompiledRegex.MatchString(query) {
 			return false
@@ -121,7 +124,10 @@ func matchAttributeForExpressionQuery(attribute validation.QueryAttribute, op st
 	return true
 }
 
-func matchAttributeForMetadataQuery(attribute validation.QueryAttribute, r *http.Request, now time.Time) bool {
+func matchAttributeForMetadataQuery(attribute validation.QueryAttribute, op string, r *http.Request, now time.Time) bool {
+	if attribute.ApiType != "" && attribute.ApiType != op {
+		return false
+	}
 	if err := r.ParseForm(); err != nil {
 		return false
 	}

pkg/querier/tripperware/query_attribute_matcher_test.go

@@ -186,11 +186,12 @@ func Test_rejectQueryOrSetPriorityShouldRejectIfMatches(t *testing.T) {
 			},
 		},
 
-		"should reject if query rejection regex matches match[] of series request and time window": {
+		"should reject if query rejection api matches and regex matches match[] of series request and time window": {
 			queryRejectionEnabled: true,
 			path:                  fmt.Sprintf("/api/v1/series?start=%d&end=%d&step=7s&match[]=%s", now.Add(-30*time.Minute).UnixMilli()/1000, now.Add(-20*time.Minute).UnixMilli()/1000, url.QueryEscape("count(sum(up))")),
 			expectedError:         httpgrpc.Errorf(http.StatusUnprocessableEntity, QueryRejectErrorMessage),
 			rejectQueryAttribute: validation.QueryAttribute{
+				ApiType:       "series",
 				Regex:         ".*sum.*",
 				CompiledRegex: regexp.MustCompile(".*sum.*"),
 				TimeWindow: validation.TimeWindow{
@@ -199,6 +200,17 @@ func Test_rejectQueryOrSetPriorityShouldRejectIfMatches(t *testing.T) {
 				},
 			},
 		},
+
+		"should not reject if query api_type doesn't match matches": {
+			queryRejectionEnabled: true,
+			path:                  fmt.Sprintf("/api/v1/series?start=%d&end=%d&step=7s&match[]=%s", now.Add(-30*time.Minute).UnixMilli()/1000, now.Add(-20*time.Minute).UnixMilli()/1000, url.QueryEscape("count(sum(up))")),
+			expectedError:         nil,
+			rejectQueryAttribute: validation.QueryAttribute{
+				ApiType:       "query",
+				Regex:         ".*sum.*",
+				CompiledRegex: regexp.MustCompile(".*sum.*"),
+			},
+		},
 	}
 
 	for testName, testData := range tests {

pkg/util/validation/limits.go

@@ -72,6 +72,7 @@ type QueryRejection struct {
 }
 
 type QueryAttribute struct {
+	ApiType                string         `yaml:"api_type" json:"api_type" doc:"nocli|description=API type for the query. Should be one of the query, query_range, series, labels, label_values. If not set, it won't be checked."`
 	Regex                  string         `yaml:"regex" json:"regex" doc:"nocli|description=Regex that the query string (or at least one of the matchers in metadata query) should match. If not set, it won't be checked."`
 	TimeWindow             TimeWindow     `yaml:"time_window" json:"time_window" doc:"nocli|description=Overall data select time window (including range selectors, modifiers and lookback delta) that the query should be within. If not set, it won't be checked."`
 	TimeRangeLimit         TimeRangeLimit `yaml:"time_range_limit" json:"time_range_limit" doc:"nocli|description=Query time range should be within this limit to match. If not set, it won't be checked."`