Query priority based on operation

CHANGELOG.md

@@ -8,6 +8,7 @@
 * [ENHANCEMENT] Distributor: Add a label references validation for remote write v2 request. #7074
 * [ENHANCEMENT] Distributor: Add count, spans, and buckets validations for native histogram. #7072
 * [ENHANCEMENT] Ruler: Add DecodingConcurrency config flag for Thanos Engine. #7118
+* [ENHANCEMENT] Query Frontend: Add query priority based on operation. #7128
 * [BUGFIX] Ring: Change DynamoDB KV to retry indefinitely for WatchKey. #7088
 * [BUGFIX] Ruler: Add XFunctions validation support. #7111
 

pkg/querier/tripperware/query_attribute_matcher.go

@@ -22,6 +22,7 @@ func rejectQueryOrSetPriority(r *http.Request, now time.Time, lookbackDelta time
 		return nil
 	}
 	op := getOperation(r)
+	reqStats := stats.FromContext(r.Context())
 
 	if op == "query" || op == "query_range" {
 		query := r.FormValue("query")
@@ -40,7 +41,6 @@ func rejectQueryOrSetPriority(r *http.Request, now time.Time, lookbackDelta time
 			}
 		}
 
-		reqStats := stats.FromContext(r.Context())
 		reqStats.SetDataSelectMaxTime(maxTime)
 		reqStats.SetDataSelectMinTime(minTime)
 
@@ -55,13 +55,24 @@ func rejectQueryOrSetPriority(r *http.Request, now time.Time, lookbackDelta time
 			}
 			reqStats.SetPriority(queryPriority.DefaultPriority)
 		}
-	}
+	} else {
+		if queryReject := limits.QueryRejection(userStr); queryReject.Enabled && (op == "series" || op == "labels" || op == "label_values") {
+			for _, attribute := range queryReject.QueryAttributes {
+				if matchAttributeForMetadataQuery(attribute, op, r, now) {
+					rejectedQueriesPerTenant.WithLabelValues(op, userStr).Inc()
+					return httpgrpc.Errorf(http.StatusUnprocessableEntity, QueryRejectErrorMessage)
+				}
+			}
+		}
 
-	if queryReject := limits.QueryRejection(userStr); queryReject.Enabled && (op == "series" || op == "labels" || op == "label_values") {
-		for _, attribute := range queryReject.QueryAttributes {
-			if matchAttributeForMetadataQuery(attribute, op, r, now) {
-				rejectedQueriesPerTenant.WithLabelValues(op, userStr).Inc()
-				return httpgrpc.Errorf(http.StatusUnprocessableEntity, QueryRejectErrorMessage)
+		if queryPriority := limits.QueryPriority(userStr); queryPriority.Enabled && len(queryPriority.Priorities) != 0 {
+			for _, priority := range queryPriority.Priorities {
+				for _, attribute := range priority.QueryAttributes {
+					if matchAttributeForMetadataQuery(attribute, op, r, now) {
+						reqStats.SetPriority(priority.Priority)
+						break
+					}
+				}
 			}
 		}
 	}
@@ -81,6 +92,8 @@ func getOperation(r *http.Request) string {
 		return "labels"
 	case strings.HasSuffix(r.URL.Path, "/values"):
 		return "label_values"
+	case strings.HasSuffix(r.URL.Path, "/metadata"):
+		return "metadata"
 	default:
 		return "other"
 	}

pkg/querier/tripperware/query_attribute_matcher_test.go

@@ -59,10 +59,11 @@ func Test_rejectQueryOrSetPriorityShouldReturnDefaultPriorityIfNotEnabledOrInval
 			path:                  "/api/v1/query?time=1536716898&query=",
 			expectedError:         httpgrpc.Errorf(http.StatusBadRequest, "unknown position: parse error: no expression found in input"),
 		},
-		"should miss if it's metadata query and only priority is enabled": {
+		"should set priority if it's metadata query and only priority is enabled": {
 			queryPriorityEnabled:  true,
 			queryRejectionEnabled: false,
 			path:                  "/api/v1/labels?match[]",
+			expectedPriority:      1,
 		},
 		"should set priority if regex match and rejection disabled": {
 			queryPriorityEnabled:  true,
@@ -706,3 +707,45 @@ func Test_matchAttributeForExpressionQueryShouldMatchUserAgentRegex(t *testing.T
 	}
 
 }
+
+func Test_rejectQueryOrSetPriorityShouldMatchOnApiType(t *testing.T) {
+
+	limits := mockLimits{queryPriority: validation.QueryPriority{
+		Priorities: []validation.PriorityDef{
+			{
+				Priority: 7,
+				QueryAttributes: []validation.QueryAttribute{
+					{
+						ApiType: "metadata",
+					},
+				},
+			},
+		},
+	},
+	}
+
+	type testCase struct {
+		path             string
+		expectedPriority int64
+	}
+
+	tests := map[string]testCase{
+		"should set priority based on api type": {
+			path:             "/api/v1/targets/metadata?limit=1",
+			expectedPriority: 7,
+		},
+	}
+
+	for testName, testData := range tests {
+		t.Run(testName, func(t *testing.T) {
+			req, err := http.NewRequest("GET", testData.path, http.NoBody)
+			require.NoError(t, err)
+			reqStats, ctx := stats.ContextWithEmptyStats(context.Background())
+			req = req.WithContext(ctx)
+			limits.queryPriority.Enabled = true
+			resultErr := rejectQueryOrSetPriority(req, time.Now(), time.Duration(1), limits, "", rejectedQueriesPerTenant)
+			assert.NoError(t, resultErr)
+			assert.Equal(t, testData.expectedPriority, reqStats.Priority)
+		})
+	}
+}