Pagination on rule groups and tests

pkg/ruler/api.go

@@ -6,16 +6,17 @@
 package ruler
 
 import (
+	"encoding/hex"
 	"encoding/json"
 	"fmt"
 	"io"
 	"net/http"
 	"net/url"
-	"sort"
 	"strconv"
 	"strings"
 	"time"
 
+	"github.com/cespare/xxhash/v2"
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
 	"github.com/gorilla/mux"
@@ -64,6 +65,7 @@ type Alert struct {
 // RuleDiscovery has info for all rules
 type RuleDiscovery struct {
 	RuleGroups []*RuleGroup `json:"groups"`
+	NextToken  string       `json:"nextToken,omitempty"`
 }
 
 // RuleGroup has info for rules which are part of a group
@@ -166,6 +168,16 @@ func (a *API) PrometheusRules(w http.ResponseWriter, req *http.Request) {
 		return
 	}
 
+	nextToken := req.URL.Query().Get("next_token")
+	var maxGroups int
+	if maxGroupsVal := req.URL.Query().Get("max_groups"); maxGroupsVal != "" {
+		maxGroups, err = strconv.Atoi(maxGroupsVal)
+		if err != nil || maxGroups < 0 {
+			respondInvalidRequest(logger, w, "invalid max groups value")
+			return
+		}
+	}
+
 	rulesReq := RulesRequest{
 		Filter:    AnyRule,
 		RuleName:  req.URL.Query()["rule_name"],
@@ -195,8 +207,21 @@ func (a *API) PrometheusRules(w http.ResponseWriter, req *http.Request) {
 	}
 
 	groups := make([]*RuleGroup, 0, len(rgs))
-
+	var newToken string
+	foundToken := false
 	for _, g := range rgs {
+		if nextToken != "" && !foundToken {
+			if nextToken != getRuleGroupNextToken(g.Group.Namespace, g.Group.Name) {
+				continue
+			}
+			foundToken = true
+		}
+
+		if maxGroups > 0 && len(groups) == maxGroups {
+			newToken = getRuleGroupNextToken(g.Group.Namespace, g.Group.Name)
+			break
+		}
+
 		grp := RuleGroup{
 			Name:           g.Group.Name,
 			File:           g.Group.Namespace,
@@ -241,17 +266,13 @@ func (a *API) PrometheusRules(w http.ResponseWriter, req *http.Request) {
 				}
 			}
 		}
+
 		groups = append(groups, &grp)
 	}
 
-	// keep data.groups are in order
-	sort.Slice(groups, func(i, j int) bool {
-		return groups[i].File < groups[j].File
-	})
-
 	b, err := json.Marshal(&response{
 		Status: "success",
-		Data:   &RuleDiscovery{RuleGroups: groups},
+		Data:   &RuleDiscovery{RuleGroups: groups, NextToken: newToken},
 	})
 	if err != nil {
 		level.Error(logger).Log("msg", "error marshaling json response", "err", err)
@@ -265,6 +286,13 @@ func (a *API) PrometheusRules(w http.ResponseWriter, req *http.Request) {
 	}
 }
 
+func getRuleGroupNextToken(file, group string) string {
+	h := xxhash.New()
+	h.Write([]byte(file + ":" + group))
+
+	return hex.EncodeToString(h.Sum(nil))
+}
+
 func (a *API) PrometheusAlerts(w http.ResponseWriter, req *http.Request) {
 	logger, ctx := spanlogger.NewWithLogger(req.Context(), a.logger, "API.PrometheusAlerts")
 	defer logger.Finish()

pkg/ruler/api_test.go

@@ -14,6 +14,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"net/url"
+	"strconv"
 	"strings"
 	"testing"
 	"time"
@@ -905,6 +906,135 @@ func TestRuler_PrometheusRules(t *testing.T) {
 	}
 }
 
+func TestRuler_PrometheusRulesPagination(t *testing.T) {
+	const (
+		userID   = "user1"
+		interval = time.Minute
+	)
+
+	ruleGroups := rulespb.RuleGroupList{}
+	for ns := 0; ns < 3; ns++ {
+		for group := 0; group < 3; group++ {
+			g := &rulespb.RuleGroupDesc{
+				Name:      fmt.Sprintf("test-group-%d", group),
+				Namespace: fmt.Sprintf("test-namespace-%d", ns),
+				User:      userID,
+				Rules: []*rulespb.RuleDesc{
+					createAlertingRule("testalertingrule", "up < 1"),
+				},
+				Interval: interval,
+			}
+			ruleGroups = append(ruleGroups, g)
+		}
+	}
+
+	cfg := defaultRulerConfig(t)
+	cfg.TenantFederation.Enabled = true
+
+	storageRules := map[string]rulespb.RuleGroupList{
+		userID: ruleGroups,
+	}
+
+	r := prepareRuler(t, cfg, newMockRuleStore(storageRules), withRulerAddrAutomaticMapping(), withLimits(validation.MockDefaultOverrides()), withStart())
+
+	// Rules will be synchronized asynchronously, so we wait until the expected number of rule groups
+	// has been synched.
+	test.Poll(t, 5*time.Second, len(ruleGroups), func() interface{} {
+		ctx := user.InjectOrgID(context.Background(), userID)
+		rls, _ := r.Rules(ctx, &RulesRequest{})
+		return len(rls.Groups)
+	})
+
+	a := NewAPI(r, r.directStore, log.NewNopLogger())
+
+	getRulesResponse := func(groupSize int, nextToken string) response {
+		queryParams := "?" + url.Values{
+			"max_groups": []string{strconv.Itoa(groupSize)},
+			"next_token": []string{nextToken},
+		}.Encode()
+		req := requestFor(t, http.MethodGet, "https://localhost:8080/prometheus/api/v1/rules"+queryParams, nil, userID)
+		w := httptest.NewRecorder()
+		a.PrometheusRules(w, req)
+
+		resp := w.Result()
+		body, _ := io.ReadAll(resp.Body)
+
+		r := response{}
+		err := json.Unmarshal(body, &r)
+		require.NoError(t, err)
+
+		return r
+	}
+
+	getRulesFromResponse := func(resp response) RuleDiscovery {
+		jsonRules, err := json.Marshal(resp.Data)
+		require.NoError(t, err)
+		returnedRules := RuleDiscovery{}
+		require.NoError(t, json.Unmarshal(jsonRules, &returnedRules))
+
+		return returnedRules
+	}
+
+	// No page size limit
+	resp := getRulesResponse(0, "")
+	require.Equal(t, "success", resp.Status)
+	rd := getRulesFromResponse(resp)
+	require.Len(t, rd.RuleGroups, len(ruleGroups))
+	require.Empty(t, rd.NextToken)
+
+	// We have 9 groups, keep fetching rules with a group page size of 2. The final
+	// page should have size 1 and an empty nextToken. Also check the groups returned
+	// in order
+	var nextToken string
+	returnedRuleGroups := make([]*RuleGroup, 0, len(ruleGroups))
+	for i := 0; i < 4; i++ {
+		resp := getRulesResponse(2, nextToken)
+		require.Equal(t, "success", resp.Status)
+
+		rd := getRulesFromResponse(resp)
+		require.Len(t, rd.RuleGroups, 2)
+		require.NotEmpty(t, rd.NextToken)
+
+		returnedRuleGroups = append(returnedRuleGroups, rd.RuleGroups[0], rd.RuleGroups[1])
+		nextToken = rd.NextToken
+	}
+	resp = getRulesResponse(2, nextToken)
+	require.Equal(t, "success", resp.Status)
+
+	rd = getRulesFromResponse(resp)
+	require.Len(t, rd.RuleGroups, 1)
+	require.Empty(t, rd.NextToken)
+	returnedRuleGroups = append(returnedRuleGroups, rd.RuleGroups[0])
+
+	// Check the returned rules match the rules written
+	require.Equal(t, len(ruleGroups), len(returnedRuleGroups))
+	for i := 0; i < len(ruleGroups); i++ {
+		require.Equal(t, ruleGroups[i].Namespace, returnedRuleGroups[i].File)
+		require.Equal(t, ruleGroups[i].Name, returnedRuleGroups[i].Name)
+		require.Equal(t, len(ruleGroups[i].Rules), len(returnedRuleGroups[i].Rules))
+		for j := 0; j < len(ruleGroups[i].Rules); j++ {
+			jsonRule, err := json.Marshal(returnedRuleGroups[i].Rules[j])
+			require.NoError(t, err)
+			rule := alertingRule{}
+			require.NoError(t, json.Unmarshal(jsonRule, &rule))
+			require.Equal(t, ruleGroups[i].Rules[j].Alert, rule.Name)
+		}
+	}
+
+	// Invalid max groups value
+	resp = getRulesResponse(-1, "")
+	require.Equal(t, "error", resp.Status)
+	require.Equal(t, v1.ErrBadData, resp.ErrorType)
+	require.Equal(t, "invalid max groups value", resp.Error)
+
+	// Bad token should return no groups
+	resp = getRulesResponse(0, "bad-token")
+	require.Equal(t, "success", resp.Status)
+
+	rd = getRulesFromResponse(resp)
+	require.Len(t, rd.RuleGroups, 0)
+}
+
 func TestRuler_PrometheusAlerts(t *testing.T) {
 	cfg := defaultRulerConfig(t)