Stats: use dashboard stats rather than list (grafana#99130)

pkg/infra/usagestats/statscollector/concurrent_users_test.go

@@ -13,6 +13,7 @@ import (
 
 	"github.com/grafana/grafana/pkg/infra/db"
 	"github.com/grafana/grafana/pkg/services/dashboards"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/folder/foldertest"
 	"github.com/grafana/grafana/pkg/services/org/orgtest"
 	"github.com/grafana/grafana/pkg/services/stats/statsimpl"
@@ -28,7 +29,7 @@ func TestMain(m *testing.M) {
 
 func TestConcurrentUsersMetrics(t *testing.T) {
 	sqlStore, cfg := db.InitTestDBWithCfg(t)
-	statsService := statsimpl.ProvideService(&setting.Cfg{}, sqlStore, &dashboards.FakeDashboardService{}, &foldertest.FakeService{}, &orgtest.FakeOrgService{})
+	statsService := statsimpl.ProvideService(&setting.Cfg{}, sqlStore, &dashboards.FakeDashboardService{}, &foldertest.FakeService{}, &orgtest.FakeOrgService{}, featuremgmt.WithFeatures())
 	s := createService(t, cfg, sqlStore, statsService)
 
 	createConcurrentTokens(t, sqlStore)
@@ -46,7 +47,7 @@ func TestConcurrentUsersMetrics(t *testing.T) {
 
 func TestConcurrentUsersStats(t *testing.T) {
 	sqlStore, cfg := db.InitTestDBWithCfg(t)
-	statsService := statsimpl.ProvideService(&setting.Cfg{}, sqlStore, &dashboards.FakeDashboardService{}, &foldertest.FakeService{}, &orgtest.FakeOrgService{})
+	statsService := statsimpl.ProvideService(&setting.Cfg{}, sqlStore, &dashboards.FakeDashboardService{}, &foldertest.FakeService{}, &orgtest.FakeOrgService{}, featuremgmt.WithFeatures())
 	s := createService(t, cfg, sqlStore, statsService)
 
 	createConcurrentTokens(t, sqlStore)

pkg/registry/apis/dashboard/legacy/storage.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"strings"
 	"time"
 
 	"github.com/grafana/authlib/claims"
@@ -340,5 +341,35 @@ func (a *dashboardSqlAccess) CountRepositoryObjects(context.Context, *resource.C
 
 // GetStats implements ResourceServer.
 func (a *dashboardSqlAccess) GetStats(ctx context.Context, req *resource.ResourceStatsRequest) (*resource.ResourceStatsResponse, error) {
-	return nil, fmt.Errorf("not yet (GetStats)")
+	info, err := claims.ParseNamespace(req.Namespace)
+	if err != nil {
+		return nil, fmt.Errorf("unable to read namespace")
+	}
+	if info.OrgID == 0 {
+		return nil, fmt.Errorf("invalid OrgID found in namespace")
+	}
+
+	if len(req.Kinds) != 1 {
+		return nil, fmt.Errorf("only can query for dashboard kind in legacy fallback")
+	}
+
+	parts := strings.SplitN(req.Kinds[0], "/", 2)
+	if len(parts) != 2 {
+		return nil, fmt.Errorf("invalid kind")
+	}
+
+	count, err := a.dashStore.CountInOrg(ctx, info.OrgID)
+	if err != nil {
+		return nil, err
+	}
+
+	return &resource.ResourceStatsResponse{
+		Stats: []*resource.ResourceStatsResponse_Stats{
+			{
+				Group:    parts[0],
+				Resource: parts[1],
+				Count:    count,
+			},
+		},
+	}, nil
 }

pkg/services/dashboards/dashboard.go

@@ -36,6 +36,7 @@ type DashboardService interface {
 	RestoreDashboard(ctx context.Context, dashboard *Dashboard, user identity.Requester, optionalFolderUID string) error
 	CleanUpDeletedDashboards(ctx context.Context) (int64, error)
 	GetSoftDeletedDashboard(ctx context.Context, orgID int64, uid string) (*Dashboard, error)
+	CountDashboardsInOrg(ctx context.Context, orgID int64) (int64, error)
 }
 
 // PluginService is a service for operating on plugin dashboards.
@@ -82,6 +83,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)
 	// CountDashboardsInFolder returns the number of dashboards associated with
 	// the given parent folder ID.
 	CountDashboardsInFolders(ctx context.Context, request *CountDashboardsInFolderRequest) (int64, error)

pkg/services/dashboards/database/database.go

@@ -266,6 +266,24 @@ func (d *dashboardStore) Count(ctx context.Context, scopeParams *quota.ScopePara
 	return u, nil
 }
 
+func (d *dashboardStore) CountInOrg(ctx context.Context, orgID int64) (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))
+		if _, err := sess.SQL(rawSQL, orgID).Get(&r); err != nil {
+			return err
+		}
+		return nil
+	}); err != nil {
+		return 0, err
+	}
+
+	return r.Count, nil
+}
+
 func getExistingDashboardByIDOrUIDForUpdate(sess *db.Session, dash *dashboards.Dashboard, overwrite bool) (bool, error) {
 	dashWithIdExists := false
 	isParentFolderChanged := false

pkg/services/dashboards/database/database_test.go

@@ -61,6 +61,7 @@ func TestIntegrationDashboardDataAccess(t *testing.T) {
 		savedDash = insertTestDashboard(t, dashboardStore, "test dash 23", 1, savedFolder.ID, savedFolder.UID, false, "prod", "webapp")
 		insertTestDashboard(t, dashboardStore, "test dash 45", 1, savedFolder.ID, savedFolder.UID, false, "prod")
 		savedDash2 = insertTestDashboard(t, dashboardStore, "test dash 67", 1, 0, "", false, "prod")
+		insertTestDashboard(t, dashboardStore, "test dash org2", 2, 0, "", false, "")
 		insertTestRule(t, sqlStore, savedFolder.OrgID, savedFolder.UID)
 	}
 
@@ -81,6 +82,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) {
+		setup()
+		count, err := dashboardStore.CountInOrg(context.Background(), 1)
+		require.NoError(t, err)
+		require.Equal(t, int64(3), count)
+
+		count, err = dashboardStore.CountInOrg(context.Background(), 2)
+		require.NoError(t, err)
+		require.Equal(t, int64(1), count)
+	})
+
 	t.Run("Should be able to get dashboard by id", func(t *testing.T) {
 		setup()
 		query := dashboards.GetDashboardQuery{

pkg/services/dashboards/service/dashboard_service.go

@@ -180,11 +180,10 @@ func (dr *DashboardServiceImpl) Count(ctx context.Context, scopeParams *quota.Sc
 		total := int64(0)
 		for _, org := range orgs {
 			ctx = identity.WithRequester(ctx, getDashboardBackgroundRequester(org.ID))
-			dashs, err := dr.listDashboardsThroughK8s(ctx, org.ID)
+			orgDashboards, err := dr.CountDashboardsInOrg(ctx, org.ID)
 			if err != nil {
-				return u, err
+				return nil, err
 			}
-			orgDashboards := int64(len(dashs))
 			total += orgDashboards
 
 			tag, err := quota.NewTag(dashboards.QuotaTargetSrv, dashboards.QuotaTarget, quota.OrgScope)
@@ -206,6 +205,28 @@ func (dr *DashboardServiceImpl) Count(ctx context.Context, scopeParams *quota.Sc
 	return dr.dashboardStore.Count(ctx, scopeParams)
 }
 
+func (dr *DashboardServiceImpl) CountDashboardsInOrg(ctx context.Context, orgID int64) (int64, error) {
+	if dr.features.IsEnabledGlobally(featuremgmt.FlagKubernetesCliDashboards) {
+		resp, err := dr.k8sclient.getSearcher().GetStats(ctx, &resource.ResourceStatsRequest{
+			Namespace: dr.k8sclient.getNamespace(orgID),
+			Kinds: []string{
+				v0alpha1.GROUP + "/" + v0alpha1.DashboardResourceInfo.GetName(),
+			},
+		})
+		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
+	}
+
+	return dr.dashboardStore.CountInOrg(ctx, orgID)
+}
+
 func readQuotaConfig(cfg *setting.Cfg) (*quota.Map, error) {
 	limits := &quota.Map{}
 

pkg/services/dashboards/service/dashboard_service_test.go

@@ -271,6 +271,11 @@ func (m *mockResourceIndexClient) Search(ctx context.Context, req *resource.Reso
 	return args.Get(0).(*resource.ResourceSearchResponse), args.Error(1)
 }
 
+func (m *mockResourceIndexClient) GetStats(ctx context.Context, in *resource.ResourceStatsRequest, opts ...grpc.CallOption) (*resource.ResourceStatsResponse, error) {
+	args := m.Called(in)
+	return args.Get(0).(*resource.ResourceStatsResponse), args.Error(1)
+}
+
 type mockResourceInterface struct {
 	mock.Mock
 	dynamic.ResourceInterface
@@ -1779,44 +1784,17 @@ func TestQuotaCount(t *testing.T) {
 		},
 	}
 
-	dashboardUnstructuredOrg1 := unstructured.UnstructuredList{
-		Items: []unstructured.Unstructured{{
-			Object: map[string]any{
-				"metadata": map[string]any{
-					"name": "uid",
-				},
-				"spec": map[string]any{
-					"test":    "test",
-					"version": int64(1),
-					"title":   "testing slugify",
-				},
-			},
-		}},
-	}
-	dashboardUnstructuredOrg2 := unstructured.UnstructuredList{
-		Items: []unstructured.Unstructured{{
-			Object: map[string]any{
-				"metadata": map[string]any{
-					"name": "uid",
-				},
-				"spec": map[string]any{
-					"test":    "test",
-					"version": int64(1),
-					"title":   "testing slugify",
-				},
+	countOrg1 := resource.ResourceStatsResponse{
+		Stats: []*resource.ResourceStatsResponse_Stats{
+			{
+				Count: 1,
 			},
 		},
+	}
+	countOrg2 := resource.ResourceStatsResponse{
+		Stats: []*resource.ResourceStatsResponse_Stats{
 			{
-				Object: map[string]any{
-					"metadata": map[string]any{
-						"name": "uid2",
-					},
-					"spec": map[string]any{
-						"test":    "test2",
-						"version": int64(1),
-						"title":   "testing slugify2",
-					},
-				},
+				Count: 2,
 			},
 		},
 	}
@@ -1833,13 +1811,11 @@ func TestQuotaCount(t *testing.T) {
 	})
 
 	t.Run("Should use Kubernetes client if feature flags are enabled", func(t *testing.T) {
-		ctx, k8sClientMock, k8sResourceMock := setupK8sDashboardTests(service)
+		ctx, k8sClientMock, _ := setupK8sDashboardTests(service)
 		orgSvc := orgtest.FakeOrgService{ExpectedOrgs: orgs}
 		service.orgService = &orgSvc
-		k8sClientMock.On("getClient", mock.Anything, int64(1)).Return(k8sResourceMock, true).Once()
-		k8sClientMock.On("getClient", mock.Anything, int64(2)).Return(k8sResourceMock, true).Once()
-		k8sResourceMock.On("List", mock.Anything, mock.Anything).Return(&dashboardUnstructuredOrg2, nil).Once()
-		k8sResourceMock.On("List", mock.Anything, mock.Anything).Return(&dashboardUnstructuredOrg1, nil).Once()
+		k8sClientMock.searcher.On("GetStats", mock.Anything).Return(&countOrg2, nil).Once()
+		k8sClientMock.searcher.On("GetStats", mock.Anything).Return(&countOrg1, nil).Once()
 
 		result, err := service.Count(ctx, query)
 		require.NoError(t, err)
@@ -1858,6 +1834,38 @@ func TestQuotaCount(t *testing.T) {
 	})
 }
 
+func TestCountDashboardsInOrg(t *testing.T) {
+	fakeStore := dashboards.FakeDashboardStore{}
+	defer fakeStore.AssertExpectations(t)
+	service := &DashboardServiceImpl{
+		cfg:            setting.NewCfg(),
+		dashboardStore: &fakeStore,
+	}
+	count := resource.ResourceStatsResponse{
+		Stats: []*resource.ResourceStatsResponse_Stats{
+			{
+				Count: 3,
+			},
+		},
+	}
+
+	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()
+		_, err := service.CountDashboardsInOrg(context.Background(), 1)
+		require.NoError(t, err)
+		fakeStore.AssertExpectations(t)
+	})
+
+	t.Run("Should use Kubernetes client if feature flags are enabled", func(t *testing.T) {
+		ctx, k8sClientMock, _ := setupK8sDashboardTests(service)
+		k8sClientMock.searcher.On("GetStats", mock.Anything).Return(&count, nil).Once()
+		result, err := service.CountDashboardsInOrg(ctx, 1)
+		require.NoError(t, err)
+		require.Equal(t, result, int64(3))
+	})
+}
+
 func TestLegacySaveCommandToUnstructured(t *testing.T) {
 	namespace := "test-namespace"
 	t.Run("successfully converts save command to unstructured", func(t *testing.T) {