diff --git a/cmd/auth/login.go b/cmd/auth/login.go index d45a77806fd..a7ac42f28c4 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -46,6 +46,11 @@ const ( defaultTimeout = 1 * time.Hour authTypeDatabricksCLI = "databricks-cli" discoveryFallbackTip = "\n\nTip: you can specify a workspace directly with: databricks auth login --host " + // discoveryHostEnvVar overrides the default https://login.databricks.com + // 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" ) // discoveryErr wraps an error (or creates a new one) and appends the @@ -618,6 +623,10 @@ func discoveryLogin(ctx context.Context, in discoveryLoginInputs) error { if len(scopesList) > 0 { opts = append(opts, u2m.WithScopes(scopesList)) } + discoveryHost := env.Get(ctx, discoveryHostEnvVar) + if discoveryHost != "" { + opts = append(opts, u2m.WithDiscoveryHost(discoveryHost)) + } // Apply timeout before creating PersistentAuth so Challenge() respects it. ctx, cancel := context.WithTimeout(ctx, in.timeout) @@ -629,7 +638,11 @@ func discoveryLogin(ctx context.Context, in discoveryLoginInputs) error { } defer persistentAuth.Close() - cmdio.LogString(ctx, "Opening login.databricks.com in your browser...") + displayHost := "login.databricks.com" + if discoveryHost != "" { + displayHost = discoveryHost + } + cmdio.LogString(ctx, fmt.Sprintf("Opening %s in your browser...", displayHost)) if err := persistentAuth.Challenge(); err != nil { return discoveryErr("login via login.databricks.com failed", err) } diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index d209c511154..86820881cf7 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -1082,6 +1082,40 @@ auth_type = databricks-cli assert.Equal(t, "222222", savedProfile.WorkspaceID, "workspace_id should be updated to fresh introspection value") } +func TestDiscoveryLogin_OverridesHostFromEnv(t *testing.T) { + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, ".databrickscfg") + err := os.WriteFile(configPath, []byte(""), 0o600) + require.NoError(t, err) + t.Setenv("DATABRICKS_CONFIG_FILE", configPath) + + oauthArg, err := u2m.NewBasicDiscoveryOAuthArgument("DISCOVERY") + require.NoError(t, err) + oauthArg.SetDiscoveredHost("https://workspace.example.com") + + dc := &fakeDiscoveryClient{ + oauthArg: oauthArg, + persistentAuth: &fakeDiscoveryPersistentAuth{ + token: &oauth2.Token{AccessToken: "test-token"}, + }, + introspectionErr: errors.New("introspection failed"), + } + + ctx, stderr := cmdio.NewTestContextWithStderr(t.Context()) + ctx = env.Set(ctx, "DATABRICKS_DISCOVERY_HOST", "https://login.staging.test") + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + browserFunc: func(string) error { return nil }, + tokenCache: newTestTokenCache(), + }) + require.NoError(t, err) + + assert.Contains(t, stderr.String(), "Opening https://login.staging.test in your browser...") + assert.NotContains(t, stderr.String(), "Opening login.databricks.com in your browser...") +} + func TestLoginRejectsPositionalArgWithHostFlag(t *testing.T) { ctx := cmdio.MockDiscard(t.Context()) authArgs := &auth.AuthArguments{Host: "https://example.com"}