Skip to content

auth: import FileTokenCache into CLI and wire DualWrite #5056

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
simonfaltum wants to merge 4 commits into
base: main
Choose a base branch
from
Open

auth: import FileTokenCache into CLI and wire DualWrite #5056

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
274a09e
auth: import FileTokenCache into CLI and wire DualWrite
simonfaltum Apr 21, 2026
ba0a2d4
Fix lint: use env.UserHomeDir and errors.Is(fs.ErrNotExist)
simonfaltum Apr 21, 2026
fb8d730
auth: wrap token cache for automatic host-key mirroring
simonfaltum Apr 22, 2026
038c1c8
auth: inject token cache into discoveryLogin and loadToken
simonfaltum Apr 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions NEXT_CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)).
Expand Down
13 changes: 11 additions & 2 deletions cmd/auth/login.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -21,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"
)
Expand Down Expand Up @@ -189,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.
Expand Down Expand Up @@ -228,6 +235,7 @@ a new profile is created.
return err
}
persistentAuthOpts := []u2m.PersistentAuthOption{
u2m.WithTokenCache(storage.NewDualWritingTokenCache(tokenCache, oauthArgument)),
u2m.WithOAuthArgument(oauthArgument),
u2m.WithBrowser(getBrowserFunc(cmd)),
}
Expand Down Expand Up @@ -562,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)
Expand All @@ -574,6 +582,7 @@ func discoveryLogin(ctx context.Context, dc discoveryClient, profileName string,
}

opts := []u2m.PersistentAuthOption{
u2m.WithTokenCache(storage.NewDualWritingTokenCache(tokenCache, arg)),
u2m.WithOAuthArgument(arg),
u2m.WithBrowser(browserFunc),
u2m.WithDiscoveryLogin(),
Expand Down
27 changes: 17 additions & 10 deletions cmd/auth/login_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand All @@ -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")
}
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
3 changes: 2 additions & 1 deletion cmd/auth/logout.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -132,7 +133,7 @@ to specify it explicitly.
profileName = selected
}

tokenCache, err := cache.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)
}
Expand Down
24 changes: 19 additions & 5 deletions cmd/auth/token.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ 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"
Expand Down Expand Up @@ -78,13 +79,19 @@ 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,
args: args,
tokenTimeout: tokenTimeout,
forceRefresh: forceRefresh,
profiler: profile.DefaultProfiler,
tokenCache: tokenCache,
persistentAuthOpts: nil,
})
if err != nil {
Expand Down Expand Up @@ -133,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
}
Expand Down Expand Up @@ -184,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
}
Expand Down Expand Up @@ -273,7 +284,9 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) {
if err != nil {
return nil, err
}
allArgs := append(args.persistentAuthOpts, u2m.WithOAuthArgument(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...)
if err != nil {
helpMsg := helpfulError(ctx, args.profileName, oauthArgument)
Expand Down Expand Up @@ -314,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
Expand Down Expand Up @@ -363,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 {
Expand Down Expand Up @@ -427,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
Expand Down Expand Up @@ -460,6 +473,7 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler) (string, *pr
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) }),
}
Expand Down
Loading
Loading