Retry token acquisition on 429, 502, 503, and 504

Expand retriable codes to include 502 (Bad Gateway) and 504 (Gateway
Timeout) alongside 429 and 503. These are infrastructure-layer transients
from load balancers and API gateways, not application errors. 500 and 501
remain non-retriable as they indicate server bugs or missing endpoints.

This matches the retry behavior of Azure SDK, AWS SDK, and Google Cloud
client libraries for token endpoints.

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

Author: renaudhartert-db

config/experimental/auth/retrying_token_source.go

@@ -3,7 +3,9 @@ package auth
 import (
 	"context"
 	"errors"
+	"net/http"
 	"net/url"
+	"slices"
 	"strings"
 	"time"
 
@@ -50,12 +52,19 @@ type httpStatusCoder interface {
 	HTTPStatusCode() int
 }
 
+var retriableCodes = []int{
+	http.StatusTooManyRequests,    // 429
+	http.StatusBadGateway,         // 502
+	http.StatusServiceUnavailable, // 503
+	http.StatusGatewayTimeout,     // 504
+}
+
 // 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 == 503
+		return slices.Contains(retriableCodes, code)
 	}
 	var urlErr *url.Error
 	if errors.As(err, &urlErr) {

config/experimental/auth/retrying_token_source_test.go

@@ -145,11 +145,21 @@ func TestIsRetriableTokenError(t *testing.T) {
 			err:  &testHTTPError{code: 500},
 			want: false,
 		},
+		{
+			name: "http 502",
+			err:  &testHTTPError{code: 502},
+			want: true,
+		},
 		{
 			name: "http 503",
 			err:  &testHTTPError{code: 503},
 			want: true,
 		},
+		{
+			name: "http 504",
+			err:  &testHTTPError{code: 504},
+			want: true,
+		},
 		{
 			name: "http 401",
 			err:  &testHTTPError{code: 401},