AuthZService: Cache folder tree (#98210)

* AuthZService: Cache folder tree

* Remove fmt

* Suggestion

* Add tests

pkg/services/authz/rbac/cache.go

@@ -1,23 +1,25 @@
 package rbac
 
-import "fmt"
-
 func userIdentifierCacheKey(namespace, userUID string) string {
-	return fmt.Sprintf("UID_%s_%s", namespace, userUID)
+	return "UID_" + namespace + "_" + userUID
 }
 
 func userIdentifierCacheKeyById(namespace, ID string) string {
-	return fmt.Sprintf("ID_%s_%s", namespace, ID)
+	return "ID_" + namespace + "_" + ID
 }
 
 func userPermCacheKey(namespace, userUID, action string) string {
-	return fmt.Sprintf("%s_%s_%s", namespace, userUID, action)
+	return namespace + "_" + userUID + "_" + action
 }
 
 func userBasicRoleCacheKey(namespace, userUID string) string {
-	return fmt.Sprintf("%s_%s", namespace, userUID)
+	return namespace + "_" + userUID
 }
 
 func userTeamCacheKey(namespace, userUID string) string {
-	return fmt.Sprintf("%s_%s", namespace, userUID)
+	return namespace + "_" + userUID
+}
+
+func folderCacheKey(namespace string) string {
+	return namespace
 }

pkg/services/authz/rbac/service.go

@@ -48,6 +48,7 @@ type Service struct {
 	permCache      *localcache.CacheService
 	teamCache      *localcache.CacheService
 	basicRoleCache *localcache.CacheService
+	folderCache    *localcache.CacheService
 }
 
 func NewService(sql legacysql.LegacyDatabaseProvider, identityStore legacy.LegacyIdentityStore, logger log.Logger, tracer tracing.Tracer) *Service {
@@ -61,6 +62,7 @@ func NewService(sql legacysql.LegacyDatabaseProvider, identityStore legacy.Legac
 		permCache:      localcache.New(shortCacheTTL, shortCleanupInterval),
 		teamCache:      localcache.New(shortCacheTTL, shortCleanupInterval),
 		basicRoleCache: localcache.New(longCacheTTL, longCleanupInterval),
+		folderCache:    localcache.New(shortCacheTTL, shortCleanupInterval),
 	}
 }
 
@@ -332,6 +334,11 @@ func (s *Service) checkInheritedPermissions(ctx context.Context, scopeMap map[st
 }
 
 func (s *Service) buildFolderTree(ctx context.Context, ns claims.NamespaceInfo) (map[string]FolderNode, error) {
+	key := folderCacheKey(ns.Value)
+	if cached, ok := s.folderCache.Get(key); ok {
+		return cached.(map[string]FolderNode), nil
+	}
+
 	folders, err := s.store.GetFolders(ctx, ns)
 	if err != nil {
 		return nil, fmt.Errorf("could not get folders: %w", err)
@@ -363,5 +370,7 @@ func (s *Service) buildFolderTree(ctx context.Context, ns claims.NamespaceInfo)
 		}
 	}
 
+	s.folderCache.Set(key, folderMap, 0)
+
 	return folderMap, nil
 }

pkg/services/authz/rbac/service_test.go

@@ -387,8 +387,95 @@ func TestService_getUserPermissions(t *testing.T) {
 	}
 }
 
+func TestService_buildFolderTree(t *testing.T) {
+	type testCase struct {
+		name         string
+		folders      []store.Folder
+		cacheHit     bool
+		expectedTree map[string]FolderNode
+	}
+
+	testCases := []testCase{
+		{
+			name: "should return folder tree from cache if available",
+			folders: []store.Folder{
+				{UID: "folder1", ParentUID: nil},
+				{UID: "folder2", ParentUID: strPtr("folder1")},
+			},
+			cacheHit: true,
+			expectedTree: map[string]FolderNode{
+				"folder1": {uid: "folder1", childrenUIDs: []string{"folder2"}},
+				"folder2": {uid: "folder2", parentUID: strPtr("folder1")},
+			},
+		},
+		{
+			name: "should return folder tree from store if not in cache",
+			folders: []store.Folder{
+				{UID: "folder1", ParentUID: nil},
+				{UID: "folder2", ParentUID: strPtr("folder1")},
+			},
+			cacheHit: false,
+			expectedTree: map[string]FolderNode{
+				"folder1": {uid: "folder1", childrenUIDs: []string{"folder2"}},
+				"folder2": {uid: "folder2", parentUID: strPtr("folder1")},
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx := context.Background()
+			ns := claims.NamespaceInfo{Value: "stacks-12", OrgID: 1, StackID: 12}
+
+			cacheService := localcache.New(shortCacheTTL, shortCleanupInterval)
+			if tc.cacheHit {
+				cacheService.Set(folderCacheKey(ns.Value), tc.expectedTree, 0)
+			}
+
+			store := &fakeStore{folders: tc.folders}
+
+			s := &Service{
+				store:       store,
+				folderCache: cacheService,
+				logger:      log.New("test"),
+			}
+
+			tree, err := s.buildFolderTree(ctx, ns)
+
+			require.NoError(t, err)
+			require.Len(t, tree, len(tc.expectedTree))
+			for _, folder := range tc.folders {
+				node, ok := tree[folder.UID]
+				require.True(t, ok)
+				// Check parent
+				if folder.ParentUID != nil {
+					require.NotNil(t, node.parentUID)
+					require.Equal(t, *folder.ParentUID, *node.parentUID)
+				} else {
+					require.Nil(t, node.parentUID)
+				}
+				// Check children
+				if len(node.childrenUIDs) > 0 {
+					epectedChildren := tc.expectedTree[folder.UID].childrenUIDs
+					require.ElementsMatch(t, node.childrenUIDs, epectedChildren)
+				}
+			}
+			if tc.cacheHit {
+				require.Zero(t, store.calls)
+			} else {
+				require.Equal(t, 1, store.calls)
+			}
+		})
+	}
+}
+
+func strPtr(s string) *string {
+	return &s
+}
+
 type fakeStore struct {
 	store.Store
+	folders         []store.Folder
 	basicRole       *store.BasicRole
 	userID          *store.UserIdentifiers
 	userPermissions []accesscontrol.Permission
@@ -420,6 +507,14 @@ func (f *fakeStore) GetUserPermissions(ctx context.Context, namespace claims.Nam
 	return f.userPermissions, nil
 }
 
+func (f *fakeStore) GetFolders(ctx context.Context, namespace claims.NamespaceInfo) ([]store.Folder, error) {
+	f.calls++
+	if f.err {
+		return nil, fmt.Errorf("store error")
+	}
+	return f.folders, nil
+}
+
 type fakeIdentityStore struct {
 	legacy.LegacyIdentityStore
 	teams []int64