Add paginated feature to list rules api

Signed-off-by: Yijie Qin <qinyijie@amazon.com>

web/api/v1/api.go

@@ -15,6 +15,7 @@ package v1
 
 import (
 	"context"
+	"crypto/sha1"
 	"errors"
 	"fmt"
 	"math"
@@ -168,6 +169,17 @@ type apiFuncResult struct {
 
 type apiFunc func(r *http.Request) apiFuncResult
 
+type listRulesPaginationRequest struct {
+	MaxAlerts     int32
+	MaxRuleGroups int32
+	NextToken     string
+}
+
+type parsePaginationError struct {
+	err       error
+	parameter string
+}
+
 // TSDBAdminStats defines the tsdb interfaces used by the v1 API for admin operations as well as statistics.
 type TSDBAdminStats interface {
 	CleanTombstones() error
@@ -1329,6 +1341,12 @@ type RuleDiscovery struct {
 	RuleGroups []*RuleGroup `json:"groups"`
 }
 
+// PaginatedRuleDiscovery has info for all rules with pagination.
+type PaginatedRuleDiscovery struct {
+	RuleGroups []*RuleGroup `json:"groups"`
+	NextToken  string       `json:"nextToken"`
+}
+
 // RuleGroup has info for rules which are part of a group.
 type RuleGroup struct {
 	Name string `json:"name"`
@@ -1363,6 +1381,29 @@ type AlertingRule struct {
 	Type string `json:"type"`
 }
 
+type AlertingRulePaginated struct {
+	// State can be "pending", "firing", "inactive".
+	State          string           `json:"state"`
+	Name           string           `json:"name"`
+	Query          string           `json:"query"`
+	Duration       float64          `json:"duration"`
+	KeepFiringFor  float64          `json:"keepFiringFor"`
+	Labels         labels.Labels    `json:"labels"`
+	Annotations    labels.Labels    `json:"annotations"`
+	AlertInfo      *AlertsPaginated `json:"alertInfo"`
+	Health         rules.RuleHealth `json:"health"`
+	LastError      string           `json:"lastError,omitempty"`
+	EvaluationTime float64          `json:"evaluationTime"`
+	LastEvaluation time.Time        `json:"lastEvaluation"`
+	// Type of an alertingRule is always "alerting".
+	Type string `json:"type"`
+}
+
+type AlertsPaginated struct {
+	Alerts  []*Alert `json:"alerts"`
+	HasMore bool     `json:"hasMore"`
+}
+
 type RecordingRule struct {
 	Name           string           `json:"name"`
 	Query          string           `json:"query"`
@@ -1408,6 +1449,15 @@ func (api *API) rules(r *http.Request) apiFuncResult {
 		return invalidParamError(err, "exclude_alerts")
 	}
 
+	paginationRequest, parseErr := parseListRulesPaginationRequest(r)
+	if parseErr != nil {
+		return invalidParamError(parseErr.err, parseErr.parameter)
+	}
+
+	if paginationRequest != nil {
+		sort.Sort(GroupStateDescs(ruleGroups))
+	}
+
 	rgs := make([]*RuleGroup, 0, len(ruleGroups))
 	for _, grp := range ruleGroups {
 		if len(rgSet) > 0 {
@@ -1422,6 +1472,12 @@ func (api *API) rules(r *http.Request) apiFuncResult {
 			}
 		}
 
+		// Skip the rule group if the next token is set and hasn't arrived the nextToken item yet.
+		groupID := getRuleGroupNextToken(grp.File(), grp.Name())
+		if paginationRequest != nil && len(paginationRequest.NextToken) > 0 && paginationRequest.NextToken >= groupID {
+			continue
+		}
+
 		apiRuleGroup := &RuleGroup{
 			Name:           grp.Name(),
 			File:           grp.File(),
@@ -1453,20 +1509,47 @@ func (api *API) rules(r *http.Request) apiFuncResult {
 				if !excludeAlerts {
 					activeAlerts = rulesAlertsToAPIAlerts(rule.ActiveAlerts())
 				}
-				enrichedRule = AlertingRule{
+				hasMore := false
-					State:          rule.State().String(),
+				if paginationRequest != nil && paginationRequest.MaxAlerts >= 0 && len(rule.ActiveAlerts()) > int(paginationRequest.MaxAlerts) {
-					Name:           rule.Name(),
+					activeAlerts = activeAlerts[:int(paginationRequest.MaxAlerts)]
-					Query:          rule.Query().String(),
+					hasMore = true
-					Duration:       rule.HoldDuration().Seconds(),
+				}
-					KeepFiringFor:  rule.KeepFiringFor().Seconds(),
+
-					Labels:         rule.Labels(),
+				if paginationRequest != nil {
-					Annotations:    rule.Annotations(),
+					enrichedRule = AlertingRulePaginated{
-					Alerts:         activeAlerts,
+						State:         rule.State().String(),
-					Health:         rule.Health(),
+						Name:          rule.Name(),
-					LastError:      lastError,
+						Query:         rule.Query().String(),
-					EvaluationTime: rule.GetEvaluationDuration().Seconds(),
+						Duration:      rule.HoldDuration().Seconds(),
-					LastEvaluation: rule.GetEvaluationTimestamp(),
+						KeepFiringFor: rule.KeepFiringFor().Seconds(),
-					Type:           "alerting",
+						Labels:        rule.Labels(),
+						Annotations:   rule.Annotations(),
+						AlertInfo: &AlertsPaginated{
+							Alerts:  activeAlerts,
+							HasMore: hasMore,
+						},
+						Health:         rule.Health(),
+						LastError:      lastError,
+						EvaluationTime: rule.GetEvaluationDuration().Seconds(),
+						LastEvaluation: rule.GetEvaluationTimestamp(),
+						Type:           "alerting",
+					}
+				} else {
+					enrichedRule = AlertingRule{
+						State:          rule.State().String(),
+						Name:           rule.Name(),
+						Query:          rule.Query().String(),
+						Duration:       rule.HoldDuration().Seconds(),
+						KeepFiringFor:  rule.KeepFiringFor().Seconds(),
+						Labels:         rule.Labels(),
+						Annotations:    rule.Annotations(),
+						Alerts:         activeAlerts,
+						Health:         rule.Health(),
+						LastError:      lastError,
+						EvaluationTime: rule.GetEvaluationDuration().Seconds(),
+						LastEvaluation: rule.GetEvaluationTimestamp(),
+						Type:           "alerting",
+					}
 				}
 			case *rules.RecordingRule:
 				if !returnRecording {
@@ -1497,6 +1580,15 @@ func (api *API) rules(r *http.Request) apiFuncResult {
 			rgs = append(rgs, apiRuleGroup)
 		}
 	}
+
+	if paginationRequest != nil {
+		paginatedRes := &PaginatedRuleDiscovery{RuleGroups: make([]*RuleGroup, 0, len(ruleGroups))}
+		returnGroups, nextToken := TruncateGroupInfos(rgs, int(paginationRequest.MaxRuleGroups))
+		paginatedRes.RuleGroups = returnGroups
+		paginatedRes.NextToken = nextToken
+		return apiFuncResult{paginatedRes, nil, nil, nil}
+	}
+
 	res.RuleGroups = rgs
 	return apiFuncResult{res, nil, nil, nil}
 }
@@ -1515,6 +1607,77 @@ func parseExcludeAlerts(r *http.Request) (bool, error) {
 	return excludeAlerts, nil
 }
 
+func parseListRulesPaginationRequest(r *http.Request) (*listRulesPaginationRequest, *parsePaginationError) {
+	var (
+		returnMaxAlert      = int32(-1)
+		returnMaxRuleGroups = int32(-1)
+		returnNextToken     = ""
+	)
+
+	if r.URL.Query().Get("maxAlerts") != "" {
+		maxAlert, err := strconv.ParseInt(r.URL.Query().Get("maxAlerts"), 10, 32)
+		if err != nil || maxAlert < 0 {
+			return nil, &parsePaginationError{
+				err:       fmt.Errorf("maxAlerts need to be a valid number and larger than 0: %w", err),
+				parameter: "maxAlerts",
+			}
+		}
+		returnMaxAlert = int32(maxAlert)
+	}
+
+	if r.URL.Query().Get("maxRuleGroups") != "" {
+		maxRuleGroups, err := strconv.ParseInt(r.URL.Query().Get("maxRuleGroups"), 10, 32)
+		if err != nil || maxRuleGroups < 0 {
+			return nil, &parsePaginationError{
+				err:       fmt.Errorf("maxRuleGroups need to be a valid number and larger than 0: %w", err),
+				parameter: "maxAlerts",
+			}
+		}
+		returnMaxRuleGroups = int32(maxRuleGroups)
+	}
+
+	if r.URL.Query().Get("nextToken") != "" {
+		returnNextToken = r.URL.Query().Get("nextToken")
+	}
+
+	if returnMaxRuleGroups >= 0 || returnMaxAlert >= 0 || returnNextToken != "" {
+		return &listRulesPaginationRequest{
+			MaxRuleGroups: returnMaxRuleGroups,
+			MaxAlerts:     returnMaxAlert,
+			NextToken:     returnNextToken,
+		}, nil
+	}
+
+	return nil, nil
+}
+
+func TruncateGroupInfos(groupInfos []*RuleGroup, maxRuleGroups int) ([]*RuleGroup, string) {
+	resultNumber := 0
+	var returnPaginationToken string
+	returnGroupDescs := make([]*RuleGroup, 0, len(groupInfos))
+	for _, groupInfo := range groupInfos {
+		// Add the rule group to the return slice if the maxRuleGroups is not hit
+		if maxRuleGroups < 0 || resultNumber < maxRuleGroups {
+			returnGroupDescs = append(returnGroupDescs, groupInfo)
+			resultNumber++
+			continue
+		}
+
+		// Return the next token if there is more aggregation group
+		if maxRuleGroups > 0 && resultNumber == maxRuleGroups {
+			returnPaginationToken = getRuleGroupNextToken(returnGroupDescs[maxRuleGroups-1].File, returnGroupDescs[maxRuleGroups-1].Name)
+			break
+		}
+	}
+	return returnGroupDescs, returnPaginationToken
+}
+
+func getRuleGroupNextToken(namespace, group string) string {
+	h := sha1.New()
+	h.Write([]byte(namespace + ";" + group))
+	return fmt.Sprintf("%x", h.Sum(nil))
+}
+
 type prometheusConfig struct {
 	YAML string `json:"yaml"`
 }
@@ -1911,3 +2074,11 @@ func parseLimitParam(limitStr string) (limit int, err error) {
 
 	return limit, nil
 }
+
+type GroupStateDescs []*rules.Group
+
+func (gi GroupStateDescs) Swap(i, j int) { gi[i], gi[j] = gi[j], gi[i] }
+func (gi GroupStateDescs) Less(i, j int) bool {
+	return getRuleGroupNextToken(gi[i].File(), gi[i].Name()) < getRuleGroupNextToken(gi[j].File(), gi[j].Name())
+}
+func (gi GroupStateDescs) Len() int { return len(gi) }

web/api/v1/api_test.go

@@ -308,7 +308,15 @@ func (m *rulesRetrieverMock) CreateRuleGroups() {
 		ShouldRestore: false,
 		Opts:          opts,
 	})
-	m.ruleGroups = []*rules.Group{group}
+	group2 := rules.NewGroup(rules.GroupOptions{
+		Name:          "grp2",
+		File:          "/path/to/file",
+		Interval:      time.Second,
+		Rules:         []rules.Rule{r[0]},
+		ShouldRestore: false,
+		Opts:          opts,
+	})
+	m.ruleGroups = []*rules.Group{group, group2}
 }
 
 func (m *rulesRetrieverMock) AlertingRules() []*rules.AlertingRule {
@@ -1066,27 +1074,53 @@ func testEndpoints(t *testing.T, api *API, tr *testTargetRetriever, es storage.E
 
 	rulesZeroFunc := func(i interface{}) {
 		if i != nil {
-			v := i.(*RuleDiscovery)
+			switch v := i.(type) {
-			for _, ruleGroup := range v.RuleGroups {
+			case *RuleDiscovery:
-				ruleGroup.EvaluationTime = float64(0)
+				for _, ruleGroup := range v.RuleGroups {
-				ruleGroup.LastEvaluation = time.Time{}
+					ruleGroup.EvaluationTime = float64(0)
-				for k, rule := range ruleGroup.Rules {
+					ruleGroup.LastEvaluation = time.Time{}
-					switch r := rule.(type) {
+					for k, rule := range ruleGroup.Rules {
-					case AlertingRule:
+						switch r := rule.(type) {
-						r.LastEvaluation = time.Time{}
+						case AlertingRule:
-						r.EvaluationTime = float64(0)
+							r.LastEvaluation = time.Time{}
-						r.LastError = ""
+							r.EvaluationTime = float64(0)
-						r.Health = "ok"
+							r.LastError = ""
-						for _, alert := range r.Alerts {
+							r.Health = "ok"
-							alert.ActiveAt = nil
+							for _, alert := range r.Alerts {
+								alert.ActiveAt = nil
+							}
+							ruleGroup.Rules[k] = r
+						case RecordingRule:
+							r.LastEvaluation = time.Time{}
+							r.EvaluationTime = float64(0)
+							r.LastError = ""
+							r.Health = "ok"
+							ruleGroup.Rules[k] = r
+						}
+					}
+				}
+			case *PaginatedRuleDiscovery:
+				for _, ruleGroup := range v.RuleGroups {
+					ruleGroup.EvaluationTime = float64(0)
+					ruleGroup.LastEvaluation = time.Time{}
+					for k, rule := range ruleGroup.Rules {
+						switch r := rule.(type) {
+						case AlertingRulePaginated:
+							r.LastEvaluation = time.Time{}
+							r.EvaluationTime = float64(0)
+							r.LastError = ""
+							r.Health = "ok"
+							for _, alert := range r.AlertInfo.Alerts {
+								alert.ActiveAt = nil
+							}
+							ruleGroup.Rules[k] = r
+						case RecordingRule:
+							r.LastEvaluation = time.Time{}
+							r.EvaluationTime = float64(0)
+							r.LastError = ""
+							r.Health = "ok"
+							ruleGroup.Rules[k] = r
 						}
-						ruleGroup.Rules[k] = r
-					case RecordingRule:
-						r.LastEvaluation = time.Time{}
-						r.EvaluationTime = float64(0)
-						r.LastError = ""
-						r.Health = "ok"
-						ruleGroup.Rules[k] = r
 					}
 				}
 			}
@@ -2127,6 +2161,25 @@ func testEndpoints(t *testing.T, api *API, tr *testTargetRetriever, es storage.E
 							},
 						},
 					},
+					{
+						Name:     "grp2",
+						File:     "/path/to/file",
+						Interval: 1,
+						Limit:    0,
+						Rules: []Rule{
+							AlertingRule{
+								State:       "inactive",
+								Name:        "test_metric3",
+								Query:       "absent(test_metric3) != 1",
+								Duration:    1,
+								Labels:      labels.Labels{},
+								Annotations: labels.Labels{},
+								Alerts:      []*Alert{},
+								Health:      "ok",
+								Type:        "alerting",
+							},
+						},
+					},
 				},
 			},
 			zeroFunc: rulesZeroFunc,
@@ -2186,6 +2239,25 @@ func testEndpoints(t *testing.T, api *API, tr *testTargetRetriever, es storage.E
 							},
 						},
 					},
+					{
+						Name:     "grp2",
+						File:     "/path/to/file",
+						Interval: 1,
+						Limit:    0,
+						Rules: []Rule{
+							AlertingRule{
+								State:       "inactive",
+								Name:        "test_metric3",
+								Query:       "absent(test_metric3) != 1",
+								Duration:    1,
+								Labels:      labels.Labels{},
+								Annotations: labels.Labels{},
+								Alerts:      nil,
+								Health:      "ok",
+								Type:        "alerting",
+							},
+						},
+					},
 				},
 			},
 			zeroFunc: rulesZeroFunc,
@@ -2245,6 +2317,25 @@ func testEndpoints(t *testing.T, api *API, tr *testTargetRetriever, es storage.E
 							},
 						},
 					},
+					{
+						Name:     "grp2",
+						File:     "/path/to/file",
+						Interval: 1,
+						Limit:    0,
+						Rules: []Rule{
+							AlertingRule{
+								State:       "inactive",
+								Name:        "test_metric3",
+								Query:       "absent(test_metric3) != 1",
+								Duration:    1,
+								Labels:      labels.Labels{},
+								Annotations: labels.Labels{},
+								Alerts:      []*Alert{},
+								Health:      "ok",
+								Type:        "alerting",
+							},
+						},
+					},
 				},
 			},
 			zeroFunc: rulesZeroFunc,
@@ -2336,6 +2427,156 @@ func testEndpoints(t *testing.T, api *API, tr *testTargetRetriever, es storage.E
 			},
 			zeroFunc: rulesZeroFunc,
 		},
+		{
+			endpoint: api.rules,
+			query: url.Values{
+				"maxRuleGroups": []string{"1"},
+			},
+			response: &PaginatedRuleDiscovery{
+				NextToken: getRuleGroupNextToken("/path/to/file", "grp"),
+				RuleGroups: []*RuleGroup{
+					{
+						Name:     "grp",
+						File:     "/path/to/file",
+						Interval: 1,
+						Limit:    0,
+						Rules: []Rule{
+							AlertingRulePaginated{
+								State:       "inactive",
+								Name:        "test_metric3",
+								Query:       "absent(test_metric3) != 1",
+								Duration:    1,
+								Labels:      labels.Labels{},
+								Annotations: labels.Labels{},
+								AlertInfo:   &AlertsPaginated{Alerts: []*Alert{}},
+								Health:      "ok",
+								Type:        "alerting",
+							},
+							AlertingRulePaginated{
+								State:       "inactive",
+								Name:        "test_metric4",
+								Query:       "up == 1",
+								Duration:    1,
+								Labels:      labels.Labels{},
+								Annotations: labels.Labels{},
+								AlertInfo:   &AlertsPaginated{Alerts: []*Alert{}},
+								Health:      "ok",
+								Type:        "alerting",
+							},
+							AlertingRulePaginated{
+								State:       "pending",
+								Name:        "test_metric5",
+								Query:       "vector(1)",
+								Duration:    1,
+								Labels:      labels.FromStrings("name", "tm5"),
+								Annotations: labels.Labels{},
+								AlertInfo: &AlertsPaginated{
+									Alerts: []*Alert{
+										{
+											Labels:      labels.FromStrings("alertname", "test_metric5", "name", "tm5"),
+											Annotations: labels.Labels{},
+											State:       "pending",
+											Value:       "1e+00",
+										},
+									},
+								},
+								Health: "ok",
+								Type:   "alerting",
+							},
+							RecordingRule{
+								Name:   "recording-rule-1",
+								Query:  "vector(1)",
+								Labels: labels.Labels{},
+								Health: "ok",
+								Type:   "recording",
+							},
+						},
+					},
+				},
+			},
+			zeroFunc: rulesZeroFunc,
+		},
+		{
+			endpoint: api.rules,
+			query: url.Values{
+				"maxAlerts": []string{"0"},
+			},
+			response: &PaginatedRuleDiscovery{
+				RuleGroups: []*RuleGroup{
+					{
+						Name:     "grp",
+						File:     "/path/to/file",
+						Interval: 1,
+						Limit:    0,
+						Rules: []Rule{
+							AlertingRulePaginated{
+								State:       "inactive",
+								Name:        "test_metric3",
+								Query:       "absent(test_metric3) != 1",
+								Duration:    1,
+								Labels:      labels.Labels{},
+								Annotations: labels.Labels{},
+								AlertInfo:   &AlertsPaginated{Alerts: []*Alert{}},
+								Health:      "ok",
+								Type:        "alerting",
+							},
+							AlertingRulePaginated{
+								State:       "inactive",
+								Name:        "test_metric4",
+								Query:       "up == 1",
+								Duration:    1,
+								Labels:      labels.Labels{},
+								Annotations: labels.Labels{},
+								AlertInfo:   &AlertsPaginated{Alerts: []*Alert{}},
+								Health:      "ok",
+								Type:        "alerting",
+							},
+							AlertingRulePaginated{
+								State:       "pending",
+								Name:        "test_metric5",
+								Query:       "vector(1)",
+								Duration:    1,
+								Labels:      labels.FromStrings("name", "tm5"),
+								Annotations: labels.Labels{},
+								AlertInfo: &AlertsPaginated{
+									Alerts:  []*Alert{},
+									HasMore: true,
+								},
+								Health: "ok",
+								Type:   "alerting",
+							},
+							RecordingRule{
+								Name:   "recording-rule-1",
+								Query:  "vector(1)",
+								Labels: labels.Labels{},
+								Health: "ok",
+								Type:   "recording",
+							},
+						},
+					},
+					{
+						Name:     "grp2",
+						File:     "/path/to/file",
+						Interval: 1,
+						Limit:    0,
+						Rules: []Rule{
+							AlertingRulePaginated{
+								State:       "inactive",
+								Name:        "test_metric3",
+								Query:       "absent(test_metric3) != 1",
+								Duration:    1,
+								Labels:      labels.Labels{},
+								Annotations: labels.Labels{},
+								AlertInfo:   &AlertsPaginated{Alerts: []*Alert{}},
+								Health:      "ok",
+								Type:        "alerting",
+							},
+						},
+					},
+				},
+			},
+			zeroFunc: rulesZeroFunc,
+		},
 		{
 			endpoint: api.queryExemplars,
 			query: url.Values{