fix(bigquery): retry getTable server errors

java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java

@@ -18,11 +18,18 @@
 import static com.google.cloud.bigquery.PolicyHelper.convertFromApiPolicy;
 import static com.google.cloud.bigquery.PolicyHelper.convertToApiPolicy;
 import static com.google.common.base.Preconditions.checkArgument;
+import static java.net.HttpURLConnection.HTTP_BAD_GATEWAY;
+import static java.net.HttpURLConnection.HTTP_GATEWAY_TIMEOUT;
+import static java.net.HttpURLConnection.HTTP_INTERNAL_ERROR;
 import static java.net.HttpURLConnection.HTTP_NOT_FOUND;
+import static java.net.HttpURLConnection.HTTP_UNAVAILABLE;
 
+import com.google.api.client.http.HttpResponseException;
 import com.google.api.core.BetaApi;
 import com.google.api.core.InternalApi;
 import com.google.api.gax.paging.Page;
+import com.google.api.gax.retrying.ResultRetryAlgorithm;
+import com.google.api.gax.retrying.TimedAttemptSettings;
 import com.google.api.services.bigquery.model.ErrorProto;
 import com.google.api.services.bigquery.model.GetQueryResultsResponse;
 import com.google.api.services.bigquery.model.QueryRequest;
@@ -65,6 +72,26 @@
 
 final class BigQueryImpl extends BaseService<BigQueryOptions> implements BigQuery {
 
+  private static final ResultRetryAlgorithm<Object> DEFAULT_GET_TABLE_RETRY_ALGORITHM =
+      new ResultRetryAlgorithm<Object>() {
+        @Override
+        public TimedAttemptSettings createNextAttempt(
+            Throwable previousThrowable,
+            Object previousResponse,
+            TimedAttemptSettings previousSettings) {
+          return null;
+        }
+
+        @Override
+        public boolean shouldRetry(Throwable previousThrowable, Object previousResponse) {
+          if (isRetryableHttpResponseException(previousThrowable)) {
+            return true;
+          }
+          return BigQueryBaseService.DEFAULT_BIGQUERY_EXCEPTION_HANDLER.shouldRetry(
+              previousThrowable, previousResponse);
+        }
+      };
+
   private static class DatasetPageFetcher implements NextPageFetcher<Dataset> {
 
     private static final long serialVersionUID = -3057564042439021278L;
@@ -1131,7 +1158,7 @@ public com.google.api.services.bigquery.model.Table call() throws IOException {
                 }
               },
               getOptions().getRetrySettings(),
-              getOptions().getResultRetryAlgorithm(),
+              getTableRetryAlgorithm(),
               getOptions().getClock(),
               EMPTY_RETRY_CONFIG,
               getOptions().isOpenTelemetryTracingEnabled(),
@@ -1152,6 +1179,25 @@ public com.google.api.services.bigquery.model.Table call() throws IOException {
     }
   }
 
