diff --git a/.nextchanges/cli/oauth-callback-port.md b/.nextchanges/cli/oauth-callback-port.md new file mode 100644 index 00000000000..52feaf4e07e --- /dev/null +++ b/.nextchanges/cli/oauth-callback-port.md @@ -0,0 +1 @@ +Fixed `databricks auth login` to honor `DATABRICKS_OAUTH_CALLBACK_PORT`. diff --git a/cmd/auth/login.go b/cmd/auth/login.go index f9f2531ac74..86c5199baf9 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -49,9 +49,22 @@ const ( // host used by the discovery login flow. Intended for testing and // development against non-production environments. See WithDiscoveryHost // in github.com/databricks/databricks-sdk-go/credentials/u2m. - discoveryHostEnvVar = "DATABRICKS_DISCOVERY_HOST" + discoveryHostEnvVar = "DATABRICKS_DISCOVERY_HOST" + oauthCallbackPortEnvVar = "DATABRICKS_OAUTH_CALLBACK_PORT" ) +func withOAuthCallbackPort(ctx context.Context, opts []u2m.PersistentAuthOption) ([]u2m.PersistentAuthOption, error) { + value := env.Get(ctx, oauthCallbackPortEnvVar) + if value == "" { + return opts, nil + } + port, err := strconv.Atoi(value) + if err != nil { + return nil, fmt.Errorf("invalid %s value %q: %w", oauthCallbackPortEnvVar, value, err) + } + return append(opts, u2m.WithPort(port)), nil +} + // discoveryErr wraps an error (or creates a new one) and appends the // discovery fallback tip so users know they can bypass login.databricks.com. func discoveryErr(msg string, err error) error { @@ -303,6 +316,10 @@ a new profile is created. u2m.WithBrowser(getBrowserFunc(cmd)), u2m.WithTokenCache(storage.WrapForOAuthArgument(ctx, tokenStore, mode, oauthArgument)), } + persistentAuthOpts, err = withOAuthCallbackPort(ctx, persistentAuthOpts) + if err != nil { + return err + } if len(scopesList) > 0 { persistentAuthOpts = append(persistentAuthOpts, u2m.WithScopes(scopesList)) } @@ -661,6 +678,10 @@ func discoveryLogin(ctx context.Context, in discoveryLoginInputs) error { u2m.WithDiscoveryLogin(), u2m.WithTokenCache(storage.WrapForOAuthArgument(ctx, in.tokenStore, in.mode, arg)), } + opts, err = withOAuthCallbackPort(ctx, opts) + if err != nil { + return discoveryErr("setting OAuth callback port", err) + } if len(scopesList) > 0 { opts = append(opts, u2m.WithScopes(scopesList)) } diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index a8eafb4be43..3a5791f7edb 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -10,6 +10,7 @@ import ( "net/http/httptest" "os" "path/filepath" + "reflect" "sync" "testing" "time" @@ -80,12 +81,13 @@ func (f *fakeDiscoveryPersistentAuth) Close() error { } type fakeDiscoveryClient struct { - oauthArg *u2m.BasicDiscoveryOAuthArgument - oauthArgErr error - persistentAuth discoveryPersistentAuth - persistentAuthErr error - introspection *auth.IntrospectionResult - introspectionErr error + oauthArg *u2m.BasicDiscoveryOAuthArgument + oauthArgErr error + persistentAuth discoveryPersistentAuth + persistentAuthErr error + persistentAuthOpts []u2m.PersistentAuthOption + introspection *auth.IntrospectionResult + introspectionErr error // For assertions introspectHost string introspectToken string @@ -99,12 +101,42 @@ func (f *fakeDiscoveryClient) NewOAuthArgument(profileName string) (*u2m.BasicDi } func (f *fakeDiscoveryClient) NewPersistentAuth(ctx context.Context, opts ...u2m.PersistentAuthOption) (discoveryPersistentAuth, error) { + f.persistentAuthOpts = opts if f.persistentAuthErr != nil { return nil, f.persistentAuthErr } return f.persistentAuth, nil } +func TestDiscoveryLogin_UsesOAuthCallbackPort(t *testing.T) { + t.Setenv("DATABRICKS_OAUTH_CALLBACK_PORT", "8030") + oauthArg, err := u2m.NewBasicDiscoveryOAuthArgument("DISCOVERY") + require.NoError(t, err) + dc := &fakeDiscoveryClient{ + oauthArg: oauthArg, + persistentAuth: &fakeDiscoveryPersistentAuth{ + challengeErr: errors.New("stop after creating persistent auth"), + }, + } + + ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + browserFunc: func(string) error { return nil }, + tokenStore: newTestStore(), + }) + require.Error(t, err) + + persistentAuth := &u2m.PersistentAuth{} + for _, opt := range dc.persistentAuthOpts { + opt(persistentAuth) + } + port := reflect.ValueOf(persistentAuth).Elem().FieldByName("port").Int() + assert.EqualValues(t, 8030, port) +} + func (f *fakeDiscoveryClient) IntrospectToken(ctx context.Context, host, accessToken string) (*auth.IntrospectionResult, error) { f.introspectHost = host f.introspectToken = accessToken