Implement k8s count handler (#97955)

* Implement k8s count handler

* Fix endpoint

* Fix type converstions

* Add tests for foldercounts

* Add more tests

* Only use sql-fallback if no values in unistore

* Update gomod

* Fix test

* Update pkg/api/folder_test.go

Co-authored-by: Bruno Abrantes <bruno.abrantes@grafana.com>

* Go.mod

---------

Co-authored-by: Bruno Abrantes <bruno.abrantes@grafana.com>

pkg/api/folder.go

@@ -50,8 +50,6 @@ func (hs *HTTPServer) registerFolderAPI(apiRoute routing.RouteRegister, authoriz
 
 		folderRoute.Group("/:uid", func(folderUidRoute routing.RouteRegister) {
 			folderUidRoute.Post("/move", authorize(accesscontrol.EvalPermission(dashboards.ActionFoldersWrite, uidScope)), routing.Wrap(hs.MoveFolder))
-			folderUidRoute.Get("/counts", authorize(accesscontrol.EvalPermission(dashboards.ActionFoldersRead, uidScope)), routing.Wrap(hs.GetFolderDescendantCounts))
-
 			folderUidRoute.Group("/permissions", func(folderPermissionRoute routing.RouteRegister) {
 				folderPermissionRoute.Get("/", authorize(accesscontrol.EvalPermission(dashboards.ActionFoldersPermissionsRead, uidScope)), routing.Wrap(hs.GetFolderPermissionList))
 				folderPermissionRoute.Post("/", authorize(accesscontrol.EvalPermission(dashboards.ActionFoldersPermissionsWrite, uidScope)), routing.Wrap(hs.UpdateFolderPermissions))
@@ -66,6 +64,7 @@ func (hs *HTTPServer) registerFolderAPI(apiRoute routing.RouteRegister, authoriz
 				folderUidRoute.Put("/", handler.updateFolder)
 				folderUidRoute.Delete("/", handler.deleteFolder)
 				folderUidRoute.Get("/", handler.getFolder)
+				folderUidRoute.Get("/counts", handler.countFolderContent)
 			})
 		} else {
 			folderRoute.Post("/", authorize(accesscontrol.EvalPermission(dashboards.ActionFoldersCreate)), routing.Wrap(hs.CreateFolder))
@@ -74,6 +73,7 @@ func (hs *HTTPServer) registerFolderAPI(apiRoute routing.RouteRegister, authoriz
 				folderUidRoute.Put("/", authorize(accesscontrol.EvalPermission(dashboards.ActionFoldersWrite, uidScope)), routing.Wrap(hs.UpdateFolder))
 				folderUidRoute.Delete("/", authorize(accesscontrol.EvalPermission(dashboards.ActionFoldersDelete, uidScope)), routing.Wrap(hs.DeleteFolder))
 				folderUidRoute.Get("/", authorize(accesscontrol.EvalPermission(dashboards.ActionFoldersRead, uidScope)), routing.Wrap(hs.GetFolderByUID))
+				folderUidRoute.Get("/counts", authorize(accesscontrol.EvalPermission(dashboards.ActionFoldersRead, uidScope)), routing.Wrap(hs.GetFolderDescendantCounts))
 			})
 		}
 	})
@@ -747,6 +747,29 @@ func (fk8s *folderK8sHandler) getFolders(c *contextmodel.ReqContext) {
 	c.JSON(http.StatusOK, hits)
 }
 
+func (fk8s *folderK8sHandler) countFolderContent(c *contextmodel.ReqContext) {
+	client, ok := fk8s.getClient(c)
+	if !ok {
+		return
+	}
+
+	uid := web.Params(c.Req)[":uid"]
+
+	counts, err := client.Get(c.Req.Context(), uid, v1.GetOptions{}, "counts")
+	if err != nil {
+		fk8s.writeError(c, err)
+		return
+	}
+
+	out, err := toFolderLegacyCounts(counts)
+	if err != nil {
+		fk8s.writeError(c, err)
+		return
+	}
+
+	c.JSON(http.StatusOK, out)
+}
+
 func (fk8s *folderK8sHandler) getFolder(c *contextmodel.ReqContext) {
 	client, ok := fk8s.getClient(c)
 	if !ok {
@@ -988,7 +1011,7 @@ func (fk8s *folderK8sHandler) getFolderACMetadata(c *contextmodel.ReqContext, f
 		return nil, nil
 	}
 
-	folderIDs, err := fk8s.getParents(f)
+	folderIDs, err := getParents(f)
 	if err != nil {
 		return nil, err
 	}
@@ -1004,7 +1027,7 @@ func (fk8s *folderK8sHandler) getFolderACMetadata(c *contextmodel.ReqContext, f
 	return metadata, nil
 }
 
-func (fk8s *folderK8sHandler) getParents(f *folder.Folder) (map[string]bool, error) {
+func getParents(f *folder.Folder) (map[string]bool, error) {
 	folderIDs := map[string]bool{f.UID: true}
 	if (f.UID == accesscontrol.GeneralFolderUID) || (f.UID == folder.SharedWithMeFolderUID) {
 		return folderIDs, nil
@@ -1022,3 +1045,24 @@ func (fk8s *folderK8sHandler) getParents(f *folder.Folder) (map[string]bool, err
 
 	return folderIDs, nil
 }
+
+func toFolderLegacyCounts(u *unstructured.Unstructured) (*folder.DescendantCounts, error) {
+	ds, err := folderalpha1.UnstructuredToDescendantCounts(u)
+	if err != nil {
+		return nil, err
+	}
+
+	var out = make(folder.DescendantCounts)
+	for _, v := range ds.Counts {
+		// if stats come from unified storage, we will use them
+		if v.Group != "sql-fallback" {
+			out[v.Resource] = v.Count
+			continue
+		}
+		// if stats are from single tenant DB and they are not in unified storage, we will use them
+		if _, ok := out[v.Resource]; !ok {
+			out[v.Resource] = v.Count
+		}
+	}
+	return &out, nil
+}

pkg/api/folder_test.go

@@ -13,6 +13,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
+	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	clientrest "k8s.io/client-go/rest"
 
 	"github.com/grafana/grafana/pkg/api/dtos"
@@ -747,3 +748,118 @@ func TestUpdateFolderLegacyAndUnifiedStorage(t *testing.T) {
 		}
 	})
 }
+
+func TestToFolderCounts(t *testing.T) {
+	var tests = []struct {
+		name        string
+		input       *unstructured.Unstructured
+		expected    *folder.DescendantCounts
+		expectError bool
+	}{
+		{
+			name: "with only counts from unified storage",
+			input: &unstructured.Unstructured{
+				Object: map[string]interface{}{
+					"apiVersion": "folder.grafana.app/v0alpha1",
+					"counts": []interface{}{
+						map[string]interface{}{
+							"group":    "alpha",
+							"resource": "folders",
+							"count":    int64(1),
+						},
+						map[string]interface{}{
+							"group":    "alpha",
+							"resource": "dashboards",
+							"count":    int64(3),
+						},
+						map[string]interface{}{
+							"group":    "alpha",
+							"resource": "alertRules",
+							"count":    int64(0),
+						},
+					},
+				},
+			},
+			expected: &folder.DescendantCounts{
+				"folders":    1,
+				"dashboards": 3,
+				"alertRules": 0,
+			},
+		},
+		{
+			name: "with counts from both storages",
+			input: &unstructured.Unstructured{
+				Object: map[string]interface{}{
+					"apiVersion": "folder.grafana.app/v0alpha1",
+					"counts": []interface{}{
+						map[string]interface{}{
+							"group":    "alpha",
+							"resource": "folders",
+							"count":    int64(1),
+						},
+						map[string]interface{}{
+							"group":    "alpha",
+							"resource": "dashboards",
+							"count":    int64(3),
+						},
+						map[string]interface{}{
+							"group":    "sql-fallback",
+							"resource": "folders",
+							"count":    int64(0),
+						},
+					},
+				},
+			},
+			expected: &folder.DescendantCounts{
+				"folders":    1,
+				"dashboards": 3,
+			},
+		},
+		{
+			name: "it uses the values from sql-fallaback if not found in unified storage",
+			input: &unstructured.Unstructured{
+				Object: map[string]interface{}{
+					"apiVersion": "folder.grafana.app/v0alpha1",
+					"counts": []interface{}{
+						map[string]interface{}{
+							"group":    "alpha",
+							"resource": "dashboards",
+							"count":    int64(3),
+						},
+						map[string]interface{}{
+							"group":    "sql-fallback",
+							"resource": "folders",
+							"count":    int64(2),
+						},
+					},
+				},
+			},
+			expected: &folder.DescendantCounts{
+				"folders":    2,
+				"dashboards": 3,
+			},
+		},
+		{
+			name: "malformed input",
+			input: &unstructured.Unstructured{
+				Object: map[string]interface{}{
+					"apiVersion": "folder.grafana.app/v0alpha1",
+					"counts":     map[string]interface{}{},
+				},
+			},
+			expectError: true,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			actual, err := toFolderLegacyCounts(tc.input)
+			if tc.expectError {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+			require.Equal(t, tc.expected, actual)
+		})
+	}
+}

pkg/apis/folder/v0alpha1/types.go

@@ -2,6 +2,8 @@ package v0alpha1
 
 import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
+	"k8s.io/apimachinery/pkg/runtime"
 )
 
 // +k8s:deepcopy-gen:interfaces=k8s.io/apimachinery/pkg/runtime.Object
@@ -76,3 +78,9 @@ type ResourceStats struct {
 	Resource string `json:"resource"`
 	Count    int64  `json:"count"`
 }
+
+func UnstructuredToDescendantCounts(obj *unstructured.Unstructured) (*DescendantCounts, error) {
+	var res DescendantCounts
+	err := runtime.DefaultUnstructuredConverter.FromUnstructured(obj.Object, &res)
+	return &res, err
+}

pkg/registry/apis/folders/register.go

@@ -128,7 +128,7 @@ func (b *FolderAPIBuilder) UpdateAPIGroupInfo(apiGroupInfo *genericapiserver.API
 	storage[resourceInfo.StoragePath()] = legacyStore
 	storage[resourceInfo.StoragePath("parents")] = &subParentsREST{b.folderSvc}
 	storage[resourceInfo.StoragePath("access")] = &subAccessREST{b.folderSvc}
-	storage[resourceInfo.StoragePath("count")] = &subCountREST{searcher: b.searcher}
+	storage[resourceInfo.StoragePath("counts")] = &subCountREST{searcher: b.searcher}
 
 	// enable dual writer
 	if optsGetter != nil && dualWriteBuilder != nil {

pkg/tests/apis/folder/folders_test.go

@@ -93,7 +93,7 @@ func TestIntegrationFoldersApp(t *testing.T) {
 				]
 			  },
 			  {
-				"name": "folders/count",
+				"name": "folders/counts",
 				"singularName": "",
 				"namespaced": true,
 				"kind": "DescendantCounts",