Allow the use of different client IDs in U2M auth

.nextchanges/cli/u2m-client-id.md

@@ -0,0 +1 @@
+* Allow OAuth U2M logins to override the CLI client ID with `--client-id`, profile `client_id`, or `DATABRICKS_CLIENT_ID`. ([#6594](https://github.com/databricks/cli/pull/6594))

acceptance/bin/discovery_browser.py

@@ -31,6 +31,11 @@
 dest_parsed = urllib.parse.urlparse(destination_url)
 dest_params = urllib.parse.parse_qs(dest_parsed.query)
 
+expected_client_id = os.environ.get("DATABRICKS_TEST_CLIENT_ID")
+if expected_client_id and dest_params.get("client_id") != [expected_client_id]:
+    sys.stderr.write(f"Expected client_id {expected_client_id!r}, got {dest_params.get('client_id')!r}\n")
+    sys.exit(1)
+
 redirect_uri = dest_params.get("redirect_uri", [None])[0]
 state = dest_params.get("state", [None])[0]
 

acceptance/cmd/auth/login/discovery/out.databrickscfg

@@ -5,6 +5,7 @@
 host         = [DATABRICKS_URL]
 account_id   = test-account-123
 workspace_id = [NUMID]
+client_id    = discovery-client-id
 auth_type    = databricks-cli
 
 [__settings__]

acceptance/cmd/auth/login/discovery/output.txt

@@ -1,5 +1,5 @@
 
->>> [CLI] auth login --profile discovery-test
+>>> [CLI] auth login --profile discovery-test --client-id discovery-client-id
 Opening login.databricks.com in your browser...
 Profile discovery-test was successfully saved
 

acceptance/cmd/auth/login/discovery/script

@@ -2,8 +2,9 @@ sethome "./home"
 
 # Use the discovery browser script that simulates login.databricks.com
 export BROWSER="discovery_browser.py"
+export DATABRICKS_TEST_CLIENT_ID="discovery-client-id"
 
-trace $CLI auth login --profile discovery-test
+trace $CLI auth login --profile discovery-test --client-id discovery-client-id
 
 trace $CLI auth profiles
 

acceptance/cmd/auth/login/host-from-profile/home/.databrickscfg.tmpl

@@ -1,2 +1,4 @@
 [existing-profile]
 host = $DATABRICKS_HOST
+auth_type = databricks-cli
+client_id = custom-client-id

acceptance/cmd/auth/login/host-from-profile/out.databrickscfg

@@ -3,5 +3,6 @@
 
 [existing-profile]
 host         = [DATABRICKS_URL]
-workspace_id = [NUMID]
 auth_type    = databricks-cli
+workspace_id = [NUMID]
+client_id    = flag-client-id

acceptance/cmd/auth/login/host-from-profile/output.txt

@@ -2,16 +2,34 @@
 === Initial profile
 [existing-profile]
 host = [DATABRICKS_URL]
+auth_type = databricks-cli
+client_id = custom-client-id
 
 === Login with existing profile (no host argument)
 >>> [CLI] auth login --profile existing-profile
 Profile existing-profile was successfully saved
 
+=== OAuth client ID
+
+>>> print_requests.py //oidc/v1/authorize --get
+json.q.client_id = "custom-client-id";
+
+=== Login with a client ID override
+
+>>> [CLI] auth login --profile existing-profile --client-id flag-client-id
+Profile existing-profile was successfully saved
+
+=== Overridden OAuth client ID
+
+>>> print_requests.py //oidc/v1/authorize --get
+json.q.client_id = "flag-client-id";
+
 === Profile after login
 ; The profile defined in the DEFAULT section is to be used as a fallback when no profile is explicitly specified.
 [DEFAULT]
 
 [existing-profile]
 host         = [DATABRICKS_URL]
-workspace_id = [NUMID]
 auth_type    = databricks-cli
+workspace_id = [NUMID]
+client_id    = flag-client-id

acceptance/cmd/auth/login/host-from-profile/script

@@ -14,6 +14,15 @@ export BROWSER="browser.py"
 title "Login with existing profile (no host argument)"
 trace $CLI auth login --profile existing-profile
 
+title "OAuth client ID\n"
+trace print_requests.py //oidc/v1/authorize --get | gron.py | grep client_id
+
+title "Login with a client ID override\n"
+trace $CLI auth login --profile existing-profile --client-id flag-client-id
+
+title "Overridden OAuth client ID\n"
+trace print_requests.py //oidc/v1/authorize --get | gron.py | grep client_id
+
 title "Profile after login\n"
 cat "./home/.databrickscfg"
 

acceptance/cmd/auth/login/host-from-profile/test.toml

@@ -1,3 +1,5 @@
 Ignore = [
     "home"
 ]
+RecordRequests = true
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]

cmd/auth/login.go

@@ -132,6 +132,7 @@ a new profile is created.
 	var configureServerless bool
 	var skipWorkspace bool
 	var scopes string
+	var clientID string
 	cmd.Flags().DurationVar(&loginTimeout, "timeout", defaultTimeout,
 		"Timeout for completing login challenge in the browser")
 	cmd.Flags().BoolVar(&configureCluster, "configure-cluster", false,
@@ -142,6 +143,8 @@ a new profile is created.
 		"Skip workspace selection for account-level access")
 	cmd.Flags().StringVar(&scopes, "scopes", "",
 		"Comma-separated list of OAuth scopes to request (defaults to 'all-apis')")
+	cmd.Flags().StringVar(&clientID, "client-id", "",
+		"OAuth client ID to use for U2M authentication")
 
 	cmd.PreRunE = profileHostConflictCheck
 
@@ -256,6 +259,9 @@ a new profile is created.
 		if err != nil {
 			return err
 		}
+		if clientID == "" {
+			clientID = u2mClientIDFromProfile(existingProfile)
+		}
 
 		// If no host is available from any source, use the discovery flow
 		// via login.databricks.com.
@@ -268,6 +274,7 @@ a new profile is created.
 				profileName:     profileName,
 				timeout:         loginTimeout,
 				scopes:          scopes,
+				clientID:        clientID,
 				existingProfile: existingProfile,
 				browserFunc:     getBrowserFunc(cmd),
 				tokenStore:      tokenStore,
@@ -302,6 +309,9 @@ a new profile is created.
 			u2m.WithBrowser(getBrowserFunc(cmd)),
 			u2m.WithTokenCache(storage.WrapForOAuthArgument(ctx, tokenStore, mode, oauthArgument)),
 		}
