Add caching of scim/me API to databricks_current_user data source a…

…nd `databricks_permissions` resource. (#2170)

* Cache scim/me API

* add mutext to protect cached user

* naming fix

* defer mu.unlock and remove unused code

* cache ws client

* inline method and rename

common/client.go

@@ -6,13 +6,35 @@ import (
  "log"
  "net/http"
  "strings"
+ "sync"
 
  "github.com/databricks/databricks-sdk-go"
  "github.com/databricks/databricks-sdk-go/client"
  "github.com/databricks/databricks-sdk-go/config"
+ "github.com/databricks/databricks-sdk-go/service/scim"
  "github.com/golang-jwt/jwt/v4"
 )
 
+type cachedMe struct {
+ internalImpl scim.CurrentUserService
+ cachedUser *scim.User
+ mu sync.Mutex
+}
+
+func (a *cachedMe) Me(ctx context.Context) (*scim.User, error) {
+ a.mu.Lock()
+ defer a.mu.Unlock()
+ if a.cachedUser != nil {
+ return a.cachedUser, nil
+ }
+ user, err := a.internalImpl.Me(ctx)
+ if err != nil {
+ return user, err
+ }
+ a.cachedUser = user
+ return user, err
+}
+
 // DatabricksClient holds properties needed for authentication and HTTP client setup
 // fields with `name` struct tags become Terraform provider attributes. `env` struct tag
 // can hold one or more coma-separated env variable names to find value, if not specified
@@ -21,11 +43,27 @@ type DatabricksClient struct {
  *client.DatabricksClient
 
  // callback used to create API1.2 call wrapper, which simplifies unit tessting
- commandFactory func(context.Context, *DatabricksClient) CommandExecutor
+ commandFactory func(context.Context, *DatabricksClient) CommandExecutor
+ cachedWorkspaceClient *databricks.WorkspaceClient
+ mu sync.Mutex
 }
 
 func (c *DatabricksClient) WorkspaceClient() (*databricks.WorkspaceClient, error) {
- return databricks.NewWorkspaceClient((*databricks.Config)(c.DatabricksClient.Config))
+ c.mu.Lock()
+ defer c.mu.Unlock()
+ if c.cachedWorkspaceClient != nil {
+ return c.cachedWorkspaceClient, nil
+ }
+ w, err := databricks.NewWorkspaceClient((*databricks.Config)(c.DatabricksClient.Config))
+ if err != nil {
+ return nil, err
+ }
+ internalImpl := w.CurrentUser.Impl()
+ w.CurrentUser.WithImpl(&cachedMe{
+ internalImpl: internalImpl,
+ })
+ c.cachedWorkspaceClient = w
+ return w, nil
 }
 
 // Get on path

exporter/context.go

@@ -157,8 +157,11 @@ func (ic *importContext) Run() error {
  } else if !info.IsDir() {
  return fmt.Errorf("the path %s is not a directory", ic.Directory)
  }
- usersAPI := scim.NewUsersAPI(ic.Context, ic.Client)
- me, err := usersAPI.Me()
+ w, err := ic.Client.WorkspaceClient()
+ if err != nil {
+ return err
+ }
+ me, err := w.CurrentUser.Me(ic.Context)
  if err != nil {
  return err
  }

internal/acceptance/secret_scope_test.go

@@ -6,7 +6,6 @@ import (
  "testing"
 
  "github.com/databricks/terraform-provider-databricks/qa"
- "github.com/databricks/terraform-provider-databricks/scim"
  "github.com/databricks/terraform-provider-databricks/secrets"
 
  "github.com/databricks/terraform-provider-databricks/common"
@@ -33,8 +32,10 @@ func TestAccSecretScopeResource(t *testing.T) {
  acls, err := secretACLAPI.List(id)
  require.NoError(t, err)
 
- usersAPI := scim.NewUsersAPI(ctx, client)
- me, err := usersAPI.Me()
+ w, err := client.WorkspaceClient()
+ require.NoError(t, err)
+
+ me, err := w.CurrentUser.Me(ctx)
  require.NoError(t, err)
  assert.Equal(t, 1, len(acls))
  assert.Equal(t, me.UserName, acls[0].Principal)

permissions/resource_permissions.go

@@ -11,8 +11,6 @@ import (
  "github.com/databricks/databricks-sdk-go"
  "github.com/databricks/databricks-sdk-go/apierr"
  "github.com/databricks/terraform-provider-databricks/common"
- "github.com/databricks/terraform-provider-databricks/scim"
-
  "github.com/hashicorp/go-cty/cty"
  "github.com/hashicorp/terraform-plugin-sdk/v2/diag"
  "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
@@ -139,7 +137,11 @@ func (a PermissionsAPI) ensureCurrentUserCanManageObject(objectID string, object
  if !a.shouldExplicitlyGrantCallingUserManagePermissions(objectID) {
  return objectACL, nil
  }
- me, err := scim.NewUsersAPI(a.context, a.client).Me()
+ w, err := a.client.WorkspaceClient()
+ if err != nil {
+ return objectACL, err
+ }
+ me, err := w.CurrentUser.Me(a.context)
  if err != nil {
  return objectACL, err
  }
@@ -182,7 +184,11 @@ func (a PermissionsAPI) Update(objectID string, objectACL AccessControlChangeLis
  }
  }
  if owners == 0 {
- me, err := scim.NewUsersAPI(a.context, a.client).Me()
+ w, err := a.client.WorkspaceClient()
+ if err != nil {
+ return err
+ }
+ me, err := w.CurrentUser.Me(a.context)
  if err != nil {
  return err
  }