diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java index 74c9ce60e84f..55cd7e4e6177 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java +++ b/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 implements BigQuery { + private static final ResultRetryAlgorithm DEFAULT_GET_TABLE_RETRY_ALGORITHM = + new ResultRetryAlgorithm() { + @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 { 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); 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..8a908e2e8011 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,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 retryThrowable = new AtomicReference<>(); + ResultRetryAlgorithm retryAlgorithm = + new ResultRetryAlgorithm() { + @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(