feat(nested folders): Add CountAlertRulesInFolder to ngalert store (#58269)

* chore: refactor CountDashboardsInFolder to use the more efficient Count() sql function

* feat(nested folders): Add CountAlertRulesInFolder to ngalert store

This commit adds CountAlertRulesInFolder and a new model for the CountAlertRulesQuery. It returns a count of alert rules associated with a given orgID and parent folder UID. (the namespace referenced inside alert rules is the parent folder).

I'm not sure where this belongs in the ngalert service, so that will come in a future commit.

pkg/services/dashboards/database/database.go

@@ -1018,17 +1018,20 @@ func (d *DashboardStore) GetDashboardTags(ctx context.Context, query *models.Get
 	})
 }
 
+// CountDashboardsInFolder returns a count of all dashboards associated with the
+// given parent folder ID.
+//
 // This will be updated to take CountDashboardsInFolderQuery as an argument and
-// lookup dashboards using the ParentFolderUID when the NestedFolder
-// implementation is complete.
+// lookup dashboards using the ParentFolderUID when dashboards are associated with a parent folder UID instead of ID.
 func (d *DashboardStore) CountDashboardsInFolder(
 	ctx context.Context, req *dashboards.CountDashboardsInFolderRequest) (int64, error) {
-	var dashboards = make([]*models.Dashboard, 0)
-	err := d.store.WithDbSession(ctx, func(sess *db.Session) error {
+	var count int64
+	var err error
+	err = d.store.WithDbSession(ctx, func(sess *db.Session) error {
 		session := sess.In("folder_id", req.FolderID).In("org_id", req.OrgID).
 			In("is_folder", d.store.GetDialect().BooleanStr(false))
-		err := session.Find(&dashboards)
+		count, err = session.Count(&models.Dashboard{})
 		return err
 	})
-	return int64(len(dashboards)), err
+	return count, err
 }

pkg/services/dashboards/models.go

@@ -32,12 +32,12 @@ type DashboardSearchProjection struct {
 
 type CountDashboardsInFolderQuery struct {
 	FolderUID string
+	OrgID     int64
 }
 
-// Note for reviewers: I wasn't sure what to name this. It's not actually a DTO
-// CountDashboardsInFolderRequest is the request passed from the service to the
-// store layer. The FolderID will be replaced with FolderUID when dashboards are
-// updated with parent folder UIDs.
+// TODO: CountDashboardsInFolderRequest is the request passed from the service
+// to the store layer. The FolderID will be replaced with FolderUID when
+// dashboards are updated with parent folder UIDs.
 type CountDashboardsInFolderRequest struct {
 	FolderID int64
 	OrgID    int64

pkg/services/ngalert/models/alert_rule.go

@@ -354,6 +354,12 @@ type ListAlertRulesQuery struct {
 	Result RulesGroup
 }
 
+// CountAlertRulesQuery is the query for counting alert rules
+type CountAlertRulesQuery struct {
+	OrgID        int64
+	NamespaceUID string
+}
+
 type GetAlertRulesForSchedulingQuery struct {
 	PopulateFolders bool
 

pkg/services/ngalert/store/alert_rule.go

@@ -226,6 +226,19 @@ func (st DBstore) UpdateAlertRules(ctx context.Context, rules []ngmodels.UpdateR
 	})
 }
 
+// CountAlertRulesInFolder is a handler for retrieving the number of alert rules of
+// specific organisation associated with a given namespace (parent folder).
+func (st DBstore) CountAlertRulesInFolder(ctx context.Context, query *ngmodels.CountAlertRulesQuery) (int64, error) {
+	var count int64
+	var err error
+	err = st.SQLStore.WithDbSession(ctx, func(sess *db.Session) error {
+		q := sess.Table("alert_rule").Where("org_id = ?", query.OrgID).Where("namespace_uid = ?", query.NamespaceUID)
+		count, err = q.Count()
+		return err
+	})
+	return count, err
+}
+
 // ListAlertRules is a handler for retrieving alert rules of specific organisation.
 func (st DBstore) ListAlertRules(ctx context.Context, query *ngmodels.ListAlertRulesQuery) error {
 	return st.SQLStore.WithDbSession(ctx, func(sess *db.Session) error {

pkg/services/ngalert/store/alert_rule_test.go

@@ -28,30 +28,9 @@ func TestIntegrationUpdateAlertRules(t *testing.T) {
 			BaseInterval: time.Duration(rand.Int63n(100)) * time.Second,
 		},
 	}
-	createRule := func(t *testing.T) *models.AlertRule {
-		rule := models.AlertRuleGen(withIntervalMatching(store.Cfg.BaseInterval))()
-		err := sqlStore.WithDbSession(context.Background(), func(sess *db.Session) error {
-			_, err := sess.Table(models.AlertRule{}).InsertOne(rule)
-			if err != nil {
-				return err
-			}
-			dbRule := &models.AlertRule{}
-			exist, err := sess.Table(models.AlertRule{}).ID(rule.ID).Get(dbRule)
-			if err != nil {
-				return err
-			}
-			if !exist {
-				return errors.New("cannot read inserted record")
-			}
-			rule = dbRule
-			return nil
-		})
-		require.NoError(t, err)
-		return rule
-	}
 
 	t.Run("should increase version", func(t *testing.T) {
-		rule := createRule(t)
+		rule := createRule(t, store)
 		newRule := models.CopyRule(rule)
 		newRule.Title = util.GenerateShortUID()
 		err := store.UpdateAlertRules(context.Background(), []models.UpdateRule{{
@@ -73,7 +52,7 @@ func TestIntegrationUpdateAlertRules(t *testing.T) {
 	})
 
 	t.Run("should fail due to optimistic locking if version does not match", func(t *testing.T) {
-		rule := createRule(t)
+		rule := createRule(t, store)
 		rule.Version-- // simulate version discrepancy
 
 		newRule := models.CopyRule(rule)
@@ -150,3 +129,70 @@ func TestIntegration_getFilterByOrgsString(t *testing.T) {
 		})
 	}
 }
+
+func TestIntegration_CountAlertRules(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping integration test")
+	}
+
+	sqlStore := db.InitTestDB(t)
+	store := DBstore{SQLStore: sqlStore}
+	rule := createRule(t, store)
+
+	tests := map[string]struct {
+		query     *models.CountAlertRulesQuery
+		expected  int64
+		expectErr bool
+	}{
+		"basic success": {
+			&models.CountAlertRulesQuery{
+				NamespaceUID: rule.NamespaceUID,
+				OrgID:        rule.OrgID,
+			},
+			1,
+			false,
+		},
+		"successfully returning no results": {
+			&models.CountAlertRulesQuery{
+				NamespaceUID: "probably not a uid we'd generate",
+				OrgID:        rule.OrgID,
+			},
+			0,
+			false,
+		},
+	}
+
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			count, err := store.CountAlertRulesInFolder(context.Background(), test.query)
+			if test.expectErr {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, test.expected, count)
+			}
+		})
+	}
+}
+
+func createRule(t *testing.T, store DBstore) *models.AlertRule {
+	rule := models.AlertRuleGen(withIntervalMatching(store.Cfg.BaseInterval))()
+	err := store.SQLStore.WithDbSession(context.Background(), func(sess *db.Session) error {
+		_, err := sess.Table(models.AlertRule{}).InsertOne(rule)
+		if err != nil {
+			return err
+		}
+		dbRule := &models.AlertRule{}
+		exist, err := sess.Table(models.AlertRule{}).ID(rule.ID).Get(dbRule)
+		if err != nil {
+			return err
+		}
+		if !exist {
+			return errors.New("cannot read inserted record")
+		}
+		rule = dbRule
+		return nil
+	})
+	require.NoError(t, err)
+	return rule
+}