Add --force-refresh flag to auth token command

mihaimitrea-db · mihaimitrea-db · commit f78303bbc8c6 · 2026-03-17T12:17:28.000Z
Wire a new --force-refresh CLI flag that delegates to the SDK's
ForceRefreshToken() method, bypassing the cached token validity
check. The default path through Token() is unchanged.

Note: this will not compile until the SDK ships ForceRefreshToken().
diff --git a/acceptance/cmd/auth/token/force-refresh-invalid-refresh-token/out.test.toml b/acceptance/cmd/auth/token/force-refresh-invalid-refresh-token/out.test.toml
diff --git a/acceptance/cmd/auth/token/force-refresh-invalid-refresh-token/output.txt b/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
diff --git a/acceptance/cmd/auth/token/force-refresh-invalid-refresh-token/script b/acceptance/cmd/auth/token/force-refresh-invalid-refresh-token/script
@@ -0,0 +1,29 @@
+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
+
+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
+
+errcode $CLI auth token --profile test-profile --force-refresh
diff --git a/acceptance/cmd/auth/token/force-refresh-invalid-refresh-token/test.toml b/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"}'
diff --git a/acceptance/cmd/auth/token/force-refresh-no-cache/out.test.toml b/acceptance/cmd/auth/token/force-refresh-no-cache/out.test.toml
diff --git a/acceptance/cmd/auth/token/force-refresh-no-cache/output.txt b/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
diff --git a/acceptance/cmd/auth/token/force-refresh-no-cache/script b/acceptance/cmd/auth/token/force-refresh-no-cache/script
@@ -0,0 +1,14 @@
+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
+
+errcode $CLI auth token --profile test-profile --force-refresh
diff --git a/acceptance/cmd/auth/token/force-refresh-no-cache/test.toml b/acceptance/cmd/auth/token/force-refresh-no-cache/test.toml
@@ -0,0 +1,3 @@
+Ignore = [
+    "home"
+]
diff --git a/acceptance/cmd/auth/token/force-refresh-success/out.test.toml b/acceptance/cmd/auth/token/force-refresh-success/out.test.toml
diff --git a/acceptance/cmd/auth/token/force-refresh-success/output.txt b/acceptance/cmd/auth/token/force-refresh-success/output.txt
@@ -0,0 +1 @@
+"oauth-token"
diff --git a/acceptance/cmd/auth/token/force-refresh-success/script b/acceptance/cmd/auth/token/force-refresh-success/script
@@ -0,0 +1,29 @@
+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
+
+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
+
+$CLI auth token --profile test-profile --force-refresh | jq .access_token
diff --git a/acceptance/cmd/auth/token/force-refresh-success/test.toml b/acceptance/cmd/auth/token/force-refresh-success/test.toml
@@ -0,0 +1,3 @@
+Ignore = [
+    "home"
+]
diff --git a/cmd/auth/token.go b/cmd/auth/token.go
@@ -64,6 +64,10 @@ using a client ID and secret is not supported.`,
 	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 +81,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 +112,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 +261,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
diff --git a/cmd/auth/token_test.go b/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) {

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	+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
	`2`	`+`
	`3`	`+Exit code: 1`