From 0354ff2ac130312da473791ab4c9015a30891832 Mon Sep 17 00:00:00 2001 From: kartikgupta-db Date: Wed, 29 Mar 2023 11:36:41 +0200 Subject: [PATCH 1/6] Cache scim/me API --- common/client.go | 37 +++++++++++++++++++++++- exporter/context.go | 7 +++-- internal/acceptance/secret_scope_test.go | 7 +++-- permissions/resource_permissions.go | 13 +++++++-- 4 files changed, 55 insertions(+), 9 deletions(-) diff --git a/common/client.go b/common/client.go index 2a6f35b59a..7d6521d6f1 100644 --- a/common/client.go +++ b/common/client.go @@ -10,9 +10,28 @@ import ( "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 CachedCurrentUserService struct { + client *client.DatabricksClient + oldImpl scim.CurrentUserService + cachedUser *scim.User +} + +func (a *CachedCurrentUserService) Me(ctx context.Context) (*scim.User, error) { + if a.cachedUser.DisplayName != "" { + return a.cachedUser, nil + } + user, err := a.oldImpl.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 @@ -22,10 +41,26 @@ type DatabricksClient struct { // callback used to create API1.2 call wrapper, which simplifies unit tessting commandFactory func(context.Context, *DatabricksClient) CommandExecutor + cachedUser scim.User } func (c *DatabricksClient) WorkspaceClient() (*databricks.WorkspaceClient, error) { - return databricks.NewWorkspaceClient((*databricks.Config)(c.DatabricksClient.Config)) + // return databricks.NewWorkspaceClient((*databricks.Config)(c.DatabricksClient.Config)) + w, err := databricks.NewWorkspaceClient((*databricks.Config)(c.DatabricksClient.Config)) + if err != nil { + return nil, err + } + return c.withCachedCurrentUserApi(w), nil +} + +func (c *DatabricksClient) withCachedCurrentUserApi(w *databricks.WorkspaceClient) *databricks.WorkspaceClient { + currentImpl := w.CurrentUser.Impl() + w.CurrentUser.WithImpl(&CachedCurrentUserService{ + client: c.DatabricksClient, + oldImpl: currentImpl, + cachedUser: &c.cachedUser, + }) + return w } // Get on path diff --git a/exporter/context.go b/exporter/context.go index b9065d5b82..de16b570fa 100644 --- a/exporter/context.go +++ b/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 } diff --git a/internal/acceptance/secret_scope_test.go b/internal/acceptance/secret_scope_test.go index efc1b3bc06..68ee0e9d6c 100644 --- a/internal/acceptance/secret_scope_test.go +++ b/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) diff --git a/permissions/resource_permissions.go b/permissions/resource_permissions.go index f8d68045c4..df488ba34d 100644 --- a/permissions/resource_permissions.go +++ b/permissions/resource_permissions.go @@ -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" @@ -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 } @@ -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 } From 5fc40415de912d597f21a809eac0ddc50a945e1b Mon Sep 17 00:00:00 2001 From: kartikgupta-db Date: Wed, 29 Mar 2023 12:03:26 +0200 Subject: [PATCH 2/6] add mutext to protect cached user --- common/client.go | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/common/client.go b/common/client.go index 7d6521d6f1..4a0e6bc689 100644 --- a/common/client.go +++ b/common/client.go @@ -6,6 +6,7 @@ import ( "log" "net/http" "strings" + "sync" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/client" @@ -18,17 +19,22 @@ type CachedCurrentUserService struct { client *client.DatabricksClient oldImpl scim.CurrentUserService cachedUser *scim.User + mu *sync.Mutex } func (a *CachedCurrentUserService) Me(ctx context.Context) (*scim.User, error) { + a.mu.Lock() if a.cachedUser.DisplayName != "" { + a.mu.Unlock() return a.cachedUser, nil } user, err := a.oldImpl.Me(ctx) if err != nil { + a.mu.Unlock() return user, err } *a.cachedUser = *user + a.mu.Unlock() return user, err } @@ -42,10 +48,10 @@ type DatabricksClient struct { // callback used to create API1.2 call wrapper, which simplifies unit tessting commandFactory func(context.Context, *DatabricksClient) CommandExecutor cachedUser scim.User + mu sync.Mutex } func (c *DatabricksClient) WorkspaceClient() (*databricks.WorkspaceClient, error) { - // return databricks.NewWorkspaceClient((*databricks.Config)(c.DatabricksClient.Config)) w, err := databricks.NewWorkspaceClient((*databricks.Config)(c.DatabricksClient.Config)) if err != nil { return nil, err @@ -59,6 +65,7 @@ func (c *DatabricksClient) withCachedCurrentUserApi(w *databricks.WorkspaceClien client: c.DatabricksClient, oldImpl: currentImpl, cachedUser: &c.cachedUser, + mu: &c.mu, }) return w } From 521d427a0d55cb0b065156a6f0beb3c0c92edfed Mon Sep 17 00:00:00 2001 From: kartikgupta-db Date: Wed, 29 Mar 2023 12:07:19 +0200 Subject: [PATCH 3/6] naming fix --- common/client.go | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/common/client.go b/common/client.go index 4a0e6bc689..1dfeb16b44 100644 --- a/common/client.go +++ b/common/client.go @@ -15,20 +15,20 @@ import ( "github.com/golang-jwt/jwt/v4" ) -type CachedCurrentUserService struct { - client *client.DatabricksClient - oldImpl scim.CurrentUserService - cachedUser *scim.User - mu *sync.Mutex +type cachedCurrentUserService struct { + client *client.DatabricksClient + internalImpl scim.CurrentUserService + cachedUser *scim.User + mu *sync.Mutex } -func (a *CachedCurrentUserService) Me(ctx context.Context) (*scim.User, error) { +func (a *cachedCurrentUserService) Me(ctx context.Context) (*scim.User, error) { a.mu.Lock() if a.cachedUser.DisplayName != "" { a.mu.Unlock() return a.cachedUser, nil } - user, err := a.oldImpl.Me(ctx) + user, err := a.internalImpl.Me(ctx) if err != nil { a.mu.Unlock() return user, err @@ -60,12 +60,12 @@ func (c *DatabricksClient) WorkspaceClient() (*databricks.WorkspaceClient, error } func (c *DatabricksClient) withCachedCurrentUserApi(w *databricks.WorkspaceClient) *databricks.WorkspaceClient { - currentImpl := w.CurrentUser.Impl() - w.CurrentUser.WithImpl(&CachedCurrentUserService{ - client: c.DatabricksClient, - oldImpl: currentImpl, - cachedUser: &c.cachedUser, - mu: &c.mu, + internalImpl := w.CurrentUser.Impl() + w.CurrentUser.WithImpl(&cachedCurrentUserService{ + client: c.DatabricksClient, + internalImpl: internalImpl, + cachedUser: &c.cachedUser, + mu: &c.mu, }) return w } From 576cd001b45a69ac010e7245bdee573d7a690be3 Mon Sep 17 00:00:00 2001 From: kartikgupta-db Date: Wed, 29 Mar 2023 12:20:42 +0200 Subject: [PATCH 4/6] defer mu.unlock and remove unused code --- common/client.go | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/common/client.go b/common/client.go index 1dfeb16b44..f1e4cd03a0 100644 --- a/common/client.go +++ b/common/client.go @@ -16,7 +16,6 @@ import ( ) type cachedCurrentUserService struct { - client *client.DatabricksClient internalImpl scim.CurrentUserService cachedUser *scim.User mu *sync.Mutex @@ -24,17 +23,15 @@ type cachedCurrentUserService struct { func (a *cachedCurrentUserService) Me(ctx context.Context) (*scim.User, error) { a.mu.Lock() + defer a.mu.Unlock() if a.cachedUser.DisplayName != "" { - a.mu.Unlock() return a.cachedUser, nil } user, err := a.internalImpl.Me(ctx) if err != nil { - a.mu.Unlock() return user, err } *a.cachedUser = *user - a.mu.Unlock() return user, err } @@ -62,7 +59,6 @@ func (c *DatabricksClient) WorkspaceClient() (*databricks.WorkspaceClient, error func (c *DatabricksClient) withCachedCurrentUserApi(w *databricks.WorkspaceClient) *databricks.WorkspaceClient { internalImpl := w.CurrentUser.Impl() w.CurrentUser.WithImpl(&cachedCurrentUserService{ - client: c.DatabricksClient, internalImpl: internalImpl, cachedUser: &c.cachedUser, mu: &c.mu, From 2ef5223ce895f631fcbeb8b6485498761362758f Mon Sep 17 00:00:00 2001 From: kartikgupta-db Date: Wed, 29 Mar 2023 12:53:24 +0200 Subject: [PATCH 5/6] cache ws client --- common/client.go | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/common/client.go b/common/client.go index f1e4cd03a0..0c06dc4dcc 100644 --- a/common/client.go +++ b/common/client.go @@ -18,20 +18,20 @@ import ( type cachedCurrentUserService struct { internalImpl scim.CurrentUserService cachedUser *scim.User - mu *sync.Mutex + mu sync.Mutex } func (a *cachedCurrentUserService) Me(ctx context.Context) (*scim.User, error) { a.mu.Lock() defer a.mu.Unlock() - if a.cachedUser.DisplayName != "" { + if a.cachedUser != nil { return a.cachedUser, nil } user, err := a.internalImpl.Me(ctx) if err != nil { return user, err } - *a.cachedUser = *user + a.cachedUser = user return user, err } @@ -43,16 +43,22 @@ type DatabricksClient struct { *client.DatabricksClient // callback used to create API1.2 call wrapper, which simplifies unit tessting - commandFactory func(context.Context, *DatabricksClient) CommandExecutor - cachedUser scim.User - mu sync.Mutex + commandFactory func(context.Context, *DatabricksClient) CommandExecutor + cachedWorkspaceClient *databricks.WorkspaceClient + mu sync.Mutex } func (c *DatabricksClient) WorkspaceClient() (*databricks.WorkspaceClient, error) { + 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 } + c.cachedWorkspaceClient = w return c.withCachedCurrentUserApi(w), nil } @@ -60,8 +66,6 @@ func (c *DatabricksClient) withCachedCurrentUserApi(w *databricks.WorkspaceClien internalImpl := w.CurrentUser.Impl() w.CurrentUser.WithImpl(&cachedCurrentUserService{ internalImpl: internalImpl, - cachedUser: &c.cachedUser, - mu: &c.mu, }) return w } From 9c1845cf6f6d4252a873ee2d5c0cc7dd6f4ccf22 Mon Sep 17 00:00:00 2001 From: kartikgupta-db Date: Wed, 29 Mar 2023 12:56:07 +0200 Subject: [PATCH 6/6] inline method and rename --- common/client.go | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/common/client.go b/common/client.go index 0c06dc4dcc..2bf31f7634 100644 --- a/common/client.go +++ b/common/client.go @@ -15,13 +15,13 @@ import ( "github.com/golang-jwt/jwt/v4" ) -type cachedCurrentUserService struct { +type cachedMe struct { internalImpl scim.CurrentUserService cachedUser *scim.User mu sync.Mutex } -func (a *cachedCurrentUserService) Me(ctx context.Context) (*scim.User, error) { +func (a *cachedMe) Me(ctx context.Context) (*scim.User, error) { a.mu.Lock() defer a.mu.Unlock() if a.cachedUser != nil { @@ -58,16 +58,12 @@ func (c *DatabricksClient) WorkspaceClient() (*databricks.WorkspaceClient, error if err != nil { return nil, err } - c.cachedWorkspaceClient = w - return c.withCachedCurrentUserApi(w), nil -} - -func (c *DatabricksClient) withCachedCurrentUserApi(w *databricks.WorkspaceClient) *databricks.WorkspaceClient { internalImpl := w.CurrentUser.Impl() - w.CurrentUser.WithImpl(&cachedCurrentUserService{ + w.CurrentUser.WithImpl(&cachedMe{ internalImpl: internalImpl, }) - return w + c.cachedWorkspaceClient = w + return w, nil } // Get on path