[Internal] Caching of group membership in databricks_group_member (#4581)

## Changes

On the installations with large number of users the provider creates a
large number of API calls to the Databricks API leading to slow
deployments. The provider also requests all members for each group
membership check.

In this change group member resource use expiring caches - in our case
we halved the default deployment times on existing installation and
eliminated Databricks API throttling.

I am also looking for some guidance on how to make this change more
generic.

## Tests

We run this provider in production with about 800 users and ~100 groups.

- [x] `make test` run locally
- [ ] relevant change in `docs/` folder
- [ ] covered with integration tests in `internal/acceptance`
- [x] using Go SDK
- [x] using TF Plugin Framework

---------

Co-authored-by: Ruslan Abdulkhalikov <[email protected]>
Co-authored-by: Copilot <[email protected]>
Co-authored-by: Alex Ott <[email protected]>
Co-authored-by: vuong-nguyen <[email protected]>
Co-authored-by: Alex Ott <[email protected]>
Co-authored-by: Miles Yucht <[email protected]>

Author: 6 people

NEXT_CHANGELOG.md

@@ -13,3 +13,5 @@
 ### Exporter
 
 ### Internal Changes
+
+* Caching group membership in `databricks_group_member` to improve performance ([#4581](https://github.com/databricks/terraform-provider-databricks/pull/4581)).

scim/groups.go

@@ -31,11 +31,9 @@ func (a GroupsAPI) Create(scimGroupRequest Group) (group Group, err error) {
 
 // Read reads and returns a Group object via SCIM api
 func (a GroupsAPI) Read(groupID, attributes string) (group Group, err error) {
-	err = a.client.Scim(a.context, http.MethodGet, fmt.Sprintf(
-		"/preview/scim/v2/Groups/%v?attributes=%s", groupID, attributes), nil, &group)
-	if err != nil {
-		return
-	}
+	key := fmt.Sprintf(
+		"/preview/scim/v2/Groups/%v?attributes=%s", groupID, attributes)
+	err = a.client.Scim(a.context, http.MethodGet, key, nil, &group)
 	return
 }
 

scim/resource_group_member.go

@@ -3,21 +3,128 @@ package scim
 import (
 	"context"
 	"fmt"
+	"sync"
 
 	"github.com/databricks/databricks-sdk-go/apierr"
 	"github.com/databricks/terraform-provider-databricks/common"
+	"github.com/hashicorp/terraform-plugin-log/tflog"
 )
 
+// groupMembersInfo is the set of members that belong to a group. Read and write access
+// to the underlying members map must be protected by the lock.
+type groupMembersInfo struct {
+	initialized bool
+	members     map[string]struct{}
+	lock        sync.Mutex
+}
+
+// groupCache is a cache of groups to their members by their ID. Access to the cache
+// must be protected by the lock.
+type groupCache struct {
+	// The mapping of cached members for this group. The cache key is the group ID.
+	// TODO: add workspace ID to the cache key when account-level and workspace-level providers are unified.
+	cache map[string]*groupMembersInfo
+	lock  sync.Mutex
+}
+
+func newGroupCache() *groupCache {
+	return &groupCache{
+		cache: make(map[string]*groupMembersInfo),
+		lock:  sync.Mutex{},
+	}
+}
+
+func (gc *groupCache) getOrCreateGroupInfo(groupID string) *groupMembersInfo {
+	gc.lock.Lock()
+	defer gc.lock.Unlock()
+
+	groupInfo, exists := gc.cache[groupID]
+	if !exists {
+		groupInfo = &groupMembersInfo{
+			initialized: false,
+			members:     make(map[string]struct{}),
+			lock:        sync.Mutex{},
+		}
+		gc.cache[groupID] = groupInfo
+	}
+	return groupInfo
+}
+
+func (gc *groupCache) getMembers(api GroupsAPI, groupID string) (map[string]struct{}, error) {
+	groupInfo := gc.getOrCreateGroupInfo(groupID)
+	groupInfo.lock.Lock()
+	defer groupInfo.lock.Unlock()
+
+	if !groupInfo.initialized {
+		tflog.Debug(api.context, fmt.Sprintf("Getting members for group %s", groupID))
+		group, err := api.Read(groupID, "members")
+		if err != nil {
+			return nil, err
+		}
+		tflog.Debug(api.context, fmt.Sprintf("Group %s has %d members", groupID, len(group.Members)))
+		for _, member := range group.Members {
+			groupInfo.members[member.Value] = struct{}{}
+		}
+		groupInfo.initialized = true
+		tflog.Debug(api.context, fmt.Sprintf("Group %s has %d members (initialized)", groupID, len(groupInfo.members)))
+	} else {
+		tflog.Debug(api.context, fmt.Sprintf("Group %s has %d members (cached)", groupID, len(groupInfo.members)))
+	}
+	membersCopy := make(map[string]struct{}, len(groupInfo.members))
+	for k, v := range groupInfo.members {
+		membersCopy[k] = v
+	}
+	return membersCopy, nil
+}
+
+func (gc *groupCache) removeMember(api GroupsAPI, groupID string, memberID string) error {
+	groupInfo := gc.getOrCreateGroupInfo(groupID)
+	groupInfo.lock.Lock()
+	defer groupInfo.lock.Unlock()
+
+	err := api.Patch(groupID, PatchRequest(
+		"remove", fmt.Sprintf(`members[value eq "%s"]`, memberID)))
+	if err != nil {
+		return err
+	}
+
+	if groupInfo.initialized {
+		delete(groupInfo.members, memberID)
+	}
+	return nil
+}
+
+func (gc *groupCache) addMember(api GroupsAPI, groupID string, memberID string) error {
+	groupInfo := gc.getOrCreateGroupInfo(groupID)
+	groupInfo.lock.Lock()
+	defer groupInfo.lock.Unlock()
+
+	err := api.Patch(groupID, PatchRequestWithValue("add", "members", memberID))
+	if err != nil {
+		return err
+	}
+	if groupInfo.initialized {
+		groupInfo.members[memberID] = struct{}{}
+	}
+	return nil
+}
+
+func hasMember(members map[string]struct{}, memberID string) bool {
+	_, ok := members[memberID]
+	return ok
+}
+
+var globalGroupsCache = newGroupCache()
+
 // ResourceGroupMember bind group with member
 func ResourceGroupMember() common.Resource {
 	return common.NewPairID("group_id", "member_id").BindResource(common.BindResource{
 		CreateContext: func(ctx context.Context, groupID, memberID string, c *common.DatabricksClient) error {
-			return NewGroupsAPI(ctx, c).Patch(groupID, PatchRequestWithValue("add", "members", memberID))
+			return globalGroupsCache.addMember(NewGroupsAPI(ctx, c), groupID, memberID)
 		},
 		ReadContext: func(ctx context.Context, groupID, memberID string, c *common.DatabricksClient) error {
-			group, err := NewGroupsAPI(ctx, c).Read(groupID, "members")
-			hasMember := ComplexValues(group.Members).HasValue(memberID)
-			if err == nil && !hasMember {
+			members, err := globalGroupsCache.getMembers(NewGroupsAPI(ctx, c), groupID)
+			if err == nil && !hasMember(members, memberID) {
 				return &apierr.APIError{
 					ErrorCode:  "NOT_FOUND",
 					StatusCode: 404,
@@ -27,8 +134,7 @@ func ResourceGroupMember() common.Resource {
 			return err
 		},
 		DeleteContext: func(ctx context.Context, groupID, memberID string, c *common.DatabricksClient) error {
-			return NewGroupsAPI(ctx, c).Patch(groupID, PatchRequest(
-				"remove", fmt.Sprintf(`members[value eq "%s"]`, memberID)))
+			return globalGroupsCache.removeMember(NewGroupsAPI(ctx, c), groupID, memberID)
 		},
 	})
 }

scim/resource_group_member_test.go

@@ -9,6 +9,8 @@ import (
 )
 
 func TestResourceGroupMemberCreate(t *testing.T) {
+	globalGroupsCache = newGroupCache()
+
 	d, err := qa.ResourceFixture{
 		Fixtures: []qa.HTTPFixture{
 			{
@@ -43,6 +45,15 @@ func TestResourceGroupMemberCreate(t *testing.T) {
 	}.Apply(t)
 	assert.NoError(t, err)
 	assert.Equal(t, "abc|bcd", d.Id())
+
+	assert.Equal(t, 1, len(globalGroupsCache.cache), "Cache should contain one entry after create")
+	groupInfo, exists := globalGroupsCache.cache["abc"]
+	assert.True(t, exists, "Cache should contain group 'abc'")
+
+	assert.True(t, groupInfo.initialized, "Group info should be initialized")
+	assert.Equal(t, 1, len(groupInfo.members), "Should have one member cached")
+	_, memberExists := groupInfo.members["bcd"]
+	assert.True(t, memberExists, "Member 'bcd' should be in cache")
 }
 
 func TestResourceGroupMemberCreate_Error(t *testing.T) {
@@ -70,6 +81,8 @@ func TestResourceGroupMemberCreate_Error(t *testing.T) {
 }
 
 func TestResourceGroupMemberRead(t *testing.T) {
+	globalGroupsCache = newGroupCache()
+
 	d, err := qa.ResourceFixture{
 		Fixtures: []qa.HTTPFixture{
 			{
@@ -93,6 +106,15 @@ func TestResourceGroupMemberRead(t *testing.T) {
 	}.Apply(t)
 	assert.NoError(t, err)
 	assert.Equal(t, "abc|bcd", d.Id(), "Id should not be empty")
+
+	assert.Equal(t, 1, len(globalGroupsCache.cache), "Cache should contain one entry after read")
+	groupInfo, exists := globalGroupsCache.cache["abc"]
+	assert.True(t, exists, "Cache should contain group 'abc'")
+
+	assert.True(t, groupInfo.initialized, "Group info should be initialized")
+	assert.Equal(t, 1, len(groupInfo.members), "Should have one member cached")
+	_, memberExists := groupInfo.members["bcd"]
+	assert.True(t, memberExists, "Member 'bcd' should be in cache")
 }
 
 func TestResourceGroupMemberRead_NoMember(t *testing.T) {
@@ -157,6 +179,20 @@ func TestResourceGroupMemberRead_Error(t *testing.T) {
 }
 
 func TestResourceGroupMemberDelete(t *testing.T) {
+	globalGroupsCache = newGroupCache()
+
+	groupInfo := globalGroupsCache.getOrCreateGroupInfo("abc")
+	groupInfo.initialized = true
+	groupInfo.members["bcd"] = struct{}{}
+
+	assert.Equal(t, 1, len(globalGroupsCache.cache), "Cache should contain one entry initially")
+	groupInfo, exists := globalGroupsCache.cache["abc"]
+	assert.True(t, exists, "Cache should contain group 'abc'")
+
+	assert.Equal(t, 1, len(groupInfo.members), "Should have one member initially")
+	_, memberExists := groupInfo.members["bcd"]
+	assert.True(t, memberExists, "Member 'bcd' should be in cache initially")
+
 	d, err := qa.ResourceFixture{
 		Fixtures: []qa.HTTPFixture{
 			{
@@ -173,6 +209,13 @@ func TestResourceGroupMemberDelete(t *testing.T) {
 	}.Apply(t)
 	assert.NoError(t, err)
 	assert.Equal(t, "abc|bcd", d.Id())
+
+	groupInfo, exists = globalGroupsCache.cache["abc"]
+	assert.True(t, exists, "Cache should still contain group 'abc'")
+
+	assert.Equal(t, 0, len(groupInfo.members), "Should have no members after delete")
+	_, memberExists = groupInfo.members["bcd"]
+	assert.False(t, memberExists, "Member 'bcd' should not be in cache after delete")
 }
 
 func TestResourceGroupMemberDelete_Error(t *testing.T) {