Add host disambiguation and positional profile support to auth token

cmd/auth/token.go

@@ -5,12 +5,16 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"strings"
 	"time"
 
 	"github.com/databricks/cli/libs/auth"
+	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/databrickscfg/profile"
+	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/databricks/databricks-sdk-go/credentials/u2m"
 	"github.com/databricks/databricks-sdk-go/credentials/u2m/cache"
+	"github.com/manifoldco/promptui"
 	"github.com/spf13/cobra"
 	"golang.org/x/oauth2"
 )
@@ -22,7 +26,7 @@ func helpfulError(ctx context.Context, profile string, persistentAuth u2m.OAuthA
 
 func newTokenCommand(authArguments *auth.AuthArguments) *cobra.Command {
 	cmd := &cobra.Command{
- Use: "token [HOST]",
+ Use: "token [HOST_OR_PROFILE]",
  Short: "Get authentication token",
  Long: `Get authentication token from the local cache in ~/.databricks/token-cache.json.
 Refresh the access token if it is expired. Note: This command only works with
@@ -93,6 +97,19 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) {
  return nil, errors.New("providing both a profile and host is not supported")
 	}
 
+	// If no --profile flag, try resolving the positional arg as a profile name.
+	// If it matches, use it. If not, fall through to host treatment.
+	if args.profileName == "" && len(args.args) == 1 {
+ candidateProfile, err := loadProfileByName(ctx, args.args[0], args.profiler)
+ if err != nil {
+ return nil, err
+ }
+ if candidateProfile != nil {
+ args.profileName = args.args[0]
+ args.args = nil
+ }
+	}
+
 	existingProfile, err := loadProfileByName(ctx, args.profileName, args.profiler)
 	if err != nil {
  return nil, err
@@ -113,6 +130,47 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) {
  return nil, err
 	}
 
+	// When no profile was specified, check if multiple profiles match the
+	// effective cache key for this host.
+	if args.profileName == "" && args.authArguments.Host != "" {
+ cfg := &config.Config{
+ Host: args.authArguments.Host,
+ AccountID: args.authArguments.AccountID,
+ Experimental_IsUnifiedHost: args.authArguments.IsUnifiedHost,
+ }
+ // Canonicalize first so HostType() can correctly identify account hosts
+ // even when the host string lacks a scheme (e.g. "accounts.cloud.databricks.com").
+ cfg.CanonicalHostName()
+ var matchFn profile.ProfileMatchFunction
+ switch cfg.HostType() {
+ case config.AccountHost, config.UnifiedHost:
+ matchFn = profile.WithHostAndAccountID(args.authArguments.Host, args.authArguments.AccountID)
+ default:
+ matchFn = profile.WithHost(args.authArguments.Host)
+ }
+
+ matchingProfiles, err := args.profiler.LoadProfiles(ctx, matchFn)
+ if err != nil && !errors.Is(err, profile.ErrNoConfiguration) {
+ return nil, err
+ }
+ if len(matchingProfiles) > 1 {
+ configPath, _ := args.profiler.GetPath(ctx)
+ if configPath == "" {
+ panic("configPath is empty but LoadProfiles returned multiple profiles")
+ }
+ if !cmdio.IsPromptSupported(ctx) {
+ names := strings.Join(matchingProfiles.Names(), " and ")
+ return nil, fmt.Errorf("%s match %s in %s. Use --profile to specify which profile to use",
+ names, args.authArguments.Host, configPath)
+ }
+ selected, err := askForMatchingProfile(ctx, matchingProfiles, args.authArguments.Host)
+ if err != nil {
+ return nil, err
+ }
+ args.profileName = selected
+ }
+	}
+
 	args.authArguments.Profile = args.profileName
 
 	ctx, cancel := context.WithTimeout(ctx, args.tokenTimeout)
@@ -149,3 +207,22 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) {
 	}
 	return t, nil
 }
+
+func askForMatchingProfile(ctx context.Context, profiles profile.Profiles, host string) (string, error) {
+	i, _, err := cmdio.RunSelect(ctx, &promptui.Select{
+ Label: "Multiple profiles match " + host,
+ Items: profiles,
+ Searcher: profiles.SearchCaseInsensitive,
+ StartInSearchMode: true,
+ Templates: &promptui.SelectTemplates{
+ Label: "{{ . | faint }}",
+ Active: `{{.Name | bold}} ({{.Host|faint}})`,
+ Inactive: `{{.Name}}`,
+ Selected: `{{ "Using profile" | faint }}: {{ .Name | bold }}`,
+ },
+	})
+	if err != nil {
+ return "", err
+	}
+	return profiles[i].Name, nil
+}

cmd/auth/token_test.go

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/databricks/cli/libs/auth"
+	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/databrickscfg/profile"
 	"github.com/databricks/databricks-sdk-go/credentials/u2m"
 	"github.com/databricks/databricks-sdk-go/httpclient/fixtures"
@@ -89,6 +90,32 @@ func TestToken_loadToken(t *testing.T) {
  Host: "https://accounts.cloud.databricks.com",
  AccountID: "active",
  },
+ {
+ Name: "workspace-a",
+ Host: "https://workspace-a.cloud.databricks.com",
+ },
+ {
+ Name: "dup1",
+ Host: "https://shared.cloud.databricks.com",
+ },
+ {
+ Name: "dup2",
+ Host: "https://shared.cloud.databricks.com",
+ },
+ {
+ Name: "acct-dup1",
+ Host: "https://accounts.cloud.databricks.com",
+ AccountID: "same-account",
+ },
+ {
+ Name: "acct-dup2",
+ Host: "https://accounts.cloud.databricks.com",
+ AccountID: "same-account",
+ },
+ {
+ Name: "default.dev",
+ Host: "https://dev.cloud.databricks.com",
+ },
  },
 	}
 	tokenCache := &inMemoryTokenCache{
@@ -107,6 +134,18 @@ func TestToken_loadToken(t *testing.T) {
  RefreshToken: "active",
  Expiry: time.Now().Add(1 * time.Hour),
  },
+ "workspace-a": {
+ RefreshToken: "workspace-a",
+ Expiry: time.Now().Add(1 * time.Hour),
+ },
+ "https://workspace-a.cloud.databricks.com": {
+ RefreshToken: "workspace-a",
+ Expiry: time.Now().Add(1 * time.Hour),
+ },
+ "default.dev": {
+ RefreshToken: "default.dev",
+ Expiry: time.Now().Add(1 * time.Hour),
+ },
  },
 	}
 	validateToken := func(resp *oauth2.Token) {
@@ -116,6 +155,7 @@ func TestToken_loadToken(t *testing.T) {
 
 	cases := []struct {
  name string
+ ctx context.Context
  args loadTokenArgs
  validateToken func(*oauth2.Token)
  wantErr string
@@ -223,10 +263,169 @@ func TestToken_loadToken(t *testing.T) {
  },
  validateToken: validateToken,
  },
+ {
+ name: "positional arg resolved as profile name",
+ args: loadTokenArgs{
+ authArguments: &auth.AuthArguments{},
+ profileName: "",
+ args: []string{"workspace-a"},
+ tokenTimeout: 1 * time.Hour,
+ profiler: profiler,
+ persistentAuthOpts: []u2m.PersistentAuthOption{
+ u2m.WithTokenCache(tokenCache),
+ u2m.WithOAuthEndpointSupplier(&MockApiClient{}),
+ u2m.WithHttpClient(&http.Client{Transport: fixtures.SliceTransport{refreshSuccessTokenResponse}}),
+ },
+ },
+ validateToken: validateToken,
+ },
+ {
+ name: "positional arg with dot treated as host when no profile matches",
+ args: loadTokenArgs{
+ authArguments: &auth.AuthArguments{},
+ profileName: "",
+ args: []string{"workspace-a.cloud.databricks.com"},
+ tokenTimeout: 1 * time.Hour,
+ profiler: profiler,
+ persistentAuthOpts: []u2m.PersistentAuthOption{
+ u2m.WithTokenCache(tokenCache),
+ u2m.WithOAuthEndpointSupplier(&MockApiClient{}),
+ u2m.WithHttpClient(&http.Client{Transport: fixtures.SliceTransport{refreshSuccessTokenResponse}}),
+ },
+ },
+ validateToken: validateToken,
+ },
+ {
+ name: "dotted profile name resolved as profile not host",
+ args: loadTokenArgs{
+ authArguments: &auth.AuthArguments{},
+ profileName: "",
+ args: []string{"default.dev"},
+ tokenTimeout: 1 * time.Hour,
+ profiler: profiler,
+ persistentAuthOpts: []u2m.PersistentAuthOption{
+ u2m.WithTokenCache(tokenCache),
+ u2m.WithOAuthEndpointSupplier(&MockApiClient{}),
+ u2m.WithHttpClient(&http.Client{Transport: fixtures.SliceTransport{refreshSuccessTokenResponse}}),
+ },
+ },
+ validateToken: validateToken,
+ },
+ {
+ name: "positional arg not a profile falls through to host",
+ args: loadTokenArgs{
+ authArguments: &auth.AuthArguments{},
+ profileName: "",
+ args: []string{"nonexistent"},
+ tokenTimeout: 1 * time.Hour,
+ profiler: profiler,
+ persistentAuthOpts: []u2m.PersistentAuthOption{
+ u2m.WithTokenCache(tokenCache),
+ u2m.WithOAuthEndpointSupplier(&MockApiClient{}),
+ },
+ },
+ wantErr: "cache: databricks OAuth is not configured for this host. " +
+ "Try logging in again with `databricks auth login --host https://nonexistent` before retrying. " +
+ "If this fails, please report this issue to the Databricks CLI maintainers at https://github.com/databricks/cli/issues/new",
+ },
+ {
+ name: "scheme-less account host ambiguity detected correctly",
+ ctx: cmdio.MockDiscard(context.Background()),
+ args: loadTokenArgs{
+ authArguments: &auth.AuthArguments{
+ Host: "accounts.cloud.databricks.com",
+ AccountID: "same-account",
+ },
+ profileName: "",
+ args: []string{},
+ tokenTimeout: 1 * time.Hour,
+ profiler: profiler,
+ persistentAuthOpts: []u2m.PersistentAuthOption{
+ u2m.WithTokenCache(tokenCache),
+ u2m.WithOAuthEndpointSupplier(&MockApiClient{}),
+ },
+ },
+ wantErr: "acct-dup1 and acct-dup2 match accounts.cloud.databricks.com in <in memory>. Use --profile to specify which profile to use",
+ },
+ {
+ name: "workspace host ambiguity — multiple profiles, non-interactive",
+ ctx: cmdio.MockDiscard(context.Background()),
+ args: loadTokenArgs{
+ authArguments: &auth.AuthArguments{
+ Host: "https://shared.cloud.databricks.com",
+ },
+ profileName: "",
+ args: []string{},
+ tokenTimeout: 1 * time.Hour,
+ profiler: profiler,
+ persistentAuthOpts: []u2m.PersistentAuthOption{
+ u2m.WithTokenCache(tokenCache),
+ u2m.WithOAuthEndpointSupplier(&MockApiClient{}),
+ },
+ },
+ wantErr: "dup1 and dup2 match https://shared.cloud.databricks.com in <in memory>. Use --profile to specify which profile to use",
+ },
+ {
+ name: "account host — same host, different account IDs — no ambiguity",
+ args: loadTokenArgs{
+ authArguments: &auth.AuthArguments{
+ Host: "https://accounts.cloud.databricks.com",
+ AccountID: "active",
+ },
+ profileName: "",
+ args: []string{},
+ tokenTimeout: 1 * time.Hour,
+ profiler: profiler,
+ persistentAuthOpts: []u2m.PersistentAuthOption{
+ u2m.WithTokenCache(tokenCache),
+ u2m.WithOAuthEndpointSupplier(&MockApiClient{}),
+ u2m.WithHttpClient(&http.Client{Transport: fixtures.SliceTransport{refreshSuccessTokenResponse}}),
+ },
+ },
+ validateToken: validateToken,
+ },
+ {
+ name: "account host — same host AND same account ID — ambiguity",
+ ctx: cmdio.MockDiscard(context.Background()),
+ args: loadTokenArgs{
+ authArguments: &auth.AuthArguments{
+ Host: "https://accounts.cloud.databricks.com",
+ AccountID: "same-account",
+ },
+ profileName: "",
+ args: []string{},
+ tokenTimeout: 1 * time.Hour,
+ profiler: profiler,
+ persistentAuthOpts: []u2m.PersistentAuthOption{
+ u2m.WithTokenCache(tokenCache),
+ u2m.WithOAuthEndpointSupplier(&MockApiClient{}),
+ },
+ },
+ wantErr: "acct-dup1 and acct-dup2 match https://accounts.cloud.databricks.com in <in memory>. Use --profile to specify which profile to use",
+ },
+ {
+ name: "profile flag + positional non-host arg still errors",
+ args: loadTokenArgs{
+ authArguments: &auth.AuthArguments{},
+ profileName: "active",
+ args: []string{"workspace-a"},
+ tokenTimeout: 1 * time.Hour,
+ profiler: profiler,
+ persistentAuthOpts: []u2m.PersistentAuthOption{
+ u2m.WithTokenCache(tokenCache),
+ u2m.WithOAuthEndpointSupplier(&MockApiClient{}),
+ },
+ },
+ wantErr: "providing both a profile and host is not supported",
+ },
 	}
 	for _, c := range cases {
  t.Run(c.name, func(t *testing.T) {
- got, err := loadToken(context.Background(), c.args)
+ ctx := c.ctx
+ if ctx == nil {
+ ctx = context.Background()
+ }
+ got, err := loadToken(ctx, c.args)
  if c.wantErr != "" {
  assert.Equal(t, c.wantErr, err.Error())
  } else {

libs/databrickscfg/profile/profiler.go

@@ -2,6 +2,8 @@ package profile
 
 import (
 	"context"
+
+	"github.com/databricks/databricks-sdk-go/config"
 )
 
 type ProfileMatchFunction func(Profile) bool
@@ -30,6 +32,29 @@ func WithName(name string) ProfileMatchFunction {
 	}
 }
 
+// WithHost returns a ProfileMatchFunction that matches profiles whose
+// canonical host equals the given host.
+func WithHost(host string) ProfileMatchFunction {
+	target := canonicalizeHost(host)
+	return func(p Profile) bool {
+ return p.Host != "" && canonicalizeHost(p.Host) == target
+	}
+}
+
+// WithHostAndAccountID returns a ProfileMatchFunction that matches profiles
+// by both canonical host and account ID.
+func WithHostAndAccountID(host, accountID string) ProfileMatchFunction {
+	target := canonicalizeHost(host)
+	return func(p Profile) bool {
+ return p.Host != "" && canonicalizeHost(p.Host) == target && p.AccountID == accountID
+	}
+}
+
+// canonicalizeHost normalizes a host using the SDK's canonical host logic.
+func canonicalizeHost(host string) string {
+	return (&config.Config{Host: host}).CanonicalHostName()
+}
+
 type Profiler interface {
 	LoadProfiles(context.Context, ProfileMatchFunction) (Profiles, error)
 	GetPath(context.Context) (string, error)