Skip to content

Commit f9b53b0

Browse files
committed
fix(bigquery): wrap default ResultRetryAlgorithm to retry transient HTTP errors
1 parent fb49fb8 commit f9b53b0

2 files changed

Lines changed: 62 additions & 0 deletions

File tree

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

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
*/
1616
package com.google.cloud.bigquery;
1717

18+
import com.google.api.client.http.HttpResponseException;
1819
import com.google.api.core.ApiClock;
1920
import com.google.api.gax.retrying.DirectRetryingExecutor;
2021
import com.google.api.gax.retrying.ExponentialRetryAlgorithm;
@@ -23,6 +24,7 @@
2324
import com.google.api.gax.retrying.RetrySettings;
2425
import com.google.api.gax.retrying.RetryingExecutor;
2526
import com.google.api.gax.retrying.RetryingFuture;
27+
import com.google.api.gax.retrying.TimedAttemptSettings;
2628
import com.google.api.gax.retrying.TimedRetryAlgorithm;
2729
import com.google.cloud.RetryHelper;
2830
import io.opentelemetry.api.trace.Span;
@@ -69,6 +71,9 @@ public static <V> V runWithRetries(
6971
// implementation does not use response at all, so ignoring its type is ok.
7072
@SuppressWarnings("unchecked")
7173
ResultRetryAlgorithm<V> algorithm = (ResultRetryAlgorithm<V>) resultRetryAlgorithm;
74+
if (algorithm == BigQueryBaseService.DEFAULT_BIGQUERY_EXCEPTION_HANDLER) {
75+
algorithm = wrapDefaultAlgorithm(algorithm);
76+
}
7277
return run(
7378
callable,
7479
new ExponentialRetryAlgorithm(retrySettings, clock),
@@ -119,6 +124,28 @@ private static <V> V run(
119124
return retryingFuture.get();
120125
}
121126

127+
private static <V> ResultRetryAlgorithm<V> wrapDefaultAlgorithm(
128+
ResultRetryAlgorithm<V> defaultAlgorithm) {
129+
return new ResultRetryAlgorithm<V>() {
130+
@Override
131+
public TimedAttemptSettings createNextAttempt(
132+
Throwable previousThrowable, V previousResponse, TimedAttemptSettings previousSettings) {
133+
return null; // Delegate timing to TimedRetryAlgorithm
134+
}
135+
136+
@Override
137+
public boolean shouldRetry(Throwable previousThrowable, V previousResponse) {
138+
if (previousThrowable instanceof HttpResponseException) {
139+
int statusCode = ((HttpResponseException) previousThrowable).getStatusCode();
140+
if (statusCode == 500 || statusCode == 502 || statusCode == 503 || statusCode == 504) {
141+
return true;
142+
}
143+
}
144+
return defaultAlgorithm.shouldRetry(previousThrowable, previousResponse);
145+
}
146+
};
147+
}
148+
122149
public static class BigQueryRetryHelperException extends RuntimeException {
123150

124151
private static final long serialVersionUID = -8519852520090965314L;

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,10 @@
3636
import static org.mockito.Mockito.verify;
3737
import static org.mockito.Mockito.when;
3838

39+
import com.google.api.client.googleapis.json.GoogleJsonError;
40+
import com.google.api.client.googleapis.json.GoogleJsonResponseException;
41+
import com.google.api.client.http.HttpHeaders;
42+
import com.google.api.client.http.HttpResponseException;
3943
import com.google.api.gax.paging.Page;
4044
import com.google.api.services.bigquery.model.ErrorProto;
4145
import com.google.api.services.bigquery.model.GetQueryResultsResponse;
@@ -935,6 +939,37 @@ void testGetTable() throws IOException {
935939
.getTableSkipExceptionTranslation(PROJECT, DATASET, TABLE, EMPTY_RPC_OPTIONS);
936940
}
937941

942+
@Test
943+
void testGetTableFailureShouldRetryServerErrors() throws IOException {
944+
GoogleJsonError error = new GoogleJsonError();
945+
error.setMessage("Visibility check was unavailable. Please retry the request");
946+
error.setCode(503);
947+
GoogleJsonError.ErrorInfo errorInfo = new GoogleJsonError.ErrorInfo();
948+
errorInfo.setReason("backendError");
949+
error.setErrors(ImmutableList.of(errorInfo));
950+
951+
when(bigqueryRpcMock.getTableSkipExceptionTranslation(
952+
PROJECT, DATASET, TABLE, EMPTY_RPC_OPTIONS))
953+
.thenThrow(new GoogleJsonResponseException(serverErrorResponse(), error))
954+
.thenReturn(TABLE_INFO_WITH_PROJECT.toPb());
955+
956+
bigquery =
957+
options.toBuilder()
958+
.setRetrySettings(ServiceOptions.getDefaultRetrySettings())
959+
.build()
960+
.getService();
961+
962+
Table table = bigquery.getTable(DATASET, TABLE);
963+
964+
assertEquals(new Table(bigquery, new TableInfo.BuilderImpl(TABLE_INFO_WITH_PROJECT)), table);
965+
verify(bigqueryRpcMock, times(2))
966+
.getTableSkipExceptionTranslation(PROJECT, DATASET, TABLE, EMPTY_RPC_OPTIONS);
967+
}
968+
969+
private static HttpResponseException.Builder serverErrorResponse() {
970+
return new HttpResponseException.Builder(503, "Service Unavailable", new HttpHeaders());
971+
}
972+
938973
@Test
939974
void testGetModel() throws IOException {
940975
when(bigqueryRpcMock.getModelSkipExceptionTranslation(

0 commit comments

Comments
 (0)