From f84654d162d61a54da764f378c5034ea9008d215 Mon Sep 17 00:00:00 2001
From: Stephanie Hingtgen <stephanie.hingtgen@grafana.com>
Date: Fri, 28 Mar 2025 19:17:50 -0600
Subject: [PATCH] Stats: Optimize getting folder stats (#103033)

---
 .../dashboard/legacysearcher/search_client.go | 10 +++++++-
 pkg/services/dashboards/dashboard.go          |  2 +-
 pkg/services/dashboards/database/database.go  |  4 ++--
 .../dashboards/database/database_test.go      | 10 +++++---
 .../dashboards/service/dashboard_service.go   |  2 +-
 .../service/dashboard_service_test.go         |  2 +-
 pkg/services/dashboards/store_mock.go         | 16 ++++++-------
 pkg/services/folder/folderimpl/folder.go      |  8 +++++++
 pkg/services/folder/folderimpl/sqlstore.go    | 18 +++++++++++++++
 .../folder/folderimpl/unifiedstore.go         | 13 +++++++++++
 .../folder/foldertest/folder_store_mock.go    | 23 +++++++++++++++++++
 pkg/services/folder/foldertest/foldertest.go  |  4 ++++
 pkg/services/folder/service.go                |  2 ++
 pkg/services/folder/store.go                  |  3 +++
 pkg/services/folder/store_fake.go             |  4 ++++
 pkg/services/stats/statsimpl/stats.go         | 15 ++++--------
 16 files changed, 109 insertions(+), 27 deletions(-)

diff --git a/pkg/registry/apis/dashboard/legacysearcher/search_client.go b/pkg/registry/apis/dashboard/legacysearcher/search_client.go
index c26f1c5c6ca..01e22c1c516 100644
--- a/pkg/registry/apis/dashboard/legacysearcher/search_client.go
+++ b/pkg/registry/apis/dashboard/legacysearcher/search_client.go
@@ -400,7 +400,15 @@ func (c *DashboardSearchClient) GetStats(ctx context.Context, req *resource.Reso
 		return nil, fmt.Errorf("invalid kind")
 	}
 
-	count, err := c.dashboardStore.CountInOrg(ctx, info.OrgID)
+	var count int64
+	switch parts[0] {
+	case dashboard.GROUP:
+		count, err = c.dashboardStore.CountInOrg(ctx, info.OrgID, false)
+	case folderv0alpha1.GROUP:
+		count, err = c.dashboardStore.CountInOrg(ctx, info.OrgID, true)
+	default:
+		return nil, fmt.Errorf("invalid group")
+	}
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/services/dashboards/dashboard.go b/pkg/services/dashboards/dashboard.go
index 3f14c6f00b0..425a3b43c2f 100644
--- a/pkg/services/dashboards/dashboard.go
+++ b/pkg/services/dashboards/dashboard.go
@@ -91,7 +91,7 @@ type Store interface {
 	ValidateDashboardBeforeSave(ctx context.Context, dashboard *Dashboard, overwrite bool) (bool, error)
 
 	Count(context.Context, *quota.ScopeParameters) (*quota.Map, error)
-	CountInOrg(ctx context.Context, orgID int64) (int64, error)
+	CountInOrg(ctx context.Context, orgID int64, isFolder bool) (int64, error)
 	// CountDashboardsInFolder returns the number of dashboards associated with
 	// the given parent folder ID.
 	CountDashboardsInFolders(ctx context.Context, request *CountDashboardsInFolderRequest) (int64, error)
diff --git a/pkg/services/dashboards/database/database.go b/pkg/services/dashboards/database/database.go
index 1ada23143d9..24bb29c7635 100644
--- a/pkg/services/dashboards/database/database.go
+++ b/pkg/services/dashboards/database/database.go
@@ -307,13 +307,13 @@ func (d *dashboardStore) Count(ctx context.Context, scopeParams *quota.ScopePara
 	return u, nil
 }
 
-func (d *dashboardStore) CountInOrg(ctx context.Context, orgID int64) (int64, error) {
+func (d *dashboardStore) CountInOrg(ctx context.Context, orgID int64, isFolder bool) (int64, error) {
 	type result struct {
 		Count int64
 	}
 	r := result{}
 	if err := d.store.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
-		rawSQL := fmt.Sprintf("SELECT COUNT(*) AS count FROM dashboard WHERE org_id=? AND is_folder=%s", d.store.GetDialect().BooleanStr(false))
+		rawSQL := fmt.Sprintf("SELECT COUNT(*) AS count FROM dashboard WHERE org_id=? AND is_folder=%s", d.store.GetDialect().BooleanStr(isFolder))
 		if _, err := sess.SQL(rawSQL, orgID).Get(&r); err != nil {
 			return err
 		}
diff --git a/pkg/services/dashboards/database/database_test.go b/pkg/services/dashboards/database/database_test.go
index 6078a30e9b6..473a4be8001 100644
--- a/pkg/services/dashboards/database/database_test.go
+++ b/pkg/services/dashboards/database/database_test.go
@@ -85,13 +85,17 @@ func TestIntegrationDashboardDataAccess(t *testing.T) {
 		require.Positive(t, len(savedFolder.UID))
 	})
 
-	t.Run("Should be able to get dashboard counts per org", func(t *testing.T) {
+	t.Run("Should be able to get counts per org", func(t *testing.T) {
 		setup()
-		count, err := dashboardStore.CountInOrg(context.Background(), 1)
+		count, err := dashboardStore.CountInOrg(context.Background(), 1, false)
 		require.NoError(t, err)
 		require.Equal(t, int64(3), count)
 
-		count, err = dashboardStore.CountInOrg(context.Background(), 2)
+		count, err = dashboardStore.CountInOrg(context.Background(), 2, false)
+		require.NoError(t, err)
+		require.Equal(t, int64(1), count)
+
+		count, err = dashboardStore.CountInOrg(context.Background(), 1, true)
 		require.NoError(t, err)
 		require.Equal(t, int64(1), count)
 	})
diff --git a/pkg/services/dashboards/service/dashboard_service.go b/pkg/services/dashboards/service/dashboard_service.go
index cba4634650e..adcbd8014be 100644
--- a/pkg/services/dashboards/service/dashboard_service.go
+++ b/pkg/services/dashboards/service/dashboard_service.go
@@ -480,7 +480,7 @@ func (dr *DashboardServiceImpl) CountDashboardsInOrg(ctx context.Context, orgID
 		return resp.Stats[0].Count, nil
 	}
 
-	return dr.dashboardStore.CountInOrg(ctx, orgID)
+	return dr.dashboardStore.CountInOrg(ctx, orgID, false)
 }
 
 func readQuotaConfig(cfg *setting.Cfg) (*quota.Map, error) {
diff --git a/pkg/services/dashboards/service/dashboard_service_test.go b/pkg/services/dashboards/service/dashboard_service_test.go
index cf35a0b1777..fcfc00c048c 100644
--- a/pkg/services/dashboards/service/dashboard_service_test.go
+++ b/pkg/services/dashboards/service/dashboard_service_test.go
@@ -2281,7 +2281,7 @@ func TestCountDashboardsInOrg(t *testing.T) {
 
 	t.Run("Should fallback to dashboard store if Kubernetes feature flags are not enabled", func(t *testing.T) {
 		service.features = featuremgmt.WithFeatures()
-		fakeStore.On("CountInOrg", mock.Anything, mock.Anything).Return(nil, nil).Once()
+		fakeStore.On("CountInOrg", mock.Anything, mock.Anything, false).Return(nil, nil).Once()
 		_, err := service.CountDashboardsInOrg(context.Background(), 1)
 		require.NoError(t, err)
 		fakeStore.AssertExpectations(t)
diff --git a/pkg/services/dashboards/store_mock.go b/pkg/services/dashboards/store_mock.go
index 4d9d94c1491..76d37ecdc8f 100644
--- a/pkg/services/dashboards/store_mock.go
+++ b/pkg/services/dashboards/store_mock.go
@@ -95,8 +95,8 @@ func (_m *FakeDashboardStore) CountDashboardsInFolders(ctx context.Context, requ
 }
 
 // CountInOrg provides a mock function with given fields: ctx, orgID
-func (_m *FakeDashboardStore) CountInOrg(ctx context.Context, orgID int64) (int64, error) {
-	ret := _m.Called(ctx, orgID)
+func (_m *FakeDashboardStore) CountInOrg(ctx context.Context, orgID int64, isFolder bool) (int64, error) {
+	ret := _m.Called(ctx, orgID, isFolder)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Count")
@@ -104,19 +104,19 @@ func (_m *FakeDashboardStore) CountInOrg(ctx context.Context, orgID int64) (int6
 
 	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, bool) (int64, error)); ok {
+		return rf(ctx, orgID, isFolder)
 	}
-	if rf, ok := ret.Get(0).(func(context.Context, int64) int64); ok {
-		r0 = rf(ctx, orgID)
+	if rf, ok := ret.Get(0).(func(context.Context, int64, bool) int64); ok {
+		r0 = rf(ctx, orgID, isFolder)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).(int64)
 		}
 	}
 
-	if rf, ok := ret.Get(1).(func(context.Context, int64) error); ok {
-		r1 = rf(ctx, orgID)
+	if rf, ok := ret.Get(1).(func(context.Context, int64, bool) error); ok {
+		r1 = rf(ctx, orgID, isFolder)
 	} else {
 		r1 = ret.Error(1)
 	}
diff --git a/pkg/services/folder/folderimpl/folder.go b/pkg/services/folder/folderimpl/folder.go
index 05cc142cf74..7aa12fd9f21 100644
--- a/pkg/services/folder/folderimpl/folder.go
+++ b/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:
diff --git a/pkg/services/folder/folderimpl/sqlstore.go b/pkg/services/folder/folderimpl/sqlstore.go
index 20c17c324b5..fd1a79e587b 100644
--- a/pkg/services/folder/folderimpl/sqlstore.go
+++ b/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")
diff --git a/pkg/services/folder/folderimpl/unifiedstore.go b/pkg/services/folder/folderimpl/unifiedstore.go
index 99a70423fe2..e1e0be576b6 100644
--- a/pkg/services/folder/folderimpl/unifiedstore.go
+++ b/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 {
diff --git a/pkg/services/folder/foldertest/folder_store_mock.go b/pkg/services/folder/foldertest/folder_store_mock.go
index 93de1422257..1637fb57e06 100644
--- a/pkg/services/folder/foldertest/folder_store_mock.go
+++ b/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 {
diff --git a/pkg/services/folder/foldertest/foldertest.go b/pkg/services/folder/foldertest/foldertest.go
index 43df5478980..aaff277fb1d 100644
--- a/pkg/services/folder/foldertest/foldertest.go
+++ b/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
+}
diff --git a/pkg/services/folder/service.go b/pkg/services/folder/service.go
index 89a123b9b45..2d72b38914d 100644
--- a/pkg/services/folder/service.go
+++ b/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.
diff --git a/pkg/services/folder/store.go b/pkg/services/folder/store.go
index e2f17986221..04350fd1443 100644
--- a/pkg/services/folder/store.go
+++ b/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)
 }
diff --git a/pkg/services/folder/store_fake.go b/pkg/services/folder/store_fake.go
index 79f1a1d6863..7d1575203f7 100644
--- a/pkg/services/folder/store_fake.go
+++ b/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
+}
diff --git a/pkg/services/stats/statsimpl/stats.go b/pkg/services/stats/statsimpl/stats.go
index b8091c8895f..a33ffbb3420 100644
--- a/pkg/services/stats/statsimpl/stats.go
+++ b/pkg/services/stats/statsimpl/stats.go
@@ -51,7 +51,6 @@ func (ss *sqlStatsService) getDashboardCount(ctx context.Context, orgs []*org.Or
 		}
 		count += dashsCount
 	}
-
 	return count, nil
 }
 
@@ -72,20 +71,16 @@ func (ss *sqlStatsService) getTagCount(ctx context.Context, orgs []*org.OrgDTO)
 }
 
 func (ss *sqlStatsService) getFolderCount(ctx context.Context, orgs []*org.OrgDTO) (int64, error) {
-	total := 0
+	total := int64(0)
 	for _, org := range orgs {
-		ctx, ident := identity.WithServiceIdentity(ctx, org.ID)
-		folders, err := ss.folderSvc.GetFolders(ctx, folder.GetFoldersQuery{
-			OrgID:        org.ID,
-			SignedInUser: ident,
-		})
+		ctx, _ = identity.WithServiceIdentity(ctx, org.ID)
+		folderCount, err := ss.folderSvc.CountFoldersInOrg(ctx, org.ID)
 		if err != nil {
 			return 0, err
 		}
-
-		total += len(folders)
+		total += folderCount
 	}
-	return int64(total), nil
+	return total, nil
 }
 
 func (ss *sqlStatsService) GetAlertNotifiersUsageStats(ctx context.Context, query *stats.GetAlertNotifierUsageStatsQuery) (result []*stats.NotifierUsageStats, err error) {