From 274a09e652c05153aefe16f27c2295029102d823 Mon Sep 17 00:00:00 2001 From: simon Date: Tue, 21 Apr 2026 18:24:54 +0200 Subject: [PATCH 1/4] auth: import FileTokenCache into CLI and wire DualWrite First of a three-PR sequence that moves file-based OAuth token cache management from the SDK to the CLI. This PR adds a CLI-local copy of FileTokenCache under libs/auth/storage, a DualWrite helper that mirrors the SDK's historical dualWrite + hostCacheKey convention, wires u2m.WithTokenCache at every NewPersistentAuth call site (including CLICredentials, which is used by every non-auth command), and switches auth logout to the CLI FileTokenCache for token removal. The SDK is unchanged so behavior is byte-for-byte identical: two redundant writes to the same file with the same keys and tokens. See documents/fy2027-q2/cli-ga/2026-04-21-move-token-cache-to-cli-plan.md. --- NEXT_CHANGELOG.md | 1 + cmd/auth/login.go | 21 +++ cmd/auth/logout.go | 3 +- cmd/auth/token.go | 19 ++- libs/auth/credentials.go | 7 + libs/auth/storage/dualwrite.go | 39 +++++ libs/auth/storage/dualwrite_test.go | 133 +++++++++++++++++ libs/auth/storage/file_cache.go | 214 +++++++++++++++++++++++++++ libs/auth/storage/file_cache_test.go | 67 +++++++++ 9 files changed, 502 insertions(+), 2 deletions(-) create mode 100644 libs/auth/storage/dualwrite.go create mode 100644 libs/auth/storage/dualwrite_test.go create mode 100644 libs/auth/storage/file_cache.go create mode 100644 libs/auth/storage/file_cache_test.go diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 16a227a237..5c924f8f73 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -9,6 +9,7 @@ * Added `--limit` flag to all paginated list commands for client-side result capping ([#4984](https://github.com/databricks/cli/pull/4984)). * 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. ### Bundles * Remove `experimental-jobs-as-code` template, superseded by `pydabs` ([#4999](https://github.com/databricks/cli/pull/4999)). diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 7157f797b7..142cfd83ba 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -10,6 +10,7 @@ import ( "time" "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/auth/storage" "github.com/databricks/cli/libs/browser" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/databrickscfg" @@ -227,7 +228,12 @@ a new profile is created. if err != nil { return err } + tc, err := storage.NewFileTokenCache() + if err != nil { + return fmt.Errorf("opening token cache: %w", err) + } persistentAuthOpts := []u2m.PersistentAuthOption{ + u2m.WithTokenCache(tc), u2m.WithOAuthArgument(oauthArgument), u2m.WithBrowser(getBrowserFunc(cmd)), } @@ -246,6 +252,11 @@ a new profile is created. if err = persistentAuth.Challenge(); err != nil { return err } + if t, lookupErr := tc.Lookup(oauthArgument.GetCacheKey()); lookupErr == nil && t != nil { + if err := storage.DualWrite(tc, oauthArgument, t); err != nil { + log.Debugf(ctx, "token cache dual-write failed: %v", err) + } + } // 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); @@ -573,7 +584,12 @@ func discoveryLogin(ctx context.Context, dc discoveryClient, profileName string, scopesList = splitScopes(existingProfile.Scopes) } + tc, err := storage.NewFileTokenCache() + if err != nil { + return discoveryErr("opening token cache", err) + } opts := []u2m.PersistentAuthOption{ + u2m.WithTokenCache(tc), u2m.WithOAuthArgument(arg), u2m.WithBrowser(browserFunc), u2m.WithDiscoveryLogin(), @@ -596,6 +612,11 @@ func discoveryLogin(ctx context.Context, dc discoveryClient, profileName string, if err := persistentAuth.Challenge(); err != nil { return discoveryErr("login via login.databricks.com failed", err) } + if t, lookupErr := tc.Lookup(arg.GetCacheKey()); lookupErr == nil && t != nil { + if err := storage.DualWrite(tc, arg, t); err != nil { + log.Debugf(ctx, "token cache dual-write failed: %v", err) + } + } discoveredHost := arg.GetDiscoveredHost() if discoveredHost == "" { diff --git a/cmd/auth/logout.go b/cmd/auth/logout.go index 3beeeefec9..078ed91d6b 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 := cache.NewFileTokenCache() + tokenCache, err := storage.NewFileTokenCache() 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) } diff --git a/cmd/auth/token.go b/cmd/auth/token.go index fbdd8811e8..86ee92bf46 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -11,12 +11,14 @@ import ( "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/auth/storage" "github.com/databricks/cli/libs/browser" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/databrickscfg" "github.com/databricks/cli/libs/databrickscfg/profile" "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/flags" + "github.com/databricks/cli/libs/log" "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" @@ -273,7 +275,12 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) { if err != nil { return nil, err } - allArgs := append(args.persistentAuthOpts, u2m.WithOAuthArgument(oauthArgument)) + tc, err := storage.NewFileTokenCache() + if err != nil { + return nil, fmt.Errorf("opening token cache: %w", err) + } + allArgs := append([]u2m.PersistentAuthOption{u2m.WithTokenCache(tc)}, args.persistentAuthOpts...) + allArgs = append(allArgs, u2m.WithOAuthArgument(oauthArgument)) persistentAuth, err := u2m.NewPersistentAuth(ctx, allArgs...) if err != nil { helpMsg := helpfulError(ctx, args.profileName, oauthArgument) @@ -459,7 +466,12 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler) (string, *pr if err != nil { return "", nil, err } + tc, err := storage.NewFileTokenCache() + if err != nil { + return "", nil, fmt.Errorf("opening token cache: %w", err) + } persistentAuthOpts := []u2m.PersistentAuthOption{ + u2m.WithTokenCache(tc), u2m.WithOAuthArgument(oauthArgument), u2m.WithBrowser(func(url string) error { return browser.Open(ctx, url) }), } @@ -478,6 +490,11 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler) (string, *pr if err = persistentAuth.Challenge(); err != nil { return "", nil, err } + if t, lookupErr := tc.Lookup(oauthArgument.GetCacheKey()); lookupErr == nil && t != nil { + if err := storage.DualWrite(tc, oauthArgument, t); err != nil { + log.Debugf(ctx, "token cache dual-write failed: %v", err) + } + } clearKeys := oauthLoginClearKeys() if !loginArgs.IsUnifiedHost { diff --git a/libs/auth/credentials.go b/libs/auth/credentials.go index 7ab6eb2a85..64fa810b75 100644 --- a/libs/auth/credentials.go +++ b/libs/auth/credentials.go @@ -3,7 +3,9 @@ package auth import ( "context" "errors" + "fmt" + "github.com/databricks/cli/libs/auth/storage" "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/config/credentials" "github.com/databricks/databricks-sdk-go/config/experimental/auth" @@ -114,6 +116,11 @@ func (c CLICredentials) persistentAuth(ctx context.Context, opts ...u2m.Persiste if c.persistentAuthFn != nil { return c.persistentAuthFn(ctx, opts...) } + tc, err := storage.NewFileTokenCache() + if err != nil { + return nil, fmt.Errorf("opening token cache: %w", err) + } + opts = append([]u2m.PersistentAuthOption{u2m.WithTokenCache(tc)}, opts...) ts, err := u2m.NewPersistentAuth(ctx, opts...) if err != nil { return nil, err diff --git a/libs/auth/storage/dualwrite.go b/libs/auth/storage/dualwrite.go new file mode 100644 index 0000000000..14ce017e23 --- /dev/null +++ b/libs/auth/storage/dualwrite.go @@ -0,0 +1,39 @@ +package storage + +import ( + "github.com/databricks/databricks-sdk-go/credentials/u2m" + u2m_cache "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" + "golang.org/x/oauth2" +) + +// DualWrite stores t under arg.GetCacheKey() (the primary key) and, if a +// legacy host-based cache key can be derived from arg, also stores t under +// that host key. Mirrors the convention used historically by +// PersistentAuth.dualWrite and hostCacheKey in the SDK. +func DualWrite(cache u2m_cache.TokenCache, arg u2m.OAuthArgument, t *oauth2.Token) error { + primaryKey := arg.GetCacheKey() + if err := cache.Store(primaryKey, t); err != nil { + return err + } + hostKey := hostCacheKey(arg) + if hostKey != "" && hostKey != primaryKey { + if err := cache.Store(hostKey, t); err != nil { + return err + } + } + return nil +} + +// hostCacheKey mirrors PersistentAuth.hostCacheKey in the SDK: discovery +// arguments learn their host from the OAuth callback, so their host key +// must be read via GetDiscoveredHost *after* Challenge(). Static arguments +// use HostCacheKeyProvider. +func hostCacheKey(arg u2m.OAuthArgument) string { + if discoveryArg, ok := arg.(u2m.DiscoveryOAuthArgument); ok { + return discoveryArg.GetDiscoveredHost() + } + if hcp, ok := arg.(u2m.HostCacheKeyProvider); ok { + return hcp.GetHostCacheKey() + } + return "" +} diff --git a/libs/auth/storage/dualwrite_test.go b/libs/auth/storage/dualwrite_test.go new file mode 100644 index 0000000000..f3da4ae8d0 --- /dev/null +++ b/libs/auth/storage/dualwrite_test.go @@ -0,0 +1,133 @@ +package storage + +import ( + "sync" + "testing" + + "github.com/databricks/databricks-sdk-go/credentials/u2m" + u2m_cache "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/oauth2" +) + +// memoryCache is a minimal in-memory TokenCache used only by DualWrite tests. +type memoryCache struct { + mu sync.Mutex + tokens map[string]*oauth2.Token +} + +func newMemoryCache() *memoryCache { + return &memoryCache{tokens: map[string]*oauth2.Token{}} +} + +func (c *memoryCache) Store(key string, t *oauth2.Token) error { + c.mu.Lock() + defer c.mu.Unlock() + if t == nil { + delete(c.tokens, key) + return nil + } + c.tokens[key] = t + return nil +} + +func (c *memoryCache) Lookup(key string) (*oauth2.Token, error) { + c.mu.Lock() + defer c.mu.Unlock() + t, ok := c.tokens[key] + if !ok { + return nil, u2m_cache.ErrNotFound + } + return t, nil +} + +// plainArg implements OAuthArgument only, exercising the "no host key" branch. +type plainArg struct { + key string +} + +func (a plainArg) GetCacheKey() string { return a.key } + +// hostArg implements HostCacheKeyProvider so DualWrite mirrors the token to +// the configured host key. +type hostArg struct { + key string + hostKey string +} + +func (a hostArg) GetCacheKey() string { return a.key } +func (a hostArg) GetHostCacheKey() string { return a.hostKey } + +func TestDualWriteNoHostKey(t *testing.T) { + cache := newMemoryCache() + arg := plainArg{key: "profile-a"} + tok := &oauth2.Token{AccessToken: "abc", RefreshToken: "r"} + + require.NoError(t, DualWrite(cache, arg, tok)) + + got, err := cache.Lookup("profile-a") + require.NoError(t, err) + assert.Equal(t, tok, got) + assert.Len(t, cache.tokens, 1) +} + +func TestDualWriteHostKeyDistinct(t *testing.T) { + cache := newMemoryCache() + arg := hostArg{key: "profile-a", hostKey: "https://example.databricks.com"} + tok := &oauth2.Token{AccessToken: "abc", RefreshToken: "r"} + + require.NoError(t, DualWrite(cache, arg, tok)) + + primary, err := cache.Lookup("profile-a") + require.NoError(t, err) + assert.Equal(t, tok, primary) + + host, err := cache.Lookup("https://example.databricks.com") + require.NoError(t, err) + assert.Equal(t, tok, host) + + assert.Len(t, cache.tokens, 2) +} + +func TestDualWriteHostKeyEqualsPrimary(t *testing.T) { + cache := newMemoryCache() + arg := hostArg{key: "https://example.databricks.com", hostKey: "https://example.databricks.com"} + tok := &oauth2.Token{AccessToken: "abc"} + + require.NoError(t, DualWrite(cache, arg, tok)) + + assert.Len(t, cache.tokens, 1) +} + +func TestDualWriteDiscoveryArgWithDiscoveredHost(t *testing.T) { + cache := newMemoryCache() + arg, err := u2m.NewBasicDiscoveryOAuthArgument("profile-a") + require.NoError(t, err) + arg.SetDiscoveredHost("https://example.databricks.com") + tok := &oauth2.Token{AccessToken: "abc"} + + require.NoError(t, DualWrite(cache, arg, tok)) + + primary, err := cache.Lookup("profile-a") + require.NoError(t, err) + assert.Equal(t, tok, primary) + + host, err := cache.Lookup("https://example.databricks.com") + require.NoError(t, err) + assert.Equal(t, tok, host) +} + +func TestDualWriteDiscoveryArgWithEmptyDiscoveredHost(t *testing.T) { + cache := newMemoryCache() + arg, err := u2m.NewBasicDiscoveryOAuthArgument("profile-a") + require.NoError(t, err) + tok := &oauth2.Token{AccessToken: "abc"} + + require.NoError(t, DualWrite(cache, arg, tok)) + + assert.Len(t, cache.tokens, 1) + primary, err := cache.Lookup("profile-a") + require.NoError(t, err) + assert.Equal(t, tok, primary) +} diff --git a/libs/auth/storage/file_cache.go b/libs/auth/storage/file_cache.go new file mode 100644 index 0000000000..d782fa73dc --- /dev/null +++ b/libs/auth/storage/file_cache.go @@ -0,0 +1,214 @@ +package storage + +import ( + "encoding/json" + "fmt" + "os" + "path/filepath" + "sync" + + u2m_cache "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" + "golang.org/x/oauth2" +) + +const ( + // tokenCacheFile is the path of the default token cache, relative to the + // user's home directory. + tokenCacheFilePath = ".databricks/token-cache.json" + + // ownerExecReadWrite is the permission for the .databricks directory. + ownerExecReadWrite = 0o700 + + // ownerReadWrite is the permission for the token-cache.json file. + ownerReadWrite = 0o600 + + // tokenCacheVersion is the version of the token cache file format. + // + // Version 1 format: + // + // { + // "version": 1, + // "tokens": { + // "": { + // "access_token": "", + // "token_type": "", + // "refresh_token": "", + // "expiry": "" + // } + // } + // } + tokenCacheVersion = 1 +) + +// tokenCacheFile is the format of the token cache file. +type tokenCacheFile struct { + Version int `json:"version"` + Tokens map[string]*oauth2.Token `json:"tokens"` +} + +type FileTokenCacheOption func(*fileTokenCache) + +func WithFileLocation(fileLocation string) FileTokenCacheOption { + return func(c *fileTokenCache) { + c.fileLocation = fileLocation + } +} + +// fileTokenCache caches tokens in "~/.databricks/token-cache.json". fileTokenCache +// implements the TokenCache interface. +type fileTokenCache struct { + fileLocation string + + // locker protects the token cache file from concurrent reads and writes. + locker sync.Mutex +} + +// NewFileTokenCache creates a new FileTokenCache. By default, the cache is +// stored in "~/.databricks/token-cache.json". The cache file is created if it +// does not already exist. The cache file is created with owner permissions +// 0600 and the directory is created with owner permissions 0700. If the cache +// file is corrupt or if its version does not match tokenCacheVersion, an error +// is returned. +func NewFileTokenCache(opts ...FileTokenCacheOption) (u2m_cache.TokenCache, error) { + c := &fileTokenCache{} + for _, opt := range opts { + opt(c) + } + if err := c.init(); err != nil { + return nil, err + } + // Fail fast if the cache is not working. + if _, err := c.load(); err != nil { + return nil, fmt.Errorf("load: %w", err) + } + return c, nil +} + +// Store implements the TokenCache interface. +func (c *fileTokenCache) Store(key string, t *oauth2.Token) error { + c.locker.Lock() + defer c.locker.Unlock() + f, err := c.load() + if err != nil { + return fmt.Errorf("load: %w", err) + } + if f.Tokens == nil { + f.Tokens = map[string]*oauth2.Token{} + } + if t == nil { + delete(f.Tokens, key) + } else { + f.Tokens[key] = t + } + raw, err := json.MarshalIndent(f, "", " ") + if err != nil { + return fmt.Errorf("marshal: %w", err) + } + if err := c.atomicWriteFile(raw); err != nil { + return fmt.Errorf("error storing token in local cache: %w", err) + } + return nil +} + +// Lookup implements the TokenCache interface. +func (c *fileTokenCache) Lookup(key string) (*oauth2.Token, error) { + c.locker.Lock() + defer c.locker.Unlock() + f, err := c.load() + if err != nil { + return nil, fmt.Errorf("load: %w", err) + } + t, ok := f.Tokens[key] + if !ok { + return nil, u2m_cache.ErrNotFound + } + return t, nil +} + +// init initializes the token cache file. It creates the file and directory if +// they do not already exist. +func (c *fileTokenCache) init() error { + // set the default file location + if c.fileLocation == "" { + home, err := os.UserHomeDir() + if err != nil { + return fmt.Errorf("failed loading home directory: %w", err) + } + c.fileLocation = filepath.Join(home, tokenCacheFilePath) + } + // Create the cache file if it does not exist. + if _, err := os.Stat(c.fileLocation); err != nil { + if !os.IsNotExist(err) { + return fmt.Errorf("stat file: %w", err) + } + // Create the parent directories if needed. + if err := os.MkdirAll(filepath.Dir(c.fileLocation), ownerExecReadWrite); err != nil { + return fmt.Errorf("mkdir: %w", err) + } + + // Create an empty cache file. + f := &tokenCacheFile{ + Version: tokenCacheVersion, + Tokens: map[string]*oauth2.Token{}, + } + raw, err := json.MarshalIndent(f, "", " ") + if err != nil { + return fmt.Errorf("marshal: %w", err) + } + if err := c.atomicWriteFile(raw); err != nil { + return fmt.Errorf("error creating token cache file: %w", err) + } + } + return nil +} + +// load loads the token cache file from disk. If the file is corrupt or if its +// version does not match tokenCacheVersion, it returns an error. +func (c *fileTokenCache) load() (*tokenCacheFile, error) { + raw, err := os.ReadFile(c.fileLocation) + if err != nil { + return nil, fmt.Errorf("read: %w", err) + } + f := &tokenCacheFile{} + if err := json.Unmarshal(raw, &f); err != nil { + return nil, fmt.Errorf("parse: %w", err) + } + if f.Version != tokenCacheVersion { + // in the later iterations we could do state upgraders, + // so that we transform token cache from v1 to v2 without + // losing the tokens and asking the user to re-authenticate. + return nil, fmt.Errorf("needs version %d, got version %d", tokenCacheVersion, f.Version) + } + return f, nil +} + +// atomicWriteFile writes data to the file atomically by first writing to a +// temporary file in the same directory and then renaming it to the target. +// This prevents corruption from interrupted writes. +func (c *fileTokenCache) atomicWriteFile(data []byte) error { + tmp, err := c.writeTmpFile(data) + if err != nil { + return err + } + defer os.Remove(tmp) + return os.Rename(tmp, c.fileLocation) +} + +func (c *fileTokenCache) writeTmpFile(data []byte) (string, error) { + tmp, err := os.CreateTemp(filepath.Dir(c.fileLocation), ".token-cache-*.tmp") + if err != nil { + return "", fmt.Errorf("create temp file: %w", err) + } + defer tmp.Close() + + if _, err := tmp.Write(data); err != nil { + return "", err + } + if err := tmp.Chmod(ownerReadWrite); err != nil { + return "", err + } + if err := tmp.Close(); err != nil { + return "", err + } + return tmp.Name(), nil +} diff --git a/libs/auth/storage/file_cache_test.go b/libs/auth/storage/file_cache_test.go new file mode 100644 index 0000000000..fdec3c5c46 --- /dev/null +++ b/libs/auth/storage/file_cache_test.go @@ -0,0 +1,67 @@ +package storage + +import ( + "os" + "path/filepath" + "testing" + + u2m_cache "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/oauth2" +) + +func setup(t *testing.T) string { + tempHomeDir := t.TempDir() + return filepath.Join(tempHomeDir, "token-cache.json") +} + +func TestStoreAndLookup(t *testing.T) { + c, err := NewFileTokenCache(WithFileLocation(setup(t))) + require.NoError(t, err) + err = c.Store("x", &oauth2.Token{ + AccessToken: "abc", + }) + require.NoError(t, err) + + err = c.Store("y", &oauth2.Token{ + AccessToken: "bcd", + }) + require.NoError(t, err) + + tok, err := c.Lookup("x") + require.NoError(t, err) + assert.Equal(t, "abc", tok.AccessToken) + + _, err = c.Lookup("z") + assert.Equal(t, u2m_cache.ErrNotFound, err) +} + +func TestNoCacheFileReturnsErrNotConfigured(t *testing.T) { + l, err := NewFileTokenCache(WithFileLocation(setup(t))) + require.NoError(t, err) + _, err = l.Lookup("x") + assert.Equal(t, u2m_cache.ErrNotFound, err) +} + +func TestLoadCorruptFile(t *testing.T) { + f := setup(t) + err := os.MkdirAll(filepath.Dir(f), ownerExecReadWrite) + require.NoError(t, err) + err = os.WriteFile(f, []byte("abc"), ownerExecReadWrite) + require.NoError(t, err) + + _, err = NewFileTokenCache(WithFileLocation(f)) + assert.EqualError(t, err, "load: parse: invalid character 'a' looking for beginning of value") +} + +func TestLoadWrongVersion(t *testing.T) { + f := setup(t) + err := os.MkdirAll(filepath.Dir(f), ownerExecReadWrite) + require.NoError(t, err) + err = os.WriteFile(f, []byte(`{"version": 823, "things": []}`), ownerExecReadWrite) + require.NoError(t, err) + + _, err = NewFileTokenCache(WithFileLocation(f)) + assert.EqualError(t, err, "load: needs version 1, got version 823") +} From ba0a2d4e894d756cf448633e5b22056c9ffe4e4d Mon Sep 17 00:00:00 2001 From: simon Date: Tue, 21 Apr 2026 20:30:38 +0200 Subject: [PATCH 2/4] Fix lint: use env.UserHomeDir and errors.Is(fs.ErrNotExist) CLI's forbidigo rules forbid os.UserHomeDir (use env.UserHomeDir) and os.IsNotExist (use errors.Is(err, fs.ErrNotExist)). Thread ctx through NewFileTokenCache so the env-based home directory lookup works. Co-authored-by: Isaac --- cmd/auth/login.go | 4 ++-- cmd/auth/logout.go | 2 +- cmd/auth/token.go | 4 ++-- libs/auth/credentials.go | 2 +- libs/auth/storage/file_cache.go | 14 +++++++++----- libs/auth/storage/file_cache_test.go | 8 ++++---- 6 files changed, 19 insertions(+), 15 deletions(-) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 142cfd83ba..adc3db3a64 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -228,7 +228,7 @@ a new profile is created. if err != nil { return err } - tc, err := storage.NewFileTokenCache() + tc, err := storage.NewFileTokenCache(ctx) if err != nil { return fmt.Errorf("opening token cache: %w", err) } @@ -584,7 +584,7 @@ func discoveryLogin(ctx context.Context, dc discoveryClient, profileName string, scopesList = splitScopes(existingProfile.Scopes) } - tc, err := storage.NewFileTokenCache() + tc, err := storage.NewFileTokenCache(ctx) if err != nil { return discoveryErr("opening token cache", err) } diff --git a/cmd/auth/logout.go b/cmd/auth/logout.go index 078ed91d6b..67829ec169 100644 --- a/cmd/auth/logout.go +++ b/cmd/auth/logout.go @@ -133,7 +133,7 @@ to specify it explicitly. profileName = selected } - tokenCache, err := storage.NewFileTokenCache() + tokenCache, err := storage.NewFileTokenCache(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) } diff --git a/cmd/auth/token.go b/cmd/auth/token.go index 86ee92bf46..d27e789b0d 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -275,7 +275,7 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) { if err != nil { return nil, err } - tc, err := storage.NewFileTokenCache() + tc, err := storage.NewFileTokenCache(ctx) if err != nil { return nil, fmt.Errorf("opening token cache: %w", err) } @@ -466,7 +466,7 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler) (string, *pr if err != nil { return "", nil, err } - tc, err := storage.NewFileTokenCache() + tc, err := storage.NewFileTokenCache(ctx) if err != nil { return "", nil, fmt.Errorf("opening token cache: %w", err) } diff --git a/libs/auth/credentials.go b/libs/auth/credentials.go index 64fa810b75..c3fb115b24 100644 --- a/libs/auth/credentials.go +++ b/libs/auth/credentials.go @@ -116,7 +116,7 @@ func (c CLICredentials) persistentAuth(ctx context.Context, opts ...u2m.Persiste if c.persistentAuthFn != nil { return c.persistentAuthFn(ctx, opts...) } - tc, err := storage.NewFileTokenCache() + tc, err := storage.NewFileTokenCache(ctx) if err != nil { return nil, fmt.Errorf("opening token cache: %w", err) } diff --git a/libs/auth/storage/file_cache.go b/libs/auth/storage/file_cache.go index d782fa73dc..f64e233b01 100644 --- a/libs/auth/storage/file_cache.go +++ b/libs/auth/storage/file_cache.go @@ -1,12 +1,16 @@ package storage import ( + "context" "encoding/json" + "errors" "fmt" + "io/fs" "os" "path/filepath" "sync" + "github.com/databricks/cli/libs/env" u2m_cache "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" "golang.org/x/oauth2" ) @@ -69,12 +73,12 @@ type fileTokenCache struct { // 0600 and the directory is created with owner permissions 0700. If the cache // file is corrupt or if its version does not match tokenCacheVersion, an error // is returned. -func NewFileTokenCache(opts ...FileTokenCacheOption) (u2m_cache.TokenCache, error) { +func NewFileTokenCache(ctx context.Context, opts ...FileTokenCacheOption) (u2m_cache.TokenCache, error) { c := &fileTokenCache{} for _, opt := range opts { opt(c) } - if err := c.init(); err != nil { + if err := c.init(ctx); err != nil { return nil, err } // Fail fast if the cache is not working. @@ -127,10 +131,10 @@ func (c *fileTokenCache) Lookup(key string) (*oauth2.Token, error) { // init initializes the token cache file. It creates the file and directory if // they do not already exist. -func (c *fileTokenCache) init() error { +func (c *fileTokenCache) init(ctx context.Context) error { // set the default file location if c.fileLocation == "" { - home, err := os.UserHomeDir() + home, err := env.UserHomeDir(ctx) if err != nil { return fmt.Errorf("failed loading home directory: %w", err) } @@ -138,7 +142,7 @@ func (c *fileTokenCache) init() error { } // Create the cache file if it does not exist. if _, err := os.Stat(c.fileLocation); err != nil { - if !os.IsNotExist(err) { + if !errors.Is(err, fs.ErrNotExist) { return fmt.Errorf("stat file: %w", err) } // Create the parent directories if needed. diff --git a/libs/auth/storage/file_cache_test.go b/libs/auth/storage/file_cache_test.go index fdec3c5c46..4df7576c69 100644 --- a/libs/auth/storage/file_cache_test.go +++ b/libs/auth/storage/file_cache_test.go @@ -17,7 +17,7 @@ func setup(t *testing.T) string { } func TestStoreAndLookup(t *testing.T) { - c, err := NewFileTokenCache(WithFileLocation(setup(t))) + c, err := NewFileTokenCache(t.Context(), WithFileLocation(setup(t))) require.NoError(t, err) err = c.Store("x", &oauth2.Token{ AccessToken: "abc", @@ -38,7 +38,7 @@ func TestStoreAndLookup(t *testing.T) { } func TestNoCacheFileReturnsErrNotConfigured(t *testing.T) { - l, err := NewFileTokenCache(WithFileLocation(setup(t))) + l, err := NewFileTokenCache(t.Context(), WithFileLocation(setup(t))) require.NoError(t, err) _, err = l.Lookup("x") assert.Equal(t, u2m_cache.ErrNotFound, err) @@ -51,7 +51,7 @@ func TestLoadCorruptFile(t *testing.T) { err = os.WriteFile(f, []byte("abc"), ownerExecReadWrite) require.NoError(t, err) - _, err = NewFileTokenCache(WithFileLocation(f)) + _, err = NewFileTokenCache(t.Context(), WithFileLocation(f)) assert.EqualError(t, err, "load: parse: invalid character 'a' looking for beginning of value") } @@ -62,6 +62,6 @@ func TestLoadWrongVersion(t *testing.T) { err = os.WriteFile(f, []byte(`{"version": 823, "things": []}`), ownerExecReadWrite) require.NoError(t, err) - _, err = NewFileTokenCache(WithFileLocation(f)) + _, err = NewFileTokenCache(t.Context(), WithFileLocation(f)) assert.EqualError(t, err, "load: needs version 1, got version 823") } From fb8d730665ff2c6eb297d87f373501ffecc1fb38 Mon Sep 17 00:00:00 2001 From: simon Date: Wed, 22 Apr 2026 12:20:36 +0200 Subject: [PATCH 3/4] auth: wrap token cache for automatic host-key mirroring Replaces the caller-side storage.DualWrite helper with a DualWritingTokenCache wrapper. Every write through the wrapper under the primary key is also mirrored under the host key, so refresh paths (Token, ForceRefreshToken) preserve cross-SDK compatibility after the SDK stops dual-writing internally, not just Challenge paths. Co-authored-by: Isaac --- cmd/auth/login.go | 14 +- cmd/auth/token.go | 11 +- libs/auth/credentials.go | 19 ++- libs/auth/storage/dual_writing_cache.go | 66 ++++++++ libs/auth/storage/dual_writing_cache_test.go | 169 +++++++++++++++++++ libs/auth/storage/dualwrite.go | 39 ----- libs/auth/storage/dualwrite_test.go | 133 --------------- 7 files changed, 251 insertions(+), 200 deletions(-) create mode 100644 libs/auth/storage/dual_writing_cache.go create mode 100644 libs/auth/storage/dual_writing_cache_test.go delete mode 100644 libs/auth/storage/dualwrite.go delete mode 100644 libs/auth/storage/dualwrite_test.go diff --git a/cmd/auth/login.go b/cmd/auth/login.go index adc3db3a64..861bc44caa 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -233,7 +233,7 @@ a new profile is created. return fmt.Errorf("opening token cache: %w", err) } persistentAuthOpts := []u2m.PersistentAuthOption{ - u2m.WithTokenCache(tc), + u2m.WithTokenCache(storage.NewDualWritingTokenCache(tc, oauthArgument)), u2m.WithOAuthArgument(oauthArgument), u2m.WithBrowser(getBrowserFunc(cmd)), } @@ -252,11 +252,6 @@ a new profile is created. if err = persistentAuth.Challenge(); err != nil { return err } - if t, lookupErr := tc.Lookup(oauthArgument.GetCacheKey()); lookupErr == nil && t != nil { - if err := storage.DualWrite(tc, oauthArgument, t); err != nil { - log.Debugf(ctx, "token cache dual-write failed: %v", err) - } - } // 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); @@ -589,7 +584,7 @@ func discoveryLogin(ctx context.Context, dc discoveryClient, profileName string, return discoveryErr("opening token cache", err) } opts := []u2m.PersistentAuthOption{ - u2m.WithTokenCache(tc), + u2m.WithTokenCache(storage.NewDualWritingTokenCache(tc, arg)), u2m.WithOAuthArgument(arg), u2m.WithBrowser(browserFunc), u2m.WithDiscoveryLogin(), @@ -612,11 +607,6 @@ func discoveryLogin(ctx context.Context, dc discoveryClient, profileName string, if err := persistentAuth.Challenge(); err != nil { return discoveryErr("login via login.databricks.com failed", err) } - if t, lookupErr := tc.Lookup(arg.GetCacheKey()); lookupErr == nil && t != nil { - if err := storage.DualWrite(tc, arg, t); err != nil { - log.Debugf(ctx, "token cache dual-write failed: %v", err) - } - } discoveredHost := arg.GetDiscoveredHost() if discoveredHost == "" { diff --git a/cmd/auth/token.go b/cmd/auth/token.go index d27e789b0d..920ff7839d 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -18,7 +18,6 @@ import ( "github.com/databricks/cli/libs/databrickscfg/profile" "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/flags" - "github.com/databricks/cli/libs/log" "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" @@ -279,7 +278,8 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) { if err != nil { return nil, fmt.Errorf("opening token cache: %w", err) } - allArgs := append([]u2m.PersistentAuthOption{u2m.WithTokenCache(tc)}, args.persistentAuthOpts...) + wrappedCache := storage.NewDualWritingTokenCache(tc, oauthArgument) + allArgs := append([]u2m.PersistentAuthOption{u2m.WithTokenCache(wrappedCache)}, args.persistentAuthOpts...) allArgs = append(allArgs, u2m.WithOAuthArgument(oauthArgument)) persistentAuth, err := u2m.NewPersistentAuth(ctx, allArgs...) if err != nil { @@ -471,7 +471,7 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler) (string, *pr return "", nil, fmt.Errorf("opening token cache: %w", err) } persistentAuthOpts := []u2m.PersistentAuthOption{ - u2m.WithTokenCache(tc), + u2m.WithTokenCache(storage.NewDualWritingTokenCache(tc, oauthArgument)), u2m.WithOAuthArgument(oauthArgument), u2m.WithBrowser(func(url string) error { return browser.Open(ctx, url) }), } @@ -490,11 +490,6 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler) (string, *pr if err = persistentAuth.Challenge(); err != nil { return "", nil, err } - if t, lookupErr := tc.Lookup(oauthArgument.GetCacheKey()); lookupErr == nil && t != nil { - if err := storage.DualWrite(tc, oauthArgument, t); err != nil { - log.Debugf(ctx, "token cache dual-write failed: %v", err) - } - } clearKeys := oauthLoginClearKeys() if !loginArgs.IsUnifiedHost { diff --git a/libs/auth/credentials.go b/libs/auth/credentials.go index c3fb115b24..a406955dd3 100644 --- a/libs/auth/credentials.go +++ b/libs/auth/credentials.go @@ -99,7 +99,7 @@ func (c CLICredentials) Configure(ctx context.Context, cfg *config.Config) (cred if err != nil { return nil, err } - ts, err := c.persistentAuth(ctx, u2m.WithOAuthArgument(oauthArg)) + ts, err := c.persistentAuth(ctx, oauthArg) if err != nil { return nil, err } @@ -109,19 +109,22 @@ func (c CLICredentials) Configure(ctx context.Context, cfg *config.Config) (cred return cp, nil } -// 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. -func (c CLICredentials) persistentAuth(ctx context.Context, opts ...u2m.PersistentAuthOption) (auth.TokenSource, error) { +// 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. +// 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, opts...) + return c.persistentAuthFn(ctx, u2m.WithOAuthArgument(arg)) } tc, err := storage.NewFileTokenCache(ctx) if err != nil { return nil, fmt.Errorf("opening token cache: %w", err) } - opts = append([]u2m.PersistentAuthOption{u2m.WithTokenCache(tc)}, opts...) - ts, err := u2m.NewPersistentAuth(ctx, opts...) + ts, err := u2m.NewPersistentAuth(ctx, + u2m.WithTokenCache(storage.NewDualWritingTokenCache(tc, arg)), + u2m.WithOAuthArgument(arg), + ) if err != nil { return nil, err } diff --git a/libs/auth/storage/dual_writing_cache.go b/libs/auth/storage/dual_writing_cache.go new file mode 100644 index 0000000000..874429cf31 --- /dev/null +++ b/libs/auth/storage/dual_writing_cache.go @@ -0,0 +1,66 @@ +package storage + +import ( + "github.com/databricks/databricks-sdk-go/credentials/u2m" + u2m_cache "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" + "golang.org/x/oauth2" +) + +// DualWritingTokenCache wraps a TokenCache so that every write under the +// primary OAuth cache key is also mirrored under the legacy host-based key. +// This preserves the cross-SDK compatibility convention historically +// implemented inside PersistentAuth.dualWrite in the SDK, now moved +// caller-side per the cache-ownership split between SDK and CLI. +// +// Mirroring happens inside Store, so every SDK-internal write (Challenge, +// refresh, discovery) dual-writes without requiring each call site to invoke +// a helper explicitly. +type DualWritingTokenCache struct { + inner u2m_cache.TokenCache + arg u2m.OAuthArgument +} + +// NewDualWritingTokenCache returns a TokenCache wrapping inner that mirrors +// writes made under arg.GetCacheKey() to the argument's host key when one +// can be derived (via DiscoveryOAuthArgument.GetDiscoveredHost or +// HostCacheKeyProvider.GetHostCacheKey). +func NewDualWritingTokenCache(inner u2m_cache.TokenCache, arg u2m.OAuthArgument) *DualWritingTokenCache { + return &DualWritingTokenCache{inner: inner, arg: arg} +} + +// Store implements [u2m_cache.TokenCache]. Writes under the primary key are +// also mirrored under the host key (when distinct); writes under any other +// key pass through unchanged so that a Store(hostKey, t) from an older SDK +// that still dual-writes internally does not recursively re-expand. +func (c *DualWritingTokenCache) Store(key string, t *oauth2.Token) error { + if err := c.inner.Store(key, t); err != nil { + return err + } + primaryKey := c.arg.GetCacheKey() + if key != primaryKey { + return nil + } + hostKey := hostCacheKey(c.arg) + if hostKey == "" || hostKey == primaryKey { + return nil + } + return c.inner.Store(hostKey, t) +} + +// Lookup implements [u2m_cache.TokenCache]; delegates to the inner cache. +func (c *DualWritingTokenCache) Lookup(key string) (*oauth2.Token, error) { + return c.inner.Lookup(key) +} + +// hostCacheKey mirrors the SDK's former PersistentAuth.hostCacheKey: +// discovery arguments expose the host via GetDiscoveredHost (populated by +// Challenge); static arguments expose it via HostCacheKeyProvider. +func hostCacheKey(arg u2m.OAuthArgument) string { + if discoveryArg, ok := arg.(u2m.DiscoveryOAuthArgument); ok { + return discoveryArg.GetDiscoveredHost() + } + if hcp, ok := arg.(u2m.HostCacheKeyProvider); ok { + return hcp.GetHostCacheKey() + } + return "" +} diff --git a/libs/auth/storage/dual_writing_cache_test.go b/libs/auth/storage/dual_writing_cache_test.go new file mode 100644 index 0000000000..884e7285e9 --- /dev/null +++ b/libs/auth/storage/dual_writing_cache_test.go @@ -0,0 +1,169 @@ +package storage + +import ( + "sync" + "testing" + + "github.com/databricks/databricks-sdk-go/credentials/u2m" + u2m_cache "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/oauth2" +) + +// memoryCache is a minimal in-memory TokenCache used only by wrapper tests. +type memoryCache struct { + mu sync.Mutex + tokens map[string]*oauth2.Token +} + +func newMemoryCache() *memoryCache { + return &memoryCache{tokens: map[string]*oauth2.Token{}} +} + +func (c *memoryCache) Store(key string, t *oauth2.Token) error { + c.mu.Lock() + defer c.mu.Unlock() + if t == nil { + delete(c.tokens, key) + return nil + } + c.tokens[key] = t + return nil +} + +func (c *memoryCache) Lookup(key string) (*oauth2.Token, error) { + c.mu.Lock() + defer c.mu.Unlock() + t, ok := c.tokens[key] + if !ok { + return nil, u2m_cache.ErrNotFound + } + return t, nil +} + +// plainArg implements OAuthArgument only, exercising the "no host key" branch. +type plainArg struct { + key string +} + +func (a plainArg) GetCacheKey() string { return a.key } + +// hostArg implements HostCacheKeyProvider so the wrapper mirrors the token +// to the configured host key. +type hostArg struct { + key string + hostKey string +} + +func (a hostArg) GetCacheKey() string { return a.key } +func (a hostArg) GetHostCacheKey() string { return a.hostKey } + +func TestDualWritingCacheStorePrimaryMirrorsHost(t *testing.T) { + inner := newMemoryCache() + arg := hostArg{key: "profile-a", hostKey: "https://example.databricks.com"} + c := NewDualWritingTokenCache(inner, arg) + tok := &oauth2.Token{AccessToken: "abc", RefreshToken: "r"} + + require.NoError(t, c.Store("profile-a", tok)) + + primary, err := inner.Lookup("profile-a") + require.NoError(t, err) + assert.Equal(t, tok, primary) + + host, err := inner.Lookup("https://example.databricks.com") + require.NoError(t, err) + assert.Equal(t, tok, host) +} + +func TestDualWritingCacheStoreNonPrimaryDoesNotMirror(t *testing.T) { + // An older SDK still running its internal dualWrite will follow up the + // primary Store with a Store(hostKey, t). The wrapper must pass that + // second write through without re-expanding into another pair. + inner := newMemoryCache() + arg := hostArg{key: "profile-a", hostKey: "https://example.databricks.com"} + c := NewDualWritingTokenCache(inner, arg) + tok := &oauth2.Token{AccessToken: "abc"} + + require.NoError(t, c.Store("https://example.databricks.com", tok)) + + host, err := inner.Lookup("https://example.databricks.com") + require.NoError(t, err) + assert.Equal(t, tok, host) + _, err = inner.Lookup("profile-a") + require.ErrorIs(t, err, u2m_cache.ErrNotFound) +} + +func TestDualWritingCacheStoreNoHostKey(t *testing.T) { + inner := newMemoryCache() + arg := plainArg{key: "profile-a"} + c := NewDualWritingTokenCache(inner, arg) + tok := &oauth2.Token{AccessToken: "abc"} + + require.NoError(t, c.Store("profile-a", tok)) + + got, err := inner.Lookup("profile-a") + require.NoError(t, err) + assert.Equal(t, tok, got) + assert.Len(t, inner.tokens, 1) +} + +func TestDualWritingCacheStoreHostKeyEqualsPrimary(t *testing.T) { + inner := newMemoryCache() + arg := hostArg{key: "https://example.databricks.com", hostKey: "https://example.databricks.com"} + c := NewDualWritingTokenCache(inner, arg) + tok := &oauth2.Token{AccessToken: "abc"} + + require.NoError(t, c.Store("https://example.databricks.com", tok)) + + assert.Len(t, inner.tokens, 1) +} + +func TestDualWritingCacheDiscoveryArgWithDiscoveredHost(t *testing.T) { + inner := newMemoryCache() + arg, err := u2m.NewBasicDiscoveryOAuthArgument("profile-a") + require.NoError(t, err) + arg.SetDiscoveredHost("https://example.databricks.com") + c := NewDualWritingTokenCache(inner, arg) + tok := &oauth2.Token{AccessToken: "abc"} + + require.NoError(t, c.Store("profile-a", tok)) + + primary, err := inner.Lookup("profile-a") + require.NoError(t, err) + assert.Equal(t, tok, primary) + + host, err := inner.Lookup("https://example.databricks.com") + require.NoError(t, err) + assert.Equal(t, tok, host) +} + +func TestDualWritingCacheDiscoveryArgWithEmptyDiscoveredHost(t *testing.T) { + inner := newMemoryCache() + arg, err := u2m.NewBasicDiscoveryOAuthArgument("profile-a") + require.NoError(t, err) + c := NewDualWritingTokenCache(inner, arg) + tok := &oauth2.Token{AccessToken: "abc"} + + require.NoError(t, c.Store("profile-a", tok)) + + assert.Len(t, inner.tokens, 1) + primary, err := inner.Lookup("profile-a") + require.NoError(t, err) + assert.Equal(t, tok, primary) +} + +func TestDualWritingCacheLookupDelegates(t *testing.T) { + inner := newMemoryCache() + arg := hostArg{key: "profile-a", hostKey: "https://example.databricks.com"} + c := NewDualWritingTokenCache(inner, arg) + tok := &oauth2.Token{AccessToken: "abc"} + require.NoError(t, inner.Store("profile-a", tok)) + + got, err := c.Lookup("profile-a") + require.NoError(t, err) + assert.Equal(t, tok, got) + + _, err = c.Lookup("missing") + require.ErrorIs(t, err, u2m_cache.ErrNotFound) +} diff --git a/libs/auth/storage/dualwrite.go b/libs/auth/storage/dualwrite.go deleted file mode 100644 index 14ce017e23..0000000000 --- a/libs/auth/storage/dualwrite.go +++ /dev/null @@ -1,39 +0,0 @@ -package storage - -import ( - "github.com/databricks/databricks-sdk-go/credentials/u2m" - u2m_cache "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" - "golang.org/x/oauth2" -) - -// DualWrite stores t under arg.GetCacheKey() (the primary key) and, if a -// legacy host-based cache key can be derived from arg, also stores t under -// that host key. Mirrors the convention used historically by -// PersistentAuth.dualWrite and hostCacheKey in the SDK. -func DualWrite(cache u2m_cache.TokenCache, arg u2m.OAuthArgument, t *oauth2.Token) error { - primaryKey := arg.GetCacheKey() - if err := cache.Store(primaryKey, t); err != nil { - return err - } - hostKey := hostCacheKey(arg) - if hostKey != "" && hostKey != primaryKey { - if err := cache.Store(hostKey, t); err != nil { - return err - } - } - return nil -} - -// hostCacheKey mirrors PersistentAuth.hostCacheKey in the SDK: discovery -// arguments learn their host from the OAuth callback, so their host key -// must be read via GetDiscoveredHost *after* Challenge(). Static arguments -// use HostCacheKeyProvider. -func hostCacheKey(arg u2m.OAuthArgument) string { - if discoveryArg, ok := arg.(u2m.DiscoveryOAuthArgument); ok { - return discoveryArg.GetDiscoveredHost() - } - if hcp, ok := arg.(u2m.HostCacheKeyProvider); ok { - return hcp.GetHostCacheKey() - } - return "" -} diff --git a/libs/auth/storage/dualwrite_test.go b/libs/auth/storage/dualwrite_test.go deleted file mode 100644 index f3da4ae8d0..0000000000 --- a/libs/auth/storage/dualwrite_test.go +++ /dev/null @@ -1,133 +0,0 @@ -package storage - -import ( - "sync" - "testing" - - "github.com/databricks/databricks-sdk-go/credentials/u2m" - u2m_cache "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - "golang.org/x/oauth2" -) - -// memoryCache is a minimal in-memory TokenCache used only by DualWrite tests. -type memoryCache struct { - mu sync.Mutex - tokens map[string]*oauth2.Token -} - -func newMemoryCache() *memoryCache { - return &memoryCache{tokens: map[string]*oauth2.Token{}} -} - -func (c *memoryCache) Store(key string, t *oauth2.Token) error { - c.mu.Lock() - defer c.mu.Unlock() - if t == nil { - delete(c.tokens, key) - return nil - } - c.tokens[key] = t - return nil -} - -func (c *memoryCache) Lookup(key string) (*oauth2.Token, error) { - c.mu.Lock() - defer c.mu.Unlock() - t, ok := c.tokens[key] - if !ok { - return nil, u2m_cache.ErrNotFound - } - return t, nil -} - -// plainArg implements OAuthArgument only, exercising the "no host key" branch. -type plainArg struct { - key string -} - -func (a plainArg) GetCacheKey() string { return a.key } - -// hostArg implements HostCacheKeyProvider so DualWrite mirrors the token to -// the configured host key. -type hostArg struct { - key string - hostKey string -} - -func (a hostArg) GetCacheKey() string { return a.key } -func (a hostArg) GetHostCacheKey() string { return a.hostKey } - -func TestDualWriteNoHostKey(t *testing.T) { - cache := newMemoryCache() - arg := plainArg{key: "profile-a"} - tok := &oauth2.Token{AccessToken: "abc", RefreshToken: "r"} - - require.NoError(t, DualWrite(cache, arg, tok)) - - got, err := cache.Lookup("profile-a") - require.NoError(t, err) - assert.Equal(t, tok, got) - assert.Len(t, cache.tokens, 1) -} - -func TestDualWriteHostKeyDistinct(t *testing.T) { - cache := newMemoryCache() - arg := hostArg{key: "profile-a", hostKey: "https://example.databricks.com"} - tok := &oauth2.Token{AccessToken: "abc", RefreshToken: "r"} - - require.NoError(t, DualWrite(cache, arg, tok)) - - primary, err := cache.Lookup("profile-a") - require.NoError(t, err) - assert.Equal(t, tok, primary) - - host, err := cache.Lookup("https://example.databricks.com") - require.NoError(t, err) - assert.Equal(t, tok, host) - - assert.Len(t, cache.tokens, 2) -} - -func TestDualWriteHostKeyEqualsPrimary(t *testing.T) { - cache := newMemoryCache() - arg := hostArg{key: "https://example.databricks.com", hostKey: "https://example.databricks.com"} - tok := &oauth2.Token{AccessToken: "abc"} - - require.NoError(t, DualWrite(cache, arg, tok)) - - assert.Len(t, cache.tokens, 1) -} - -func TestDualWriteDiscoveryArgWithDiscoveredHost(t *testing.T) { - cache := newMemoryCache() - arg, err := u2m.NewBasicDiscoveryOAuthArgument("profile-a") - require.NoError(t, err) - arg.SetDiscoveredHost("https://example.databricks.com") - tok := &oauth2.Token{AccessToken: "abc"} - - require.NoError(t, DualWrite(cache, arg, tok)) - - primary, err := cache.Lookup("profile-a") - require.NoError(t, err) - assert.Equal(t, tok, primary) - - host, err := cache.Lookup("https://example.databricks.com") - require.NoError(t, err) - assert.Equal(t, tok, host) -} - -func TestDualWriteDiscoveryArgWithEmptyDiscoveredHost(t *testing.T) { - cache := newMemoryCache() - arg, err := u2m.NewBasicDiscoveryOAuthArgument("profile-a") - require.NoError(t, err) - tok := &oauth2.Token{AccessToken: "abc"} - - require.NoError(t, DualWrite(cache, arg, tok)) - - assert.Len(t, cache.tokens, 1) - primary, err := cache.Lookup("profile-a") - require.NoError(t, err) - assert.Equal(t, tok, primary) -} From 038c1c87cf926a4ce042f0b2553c831c49eecdf3 Mon Sep 17 00:00:00 2001 From: simon Date: Wed, 22 Apr 2026 13:24:05 +0200 Subject: [PATCH 4/4] auth: inject token cache into discoveryLogin and loadToken Callers (cmd.RunE closures) now construct the FileTokenCache and pass it to discoveryLogin, runInlineLogin, and loadToken. Previously each of those helpers built the file cache internally, which meant unit tests hitting discoveryLogin or loadToken would create/touch ~/.databricks/token-cache.json on the developer's machine. Tests now pass the in-memory cache helper, so the real file is no longer a side effect of running the suite. Addresses review feedback from Mihai on login.go and token.go. Co-authored-by: Isaac --- cmd/auth/login.go | 22 ++++++++++------------ cmd/auth/login_test.go | 27 +++++++++++++++++---------- cmd/auth/token.go | 30 ++++++++++++++++-------------- cmd/auth/token_test.go | 30 ++++++++++++++++++++++++++++++ 4 files changed, 73 insertions(+), 36 deletions(-) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 861bc44caa..21c3aa3608 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -22,6 +22,7 @@ import ( "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/config/experimental/auth/authconv" "github.com/databricks/databricks-sdk-go/credentials/u2m" + "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" "github.com/spf13/cobra" "golang.org/x/oauth2" ) @@ -190,13 +191,18 @@ 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{}, profileName, loginTimeout, scopes, existingProfile, getBrowserFunc(cmd)) + return discoveryLogin(ctx, &defaultDiscoveryClient{}, tokenCache, profileName, loginTimeout, scopes, existingProfile, getBrowserFunc(cmd)) } // Load unified host flag from the profile if not explicitly set via CLI flag. @@ -228,12 +234,8 @@ a new profile is created. if err != nil { return err } - tc, err := storage.NewFileTokenCache(ctx) - if err != nil { - return fmt.Errorf("opening token cache: %w", err) - } persistentAuthOpts := []u2m.PersistentAuthOption{ - u2m.WithTokenCache(storage.NewDualWritingTokenCache(tc, oauthArgument)), + u2m.WithTokenCache(storage.NewDualWritingTokenCache(tokenCache, oauthArgument)), u2m.WithOAuthArgument(oauthArgument), u2m.WithBrowser(getBrowserFunc(cmd)), } @@ -568,7 +570,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, profileName string, timeout time.Duration, scopes string, existingProfile *profile.Profile, browserFunc func(string) error) error { +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 { arg, err := dc.NewOAuthArgument(profileName) if err != nil { return discoveryErr("setting up login.databricks.com", err) @@ -579,12 +581,8 @@ func discoveryLogin(ctx context.Context, dc discoveryClient, profileName string, scopesList = splitScopes(existingProfile.Scopes) } - tc, err := storage.NewFileTokenCache(ctx) - if err != nil { - return discoveryErr("opening token cache", err) - } opts := []u2m.PersistentAuthOption{ - u2m.WithTokenCache(storage.NewDualWritingTokenCache(tc, arg)), + u2m.WithTokenCache(storage.NewDualWritingTokenCache(tokenCache, arg)), u2m.WithOAuthArgument(arg), u2m.WithBrowser(browserFunc), u2m.WithDiscoveryLogin(), diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index 81924f027a..2b8d473f51 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -20,12 +20,19 @@ import ( "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/log" "github.com/databricks/databricks-sdk-go/credentials/u2m" + "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" "github.com/spf13/cobra" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/oauth2" ) +// newTestTokenCache returns an in-memory token cache for tests so that +// discoveryLogin and other login helpers don't touch ~/.databricks/token-cache.json. +func newTestTokenCache() cache.TokenCache { + return &inMemoryTokenCache{Tokens: map[string]*oauth2.Token{}} +} + // logBuffer is a thread-safe bytes.Buffer for capturing log output in tests. type logBuffer struct { mu sync.Mutex @@ -623,7 +630,7 @@ 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 }) + err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "all-apis, ,sql,", nil, func(string) error { return nil }) require.NoError(t, err) assert.Equal(t, "https://workspace.example.com", dc.introspectHost) @@ -671,7 +678,7 @@ func TestDiscoveryLogin_AccountIDMismatchWarning(t *testing.T) { AccountID: "old-account-id", } - err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) require.NoError(t, err) // Verify warning about mismatched account IDs was logged. @@ -719,7 +726,7 @@ func TestDiscoveryLogin_NoWarningWhenAccountIDsMatch(t *testing.T) { AccountID: "same-account-id", } - err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) require.NoError(t, err) // No warning should be logged when account IDs match. @@ -739,7 +746,7 @@ func TestDiscoveryLogin_EmptyDiscoveredHostReturnsError(t *testing.T) { } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", nil, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", nil, func(string) error { return nil }) require.Error(t, err) assert.Contains(t, err.Error(), "no workspace host was discovered") } @@ -771,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, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -808,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, "DISCOVERY", time.Second, "all-apis", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "all-apis", existingProfile, func(string) error { return nil }) require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -848,7 +855,7 @@ func TestDiscoveryLogin_SPOGHostPopulatesAccountIDFromDiscovery(t *testing.T) { } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", nil, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", nil, func(string) error { return nil }) require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -883,7 +890,7 @@ func TestDiscoveryLogin_IntrospectionFallsBackWhenDiscoveryFails(t *testing.T) { } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", nil, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", nil, func(string) error { return nil }) require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -932,7 +939,7 @@ auth_type = databricks-cli } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -982,7 +989,7 @@ auth_type = databricks-cli } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) diff --git a/cmd/auth/token.go b/cmd/auth/token.go index 920ff7839d..d82bcb2c9a 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -79,6 +79,11 @@ and secret is not supported.`, ctx := cmd.Context() profileName := cmd.Flag("profile").Value.String() + tokenCache, err := storage.NewFileTokenCache(ctx) + if err != nil { + return fmt.Errorf("opening token cache: %w", err) + } + t, err := loadToken(ctx, loadTokenArgs{ authArguments: authArguments, profileName: profileName, @@ -86,6 +91,7 @@ and secret is not supported.`, tokenTimeout: tokenTimeout, forceRefresh: forceRefresh, profiler: profile.DefaultProfiler, + tokenCache: tokenCache, persistentAuthOpts: nil, }) if err != nil { @@ -134,6 +140,10 @@ type loadTokenArgs struct { // profiler is the profiler to use for reading the host and account ID from the .databrickscfg file. profiler profile.Profiler + // tokenCache is the underlying TokenCache used for OAuth tokens. The caller is + // responsible for construction so that tests can substitute an in-memory cache. + tokenCache cache.TokenCache + // persistentAuthOpts are the options to pass to the persistent auth client. persistentAuthOpts []u2m.PersistentAuthOption } @@ -185,7 +195,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) + resolvedProfile, existingProfile, err = resolveNoArgsToken(ctx, args.profiler, args.authArguments, args.tokenCache) if err != nil { return nil, err } @@ -274,11 +284,7 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) { if err != nil { return nil, err } - tc, err := storage.NewFileTokenCache(ctx) - if err != nil { - return nil, fmt.Errorf("opening token cache: %w", err) - } - wrappedCache := storage.NewDualWritingTokenCache(tc, oauthArgument) + wrappedCache := storage.NewDualWritingTokenCache(args.tokenCache, oauthArgument) allArgs := append([]u2m.PersistentAuthOption{u2m.WithTokenCache(wrappedCache)}, args.persistentAuthOpts...) allArgs = append(allArgs, u2m.WithOAuthArgument(oauthArgument)) persistentAuth, err := u2m.NewPersistentAuth(ctx, allArgs...) @@ -321,7 +327,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) (string, *profile.Profile, error) { +func resolveNoArgsToken(ctx context.Context, profiler profile.Profiler, authArgs *auth.AuthArguments, tokenCache cache.TokenCache) (string, *profile.Profile, error) { // Step 1: Try DATABRICKS_HOST env var (highest priority). if envHost := env.Get(ctx, "DATABRICKS_HOST"); envHost != "" { authArgs.Host = envHost @@ -370,7 +376,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) + return runInlineLogin(ctx, profiler, tokenCache) default: p, err := loadProfileByName(ctx, selectedName, profiler) if err != nil { @@ -434,7 +440,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) (string, *profile.Profile, error) { +func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenCache cache.TokenCache) (string, *profile.Profile, error) { profileName, err := promptForProfile(ctx, "DEFAULT") if err != nil { return "", nil, err @@ -466,12 +472,8 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler) (string, *pr if err != nil { return "", nil, err } - tc, err := storage.NewFileTokenCache(ctx) - if err != nil { - return "", nil, fmt.Errorf("opening token cache: %w", err) - } persistentAuthOpts := []u2m.PersistentAuthOption{ - u2m.WithTokenCache(storage.NewDualWritingTokenCache(tc, oauthArgument)), + u2m.WithTokenCache(storage.NewDualWritingTokenCache(tokenCache, oauthArgument)), u2m.WithOAuthArgument(oauthArgument), u2m.WithBrowser(func(url string) error { return browser.Open(ctx, url) }), } diff --git a/cmd/auth/token_test.go b/cmd/auth/token_test.go index ec7fe2004a..3dfa4e5d21 100644 --- a/cmd/auth/token_test.go +++ b/cmd/auth/token_test.go @@ -221,6 +221,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -241,6 +242,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -258,6 +260,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -275,6 +278,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -292,6 +296,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -308,6 +313,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -324,6 +330,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -340,6 +347,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{"workspace-a"}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -356,6 +364,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{"workspace-a.cloud.databricks.com"}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -372,6 +381,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{"default.dev"}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -388,6 +398,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{"nonexistent.cloud.databricks.com"}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -419,6 +430,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -436,6 +448,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -454,6 +467,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -473,6 +487,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -490,6 +505,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -508,6 +524,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -526,6 +543,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -542,6 +560,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{"workspace-a"}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -557,6 +576,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: nil, }, wantErr: "no profile specified. Use --profile to specify which profile to use", @@ -569,6 +589,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{}, tokenTimeout: 1 * time.Hour, profiler: profile.InMemoryProfiler{}, + tokenCache: tokenCache, persistentAuthOpts: nil, }, wantErr: "no profiles configured. Run 'databricks auth login' to create a profile", @@ -581,6 +602,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{}, tokenTimeout: 1 * time.Hour, profiler: errProfiler{err: profile.ErrNoConfiguration}, + tokenCache: tokenCache, persistentAuthOpts: nil, }, wantErr: "no profiles configured. Run 'databricks auth login' to create a profile", @@ -638,6 +660,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -658,6 +681,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -678,6 +702,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -699,6 +724,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -734,6 +760,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -750,6 +777,7 @@ func TestToken_loadToken(t *testing.T) { args: []string{}, tokenTimeout: 1 * time.Hour, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -769,6 +797,7 @@ func TestToken_loadToken(t *testing.T) { tokenTimeout: 1 * time.Hour, forceRefresh: true, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}), @@ -786,6 +815,7 @@ func TestToken_loadToken(t *testing.T) { tokenTimeout: 1 * time.Hour, forceRefresh: true, profiler: profiler, + tokenCache: tokenCache, persistentAuthOpts: []u2m.PersistentAuthOption{ u2m.WithTokenCache(tokenCache), u2m.WithOAuthEndpointSupplier(&MockApiClient{}),