Skip to content

Commit 83245b9

Browse files
PhongChuongPhongChuong
authored
fix: retry ExceptionHandler not retrying on IOException (#3668)
* fix: retry ExceptionHandler not retrying on IOException due to exception translation * fix: retry ExceptionHandler not retrying on IOException due to exception translation * fix: retry ExceptionHandler not retrying on IOException due to exception translation * fix: retry ExceptionHandler not retrying on IOException due to exception translation * fix: retry ExceptionHandler not retrying on IOException due to exception translation * cast to HttpBigQueryRpc * cast to HttpBigQueryRpc * fix BigQueryRpc mocks in unit tests * fix clirr check and format * fix clirr * fix clirr * refactor HttpBigQueryRpc.write * refactor HttpBigQueryRpc.write * refactor ConnectionImpl HttpBigQueryRpc call * refactor ConnectionImpl HttpBigQueryRpc call * refactor create table * refactor create table * refactor BigQueryImpl * refactor ConnectionImpl * add missing unit test for BigQueryImpl deleteJob * add tests to validate retry on BIGQUERY_EXCEPTION_HANDLER * handle exception wrapping/unwrapping with new methods * update runWithRetries to use BigQuery ruWithRetries to correctly handle thrown IOExceptions * fix unit test * add unit tests for retry in TableDataWriteChannel * remove unnecessary null check * fix unit test where null is returned in mocks causing errors
1 parent 11b5809 commit 83245b9

11 files changed

+1791
-1138
lines changed

‎google-cloud-bigquery/clirr-ignored-differences.xml

+7
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,13 @@
1414
<method>com.google.api.services.bigquery.model.GetQueryResultsResponse getQueryResultsWithRowLimit(java.lang.String, java.lang.String, java.lang.String, java.lang.Integer)</method>
1515
<justification>getQueryResultsWithRowLimit is just used by ConnectionImpl at the moment so it should be fine to update the signature instead of writing an overloaded method</justification>
1616
</difference>
17+
<difference>
18+
<differenceType>7006</differenceType>
19+
<className>com/google/cloud/bigquery/BigQueryOptions*</className>
20+
<method>*getBigQueryRpcV2(*)</method>
21+
<to>com.google.cloud.bigquery.spi.v2.HttpBigQueryRpc</to>
22+
<justification>getBigQueryRpcV2 is protected and is only used within the BigQuery package</justification>
23+
</difference>
1724
<difference>
1825
<differenceType>7013</differenceType>
1926
<className>com/google/cloud/bigquery/ExternalTableDefinition*</className>

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

+1
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ protected BigQueryBaseService(ServiceOptions options) {
3131
.abortOn(RuntimeException.class)
3232
.retryOn(java.net.ConnectException.class) // retry on Connection Exception
3333
.retryOn(java.net.UnknownHostException.class) // retry on UnknownHostException
34+
.retryOn(java.net.SocketException.class) // retry on SocketException
3435
.addInterceptors(EXCEPTION_HANDLER_INTERCEPTOR)
3536
.build();
3637
}

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

+273-188
Large diffs are not rendered by default.

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

+2-3
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121
import com.google.cloud.ServiceRpc;
2222
import com.google.cloud.TransportOptions;
2323
import com.google.cloud.bigquery.spi.BigQueryRpcFactory;
24-
import com.google.cloud.bigquery.spi.v2.BigQueryRpc;
2524
import com.google.cloud.bigquery.spi.v2.HttpBigQueryRpc;
2625
import com.google.cloud.http.HttpTransportOptions;
2726
import com.google.common.annotations.VisibleForTesting;
@@ -132,8 +131,8 @@ protected Set<String> getScopes() {
132131
return SCOPES;
133132
}
134133

135-
protected BigQueryRpc getBigQueryRpcV2() {
136-
return (BigQueryRpc) getRpc();
134+
protected HttpBigQueryRpc getBigQueryRpcV2() {
135+
return (HttpBigQueryRpc) getRpc();
137136
}
138137

139138
public String getLocation() {

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

+6
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import com.google.api.gax.retrying.RetryingFuture;
2626
import com.google.api.gax.retrying.TimedRetryAlgorithm;
2727
import com.google.cloud.RetryHelper;
28+
import java.io.IOException;
2829
import java.util.concurrent.Callable;
2930
import java.util.concurrent.ExecutionException;
3031
import java.util.logging.Level;
@@ -52,6 +53,11 @@ public static <V> V runWithRetries(
5253
algorithm,
5354
bigQueryRetryConfig);
5455
} catch (Exception e) {
56+
// Checks for IOException and translate it into BigQueryException. The BigQueryException
57+
// constructor parses the IOException and translate it into internal code.
58+
if (e.getCause() instanceof IOException) {
59+
throw new BigQueryRetryHelperException(new BigQueryException((IOException) e.getCause()));
60+
}
5561
throw new BigQueryRetryHelperException(e.getCause());
5662
}
5763
}

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

+48-26
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616

1717
package com.google.cloud.bigquery;
1818

19-
import static com.google.cloud.RetryHelper.runWithRetries;
2019
import static java.net.HttpURLConnection.HTTP_NOT_FOUND;
2120

2221
import com.google.api.core.BetaApi;
@@ -28,8 +27,8 @@
2827
import com.google.api.services.bigquery.model.QueryRequest;
2928
import com.google.api.services.bigquery.model.TableDataList;
3029
import com.google.api.services.bigquery.model.TableRow;
31-
import com.google.cloud.RetryHelper;
3230
import com.google.cloud.Tuple;
31+
import com.google.cloud.bigquery.BigQueryRetryHelper.BigQueryRetryHelperException;
3332
import com.google.cloud.bigquery.JobStatistics.QueryStatistics;
3433
import com.google.cloud.bigquery.JobStatistics.SessionInfo;
3534
import com.google.cloud.bigquery.spi.v2.BigQueryRpc;
@@ -102,6 +101,8 @@ class ConnectionImpl implements Connection {
102101
bufferFvl; // initialized lazily iff we end up using the tabledata.list end point
103102
private BlockingQueue<BigQueryResultImpl.Row>
104103
bufferRow; // initialized lazily iff we end up using Read API
104+
private static final BigQueryRetryConfig EMPTY_RETRY_CONFIG =
105+
BigQueryRetryConfig.newBuilder().build();
105106

106107
ConnectionImpl(
107108
ConnectionSettings connectionSettings,
@@ -466,12 +467,15 @@ private BigQueryResult queryRpc(
466467
try {
467468
results =
468469
BigQueryRetryHelper.runWithRetries(
469-
() -> bigQueryRpc.queryRpc(projectId, queryRequest),
470+
() ->
471+
bigQueryOptions
472+
.getBigQueryRpcV2()
473+
.queryRpcSkipExceptionTranslation(projectId, queryRequest),
470474
bigQueryOptions.getRetrySettings(),
471475
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
472476
bigQueryOptions.getClock(),
473477
retryConfig);
474-
} catch (BigQueryRetryHelper.BigQueryRetryHelperException e) {
478+
} catch (BigQueryRetryHelperException e) {
475479
throw BigQueryException.translateAndThrow(e);
476480
}
477481

@@ -914,21 +918,30 @@ private Job getQueryJobRpc(JobId jobId) {
914918
com.google.api.services.bigquery.model.Job jobPb;
915919
try {
916920
jobPb =
917-
runWithRetries(
921+
BigQueryRetryHelper.runWithRetries(
918922
() ->
919-
bigQueryRpc.getQueryJob(
920-
completeJobId.getProject(),
921-
completeJobId.getJob(),
922-
completeJobId.getLocation()),
923+
bigQueryOptions
924+
.getBigQueryRpcV2()
925+
.getQueryJobSkipExceptionTranslation(
926+
completeJobId.getProject(),
927+
completeJobId.getJob(),
928+
completeJobId.getLocation()),
923929
bigQueryOptions.getRetrySettings(),
924930
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
925-
bigQueryOptions.getClock());
926-
if (bigQueryOptions.getThrowNotFound() && jobPb == null) {
927-
throw new BigQueryException(HTTP_NOT_FOUND, "Query job not found");
931+
bigQueryOptions.getClock(),
932+
EMPTY_RETRY_CONFIG);
933+
} catch (BigQueryRetryHelperException e) {
934+
if (e.getCause() instanceof BigQueryException) {
935+
if (((BigQueryException) e.getCause()).getCode() == HTTP_NOT_FOUND) {
936+
if (bigQueryOptions.getThrowNotFound()) {
937+
throw new BigQueryException(HTTP_NOT_FOUND, "Query job not found");
938+
}
939+
return null;
940+
}
928941
}
929-
} catch (RetryHelper.RetryHelperException e) {
930942
throw BigQueryException.translateAndThrow(e);
931943
}
944+
// getQueryJobSkipExceptionTranslation will never return null so this is safe.
932945
return Job.fromPb(bigQueryOptions.getService(), jobPb);
933946
}
934947

@@ -948,22 +961,23 @@ TableDataList tableDataListRpc(TableId destinationTable, String pageToken) {
948961
? bigQueryOptions.getProjectId()
949962
: destinationTable.getProject());
950963
TableDataList results =
951-
runWithRetries(
964+
BigQueryRetryHelper.runWithRetries(
952965
() ->
953966
bigQueryOptions
954967
.getBigQueryRpcV2()
955-
.listTableDataWithRowLimit(
968+
.listTableDataWithRowLimitSkipExceptionTranslation(
956969
completeTableId.getProject(),
957970
completeTableId.getDataset(),
958971
completeTableId.getTable(),
959972
connectionSettings.getMaxResultPerPage(),
960973
pageToken),
961974
bigQueryOptions.getRetrySettings(),
962975
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
963-
bigQueryOptions.getClock());
976+
bigQueryOptions.getClock(),
977+
EMPTY_RETRY_CONFIG);
964978

965979
return results;
966-
} catch (RetryHelper.RetryHelperException e) {
980+
} catch (BigQueryRetryHelperException e) {
967981
throw BigQueryException.translateAndThrow(e);
968982
}
969983
}
@@ -1177,12 +1191,14 @@ GetQueryResultsResponse getQueryResultsFirstPage(JobId jobId) {
11771191
results =
11781192
BigQueryRetryHelper.runWithRetries(
11791193
() ->
1180-
bigQueryRpc.getQueryResultsWithRowLimit(
1181-
completeJobId.getProject(),
1182-
completeJobId.getJob(),
1183-
completeJobId.getLocation(),
1184-
connectionSettings.getMaxResultPerPage(),
1185-
timeoutMs),
1194+
bigQueryOptions
1195+
.getBigQueryRpcV2()
1196+
.getQueryResultsWithRowLimitSkipExceptionTranslation(
1197+
completeJobId.getProject(),
1198+
completeJobId.getJob(),
1199+
completeJobId.getLocation(),
1200+
connectionSettings.getMaxResultPerPage(),
1201+
timeoutMs),
11861202
bigQueryOptions.getRetrySettings(),
11871203
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
11881204
bigQueryOptions.getClock(),
@@ -1197,7 +1213,7 @@ GetQueryResultsResponse getQueryResultsFirstPage(JobId jobId) {
11971213
// with the case where there is a HTTP error
11981214
throw new BigQueryException(bigQueryErrors);
11991215
}
1200-
} catch (BigQueryRetryHelper.BigQueryRetryHelperException e) {
1216+
} catch (BigQueryRetryHelperException e) {
12011217
logger.log(Level.WARNING, "\n Error occurred while calling getQueryResultsWithRowLimit", e);
12021218
throw BigQueryException.translateAndThrow(e);
12031219
}
@@ -1442,7 +1458,10 @@ com.google.api.services.bigquery.model.Job createQueryJob(
14421458
try {
14431459
queryJob =
14441460
BigQueryRetryHelper.runWithRetries(
1445-
() -> bigQueryRpc.createJobForQuery(jobPb),
1461+
() ->
1462+
bigQueryOptions
1463+
.getBigQueryRpcV2()
1464+
.createJobForQuerySkipExceptionTranslation(jobPb),
14461465
bigQueryOptions.getRetrySettings(),
14471466
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
14481467
bigQueryOptions.getClock(),
@@ -1482,7 +1501,10 @@ com.google.api.services.bigquery.model.Job createDryRunJob(String sql) {
14821501
try {
14831502
dryRunJob =
14841503
BigQueryRetryHelper.runWithRetries(
1485-
() -> bigQueryRpc.createJobForQuery(jobPb),
1504+
() ->
1505+
bigQueryOptions
1506+
.getBigQueryRpcV2()
1507+
.createJobForQuerySkipExceptionTranslation(jobPb),
14861508
bigQueryOptions.getRetrySettings(),
14871509
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
14881510
bigQueryOptions.getClock(),

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

+18-13
Original file line numberDiff line numberDiff line change
@@ -16,12 +16,11 @@
1616

1717
package com.google.cloud.bigquery;
1818

19-
import static com.google.cloud.RetryHelper.runWithRetries;
20-
2119
import com.google.cloud.BaseWriteChannel;
2220
import com.google.cloud.RestorableState;
23-
import com.google.cloud.RetryHelper;
2421
import com.google.cloud.WriteChannel;
22+
import com.google.cloud.bigquery.BigQueryRetryHelper.BigQueryRetryHelperException;
23+
import java.io.IOException;
2524
import java.util.List;
2625
import java.util.Objects;
2726
import java.util.concurrent.Callable;
@@ -34,6 +33,9 @@
3433
public class TableDataWriteChannel
3534
extends BaseWriteChannel<BigQueryOptions, WriteChannelConfiguration> {
3635

36+
private static final BigQueryRetryConfig EMPTY_RETRY_CONFIG =
37+
BigQueryRetryConfig.newBuilder().build();
38+
3739
private Job job;
3840

3941
TableDataWriteChannel(
@@ -50,20 +52,22 @@ public class TableDataWriteChannel
5052
protected void flushBuffer(final int length, final boolean last) {
5153
try {
5254
com.google.api.services.bigquery.model.Job jobPb =
53-
runWithRetries(
55+
BigQueryRetryHelper.runWithRetries(
5456
new Callable<com.google.api.services.bigquery.model.Job>() {
5557
@Override
56-
public com.google.api.services.bigquery.model.Job call() {
58+
public com.google.api.services.bigquery.model.Job call() throws IOException {
5759
return getOptions()
5860
.getBigQueryRpcV2()
59-
.write(getUploadId(), getBuffer(), 0, getPosition(), length, last);
61+
.writeSkipExceptionTranslation(
62+
getUploadId(), getBuffer(), 0, getPosition(), length, last);
6063
}
6164
},
6265
getOptions().getRetrySettings(),
6366
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
64-
getOptions().getClock());
67+
getOptions().getClock(),
68+
EMPTY_RETRY_CONFIG);
6569
job = jobPb != null ? Job.fromPb(getOptions().getService(), jobPb) : null;
66-
} catch (RetryHelper.RetryHelperException e) {
70+
} catch (BigQueryRetryHelperException e) {
6771
throw BigQueryException.translateAndThrow(e);
6872
}
6973
}
@@ -78,22 +82,23 @@ private static String open(
7882
final JobId jobId,
7983
final WriteChannelConfiguration writeChannelConfiguration) {
8084
try {
81-
return runWithRetries(
85+
return BigQueryRetryHelper.runWithRetries(
8286
new Callable<String>() {
8387
@Override
84-
public String call() {
88+
public String call() throws IOException {
8589
return options
8690
.getBigQueryRpcV2()
87-
.open(
91+
.openSkipExceptionTranslation(
8892
new com.google.api.services.bigquery.model.Job()
8993
.setConfiguration(writeChannelConfiguration.toPb())
9094
.setJobReference(jobId.toPb()));
9195
}
9296
},
9397
options.getRetrySettings(),
9498
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
95-
options.getClock());
96-
} catch (RetryHelper.RetryHelperException e) {
99+
options.getClock(),
100+
EMPTY_RETRY_CONFIG);
101+
} catch (BigQueryRetryHelperException e) {
97102
throw BigQueryException.translateAndThrow(e);
98103
}
99104
}

0 commit comments

Comments
 (0)