Apply auth token --timeout to the inline login

cmd/auth/token.go

@@ -50,7 +50,7 @@ and secret is not supported.`,
 
 	var tokenTimeout time.Duration
 	cmd.Flags().DurationVar(&tokenTimeout, "timeout", defaultTimeout,
-		"Timeout for acquiring a token.")
+		"Timeout for acquiring a token, including any browser login.")
 
 	var forceRefresh bool
 	cmd.Flags().BoolVar(&forceRefresh, "force-refresh", false,
@@ -199,7 +199,7 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) {
 	// resolve the target through environment variables or interactive profile selection.
 	if args.profileName == "" && args.authArguments.Host == "" && len(args.args) == 0 {
 		var resolvedProfile string
-		resolvedProfile, existingProfile, err = resolveNoArgsToken(ctx, args.profiler, args.authArguments, args.tokenStore, args.mode)
+		resolvedProfile, existingProfile, err = resolveNoArgsToken(ctx, args.profiler, args.authArguments, args.tokenStore, args.mode, args.tokenTimeout)
 		if err != nil {
 			return nil, err
 		}
@@ -347,7 +347,7 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) {
 //
 // Returns the resolved profile name and profile (if any). The host and related
 // fields on authArgs are updated in place when resolved via environment variables.
-func resolveNoArgsToken(ctx context.Context, profiler profile.Profiler, authArgs *auth.AuthArguments, tokenStore storage.Store, mode storage.StorageMode) (string, *profile.Profile, error) {
+func resolveNoArgsToken(ctx context.Context, profiler profile.Profiler, authArgs *auth.AuthArguments, tokenStore storage.Store, mode storage.StorageMode, timeout time.Duration) (string, *profile.Profile, error) {
 	// Step 1: Try DATABRICKS_HOST env var (highest priority).
 	if envHost := env.Get(ctx, "DATABRICKS_HOST"); envHost != "" {
 		authArgs.Host = envHost
@@ -410,7 +410,7 @@ func resolveNoArgsToken(ctx context.Context, profiler profile.Profiler, authArgs
 		// Fall through — setHostAndAccountId will prompt for the host.
 		return "", nil, nil
 	case profilePickerCreateNew:
-		return runInlineLogin(ctx, profiler, tokenStore, mode)
+		return runInlineLogin(ctx, profiler, tokenStore, mode, timeout)
 	default:
 		p, err := loadProfileByName(ctx, selectedName, profiler)
 		if err != nil {
@@ -425,7 +425,7 @@ func resolveNoArgsToken(ctx context.Context, profiler profile.Profiler, authArgs
 // .databrickscfg, and returns the new profile name and profile. It rejects
 // the names of existing sections so that it cannot log in to an existing
 // profile.
-func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenStore storage.Store, mode storage.StorageMode) (string, *profile.Profile, error) {
+func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenStore storage.Store, mode storage.StorageMode, timeout time.Duration) (string, *profile.Profile, error) {
 	takenNames, err := databrickscfg.GetSectionNames(ctx, env.Get(ctx, "DATABRICKS_CONFIG_FILE"))
 	if err != nil {
 		return "", nil, err
@@ -455,7 +455,7 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenStore s
 
 	// Challenge uses the context passed to NewPersistentAuth, so the timeout
 	// must be set before it.
-	ctx, cancel := context.WithTimeout(ctx, defaultTimeout)
+	ctx, cancel := context.WithTimeout(ctx, timeout)
 	defer cancel()
 
 	persistentAuth, err := u2m.NewPersistentAuth(ctx, persistentAuthOpts...)

cmd/auth/token_test.go

@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"errors"
 	"net/http"
+	"net/http/httptest"
 	"net/url"
 	"runtime"
 	"strings"
@@ -984,7 +985,7 @@ host = https://prod.cloud.databricks.test
 cluster_id = abc
 `))
 	tm := termtest.Run(t, func(ctx context.Context) (string, error) {
-		name, _, err := resolveNoArgsToken(ctx, profile.DefaultProfiler, &auth.AuthArguments{}, nil, storage.StorageModePlaintext)
+		name, _, err := resolveNoArgsToken(ctx, profile.DefaultProfiler, &auth.AuthArguments{}, nil, storage.StorageModePlaintext, defaultTimeout)
 		return name, err
 	})
 	tm.WaitFor(profilePickerCreateNewLabel)
@@ -998,3 +999,47 @@ cluster_id = abc
 	_, err := tm.Result()
 	assert.EqualError(t, err, "^C")
 }
+
+func TestLoadTokenCreateNewProfileAppliesTimeout(t *testing.T) {
+	// See TestLoginDeviceCodePickerEmptyHost.
+	if runtime.GOOS == "windows" {
+		t.Skip("bubbletea cannot cancel reads on non-console input on Windows")
+	}
+	// Either one would resolve the profile before the picker is shown.
+	t.Setenv("DATABRICKS_HOST", "")
+	t.Setenv("DATABRICKS_CONFIG_PROFILE", "")
+	t.Setenv("DATABRICKS_CONFIG_FILE", writeTempConfig(t, `[prod]
+host = https://prod.cloud.databricks.test
+`))
+	t.Setenv("BROWSER", "none")
+
+	// Host discovery runs before the login timeout starts and has its own 5s
+	// limit, so answer it immediately. Every other request blocks until the
+	// login's context is done.
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "/.well-known/databricks-config" {
+			w.WriteHeader(http.StatusNotFound)
+			return
+		}
+		<-r.Context().Done()
+	}))
+	t.Cleanup(server.Close)
+
+	tm := termtest.Run(t, func(ctx context.Context) (*oauth2.Token, error) {
+		return loadToken(ctx, loadTokenArgs{
+			authArguments: &auth.AuthArguments{},
+			tokenTimeout:  100 * time.Millisecond,
+			profiler:      profile.DefaultProfiler,
+			tokenStore:    storage.NewMemoryStore(),
+			mode:          storage.StorageModePlaintext,
+		})
+	})
+	tm.WaitFor(profilePickerCreateNewLabel)
+	tm.Type(termtest.KeyDown + termtest.KeyEnter)
+	tm.WaitFor("Databricks profile name [DEFAULT]")
+	tm.Type("new" + termtest.KeyEnter)
+	tm.WaitFor("Databricks host")
+	tm.Type(server.URL + termtest.KeyEnter)
+	_, err := tm.Result()
+	assert.ErrorIs(t, err, context.DeadlineExceeded)
+}