Alerting: Allow Rules to Schedule to be filtered by Rule Group (#59990)

* Alerting: Allow Rules to Schedule to be filtered by Rule Group

pkg/services/ngalert/eval/eval.go

@@ -29,7 +29,7 @@ var logger = log.New("ngalert.eval")
 type EvaluatorFactory interface {
 	// Validate validates that the condition is correct. Returns nil if the condition is correct. Otherwise, error that describes the failure
 	Validate(ctx EvaluationContext, condition models.Condition) error
-	// BuildRuleEvaluator build an evaluator pipeline ready to evaluate a rule's query
+	// Create builds an evaluator pipeline ready to evaluate a rule's query
 	Create(ctx EvaluationContext, condition models.Condition) (ConditionEvaluator, error)
 }
 

pkg/services/ngalert/models/alert_rule.go

@@ -398,6 +398,7 @@ type CountAlertRulesQuery struct {
 
 type GetAlertRulesForSchedulingQuery struct {
 	PopulateFolders bool
+	RuleGroups      []string
 
 	ResultRules         []*AlertRule
 	ResultFoldersTitles map[string]string

pkg/services/ngalert/store/alert_rule.go

@@ -421,38 +421,24 @@ func (st DBstore) GetNamespaceByUID(ctx context.Context, uid string, orgID int64
 	return folder, nil
 }
 
-func (st DBstore) getFilterByOrgsString() (string, []interface{}) {
-	if len(st.Cfg.DisabledOrgs) == 0 {
-		return "", nil
-	}
-	builder := strings.Builder{}
-	builder.WriteString("org_id NOT IN(")
-	idx := len(st.Cfg.DisabledOrgs)
-	args := make([]interface{}, 0, len(st.Cfg.DisabledOrgs))
-	for orgId := range st.Cfg.DisabledOrgs {
-		args = append(args, orgId)
-		builder.WriteString("?")
-		idx--
-		if idx == 0 {
-			builder.WriteString(")")
-			break
-		}
-		builder.WriteString(",")
-	}
-	return builder.String(), args
-}
-
 func (st DBstore) GetAlertRulesKeysForScheduling(ctx context.Context) ([]ngmodels.AlertRuleKeyWithVersion, error) {
 	var result []ngmodels.AlertRuleKeyWithVersion
 	err := st.SQLStore.WithDbSession(ctx, func(sess *db.Session) error {
-		alertRulesSql := "SELECT org_id, uid, version FROM alert_rule"
+		alertRulesSql := sess.Table("alert_rule").Select("org_id, uid, version")
-		filter, args := st.getFilterByOrgsString()
+		var disabledOrgs []int64
-		if filter != "" {
+
-			alertRulesSql += " WHERE " + filter
+		for orgID := range st.Cfg.DisabledOrgs {
+			disabledOrgs = append(disabledOrgs, orgID)
 		}
-		if err := sess.SQL(alertRulesSql, args...).Find(&result); err != nil {
+
+		if len(disabledOrgs) > 0 {
+			alertRulesSql = alertRulesSql.NotIn("org_id", disabledOrgs)
+		}
+
+		if err := alertRulesSql.Find(&result); err != nil {
 			return err
 		}
+
 		return nil
 	})
 	return result, err
@@ -466,21 +452,29 @@ func (st DBstore) GetAlertRulesForScheduling(ctx context.Context, query *ngmodel
 	}
 	var rules []*ngmodels.AlertRule
 	return st.SQLStore.WithDbSession(ctx, func(sess *db.Session) error {
-		foldersSql := "SELECT D.uid, D.title FROM dashboard AS D WHERE is_folder IS TRUE AND EXISTS (SELECT 1 FROM alert_rule AS A WHERE D.uid = A.namespace_uid)"
+		var disabledOrgs []int64
-		alertRulesSql := "SELECT * FROM alert_rule"
+		for orgID := range st.Cfg.DisabledOrgs {
-		filter, args := st.getFilterByOrgsString()
+			disabledOrgs = append(disabledOrgs, orgID)
-		if filter != "" {
+		}
-			foldersSql += " AND " + filter
+
-			alertRulesSql += " WHERE " + filter
+		alertRulesSql := sess.Table("alert_rule")
+		if len(disabledOrgs) > 0 {
+			alertRulesSql.NotIn("org_id", disabledOrgs)
+		}
+
+		if len(query.RuleGroups) > 0 {
+			alertRulesSql.In("rule_group", query.RuleGroups)
 		}
 
 		rule := new(ngmodels.AlertRule)
-		rows, err := sess.SQL(alertRulesSql, args...).Rows(rule)
+		rows, err := alertRulesSql.Rows(rule)
 		if err != nil {
 			return fmt.Errorf("failed to fetch alert rules: %w", err)
 		}
 		defer func() {
-			_ = rows.Close()
+			if err := rows.Close(); err != nil {
+				st.Logger.Error("unable to close rows session", "error", err)
+			}
 		}()
 
 		// Deserialize each rule separately in case any of them contain invalid JSON.
@@ -495,8 +489,16 @@ func (st DBstore) GetAlertRulesForScheduling(ctx context.Context, query *ngmodel
 		}
 
 		query.ResultRules = rules
+
 		if query.PopulateFolders {
-			if err := sess.SQL(foldersSql, args...).Find(&folders); err != nil {
+			foldersSql := sess.Table("dashboard").Alias("d").Select("d.uid, d.title").
+				Where("is_folder = ?", st.SQLStore.GetDialect().BooleanStr(true)).
+				And(`EXISTS (SELECT 1 FROM alert_rule a WHERE d.uid = a.namespace_uid)`)
+			if len(disabledOrgs) > 0 {
+				foldersSql.NotIn("org_id", disabledOrgs)
+			}
+
+			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))

pkg/services/ngalert/store/alert_rule_test.go

@@ -7,7 +7,14 @@ import (
 	"testing"
 	"time"
 
-	"github.com/stretchr/testify/assert"
+	"github.com/grafana/grafana/pkg/bus"
+	"github.com/grafana/grafana/pkg/infra/tracing"
+	"github.com/grafana/grafana/pkg/services/folder"
+	"github.com/grafana/grafana/pkg/services/folder/folderimpl"
+	"github.com/grafana/grafana/pkg/services/org"
+	"github.com/grafana/grafana/pkg/services/sqlstore"
+	"github.com/grafana/grafana/pkg/services/user"
+
 	"github.com/stretchr/testify/require"
 	"golang.org/x/exp/rand"
 
@@ -21,12 +28,13 @@ func TestIntegrationUpdateAlertRules(t *testing.T) {
 	if testing.Short() {
 		t.Skip("skipping integration test")
 	}
+	cfg := setting.NewCfg()
+	cfg.UnifiedAlerting = setting.UnifiedAlertingSettings{BaseInterval: time.Duration(rand.Int63n(100)) * time.Second}
 	sqlStore := db.InitTestDB(t)
 	store := &DBstore{
-		SQLStore: sqlStore,
+		SQLStore:      sqlStore,
-		Cfg: setting.UnifiedAlertingSettings{
+		Cfg:           cfg.UnifiedAlerting,
-			BaseInterval: time.Duration(rand.Int63n(100)) * time.Second,
+		FolderService: setupFolderService(t, sqlStore, cfg),
-		},
 	}
 
 	t.Run("should increase version", func(t *testing.T) {
@@ -68,75 +76,112 @@ func TestIntegrationUpdateAlertRules(t *testing.T) {
 	})
 }
 
-func withIntervalMatching(baseInterval time.Duration) func(*models.AlertRule) {
+func TestIntegration_GetAlertRulesForScheduling(t *testing.T) {
-	return func(rule *models.AlertRule) {
-		rule.IntervalSeconds = int64(baseInterval.Seconds()) * rand.Int63n(10)
-		rule.For = time.Duration(rule.IntervalSeconds*rand.Int63n(9)+1) * time.Second
-	}
-}
-
-func TestIntegration_getFilterByOrgsString(t *testing.T) {
 	if testing.Short() {
 		t.Skip("skipping integration test")
 	}
-	testCases := []struct {
+
-		testName       string
+	cfg := setting.NewCfg()
-		orgs           map[int64]struct{}
+	cfg.UnifiedAlerting = setting.UnifiedAlertingSettings{
-		expectedFilter string
+		BaseInterval: time.Duration(rand.Int63n(100)) * time.Second,
-		expectedArgs   []interface{}
+	}
+
+	sqlStore := db.InitTestDB(t)
+	store := &DBstore{
+		SQLStore: sqlStore,
+		Cfg: setting.UnifiedAlertingSettings{
+			BaseInterval: time.Duration(rand.Int63n(100)) * time.Second,
+		},
+		FolderService: setupFolderService(t, sqlStore, cfg),
+	}
+
+	rule1 := createRule(t, store)
+	rule2 := createRule(t, store)
+
+	tc := []struct {
+		name         string
+		rules        []string
+		ruleGroups   []string
+		disabledOrgs []int64
+		folders      map[string]string
 	}{
 		{
-			testName:       "should return empty string if map is empty",
+			name:  "without a rule group filter, it returns all created rules",
-			orgs:           map[int64]struct{}{},
+			rules: []string{rule1.Title, rule2.Title},
-			expectedFilter: "",
-			expectedArgs:   nil,
 		},
 		{
-			testName:       "should return empty string if map is nil",
+			name:       "with a rule group filter, it only returns the rules that match on rule group",
-			orgs:           nil,
+			ruleGroups: []string{rule1.RuleGroup},
-			expectedFilter: "",
+			rules:      []string{rule1.Title},
-			expectedArgs:   nil,
 		},
 		{
-			testName: "should return correct filter if single element",
+			name:         "with a filter on orgs, it returns rules that do not belong to that org",
-			orgs: map[int64]struct{}{
+			rules:        []string{rule1.Title},
-				1: {},
+			disabledOrgs: []int64{rule2.OrgID},
-			},
-			expectedFilter: "org_id NOT IN(?)",
-			expectedArgs:   []interface{}{int64(1)},
 		},
 		{
-			testName: "should return correct filter if many elements",
+			name:    "with populate folders enabled, it returns them",
-			orgs: map[int64]struct{}{
+			rules:   []string{rule1.Title, rule2.Title},
-				1: {},
+			folders: map[string]string{rule1.NamespaceUID: rule1.Title, rule2.NamespaceUID: rule2.Title},
-				2: {},
+		},
-				3: {},
+		{
-			},
+			name:         "with populate folders enabled and a filter on orgs, it only returns selected information",
-			expectedFilter: "org_id NOT IN(?,?,?)",
+			rules:        []string{rule1.Title},
-			expectedArgs:   []interface{}{int64(1), int64(2), int64(3)},
+			disabledOrgs: []int64{rule2.OrgID},
+			folders:      map[string]string{rule1.NamespaceUID: rule1.Title},
 		},
 	}
-	for _, testCase := range testCases {
+
-		t.Run(testCase.testName, func(t *testing.T) {
+	for _, tt := range tc {
-			store := &DBstore{
+		t.Run(tt.name, func(t *testing.T) {
-				Cfg: setting.UnifiedAlertingSettings{
+			if len(tt.disabledOrgs) > 0 {
-					DisabledOrgs: testCase.orgs,
+				store.Cfg.DisabledOrgs = map[int64]struct{}{}
-				},
+
+				for _, orgID := range tt.disabledOrgs {
+					store.Cfg.DisabledOrgs[orgID] = struct{}{}
+					t.Cleanup(func() {
+						delete(store.Cfg.DisabledOrgs, orgID)
+					})
+				}
+			}
+
+			populateFolders := len(tt.folders) > 0
+			query := &models.GetAlertRulesForSchedulingQuery{
+				RuleGroups:      tt.ruleGroups,
+				PopulateFolders: populateFolders,
+			}
+			require.NoError(t, store.GetAlertRulesForScheduling(context.Background(), query))
+			require.Len(t, query.ResultRules, len(tt.rules))
+
+			r := make([]string, 0, len(query.ResultRules))
+			for _, rule := range query.ResultRules {
+				r = append(r, rule.Title)
+			}
+
+			require.ElementsMatch(t, r, tt.rules)
+
+			if populateFolders {
+				require.Equal(t, tt.folders, query.ResultFoldersTitles)
 			}
-			filter, args := store.getFilterByOrgsString()
-			assert.Equal(t, testCase.expectedFilter, filter)
-			assert.ElementsMatch(t, testCase.expectedArgs, args)
 		})
 	}
 }
 
+func withIntervalMatching(baseInterval time.Duration) func(*models.AlertRule) {
+	return func(rule *models.AlertRule) {
+		rule.IntervalSeconds = int64(baseInterval.Seconds()) * rand.Int63n(10)
+		rule.For = time.Duration(rule.IntervalSeconds*rand.Int63n(9)+1) * time.Second
+	}
+}
+
 func TestIntegration_CountAlertRules(t *testing.T) {
 	if testing.Short() {
 		t.Skip("skipping integration test")
 	}
 
 	sqlStore := db.InitTestDB(t)
-	store := &DBstore{SQLStore: sqlStore}
+	cfg := setting.NewCfg()
+	store := &DBstore{SQLStore: sqlStore, FolderService: setupFolderService(t, sqlStore, cfg)}
 	rule := createRule(t, store)
 
 	tests := map[string]struct {
@@ -176,7 +221,9 @@ func TestIntegration_CountAlertRules(t *testing.T) {
 }
 
 func createRule(t *testing.T, store *DBstore) *models.AlertRule {
+	t.Helper()
 	rule := models.AlertRuleGen(withIntervalMatching(store.Cfg.BaseInterval), models.WithUniqueID())()
+	createFolder(t, store, rule.NamespaceUID, rule.Title, rule.OrgID)
 	err := store.SQLStore.WithDbSession(context.Background(), func(sess *db.Session) error {
 		_, err := sess.Table(models.AlertRule{}).InsertOne(rule)
 		if err != nil {
@@ -191,8 +238,41 @@ func createRule(t *testing.T, store *DBstore) *models.AlertRule {
 			return errors.New("cannot read inserted record")
 		}
 		rule = dbRule
+
+		require.NoError(t, err)
+
 		return nil
 	})
 	require.NoError(t, err)
+
 	return rule
 }
+
+func createFolder(t *testing.T, store *DBstore, namespace, title string, orgID int64) {
+	t.Helper()
+	u := &user.SignedInUser{
+		UserID:         1,
+		OrgID:          orgID,
+		OrgRole:        org.RoleAdmin,
+		IsGrafanaAdmin: true,
+	}
+
+	_, err := store.FolderService.Create(context.Background(), &folder.CreateFolderCommand{
+		UID:          namespace,
+		OrgID:        orgID,
+		Title:        title,
+		Description:  "",
+		SignedInUser: u,
+	})
+
+	require.NoError(t, err)
+}
+
+func setupFolderService(t *testing.T, sqlStore *sqlstore.SQLStore, cfg *setting.Cfg) folder.Service {
+	tracer := tracing.InitializeTracerForTest()
+	inProcBus := bus.ProvideBus(tracer)
+	folderStore := folderimpl.ProvideDashboardFolderStore(sqlStore)
+	_, dashboardStore := SetupDashboardService(t, sqlStore, folderStore, cfg)
+
+	return SetupFolderService(t, cfg, dashboardStore, folderStore, inProcBus)
+}

pkg/services/ngalert/store/testing.go

@@ -6,7 +6,25 @@ import (
 	"sync"
 	"testing"
 
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/grafana/pkg/bus"
+	"github.com/grafana/grafana/pkg/services/accesscontrol"
+	acmock "github.com/grafana/grafana/pkg/services/accesscontrol/mock"
+	"github.com/grafana/grafana/pkg/services/dashboards"
+	"github.com/grafana/grafana/pkg/services/dashboards/database"
+	dashboardservice "github.com/grafana/grafana/pkg/services/dashboards/service"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
+	"github.com/grafana/grafana/pkg/services/folder"
+	"github.com/grafana/grafana/pkg/services/folder/folderimpl"
+	"github.com/grafana/grafana/pkg/services/folder/foldertest"
+	"github.com/grafana/grafana/pkg/services/guardian"
 	"github.com/grafana/grafana/pkg/services/ngalert/models"
+	"github.com/grafana/grafana/pkg/services/quota/quotatest"
+	"github.com/grafana/grafana/pkg/services/sqlstore"
+	"github.com/grafana/grafana/pkg/services/tag/tagimpl"
+	"github.com/grafana/grafana/pkg/setting"
 )
 
 func NewFakeImageStore(t *testing.T) *FakeImageStore {
@@ -100,3 +118,45 @@ func (f *FakeAdminConfigStore) UpdateAdminConfiguration(cmd UpdateAdminConfigura
 
 	return nil
 }
+
+func SetupFolderService(tb testing.TB, cfg *setting.Cfg, dashboardStore dashboards.Store, folderStore *folderimpl.DashboardFolderStoreImpl, bus *bus.InProcBus) folder.Service {
+	tb.Helper()
+
+	ac := acmock.New()
+	features := featuremgmt.WithFeatures()
+
+	return folderimpl.ProvideService(ac, bus, cfg, dashboardStore, folderStore, nil, features)
+}
+
+func SetupDashboardService(tb testing.TB, sqlStore *sqlstore.SQLStore, fs *folderimpl.DashboardFolderStoreImpl, cfg *setting.Cfg) (*dashboardservice.DashboardServiceImpl, dashboards.Store) {
+	tb.Helper()
+
+	origNewGuardian := guardian.New
+	guardian.MockDashboardGuardian(&guardian.FakeDashboardGuardian{
+		CanSaveValue:  true,
+		CanViewValue:  true,
+		CanAdminValue: true,
+	})
+	tb.Cleanup(func() {
+		guardian.New = origNewGuardian
+	})
+
+	ac := acmock.New()
+	dashboardPermissions := acmock.NewMockedPermissionsService()
+	folderPermissions := acmock.NewMockedPermissionsService()
+	folderPermissions.On("SetPermissions", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return([]accesscontrol.ResourcePermission{}, nil)
+
+	features := featuremgmt.WithFeatures()
+	quotaService := quotatest.New(false, nil)
+
+	dashboardStore, err := database.ProvideDashboardStore(sqlStore, sqlStore.Cfg, features, tagimpl.ProvideService(sqlStore, sqlStore.Cfg), quotaService)
+	dashboardService := dashboardservice.ProvideDashboardServiceImpl(
+		cfg, dashboardStore, fs, nil,
+		features, folderPermissions, dashboardPermissions, ac,
+		foldertest.NewFakeService(),
+	)
+
+	require.NoError(tb, err)
+
+	return dashboardService, dashboardStore
+}

pkg/services/ngalert/tests/util.go

@@ -9,7 +9,6 @@ import (
 	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
-	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 
 	"github.com/grafana/grafana/pkg/api/routing"
@@ -19,17 +18,12 @@ import (
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/infra/tracing"
 	"github.com/grafana/grafana/pkg/plugins"
-	"github.com/grafana/grafana/pkg/services/accesscontrol"
 	acmock "github.com/grafana/grafana/pkg/services/accesscontrol/mock"
 	"github.com/grafana/grafana/pkg/services/annotations/annotationstest"
 	"github.com/grafana/grafana/pkg/services/dashboards"
-	databasestore "github.com/grafana/grafana/pkg/services/dashboards/database"
-	dashboardservice "github.com/grafana/grafana/pkg/services/dashboards/service"
 	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/folder"
 	"github.com/grafana/grafana/pkg/services/folder/folderimpl"
-	"github.com/grafana/grafana/pkg/services/folder/foldertest"
-	"github.com/grafana/grafana/pkg/services/guardian"
 	"github.com/grafana/grafana/pkg/services/ngalert"
 	"github.com/grafana/grafana/pkg/services/ngalert/metrics"
 	"github.com/grafana/grafana/pkg/services/ngalert/models"
@@ -38,7 +32,6 @@ import (
 	"github.com/grafana/grafana/pkg/services/quota/quotatest"
 	"github.com/grafana/grafana/pkg/services/secrets/database"
 	secretsManager "github.com/grafana/grafana/pkg/services/secrets/manager"
-	"github.com/grafana/grafana/pkg/services/tag/tagimpl"
 	"github.com/grafana/grafana/pkg/services/user"
 	"github.com/grafana/grafana/pkg/setting"
 	"github.com/grafana/grafana/pkg/util"
@@ -47,15 +40,6 @@ import (
 // SetupTestEnv initializes a store to used by the tests.
 func SetupTestEnv(tb testing.TB, baseInterval time.Duration) (*ngalert.AlertNG, *store.DBstore) {
 	tb.Helper()
-	origNewGuardian := guardian.New
-	guardian.MockDashboardGuardian(&guardian.FakeDashboardGuardian{
-		CanSaveValue:  true,
-		CanViewValue:  true,
-		CanAdminValue: true,
-	})
-	tb.Cleanup(func() {
-		guardian.New = origNewGuardian
-	})
 
 	cfg := setting.NewCfg()
 	cfg.IsFeatureToggleEnabled = featuremgmt.WithFeatures().IsEnabled
@@ -69,27 +53,14 @@ func SetupTestEnv(tb testing.TB, baseInterval time.Duration) (*ngalert.AlertNG,
 	m := metrics.NewNGAlert(prometheus.NewRegistry())
 	sqlStore := db.InitTestDB(tb)
 	secretsService := secretsManager.SetupTestService(tb, database.ProvideSecretsStore(sqlStore))
-	quotaService := quotatest.New(false, nil)
-	dashboardStore, err := databasestore.ProvideDashboardStore(sqlStore, sqlStore.Cfg, featuremgmt.WithFeatures(), tagimpl.ProvideService(sqlStore, sqlStore.Cfg), quotaService)
-	require.NoError(tb, err)
 
 	ac := acmock.New()
-	features := featuremgmt.WithFeatures()
-	folderPermissions := acmock.NewMockedPermissionsService()
-	folderPermissions.On("SetPermissions", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return([]accesscontrol.ResourcePermission{}, nil)
-	dashboardPermissions := acmock.NewMockedPermissionsService()
-
-	folderStore := folderimpl.ProvideDashboardFolderStore(sqlStore)
-
-	dashboardService := dashboardservice.ProvideDashboardServiceImpl(
-		cfg, dashboardStore, folderStore, nil,
-		features, folderPermissions, dashboardPermissions, ac,
-		foldertest.NewFakeService(),
-	)
 
 	tracer := tracing.InitializeTracerForTest()
 	bus := bus.ProvideBus(tracer)
-	folderService := folderimpl.ProvideService(ac, bus, cfg, dashboardStore, folderStore, nil, features)
+	folderStore := folderimpl.ProvideDashboardFolderStore(sqlStore)
+	dashboardService, dashboardStore := store.SetupDashboardService(tb, sqlStore, folderStore, cfg)
+	folderService := store.SetupFolderService(tb, cfg, dashboardStore, folderStore, bus)
 
 	ng, err := ngalert.ProvideService(
 		cfg, featuremgmt.WithFeatures(), nil, nil, routing.NewRouteRegister(), sqlStore, nil, nil, nil, quotatest.New(false, nil),