From d4af6f7707b3c5ee12cde53c7485a9b743034119 Mon Sep 17 00:00:00 2001
From: shollyman <shollyman@google.com>
Date: Mon, 3 May 2021 16:38:44 -0700
Subject: [PATCH] feat(bigquery): augment retry predicate to support additional
 errors (#4046)

* feat(bigquery): augment retry predicate to support additional errors.

Based on similar refactor in storage:
https://github.com/googleapis/google-cloud-go/pull/4019

Fixes: https://github.com/googleapis/google-cloud-go/issues/4021
---
 bigquery/bigquery.go      | 38 +++++++++++++++++++++++++++++++-------
 bigquery/bigquery_test.go | 34 ++++++++++++++++++++++++++++++++++
 bigquery/go.mod           |  1 +
 3 files changed, 66 insertions(+), 7 deletions(-)

diff --git a/bigquery/bigquery.go b/bigquery/bigquery.go
index 1beb248b1fb..0a32f02c3d5 100644
--- a/bigquery/bigquery.go
+++ b/bigquery/bigquery.go
@@ -19,6 +19,8 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"net/url"
+	"strings"
 	"time"
 
 	"cloud.google.com/go/internal"
@@ -169,6 +171,9 @@ func retryableError(err error) bool {
 	if err == nil {
 		return false
 	}
+	if err == io.ErrUnexpectedEOF {
+		return true
+	}
 	// Special case due to http2: https://github.com/googleapis/google-cloud-go/issues/1793
 	// Due to Go's default being higher for streams-per-connection than is accepted by the
 	// BQ backend, it's possible to get streams refused immediately after a connection is
@@ -177,13 +182,32 @@ func retryableError(err error) bool {
 	if err.Error() == "http2: stream closed" {
 		return true
 	}
-	e, ok := err.(*googleapi.Error)
-	if !ok {
-		return false
+
+	switch e := err.(type) {
+	case *googleapi.Error:
+		// We received a structured error from backend.
+		var reason string
+		if len(e.Errors) > 0 {
+			reason = e.Errors[0].Reason
+		}
+		if e.Code == http.StatusServiceUnavailable || e.Code == http.StatusBadGateway || reason == "backendError" || reason == "rateLimitExceeded" {
+			return true
+		}
+	case *url.Error:
+		retryable := []string{"connection refused", "connection reset"}
+		for _, s := range retryable {
+			if strings.Contains(e.Error(), s) {
+				return true
+			}
+		}
+	case interface{ Temporary() bool }:
+		if e.Temporary() {
+			return true
+		}
 	}
-	var reason string
-	if len(e.Errors) > 0 {
-		reason = e.Errors[0].Reason
+	// Unwrap is only supported in go1.13.x+
+	if e, ok := err.(interface{ Unwrap() error }); ok {
+		return retryableError(e.Unwrap())
 	}
-	return e.Code == http.StatusServiceUnavailable || e.Code == http.StatusBadGateway || reason == "backendError" || reason == "rateLimitExceeded"
+	return false
 }
diff --git a/bigquery/bigquery_test.go b/bigquery/bigquery_test.go
index c5a05aadbc2..2a7d58a29a6 100644
--- a/bigquery/bigquery_test.go
+++ b/bigquery/bigquery_test.go
@@ -16,9 +16,12 @@ package bigquery
 
 import (
 	"errors"
+	"io"
 	"net/http"
+	"net/url"
 	"testing"
 
+	"golang.org/x/xerrors"
 	"google.golang.org/api/googleapi"
 )
 
@@ -38,6 +41,11 @@ func TestRetryableErrors(t *testing.T) {
 			errors.New("http2: stream closed"),
 			true,
 		},
+		{
+			"io ErrUnexpectedEOF",
+			io.ErrUnexpectedEOF,
+			true,
+		},
 		{
 			"unavailable",
 			&googleapi.Error{
@@ -46,6 +54,32 @@ func TestRetryableErrors(t *testing.T) {
 			},
 			true,
 		},
+		{
+			"url connection error",
+			&url.Error{Op: "blah", URL: "blah", Err: errors.New("connection refused")},
+			true,
+		},
+		{
+			"url other error",
+			&url.Error{Op: "blah", URL: "blah", Err: errors.New("blah")},
+			false,
+		},
+		{
+			"wrapped retryable",
+			xerrors.Errorf("test of wrapped retryable: %w", &googleapi.Error{
+				Code:    http.StatusServiceUnavailable,
+				Message: "foo",
+				Errors: []googleapi.ErrorItem{
+					{Reason: "backendError", Message: "foo"},
+				},
+			}),
+			true,
+		},
+		{
+			"wrapped non-retryable",
+			xerrors.Errorf("test of wrapped retryable: %w", errors.New("blah")),
+			false,
+		},
 		{
 			// not retried per https://google.aip.dev/194
 			"internal error",
diff --git a/bigquery/go.mod b/bigquery/go.mod
index 357175510d2..2f15230fa17 100644
--- a/bigquery/go.mod
+++ b/bigquery/go.mod
@@ -8,6 +8,7 @@ require (
 	github.com/golang/protobuf v1.5.2
 	github.com/google/go-cmp v0.5.5
 	github.com/googleapis/gax-go/v2 v2.0.5
+	golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1
 	google.golang.org/api v0.46.0
 	google.golang.org/genproto v0.0.0-20210503173045-b96a97608f20
 	google.golang.org/grpc v1.37.0