From ab7ed04dbfbf9a74a6ea69c39b3566accd8d4d30 Mon Sep 17 00:00:00 2001
From: Miles Yucht <miles@databricks.com>
Date: Mon, 23 Dec 2024 13:57:43 +0100
Subject: [PATCH] fix

---
 httpclient/retries.go | 1 +
 retries/retries.go    | 6 +++---
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/httpclient/retries.go b/httpclient/retries.go
index aec8e9d7..0f150a3a 100644
--- a/httpclient/retries.go
+++ b/httpclient/retries.go
@@ -7,6 +7,7 @@ import (
 func makeRetrier[T any](c ClientConfig) retries.Retrier[T] {
 	return retries.New[T](
 		retries.WithTimeout(c.RetryTimeout),
+		retries.WithRetryFunc(retries.DefaultShouldRetry),
 		retries.WithBackoffFunc(c.RetryBackoff),
 	)
 }
diff --git a/retries/retries.go b/retries/retries.go
index 9fafa960..dfec6ac0 100644
--- a/retries/retries.go
+++ b/retries/retries.go
@@ -240,7 +240,7 @@ func (r Retrier[T]) Run(ctx context.Context, fn func(context.Context) (*T, error
 	}
 }
 
-func shouldRetry(err error) bool {
+func DefaultShouldRetry(err error) bool {
 	if err == nil {
 		return false
 	}
@@ -252,7 +252,7 @@ func shouldRetry(err error) bool {
 }
 
 func Wait(ctx context.Context, timeout time.Duration, fn func() *Err) error {
-	return New[struct{}](WithTimeout(timeout), WithRetryFunc(shouldRetry)).Wait(ctx, func(_ context.Context) error {
+	return New[struct{}](WithTimeout(timeout), WithRetryFunc(DefaultShouldRetry)).Wait(ctx, func(_ context.Context) error {
 		err := fn()
 		if err != nil {
 			return err
@@ -262,7 +262,7 @@ func Wait(ctx context.Context, timeout time.Duration, fn func() *Err) error {
 }
 
 func Poll[T any](ctx context.Context, timeout time.Duration, fn func() (*T, *Err)) (*T, error) {
-	return New[T](WithTimeout(timeout), WithRetryFunc(shouldRetry)).Run(ctx, func(_ context.Context) (*T, error) {
+	return New[T](WithTimeout(timeout), WithRetryFunc(DefaultShouldRetry)).Run(ctx, func(_ context.Context) (*T, error) {
 		res, err := fn()
 		if err != nil {
 			return res, err