Skip to content

Commit

Permalink
Add caching of scim/me API to databricks_current_user data source a…
Browse files Browse the repository at this point in the history
…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
  • Loading branch information
kartikgupta-db authored Mar 29, 2023
1 parent eb4fe2a commit 9e1eb13
Show file tree
Hide file tree
Showing 4 changed files with 59 additions and 10 deletions.
42 changes: 40 additions & 2 deletions common/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
7 changes: 5 additions & 2 deletions exporter/context.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down
7 changes: 4 additions & 3 deletions internal/acceptance/secret_scope_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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)
Expand Down
13 changes: 10 additions & 3 deletions permissions/resource_permissions.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@ import (
"github.com/databricks/terraform-provider-databricks/common"
"github.com/databricks/terraform-provider-databricks/jobs"
"github.com/databricks/terraform-provider-databricks/pipelines"
"github.com/databricks/terraform-provider-databricks/scim"

"github.com/hashicorp/go-cty/cty"
"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
Expand Down Expand Up @@ -141,7 +140,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
}
Expand Down Expand Up @@ -184,7 +187,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
}
Expand Down

0 comments on commit 9e1eb13

Please sign in to comment.