+		if clientID != "" {
+			persistentAuthOpts = append(persistentAuthOpts, u2m.WithClientID(clientID))
+		}
 		if len(scopesList) > 0 {
 			persistentAuthOpts = append(persistentAuthOpts, u2m.WithScopes(scopesList))
 		}
@@ -393,6 +403,7 @@ a new profile is created.
 				ConfigFile:          env.Get(ctx, "DATABRICKS_CONFIG_FILE"),
 				ServerlessComputeID: serverlessComputeID,
 				Scopes:              scopesList,
+				ClientID:            clientID,
 			}, clearKeys...)
 			if err != nil {
 				return err
@@ -634,6 +645,7 @@ type discoveryLoginInputs struct {
 	profileName     string
 	timeout         time.Duration
 	scopes          string
+	clientID        string
 	existingProfile *profile.Profile
 	browserFunc     func(string) error
 	tokenStore      storage.Store
@@ -660,6 +672,9 @@ func discoveryLogin(ctx context.Context, in discoveryLoginInputs) error {
 		u2m.WithDiscoveryLogin(),
 		u2m.WithTokenCache(storage.WrapForOAuthArgument(ctx, in.tokenStore, in.mode, arg)),
 	}
+	if in.clientID != "" {
+		opts = append(opts, u2m.WithClientID(in.clientID))
+	}
 	if len(scopesList) > 0 {
 		opts = append(opts, u2m.WithScopes(scopesList))
 	}
@@ -750,6 +765,7 @@ func discoveryLogin(ctx context.Context, in discoveryLoginInputs) error {
 		WorkspaceID: workspaceID,
 		Scopes:      scopesList,
 		ConfigFile:  configFile,
+		ClientID:    in.clientID,
 	}, clearKeys...)
 	if err != nil {
 		if configFile != "" {
@@ -785,6 +801,14 @@ func oauthLoginClearKeys() []string {
 	return databrickscfg.AuthCredentialKeys()
 }
 
+// u2mClientIDFromProfile excludes client IDs belonging to other auth types.
+func u2mClientIDFromProfile(p *profile.Profile) string {
+	if p == nil || p.AuthType != authTypeDatabricksCLI {
+		return ""
+	}
+	return p.ClientID
+}
+
 // promptForWorkspaceSelection lists workspaces for a SPOG account and lets the
 // user pick one. Returns the selected workspace ID or empty string if skipped.
 // This is best-effort: errors are returned to the caller for logging, not shown

cmd/auth/login_test.go

@@ -458,6 +458,35 @@ func TestSplitScopes(t *testing.T) {
 	}
 }
 
+func TestU2MClientIDFromProfile(t *testing.T) {
+	tests := []struct {
+		name    string
+		profile *profile.Profile
+		want    string
+	}{
+		{name: "no profile"},
+		{
+			name:    "implicit auth type",
+			profile: &profile.Profile{ClientID: "custom-client-id"},
+		},
+		{
+			name:    "M2M auth type",
+			profile: &profile.Profile{AuthType: "oauth-m2m", ClientID: "custom-client-id"},
+		},
+		{
+			name:    "U2M auth type",
+			profile: &profile.Profile{AuthType: authTypeDatabricksCLI, ClientID: "custom-client-id"},
+			want:    "custom-client-id",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equal(t, tt.want, u2mClientIDFromProfile(tt.profile))
+		})
+	}
+}
+
 func TestRunHostDiscovery_NoHost(t *testing.T) {
 	ctx := t.Context()
 	args := &auth.AuthArguments{}
@@ -952,9 +981,11 @@ func TestDiscoveryLogin_ReloginPreservesExistingProfileScopes(t *testing.T) {
 	}
 
 	existingProfile := &profile.Profile{
-		Name:   "DISCOVERY",
-		Host:   "https://old-workspace.example.com",
-		Scopes: "sql,clusters",
+		Name:     "DISCOVERY",
+		Host:     "https://old-workspace.example.com",
+		Scopes:   "sql,clusters",
+		AuthType: authTypeDatabricksCLI,
+		ClientID: "custom-client-id",
 	}
 
 	// No --scopes flag (empty string), should fall back to existing profile scopes.
@@ -963,6 +994,7 @@ func TestDiscoveryLogin_ReloginPreservesExistingProfileScopes(t *testing.T) {
 		dc:              dc,
 		profileName:     "DISCOVERY",
 		timeout:         time.Second,
+		clientID:        existingProfile.ClientID,
 		existingProfile: existingProfile,
 		browserFunc:     func(string) error { return nil },
 		tokenStore:      newTestStore(),
@@ -974,6 +1006,51 @@ func TestDiscoveryLogin_ReloginPreservesExistingProfileScopes(t *testing.T) {
 	require.NotNil(t, savedProfile)
 	assert.Equal(t, "https://workspace.example.com", savedProfile.Host)
 	assert.Equal(t, "sql,clusters", savedProfile.Scopes)
+	assert.Equal(t, "custom-client-id", savedProfile.ClientID)
+}
+
+func TestDiscoveryLogin_ExplicitClientIDOverridesExistingProfile(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"),
+	}
+
+	existingProfile := &profile.Profile{
+		Name:     "DISCOVERY",
+		Host:     "https://old-workspace.example.com",
+		AuthType: authTypeDatabricksCLI,
+		ClientID: "profile-client-id",
+	}
+
+	ctx, _ := cmdio.NewTestContextWithStdout(t.Context())
+	err = discoveryLogin(ctx, discoveryLoginInputs{
+		dc:              dc,
+		profileName:     "DISCOVERY",
+		timeout:         time.Second,
+		clientID:        "flag-client-id",
+		existingProfile: existingProfile,
+		browserFunc:     func(string) error { return nil },
+		tokenStore:      newTestStore(),
+	})
+	require.NoError(t, err)
+
+	savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler)
+	require.NoError(t, err)
+	require.NotNil(t, savedProfile)
+	assert.Equal(t, "flag-client-id", savedProfile.ClientID)
 }
 
 func TestDiscoveryLogin_ExplicitScopesOverrideExistingProfile(t *testing.T) {

cmd/auth/token.go

@@ -265,6 +265,9 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) {
 		return nil, err
 	}
 	allArgs := append([]u2m.PersistentAuthOption{u2m.WithTokenCache(storage.OAuthTokenCache(ctx, args.tokenStore, args.mode))}, args.persistentAuthOpts...)
+	if clientID := u2mClientIDFromProfile(existingProfile); clientID != "" {
+		allArgs = append(allArgs, u2m.WithClientID(clientID))
+	}
 	allArgs = append(allArgs, u2m.WithOAuthArgument(oauthArgument))
 	persistentAuth, err := u2m.NewPersistentAuth(ctx, allArgs...)
 	if err != nil {
@@ -430,6 +433,9 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenStore s
 		u2m.WithBrowser(func(url string) error { return browser.Open(ctx, url) }),
 		u2m.WithTokenCache(storage.WrapForOAuthArgument(ctx, tokenStore, mode, oauthArgument)),
 	}
+	if clientID := u2mClientIDFromProfile(existingProfile); clientID != "" {
+		persistentAuthOpts = append(persistentAuthOpts, u2m.WithClientID(clientID))
+	}
 	if len(scopesList) > 0 {
 		persistentAuthOpts = append(persistentAuthOpts, u2m.WithScopes(scopesList))
 	}
@@ -458,6 +464,7 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenStore s
 		WorkspaceID: loginArgs.WorkspaceID,
 		ConfigFile:  env.Get(ctx, "DATABRICKS_CONFIG_FILE"),
 		Scopes:      scopesList,
+		ClientID:    u2mClientIDFromProfile(existingProfile),
 	}, clearKeys...)
 	if err != nil {
 		return "", nil, err