Stats: Optimize getting folder stats (#103033)

pkg/services/folder/folderimpl/folder.go

@@ -202,6 +202,14 @@ func (s *Service) DBMigration(db db.DB) {
 	s.log.Debug("syncing dashboard and folder tables finished")
 }
 
+func (s *Service) CountFoldersInOrg(ctx context.Context, orgID int64) (int64, error) {
+	if s.features.IsEnabledGlobally(featuremgmt.FlagKubernetesClientDashboardsFolders) {
+		return s.unifiedStore.CountInOrg(ctx, orgID)
+	}
+
+	return s.store.CountInOrg(ctx, orgID)
+}
+
 func (s *Service) SearchFolders(ctx context.Context, q folder.SearchFoldersQuery) (model.HitList, error) {
 	if s.features.IsEnabledGlobally(featuremgmt.FlagKubernetesClientDashboardsFolders) {
 		// TODO:

pkg/services/folder/folderimpl/sqlstore.go

@@ -16,6 +16,7 @@ import (
 	"github.com/grafana/grafana/pkg/services/accesscontrol"
 	"github.com/grafana/grafana/pkg/services/dashboards"
 	"github.com/grafana/grafana/pkg/services/folder"
+	"github.com/grafana/grafana/pkg/services/sqlstore"
 	"github.com/grafana/grafana/pkg/services/sqlstore/migrator"
 	"github.com/grafana/grafana/pkg/util"
 )
@@ -34,6 +35,23 @@ func ProvideStore(db db.DB) *FolderStoreImpl {
 	return &FolderStoreImpl{db: db, log: log.New("folder-store")}
 }
 
+func (ss *FolderStoreImpl) CountInOrg(ctx context.Context, orgID int64) (int64, error) {
+	type result struct {
+		Count int64
+	}
+	r := result{}
+	if err := ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
+		if _, err := sess.SQL("SELECT COUNT(*) AS count FROM folder WHERE org_id=?", orgID).Get(&r); err != nil {
+			return err
+		}
+		return nil
+	}); err != nil {
+		return 0, err
+	}
+
+	return r.Count, nil
+}
+
 func (ss *FolderStoreImpl) Create(ctx context.Context, cmd folder.CreateFolderCommand) (*folder.Folder, error) {
 	if cmd.UID == "" {
 		return nil, folder.ErrBadRequest.Errorf("missing UID")

pkg/services/folder/folderimpl/unifiedstore.go

@@ -452,6 +452,19 @@ func (ss *FolderUnifiedStoreImpl) CountFolderContent(ctx context.Context, orgID
 	return *res, err
 }
 
+func (ss *FolderUnifiedStoreImpl) CountInOrg(ctx context.Context, orgID int64) (int64, error) {
+	resp, err := ss.k8sclient.GetStats(ctx, orgID)
+	if err != nil {
+		return 0, err
+	}
+
+	if len(resp.Stats) != 1 {
+		return 0, fmt.Errorf("expected 1 stat, got %d", len(resp.Stats))
+	}
+
+	return resp.Stats[0].Count, nil
+}
+
 func toFolderLegacyCounts(u *unstructured.Unstructured) (*folder.DescendantCounts, error) {
 	ds, err := v0alpha1.UnstructuredToDescendantCounts(u)
 	if err != nil {

pkg/services/folder/foldertest/folder_store_mock.go

@@ -118,6 +118,29 @@ func (_m *FakeFolderStore) GetFolders(ctx context.Context, orgID int64, uids []s
 	return r0, r1
 }
 
+func (_m *FakeFolderStore) CountInOrg(ctx context.Context, orgID int64) (int64, error) {
+	ret := _m.Called(ctx, orgID)
+
+	var r0 int64
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, int64) (int64, error)); ok {
+		return rf(ctx, orgID)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, int64) int64); ok {
+		r0 = rf(ctx, orgID)
+	} else {
+		r0 = ret.Get(0).(int64)
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, int64) error); ok {
+		r1 = rf(ctx, orgID)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
 // NewFakeFolderStore creates a new instance of FakeFolderStore. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
 // The first argument is typically a *testing.T value.
 func NewFakeFolderStore(t interface {

pkg/services/folder/foldertest/foldertest.go

@@ -94,3 +94,7 @@ func (s *FakeService) GetFoldersLegacy(ctx context.Context, q folder.GetFoldersQ
 	s.LastQuery = q
 	return s.ExpectedFolders, s.ExpectedError
 }
+
+func (s *FakeService) CountFoldersInOrg(ctx context.Context, orgID int64) (int64, error) {
+	return int64(len(s.ExpectedFolders)), s.ExpectedError
+}

pkg/services/folder/service.go

@@ -56,6 +56,8 @@ type Service interface {
 
 	GetDescendantCounts(ctx context.Context, q *GetDescendantCountsQuery) (DescendantCounts, error)
 	GetDescendantCountsLegacy(ctx context.Context, q *GetDescendantCountsQuery) (DescendantCounts, error)
+
+	CountFoldersInOrg(ctx context.Context, orgID int64) (int64, error)
 }
 
 // FolderStore is a folder store.

pkg/services/folder/store.go

@@ -48,4 +48,7 @@ type Store interface {
 	GetFolders(ctx context.Context, q GetFoldersFromStoreQuery) ([]*Folder, error)
 	// GetDescendants returns all descendants of a folder
 	GetDescendants(ctx context.Context, orgID int64, anchestor_uid string) ([]*Folder, error)
+
+	// CountInOrg returns the number of folders in the given org
+	CountInOrg(ctx context.Context, orgID int64) (int64, error)
 }

pkg/services/folder/store_fake.go

@@ -62,3 +62,7 @@ func (f *fakeStore) GetFolders(ctx context.Context, q GetFoldersFromStoreQuery)
 func (f *fakeStore) GetDescendants(ctx context.Context, orgID int64, ancestor_uid string) ([]*Folder, error) {
 	return f.ExpectedFolders, f.ExpectedError
 }
+
+func (f *fakeStore) CountInOrg(ctx context.Context, orgID int64) (int64, error) {
+	return int64(len(f.ExpectedFolders)), f.ExpectedError
+}