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,8 @@
+Ignore = [
+    "home"
+]
+
+[[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,4 @@
+setup_test_profile
+setup_test_token_cache
+
+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/token.go

@@ -55,15 +55,20 @@
 		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 always refresh regardless of the token's remaining lifetime. 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 @@
 			profileName:        profileName,
 			args:               args,
 			tokenTimeout:       tokenTimeout,
+			forceRefresh:       forceRefresh,
 			profiler:           profile.DefaultProfiler,
 			persistentAuthOpts: nil,
 		})
@@ -107,6 +113,9 @@
 	// 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 @@
 		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

@@ -135,6 +135,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,6 +185,11 @@ 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) {
@@ -699,6 +708,58 @@ 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{}),
+				},
+			},
+			validateToken: func(resp *oauth2.Token) {
+				assert.Equal(t, "cached-access-token", resp.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) {