Alerting: Add pagination to /api/prometheus/grafana/api/v1/rules (#95959)

* Intermediate step before refactoring

* Sort groups to paginate on them

* Formatting and improved test

* Address comments

* Update tests

pkg/services/ngalert/api/api_prometheus.go

@@ -2,10 +2,12 @@ package api
 
 import (
 	"context"
+	"encoding/base64"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"net/url"
+	"slices"
 	"sort"
 	"strconv"
 	"strings"
@@ -264,7 +266,6 @@ func PrepareRuleGroupStatuses(log log.Logger, manager state.AlertInstanceManager
 		return ruleResponse
 	}
 
-	limitGroups := getInt64WithDefault(opts.Query, "limit", -1)
 	limitRulesPerGroup := getInt64WithDefault(opts.Query, "limit_rules", -1)
 	limitAlertsPerRule := getInt64WithDefault(opts.Query, "limit_alerts", -1)
 	matchers, err := getMatchersFromQuery(opts.Query)
@@ -331,15 +332,15 @@ func PrepareRuleGroupStatuses(log log.Logger, manager state.AlertInstanceManager
 		ruleNamesSet[rn] = struct{}{}
 	}
 
-	groupedRules := getGroupedRules(ruleList, ruleNamesSet)
+	maxGroups := getInt64WithDefault(opts.Query, "group_limit", -1)
+	nextToken := opts.Query.Get("group_next_token")
+
+	groupedRules := getGroupedRules(log, ruleList, ruleNamesSet, opts.Namespaces)
 	rulesTotals := make(map[string]int64, len(groupedRules))
-	for groupKey, rules := range groupedRules {
-		folder, ok := opts.Namespaces[groupKey.NamespaceUID]
-		if !ok {
-			log.Warn("Query returned rules that belong to folder the user does not have access to. All rules that belong to that namespace will not be added to the response", "folder_uid", groupKey.NamespaceUID)
-			continue
-		}
-		ok, err := opts.AuthorizeRuleGroup(rules)
+	var newToken string
+	foundToken := false
+	for _, rg := range groupedRules {
+		ok, err := opts.AuthorizeRuleGroup(rg.Rules)
 		if err != nil {
 			ruleResponse.DiscoveryBase.Status = "error"
 			ruleResponse.DiscoveryBase.Error = fmt.Sprintf("cannot authorize access to rule group: %s", err.Error())
@@ -350,7 +351,19 @@ func PrepareRuleGroupStatuses(log log.Logger, manager state.AlertInstanceManager
 			continue
 		}
 
-		ruleGroup, totals := toRuleGroup(log, manager, status, groupKey, folder, rules, limitAlertsPerRule, withStatesFast, matchers, labelOptions)
+		if nextToken != "" && !foundToken {
+			if nextToken != getRuleGroupNextToken(rg.Folder, rg.GroupKey.RuleGroup) {
+				continue
+			}
+			foundToken = true
+		}
+
+		if maxGroups > -1 && len(ruleResponse.Data.RuleGroups) == int(maxGroups) {
+			newToken = getRuleGroupNextToken(rg.Folder, rg.GroupKey.RuleGroup)
+			break
+		}
+
+		ruleGroup, totals := toRuleGroup(log, manager, status, rg.GroupKey, rg.Folder, rg.Rules, limitAlertsPerRule, withStatesFast, matchers, labelOptions)
 		ruleGroup.Totals = totals
 		for k, v := range totals {
 			rulesTotals[k] += v
@@ -367,18 +380,28 @@ func PrepareRuleGroupStatuses(log log.Logger, manager state.AlertInstanceManager
 		ruleResponse.Data.RuleGroups = append(ruleResponse.Data.RuleGroups, *ruleGroup)
 	}
 
-	ruleResponse.Data.Totals = rulesTotals
+	ruleResponse.Data.NextToken = newToken
 
-	// Sort Rule Groups before checking limits
-	apimodels.RuleGroupsBy(apimodels.RuleGroupsByFileAndName).Sort(ruleResponse.Data.RuleGroups)
-	if limitGroups > -1 && int64(len(ruleResponse.Data.RuleGroups)) >= limitGroups {
-		ruleResponse.Data.RuleGroups = ruleResponse.Data.RuleGroups[0:limitGroups]
+	// Only return Totals if there is no pagination
+	if maxGroups == -1 {
+		ruleResponse.Data.Totals = rulesTotals
 	}
 
 	return ruleResponse
 }
 
-func getGroupedRules(ruleList ngmodels.RulesGroup, ruleNamesSet map[string]struct{}) map[ngmodels.AlertRuleGroupKey][]*ngmodels.AlertRule {
+func getRuleGroupNextToken(namespace, group string) string {
+	return base64.URLEncoding.EncodeToString([]byte(namespace + "/" + group))
+}
+
+type ruleGroup struct {
+	Folder   string
+	GroupKey ngmodels.AlertRuleGroupKey
+	Rules    []*ngmodels.AlertRule
+}
+
+// Returns a slice of rule groups ordered by namespace and group name
+func getGroupedRules(log log.Logger, ruleList ngmodels.RulesGroup, ruleNamesSet map[string]struct{}, namespaceMap map[string]string) []*ruleGroup {
 	// Group rules together by Namespace and Rule Group. Rules are also grouped by Org ID,
 	// but in this API all rules belong to the same organization. Also filter by rule name if
 	// it was provided as a query param.
@@ -394,12 +417,38 @@ func getGroupedRules(ruleList ngmodels.RulesGroup, ruleNamesSet map[string]struc
 		ruleGroup = append(ruleGroup, rule)
 		groupedRules[groupKey] = ruleGroup
 	}
-	// Sort the rules in each rule group by index. We do this at the end instead of
-	// after each append to avoid having to sort each group multiple times.
-	for _, groupRules := range groupedRules {
+
+	ruleGroups := make([]*ruleGroup, 0, len(groupedRules))
+	for groupKey, groupRules := range groupedRules {
+		folder, ok := namespaceMap[groupKey.NamespaceUID]
+		if !ok {
+			log.Warn("Query returned rules that belong to folder the user does not have access to. All rules that belong to that namespace will not be added to the response", "folder_uid", groupKey.NamespaceUID)
+			continue
+		}
+
+		// Sort the rules in each rule group by index. We do this at the end instead of
+		// after each append to avoid having to sort each group multiple times.
 		ngmodels.AlertRulesBy(ngmodels.AlertRulesByIndex).Sort(groupRules)
+
+		ruleGroups = append(ruleGroups, &ruleGroup{
+			Folder:   folder,
+			GroupKey: groupKey,
+			Rules:    groupRules,
+		})
 	}
-	return groupedRules
+
+	// Sort the groups first by namespace, then group name
+	slices.SortFunc(ruleGroups, func(a, b *ruleGroup) int {
+		nsCmp := strings.Compare(a.Folder, b.Folder)
+		if nsCmp != 0 {
+			return nsCmp
+		}
+
+		// If Namespaces are equal, check the group names
+		return strings.Compare(a.GroupKey.RuleGroup, b.GroupKey.RuleGroup)
+	})
+
+	return ruleGroups
 }
 
 func filterRules(ruleGroup *apimodels.RuleGroup, withStatesFast map[eval.State]struct{}) {