Add --force-refresh flag to auth token command

NEXT_CHANGELOG.md

@@ -7,6 +7,7 @@
 - Add `bundle.engine` config setting to select the deployment engine (`terraform` or `direct`). The `DATABRICKS_BUNDLE_ENGINE` environment variable takes precedence over this setting. When the configured engine doesn't match existing deployment state, a warning is issued and the existing engine is used ([#4749](https://github.com/databricks/cli/pull/4749)).
 
 ### CLI
+* Add `--force-refresh` flag to `databricks auth token` to force a token refresh even when the cached token is still valid ([#4767](https://github.com/databricks/cli/pull/4767)).
 
 ### Bundles
 * engine/direct: Fix permanent drift on experiment name field ([#4627](https://github.com/databricks/cli/pull/4627))

acceptance/cmd/auth/token/force-refresh-invalid-refresh-token/output.txt

@@ -0,0 +1,4 @@
+Error: A new access token could not be retrieved because the refresh token is invalid. To reauthenticate, run the following command:
+  $ databricks auth login --profile test-profile
+
+Exit code: 1

acceptance/cmd/auth/token/force-refresh-invalid-refresh-token/script

@@ -0,0 +1,4 @@
+setup_test_profile
+setup_test_token_cache
+
+errcode $CLI auth token --profile test-profile --force-refresh

acceptance/cmd/auth/token/force-refresh-invalid-refresh-token/test.toml

@@ -0,0 +1,4 @@
+[[Server]]
+Pattern = "POST /oidc/v1/token"
+Response.StatusCode = 401
+Response.Body = '{"error": "invalid_request", "error_description": "Refresh token is invalid"}'

acceptance/cmd/auth/token/force-refresh-no-cache/output.txt

@@ -0,0 +1,3 @@
+Error: cache: databricks OAuth is not configured for this host. Try logging in again with `databricks auth login --profile test-profile` before retrying. If this fails, please report this issue to the Databricks CLI maintainers at https://github.com/databricks/cli/issues/new
+
+Exit code: 1

acceptance/cmd/auth/token/force-refresh-no-cache/script

@@ -0,0 +1,3 @@
+setup_test_profile
+
+errcode $CLI auth token --profile test-profile --force-refresh

acceptance/cmd/auth/token/force-refresh-success/output.txt

@@ -0,0 +1,3 @@
+
+>>> [CLI] auth token --profile test-profile --force-refresh
+"oauth-token"

acceptance/cmd/auth/token/force-refresh-success/script

@@ -0,0 +1,6 @@
+setup_test_profile
+setup_test_token_cache
+
+# The cached token is "cached-access-token". The mock OIDC server returns
+# "oauth-token", so this confirms --force-refresh fetched a fresh token.
+trace $CLI auth token --profile test-profile --force-refresh | jq .access_token

acceptance/cmd/auth/token/script.prepare

@@ -0,0 +1,31 @@
+setup_test_profile() {
+    export DATABRICKS_HOST_ORIG="$DATABRICKS_HOST"
+
+    sethome "./home"
+    unset DATABRICKS_HOST
+    unset DATABRICKS_TOKEN
+    unset DATABRICKS_CONFIG_PROFILE
+
+    cat > "./home/.databrickscfg" <<ENDCFG
+[test-profile]
+host = $DATABRICKS_HOST_ORIG
+auth_type = databricks-cli
+ENDCFG
+}
+
+setup_test_token_cache() {
+    mkdir -p "./home/.databricks"
+    cat > "./home/.databricks/token-cache.json" <<ENDCACHE
+{
+  "version": 1,
+  "tokens": {
+    "test-profile": {
+      "access_token": "cached-access-token",
+      "token_type": "Bearer",
+      "refresh_token": "test-refresh-token",
+      "expiry": "2099-01-01T00:00:00Z"
+    }
+  }
+}
+ENDCACHE
+}

cmd/auth/in_memory_test.go

@@ -10,20 +10,28 @@ type inMemoryTokenCache struct {
 }
 
 // Lookup implements TokenCache.
+// Returns a copy to match real (file-backed) cache behavior, where each
+// Lookup deserializes a fresh struct. Without the copy, callers that
+// mutate the returned token (e.g. clearing RefreshToken) would corrupt
+// entries shared across test cases.
 func (i *inMemoryTokenCache) Lookup(key string) (*oauth2.Token, error) {
 	token, ok := i.Tokens[key]
 	if !ok {
 		return nil, cache.ErrNotFound
 	}
-	return token, nil
+	cp := *token
+	return &cp, nil
 }
 
 // Store implements TokenCache.
+// Stores a copy to prevent callers from mutating cached entries after Store
+// returns (mirrors file-backed cache semantics).
 func (i *inMemoryTokenCache) Store(key string, t *oauth2.Token) error {
 	if t == nil {
 		delete(i.Tokens, key)
 	} else {
-		i.Tokens[key] = t
+		cp := *t
+		i.Tokens[key] = &cp
 	}
 	return nil
 }

cmd/auth/token.go

@@ -55,15 +55,20 @@ func newTokenCommand(authArguments *auth.AuthArguments) *cobra.Command {
 		Use:   "token [HOST_OR_PROFILE]",
 		Short: "Get authentication token",
 		Long: `Get authentication token from the local cache in ~/.databricks/token-cache.json.
-Refresh the access token if it is expired. Note: This command only works with
-U2M authentication (using the 'databricks auth login' command). M2M authentication
-using a client ID and secret is not supported.`,
+Refresh the access token if it is expired or close to expiry. Use --force-refresh
+to bypass expiry checks. Note: This command only works with U2M authentication
+(using the 'databricks auth login' command). M2M authentication using a client ID
+and secret is not supported.`,
 	}
 
 	var tokenTimeout time.Duration
 	cmd.Flags().DurationVar(&tokenTimeout, "timeout", defaultTimeout,
 		"Timeout for acquiring a token.")
 
+	var forceRefresh bool
+	cmd.Flags().BoolVar(&forceRefresh, "force-refresh", false,
+		"Force a token refresh even if the cached token is still valid.")
+
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
 		profileName := ""
@@ -77,6 +82,7 @@ using a client ID and secret is not supported.`,
 			profileName:        profileName,
 			args:               args,
 			tokenTimeout:       tokenTimeout,
+			forceRefresh:       forceRefresh,
 			profiler:           profile.DefaultProfiler,
 			persistentAuthOpts: nil,
 		})
@@ -107,6 +113,9 @@ type loadTokenArgs struct {
 	// tokenTimeout is the timeout for retrieving (and potentially refreshing) an OAuth token.
 	tokenTimeout time.Duration
 
+	// forceRefresh forces a token refresh even if the cached token is still valid.
+	forceRefresh bool
+
 	// profiler is the profiler to use for reading the host and account ID from the .databrickscfg file.
 	profiler profile.Profiler
 
@@ -253,7 +262,12 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) {
 		helpMsg := helpfulError(ctx, args.profileName, oauthArgument)
 		return nil, fmt.Errorf("%w. %s", err, helpMsg)
 	}
-	t, err := persistentAuth.Token()
+	var t *oauth2.Token
+	if args.forceRefresh {
+		t, err = persistentAuth.ForceRefreshToken()
+	} else {
+		t, err = persistentAuth.Token()
+	}
 	if err != nil {
 		if errors.Is(err, cache.ErrNotFound) {
 			// The error returned by the SDK when the token cache doesn't exist or doesn't contain a token

cmd/auth/token_test.go

@@ -2,6 +2,7 @@ package auth
 
 import (
 	"context"
+	"errors"
 	"net/http"
 	"testing"
 	"time"
@@ -16,6 +17,12 @@ import (
 	"golang.org/x/oauth2"
 )
 
+type failOnCallTransport struct{}
+
+func (failOnCallTransport) RoundTrip(*http.Request) (*http.Response, error) {
+	return nil, errors.New("unexpected HTTP call")
+}
+
 var refreshFailureTokenResponse = fixtures.HTTPFixture{
 	MatchAny: true,
 	Status:   401,
@@ -135,6 +142,10 @@ func TestToken_loadToken(t *testing.T) {
 				Host:                 "https://m2m.cloud.databricks.com",
 				HasClientCredentials: true,
 			},
+			{
+				Name: "valid-token",
+				Host: "https://valid-token.cloud.databricks.com",
+			},
 		},
 	}
 	tokenCache := &inMemoryTokenCache{
@@ -181,11 +192,16 @@ func TestToken_loadToken(t *testing.T) {
 				RefreshToken: "dup1",
 				Expiry:       time.Now().Add(1 * time.Hour),
 			},
+			"valid-token": {
+				AccessToken:  "cached-access-token",
+				RefreshToken: "valid-token",
+				Expiry:       time.Now().Add(1 * time.Hour),
+			},
 		},
 	}
-	validateToken := func(resp *oauth2.Token) {
-		assert.Equal(t, "new-access-token", resp.AccessToken)
-		assert.Equal(t, "Bearer", resp.TokenType)
+	validateToken := func(got *oauth2.Token) {
+		assert.Equal(t, "new-access-token", got.AccessToken)
+		assert.Equal(t, "Bearer", got.TokenType)
 	}
 
 	cases := []struct {
@@ -699,6 +715,59 @@ func TestToken_loadToken(t *testing.T) {
 			},
 			validateToken: validateToken,
 		},
+		{
+			name: "default path reuses valid cached token without refresh",
+			args: loadTokenArgs{
+				authArguments: &auth.AuthArguments{},
+				profileName:   "valid-token",
+				args:          []string{},
+				tokenTimeout:  1 * time.Hour,
+				profiler:      profiler,
+				persistentAuthOpts: []u2m.PersistentAuthOption{
+					u2m.WithTokenCache(tokenCache),
+					u2m.WithOAuthEndpointSupplier(&MockApiClient{}),
+					u2m.WithHttpClient(&http.Client{Transport: failOnCallTransport{}}),
+				},
+			},
+			validateToken: func(got *oauth2.Token) {
+				assert.Equal(t, "cached-access-token", got.AccessToken)
+			},
+		},
+		{
+			name: "force refresh refreshes valid cached token",
+			args: loadTokenArgs{
+				authArguments: &auth.AuthArguments{},
+				profileName:   "valid-token",
+				args:          []string{},
+				tokenTimeout:  1 * time.Hour,
+				forceRefresh:  true,
+				profiler:      profiler,
+				persistentAuthOpts: []u2m.PersistentAuthOption{
+					u2m.WithTokenCache(tokenCache),
+					u2m.WithOAuthEndpointSupplier(&MockApiClient{}),
+					u2m.WithHttpClient(&http.Client{Transport: fixtures.SliceTransport{refreshSuccessTokenResponse}}),
+				},
+			},
+			validateToken: validateToken,
+		},
+		{
+			name: "force refresh preserves error handling on refresh failure",
+			args: loadTokenArgs{
+				authArguments: &auth.AuthArguments{},
+				profileName:   "valid-token",
+				args:          []string{},
+				tokenTimeout:  1 * time.Hour,
+				forceRefresh:  true,
+				profiler:      profiler,
+				persistentAuthOpts: []u2m.PersistentAuthOption{
+					u2m.WithTokenCache(tokenCache),
+					u2m.WithOAuthEndpointSupplier(&MockApiClient{}),
+					u2m.WithHttpClient(&http.Client{Transport: fixtures.SliceTransport{refreshFailureTokenResponse}}),
+				},
+			},
+			wantErr: `A new access token could not be retrieved because the refresh token is invalid. To reauthenticate, run the following command:
+  $ databricks auth login --profile valid-token`,
+		},
 	}
 	for _, c := range cases {
 		t.Run(c.name, func(t *testing.T) {

go.mod

@@ -107,3 +107,5 @@ require (
 	google.golang.org/grpc v1.78.0 // indirect
 	google.golang.org/protobuf v1.36.11 // indirect
 )
+
+replace github.com/databricks/databricks-sdk-go => ../databricks-sdk-go

go.sum

@@ -75,8 +75,6 @@ github.com/creack/pty v1.1.24 h1:bJrF4RRfyJnbTJqzRLHzcGaZK1NeM5kTC9jGgovnR1s=
 github.com/creack/pty v1.1.24/go.mod h1:08sCNb52WyoAwi2QDyzUCTgcvVFhUzewun7wtTfvcwE=
 github.com/cyphar/filepath-securejoin v0.4.1 h1:JyxxyPEaktOD+GAnqIqTf9A8tHyAG22rowi7HkoSU1s=
 github.com/cyphar/filepath-securejoin v0.4.1/go.mod h1:Sdj7gXlvMcPZsbhwhQ33GguGLDGQL7h7bg04C/+u9jI=
-github.com/databricks/databricks-sdk-go v0.120.0 h1:XLEoLeVUB/MFygyklLiB2HtQTeaULnfr1RyGtYcl2gQ=
-github.com/databricks/databricks-sdk-go v0.120.0/go.mod h1:hWoHnHbNLjPKiTm5K/7bcIv3J3Pkgo5x9pPzh8K3RVE=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=