Add retry logic to token acquisition for auth credentials

Token acquisition for OIDC, M2M OAuth, and Azure client secret
credentials now retries on transient failures (429, 5xx, network errors).

Retries happen at the token source level rather than the HTTP transport
level. The existing RoundTrip implementation on ApiClient violates the
standard http.RoundTripper contract by retrying requests and interpreting
HTTP status codes -- an anti-pattern that arguably led to #1398 in the
first place, since the transport-level retry cannot rewind request bodies
created by the oauth2 library. By retrying at the application level,
each attempt creates a fresh HTTP request, sidestepping the body-reset
problem entirely.

Fixes #1398
Fixes #1072

Co-authored-by: Isaac
Signed-off-by: Ubuntu <renaud.hartert@databricks.com>

Author: renaudhartert-db

NEXT_CHANGELOG.md

@@ -13,6 +13,7 @@
 
 ### Bug Fixes
 
+ * Add retry logic to token acquisition for OIDC, M2M, and Azure client secret credentials ([#1398](https://github.com/databricks/databricks-sdk-go/issues/1398), [#1072](https://github.com/databricks/databricks-sdk-go/issues/1072)).
  * Fix double-caching of OAuth tokens in Azure client secret credentials ([#1549](https://github.com/databricks/databricks-sdk-go/issues/1549)).
  * Disable async token refresh for GCP credential providers to avoid wasted refresh attempts caused by double-caching with Google's internal `oauth2.ReuseTokenSource` ([#1549](https://github.com/databricks/databricks-sdk-go/issues/1549)).
  * Fixed double-caching in M2M OAuth that prevented the proactive async token refresh from reaching the HTTP endpoint until ~10s before expiry, causing bursts of 401 errors at token rotation boundaries ([#1549](https://github.com/databricks/databricks-sdk-go/issues/1549)).

config/auth_azure_client_secret.go

@@ -64,8 +64,8 @@ func (c AzureClientSecretCredentials) Configure(ctx context.Context, cfg *Config
 	aadEndpoint := env.AzureActiveDirectoryEndpoint()
 	managementEndpoint := env.AzureServiceManagementEndpoint()
 	opts := cacheOptions(cfg)
-	inner := azureReuseTokenSource(nil, c.tokenSourceFor(ctx, cfg, aadEndpoint, env.AzureApplicationID), opts...)
-	management := azureReuseTokenSource(nil, c.tokenSourceFor(ctx, cfg, aadEndpoint, managementEndpoint), opts...)
+	inner := azureReuseTokenSource(nil, auth.NewRetryingTokenSource(c.tokenSourceFor(ctx, cfg, aadEndpoint, env.AzureApplicationID)), opts...)
+	management := azureReuseTokenSource(nil, auth.NewRetryingTokenSource(c.tokenSourceFor(ctx, cfg, aadEndpoint, managementEndpoint)), opts...)
 	visitor := azureVisitor(cfg, serviceToServiceVisitor(inner, management, xDatabricksAzureSpManagementToken, false, opts...))
 	return newVisitorOAuthCredentials(visitor, inner), nil
 }

config/auth_m2m.go

@@ -50,6 +50,6 @@ func (c M2mCredentials) Configure(ctx context.Context, cfg *Config) (credentials
 		return ccfg.Token(ctx)
 	})
 	return credentials.NewOAuthCredentialsProviderFromTokenSource(
-		auth.NewCachedTokenSource(directTS, cacheOptions(cfg)...),
+		auth.NewCachedTokenSource(auth.NewRetryingTokenSource(directTS), cacheOptions(cfg)...),
 	), nil
 }

config/auth_oidc.go

@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"github.com/databricks/databricks-sdk-go/config/credentials"
+	"github.com/databricks/databricks-sdk-go/config/experimental/auth"
 	"github.com/databricks/databricks-sdk-go/config/experimental/auth/oidc"
 )
 
@@ -100,7 +101,7 @@ func oidcStrategy(cfg *Config, name string, ts oidc.IDTokenSource) CredentialsSt
 	}
 	oidcConfig.SetScopes(cfg.GetScopes())
 	tokenSource := oidc.NewDatabricksOIDCTokenSource(oidcConfig)
-	return NewTokenSourceStrategy(name, tokenSource)
+	return NewTokenSourceStrategy(name, auth.NewRetryingTokenSource(tokenSource))
 }
 
 // failedStrategy is a CredentialsStrategy that always fails.

config/experimental/auth/retrying_token_source.go

@@ -0,0 +1,95 @@
+package auth
+
+import (
+	"context"
+	"errors"
+	"net/url"
+	"strings"
+	"time"
+
+	"github.com/databricks/databricks-sdk-go/experimental/api"
+	"golang.org/x/oauth2"
+)
+
+// retryingTokenSource wraps a TokenSource with retry logic for transient
+// failures during token acquisition. Each retry calls the underlying Token()
+// method, which creates a fresh HTTP request -- avoiding the body-rewinding
+// problems that occur when retrying at the transport level.
+type retryingTokenSource struct {
+	inner TokenSource
+	opts  []api.Option
+}
+
+var defaultOptions = []api.Option{
+	api.WithTimeout(1 * time.Minute),
+	api.WithRetrier(func() api.Retrier {
+		return api.RetryOn(api.BackoffPolicy{}, isRetriableTokenError)
+	}),
+}
+
+// NewRetryingTokenSource wraps inner with retry logic for transient failures.
+// By default it uses exponential backoff with a 1-minute timeout and a 30-second maximum delay.
+// Additional api.Option values can be provided to override the defaults.
+func NewRetryingTokenSource(inner TokenSource, opts ...api.Option) TokenSource {
+	return &retryingTokenSource{
+		inner: inner,
+		opts:  append(defaultOptions, opts...),
+	}
+}
+
+// Token returns a token from the underlying source, retrying on transient
+// errors.
+func (r *retryingTokenSource) Token(ctx context.Context) (*oauth2.Token, error) {
+	return api.ExecuteWithResult(ctx, r.inner.Token, r.opts...)
+}
+
+// httpStatusCoder is implemented by errors that carry an HTTP status code.
+// This interface avoids importing httpclient (which would create a cycle)
+// while still allowing to classify httpclient.HttpError by status code.
+type httpStatusCoder interface {
+	HTTPStatusCode() int
+}
+
+// isRetriableTokenError returns true if the error is a transient failure that
+// should be retried. This covers HTTP errors from the SDK's transport layer,
+// OAuth2 token endpoint errors, and transient network errors.
+func isRetriableTokenError(err error) bool {
+	if code := httpStatusCode(err); code != 0 {
+		return code == 429 || code >= 500
+	}
+	var urlErr *url.Error
+	if errors.As(err, &urlErr) {
+		msg := urlErr.Error()
+		for _, s := range transientNetworkErrors {
+			if strings.Contains(msg, s) {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+// httpStatusCode extracts the HTTP status code from an error, if available.
+func httpStatusCode(err error) int {
+	// Check oauth2.RetrieveError (has Response field with StatusCode).
+	var retrieveErr *oauth2.RetrieveError
+	if errors.As(err, &retrieveErr) && retrieveErr.Response != nil {
+		return retrieveErr.Response.StatusCode
+	}
+	// Check for any error that exposes a StatusCode() method.
+	var sc httpStatusCoder
+	if errors.As(err, &sc) {
+		return sc.HTTPStatusCode()
+	}
+	return 0
+}
+
+// transientNetworkErrors is the list of error substrings that indicate a
+// transient network failure. These mirror the checks in
+// httpclient/errors.go isRetriableUrlError.
+var transientNetworkErrors = []string{
+	"connection reset by peer",
+	"TLS handshake timeout",
+	"connection refused",
+	"i/o timeout",
+}

config/experimental/auth/retrying_token_source_test.go

@@ -0,0 +1,218 @@
+package auth
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/http"
+	"net/url"
+	"testing"
+	"time"
+
+	"golang.org/x/oauth2"
+)
+
+// testHTTPError is a test error that implements the httpStatusCoder interface,
+// mirroring httpclient.HttpError without importing it.
+type testHTTPError struct {
+	code    int
+	message string
+}
+
+func (e *testHTTPError) Error() string       { return fmt.Sprintf("http %d: %s", e.code, e.message) }
+func (e *testHTTPError) HTTPStatusCode() int { return e.code }
+
+func TestRetryingTokenSource(t *testing.T) {
+	validToken := &oauth2.Token{
+		AccessToken: "test-token",
+		Expiry:      time.Now().Add(time.Hour),
+	}
+
+	testCases := []struct {
+		name       string
+		callErrors []error
+		wantToken  bool
+		wantErr    bool
+		wantCalls  int
+	}{
+		{
+			name:       "success on first call",
+			callErrors: []error{nil},
+			wantToken:  true,
+			wantCalls:  1,
+		},
+		{
+			name: "retry on http 429 then succeed",
+			callErrors: []error{
+				&testHTTPError{code: 429, message: "rate limited"},
+				nil,
+			},
+			wantToken: true,
+			wantCalls: 2,
+		},
+		{
+			name: "retry on http 500 then succeed",
+			callErrors: []error{
+				&testHTTPError{code: 500, message: "server error"},
+				nil,
+			},
+			wantToken: true,
+			wantCalls: 2,
+		},
+		{
+			name: "retry on oauth2 retrieve error 429",
+			callErrors: []error{
+				&oauth2.RetrieveError{Response: &http.Response{StatusCode: 429}},
+				nil,
+			},
+			wantToken: true,
+			wantCalls: 2,
+		},
+		{
+			name: "retry on transient network error",
+			callErrors: []error{
+				&url.Error{Op: "Post", URL: "https://host/token", Err: fmt.Errorf("connection reset by peer")},
+				nil,
+			},
+			wantToken: true,
+			wantCalls: 2,
+		},
+		{
+			name: "no retry on http 401",
+			callErrors: []error{
+				&testHTTPError{code: 401, message: "unauthorized"},
+			},
+			wantErr:   true,
+			wantCalls: 1,
+		},
+		{
+			name: "no retry on http 400",
+			callErrors: []error{
+				&testHTTPError{code: 400, message: "bad request"},
+			},
+			wantErr:   true,
+			wantCalls: 1,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			callCount := 0
+			inner := TokenSourceFn(func(ctx context.Context) (*oauth2.Token, error) {
+				err := tc.callErrors[callCount]
+				callCount++
+				if err != nil {
+					return nil, err
+				}
+				return validToken, nil
+			})
+
+			ts := NewRetryingTokenSource(inner)
+			tok, err := ts.Token(context.Background())
+
+			if callCount != tc.wantCalls {
+				t.Errorf("got %d calls, want %d", callCount, tc.wantCalls)
+			}
+			if tc.wantErr && err == nil {
+				t.Error("got nil error, want error")
+			}
+			if !tc.wantErr && err != nil {
+				t.Errorf("got error %v, want nil", err)
+			}
+			if tc.wantToken && tok == nil {
+				t.Error("got nil token, want token")
+			}
+			if !tc.wantToken && tok != nil {
+				t.Errorf("got token %v, want nil", tok)
+			}
+		})
+	}
+}
+
+func TestIsRetriableTokenError(t *testing.T) {
+	testCases := []struct {
+		name string
+		err  error
+		want bool
+	}{
+		{
+			name: "http 429",
+			err:  &testHTTPError{code: 429},
+			want: true,
+		},
+		{
+			name: "http 500",
+			err:  &testHTTPError{code: 500},
+			want: true,
+		},
+		{
+			name: "http 503",
+			err:  &testHTTPError{code: 503},
+			want: true,
+		},
+		{
+			name: "http 401",
+			err:  &testHTTPError{code: 401},
+			want: false,
+		},
+		{
+			name: "http 403",
+			err:  &testHTTPError{code: 403},
+			want: false,
+		},
+		{
+			name: "oauth2 retrieve error 429",
+			err:  &oauth2.RetrieveError{Response: &http.Response{StatusCode: 429}},
+			want: true,
+		},
+		{
+			name: "oauth2 retrieve error 500",
+			err:  &oauth2.RetrieveError{Response: &http.Response{StatusCode: 500}},
+			want: true,
+		},
+		{
+			name: "oauth2 retrieve error 400",
+			err:  &oauth2.RetrieveError{Response: &http.Response{StatusCode: 400}},
+			want: false,
+		},
+		{
+			name: "connection reset",
+			err:  &url.Error{Op: "Post", URL: "https://host/token", Err: fmt.Errorf("connection reset by peer")},
+			want: true,
+		},
+		{
+			name: "tls handshake timeout",
+			err:  &url.Error{Op: "Post", URL: "https://host/token", Err: fmt.Errorf("TLS handshake timeout")},
+			want: true,
+		},
+		{
+			name: "connection refused",
+			err:  &url.Error{Op: "Post", URL: "https://host/token", Err: fmt.Errorf("connection refused")},
+			want: true,
+		},
+		{
+			name: "i/o timeout",
+			err:  &url.Error{Op: "Post", URL: "https://host/token", Err: fmt.Errorf("i/o timeout")},
+			want: true,
+		},
+		{
+			name: "url error non-transient",
+			err:  &url.Error{Op: "Post", URL: "https://host/token", Err: fmt.Errorf("no such host")},
+			want: false,
+		},
+		{
+			name: "generic error",
+			err:  errors.New("something went wrong"),
+			want: false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			got := isRetriableTokenError(tc.err)
+			if got != tc.want {
+				t.Errorf("isRetriableTokenError(%v) = %v, want %v", tc.err, got, tc.want)
+			}
+		})
+	}
+}

experimental/api/api.go

@@ -22,6 +22,23 @@ func Execute(ctx context.Context, call Call, opts ...Option) error {
 	return execute(ctx, call, options, sleep)
 }
 
+// ExecuteWithResult returns the result of calling op with the given options.
+// It is a convenience wrapper around Execute for operations that return a
+// value. In case of error, the zero value of T is returned.
+func ExecuteWithResult[T any](ctx context.Context, op func(context.Context) (T, error), opts ...Option) (T, error) {
+	var result T
+	err := Execute(ctx, func(ctx context.Context) error {
+		var err error
+		result, err = op(ctx)
+		return err
+	}, opts...)
+	if err != nil {
+		var zero T // guarantee zero value on error
+		return zero, err
+	}
+	return result, nil
+}
+
 // sleep sleeps for the given duration. It is mostly equivalent to time.Sleep,
 // but can be interrupted by ctx.Done() if the context completes before the
 // duration elapses.