query rejection - address comments

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

Author: erlan-z

docs/configuration/config-file-reference.md

@@ -3287,7 +3287,9 @@ query_rejection:
   # CLI flag: -frontend.query-rejection.enabled
   [enabled: <boolean> | default = false]
 
-  # List of query attributes for rejection.
+  # List of query attributes that queries will be matched against. Query will be
+  # matched against each of those query attributes separately and if query
+  # matches any of them then it will be rejected.
   [query_attributes: <list of QueryAttribute> | default = []]
 
 # Duration to delay the evaluation of rules to ensure the underlying metrics
@@ -5350,14 +5352,17 @@ limits:
 # priority level. Value between 0 and 1 will be used as a percentage.
 [reserved_queriers: <float> | default = 0]
 
-# List of query attributes to assign the priority.
+# List of query attributes that queries will be matched against. Query will be
+# matched against each of those query attributes separately and if query matches
+# any of them then it will be assigned to this priority.
 [query_attributes: <list of QueryAttribute> | default = []]
 ```
 
 ### `QueryAttribute`
 
 ```yaml
-# Regex that the query string should match. If not set, it won't be checked.
+# Regex that the query string(or any of the matchers in metadata query) should
+# match. If not set, it won't be checked.
 [regex: <string> | default = ""]
 
 # Overall data select time window (including range selectors, modifiers and
@@ -5374,7 +5379,8 @@ time_window:
   # checked.
   [end: <int> | default = 0]
 
-# Limit that query time range should be within. If not set, it won't be checked.
+# Queries with time range that is within this limits will match. If not set, it
+# won't be checked.
 time_range_limit:
   # Query time range should be above or equal to this value to match. If set to
   # 0, it won't be checked.
@@ -5395,13 +5401,18 @@ query_step_limit:
   # won't be checked.
   [max: <int> | default = 0]
 
-# User agent for the query. If not set, it won't be checked.
+# Regex that User-Agent header of the request should match. If not set, it won't
+# be checked.
 [user_agent: <string> | default = ""]
 
-# Dashboard UID for the query. If not set, it won't be checked.
+# Grafana includes X-Dashboard-Uid header in query requests. If this field is
+# provided then X-Dashboard-Uid header of request should match this value. If
+# not set, it won't be checked.
 [dashboard_uid: <string> | default = ""]
 
