Skip to content

Commit b11388b

Browse files
goutamadwantgoutamadwant
committed
fix(bigquery): preserve custom getTable retry handling
1 parent ef4c18a commit b11388b

2 files changed

Lines changed: 110 additions & 17 deletions

File tree

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

Lines changed: 51 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,18 @@
1818
import static com.google.cloud.bigquery.PolicyHelper.convertFromApiPolicy;
1919
import static com.google.cloud.bigquery.PolicyHelper.convertToApiPolicy;
2020
import static com.google.common.base.Preconditions.checkArgument;
21+
import static java.net.HttpURLConnection.HTTP_BAD_GATEWAY;
22+
import static java.net.HttpURLConnection.HTTP_GATEWAY_TIMEOUT;
23+
import static java.net.HttpURLConnection.HTTP_INTERNAL_ERROR;
2124
import static java.net.HttpURLConnection.HTTP_NOT_FOUND;
25+
import static java.net.HttpURLConnection.HTTP_UNAVAILABLE;
2226

2327
import com.google.api.client.http.HttpResponseException;
2428
import com.google.api.core.BetaApi;
2529
import com.google.api.core.InternalApi;
2630
import com.google.api.gax.paging.Page;
31+
import com.google.api.gax.retrying.ResultRetryAlgorithm;
32+
import com.google.api.gax.retrying.TimedAttemptSettings;
2733
import com.google.api.services.bigquery.model.ErrorProto;
2834
import com.google.api.services.bigquery.model.GetQueryResultsResponse;
2935
import com.google.api.services.bigquery.model.QueryRequest;
@@ -66,6 +72,26 @@
6672