+  private ResultRetryAlgorithm<?> getTableRetryAlgorithm() {
+    ResultRetryAlgorithm<?> configuredAlgorithm = getOptions().getResultRetryAlgorithm();
+    if (configuredAlgorithm != BigQueryBaseService.DEFAULT_BIGQUERY_EXCEPTION_HANDLER) {
+      return configuredAlgorithm;
+    }
+    return DEFAULT_GET_TABLE_RETRY_ALGORITHM;
+  }
+
+  private static boolean isRetryableHttpResponseException(Throwable previousThrowable) {
+    if (!(previousThrowable instanceof HttpResponseException)) {
+      return false;
+    }
+    int statusCode = ((HttpResponseException) previousThrowable).getStatusCode();
+    return statusCode == HTTP_INTERNAL_ERROR
+        || statusCode == HTTP_BAD_GATEWAY
+        || statusCode == HTTP_UNAVAILABLE
+        || statusCode == HTTP_GATEWAY_TIMEOUT;
+  }
+
   @Override
   public Model getModel(String datasetId, String modelId, ModelOption... options) {
     return getModel(ModelId.of(datasetId, modelId), options);

java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java

@@ -36,7 +36,13 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import com.google.api.client.googleapis.json.GoogleJsonError;
+import com.google.api.client.googleapis.json.GoogleJsonResponseException;
+import com.google.api.client.http.HttpHeaders;
+import com.google.api.client.http.HttpResponseException;
 import com.google.api.gax.paging.Page;
+import com.google.api.gax.retrying.ResultRetryAlgorithm;
+import com.google.api.gax.retrying.TimedAttemptSettings;
 import com.google.api.services.bigquery.model.ErrorProto;
 import com.google.api.services.bigquery.model.GetQueryResultsResponse;
 import com.google.api.services.bigquery.model.JobConfigurationQuery;
@@ -72,6 +78,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -935,6 +942,86 @@ void testGetTable() throws IOException {
         .getTableSkipExceptionTranslation(PROJECT, DATASET, TABLE, EMPTY_RPC_OPTIONS);
   }
 
+  @Test
+  void testGetTableFailureShouldRetryServerErrors() throws IOException {
+    when(bigqueryRpcMock.getTableSkipExceptionTranslation(
+            PROJECT, DATASET, TABLE, EMPTY_RPC_OPTIONS))
+        .thenThrow(serviceUnavailableException())
+        .thenReturn(TABLE_INFO_WITH_PROJECT.toPb());
+
+    bigquery =
+        options.toBuilder()
+            .setRetrySettings(ServiceOptions.getDefaultRetrySettings())
+            .build()
+            .getService();
+
+    Table table = bigquery.getTable(DATASET, TABLE);
+
+    assertEquals(new Table(bigquery, new TableInfo.BuilderImpl(TABLE_INFO_WITH_PROJECT)), table);
+    verify(bigqueryRpcMock, times(2))
+        .getTableSkipExceptionTranslation(PROJECT, DATASET, TABLE, EMPTY_RPC_OPTIONS);
+  }
+
+  @Test
+  void testGetTableFailureUsesCustomRetryAlgorithm() throws IOException {
+    AtomicReference<Throwable> retryThrowable = new AtomicReference<>();
+    ResultRetryAlgorithm<Object> retryAlgorithm =
+        new ResultRetryAlgorithm<Object>() {
+          @Override
+          public TimedAttemptSettings createNextAttempt(
+              Throwable previousThrowable,
+              Object previousResponse,
+              TimedAttemptSettings previousSettings) {
+            if (previousThrowable != null) {
+              retryThrowable.set(previousThrowable);
+            }
+            return null;
+          }
+
+          @Override
+          public boolean shouldRetry(Throwable previousThrowable, Object previousResponse) {
+            if (previousThrowable != null) {
+              retryThrowable.set(previousThrowable);
+            }
+            return previousThrowable instanceof HttpResponseException;
+          }
+        };
+
+    when(bigqueryRpcMock.getTableSkipExceptionTranslation(
+            PROJECT, DATASET, TABLE, EMPTY_RPC_OPTIONS))
+        .thenThrow(serviceUnavailableException())
+        .thenReturn(TABLE_INFO_WITH_PROJECT.toPb());
+
+    bigquery =
+        options.toBuilder()
+            .setRetrySettings(ServiceOptions.getDefaultRetrySettings())
+            .setResultRetryAlgorithm(retryAlgorithm)
+            .build()
+            .getService();
+
+    assertSame(retryAlgorithm, bigquery.getOptions().getResultRetryAlgorithm());
+    Table table = bigquery.getTable(DATASET, TABLE);
+
+    assertEquals(new Table(bigquery, new TableInfo.BuilderImpl(TABLE_INFO_WITH_PROJECT)), table);
+    assertThat(retryThrowable.get()).isInstanceOf(HttpResponseException.class);
+    verify(bigqueryRpcMock, times(2))
+        .getTableSkipExceptionTranslation(PROJECT, DATASET, TABLE, EMPTY_RPC_OPTIONS);
+  }
+
+  private static GoogleJsonResponseException serviceUnavailableException() {
+    GoogleJsonError error = new GoogleJsonError();
+    error.setMessage("Visibility check was unavailable. Please retry the request");
+    error.setCode(503);
+    GoogleJsonError.ErrorInfo errorInfo = new GoogleJsonError.ErrorInfo();
+    errorInfo.setReason("backendError");
+    error.setErrors(ImmutableList.of(errorInfo));
+    return new GoogleJsonResponseException(serverErrorResponse(), error);
+  }
+
+  private static HttpResponseException.Builder serverErrorResponse() {
+    return new HttpResponseException.Builder(503, "Service Unavailable", new HttpHeaders());
+  }
+
   @Test
   void testGetModel() throws IOException {
     when(bigqueryRpcMock.getModelSkipExceptionTranslation(