-# Panel Id for the query. If not set, it won't be checked.
+# Grafana includes X-Panel-Id header in query requests. If this field is
+# provided then X-Panel-Id header of request should match this value. If not
+# set, it won't be checked.
 [panel_id: <string> | default = ""]
 ```
 

docs/configuration/v1-guarantees.md

@@ -113,3 +113,4 @@ Currently experimental features are:
   - `-ruler.ring.tokens-file-path` (path) CLI flag
 - Native Histograms
   - Ingestion can be enabled by setting `-blocks-storage.tsdb.enable-native-histograms=true` on Ingester.
+- Query-frontend: query rejection (`-frontend.query-rejection.enabled`)

pkg/querier/tripperware/query_attribute_matcher.go

@@ -21,11 +21,9 @@ func rejectQueryOrSetPriority(r *http.Request, now time.Time, lookbackDelta time
 	if limits == nil || !(limits.QueryPriority(userStr).Enabled || limits.QueryRejection(userStr).Enabled) {
 		return nil
 	}
+	op := getOperation(r)
 
-	isExpressionQuery := strings.HasSuffix(r.URL.Path, "/query") || strings.HasSuffix(r.URL.Path, "/query_range")
-	isMetadataQuery := strings.HasSuffix(r.URL.Path, "/series") || strings.HasSuffix(r.URL.Path, "/labels") || strings.HasSuffix(r.URL.Path, "/values")
-
-	if isExpressionQuery {
+	if op == "query" || op == "query_range" {
 		query := r.FormValue("query")
 		expr, err := parser.ParseExpr(query)
 		if err != nil {
@@ -36,7 +34,7 @@ func rejectQueryOrSetPriority(r *http.Request, now time.Time, lookbackDelta time
 		if queryReject := limits.QueryRejection(userStr); queryReject.Enabled && query != "" {
 			for _, attribute := range queryReject.QueryAttributes {
 				if matchAttributeForExpressionQuery(attribute, r, query, expr, now, minTime, maxTime) {
-					rejectedQueriesPerTenant.WithLabelValues(userStr).Inc()
+					rejectedQueriesPerTenant.WithLabelValues(op, userStr).Inc()
 					return httpgrpc.Errorf(http.StatusUnprocessableEntity, queryRejectErrorMessage)
 				}
 			}
@@ -59,11 +57,10 @@ func rejectQueryOrSetPriority(r *http.Request, now time.Time, lookbackDelta time
 		}
 	}
 
-	if queryReject := limits.QueryRejection(userStr); queryReject.Enabled && isMetadataQuery {
+	if queryReject := limits.QueryRejection(userStr); queryReject.Enabled && (op == "series" || op == "labels" || op == "label_values") {
 		for _, attribute := range queryReject.QueryAttributes {
-
 			if matchAttributeForMetadataQuery(attribute, r, now) {
-				rejectedQueriesPerTenant.WithLabelValues(userStr).Inc()
+				rejectedQueriesPerTenant.WithLabelValues(op, userStr).Inc()
 				return httpgrpc.Errorf(http.StatusUnprocessableEntity, queryRejectErrorMessage)
 			}
 		}
@@ -72,6 +69,23 @@ func rejectQueryOrSetPriority(r *http.Request, now time.Time, lookbackDelta time
 	return nil
 }
 
+func getOperation(r *http.Request) string {
+	switch {
+	case strings.HasSuffix(r.URL.Path, "/query"):
+		return "query"
+	case strings.HasSuffix(r.URL.Path, "/query_range"):
+		return "query_range"
+	case strings.HasSuffix(r.URL.Path, "/series"):
+		return "series"
+	case strings.HasSuffix(r.URL.Path, "/labels"):
+		return "labels"
+	case strings.HasSuffix(r.URL.Path, "/values"):
+		return "label_values"
+	default:
+		return "other"
+	}
+}
+
 func matchAttributeForExpressionQuery(attribute validation.QueryAttribute, r *http.Request, query string, expr parser.Expr, now time.Time, minTime, maxTime int64) bool {
 	if attribute.Regex != "" && attribute.Regex != ".*" && attribute.Regex != ".+" {
 		if attribute.CompiledRegex != nil && !attribute.CompiledRegex.MatchString(query) {
@@ -83,16 +97,18 @@ func matchAttributeForExpressionQuery(attribute validation.QueryAttribute, r *ht
 		return false
 	}
 
-	if !isWithinTimeRangeAttribute(attribute.TimeRangeLimit, maxTime-minTime) {
+	if !isWithinTimeRangeAttribute(attribute.TimeRangeLimit, minTime, maxTime) {
 		return false
 	}
 
 	if !isWithinQueryStepLimit(attribute.QueryStepLimit, r, expr) {
 		return false
 	}
 
-	if attribute.UserAgent != "" && attribute.UserAgent != r.Header.Get("User-Agent") {
-		return false
+	if attribute.UserAgentRegex != "" && attribute.UserAgentRegex != ".*" && attribute.CompiledUserAgentRegex != nil {
+		if !attribute.CompiledUserAgentRegex.MatchString(r.Header.Get("User-Agent")) {
+			return false
+		}
 	}
 
 	if attribute.DashboardUID != "" && attribute.DashboardUID != r.Header.Get("X-Dashboard-Uid") {
@@ -107,13 +123,13 @@ func matchAttributeForExpressionQuery(attribute validation.QueryAttribute, r *ht
 }
 
 func matchAttributeForMetadataQuery(attribute validation.QueryAttribute, r *http.Request, now time.Time) bool {
-	if attribute.Regex != "" && attribute.Regex != ".*" && attribute.Regex != ".+" && attribute.CompiledRegex != nil {
+	if attribute.Regex != "" && attribute.Regex != ".*" && attribute.CompiledRegex != nil {
 		for _, matcher := range r.Form["match[]"] {
 			if attribute.CompiledRegex.MatchString(matcher) {
-				continue
+				break
 			}
-			return false
 		}
+		return false
 	}
 
 	startTime, _ := util.ParseTime(r.FormValue("start"))
@@ -123,20 +139,14 @@ func matchAttributeForMetadataQuery(attribute validation.QueryAttribute, r *http
 		return false
 	}
 
-	if startTime != 0 && endTime != 0 && !isWithinTimeRangeAttribute(attribute.TimeRangeLimit, endTime-startTime) {
-		return false
-	}
-
-	if attribute.UserAgent != "" && attribute.UserAgent != r.Header.Get("User-Agent") {
+	if !isWithinTimeRangeAttribute(attribute.TimeRangeLimit, startTime, endTime) {
 		return false
 	}
 
-	if attribute.DashboardUID != "" && attribute.DashboardUID != r.Header.Get("X-Dashboard-Uid") {
-		return false
-	}
-
-	if attribute.PanelID != "" && attribute.PanelID != r.Header.Get("X-Panel-Id") {
-		return false
+	if attribute.UserAgentRegex != "" && attribute.UserAgentRegex != ".*" && attribute.CompiledUserAgentRegex != nil {
+		if !attribute.CompiledUserAgentRegex.MatchString(r.Header.Get("User-Agent")) {
+			return false
+		}
 	}
 
 	return true
@@ -164,11 +174,17 @@ func isWithinTimeAttributes(timeWindow validation.TimeWindow, now time.Time, sta
 	return true
 }
 
-func isWithinTimeRangeAttribute(limit validation.TimeRangeLimit, timeRangeInMillis int64) bool {
+func isWithinTimeRangeAttribute(limit validation.TimeRangeLimit, startTime, endTime int64) bool {
 	if limit.Min == 0 && limit.Max == 0 {
 		return true
 	}
 
+	if startTime == 0 || endTime == 0 {
+		return false
+	}
+
+	timeRangeInMillis := endTime - startTime
+
 	if limit.Min != 0 && time.Duration(limit.Min).Milliseconds() > timeRangeInMillis {
 		return false
 	}