Alerting: Fix scheduler to group folders by the unique key (orgID and UID) (#81303)

pkg/services/ngalert/models/alert_rule.go

@@ -488,6 +488,13 @@ func (alertRule *AlertRule) ResourceOrgID() int64 {
 	return alertRule.OrgID
 }
 
+func (alertRule *AlertRule) GetFolderKey() FolderKey {
+	return FolderKey{
+		OrgID: alertRule.OrgID,
+		UID:   alertRule.NamespaceUID,
+	}
+}
+
 // AlertRuleVersion is the model for alert rule versions in unified alerting.
 type AlertRuleVersion struct {
 	ID               int64  `xorm:"pk autoincr 'id'"`
@@ -546,13 +553,22 @@ type CountAlertRulesQuery struct {
 	NamespaceUID string
 }
 
+type FolderKey struct {
+	OrgID int64
+	UID   string
+}
+
+func (f FolderKey) String() string {
+	return fmt.Sprintf("%d:%s", f.OrgID, f.UID)
+}
+
 type GetAlertRulesForSchedulingQuery struct {
 	PopulateFolders bool
 	RuleGroups      []string
 
 	ResultRules []*AlertRule
 	// A map of folder UID to folder Title in NamespaceKey format (see GetNamespaceKey)
-	ResultFoldersTitles map[string]string
+	ResultFoldersTitles map[FolderKey]string
 }
 
 // ListNamespaceAlertRulesQuery is the query for listing namespace alert rules

pkg/services/ngalert/schedule/registry.go

@@ -134,12 +134,12 @@ type evaluation struct {
 
 type alertRulesRegistry struct {
 	rules        map[models.AlertRuleKey]*models.AlertRule
-	folderTitles map[string]string
+	folderTitles map[models.FolderKey]string
 	mu           sync.Mutex
 }
 
 // all returns all rules in the registry.
-func (r *alertRulesRegistry) all() ([]*models.AlertRule, map[string]string) {
+func (r *alertRulesRegistry) all() ([]*models.AlertRule, map[models.FolderKey]string) {
 	r.mu.Lock()
 	defer r.mu.Unlock()
 	result := make([]*models.AlertRule, 0, len(r.rules))
@@ -156,7 +156,7 @@ func (r *alertRulesRegistry) get(k models.AlertRuleKey) *models.AlertRule {
 }
 
 // set replaces all rules in the registry. Returns difference between previous and the new current version of the registry
-func (r *alertRulesRegistry) set(rules []*models.AlertRule, folders map[string]string) diff {
+func (r *alertRulesRegistry) set(rules []*models.AlertRule, folders map[models.FolderKey]string) diff {
 	r.mu.Lock()
 	defer r.mu.Unlock()
 	rulesMap := make(map[models.AlertRuleKey]*models.AlertRule)

pkg/services/ngalert/schedule/registry_test.go

@@ -239,7 +239,7 @@ func TestSchedulableAlertRulesRegistry(t *testing.T) {
 	assert.Len(t, rules, 0)
 	assert.Len(t, folders, 0)
 
-	expectedFolders := map[string]string{"test-uid": "test-title"}
+	expectedFolders := map[models.FolderKey]string{models.FolderKey{OrgID: 1, UID: "test-uid"}: "test-title"}
 	// replace all rules in the registry with foo
 	r.set([]*models.AlertRule{{OrgID: 1, UID: "foo", Version: 1}}, expectedFolders)
 	rules, folders = r.all()
@@ -308,7 +308,7 @@ func TestSchedulableAlertRulesRegistry_set(t *testing.T) {
 		for _, rule := range initialRules {
 			newRules = append(newRules, models.CopyRule(rule))
 		}
-		diff := r.set(newRules, map[string]string{})
+		diff := r.set(newRules, map[models.FolderKey]string{})
 		require.Truef(t, diff.IsEmpty(), "Diff is not empty. Probably we check something else than key + version")
 	})
 	t.Run("should return empty diff if version does not change", func(t *testing.T) {
@@ -324,7 +324,7 @@ func TestSchedulableAlertRulesRegistry_set(t *testing.T) {
 			newRules = append(newRules, rule)
 		}
 
-		diff := r.set(newRules, map[string]string{})
+		diff := r.set(newRules, map[models.FolderKey]string{})
 		require.Truef(t, diff.IsEmpty(), "Diff is not empty. Probably we check something else than key + version")
 	})
 	t.Run("should return key in diff if version changes", func(t *testing.T) {
@@ -340,7 +340,7 @@ func TestSchedulableAlertRulesRegistry_set(t *testing.T) {
 		}
 		require.NotEmptyf(t, expectedUpdated, "Input parameters have changed. Nothing to assert")
 
-		diff := r.set(newRules, map[string]string{})
+		diff := r.set(newRules, map[models.FolderKey]string{})
 		require.Falsef(t, diff.IsEmpty(), "Diff is empty but should not be")
 		require.Equal(t, expectedUpdated, diff.updated)
 	})

pkg/services/ngalert/schedule/schedule.go

@@ -301,7 +301,7 @@ func (sch *schedule) processTick(ctx context.Context, dispatcherGroup *errgroup.
 
 		var folderTitle string
 		if !sch.disableGrafanaFolder {
-			title, ok := folderTitles[item.NamespaceUID]
+			title, ok := folderTitles[item.GetFolderKey()]
 			if ok {
 				folderTitle = title
 			} else {

pkg/services/ngalert/schedule/schedule_unit_test.go

@@ -586,7 +586,7 @@ func TestSchedule_ruleRoutine(t *testing.T) {
 
 		sch, ruleStore, _, _ := createSchedule(evalAppliedChan, &sender)
 		ruleStore.PutRule(context.Background(), rule)
-		sch.schedulableAlertRules.set([]*models.AlertRule{rule}, map[string]string{rule.NamespaceUID: folderTitle})
+		sch.schedulableAlertRules.set([]*models.AlertRule{rule}, map[models.FolderKey]string{rule.GetFolderKey(): folderTitle})
 
 		go func() {
 			ctx, cancel := context.WithCancel(context.Background())

pkg/services/ngalert/schedule/testing.go

@@ -57,10 +57,11 @@ func (f *fakeRulesStore) GetAlertRulesKeysForScheduling(ctx context.Context) ([]
 }
 
 func (f *fakeRulesStore) GetAlertRulesForScheduling(ctx context.Context, query *models.GetAlertRulesForSchedulingQuery) error {
-	query.ResultFoldersTitles = map[string]string{}
+	query.ResultFoldersTitles = map[models.FolderKey]string{}
 	for _, rule := range f.rules {
 		query.ResultRules = append(query.ResultRules, rule)
-		query.ResultFoldersTitles[rule.NamespaceUID] = f.getNamespaceTitle(rule.NamespaceUID)
+		key := models.FolderKey{OrgID: rule.OrgID, UID: rule.UID}
+		query.ResultFoldersTitles[key] = f.getNamespaceTitle(rule.NamespaceUID)
 	}
 	return nil
 }

pkg/services/ngalert/store/alert_rule.go

@@ -514,6 +514,7 @@ func (st DBstore) GetAlertRulesKeysForScheduling(ctx context.Context) ([]ngmodel
 // GetAlertRulesForScheduling returns a short version of all alert rules except those that belong to an excluded list of organizations
 func (st DBstore) GetAlertRulesForScheduling(ctx context.Context, query *ngmodels.GetAlertRulesForSchedulingQuery) error {
 	var folders []struct {
+		OrgId     int64
 		Uid       string
 		Title     string
 		ParentUid string
@@ -565,7 +566,7 @@ func (st DBstore) GetAlertRulesForScheduling(ctx context.Context, query *ngmodel
 		query.ResultRules = rules
 
 		if query.PopulateFolders {
-			foldersSql := sess.Table("folder").Alias("d").Select("d.uid, d.title, d.parent_uid").
+			foldersSql := sess.Table("folder").Alias("d").Select("d.org_id, d.uid, d.title, d.parent_uid").
 				Where(`EXISTS (SELECT 1 FROM alert_rule a WHERE d.uid = a.namespace_uid AND d.org_id = a.org_id)`)
 			if len(disabledOrgs) > 0 {
 				foldersSql.NotIn("org_id", disabledOrgs)
@@ -574,9 +575,13 @@ func (st DBstore) GetAlertRulesForScheduling(ctx context.Context, query *ngmodel
 			if err := foldersSql.Find(&folders); err != nil {
 				return fmt.Errorf("failed to fetch a list of folders that contain alert rules: %w", err)
 			}
-			query.ResultFoldersTitles = make(map[string]string, len(folders))
+			query.ResultFoldersTitles = make(map[ngmodels.FolderKey]string, len(folders))
-			for _, folder := range folders {
+			for _, f := range folders {
-				query.ResultFoldersTitles[folder.Uid] = ngmodels.GetNamespaceKey(folder.ParentUid, folder.Title)
+				key := ngmodels.FolderKey{
+					OrgID: f.OrgId,
+					UID:   f.Uid,
+				}
+				query.ResultFoldersTitles[key] = ngmodels.GetNamespaceKey(f.ParentUid, f.Title)
 			}
 		}
 		return nil

pkg/services/ngalert/store/alert_rule_test.go

@@ -352,7 +352,7 @@ func TestIntegration_GetAlertRulesForScheduling(t *testing.T) {
 		rules        []string
 		ruleGroups   []string
 		disabledOrgs []int64
-		folders      map[string]string
+		folders      map[models.FolderKey]string
 	}{
 		{
 			name:  "without a rule group filter, it returns all created rules",
@@ -371,13 +371,13 @@ func TestIntegration_GetAlertRulesForScheduling(t *testing.T) {
 		{
 			name:    "with populate folders enabled, it returns them",
 			rules:   []string{rule1.Title, rule2.Title},
-			folders: map[string]string{rule1.NamespaceUID: models.GetNamespaceKey(parentFolderUid, rule1.Title), rule2.NamespaceUID: rule2.Title},
+			folders: map[models.FolderKey]string{rule1.GetFolderKey(): models.GetNamespaceKey(parentFolderUid, rule1.Title), rule2.GetFolderKey(): rule2.Title},
 		},
 		{
 			name:         "with populate folders enabled and a filter on orgs, it only returns selected information",
 			rules:        []string{rule1.Title},
 			disabledOrgs: []int64{rule2.OrgID},
-			folders:      map[string]string{rule1.NamespaceUID: models.GetNamespaceKey(parentFolderUid, rule1.Title)},
+			folders:      map[models.FolderKey]string{rule1.GetFolderKey(): models.GetNamespaceKey(parentFolderUid, rule1.Title)},
 		},
 	}