6773
final class BigQueryImpl extends BaseService<BigQueryOptions> implements BigQuery {
6874

75+
private static final ResultRetryAlgorithm<Object> DEFAULT_GET_TABLE_RETRY_ALGORITHM =
76+
new ResultRetryAlgorithm<Object>() {
77+
@Override
78+
public TimedAttemptSettings createNextAttempt(
79+
Throwable previousThrowable,
80+
Object previousResponse,
81+
TimedAttemptSettings previousSettings) {
82+
return null;
83+
}
84+
85+
@Override
86+
public boolean shouldRetry(Throwable previousThrowable, Object previousResponse) {
87+
if (isRetryableHttpResponseException(previousThrowable)) {
88+
return true;
89+
}
90+
return BigQueryBaseService.DEFAULT_BIGQUERY_EXCEPTION_HANDLER.shouldRetry(
91+
previousThrowable, previousResponse);
92+
}
93+
};
94+
6995
private static class DatasetPageFetcher implements NextPageFetcher<Dataset> {
7096

7197
private static final long serialVersionUID = -3057564042439021278L;
@@ -1124,19 +1150,15 @@ && getOptions().getOpenTelemetryTracer() != null) {
11241150
new Callable<com.google.api.services.bigquery.model.Table>() {
11251151
@Override
11261152
public com.google.api.services.bigquery.model.Table call() throws IOException {
1127-
try {
1128-
return bigQueryRpc.getTableSkipExceptionTranslation(
1129-
completeTableId.getProject(),
1130-
completeTableId.getDataset(),
1131-
completeTableId.getTable(),
1132-
optionsMap);
1133-
} catch (HttpResponseException e) {
1134-
throw new BigQueryException(e);
1135-
}
1153+
return bigQueryRpc.getTableSkipExceptionTranslation(
1154+
completeTableId.getProject(),
1155+
completeTableId.getDataset(),
1156+
completeTableId.getTable(),
1157+
optionsMap);
11361158
}
11371159
},
11381160
getOptions().getRetrySettings(),
1139-
getOptions().getResultRetryAlgorithm(),
1161+
getTableRetryAlgorithm(),
11401162
getOptions().getClock(),
11411163
EMPTY_RETRY_CONFIG,
11421164
getOptions().isOpenTelemetryTracingEnabled(),
@@ -1157,6 +1179,25 @@ public com.google.api.services.bigquery.model.Table call() throws IOException {
11571179
}
11581180
}
11591181

1182+
private ResultRetryAlgorithm<?> getTableRetryAlgorithm() {
1183+
ResultRetryAlgorithm<?> configuredAlgorithm = getOptions().getResultRetryAlgorithm();
1184+
if (configuredAlgorithm != BigQueryBaseService.DEFAULT_BIGQUERY_EXCEPTION_HANDLER) {
1185+
return configuredAlgorithm;
1186+
}
1187+
return DEFAULT_GET_TABLE_RETRY_ALGORITHM;
1188+
}
1189+
1190+
private static boolean isRetryableHttpResponseException(Throwable previousThrowable) {
1191+
if (!(previousThrowable instanceof HttpResponseException)) {
1192+
return false;
1193+
}
1194+
int statusCode = ((HttpResponseException) previousThrowable).getStatusCode();
1195+
return statusCode == HTTP_INTERNAL_ERROR
1196+
|| statusCode == HTTP_BAD_GATEWAY
1197+
|| statusCode == HTTP_UNAVAILABLE
1198+
|| statusCode == HTTP_GATEWAY_TIMEOUT;
1199+
}
1200+
11601201
@Override
11611202
public Model getModel(String datasetId, String modelId, ModelOption... options) {
11621203
return getModel(ModelId.of(datasetId, modelId), options);

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

Lines changed: 59 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,8 @@
4141
import com.google.api.client.http.HttpHeaders;
4242
import com.google.api.client.http.HttpResponseException;
4343
import com.google.api.gax.paging.Page;
44+
import com.google.api.gax.retrying.ResultRetryAlgorithm;
45+
import com.google.api.gax.retrying.TimedAttemptSettings;
4446
import com.google.api.services.bigquery.model.ErrorProto;
4547
import com.google.api.services.bigquery.model.GetQueryResultsResponse;
4648
import com.google.api.services.bigquery.model.JobConfigurationQuery;
@@ -76,6 +78,7 @@
7678
import java.util.Collections;
7779
import java.util.List;
7880
import java.util.Map;
81+
import java.util.concurrent.atomic.AtomicReference;
7982
import org.junit.jupiter.api.Assertions;
8083
import org.junit.jupiter.api.BeforeEach;
8184
import org.junit.jupiter.api.Test;
@@ -941,31 +944,80 @@ void testGetTable() throws IOException {
941944

942945
@Test
943946
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));
947+
when(bigqueryRpcMock.getTableSkipExceptionTranslation(
948+
PROJECT, DATASET, TABLE, EMPTY_RPC_OPTIONS))
949+
.thenThrow(serviceUnavailableException())
950+
.thenReturn(TABLE_INFO_WITH_PROJECT.toPb());
951+
952+
bigquery =
953+
options.toBuilder()
954+
.setRetrySettings(ServiceOptions.getDefaultRetrySettings())
955+
.build()
956+
.getService();
957+
958+
Table table = bigquery.getTable(DATASET, TABLE);
959+
960+
assertEquals(new Table(bigquery, new TableInfo.BuilderImpl(TABLE_INFO_WITH_PROJECT)), table);
961+
verify(bigqueryRpcMock, times(2))
962+
.getTableSkipExceptionTranslation(PROJECT, DATASET, TABLE, EMPTY_RPC_OPTIONS);
963+
}
964+
965+
@Test
966+
void testGetTableFailureUsesCustomRetryAlgorithm() throws IOException {
967+
AtomicReference<Throwable> retryThrowable = new AtomicReference<>();
968+
ResultRetryAlgorithm<Object> retryAlgorithm =
969+
new ResultRetryAlgorithm<Object>() {
970+
@Override
971+
public TimedAttemptSettings createNextAttempt(
972+
Throwable previousThrowable,
973+
Object previousResponse,
974+
TimedAttemptSettings previousSettings) {
975+
if (previousThrowable != null) {
976+
retryThrowable.set(previousThrowable);
977+
}
978+
return null;
979+
}
980+
981+
@Override
982+
public boolean shouldRetry(Throwable previousThrowable, Object previousResponse) {
983+
if (previousThrowable != null) {
984+
retryThrowable.set(previousThrowable);
985+
}
986+
return previousThrowable instanceof HttpResponseException;
987+
}
988+
};
950989

951990
when(bigqueryRpcMock.getTableSkipExceptionTranslation(
952991
PROJECT, DATASET, TABLE, EMPTY_RPC_OPTIONS))
953-
.thenThrow(new GoogleJsonResponseException(serverErrorResponse(), error))
992+
.thenThrow(serviceUnavailableException())
954993
.thenReturn(TABLE_INFO_WITH_PROJECT.toPb());
955994

956995
bigquery =
957996
options.toBuilder()
958997
.setRetrySettings(ServiceOptions.getDefaultRetrySettings())
998+
.setResultRetryAlgorithm(retryAlgorithm)
959999
.build()
9601000
.getService();
9611001

1002+
assertSame(retryAlgorithm, bigquery.getOptions().getResultRetryAlgorithm());
9621003
Table table = bigquery.getTable(DATASET, TABLE);
9631004

9641005
assertEquals(new Table(bigquery, new TableInfo.BuilderImpl(TABLE_INFO_WITH_PROJECT)), table);
1006+
assertThat(retryThrowable.get()).isInstanceOf(HttpResponseException.class);
9651007
verify(bigqueryRpcMock, times(2))
9661008
.getTableSkipExceptionTranslation(PROJECT, DATASET, TABLE, EMPTY_RPC_OPTIONS);
9671009
}
9681010

1011+
private static GoogleJsonResponseException serviceUnavailableException() {
1012+
GoogleJsonError error = new GoogleJsonError();
1013+
error.setMessage("Visibility check was unavailable. Please retry the request");
1014+
error.setCode(503);
1015+
GoogleJsonError.ErrorInfo errorInfo = new GoogleJsonError.ErrorInfo();
1016+
errorInfo.setReason("backendError");
1017+
error.setErrors(ImmutableList.of(errorInfo));
1018+
return new GoogleJsonResponseException(serverErrorResponse(), error);
1019+
}
1020+
9691021
private static HttpResponseException.Builder serverErrorResponse() {
9701022
return new HttpResponseException.Builder(503, "Service Unavailable", new HttpHeaders());
9711023
}

0 commit comments

Comments
 (0)