Retry bq upload with error code 503/500

This CL makes bigquery uploader retry the upload for error code 503 and 500 (re-triable backend error)

This CL also fix a nil pointer issue in the retry logic by adding a check to see if an error is able to cast into a googleAPI error.

Test: Unit Test
Bug: b/358569375
Change-Id: I955d474ce5d5a23087ffcbcd655712b273278d9e
GitOrigin-RevId: 2f13cf9b9b301b79e3d7600a11fbe495333980bc

internal/pkg/bigquery/BUILD.bazel

@@ -21,4 +21,8 @@ go_test(
     name = "bigquery_test",
     srcs = ["bigquery_test.go"],
     embed = [":bigquery"],
+    deps = [
+        "@com_github_eapache_go_resiliency//retrier",
+        "@org_golang_google_api//googleapi",
+    ],
 )

internal/pkg/bigquery/bigquery.go

@@ -184,15 +184,16 @@ func (b BQClassifier) Classify(err error) retrier.Action {
 		return retrier.Succeed
 	}
 	var apiError *googleapi.Error
-	errors.As(err, &apiError)
-	// For errors, such as no permission, wrong project name, wrong dataset/table name,
-	// they are considered as non-retryable errors.
-	if apiError.Code >= 400 {
-		fatalErr := fmt.Errorf("non retryable error occurred in uploading LogRecords to BigQuery: %v", err)
+	if isAPIError := errors.As(err, &apiError); isAPIError && (apiError.Code == 500 || apiError.Code == 503) {
+		return retrier.Retry
+	} else {
+		// For errors, such as no permission, wrong project name, wrong dataset/table
+		// name, they are considered as non-retryable errors. For instructions about
+		// what errors could be retried: see here: https://cloud.google.com/bigquery/docs/error-messages
+		fatalErr := fmt.Errorf("fatal error occurred in uploading LogRecords to BigQuery: %v", err)
 		b.bQSpec.Err.Store(&fatalErr)
 		return retrier.Fail
 	}
-	return retrier.Retry
 }
 
 func uploadWithRetry(bqSpec *BQSpec, items []*bigquerytranslator.Item, successful *int32, failed *int32) {

internal/pkg/bigquery/bigquery_test.go

@@ -14,7 +14,14 @@
 
 package bigquery
 
-import "testing"
+import (
+	"errors"
+	"sync/atomic"
+	"testing"
+
+	"github.com/eapache/go-resiliency/retrier"
+	"google.golang.org/api/googleapi"
+)
 
 func TestParseResourceSpec(t *testing.T) {
 	tests := []struct {
@@ -77,3 +84,40 @@ func TestParseResourceSpec(t *testing.T) {
 		})
 	}
 }
+
+func TestBQClassifier_Classify(t *testing.T) {
+	bs := &BQSpec{
+		Err: atomic.Pointer[error]{},
+	}
+	tests := []struct {
+		name  string
+		bqErr error
+		want  retrier.Action
+	}{
+		{name: "no error", bqErr: nil, want: retrier.Succeed},
+		{name: "non googleapi error", bqErr: errors.New("a non-googleapi error"), want: retrier.Fail},
+		{name: "non-retriable googleapi error", bqErr: &googleapi.Error{
+			Code:    403,
+			Message: "billingNotEnabled",
+		}, want: retrier.Fail},
+		{name: "retryiable googleapi error: internalError", bqErr: &googleapi.Error{
+			Code:    500,
+			Message: "internalError",
+		}, want: retrier.Retry},
+		{name: "retryiable googleapi error: backendError", bqErr: &googleapi.Error{
+			Code:    503,
+			Message: "backendError",
+		}, want: retrier.Retry},
+	}
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			classifier := BQClassifier{bQSpec: bs}
+			got := classifier.Classify(tc.bqErr)
+			if got != tc.want {
+				t.Errorf("Classify() = %v, want %v", got, tc.want)
+			}
+		})
+	}
+}