auth: add DATABRICKS_DISCOVERY_HOST env var to override login.databricks.com

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 <url>"
+	// 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)
 	}

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"}