steve/databricks-cli
steve
/
databricks-cli
mirror of https://github.com/databricks/cli.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Replace API call to test configuration with dummy authenticate call (#728) 

## Changes

This reduces the latency of every workspace command by the duration of a
single API call to retrieve the current user (which can take up to a
full second).

Note: the better place to verify that a request can be authenticated is
the SDK itself.

## Tests

* Unit test to confirm an the empty `*http.Request` can be constructed
* Manually confirmed that the additional API call no longer happens
This commit is contained in:
Pieter Noordhuis 2023-09-05 13:10:37 +02:00 committed by GitHub
parent bbbeabf98c
commit f62def3e77
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 27 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
cmd/root/auth.go
View File

@ -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,7 @@ 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)
err = w.Config.Authenticate(emptyHttpRequest(ctx))
if cmdio.IsInteractive(ctx) && errors.Is(err, config.ErrCannotConfigureAuth) {
profile, err := askForWorkspaceProfile()
if err != nil {
@ -107,7 +105,6 @@ TRY_AUTH: // or try picking a config profile dynamically
if err != nil {
return err
}
ctx = context.WithValue(ctx, &currentUser, me)
ctx = context.WithValue(ctx, &workspaceClient, w)
cmd.SetContext(ctx)
return nil
@ -194,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 {
@ -209,11 +217,3 @@ func AccountClient(ctx context.Context) *databricks.AccountClient {
}
return a
}
func Me(ctx context.Context) *iam.User {
me, ok := ctx.Value(&currentUser).(*iam.User)
if !ok {
panic("cannot get current user. Please report it as a bug")
}
return me
}

14
cmd/root/auth_test.go Normal file
View File

@ -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)
}