From ba1857e86c5b71894d8a0f4656ee4c3fbbcf3d5e Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 17 Apr 2026 13:33:38 +0200 Subject: [PATCH 01/12] cmd/auth: add newAuthCache helper for storage-mode resolution Introduces a thin helper that resolves the CLI's token storage mode via libs/auth/storage.ResolveStorageMode and returns a cache.TokenCache. The helper is split into a public convenience form and an injectable core (newAuthCacheWith) so unit tests exercise the secure path with a fake cache without touching the real OS keyring. Follow-up commits wire this into the auth command call sites. --- cmd/auth/cache.go | 59 ++++++++++++++++++++++++ cmd/auth/cache_test.go | 100 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 159 insertions(+) create mode 100644 cmd/auth/cache.go create mode 100644 cmd/auth/cache_test.go diff --git a/cmd/auth/cache.go b/cmd/auth/cache.go new file mode 100644 index 0000000000..2fac817278 --- /dev/null +++ b/cmd/auth/cache.go @@ -0,0 +1,59 @@ +package auth + +import ( + "context" + "fmt" + + "github.com/databricks/cli/libs/auth/storage" + "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" +) + +// cacheFactories bundles the constructors newAuthCache depends on. Extracted +// so unit tests can inject stubs without hitting the real OS keyring or +// filesystem. Production code uses defaultCacheFactories(). +type cacheFactories struct { + newFile func() (cache.TokenCache, error) + newKeyring func() cache.TokenCache +} + +// defaultCacheFactories returns the production factory set. +func defaultCacheFactories() cacheFactories { + return cacheFactories{ + newFile: func() (cache.TokenCache, error) { return cache.NewFileTokenCache() }, + newKeyring: storage.NewKeyringCache, + } +} + +// newAuthCache resolves the storage mode for this invocation and returns +// the corresponding token cache plus the resolved mode (so callers can log +// or surface it). +// +// override is usually the command-level flag value (e.g. the result of +// --secure-storage). Pass "" when the command has no flag. +func newAuthCache(ctx context.Context, override storage.StorageMode) (cache.TokenCache, storage.StorageMode, error) { + return newAuthCacheWith(ctx, override, defaultCacheFactories()) +} + +// newAuthCacheWith is the pure form of newAuthCache. It takes the factory +// set as a parameter so tests can inject stubs. +func newAuthCacheWith(ctx context.Context, override storage.StorageMode, f cacheFactories) (cache.TokenCache, storage.StorageMode, error) { + mode, err := storage.ResolveStorageMode(ctx, override) + if err != nil { + return nil, "", err + } + switch mode { + case storage.StorageModeSecure: + return f.newKeyring(), mode, nil + case storage.StorageModeLegacy, storage.StorageModePlaintext: + // MS1 ships no dedicated plaintext implementation; the switch will + // be added in MS3. Until then the file cache is the safest existing + // behavior and the resolver still accepts the value. + c, err := f.newFile() + if err != nil { + return nil, "", fmt.Errorf("open file token cache: %w", err) + } + return c, mode, nil + default: + return nil, "", fmt.Errorf("unsupported storage mode %q", string(mode)) + } +} diff --git a/cmd/auth/cache_test.go b/cmd/auth/cache_test.go new file mode 100644 index 0000000000..b8d30d2f23 --- /dev/null +++ b/cmd/auth/cache_test.go @@ -0,0 +1,100 @@ +package auth + +import ( + "errors" + "testing" + + "github.com/databricks/cli/libs/auth/storage" + "github.com/databricks/cli/libs/env" + "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/oauth2" +) + +// stubCache is a test double for cache.TokenCache that records the source +// it was constructed from. It lets the tests confirm which factory ran. +type stubCache struct{ source string } + +func (stubCache) Store(string, *oauth2.Token) error { return nil } +func (stubCache) Lookup(string) (*oauth2.Token, error) { return nil, cache.ErrNotFound } + +func fakeFactories(t *testing.T) cacheFactories { + t.Helper() + return cacheFactories{ + newFile: func() (cache.TokenCache, error) { return stubCache{source: "file"}, nil }, + newKeyring: func() cache.TokenCache { return stubCache{source: "keyring"} }, + } +} + +func TestNewAuthCache_DefaultsToLegacyFile(t *testing.T) { + ctx := t.Context() + + got, mode, err := newAuthCacheWith(ctx, "", fakeFactories(t)) + + require.NoError(t, err) + assert.Equal(t, storage.StorageModeLegacy, mode) + assert.Equal(t, "file", got.(stubCache).source) +} + +func TestNewAuthCache_OverrideSecureUsesKeyring(t *testing.T) { + ctx := t.Context() + + got, mode, err := newAuthCacheWith(ctx, storage.StorageModeSecure, fakeFactories(t)) + + require.NoError(t, err) + assert.Equal(t, storage.StorageModeSecure, mode) + assert.Equal(t, "keyring", got.(stubCache).source) +} + +func TestNewAuthCache_EnvVarSelectsSecure(t *testing.T) { + ctx := env.Set(t.Context(), storage.EnvVar, "secure") + + got, mode, err := newAuthCacheWith(ctx, "", fakeFactories(t)) + + require.NoError(t, err) + assert.Equal(t, storage.StorageModeSecure, mode) + assert.Equal(t, "keyring", got.(stubCache).source) +} + +func TestNewAuthCache_PlaintextFallsBackToFile(t *testing.T) { + ctx := t.Context() + + got, mode, err := newAuthCacheWith(ctx, storage.StorageModePlaintext, fakeFactories(t)) + + require.NoError(t, err) + assert.Equal(t, storage.StorageModePlaintext, mode) + assert.Equal(t, "file", got.(stubCache).source) +} + +func TestNewAuthCache_InvalidOverrideReturnsError(t *testing.T) { + ctx := t.Context() + + _, _, err := newAuthCacheWith(ctx, storage.StorageMode("bogus"), fakeFactories(t)) + + require.Error(t, err) + assert.Contains(t, err.Error(), `unknown storage mode "bogus"`) +} + +func TestNewAuthCache_InvalidEnvReturnsError(t *testing.T) { + ctx := env.Set(t.Context(), storage.EnvVar, "bogus") + + _, _, err := newAuthCacheWith(ctx, "", fakeFactories(t)) + + require.Error(t, err) + assert.Contains(t, err.Error(), "DATABRICKS_AUTH_STORAGE") +} + +func TestNewAuthCache_FileFactoryErrorPropagates(t *testing.T) { + ctx := t.Context() + boom := errors.New("disk full") + factories := cacheFactories{ + newFile: func() (cache.TokenCache, error) { return nil, boom }, + newKeyring: func() cache.TokenCache { return stubCache{source: "keyring"} }, + } + + _, _, err := newAuthCacheWith(ctx, storage.StorageModeLegacy, factories) + + require.Error(t, err) + assert.ErrorIs(t, err, boom) +} From ca4441763a032a9e503e4375510862035a9723c4 Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 17 Apr 2026 13:37:07 +0200 Subject: [PATCH 02/12] cmd/auth: document why newFile is wrapped in a closure cache.NewFileTokenCache has a variadic signature (opts ...FileTokenCacheOption), which cannot be assigned directly to a non-variadic field. Adds a one-line explanation so future readers do not 'simplify' the closure back into a direct function reference and break the build. --- cmd/auth/cache.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/cmd/auth/cache.go b/cmd/auth/cache.go index 2fac817278..29022e4585 100644 --- a/cmd/auth/cache.go +++ b/cmd/auth/cache.go @@ -17,6 +17,9 @@ type cacheFactories struct { } // defaultCacheFactories returns the production factory set. +// newFile is wrapped in a closure because cache.NewFileTokenCache is variadic +// (func(...FileTokenCacheOption)) and cannot satisfy the non-variadic field type +// by direct reference. The closure calls it with no options (SDK defaults). func defaultCacheFactories() cacheFactories { return cacheFactories{ newFile: func() (cache.TokenCache, error) { return cache.NewFileTokenCache() }, From d419b451b02541fc53af933f456ab28c6368fcd8 Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 17 Apr 2026 13:42:21 +0200 Subject: [PATCH 03/12] cmd/auth: wire secure-storage cache into auth login Adds a hidden --secure-storage BoolVar to auth login and routes the resolved token cache into both NewPersistentAuth call sites (main flow and discoveryLogin). Users opt into OS-native secure storage either per invocation via the flag or for the whole shell via DATABRICKS_AUTH_STORAGE=secure. The default storage mode remains legacy. Flag is MarkHidden because MS1 is experimental; release notes document the env var for discovery. --- cmd/auth/login.go | 49 ++++++++++++++++++++++++++++++++++-------- cmd/auth/login_test.go | 20 ++++++++--------- 2 files changed, 50 insertions(+), 19 deletions(-) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 21c3aa3608..301df2601b 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -60,6 +60,23 @@ func discoveryErr(msg string, err error) error { return fmt.Errorf("%s%s", msg, discoveryFallbackTip) } +// dualWriteLegacyHostKey mirrors the freshly minted token under the legacy +// host-based cache key so users alternating between CLI and SDK find it. +// Skipped for secure mode to avoid multiplying keyring entries. +func dualWriteLegacyHostKey(ctx context.Context, tokenCache cache.TokenCache, arg u2m.OAuthArgument, mode storage.StorageMode) { + if mode != storage.StorageModeLegacy { + return + } + t, err := tokenCache.Lookup(arg.GetCacheKey()) + if err != nil || t == nil { + return + } + dual := storage.NewDualWritingTokenCache(tokenCache, arg) + if err := dual.Store(arg.GetCacheKey(), t); err != nil { + log.Debugf(ctx, "token cache dual-write failed: %v", err) + } +} + type discoveryPersistentAuth interface { Challenge() error Token() (*oauth2.Token, error) @@ -141,12 +158,29 @@ a new profile is created. cmd.Flags().StringVar(&scopes, "scopes", "", "Comma-separated list of OAuth scopes to request (defaults to 'all-apis')") + var secureStorage bool + cmd.Flags().BoolVar(&secureStorage, "secure-storage", false, + "Experimental: write OAuth tokens to the OS-native secure store") + // Hidden during MS1; discovery is via release notes and the + // DATABRICKS_AUTH_STORAGE env var. See + // documents/fy2027-q2/cli-ga/2026-04-13-cli-ga-rollout-contract.md. + _ = cmd.Flags().MarkHidden("secure-storage") + cmd.PreRunE = profileHostConflictCheck cmd.RunE = func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() profileName := cmd.Flag("profile").Value.String() + var storageOverride storage.StorageMode + if secureStorage { + storageOverride = storage.StorageModeSecure + } + tokenCache, mode, err := newAuthCache(ctx, storageOverride) + if err != nil { + return err + } + // Cluster and Serverless are mutually exclusive. if configureCluster && configureServerless { return errors.New("please either configure serverless or cluster, not both") @@ -191,18 +225,13 @@ a new profile is created. return err } - tokenCache, err := storage.NewFileTokenCache(ctx) - if err != nil { - return fmt.Errorf("opening token cache: %w", err) - } - // If no host is available from any source, use the discovery flow // via login.databricks.com. if shouldUseDiscovery(authArguments.Host, args, existingProfile) { if err := validateDiscoveryFlagCompatibility(cmd); err != nil { return err } - return discoveryLogin(ctx, &defaultDiscoveryClient{}, tokenCache, profileName, loginTimeout, scopes, existingProfile, getBrowserFunc(cmd)) + return discoveryLogin(ctx, &defaultDiscoveryClient{}, profileName, loginTimeout, scopes, existingProfile, getBrowserFunc(cmd), tokenCache, mode) } // Load unified host flag from the profile if not explicitly set via CLI flag. @@ -235,9 +264,9 @@ a new profile is created. return err } persistentAuthOpts := []u2m.PersistentAuthOption{ - u2m.WithTokenCache(storage.NewDualWritingTokenCache(tokenCache, oauthArgument)), u2m.WithOAuthArgument(oauthArgument), u2m.WithBrowser(getBrowserFunc(cmd)), + u2m.WithTokenCache(tokenCache), } if len(scopesList) > 0 { persistentAuthOpts = append(persistentAuthOpts, u2m.WithScopes(scopesList)) @@ -254,6 +283,7 @@ a new profile is created. if err = persistentAuth.Challenge(); err != nil { return err } + dualWriteLegacyHostKey(ctx, tokenCache, oauthArgument, mode) // At this point, an OAuth token has been successfully minted and stored // in the CLI cache. The rest of the command focuses on: // 1. Workspace selection for SPOG hosts (best-effort); @@ -570,7 +600,7 @@ func validateDiscoveryFlagCompatibility(cmd *cobra.Command) error { // discoveryLogin runs the login.databricks.com discovery flow. The user // authenticates in the browser, selects a workspace, and the CLI receives // the workspace host from the OAuth callback's iss parameter. -func discoveryLogin(ctx context.Context, dc discoveryClient, tokenCache cache.TokenCache, profileName string, timeout time.Duration, scopes string, existingProfile *profile.Profile, browserFunc func(string) error) error { +func discoveryLogin(ctx context.Context, dc discoveryClient, profileName string, timeout time.Duration, scopes string, existingProfile *profile.Profile, browserFunc func(string) error, tokenCache cache.TokenCache, mode storage.StorageMode) error { arg, err := dc.NewOAuthArgument(profileName) if err != nil { return discoveryErr("setting up login.databricks.com", err) @@ -582,10 +612,10 @@ func discoveryLogin(ctx context.Context, dc discoveryClient, tokenCache cache.To } opts := []u2m.PersistentAuthOption{ - u2m.WithTokenCache(storage.NewDualWritingTokenCache(tokenCache, arg)), u2m.WithOAuthArgument(arg), u2m.WithBrowser(browserFunc), u2m.WithDiscoveryLogin(), + u2m.WithTokenCache(tokenCache), } if len(scopesList) > 0 { opts = append(opts, u2m.WithScopes(scopesList)) @@ -605,6 +635,7 @@ func discoveryLogin(ctx context.Context, dc discoveryClient, tokenCache cache.To if err := persistentAuth.Challenge(); err != nil { return discoveryErr("login via login.databricks.com failed", err) } + dualWriteLegacyHostKey(ctx, tokenCache, arg, mode) discoveredHost := arg.GetDiscoveredHost() if discoveredHost == "" { diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index 2b8d473f51..c57efd352f 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -630,7 +630,7 @@ func TestDiscoveryLogin_IntrospectionFailureStillSavesProfile(t *testing.T) { } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "all-apis, ,sql,", nil, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "all-apis, ,sql,", nil, func(string) error { return nil }, newTestTokenCache(), "") require.NoError(t, err) assert.Equal(t, "https://workspace.example.com", dc.introspectHost) @@ -678,7 +678,7 @@ func TestDiscoveryLogin_AccountIDMismatchWarning(t *testing.T) { AccountID: "old-account-id", } - err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }, newTestTokenCache(), "") require.NoError(t, err) // Verify warning about mismatched account IDs was logged. @@ -726,7 +726,7 @@ func TestDiscoveryLogin_NoWarningWhenAccountIDsMatch(t *testing.T) { AccountID: "same-account-id", } - err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }, newTestTokenCache(), "") require.NoError(t, err) // No warning should be logged when account IDs match. @@ -746,7 +746,7 @@ func TestDiscoveryLogin_EmptyDiscoveredHostReturnsError(t *testing.T) { } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", nil, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", nil, func(string) error { return nil }, newTestTokenCache(), "") require.Error(t, err) assert.Contains(t, err.Error(), "no workspace host was discovered") } @@ -778,7 +778,7 @@ func TestDiscoveryLogin_ReloginPreservesExistingProfileScopes(t *testing.T) { // No --scopes flag (empty string), should fall back to existing profile scopes. ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }, newTestTokenCache(), "") require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -815,7 +815,7 @@ func TestDiscoveryLogin_ExplicitScopesOverrideExistingProfile(t *testing.T) { // Explicit --scopes flag should override existing profile scopes. ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "all-apis", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "all-apis", existingProfile, func(string) error { return nil }, newTestTokenCache(), "") require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -855,7 +855,7 @@ func TestDiscoveryLogin_SPOGHostPopulatesAccountIDFromDiscovery(t *testing.T) { } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", nil, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", nil, func(string) error { return nil }, newTestTokenCache(), "") require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -890,7 +890,7 @@ func TestDiscoveryLogin_IntrospectionFallsBackWhenDiscoveryFails(t *testing.T) { } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", nil, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", nil, func(string) error { return nil }, newTestTokenCache(), "") require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -939,7 +939,7 @@ auth_type = databricks-cli } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }, newTestTokenCache(), "") require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -989,7 +989,7 @@ auth_type = databricks-cli } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }, newTestTokenCache(), "") require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) From 5d3a9f2815419c14ee224643e08c10eb07f6cd29 Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 17 Apr 2026 13:48:15 +0200 Subject: [PATCH 04/12] cmd/auth: wire secure-storage cache into auth token Resolves the token cache via newAuthCache at the two NewPersistentAuth call sites inside cmd/auth/token.go (loadToken via newTokenCommand and runInlineLogin). No flag is added: per the rollout contract, auth token reads from whichever mode the resolver selects via DATABRICKS_AUTH_STORAGE or [__settings__].auth_storage. --- cmd/auth/token.go | 24 +++++++++++++++--------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/cmd/auth/token.go b/cmd/auth/token.go index d82bcb2c9a..ce1b0a04b3 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -79,9 +79,9 @@ and secret is not supported.`, ctx := cmd.Context() profileName := cmd.Flag("profile").Value.String() - tokenCache, err := storage.NewFileTokenCache(ctx) + tokenCache, mode, err := newAuthCache(ctx, "") if err != nil { - return fmt.Errorf("opening token cache: %w", err) + return err } t, err := loadToken(ctx, loadTokenArgs{ @@ -92,6 +92,7 @@ and secret is not supported.`, forceRefresh: forceRefresh, profiler: profile.DefaultProfiler, tokenCache: tokenCache, + mode: mode, persistentAuthOpts: nil, }) if err != nil { @@ -144,6 +145,11 @@ type loadTokenArgs struct { // responsible for construction so that tests can substitute an in-memory cache. tokenCache cache.TokenCache + // mode is the resolved storage mode. When set to StorageModeLegacy, login + // paths mirror freshly minted tokens under the legacy host-based key so + // older SDKs that still look up by host continue to find them. + mode storage.StorageMode + // persistentAuthOpts are the options to pass to the persistent auth client. persistentAuthOpts []u2m.PersistentAuthOption } @@ -195,7 +201,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.tokenCache) + resolvedProfile, existingProfile, err = resolveNoArgsToken(ctx, args.profiler, args.authArguments, args.tokenCache, args.mode) if err != nil { return nil, err } @@ -284,8 +290,7 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) { if err != nil { return nil, err } - wrappedCache := storage.NewDualWritingTokenCache(args.tokenCache, oauthArgument) - allArgs := append([]u2m.PersistentAuthOption{u2m.WithTokenCache(wrappedCache)}, args.persistentAuthOpts...) + allArgs := append([]u2m.PersistentAuthOption{u2m.WithTokenCache(args.tokenCache)}, args.persistentAuthOpts...) allArgs = append(allArgs, u2m.WithOAuthArgument(oauthArgument)) persistentAuth, err := u2m.NewPersistentAuth(ctx, allArgs...) if err != nil { @@ -327,7 +332,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, tokenCache cache.TokenCache) (string, *profile.Profile, error) { +func resolveNoArgsToken(ctx context.Context, profiler profile.Profiler, authArgs *auth.AuthArguments, tokenCache cache.TokenCache, mode storage.StorageMode) (string, *profile.Profile, error) { // Step 1: Try DATABRICKS_HOST env var (highest priority). if envHost := env.Get(ctx, "DATABRICKS_HOST"); envHost != "" { authArgs.Host = envHost @@ -376,7 +381,7 @@ func resolveNoArgsToken(ctx context.Context, profiler profile.Profiler, authArgs // Fall through — setHostAndAccountId will prompt for the host. return "", nil, nil case createNewSelected: - return runInlineLogin(ctx, profiler, tokenCache) + return runInlineLogin(ctx, profiler, tokenCache, mode) default: p, err := loadProfileByName(ctx, selectedName, profiler) if err != nil { @@ -440,7 +445,7 @@ func promptForProfileSelection(ctx context.Context, profiles profile.Profiles) ( // runInlineLogin runs a minimal interactive login flow: prompts for a profile // name and host, performs the OAuth challenge, saves the profile to // .databrickscfg, and returns the new profile name and profile. -func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenCache cache.TokenCache) (string, *profile.Profile, error) { +func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenCache cache.TokenCache, mode storage.StorageMode) (string, *profile.Profile, error) { profileName, err := promptForProfile(ctx, "DEFAULT") if err != nil { return "", nil, err @@ -473,9 +478,9 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenCache c return "", nil, err } persistentAuthOpts := []u2m.PersistentAuthOption{ - u2m.WithTokenCache(storage.NewDualWritingTokenCache(tokenCache, oauthArgument)), u2m.WithOAuthArgument(oauthArgument), u2m.WithBrowser(func(url string) error { return browser.Open(ctx, url) }), + u2m.WithTokenCache(tokenCache), } if len(scopesList) > 0 { persistentAuthOpts = append(persistentAuthOpts, u2m.WithScopes(scopesList)) @@ -492,6 +497,7 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenCache c if err = persistentAuth.Challenge(); err != nil { return "", nil, err } + dualWriteLegacyHostKey(ctx, tokenCache, oauthArgument, mode) clearKeys := oauthLoginClearKeys() if !loginArgs.IsUnifiedHost { From ed83359faaa9a0f4ebf6e0f067fb69cb901a8904 Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 17 Apr 2026 13:52:29 +0200 Subject: [PATCH 05/12] cmd/auth: route auth logout through newAuthCache Replaces the direct cache.NewFileTokenCache() call with newAuthCache so that logout targets whichever backend ResolveStorageMode selects. The resolver still defaults to legacy, so existing flows are unchanged; once a user opts into secure storage, logout clears the keyring entry instead of the file cache. --- cmd/auth/logout.go | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/cmd/auth/logout.go b/cmd/auth/logout.go index 67829ec169..81c10a4f8b 100644 --- a/cmd/auth/logout.go +++ b/cmd/auth/logout.go @@ -22,7 +22,6 @@ import ( "strings" "github.com/databricks/cli/libs/auth" - "github.com/databricks/cli/libs/auth/storage" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/databrickscfg" "github.com/databricks/cli/libs/databrickscfg/profile" @@ -133,9 +132,9 @@ to specify it explicitly. profileName = selected } - tokenCache, err := storage.NewFileTokenCache(ctx) + tokenCache, _, err := newAuthCache(ctx, "") if err != nil { - return fmt.Errorf("failed to open token cache, please check if the file version is up-to-date and that the file is not corrupted: %w", err) + return fmt.Errorf("failed to open token cache: %w", err) } return runLogout(ctx, logoutArgs{ From 13e138cd235dbdac850d8d200851cca957df61df Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 17 Apr 2026 14:37:33 +0200 Subject: [PATCH 06/12] acceptance: cover storage-mode wiring for auth commands Adds three acceptance tests under acceptance/cmd/auth/storage-modes/: - invalid-env: DATABRICKS_AUTH_STORAGE=bogus surfaces a clear error - legacy-env-default: explicit legacy mode behaves like the default path - hidden-flag: --secure-storage is hidden from auth login --help Secure-mode end-to-end behavior is covered by unit tests in cmd/auth/cache_test.go because acceptance runners cannot safely exercise the real OS keyring on macOS (would write to Keychain) or Linux (no D-Bus session in CI). --- .../storage-modes/hidden-flag/out.test.toml | 5 ++ .../auth/storage-modes/hidden-flag/output.txt | 47 +++++++++++++++++++ .../cmd/auth/storage-modes/hidden-flag/script | 4 ++ .../storage-modes/invalid-env/out.test.toml | 5 ++ .../auth/storage-modes/invalid-env/output.txt | 5 ++ .../cmd/auth/storage-modes/invalid-env/script | 6 +++ .../legacy-env-default/out.test.toml | 5 ++ .../legacy-env-default/output.txt | 11 +++++ .../storage-modes/legacy-env-default/script | 29 ++++++++++++ .../cmd/auth/storage-modes/script.prepare | 8 ++++ acceptance/cmd/auth/storage-modes/test.toml | 3 ++ 11 files changed, 128 insertions(+) create mode 100644 acceptance/cmd/auth/storage-modes/hidden-flag/out.test.toml create mode 100644 acceptance/cmd/auth/storage-modes/hidden-flag/output.txt create mode 100644 acceptance/cmd/auth/storage-modes/hidden-flag/script create mode 100644 acceptance/cmd/auth/storage-modes/invalid-env/out.test.toml create mode 100644 acceptance/cmd/auth/storage-modes/invalid-env/output.txt create mode 100644 acceptance/cmd/auth/storage-modes/invalid-env/script create mode 100644 acceptance/cmd/auth/storage-modes/legacy-env-default/out.test.toml create mode 100644 acceptance/cmd/auth/storage-modes/legacy-env-default/output.txt create mode 100644 acceptance/cmd/auth/storage-modes/legacy-env-default/script create mode 100644 acceptance/cmd/auth/storage-modes/script.prepare create mode 100644 acceptance/cmd/auth/storage-modes/test.toml diff --git a/acceptance/cmd/auth/storage-modes/hidden-flag/out.test.toml b/acceptance/cmd/auth/storage-modes/hidden-flag/out.test.toml new file mode 100644 index 0000000000..d560f1de04 --- /dev/null +++ b/acceptance/cmd/auth/storage-modes/hidden-flag/out.test.toml @@ -0,0 +1,5 @@ +Local = true +Cloud = false + +[EnvMatrix] + DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/cmd/auth/storage-modes/hidden-flag/output.txt b/acceptance/cmd/auth/storage-modes/hidden-flag/output.txt new file mode 100644 index 0000000000..f6a6daf147 --- /dev/null +++ b/acceptance/cmd/auth/storage-modes/hidden-flag/output.txt @@ -0,0 +1,47 @@ +Log into a Databricks workspace or account. + +This command authenticates via OAuth in the browser and saves the result +to a configuration profile (in ~/.databrickscfg by default). Other Databricks CLI +commands and SDKs can use this profile via the --profile flag. For more +information, see: + AWS: https://docs.databricks.com/dev-tools/auth/index.html + Azure: https://learn.microsoft.com/azure/databricks/dev-tools/auth + GCP: https://docs.gcp.databricks.com/dev-tools/auth/index.html + +If no host is provided, the CLI opens login.databricks.com where you can +authenticate and select a workspace. + +The positional argument is resolved as a profile name first. If no profile +with that name exists and the argument looks like a URL, it is used as a +host. The positional argument cannot be combined with --host or --profile; +use the flags directly to specify both. + +The host URL may include query parameters to set the workspace and account ID: + + databricks auth login --host "https://?o=&account_id=" + +Note: URLs containing "?" must be quoted to prevent shell interpretation. + +If a profile with the given name already exists, it is updated. Otherwise +a new profile is created. + +Usage: + databricks auth login [PROFILE] [flags] + +Flags: + --configure-cluster Prompts to configure cluster + --configure-serverless Prompts to configure serverless + -h, --help help for login + --scopes string Comma-separated list of OAuth scopes to request (defaults to 'all-apis') + --skip-workspace Skip workspace selection for account-level access + --timeout duration Timeout for completing login challenge in the browser (default 1h0m0s) + +Global Flags: + --account-id string Databricks Account ID + --debug enable debug logging + --experimental-is-unified-host Flag to indicate if the host is a unified host + --host string Databricks Host + -o, --output type output type: text or json (default text) + -p, --profile string ~/.databrickscfg profile + -t, --target string bundle target to use (if applicable) + --workspace-id string Databricks Workspace ID diff --git a/acceptance/cmd/auth/storage-modes/hidden-flag/script b/acceptance/cmd/auth/storage-modes/hidden-flag/script new file mode 100644 index 0000000000..3d36f2d4d0 --- /dev/null +++ b/acceptance/cmd/auth/storage-modes/hidden-flag/script @@ -0,0 +1,4 @@ +# The --secure-storage flag is hidden during MS1. It must not appear in +# --help output. Route through contains.py so a regression fails loudly +# regardless of help-text reordering. +$CLI auth login --help | contains.py !--secure-storage diff --git a/acceptance/cmd/auth/storage-modes/invalid-env/out.test.toml b/acceptance/cmd/auth/storage-modes/invalid-env/out.test.toml new file mode 100644 index 0000000000..d560f1de04 --- /dev/null +++ b/acceptance/cmd/auth/storage-modes/invalid-env/out.test.toml @@ -0,0 +1,5 @@ +Local = true +Cloud = false + +[EnvMatrix] + DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/cmd/auth/storage-modes/invalid-env/output.txt b/acceptance/cmd/auth/storage-modes/invalid-env/output.txt new file mode 100644 index 0000000000..5723269bb1 --- /dev/null +++ b/acceptance/cmd/auth/storage-modes/invalid-env/output.txt @@ -0,0 +1,5 @@ + +>>> [CLI] auth token --profile nonexistent +Error: DATABRICKS_AUTH_STORAGE: unknown storage mode "bogus" (want legacy, secure, or plaintext) + +Exit code: 1 diff --git a/acceptance/cmd/auth/storage-modes/invalid-env/script b/acceptance/cmd/auth/storage-modes/invalid-env/script new file mode 100644 index 0000000000..8b00e9d54d --- /dev/null +++ b/acceptance/cmd/auth/storage-modes/invalid-env/script @@ -0,0 +1,6 @@ +export DATABRICKS_AUTH_STORAGE=bogus + +# Any auth command that resolves the storage mode must surface the error. +# auth token is the smallest reproducer because it doesn't perform any +# network I/O before resolving the mode. +trace $CLI auth token --profile nonexistent diff --git a/acceptance/cmd/auth/storage-modes/legacy-env-default/out.test.toml b/acceptance/cmd/auth/storage-modes/legacy-env-default/out.test.toml new file mode 100644 index 0000000000..d560f1de04 --- /dev/null +++ b/acceptance/cmd/auth/storage-modes/legacy-env-default/out.test.toml @@ -0,0 +1,5 @@ +Local = true +Cloud = false + +[EnvMatrix] + DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/cmd/auth/storage-modes/legacy-env-default/output.txt b/acceptance/cmd/auth/storage-modes/legacy-env-default/output.txt new file mode 100644 index 0000000000..8994c41ebb --- /dev/null +++ b/acceptance/cmd/auth/storage-modes/legacy-env-default/output.txt @@ -0,0 +1,11 @@ + +=== Token cache keys before logout +[ + "dev" +] + +>>> [CLI] auth logout --profile dev --auto-approve +Logged out of profile "dev". Use --delete to also remove it from the config file. + +=== Token cache keys after logout (should be empty) +[] diff --git a/acceptance/cmd/auth/storage-modes/legacy-env-default/script b/acceptance/cmd/auth/storage-modes/legacy-env-default/script new file mode 100644 index 0000000000..37f367fd73 --- /dev/null +++ b/acceptance/cmd/auth/storage-modes/legacy-env-default/script @@ -0,0 +1,29 @@ +export DATABRICKS_AUTH_STORAGE=legacy + +cat > "./home/.databrickscfg" < "./home/.databricks/token-cache.json" < Date: Fri, 17 Apr 2026 14:42:11 +0200 Subject: [PATCH 07/12] NEXT_CHANGELOG: announce experimental secure token storage --- NEXT_CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index b9ae889e7a..eb057193ff 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -10,6 +10,7 @@ * Accept `yes` in addition to `y` for confirmation prompts, and show `[y/N]` to indicate that no is the default. * Deprecated `auth env`. The command is hidden from help listings and prints a deprecation warning to stderr; it will be removed in a future release. * Moved file-based OAuth token cache management from the SDK to the CLI. No user-visible change; part of a three-PR sequence that makes the CLI the sole owner of its token cache. +* Added experimental OS-native secure token storage behind the `--secure-storage` flag on `databricks auth login` and the `DATABRICKS_AUTH_STORAGE=secure` environment variable. Hidden from help during MS1. Legacy file-backed token storage remains the default. ### Bundles * Remove `experimental-jobs-as-code` template, superseded by `pydabs` ([#4999](https://github.com/databricks/cli/pull/4999)). From fb67d7fbdec3d9d4cee3a4b28a384c92dfb88de3 Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 17 Apr 2026 15:14:43 +0200 Subject: [PATCH 08/12] cmd/auth: drop --secure-storage flag; resolve mode via env/config only Storage backend is a per-machine setting, not a per-invocation choice. Keeping a write-time flag on login while token/logout rely on the resolver creates an asymmetry: login --secure-storage writes to the keyring, but a later auth token without the env var defaults to the legacy file cache and can't find the token. Resolution is now always env -> [__settings__].auth_storage -> default. login.go matches token.go and logout.go: newAuthCache(ctx, "") with no per-call override from a flag. Also drops the hidden-flag acceptance test (no flag to hide) and updates NEXT_CHANGELOG to advertise the env var and config key instead of the flag. --- NEXT_CHANGELOG.md | 1 + .../storage-modes/hidden-flag/out.test.toml | 5 -- .../auth/storage-modes/hidden-flag/output.txt | 47 ------------------- .../cmd/auth/storage-modes/hidden-flag/script | 4 -- cmd/auth/login.go | 14 +----- 5 files changed, 2 insertions(+), 69 deletions(-) delete mode 100644 acceptance/cmd/auth/storage-modes/hidden-flag/out.test.toml delete mode 100644 acceptance/cmd/auth/storage-modes/hidden-flag/output.txt delete mode 100644 acceptance/cmd/auth/storage-modes/hidden-flag/script diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index eb057193ff..27ce418ded 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -11,6 +11,7 @@ * Deprecated `auth env`. The command is hidden from help listings and prints a deprecation warning to stderr; it will be removed in a future release. * Moved file-based OAuth token cache management from the SDK to the CLI. No user-visible change; part of a three-PR sequence that makes the CLI the sole owner of its token cache. * Added experimental OS-native secure token storage behind the `--secure-storage` flag on `databricks auth login` and the `DATABRICKS_AUTH_STORAGE=secure` environment variable. Hidden from help during MS1. Legacy file-backed token storage remains the default. +* Added experimental OS-native secure token storage opt-in via `DATABRICKS_AUTH_STORAGE=secure` or `[__settings__].auth_storage = secure` in `.databrickscfg`. Legacy file-backed token storage remains the default. ### Bundles * Remove `experimental-jobs-as-code` template, superseded by `pydabs` ([#4999](https://github.com/databricks/cli/pull/4999)). diff --git a/acceptance/cmd/auth/storage-modes/hidden-flag/out.test.toml b/acceptance/cmd/auth/storage-modes/hidden-flag/out.test.toml deleted file mode 100644 index d560f1de04..0000000000 --- a/acceptance/cmd/auth/storage-modes/hidden-flag/out.test.toml +++ /dev/null @@ -1,5 +0,0 @@ -Local = true -Cloud = false - -[EnvMatrix] - DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/cmd/auth/storage-modes/hidden-flag/output.txt b/acceptance/cmd/auth/storage-modes/hidden-flag/output.txt deleted file mode 100644 index f6a6daf147..0000000000 --- a/acceptance/cmd/auth/storage-modes/hidden-flag/output.txt +++ /dev/null @@ -1,47 +0,0 @@ -Log into a Databricks workspace or account. - -This command authenticates via OAuth in the browser and saves the result -to a configuration profile (in ~/.databrickscfg by default). Other Databricks CLI -commands and SDKs can use this profile via the --profile flag. For more -information, see: - AWS: https://docs.databricks.com/dev-tools/auth/index.html - Azure: https://learn.microsoft.com/azure/databricks/dev-tools/auth - GCP: https://docs.gcp.databricks.com/dev-tools/auth/index.html - -If no host is provided, the CLI opens login.databricks.com where you can -authenticate and select a workspace. - -The positional argument is resolved as a profile name first. If no profile -with that name exists and the argument looks like a URL, it is used as a -host. The positional argument cannot be combined with --host or --profile; -use the flags directly to specify both. - -The host URL may include query parameters to set the workspace and account ID: - - databricks auth login --host "https://?o=&account_id=" - -Note: URLs containing "?" must be quoted to prevent shell interpretation. - -If a profile with the given name already exists, it is updated. Otherwise -a new profile is created. - -Usage: - databricks auth login [PROFILE] [flags] - -Flags: - --configure-cluster Prompts to configure cluster - --configure-serverless Prompts to configure serverless - -h, --help help for login - --scopes string Comma-separated list of OAuth scopes to request (defaults to 'all-apis') - --skip-workspace Skip workspace selection for account-level access - --timeout duration Timeout for completing login challenge in the browser (default 1h0m0s) - -Global Flags: - --account-id string Databricks Account ID - --debug enable debug logging - --experimental-is-unified-host Flag to indicate if the host is a unified host - --host string Databricks Host - -o, --output type output type: text or json (default text) - -p, --profile string ~/.databrickscfg profile - -t, --target string bundle target to use (if applicable) - --workspace-id string Databricks Workspace ID diff --git a/acceptance/cmd/auth/storage-modes/hidden-flag/script b/acceptance/cmd/auth/storage-modes/hidden-flag/script deleted file mode 100644 index 3d36f2d4d0..0000000000 --- a/acceptance/cmd/auth/storage-modes/hidden-flag/script +++ /dev/null @@ -1,4 +0,0 @@ -# The --secure-storage flag is hidden during MS1. It must not appear in -# --help output. Route through contains.py so a regression fails loudly -# regardless of help-text reordering. -$CLI auth login --help | contains.py !--secure-storage diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 301df2601b..17751a3a8e 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -158,25 +158,13 @@ a new profile is created. cmd.Flags().StringVar(&scopes, "scopes", "", "Comma-separated list of OAuth scopes to request (defaults to 'all-apis')") - var secureStorage bool - cmd.Flags().BoolVar(&secureStorage, "secure-storage", false, - "Experimental: write OAuth tokens to the OS-native secure store") - // Hidden during MS1; discovery is via release notes and the - // DATABRICKS_AUTH_STORAGE env var. See - // documents/fy2027-q2/cli-ga/2026-04-13-cli-ga-rollout-contract.md. - _ = cmd.Flags().MarkHidden("secure-storage") - cmd.PreRunE = profileHostConflictCheck cmd.RunE = func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() profileName := cmd.Flag("profile").Value.String() - var storageOverride storage.StorageMode - if secureStorage { - storageOverride = storage.StorageModeSecure - } - tokenCache, mode, err := newAuthCache(ctx, storageOverride) + tokenCache, mode, err := newAuthCache(ctx, "") if err != nil { return err } From ee7a01f81a59bc6aadf764a681a2b9a9be3833ae Mon Sep 17 00:00:00 2001 From: simon Date: Mon, 20 Apr 2026 08:28:58 +0200 Subject: [PATCH 09/12] cmd/auth: apply review feedback Drop MS1/MS3 comment markers from cache.go and convert discoveryLogin to take a discoveryLoginInputs struct for readability (8 positional args was over the threshold). No behavior change. Co-authored-by: Isaac --- cmd/auth/cache.go | 5 +-- cmd/auth/login.go | 60 +++++++++++++++++++--------- cmd/auth/login_test.go | 88 +++++++++++++++++++++++++++++++++++++----- 3 files changed, 121 insertions(+), 32 deletions(-) diff --git a/cmd/auth/cache.go b/cmd/auth/cache.go index 29022e4585..3487573a16 100644 --- a/cmd/auth/cache.go +++ b/cmd/auth/cache.go @@ -48,9 +48,8 @@ func newAuthCacheWith(ctx context.Context, override storage.StorageMode, f cache case storage.StorageModeSecure: return f.newKeyring(), mode, nil case storage.StorageModeLegacy, storage.StorageModePlaintext: - // MS1 ships no dedicated plaintext implementation; the switch will - // be added in MS3. Until then the file cache is the safest existing - // behavior and the resolver still accepts the value. + // Plaintext currently maps to the file cache; a dedicated + // plaintext backend (no host-keyed dual-writes) is a follow-up. c, err := f.newFile() if err != nil { return nil, "", fmt.Errorf("open file token cache: %w", err) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 17751a3a8e..34dca1bc30 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -219,7 +219,16 @@ a new profile is created. if err := validateDiscoveryFlagCompatibility(cmd); err != nil { return err } - return discoveryLogin(ctx, &defaultDiscoveryClient{}, profileName, loginTimeout, scopes, existingProfile, getBrowserFunc(cmd), tokenCache, mode) + return discoveryLogin(ctx, discoveryLoginInputs{ + dc: &defaultDiscoveryClient{}, + profileName: profileName, + timeout: loginTimeout, + scopes: scopes, + existingProfile: existingProfile, + browserFunc: getBrowserFunc(cmd), + tokenCache: tokenCache, + mode: mode, + }) } // Load unified host flag from the profile if not explicitly set via CLI flag. @@ -585,35 +594,48 @@ func validateDiscoveryFlagCompatibility(cmd *cobra.Command) error { return nil } +// discoveryLoginInputs groups the dependencies of discoveryLogin. +// See https://google.github.io/styleguide/go/best-practices#option-structure. +type discoveryLoginInputs struct { + dc discoveryClient + profileName string + timeout time.Duration + scopes string + existingProfile *profile.Profile + browserFunc func(string) error + tokenCache cache.TokenCache + mode storage.StorageMode +} + // discoveryLogin runs the login.databricks.com discovery flow. The user // authenticates in the browser, selects a workspace, and the CLI receives // the workspace host from the OAuth callback's iss parameter. -func discoveryLogin(ctx context.Context, dc discoveryClient, profileName string, timeout time.Duration, scopes string, existingProfile *profile.Profile, browserFunc func(string) error, tokenCache cache.TokenCache, mode storage.StorageMode) error { - arg, err := dc.NewOAuthArgument(profileName) +func discoveryLogin(ctx context.Context, in discoveryLoginInputs) error { + arg, err := in.dc.NewOAuthArgument(in.profileName) if err != nil { return discoveryErr("setting up login.databricks.com", err) } - scopesList := splitScopes(scopes) - if len(scopesList) == 0 && existingProfile != nil && existingProfile.Scopes != "" { - scopesList = splitScopes(existingProfile.Scopes) + scopesList := splitScopes(in.scopes) + if len(scopesList) == 0 && in.existingProfile != nil && in.existingProfile.Scopes != "" { + scopesList = splitScopes(in.existingProfile.Scopes) } opts := []u2m.PersistentAuthOption{ u2m.WithOAuthArgument(arg), - u2m.WithBrowser(browserFunc), + u2m.WithBrowser(in.browserFunc), u2m.WithDiscoveryLogin(), - u2m.WithTokenCache(tokenCache), + u2m.WithTokenCache(in.tokenCache), } if len(scopesList) > 0 { opts = append(opts, u2m.WithScopes(scopesList)) } // Apply timeout before creating PersistentAuth so Challenge() respects it. - ctx, cancel := context.WithTimeout(ctx, timeout) + ctx, cancel := context.WithTimeout(ctx, in.timeout) defer cancel() - persistentAuth, err := dc.NewPersistentAuth(ctx, opts...) + persistentAuth, err := in.dc.NewPersistentAuth(ctx, opts...) if err != nil { return discoveryErr("setting up login.databricks.com", err) } @@ -623,7 +645,7 @@ func discoveryLogin(ctx context.Context, dc discoveryClient, profileName string, if err := persistentAuth.Challenge(); err != nil { return discoveryErr("login via login.databricks.com failed", err) } - dualWriteLegacyHostKey(ctx, tokenCache, arg, mode) + dualWriteLegacyHostKey(ctx, in.tokenCache, arg, in.mode) discoveredHost := arg.GetDiscoveredHost() if discoveredHost == "" { @@ -646,7 +668,7 @@ func discoveryLogin(ctx context.Context, dc discoveryClient, profileName string, return fmt.Errorf("retrieving token after login: %w", err) } - introspection, err := dc.IntrospectToken(ctx, discoveredHost, tok.AccessToken) + introspection, err := in.dc.IntrospectToken(ctx, discoveredHost, tok.AccessToken) if err != nil { log.Debugf(ctx, "token introspection failed (non-fatal): %v", err) } else { @@ -657,10 +679,10 @@ func discoveryLogin(ctx context.Context, dc discoveryClient, profileName string, accountID = introspection.AccountID } - if existingProfile != nil && existingProfile.AccountID != "" && introspection.AccountID != "" && - existingProfile.AccountID != introspection.AccountID { + if in.existingProfile != nil && in.existingProfile.AccountID != "" && introspection.AccountID != "" && + in.existingProfile.AccountID != introspection.AccountID { log.Warnf(ctx, "detected account ID %q differs from existing profile account ID %q", - introspection.AccountID, existingProfile.AccountID) + introspection.AccountID, in.existingProfile.AccountID) } } @@ -679,7 +701,7 @@ func discoveryLogin(ctx context.Context, dc discoveryClient, profileName string, "serverless_compute_id", ) err = databrickscfg.SaveToProfile(ctx, &config.Config{ - Profile: profileName, + Profile: in.profileName, Host: discoveredHost, AuthType: authTypeDatabricksCLI, AccountID: accountID, @@ -689,12 +711,12 @@ func discoveryLogin(ctx context.Context, dc discoveryClient, profileName string, }, clearKeys...) if err != nil { if configFile != "" { - return fmt.Errorf("saving profile %q to %s: %w", profileName, configFile, err) + return fmt.Errorf("saving profile %q to %s: %w", in.profileName, configFile, err) } - return fmt.Errorf("saving profile %q: %w", profileName, err) + return fmt.Errorf("saving profile %q: %w", in.profileName, err) } - cmdio.LogString(ctx, fmt.Sprintf("Profile %s was successfully saved", profileName)) + cmdio.LogString(ctx, fmt.Sprintf("Profile %s was successfully saved", in.profileName)) return nil } diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index c57efd352f..b5d8a39f43 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -630,7 +630,14 @@ func TestDiscoveryLogin_IntrospectionFailureStillSavesProfile(t *testing.T) { } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "all-apis, ,sql,", nil, func(string) error { return nil }, newTestTokenCache(), "") + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + scopes: "all-apis, ,sql,", + browserFunc: func(string) error { return nil }, + tokenCache: newTestTokenCache(), + }) require.NoError(t, err) assert.Equal(t, "https://workspace.example.com", dc.introspectHost) @@ -678,7 +685,14 @@ func TestDiscoveryLogin_AccountIDMismatchWarning(t *testing.T) { AccountID: "old-account-id", } - err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }, newTestTokenCache(), "") + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + existingProfile: existingProfile, + browserFunc: func(string) error { return nil }, + tokenCache: newTestTokenCache(), + }) require.NoError(t, err) // Verify warning about mismatched account IDs was logged. @@ -726,7 +740,14 @@ func TestDiscoveryLogin_NoWarningWhenAccountIDsMatch(t *testing.T) { AccountID: "same-account-id", } - err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }, newTestTokenCache(), "") + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + existingProfile: existingProfile, + browserFunc: func(string) error { return nil }, + tokenCache: newTestTokenCache(), + }) require.NoError(t, err) // No warning should be logged when account IDs match. @@ -746,7 +767,13 @@ func TestDiscoveryLogin_EmptyDiscoveredHostReturnsError(t *testing.T) { } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", nil, func(string) error { return nil }, newTestTokenCache(), "") + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + browserFunc: func(string) error { return nil }, + tokenCache: newTestTokenCache(), + }) require.Error(t, err) assert.Contains(t, err.Error(), "no workspace host was discovered") } @@ -778,7 +805,14 @@ func TestDiscoveryLogin_ReloginPreservesExistingProfileScopes(t *testing.T) { // No --scopes flag (empty string), should fall back to existing profile scopes. ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }, newTestTokenCache(), "") + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + existingProfile: existingProfile, + browserFunc: func(string) error { return nil }, + tokenCache: newTestTokenCache(), + }) require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -815,7 +849,15 @@ func TestDiscoveryLogin_ExplicitScopesOverrideExistingProfile(t *testing.T) { // Explicit --scopes flag should override existing profile scopes. ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "all-apis", existingProfile, func(string) error { return nil }, newTestTokenCache(), "") + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + scopes: "all-apis", + existingProfile: existingProfile, + browserFunc: func(string) error { return nil }, + tokenCache: newTestTokenCache(), + }) require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -855,7 +897,13 @@ func TestDiscoveryLogin_SPOGHostPopulatesAccountIDFromDiscovery(t *testing.T) { } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", nil, func(string) error { return nil }, newTestTokenCache(), "") + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + browserFunc: func(string) error { return nil }, + tokenCache: newTestTokenCache(), + }) require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -890,7 +938,13 @@ func TestDiscoveryLogin_IntrospectionFallsBackWhenDiscoveryFails(t *testing.T) { } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", nil, func(string) error { return nil }, newTestTokenCache(), "") + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + browserFunc: func(string) error { return nil }, + tokenCache: newTestTokenCache(), + }) require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -939,7 +993,14 @@ auth_type = databricks-cli } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }, newTestTokenCache(), "") + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + existingProfile: existingProfile, + browserFunc: func(string) error { return nil }, + tokenCache: newTestTokenCache(), + }) require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -989,7 +1050,14 @@ auth_type = databricks-cli } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }, newTestTokenCache(), "") + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + existingProfile: existingProfile, + browserFunc: func(string) error { return nil }, + tokenCache: newTestTokenCache(), + }) require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) From 09f6843e51964c34a0654f00bcfdc0737ba14929 Mon Sep 17 00:00:00 2001 From: simon Date: Tue, 21 Apr 2026 16:44:22 +0200 Subject: [PATCH 10/12] libs/auth: move cache resolver into storage package; route CLICredentials through it --- cmd/auth/login.go | 2 +- cmd/auth/logout.go | 3 +- cmd/auth/token.go | 2 +- libs/auth/credentials.go | 22 +++-- libs/auth/credentials_test.go | 97 +++++++++++++++++++ {cmd/auth => libs/auth/storage}/cache.go | 31 +++--- {cmd/auth => libs/auth/storage}/cache_test.go | 59 ++++++----- 7 files changed, 167 insertions(+), 49 deletions(-) rename {cmd/auth => libs/auth/storage}/cache.go (58%) rename {cmd/auth => libs/auth/storage}/cache_test.go (52%) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 34dca1bc30..a5a72b0a58 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -164,7 +164,7 @@ a new profile is created. ctx := cmd.Context() profileName := cmd.Flag("profile").Value.String() - tokenCache, mode, err := newAuthCache(ctx, "") + tokenCache, mode, err := storage.ResolveCache(ctx, "") if err != nil { return err } diff --git a/cmd/auth/logout.go b/cmd/auth/logout.go index 81c10a4f8b..bdd0f75430 100644 --- a/cmd/auth/logout.go +++ b/cmd/auth/logout.go @@ -22,6 +22,7 @@ import ( "strings" "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/auth/storage" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/databrickscfg" "github.com/databricks/cli/libs/databrickscfg/profile" @@ -132,7 +133,7 @@ to specify it explicitly. profileName = selected } - tokenCache, _, err := newAuthCache(ctx, "") + tokenCache, _, err := storage.ResolveCache(ctx, "") if err != nil { return fmt.Errorf("failed to open token cache: %w", err) } diff --git a/cmd/auth/token.go b/cmd/auth/token.go index ce1b0a04b3..da954b6318 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -79,7 +79,7 @@ and secret is not supported.`, ctx := cmd.Context() profileName := cmd.Flag("profile").Value.String() - tokenCache, mode, err := newAuthCache(ctx, "") + tokenCache, mode, err := storage.ResolveCache(ctx, "") if err != nil { return err } diff --git a/libs/auth/credentials.go b/libs/auth/credentials.go index a406955dd3..a756eeb462 100644 --- a/libs/auth/credentials.go +++ b/libs/auth/credentials.go @@ -109,22 +109,24 @@ func (c CLICredentials) Configure(ctx context.Context, cfg *config.Config) (cred return cp, nil } -// persistentAuth returns a token source. It wraps the file-backed token -// cache with a dual-writing cache so every token write (Challenge, refresh, -// discovery) mirrors to the legacy host key for cross-SDK compatibility. +// persistentAuth returns a token source. It resolves the configured storage +// mode and plumbs the resulting token cache into the SDK so every read is +// backed by the same storage the CLI wrote on login, not the SDK's default +// file cache (which would split tokens for secure-storage users). // The persistentAuthFn override is used in tests. func (c CLICredentials) persistentAuth(ctx context.Context, arg u2m.OAuthArgument) (auth.TokenSource, error) { - if c.persistentAuthFn != nil { - return c.persistentAuthFn(ctx, u2m.WithOAuthArgument(arg)) - } - tc, err := storage.NewFileTokenCache(ctx) + tokenCache, _, err := storage.ResolveCache(ctx, "") if err != nil { return nil, fmt.Errorf("opening token cache: %w", err) } - ts, err := u2m.NewPersistentAuth(ctx, - u2m.WithTokenCache(storage.NewDualWritingTokenCache(tc, arg)), + opts := []u2m.PersistentAuthOption{ u2m.WithOAuthArgument(arg), - ) + u2m.WithTokenCache(tokenCache), + } + if c.persistentAuthFn != nil { + return c.persistentAuthFn(ctx, opts...) + } + ts, err := u2m.NewPersistentAuth(ctx, opts...) if err != nil { return nil, err } diff --git a/libs/auth/credentials_test.go b/libs/auth/credentials_test.go index 1bc70b63ab..9412d16418 100644 --- a/libs/auth/credentials_test.go +++ b/libs/auth/credentials_test.go @@ -4,15 +4,28 @@ import ( "context" "errors" "net/http" + "os" + "path/filepath" "slices" "testing" + "github.com/databricks/cli/libs/auth/storage" "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/config/experimental/auth" "github.com/databricks/databricks-sdk-go/credentials/u2m" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "golang.org/x/oauth2" ) +// hermeticAuthStorage isolates the test from the caller's real env vars and +// .databrickscfg so storage.ResolveCache sees a clean default. +func hermeticAuthStorage(t *testing.T) { + t.Helper() + t.Setenv(storage.EnvVar, "") + t.Setenv("DATABRICKS_CONFIG_FILE", filepath.Join(t.TempDir(), "databrickscfg")) +} + // TestCredentialChainOrder purely exists as an extra measure to catch // accidental change in the ordering. func TestCredentialChainOrder(t *testing.T) { @@ -163,6 +176,7 @@ func TestCLICredentialsConfigure(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { + hermeticAuthStorage(t) ctx := t.Context() c := CLICredentials{persistentAuthFn: tt.persistentAuthFn} @@ -190,3 +204,86 @@ func TestCLICredentialsConfigure(t *testing.T) { }) } } + +// TestCLICredentialsConfigure_ThreadsResolvedTokenCache guards against a +// regression where Configure forgot to pass u2m.WithTokenCache. Without it, +// the SDK's NewPersistentAuth silently defaulted to the file cache, so users +// who opted into secure storage saw "cache: token not found" on every command +// other than auth login/token/logout. +func TestCLICredentialsConfigure_ThreadsResolvedTokenCache(t *testing.T) { + hermeticAuthStorage(t) + + var receivedOpts []u2m.PersistentAuthOption + c := CLICredentials{ + persistentAuthFn: func(_ context.Context, opts ...u2m.PersistentAuthOption) (auth.TokenSource, error) { + receivedOpts = opts + return auth.TokenSourceFn(func(_ context.Context) (*oauth2.Token, error) { + return &oauth2.Token{AccessToken: "tok"}, nil + }), nil + }, + } + + _, err := c.Configure(t.Context(), &config.Config{Host: "https://x.cloud.databricks.com"}) + require.NoError(t, err) + + // Two opts expected: WithOAuthArgument and WithTokenCache. The length + // check is the most resilient way to assert both were passed without + // poking at u2m's unexported state. + assert.Len(t, receivedOpts, 2) +} + +// TestCLICredentialsConfigure_PropagatesStorageResolutionError confirms +// Configure surfaces invalid DATABRICKS_AUTH_STORAGE values instead of +// silently falling back to the file cache. If Configure ever stops calling +// storage.ResolveCache, this test will catch it. +func TestCLICredentialsConfigure_PropagatesStorageResolutionError(t *testing.T) { + hermeticAuthStorage(t) + t.Setenv(storage.EnvVar, "bogus") + + c := CLICredentials{ + persistentAuthFn: func(_ context.Context, _ ...u2m.PersistentAuthOption) (auth.TokenSource, error) { + t.Fatal("persistentAuthFn must not be called when cache resolution fails") + return nil, nil + }, + } + + _, err := c.Configure(t.Context(), &config.Config{Host: "https://x.cloud.databricks.com"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "DATABRICKS_AUTH_STORAGE") +} + +// Writing a throwaway config file is verbose enough that future tests may +// want it too. Keeping the helper scoped here so it stays close to use. +func writeAuthStorageConfig(t *testing.T, mode string) { + t.Helper() + dir := t.TempDir() + configPath := filepath.Join(dir, "databrickscfg") + body := "[__settings__]\nauth_storage = " + mode + "\n" + require.NoError(t, os.WriteFile(configPath, []byte(body), 0o600)) + t.Setenv("DATABRICKS_CONFIG_FILE", configPath) + t.Setenv(storage.EnvVar, "") +} + +// TestCLICredentialsConfigure_HonorsConfigFileSecureMode proves that +// Configure picks up auth_storage = secure from .databrickscfg, not just +// from DATABRICKS_AUTH_STORAGE. Both sources flow through the same resolver, +// but the PR's user-facing docs promise both work and nothing was asserting +// that for this call site. +func TestCLICredentialsConfigure_HonorsConfigFileSecureMode(t *testing.T) { + writeAuthStorageConfig(t, "secure") + + c := CLICredentials{ + persistentAuthFn: func(_ context.Context, opts ...u2m.PersistentAuthOption) (auth.TokenSource, error) { + // The presence of the second opt is verified by the sibling + // test; here we just need Configure to succeed end-to-end when + // the config file selects secure storage. + assert.Len(t, opts, 2) + return auth.TokenSourceFn(func(_ context.Context) (*oauth2.Token, error) { + return &oauth2.Token{AccessToken: "tok"}, nil + }), nil + }, + } + + _, err := c.Configure(t.Context(), &config.Config{Host: "https://x.cloud.databricks.com"}) + require.NoError(t, err) +} diff --git a/cmd/auth/cache.go b/libs/auth/storage/cache.go similarity index 58% rename from cmd/auth/cache.go rename to libs/auth/storage/cache.go index 3487573a16..a0c3e9fa04 100644 --- a/cmd/auth/cache.go +++ b/libs/auth/storage/cache.go @@ -1,14 +1,13 @@ -package auth +package storage import ( "context" "fmt" - "github.com/databricks/cli/libs/auth/storage" "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" ) -// cacheFactories bundles the constructors newAuthCache depends on. Extracted +// cacheFactories bundles the constructors ResolveCache depends on. Extracted // so unit tests can inject stubs without hitting the real OS keyring or // filesystem. Production code uses defaultCacheFactories(). type cacheFactories struct { @@ -23,31 +22,35 @@ type cacheFactories struct { func defaultCacheFactories() cacheFactories { return cacheFactories{ newFile: func() (cache.TokenCache, error) { return cache.NewFileTokenCache() }, - newKeyring: storage.NewKeyringCache, + newKeyring: NewKeyringCache, } } -// newAuthCache resolves the storage mode for this invocation and returns +// ResolveCache resolves the storage mode for this invocation and returns // the corresponding token cache plus the resolved mode (so callers can log // or surface it). // -// override is usually the command-level flag value (e.g. the result of -// --secure-storage). Pass "" when the command has no flag. -func newAuthCache(ctx context.Context, override storage.StorageMode) (cache.TokenCache, storage.StorageMode, error) { - return newAuthCacheWith(ctx, override, defaultCacheFactories()) +// override is usually the command-level flag value. Pass "" when the command +// has no flag; precedence then falls through to env -> config -> default. +// +// Every CLI code path that calls u2m.NewPersistentAuth must route the result +// through u2m.WithTokenCache, otherwise the SDK defaults to the file cache +// and splits the user's tokens across two backends. +func ResolveCache(ctx context.Context, override StorageMode) (cache.TokenCache, StorageMode, error) { + return resolveCacheWith(ctx, override, defaultCacheFactories()) } -// newAuthCacheWith is the pure form of newAuthCache. It takes the factory +// resolveCacheWith is the pure form of ResolveCache. It takes the factory // set as a parameter so tests can inject stubs. -func newAuthCacheWith(ctx context.Context, override storage.StorageMode, f cacheFactories) (cache.TokenCache, storage.StorageMode, error) { - mode, err := storage.ResolveStorageMode(ctx, override) +func resolveCacheWith(ctx context.Context, override StorageMode, f cacheFactories) (cache.TokenCache, StorageMode, error) { + mode, err := ResolveStorageMode(ctx, override) if err != nil { return nil, "", err } switch mode { - case storage.StorageModeSecure: + case StorageModeSecure: return f.newKeyring(), mode, nil - case storage.StorageModeLegacy, storage.StorageModePlaintext: + case StorageModeLegacy, StorageModePlaintext: // Plaintext currently maps to the file cache; a dedicated // plaintext backend (no host-keyed dual-writes) is a follow-up. c, err := f.newFile() diff --git a/cmd/auth/cache_test.go b/libs/auth/storage/cache_test.go similarity index 52% rename from cmd/auth/cache_test.go rename to libs/auth/storage/cache_test.go index b8d30d2f23..58b0032e01 100644 --- a/cmd/auth/cache_test.go +++ b/libs/auth/storage/cache_test.go @@ -1,10 +1,10 @@ -package auth +package storage import ( "errors" + "path/filepath" "testing" - "github.com/databricks/cli/libs/auth/storage" "github.com/databricks/cli/libs/env" "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" "github.com/stretchr/testify/assert" @@ -27,65 +27,80 @@ func fakeFactories(t *testing.T) cacheFactories { } } -func TestNewAuthCache_DefaultsToLegacyFile(t *testing.T) { +// hermetic isolates the test from the caller's real env vars and +// .databrickscfg so ResolveStorageMode starts from a clean default. +func hermetic(t *testing.T) { + t.Helper() + t.Setenv(EnvVar, "") + t.Setenv("DATABRICKS_CONFIG_FILE", filepath.Join(t.TempDir(), "databrickscfg")) +} + +func TestResolveCache_DefaultsToLegacyFile(t *testing.T) { + hermetic(t) ctx := t.Context() - got, mode, err := newAuthCacheWith(ctx, "", fakeFactories(t)) + got, mode, err := resolveCacheWith(ctx, "", fakeFactories(t)) require.NoError(t, err) - assert.Equal(t, storage.StorageModeLegacy, mode) + assert.Equal(t, StorageModeLegacy, mode) assert.Equal(t, "file", got.(stubCache).source) } -func TestNewAuthCache_OverrideSecureUsesKeyring(t *testing.T) { +func TestResolveCache_OverrideSecureUsesKeyring(t *testing.T) { + hermetic(t) ctx := t.Context() - got, mode, err := newAuthCacheWith(ctx, storage.StorageModeSecure, fakeFactories(t)) + got, mode, err := resolveCacheWith(ctx, StorageModeSecure, fakeFactories(t)) require.NoError(t, err) - assert.Equal(t, storage.StorageModeSecure, mode) + assert.Equal(t, StorageModeSecure, mode) assert.Equal(t, "keyring", got.(stubCache).source) } -func TestNewAuthCache_EnvVarSelectsSecure(t *testing.T) { - ctx := env.Set(t.Context(), storage.EnvVar, "secure") +func TestResolveCache_EnvVarSelectsSecure(t *testing.T) { + hermetic(t) + ctx := env.Set(t.Context(), EnvVar, "secure") - got, mode, err := newAuthCacheWith(ctx, "", fakeFactories(t)) + got, mode, err := resolveCacheWith(ctx, "", fakeFactories(t)) require.NoError(t, err) - assert.Equal(t, storage.StorageModeSecure, mode) + assert.Equal(t, StorageModeSecure, mode) assert.Equal(t, "keyring", got.(stubCache).source) } -func TestNewAuthCache_PlaintextFallsBackToFile(t *testing.T) { +func TestResolveCache_PlaintextFallsBackToFile(t *testing.T) { + hermetic(t) ctx := t.Context() - got, mode, err := newAuthCacheWith(ctx, storage.StorageModePlaintext, fakeFactories(t)) + got, mode, err := resolveCacheWith(ctx, StorageModePlaintext, fakeFactories(t)) require.NoError(t, err) - assert.Equal(t, storage.StorageModePlaintext, mode) + assert.Equal(t, StorageModePlaintext, mode) assert.Equal(t, "file", got.(stubCache).source) } -func TestNewAuthCache_InvalidOverrideReturnsError(t *testing.T) { +func TestResolveCache_InvalidOverrideReturnsError(t *testing.T) { + hermetic(t) ctx := t.Context() - _, _, err := newAuthCacheWith(ctx, storage.StorageMode("bogus"), fakeFactories(t)) + _, _, err := resolveCacheWith(ctx, StorageMode("bogus"), fakeFactories(t)) require.Error(t, err) assert.Contains(t, err.Error(), `unknown storage mode "bogus"`) } -func TestNewAuthCache_InvalidEnvReturnsError(t *testing.T) { - ctx := env.Set(t.Context(), storage.EnvVar, "bogus") +func TestResolveCache_InvalidEnvReturnsError(t *testing.T) { + hermetic(t) + ctx := env.Set(t.Context(), EnvVar, "bogus") - _, _, err := newAuthCacheWith(ctx, "", fakeFactories(t)) + _, _, err := resolveCacheWith(ctx, "", fakeFactories(t)) require.Error(t, err) assert.Contains(t, err.Error(), "DATABRICKS_AUTH_STORAGE") } -func TestNewAuthCache_FileFactoryErrorPropagates(t *testing.T) { +func TestResolveCache_FileFactoryErrorPropagates(t *testing.T) { + hermetic(t) ctx := t.Context() boom := errors.New("disk full") factories := cacheFactories{ @@ -93,7 +108,7 @@ func TestNewAuthCache_FileFactoryErrorPropagates(t *testing.T) { newKeyring: func() cache.TokenCache { return stubCache{source: "keyring"} }, } - _, _, err := newAuthCacheWith(ctx, storage.StorageModeLegacy, factories) + _, _, err := resolveCacheWith(ctx, StorageModeLegacy, factories) require.Error(t, err) assert.ErrorIs(t, err, boom) From 5e8034bd62b9390c99967547021992eb2b8626ef Mon Sep 17 00:00:00 2001 From: simon Date: Wed, 22 Apr 2026 11:33:28 +0200 Subject: [PATCH 11/12] libs/auth/storage: use CLI's FileTokenCache in ResolveCache The file factory was calling the SDK's cache.NewFileTokenCache(), which is being phased out in favor of the CLI's own storage.NewFileTokenCache(ctx) imported in #5056. Route ResolveCache through it so that legacy and plaintext modes share a single file cache implementation owned by the CLI. Co-authored-by: Isaac --- libs/auth/storage/cache.go | 9 +++------ libs/auth/storage/cache_test.go | 5 +++-- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/libs/auth/storage/cache.go b/libs/auth/storage/cache.go index a0c3e9fa04..7a8bb775ab 100644 --- a/libs/auth/storage/cache.go +++ b/libs/auth/storage/cache.go @@ -11,17 +11,14 @@ import ( // so unit tests can inject stubs without hitting the real OS keyring or // filesystem. Production code uses defaultCacheFactories(). type cacheFactories struct { - newFile func() (cache.TokenCache, error) + newFile func(context.Context) (cache.TokenCache, error) newKeyring func() cache.TokenCache } // defaultCacheFactories returns the production factory set. -// newFile is wrapped in a closure because cache.NewFileTokenCache is variadic -// (func(...FileTokenCacheOption)) and cannot satisfy the non-variadic field type -// by direct reference. The closure calls it with no options (SDK defaults). func defaultCacheFactories() cacheFactories { return cacheFactories{ - newFile: func() (cache.TokenCache, error) { return cache.NewFileTokenCache() }, + newFile: func(ctx context.Context) (cache.TokenCache, error) { return NewFileTokenCache(ctx) }, newKeyring: NewKeyringCache, } } @@ -53,7 +50,7 @@ func resolveCacheWith(ctx context.Context, override StorageMode, f cacheFactorie case StorageModeLegacy, StorageModePlaintext: // Plaintext currently maps to the file cache; a dedicated // plaintext backend (no host-keyed dual-writes) is a follow-up. - c, err := f.newFile() + c, err := f.newFile(ctx) if err != nil { return nil, "", fmt.Errorf("open file token cache: %w", err) } diff --git a/libs/auth/storage/cache_test.go b/libs/auth/storage/cache_test.go index 58b0032e01..6a967891b9 100644 --- a/libs/auth/storage/cache_test.go +++ b/libs/auth/storage/cache_test.go @@ -1,6 +1,7 @@ package storage import ( + "context" "errors" "path/filepath" "testing" @@ -22,7 +23,7 @@ func (stubCache) Lookup(string) (*oauth2.Token, error) { return nil, cache.ErrNo func fakeFactories(t *testing.T) cacheFactories { t.Helper() return cacheFactories{ - newFile: func() (cache.TokenCache, error) { return stubCache{source: "file"}, nil }, + newFile: func(context.Context) (cache.TokenCache, error) { return stubCache{source: "file"}, nil }, newKeyring: func() cache.TokenCache { return stubCache{source: "keyring"} }, } } @@ -104,7 +105,7 @@ func TestResolveCache_FileFactoryErrorPropagates(t *testing.T) { ctx := t.Context() boom := errors.New("disk full") factories := cacheFactories{ - newFile: func() (cache.TokenCache, error) { return nil, boom }, + newFile: func(context.Context) (cache.TokenCache, error) { return nil, boom }, newKeyring: func() cache.TokenCache { return stubCache{source: "keyring"} }, } From 4d6f006a0057630f1cc6e01118d2f60808ec8466 Mon Sep 17 00:00:00 2001 From: simon Date: Wed, 22 Apr 2026 11:44:17 +0200 Subject: [PATCH 12/12] libs/auth/credentials: drop redundant file cache in persistentAuth Configure already resolves the correct cache via storage.ResolveCache and passes it via u2m.WithTokenCache. The persistentAuth helper opened a second storage.FileTokenCache and prepended its own WithTokenCache to opts. Because u2m options are last-one-wins, the keyring still won, but we were opening a file cache on every workspace client just to have it immediately overridden. Remove the hardcoded file cache and rely on the caller to supply one. Now there is a single place that manages the token cache (ResolveCache in Configure) and the helper is a thin wrapper around u2m.NewPersistentAuth. Co-authored-by: Isaac --- libs/auth/credentials.go | 36 ++++++++++++++++++++---------------- 1 file changed, 20 insertions(+), 16 deletions(-) diff --git a/libs/auth/credentials.go b/libs/auth/credentials.go index a756eeb462..6b95177353 100644 --- a/libs/auth/credentials.go +++ b/libs/auth/credentials.go @@ -3,7 +3,6 @@ package auth import ( "context" "errors" - "fmt" "github.com/databricks/cli/libs/auth/storage" "github.com/databricks/databricks-sdk-go/config" @@ -99,7 +98,19 @@ func (c CLICredentials) Configure(ctx context.Context, cfg *config.Config) (cred if err != nil { return nil, err } - ts, err := c.persistentAuth(ctx, oauthArg) + // Without WithTokenCache, u2m.NewPersistentAuth falls back to the SDK's + // default file cache. For secure-storage users that would split tokens + // across two backends: login writes to the keyring, but every workspace + // client built through this strategy would read an empty file cache and + // fail with "cache: token not found". + tokenCache, _, err := storage.ResolveCache(ctx, "") + if err != nil { + return nil, err + } + ts, err := c.persistentAuth(ctx, + u2m.WithOAuthArgument(oauthArg), + u2m.WithTokenCache(tokenCache), + ) if err != nil { return nil, err } @@ -109,20 +120,13 @@ func (c CLICredentials) Configure(ctx context.Context, cfg *config.Config) (cred return cp, nil } -// persistentAuth returns a token source. It resolves the configured storage -// mode and plumbs the resulting token cache into the SDK so every read is -// backed by the same storage the CLI wrote on login, not the SDK's default -// file cache (which would split tokens for secure-storage users). -// The persistentAuthFn override is used in tests. -func (c CLICredentials) persistentAuth(ctx context.Context, arg u2m.OAuthArgument) (auth.TokenSource, error) { - tokenCache, _, err := storage.ResolveCache(ctx, "") - if err != nil { - return nil, fmt.Errorf("opening token cache: %w", err) - } - opts := []u2m.PersistentAuthOption{ - u2m.WithOAuthArgument(arg), - u2m.WithTokenCache(tokenCache), - } +// persistentAuth returns a token source. It is a convenience function that +// overrides the default implementation of the persistent auth client if +// an alternative implementation is provided for testing. The caller is +// responsible for supplying the token cache via u2m.WithTokenCache; Configure +// does this via storage.ResolveCache so login, refresh, and all workspace +// clients share the same backend. +func (c CLICredentials) persistentAuth(ctx context.Context, opts ...u2m.PersistentAuthOption) (auth.TokenSource, error) { if c.persistentAuthFn != nil { return c.persistentAuthFn(ctx, opts...) }