Remove redundant retries for token refreshes as they're already handled in httpclient

config/config.go

@@ -229,6 +229,13 @@ func (c *Config) EnsureResolved() error {
 		HTTPTimeout:        time.Duration(c.HTTPTimeoutSeconds) * time.Second,
 		Transport:          c.HTTPTransport,
 		ErrorMapper:        c.refreshTokenErrorMapper,
+		TransientErrors: []string{
+			"throttled",
+			"too many requests",
+			"429",
+			"request limit exceeded",
+			"rate limit",
+		},
 	})
 	c.resolved = true
 	return nil

config/oauth_visitors.go

@@ -1,45 +1,26 @@
 package config
 
 import (
-	"context"
 	"fmt"
 	"net/http"
-	"strings"
 	"time"
 
-	"github.com/databricks/databricks-sdk-go/retries"
 	"golang.org/x/oauth2"
 )
 
-func retriableTokenSource(ctx context.Context, ts oauth2.TokenSource) (*oauth2.Token, error) {
-	return retries.Poll(ctx, 1*time.Minute, func() (*oauth2.Token, *retries.Err) {
-		token, err := ts.Token()
-		if err == nil {
-			return token, nil
-		}
-		var retryKeywords = []string{"throttled", "too many requests", "429", "request limit exceeded", "rate limit"}
-		for _, retryKeyword := range retryKeywords {
-			if strings.Contains(err.Error(), retryKeyword) {
-				return nil, retries.Continue(err)
-			}
-		}
-		return nil, retries.Halt(err)
-	})
-}
-
 // serviceToServiceVisitor returns a visitor that sets the Authorization header to the token from the auth token source
 // and the provided secondary header to the token from the secondary token source.
 func serviceToServiceVisitor(auth, secondary oauth2.TokenSource, secondaryHeader string) func(r *http.Request) error {
 	refreshableAuth := oauth2.ReuseTokenSource(nil, auth)
 	refreshableSecondary := oauth2.ReuseTokenSource(nil, secondary)
 	return func(r *http.Request) error {
-		inner, err := retriableTokenSource(r.Context(), refreshableAuth)
+		inner, err := refreshableAuth.Token()
 		if err != nil {
 			return fmt.Errorf("inner token: %w", err)
 		}
 		inner.SetAuthHeader(r)
 
-		cloud, err := retriableTokenSource(r.Context(), refreshableSecondary)
+		cloud, err := refreshableSecondary.Token()
 		if err != nil {
 			return fmt.Errorf("cloud token: %w", err)
 		}
@@ -52,7 +33,7 @@ func serviceToServiceVisitor(auth, secondary oauth2.TokenSource, secondaryHeader
 func refreshableVisitor(inner oauth2.TokenSource) func(r *http.Request) error {
 	refreshableAuth := oauth2.ReuseTokenSource(nil, inner)
 	return func(r *http.Request) error {
-		inner, err := retriableTokenSource(r.Context(), refreshableAuth)
+		inner, err := refreshableAuth.Token()
 		if err != nil {
 			return fmt.Errorf("inner token: %w", err)
 		}

config/oauth_visitors_test.go

@@ -1,8 +1,6 @@
 package config
 
 import (
-	"context"
-	"errors"
 	"testing"
 	"time"
 
@@ -18,46 +16,6 @@ func (m mockTokenSource) Token() (*oauth2.Token, error) {
 	return m.mockedTokenFunc()
 }
 
-func TestOAuthWithRetry(t *testing.T) {
-	triedOnce := false
-	mockSource := mockTokenSource{
-		mockedTokenFunc: func() (*oauth2.Token, error) {
-			if triedOnce == true {
-				return nil, errors.New("retried")
-			}
-			triedOnce = true
-			return nil, errors.New("throttled")
-		},
-	}
-	token, err := retriableTokenSource(context.Background(), mockSource)
-	assert.Nil(t, token)
-	assert.Contains(t, err.Error(), "retried")
-}
-
-func TestOAuthWithoutRetry(t *testing.T) {
-	mockSource := mockTokenSource{
-		mockedTokenFunc: func() (*oauth2.Token, error) {
-			return nil, errors.New("halt")
-		},
-	}
-	token, err := retriableTokenSource(context.Background(), mockSource)
-	assert.Nil(t, token)
-	assert.Contains(t, err.Error(), "halt")
-}
-
-func TestOAuthWithValidToken(t *testing.T) {
-	mockSource := mockTokenSource{
-		mockedTokenFunc: func() (*oauth2.Token, error) {
-			testToken := oauth2.Token{}
-			testToken.TokenType = "test"
-			return &testToken, nil
-		},
-	}
-	token, err := retriableTokenSource(context.Background(), mockSource)
-	assert.Equal(t, "test", token.TokenType)
-	assert.NoError(t, err)
-}
-
 func TestAzureReuseTokenSource(t *testing.T) {
 	mockSource := mockTokenSource{
 		mockedTokenFunc: func() (*oauth2.Token, error) {