diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryHelper.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryHelper.java index 98adb0b273e1..04a3325072ca 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryHelper.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryHelper.java @@ -15,6 +15,7 @@ */ package com.google.cloud.bigquery; +import com.google.api.client.http.HttpResponseException; import com.google.api.core.ApiClock; import com.google.api.gax.retrying.DirectRetryingExecutor; import com.google.api.gax.retrying.ExponentialRetryAlgorithm; @@ -23,6 +24,7 @@ import com.google.api.gax.retrying.RetrySettings; import com.google.api.gax.retrying.RetryingExecutor; import com.google.api.gax.retrying.RetryingFuture; +import com.google.api.gax.retrying.TimedAttemptSettings; import com.google.api.gax.retrying.TimedRetryAlgorithm; import com.google.cloud.RetryHelper; import io.opentelemetry.api.trace.Span; @@ -69,6 +71,9 @@ public static V runWithRetries( // implementation does not use response at all, so ignoring its type is ok. @SuppressWarnings("unchecked") ResultRetryAlgorithm algorithm = (ResultRetryAlgorithm) resultRetryAlgorithm; + if (algorithm == BigQueryBaseService.DEFAULT_BIGQUERY_EXCEPTION_HANDLER) { + algorithm = wrapDefaultAlgorithm(algorithm); + } return run( callable, new ExponentialRetryAlgorithm(retrySettings, clock), @@ -119,6 +124,28 @@ private static V run( return retryingFuture.get(); } + private static ResultRetryAlgorithm wrapDefaultAlgorithm( + ResultRetryAlgorithm defaultAlgorithm) { + return new ResultRetryAlgorithm() { + @Override + public TimedAttemptSettings createNextAttempt( + Throwable previousThrowable, V previousResponse, TimedAttemptSettings previousSettings) { + return null; // Delegate timing to TimedRetryAlgorithm + } + + @Override + public boolean shouldRetry(Throwable previousThrowable, V previousResponse) { + if (previousThrowable instanceof HttpResponseException) { + int statusCode = ((HttpResponseException) previousThrowable).getStatusCode(); + if (statusCode == 500 || statusCode == 502 || statusCode == 503 || statusCode == 504) { + return true; + } + } + return defaultAlgorithm.shouldRetry(previousThrowable, previousResponse); + } + }; + } + public static class BigQueryRetryHelperException extends RuntimeException { private static final long serialVersionUID = -8519852520090965314L; diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java index 20a6ef679e89..7fea041b4025 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java @@ -36,6 +36,10 @@ 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.services.bigquery.model.ErrorProto; import com.google.api.services.bigquery.model.GetQueryResultsResponse; @@ -935,6 +939,37 @@ void testGetTable() throws IOException { .getTableSkipExceptionTranslation(PROJECT, DATASET, TABLE, EMPTY_RPC_OPTIONS); } + @Test + void testGetTableFailureShouldRetryServerErrors() throws IOException { + 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)); + + when(bigqueryRpcMock.getTableSkipExceptionTranslation( + PROJECT, DATASET, TABLE, EMPTY_RPC_OPTIONS)) + .thenThrow(new GoogleJsonResponseException(serverErrorResponse(), error)) + .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); + } + + private static HttpResponseException.Builder serverErrorResponse() { + return new HttpResponseException.Builder(503, "Service Unavailable", new HttpHeaders()); + } + @Test void testGetModel() throws IOException { when(bigqueryRpcMock.getModelSkipExceptionTranslation(