Skip to content
This repository was archived by the owner on Mar 23, 2026. It is now read-only.

Commit 44d9795

Browse files
prash-miprash-mi
authored
feat: Implemented BigQueryRetryAlgorithm to retry on the basis of the configured re-triable error messages (#1426)
* Updated BigQueryImpl * Initial Commit * Using BigQueryRetryAlgorithm as the retry algorithm * Created BigQueryRetryAlgorithm as a subclass of RetryAlgorithm * BigQueryErrorMessages property file * Implemented Builder Logic for BigQueryRetryConfig * Using BigQueryRetryConfig for getQueryResults * Updated shouldRetry with the logic to retry based on error messages * Implemented null checks on shouldRetryBasedOnBigQueryRetryConfig * Removed `Status` from shouldRetryBasedOnBigQueryRetryConfig implementation * Removed unused imports * created DEFAULT_RATE_LIMIT_EXCEEDED_RETRY_CONFIG for getQueryResults * Added testGetQueryResultsRetry test for testing getQueryResults Retry * Overriding createNextAttempt method so that it generates an attempt based on the error message * Linted BigQueryRetryHelper * Linted testGetQueryResultsRetry * Linted BigQueryRetryAlgorithm * Linted BigQueryErrorMessages * Linted BigQueryRetryConfig * Fixed Linting * Fixed Linting * Fixed Linting * Created translateAndThrow(BigQueryRetryHelper.BigQueryRetryHelperException ex) method to handle BigQueryRetryHelperException * Handling BigQueryRetryHelper.BigQueryRetryHelperException for getQueryResults * Implementing BigQueryRetryHelper.runWithRetries from TableResult.queryRPC method * Implementing testFastQueryRateLimitIdempotency Method to test Idempotency of the BigQueryRetryHelper.runWithRetries for TableResult.query(...) * Changed DEFAULT_RATE_LIMIT_EXCEEDED_RETRY_CONFIG to DEFAULT_RETRY_CONFIG * Implemented `BigQueryRetryHelper.runWithRetries` on `QueryResponse waitForQueryResults` method, which is used by `TableResult getQueryResults` method * Revert "Implemented `BigQueryRetryHelper.runWithRetries` on `QueryResponse waitForQueryResults` method, which is used by `TableResult getQueryResults` method" This reverts commit 84a3418. * Revert "Changed DEFAULT_RATE_LIMIT_EXCEEDED_RETRY_CONFIG to DEFAULT_RETRY_CONFIG" This reverts commit 22b1706. * Renamed DEFAULT_RATE_LIMIT_EXCEEDED_RETRY_CONFIG to DEFAULT_RETRY_CONFIG * Revert "Renamed DEFAULT_RATE_LIMIT_EXCEEDED_RETRY_CONFIG to DEFAULT_RETRY_CONFIG" This reverts commit 2d21e11. * Renamed DEFAULT_RATE_LIMIT_EXCEEDED_RETRY_CONFIG to DEFAULT_RETRY_CONFIG * Implemented BigQueryRetryHelper.runWithRetries on `QueryResponse waitForQueryResults` method, which is used by `TableResult getQueryResults` method
1 parent 5bee589 commit 44d9795

8 files changed

Lines changed: 433 additions & 13 deletions

File tree

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

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
/*
2+
* Copyright 2021 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.bigquery;
18+
19+
public class BigQueryErrorMessages {
20+
public static final String RATE_LIMIT_EXCEEDED_MSG =
21+
"Exceeded rate limits:"; // Error Message for RateLimitExceeded Error
22+
}

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,14 @@ static BaseServiceException translateAndThrow(RetryHelperException ex) {
122122
throw new BigQueryException(UNKNOWN_CODE, ex.getMessage(), ex.getCause());
123123
}
124124

125+
static BaseServiceException translateAndThrow(
126+
BigQueryRetryHelper.BigQueryRetryHelperException ex) {
127+
if (ex.getCause() instanceof BaseServiceException) {
128+
throw (BaseServiceException) ex.getCause();
129+
}
130+
throw new BigQueryException(UNKNOWN_CODE, ex.getMessage(), ex.getCause());
131+
}
132+
125133
static BaseServiceException translateAndThrow(ExecutionException ex) {
126134
BaseServiceException.translate(ex);
127135
throw new BigQueryException(UNKNOWN_CODE, ex.getMessage(), ex.getCause());

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

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -237,6 +237,10 @@ public Page<FieldValueList> getNextPage() {
237237
}
238238

239239
private final BigQueryRpc bigQueryRpc;
240+
private static final BigQueryRetryConfig DEFAULT_RETRY_CONFIG =
241+
BigQueryRetryConfig.newBuilder()
242+
.retryOnMessage(BigQueryErrorMessages.RATE_LIMIT_EXCEEDED_MSG)
243+
.build(); // retry config with Error Message for RateLimitExceeded Error
240244

241245
BigQueryImpl(BigQueryOptions options) {
242246
super(options);
@@ -1271,7 +1275,7 @@ private TableResult queryRpc(
12711275
com.google.api.services.bigquery.model.QueryResponse results;
12721276
try {
12731277
results =
1274-
runWithRetries(
1278+
BigQueryRetryHelper.runWithRetries(
12751279
new Callable<com.google.api.services.bigquery.model.QueryResponse>() {
12761280
@Override
12771281
public com.google.api.services.bigquery.model.QueryResponse call() {
@@ -1280,8 +1284,9 @@ public com.google.api.services.bigquery.model.QueryResponse call() {
12801284
},
12811285
getOptions().getRetrySettings(),
12821286
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
1283-
getOptions().getClock());
1284-
} catch (RetryHelperException e) {
1287+
getOptions().getClock(),
1288+
DEFAULT_RETRY_CONFIG);
1289+
} catch (BigQueryRetryHelper.BigQueryRetryHelperException e) {
12851290
throw BigQueryException.translateAndThrow(e);
12861291
}
12871292

@@ -1362,7 +1367,7 @@ private static QueryResponse getQueryResults(
13621367
: jobId.getLocation());
13631368
try {
13641369
GetQueryResultsResponse results =
1365-
runWithRetries(
1370+
BigQueryRetryHelper.runWithRetries(
13661371
new Callable<GetQueryResultsResponse>() {
13671372
@Override
13681373
public GetQueryResultsResponse call() {
@@ -1376,8 +1381,10 @@ public GetQueryResultsResponse call() {
13761381
}
13771382
},
13781383
serviceOptions.getRetrySettings(),
1379-
EXCEPTION_HANDLER,
1380-
serviceOptions.getClock());
1384+
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
1385+
serviceOptions.getClock(),
1386+
DEFAULT_RETRY_CONFIG);
1387+
13811388
TableSchema schemaPb = results.getSchema();
13821389

13831390
ImmutableList.Builder<BigQueryError> errors = ImmutableList.builder();
@@ -1393,7 +1400,7 @@ public GetQueryResultsResponse call() {
13931400
.setTotalRows(results.getTotalRows() == null ? 0 : results.getTotalRows().longValue())
13941401
.setErrors(errors.build())
13951402
.build();
1396-
} catch (RetryHelper.RetryHelperException e) {
1403+
} catch (BigQueryRetryHelper.BigQueryRetryHelperException e) {
13971404
throw BigQueryException.translateAndThrow(e);
13981405
}
13991406
}

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

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
/*
2+
* Copyright 2021 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.bigquery;
18+
19+
import static com.google.common.base.Preconditions.checkNotNull;
20+
21+
import com.google.api.gax.retrying.ResultRetryAlgorithm;
22+
import com.google.api.gax.retrying.ResultRetryAlgorithmWithContext;
23+
import com.google.api.gax.retrying.RetryAlgorithm;
24+
import com.google.api.gax.retrying.RetryingContext;
25+
import com.google.api.gax.retrying.TimedAttemptSettings;
26+
import com.google.api.gax.retrying.TimedRetryAlgorithm;
27+
import com.google.api.gax.retrying.TimedRetryAlgorithmWithContext;
28+
import java.util.Iterator;
29+
import java.util.concurrent.CancellationException;
30+
31+
public class BigQueryRetryAlgorithm<ResponseT> extends RetryAlgorithm<ResponseT> {
32+
private final BigQueryRetryConfig bigQueryRetryConfig;
33+
private final ResultRetryAlgorithm<ResponseT> resultAlgorithm;
34+
private final TimedRetryAlgorithm timedAlgorithm;
35+
private final ResultRetryAlgorithmWithContext<ResponseT> resultAlgorithmWithContext;
36+
private final TimedRetryAlgorithmWithContext timedAlgorithmWithContext;
37+
38+
public BigQueryRetryAlgorithm(
39+
ResultRetryAlgorithm<ResponseT> resultAlgorithm,
40+
TimedRetryAlgorithm timedAlgorithm,
41+
BigQueryRetryConfig bigQueryRetryConfig) {
42+
super(resultAlgorithm, timedAlgorithm);
43+
this.bigQueryRetryConfig = checkNotNull(bigQueryRetryConfig);
44+
this.resultAlgorithm = checkNotNull(resultAlgorithm);
45+
this.timedAlgorithm = checkNotNull(timedAlgorithm);
46+
this.resultAlgorithmWithContext = null;
47+
this.timedAlgorithmWithContext = null;
48+
}
49+
50+
@Override
51+
public boolean shouldRetry(
52+
RetryingContext context,
53+
Throwable previousThrowable,
54+
ResponseT previousResponse,
55+
TimedAttemptSettings nextAttemptSettings)
56+
throws CancellationException {
57+
// Implementing shouldRetryBasedOnBigQueryRetryConfig so that we can retry exceptions based on
58+
// the exception messages
59+
return (shouldRetryBasedOnResult(context, previousThrowable, previousResponse)
60+
|| shouldRetryBasedOnBigQueryRetryConfig(previousThrowable, bigQueryRetryConfig))
61+
&& shouldRetryBasedOnTiming(context, nextAttemptSettings);
62+
}
63+
64+
private boolean shouldRetryBasedOnBigQueryRetryConfig(
65+
Throwable previousThrowable, BigQueryRetryConfig bigQueryRetryConfig) {
66+
/*
67+
We are deciding if a given error should be retried on the basis of error message.
68+
Cannot rely on Error/Status code as for example error code 400 (which is not retriable) could be thrown due to rateLimitExceed, which is retriable
69+
*/
70+
String errorDesc;
71+
if (previousThrowable != null && (errorDesc = previousThrowable.getMessage()) != null) {
72+
for (Iterator<String> retriableMessages =
73+
bigQueryRetryConfig.getRetriableErrorMessages().iterator();
74+
retriableMessages.hasNext(); ) {
75+
if (errorDesc.contains(retriableMessages.next())) { // Error message should be retried
76+
return true;
77+
}
78+
}
79+
}
80+
return false;
81+
}
82+
83+
/*Duplicating this method as it can not be inherited from the RetryAlgorithm due to the default access modifier*/
84+
boolean shouldRetryBasedOnResult(
85+
RetryingContext context, Throwable previousThrowable, ResponseT previousResponse) {
86+
if (resultAlgorithmWithContext != null && context != null) {
87+
return resultAlgorithmWithContext.shouldRetry(context, previousThrowable, previousResponse);
88+
}
89+
return getResultAlgorithm().shouldRetry(previousThrowable, previousResponse);
90+
}
91+
92+
/*Duplicating this method as it can not be inherited from the RetryAlgorithm due to the private access modifier*/
93+
private boolean shouldRetryBasedOnTiming(
94+
RetryingContext context, TimedAttemptSettings nextAttemptSettings) {
95+
if (nextAttemptSettings == null) {
96+
return false;
97+
}
98+
if (timedAlgorithmWithContext != null && context != null) {
99+
return timedAlgorithmWithContext.shouldRetry(context, nextAttemptSettings);
100+
}
101+
return getTimedAlgorithm().shouldRetry(nextAttemptSettings);
102+
}
103+
104+
@Override
105+
public TimedAttemptSettings createNextAttempt(
106+
RetryingContext context,
107+
Throwable previousThrowable,
108+
ResponseT previousResponse,
109+
TimedAttemptSettings previousSettings) {
110+
// a small optimization that avoids calling relatively heavy methods
111+
// like timedAlgorithm.createNextAttempt(), when it is not necessary.
112+
113+
if (!((shouldRetryBasedOnResult(context, previousThrowable, previousResponse)
114+
|| shouldRetryBasedOnBigQueryRetryConfig(
115+
previousThrowable,
116+
bigQueryRetryConfig)))) { // Calling shouldRetryBasedOnBigQueryRetryConfig to check if
117+
// the error message could be retried
118+
return null;
119+
}
120+
121+
TimedAttemptSettings newSettings =
122+
createNextAttemptBasedOnResult(
123+
context, previousThrowable, previousResponse, previousSettings);
124+
if (newSettings == null) {
125+
newSettings = createNextAttemptBasedOnTiming(context, previousSettings);
126+
}
127+
return newSettings;
128+
}
129+
130+
/*Duplicating this method as it can not be inherited from the RetryAlgorithm due to the private access modifier*/
131+
private TimedAttemptSettings createNextAttemptBasedOnResult(
132+
RetryingContext context,
133+
Throwable previousThrowable,
134+
ResponseT previousResponse,
135+
TimedAttemptSettings previousSettings) {
136+
if (resultAlgorithmWithContext != null && context != null) {
137+
return resultAlgorithmWithContext.createNextAttempt(
138+
context, previousThrowable, previousResponse, previousSettings);
139+
}
140+
return getResultAlgorithm()
141+
.createNextAttempt(previousThrowable, previousResponse, previousSettings);
142+
}
143+
144+
/*Duplicating this method as it can not be inherited from the RetryAlgorithm due to the private access modifier*/
145+
private TimedAttemptSettings createNextAttemptBasedOnTiming(
146+
RetryingContext context, TimedAttemptSettings previousSettings) {
147+
if (timedAlgorithmWithContext != null && context != null) {
148+
return timedAlgorithmWithContext.createNextAttempt(context, previousSettings);
149+
}
150+
return getTimedAlgorithm().createNextAttempt(previousSettings);
151+
}
152+
}

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

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
/*
2+
* Copyright 2021 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.bigquery;
17+
18+
import static com.google.common.base.Preconditions.checkNotNull;
19+
20+
import com.google.common.collect.ImmutableSet;
21+
22+
public class BigQueryRetryConfig {
23+
private final ImmutableSet<String> retriableErrorMessages;
24+
25+
private BigQueryRetryConfig(Builder builder) {
26+
retriableErrorMessages = builder.retriableErrorMessages.build();
27+
}
28+
29+
public ImmutableSet<String> getRetriableErrorMessages() {
30+
return retriableErrorMessages;
31+
}
32+
33+
// BigQueryRetryConfig builder
34+
public static class Builder {
35+
private final ImmutableSet.Builder<String> retriableErrorMessages = ImmutableSet.builder();
36+
37+
private Builder() {}
38+
39+
public final Builder retryOnMessage(String... errorMessages) {
40+
for (String errorMessage : errorMessages) {
41+
retriableErrorMessages.add(checkNotNull(errorMessage));
42+
}
43+
return this;
44+
}
45+
46+
public BigQueryRetryConfig build() {
47+
return new BigQueryRetryConfig(this);
48+
}
49+
}
50+
51+
public static Builder newBuilder() {
52+
return new Builder();
53+
}
54+
}

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

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
/*
2+
* Copyright 2021 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.bigquery;
17+
18+
import com.google.api.core.ApiClock;
19+
import com.google.api.gax.retrying.*;
20+
import com.google.cloud.RetryHelper;
21+
import java.util.concurrent.Callable;
22+
import java.util.concurrent.ExecutionException;
23+
24+
public class BigQueryRetryHelper extends RetryHelper {
25+
26+
public static <V> V runWithRetries(
27+
Callable<V> callable,
28+
RetrySettings retrySettings,
29+
ResultRetryAlgorithm<?> resultRetryAlgorithm,
30+
ApiClock clock,
31+
BigQueryRetryConfig bigQueryRetryConfig)
32+
throws RetryHelperException {
33+
try {
34+
// Suppressing should be ok as a workaraund. Current and only ResultRetryAlgorithm
35+
// implementation does not use response at all, so ignoring its type is ok.
36+
@SuppressWarnings("unchecked")
37+
ResultRetryAlgorithm<V> algorithm = (ResultRetryAlgorithm<V>) resultRetryAlgorithm;
38+
return run(
39+
callable,
40+
new ExponentialRetryAlgorithm(retrySettings, clock),
41+
algorithm,
42+
bigQueryRetryConfig);
43+
} catch (Exception e) {
44+
throw new BigQueryRetryHelperException(e.getCause());
45+
}
46+
}
47+
48+
private static <V> V run(
49+
Callable<V> callable,
50+
TimedRetryAlgorithm timedAlgorithm,
51+
ResultRetryAlgorithm<V> resultAlgorithm,
52+
BigQueryRetryConfig bigQueryRetryConfig)
53+
throws ExecutionException, InterruptedException {
54+
RetryAlgorithm<V> retryAlgorithm =
55+
new BigQueryRetryAlgorithm<>(
56+
resultAlgorithm,
57+
timedAlgorithm,
58+
bigQueryRetryConfig); // using BigQueryRetryAlgorithm in place of
59+
// com.google.api.gax.retrying.RetryAlgorithm, as
60+
// BigQueryRetryAlgorithm retries considering bigQueryRetryConfig
61+
RetryingExecutor<V> executor = new DirectRetryingExecutor<>(retryAlgorithm);
62+
63+
RetryingFuture<V> retryingFuture = executor.createFuture(callable);
64+
executor.submit(retryingFuture);
65+
return retryingFuture.get();
66+
}
67+
68+
public static class BigQueryRetryHelperException extends RuntimeException {
69+
70+
private static final long serialVersionUID = -8519852520090965314L;
71+
72+
BigQueryRetryHelperException(Throwable cause) {
73+
super(cause);
74+
}
75+
}
76+
}

0 commit comments

Comments
 (0)