From 0c23b426d1d28036c4bc6c40f6ec45d97518adc3 Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Mon, 4 Sep 2023 23:13:40 +0200 Subject: [PATCH 1/3] Replace API call to test configuration with authenticate call This reduces the latency of every workspace command by the duration of a single `scim/Me` API call (which can take up to a full second). --- cmd/root/auth.go | 23 ++++++++++------------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/cmd/root/auth.go b/cmd/root/auth.go index e56074ef4d..3961ac9b0e 100644 --- a/cmd/root/auth.go +++ b/cmd/root/auth.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "net/http" "os" "github.com/databricks/cli/bundle" @@ -11,7 +12,6 @@ import ( "github.com/databricks/cli/libs/databrickscfg" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/config" - "github.com/databricks/databricks-sdk-go/service/iam" "github.com/manifoldco/promptui" "github.com/spf13/cobra" ) @@ -19,7 +19,6 @@ import ( // Placeholders to use as unique keys in context.Context. var workspaceClient int var accountClient int -var currentUser int func initProfileFlag(cmd *cobra.Command) { cmd.PersistentFlags().StringP("profile", "p", "", "~/.databrickscfg profile") @@ -94,8 +93,15 @@ TRY_AUTH: // or try picking a config profile dynamically if err != nil { return err } - // get current user identity also to verify validity of configuration - me, err := w.CurrentUser.Me(ctx) + + // Verify that the client is configured correctly by authenticating a dummy request. + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://localhost/doesntmatter", nil) + if err != nil { + // This can only fail if the context is nil, which would be a bug. + panic(err) + } + + err = w.Config.Authenticate(req) if cmdio.IsInteractive(ctx) && errors.Is(err, config.ErrCannotConfigureAuth) { profile, err := askForWorkspaceProfile() if err != nil { @@ -107,7 +113,6 @@ TRY_AUTH: // or try picking a config profile dynamically if err != nil { return err } - ctx = context.WithValue(ctx, ¤tUser, me) ctx = context.WithValue(ctx, &workspaceClient, w) cmd.SetContext(ctx) return nil @@ -209,11 +214,3 @@ func AccountClient(ctx context.Context) *databricks.AccountClient { } return a } - -func Me(ctx context.Context) *iam.User { - me, ok := ctx.Value(¤tUser).(*iam.User) - if !ok { - panic("cannot get current user. Please report it as a bug") - } - return me -} From 46e64160ba81f5e8e5179d921e9126fa35678994 Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Mon, 4 Sep 2023 23:28:09 +0200 Subject: [PATCH 2/3] Use helper to produce empty http.Request --- cmd/root/auth.go | 20 ++++++++++++-------- cmd/root/auth_test.go | 14 ++++++++++++++ 2 files changed, 26 insertions(+), 8 deletions(-) create mode 100644 cmd/root/auth_test.go diff --git a/cmd/root/auth.go b/cmd/root/auth.go index 3961ac9b0e..56a6bbc025 100644 --- a/cmd/root/auth.go +++ b/cmd/root/auth.go @@ -94,14 +94,7 @@ TRY_AUTH: // or try picking a config profile dynamically return err } - // Verify that the client is configured correctly by authenticating a dummy request. - req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://localhost/doesntmatter", nil) - if err != nil { - // This can only fail if the context is nil, which would be a bug. - panic(err) - } - - err = w.Config.Authenticate(req) + err = w.Config.Authenticate(emptyHttpRequest(ctx)) if cmdio.IsInteractive(ctx) && errors.Is(err, config.ErrCannotConfigureAuth) { profile, err := askForWorkspaceProfile() if err != nil { @@ -214,3 +207,14 @@ func AccountClient(ctx context.Context) *databricks.AccountClient { } return a } + +// To verify that a client is configured correctly, we pass an empty HTTP request +// to a client's `config.Authenticate` function. Note: this functionality +// should be supported by the SDK itself. +func emptyHttpRequest(ctx context.Context) *http.Request { + req, err := http.NewRequestWithContext(ctx, "", "", nil) + if err != nil { + panic(err) + } + return req +} diff --git a/cmd/root/auth_test.go b/cmd/root/auth_test.go new file mode 100644 index 0000000000..75d255b58b --- /dev/null +++ b/cmd/root/auth_test.go @@ -0,0 +1,14 @@ +package root + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestEmptyHttpRequest(t *testing.T) { + ctx, _ := context.WithCancel(context.Background()) + req := emptyHttpRequest(ctx) + assert.Equal(t, req.Context(), ctx) +} From 27a5425bb629f0ed400c982c5f57c03a8dfd55af Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Mon, 4 Sep 2023 23:29:01 +0200 Subject: [PATCH 3/3] Move around --- cmd/root/auth.go | 23 +++++++++++------------ 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/cmd/root/auth.go b/cmd/root/auth.go index 56a6bbc025..d4c9a31b92 100644 --- a/cmd/root/auth.go +++ b/cmd/root/auth.go @@ -93,7 +93,6 @@ TRY_AUTH: // or try picking a config profile dynamically if err != nil { return err } - err = w.Config.Authenticate(emptyHttpRequest(ctx)) if cmdio.IsInteractive(ctx) && errors.Is(err, config.ErrCannotConfigureAuth) { profile, err := askForWorkspaceProfile() @@ -192,6 +191,17 @@ func askForAccountProfile() (string, error) { return profiles[i].Name, nil } +// To verify that a client is configured correctly, we pass an empty HTTP request +// to a client's `config.Authenticate` function. Note: this functionality +// should be supported by the SDK itself. +func emptyHttpRequest(ctx context.Context) *http.Request { + req, err := http.NewRequestWithContext(ctx, "", "", nil) + if err != nil { + panic(err) + } + return req +} + func WorkspaceClient(ctx context.Context) *databricks.WorkspaceClient { w, ok := ctx.Value(&workspaceClient).(*databricks.WorkspaceClient) if !ok { @@ -207,14 +217,3 @@ func AccountClient(ctx context.Context) *databricks.AccountClient { } return a } - -// To verify that a client is configured correctly, we pass an empty HTTP request -// to a client's `config.Authenticate` function. Note: this functionality -// should be supported by the SDK itself. -func emptyHttpRequest(ctx context.Context) *http.Request { - req, err := http.NewRequestWithContext(ctx, "", "", nil) - if err != nil { - panic(err) - } - return req -}