From 2f0190d775062b1f1e33a8b6323f5fad78ddacb0 Mon Sep 17 00:00:00 2001 From: William Wernert Date: Fri, 1 Aug 2025 12:54:13 -0400 Subject: [PATCH] Alerting: Add store level pagination of rules (#108633) --- pkg/services/ngalert/api/persist.go | 1 + .../ngalert/api/prometheus/api_prometheus.go | 66 ++---- pkg/services/ngalert/models/alert_rule.go | 56 +++++ pkg/services/ngalert/store/alert_rule.go | 180 ++++++++++++++++ pkg/services/ngalert/store/alert_rule_test.go | 197 +++++++++++++++++- pkg/services/ngalert/tests/fakes/rules.go | 87 ++++++++ 6 files changed, 536 insertions(+), 51 deletions(-) diff --git a/pkg/services/ngalert/api/persist.go b/pkg/services/ngalert/api/persist.go index 2150edbe34c..5a60c2e0637 100644 --- a/pkg/services/ngalert/api/persist.go +++ b/pkg/services/ngalert/api/persist.go @@ -23,6 +23,7 @@ type RuleStore interface { GetAlertRuleByUID(ctx context.Context, query *ngmodels.GetAlertRuleByUIDQuery) (*ngmodels.AlertRule, error) GetAlertRulesGroupByRuleUID(ctx context.Context, query *ngmodels.GetAlertRulesGroupByRuleUIDQuery) ([]*ngmodels.AlertRule, error) ListAlertRules(ctx context.Context, query *ngmodels.ListAlertRulesQuery) (ngmodels.RulesGroup, error) + ListAlertRulesByGroup(ctx context.Context, query *ngmodels.ListAlertRulesByGroupQuery) (ngmodels.RulesGroup, string, error) ListDeletedRules(ctx context.Context, orgID int64) ([]*ngmodels.AlertRule, error) // InsertAlertRules will insert all alert rules passed into the function diff --git a/pkg/services/ngalert/api/prometheus/api_prometheus.go b/pkg/services/ngalert/api/prometheus/api_prometheus.go index 2e753a6db1d..c223bce9aac 100644 --- a/pkg/services/ngalert/api/prometheus/api_prometheus.go +++ b/pkg/services/ngalert/api/prometheus/api_prometheus.go @@ -2,7 +2,6 @@ package api import ( "context" - "encoding/base64" "encoding/json" "errors" "fmt" @@ -30,7 +29,7 @@ import ( type RuleStoreReader interface { GetUserVisibleNamespaces(context.Context, int64, identity.Requester) (map[string]*folder.Folder, error) - ListAlertRules(ctx context.Context, query *ngmodels.ListAlertRulesQuery) (ngmodels.RulesGroup, error) + ListAlertRulesStore } type RuleGroupAccessControlService interface { @@ -241,7 +240,7 @@ type RuleGroupStatusesOptions struct { } type ListAlertRulesStore interface { - ListAlertRules(ctx context.Context, query *ngmodels.ListAlertRulesQuery) (ngmodels.RulesGroup, error) + ListAlertRulesByGroup(ctx context.Context, query *ngmodels.ListAlertRulesByGroupQuery) (ngmodels.RulesGroup, string, error) } func (srv PrometheusSrv) RouteGetRuleStatuses(c *contextmodel.ReqContext) response.Response { @@ -476,15 +475,24 @@ func PrepareRuleGroupStatuses(log log.Logger, store ListAlertRulesStore, opts Ru receiverName := opts.Query.Get("receiver_name") - alertRuleQuery := ngmodels.ListAlertRulesQuery{ - OrgID: opts.OrgID, - NamespaceUIDs: namespaceUIDs, - DashboardUID: dashboardUID, - PanelID: panelID, - RuleGroups: ruleGroups, - ReceiverName: receiverName, + maxGroups := getInt64WithDefault(opts.Query, "group_limit", -1) + nextToken := opts.Query.Get("group_next_token") + + if maxGroups == 0 { + return ruleResponse } - ruleList, err := store.ListAlertRules(opts.Ctx, &alertRuleQuery) + + byGroupQuery := ngmodels.ListAlertRulesByGroupQuery{ + OrgID: opts.OrgID, + GroupLimit: maxGroups, + GroupContinueToken: nextToken, + NamespaceUIDs: namespaceUIDs, + DashboardUID: dashboardUID, + PanelID: panelID, + RuleGroups: ruleGroups, + ReceiverName: receiverName, + } + ruleList, continueToken, err := store.ListAlertRulesByGroup(opts.Ctx, &byGroupQuery) if err != nil { ruleResponse.Status = "error" ruleResponse.Error = fmt.Sprintf("failure getting rules: %s", err.Error()) @@ -498,31 +506,9 @@ func PrepareRuleGroupStatuses(log log.Logger, store ListAlertRulesStore, opts Ru ruleNamesSet[rn] = struct{}{} } - maxGroups := getInt64WithDefault(opts.Query, "group_limit", -1) - nextToken := opts.Query.Get("group_next_token") - if nextToken != "" { - if _, err := base64.URLEncoding.DecodeString(nextToken); err != nil { - nextToken = "" - } - } - groupedRules := getGroupedRules(log, ruleList, ruleNamesSet, opts.AllowedNamespaces) rulesTotals := make(map[string]int64, len(groupedRules)) - var newToken string - foundToken := false for _, rg := range groupedRules { - if nextToken != "" && !foundToken { - if !tokenGreaterThanOrEqual(getRuleGroupNextToken(rg.Folder, rg.GroupKey.RuleGroup), nextToken) { - continue - } - foundToken = true - } - - if maxGroups > -1 && len(ruleResponse.Data.RuleGroups) == int(maxGroups) { - newToken = getRuleGroupNextToken(rg.Folder, rg.GroupKey.RuleGroup) - break - } - ruleGroup, totals := toRuleGroup(log, rg.GroupKey, rg.Folder, rg.Rules, provenanceRecords, limitAlertsPerRule, stateFilterSet, matchers, labelOptions, ruleStatusMutator, alertStateMutator) ruleGroup.Totals = totals for k, v := range totals { @@ -546,7 +532,7 @@ func PrepareRuleGroupStatuses(log log.Logger, store ListAlertRulesStore, opts Ru } } - ruleResponse.Data.NextToken = newToken + ruleResponse.Data.NextToken = continueToken // Only return Totals if there is no pagination if maxGroups == -1 { @@ -556,18 +542,6 @@ func PrepareRuleGroupStatuses(log log.Logger, store ListAlertRulesStore, opts Ru return ruleResponse } -func getRuleGroupNextToken(namespace, group string) string { - return base64.URLEncoding.EncodeToString([]byte(namespace + "/" + group)) -} - -// Returns true if tokenA >= tokenB -func tokenGreaterThanOrEqual(tokenA string, tokenB string) bool { - decodedTokenA, _ := base64.URLEncoding.DecodeString(tokenA) - decodedTokenB, _ := base64.URLEncoding.DecodeString(tokenB) - - return string(decodedTokenA) >= string(decodedTokenB) -} - type ruleGroup struct { Folder string GroupKey ngmodels.AlertRuleGroupKey diff --git a/pkg/services/ngalert/models/alert_rule.go b/pkg/services/ngalert/models/alert_rule.go index b4a4a94eaba..b1f415f929d 100644 --- a/pkg/services/ngalert/models/alert_rule.go +++ b/pkg/services/ngalert/models/alert_rule.go @@ -2,6 +2,7 @@ package models import ( "context" + "encoding/base64" "encoding/json" "errors" "fmt" @@ -848,6 +849,61 @@ type GetAlertRulesGroupByRuleUIDQuery struct { OrgID int64 } +type RuleTypeFilter int + +const ( + RuleTypeFilterAll RuleTypeFilter = iota + RuleTypeFilterAlerting + RuleTypeFilterRecording +) + +type ListAlertRulesByGroupQuery struct { + OrgID int64 + RuleUIDs []string + NamespaceUIDs []string + ExcludeOrgs []int64 + RuleGroups []string + + // DashboardUID and PanelID are optional and allow filtering rules + // to return just those for a dashboard and panel. + DashboardUID string + PanelID int64 + + ReceiverName string + TimeIntervalName string + + HasPrometheusRuleDefinition *bool + + RuleType RuleTypeFilter + + GroupLimit int64 // Number of groups to fetch + GroupContinueToken string // Token for per-group pagination +} + +type GroupCursor struct { + NamespaceUID string `json:"n"` + RuleGroup string `json:"g"` +} + +func EncodeGroupCursor(c GroupCursor) string { + data, _ := json.Marshal(c) + return base64.URLEncoding.EncodeToString(data) +} + +func DecodeGroupCursor(token string) (GroupCursor, error) { + var c GroupCursor + data, err := base64.URLEncoding.DecodeString(token) + if err != nil { + return c, fmt.Errorf("failed to decode group token: %w", err) + } + + if err := json.Unmarshal(data, &c); err != nil { + return c, fmt.Errorf("failed to unmarshal group cursor: %w", err) + } + + return c, nil +} + // ListAlertRulesQuery is the query for listing alert rules type ListAlertRulesQuery struct { OrgID int64 diff --git a/pkg/services/ngalert/store/alert_rule.go b/pkg/services/ngalert/store/alert_rule.go index c535bbe4379..ab8676c754d 100644 --- a/pkg/services/ngalert/store/alert_rule.go +++ b/pkg/services/ngalert/store/alert_rule.go @@ -582,6 +582,186 @@ func (st DBstore) CountInFolders(ctx context.Context, orgID int64, folderUIDs [] return count, err } +func (st DBstore) ListAlertRulesByGroup(ctx context.Context, query *ngmodels.ListAlertRulesByGroupQuery) (result ngmodels.RulesGroup, nextToken string, err error) { + err = st.SQLStore.WithDbSession(ctx, func(sess *db.Session) error { + q := sess.Table("alert_rule") + + if query.OrgID >= 0 { + q = q.Where("org_id = ?", query.OrgID) + } + + if query.DashboardUID != "" { + q = q.Where("dashboard_uid = ?", query.DashboardUID) + if query.PanelID != 0 { + q = q.Where("panel_id = ?", query.PanelID) + } + } + + if len(query.NamespaceUIDs) > 0 { + args, in := getINSubQueryArgs(query.NamespaceUIDs) + q = q.Where(fmt.Sprintf("namespace_uid IN (%s)", strings.Join(in, ",")), args...) + } + + if len(query.RuleUIDs) > 0 { + args, in := getINSubQueryArgs(query.RuleUIDs) + q = q.Where(fmt.Sprintf("uid IN (%s)", strings.Join(in, ",")), args...) + } + + var groupsMap map[string]struct{} + if len(query.RuleGroups) > 0 { + groupsMap = make(map[string]struct{}) + args, in := getINSubQueryArgs(query.RuleGroups) + q = q.Where(fmt.Sprintf("rule_group IN (%s)", strings.Join(in, ",")), args...) + for _, group := range query.RuleGroups { + groupsMap[group] = struct{}{} + } + } + + if query.ReceiverName != "" { + q, err = st.filterByContentInNotificationSettings(query.ReceiverName, q) + if err != nil { + return err + } + } + + if query.TimeIntervalName != "" { + q, err = st.filterByContentInNotificationSettings(query.TimeIntervalName, q) + if err != nil { + return err + } + } + + if query.HasPrometheusRuleDefinition != nil { + q, err = st.filterWithPrometheusRuleDefinition(*query.HasPrometheusRuleDefinition, q) + if err != nil { + return err + } + } + + switch query.RuleType { + case ngmodels.RuleTypeFilterAlerting: + q = q.Where("record = ''") + case ngmodels.RuleTypeFilterRecording: + q = q.Where("record != ''") + case ngmodels.RuleTypeFilterAll: + // no additional filter + default: + return fmt.Errorf("unknown rule type filter %q", query.RuleType) + } + + // Order by group first, then by rule index within group + q = q.Asc("namespace_uid", "rule_group", "rule_group_idx", "id") + + var cursor ngmodels.GroupCursor + if query.GroupContinueToken != "" { + // only set the cursor if it's valid, otherwise we'll start from the beginning + if cur, err := ngmodels.DecodeGroupCursor(query.GroupContinueToken); err == nil { + cursor = cur + } + } + + // Build group cursor condition + if cursor.NamespaceUID != "" { + q = buildGroupCursorCondition(q, cursor) + } + + // No arbitrary fetch limit - let the loop control pagination + alertRules := make([]*ngmodels.AlertRule, 0) + rule := new(alertRule) + rows, err := q.Rows(rule) + if err != nil { + return err + } + defer func() { + _ = rows.Close() + }() + + // Process rules and implement per-group pagination + var groupsFetched int64 + for rows.Next() { + rule := new(alertRule) + err = rows.Scan(rule) + if err != nil { + st.Logger.Error("Invalid rule found in DB store, ignoring it", "func", "ListAlertRulesByGroup", "error", err) + continue + } + + converted, err := alertRuleToModelsAlertRule(*rule, st.Logger) + if err != nil { + st.Logger.Error("Invalid rule found in DB store, cannot convert, ignoring it", "func", "ListAlertRulesByGroup", "error", err) + continue + } + + // Check if we've moved to a new group + key := ngmodels.GroupCursor{ + NamespaceUID: converted.NamespaceUID, + RuleGroup: converted.RuleGroup, + } + if key != cursor { + // Check if we've reached the group limit + if query.GroupLimit > 0 && groupsFetched == query.GroupLimit { + // Generate next token for the next group + nextToken = ngmodels.EncodeGroupCursor(cursor) + break + } + + // Reset for new group + cursor = key + groupsFetched++ + } + + // Apply post-query filters + if !shouldIncludeRule(&converted, query, groupsMap) { + continue + } + + alertRules = append(alertRules, &converted) + } + + result = alertRules + return nil + }) + return result, nextToken, err +} + +func buildGroupCursorCondition(sess *xorm.Session, c ngmodels.GroupCursor) *xorm.Session { + return sess.Where("(namespace_uid > ?)", c.NamespaceUID). + Or("(namespace_uid = ? AND rule_group > ?)", c.NamespaceUID, c.RuleGroup) +} + +func shouldIncludeRule(rule *ngmodels.AlertRule, query *ngmodels.ListAlertRulesByGroupQuery, groupsMap map[string]struct{}) bool { + if query.ReceiverName != "" { + if !slices.ContainsFunc(rule.NotificationSettings, func(settings ngmodels.NotificationSettings) bool { + return settings.Receiver == query.ReceiverName + }) { + return false + } + } + + if query.TimeIntervalName != "" { + if !slices.ContainsFunc(rule.NotificationSettings, func(settings ngmodels.NotificationSettings) bool { + return slices.Contains(settings.MuteTimeIntervals, query.TimeIntervalName) || + slices.Contains(settings.ActiveTimeIntervals, query.TimeIntervalName) + }) { + return false + } + } + + if query.HasPrometheusRuleDefinition != nil { + if *query.HasPrometheusRuleDefinition != rule.HasPrometheusRuleDefinition() { + return false + } + } + + if groupsMap != nil { + if _, ok := groupsMap[rule.RuleGroup]; !ok { + return false + } + } + + return true +} + // ListAlertRules is a handler for retrieving alert rules of specific organisation. func (st DBstore) ListAlertRules(ctx context.Context, query *ngmodels.ListAlertRulesQuery) (result ngmodels.RulesGroup, err error) { err = st.SQLStore.WithDbSession(ctx, func(sess *db.Session) error { diff --git a/pkg/services/ngalert/store/alert_rule_test.go b/pkg/services/ngalert/store/alert_rule_test.go index 74f51fd305a..a512b5949b4 100644 --- a/pkg/services/ngalert/store/alert_rule_test.go +++ b/pkg/services/ngalert/store/alert_rule_test.go @@ -1596,14 +1596,14 @@ func TestIntegrationGetRuleVersions(t *testing.T) { // createAlertRule creates an alert rule in the database and returns it. // If a generator is not specified, uniqueness of primary key is not guaranteed. -func createRule(t *testing.T, store *DBstore, generator *models.AlertRuleGenerator) *models.AlertRule { - t.Helper() +func createRule(tb testing.TB, store *DBstore, generator *models.AlertRuleGenerator) *models.AlertRule { + tb.Helper() if generator == nil { generator = models.RuleGen.With(models.RuleMuts.WithIntervalMatching(store.Cfg.BaseInterval)) } rule := generator.GenerateRef() converted, err := alertRuleFromModelsAlertRule(*rule) - require.NoError(t, err) + require.NoError(tb, err) err = store.SQLStore.WithDbSession(context.Background(), func(sess *db.Session) error { converted.ID = 0 _, err := sess.Table(alertRule{}).InsertOne(&converted) @@ -1622,12 +1622,12 @@ func createRule(t *testing.T, store *DBstore, generator *models.AlertRuleGenerat rule = &r return err }) - require.NoError(t, err) + require.NoError(tb, err) return rule } -func setupFolderService(t *testing.T, sqlStore db.DB, cfg *setting.Cfg, features featuremgmt.FeatureToggles) folder.Service { +func setupFolderService(t testing.TB, sqlStore db.DB, cfg *setting.Cfg, features featuremgmt.FeatureToggles) folder.Service { tracer := tracing.InitializeTracerForTest() inProcBus := bus.ProvideBus(tracer) folderStore := folderimpl.ProvideDashboardFolderStore(sqlStore) @@ -1735,6 +1735,169 @@ func TestIntegration_AlertRuleVersionsCleanup(t *testing.T) { }) } +func TestIntegration_ListAlertRulesByGroup(t *testing.T) { + if testing.Short() { + t.Skip("skipping integration test") + } + + sqlStore := db.InitTestDB(t) + cfg := setting.NewCfg() + cfg.UnifiedAlerting = setting.UnifiedAlertingSettings{ + BaseInterval: time.Duration(rand.Int64N(100)+1) * time.Second, + } + folderService := setupFolderService(t, sqlStore, cfg, featuremgmt.WithFeatures()) + bus := &fakeBus{} + orgID := int64(1) + ruleGen := models.RuleGen.With( + models.RuleMuts.WithIntervalMatching(cfg.UnifiedAlerting.BaseInterval), + models.RuleMuts.WithOrgID(orgID), + ) + store := createTestStore(sqlStore, folderService, &logtest.Fake{}, cfg.UnifiedAlerting, bus) + + // set test params + numFolders := 10 + numRules := 50 + rulesPerGroup := 5 + totalGroups := numRules / rulesPerGroup // 10 + + // create rules with different group names + rules, _ := createManyRules(t, + store, + ruleGen, + numFolders, + numRules, + rulesPerGroup, + ) + // sort rules by folder, then group, then group index + slices.SortStableFunc(rules, func(a, b *models.AlertRule) int { + if a.NamespaceUID != b.NamespaceUID { + return strings.Compare(a.NamespaceUID, b.NamespaceUID) + } + if a.RuleGroup != b.RuleGroup { + return strings.Compare(a.RuleGroup, b.RuleGroup) + } + return a.RuleGroupIndex - b.RuleGroupIndex + }) + + t.Run("should return all rules when no limit passed", func(t *testing.T) { + result, continueToken, err := store.ListAlertRulesByGroup(context.Background(), &models.ListAlertRulesByGroupQuery{ + OrgID: orgID, + }) + require.NoError(t, err) + require.Len(t, result, 50, "should return all rules when no limit is set") + require.Empty(t, continueToken, "continue token should be empty when no limit is set") + }) + + t.Run("should return paginated results when group limit is set", func(t *testing.T) { + // random number from 1 to totalGroups - 1 (to ensure we always receive less than totalGroups) + groupLimit := rand.Int64N(int64(totalGroups)-1) + 1 + result, continueToken, err := store.ListAlertRulesByGroup(context.Background(), &models.ListAlertRulesByGroupQuery{ + OrgID: orgID, + GroupLimit: groupLimit, + }) + require.NoError(t, err) + expectedRuleCount := groupLimit * int64(rulesPerGroup) + require.Len(t, result, int(expectedRuleCount), fmt.Sprintf("should return %d rules when group limit is set", expectedRuleCount)) + require.NotEmpty(t, continueToken, "continue token should not be empty when limit is set") + }) + + t.Run("pagination should all for continuation", func(t *testing.T) { + groupLimit := int64(2) // fixed group limit for this test + result, continueToken, err := store.ListAlertRulesByGroup(context.Background(), &models.ListAlertRulesByGroupQuery{ + OrgID: orgID, + GroupLimit: groupLimit, + }) + require.NoError(t, err) + require.Len(t, result, int(groupLimit*int64(rulesPerGroup)), "should return rules for the first two groups") + require.NotEmpty(t, continueToken, "continue token should not be empty") + + for i, rule := range result { + expected := rules[i].RuleGroup + actual := rule.RuleGroup + require.Equal(t, expected, actual, "rules should be ordered by group name") + } + + resultRules := make([]*models.AlertRule, 0, len(result)) + resultRules = append(resultRules, result...) + + // Continue from previous, fetching the rest of the rules + result, continueToken, err = store.ListAlertRulesByGroup(context.Background(), &models.ListAlertRulesByGroupQuery{ + OrgID: orgID, + GroupContinueToken: continueToken, + }) + require.NoError(t, err) + resultRules = append(resultRules, result...) + require.Len(t, resultRules, numRules, "should return all rules when continuing from the last token") + require.Empty(t, continueToken, "continue token should be empty when all rules are fetched") + for i, rule := range resultRules { + expected := rules[i].RuleGroup + actual := rule.RuleGroup + require.Equal(t, expected, actual, "rules should be ordered by group name") + } + }) +} + +func Benchmark_ListAlertRules(b *testing.B) { + orgID := int64(1) + ruleGen := models.RuleGen + + // init + sqlStore := db.InitTestDB(b) + cfg := setting.NewCfg() + cfg.UnifiedAlerting = setting.UnifiedAlertingSettings{ + BaseInterval: time.Duration(rand.Int64N(100)) * time.Second, + } + folderService := setupFolderService(b, sqlStore, cfg, featuremgmt.WithFeatures()) + bus := &fakeBus{} + store := createTestStore(sqlStore, folderService, &logtest.Fake{}, cfg.UnifiedAlerting, bus) + + ruleGen = ruleGen.With( + ruleGen.WithIntervalMatching(cfg.UnifiedAlerting.BaseInterval), + ruleGen.WithOrgID(orgID), + ) + + // define benchmark parameters + numFolders := 5 + numRules := 10000 + rulesPerGroup := 100 + assert.Greater(b, numRules, rulesPerGroup, "n must be greater than rulesPerGroup") + assert.Equal(b, 0, numRules%rulesPerGroup, "n % rulesPerGroup must be zero to create equal groups") + + // create rules and folders (5 folders, each with n/rulesPerGroup rules) + _, _ = createManyRules(b, + store, + ruleGen, + numFolders, // number of folders + numRules, // total number of rules + rulesPerGroup, // rules per group + ) + + b.Run(fmt.Sprintf("list %d rules unpaginated", numRules), func(b *testing.B) { + for b.Loop() { + _, err := store.ListAlertRules(context.Background(), &models.ListAlertRulesQuery{ + OrgID: orgID, + }) + if err != nil { + b.Fatal(err) + } + } + }) + + for _, groupLimit := range []int{1, 2, 5, 10, 50, 100} { + b.Run(fmt.Sprintf("list %d groups paginated", groupLimit), func(b *testing.B) { + for b.Loop() { + _, _, err := store.ListAlertRulesByGroup(context.Background(), &models.ListAlertRulesByGroupQuery{ + OrgID: orgID, + GroupLimit: int64(groupLimit), + }) + if err != nil { + b.Fatal(err) + } + } + }) + } +} + func TestIntegration_ListAlertRules(t *testing.T) { if testing.Short() { t.Skip("skipping integration test") @@ -1975,3 +2138,27 @@ func (f *fakeBus) Publish(ctx context.Context, msg bus.Msg) error { return nil } + +func createManyRules(tb testing.TB, store *DBstore, ruleGen *models.AlertRuleGenerator, numFolders, numRules, rulesPerGroup int) ([]*models.AlertRule, []string) { + tb.Helper() + + require.Greater(tb, numRules, 0, "numRules must be greater than 0") + require.Greater(tb, numFolders, 0, "numFolders must be greater than 0") + require.Greater(tb, numRules, rulesPerGroup, "numRules must be greater than rulesPerGroup") + require.Greater(tb, rulesPerGroup, 0, "rulesPerGroup must be greater than 0") + require.Equal(tb, numRules%rulesPerGroup, 0, "numRules % rulesPerGroup must be zero to create equal groups") + + rules := make([]*models.AlertRule, 0, numRules) + namespaceUIDs := make([]string, numFolders) + for i := range namespaceUIDs { + namespaceUIDs[i] = fmt.Sprintf("ns-%d", i) + } + for i := 0; i < numRules; i++ { + gen := ruleGen.With( + ruleGen.WithNamespaceUID(namespaceUIDs[i%numFolders]), + ruleGen.WithGroupName(fmt.Sprintf("group_%d", i%(numRules/rulesPerGroup))), + ) + rules = append(rules, createRule(tb, store, gen)) + } + return rules, namespaceUIDs +} diff --git a/pkg/services/ngalert/tests/fakes/rules.go b/pkg/services/ngalert/tests/fakes/rules.go index 046522f3214..c17c70bf6c9 100644 --- a/pkg/services/ngalert/tests/fakes/rules.go +++ b/pkg/services/ngalert/tests/fakes/rules.go @@ -4,6 +4,7 @@ import ( "context" "math/rand" "slices" + "strings" "sync" "testing" "time" @@ -185,6 +186,88 @@ func (f *RuleStore) GetAlertRulesGroupByRuleUID(_ context.Context, q *models.Get return ruleList, nil } +func (f *RuleStore) ListAlertRulesByGroup(_ context.Context, q *models.ListAlertRulesByGroupQuery) (models.RulesGroup, string, error) { + f.mtx.Lock() + defer f.mtx.Unlock() + f.RecordedOps = append(f.RecordedOps, *q) + + if err := f.Hook(*q); err != nil { + return nil, "", err + } + + query := &models.ListAlertRulesQuery{ + OrgID: q.OrgID, + NamespaceUIDs: q.NamespaceUIDs, + DashboardUID: q.DashboardUID, + PanelID: q.PanelID, + RuleGroups: q.RuleGroups, + RuleUIDs: q.RuleUIDs, + ReceiverName: q.ReceiverName, + HasPrometheusRuleDefinition: q.HasPrometheusRuleDefinition, + } + + ruleList, err := f.listAlertRules(query) + if err != nil { + return nil, "", err + } + + // < group limit logic > + + // sort rules to ensure order is consistent, pagination depends on this + slices.SortFunc(ruleList, func(a, b *models.AlertRule) int { + nsCmp := strings.Compare(a.NamespaceUID, b.NamespaceUID) + if nsCmp != 0 { + return nsCmp + } + rgCmp := strings.Compare(a.RuleGroup, b.RuleGroup) + if rgCmp != 0 { + return rgCmp + } + return models.RulesGroupComparer(a, b) + }) + + var nextToken string + var cursor models.GroupCursor + if q.GroupContinueToken != "" { + if cur, err := models.DecodeGroupCursor(q.GroupContinueToken); err == nil { + cursor = cur + } + } + + if q.GroupLimit < 0 { + return ruleList, "", nil + } + + outputRules := make([]*models.AlertRule, 0, len(ruleList)) + var groupsFetched int64 + initialCursor := cursor + for _, r := range ruleList { + // skip rules before the initial cursor + if initialCursor.NamespaceUID != "" && + (strings.Compare(r.NamespaceUID, initialCursor.NamespaceUID) < 0 || + (strings.Compare(r.NamespaceUID, initialCursor.NamespaceUID) == 0 && strings.Compare(r.RuleGroup, initialCursor.RuleGroup) <= 0)) { + continue + } + + key := models.GroupCursor{ + NamespaceUID: r.NamespaceUID, + RuleGroup: r.RuleGroup, + } + if key != cursor { + if q.GroupLimit > 0 && groupsFetched == q.GroupLimit { + nextToken = models.EncodeGroupCursor(cursor) + break + } + cursor = key + groupsFetched++ + } + + outputRules = append(outputRules, r) + } + + return outputRules, nextToken, nil +} + func (f *RuleStore) ListAlertRules(_ context.Context, q *models.ListAlertRulesQuery) (models.RulesGroup, error) { f.mtx.Lock() defer f.mtx.Unlock() @@ -194,6 +277,10 @@ func (f *RuleStore) ListAlertRules(_ context.Context, q *models.ListAlertRulesQu return nil, err } + return f.listAlertRules(q) +} + +func (f *RuleStore) listAlertRules(q *models.ListAlertRulesQuery) (models.RulesGroup, error) { hasDashboard := func(r *models.AlertRule, dashboardUID string, panelID int64) bool { if dashboardUID != "" { if r.DashboardUID == nil || *r.DashboardUID != dashboardUID {