Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 97a61dc

Browse files
feat: create Job retry for rate limit exceeded with status code 200 (#1744)
* retrying job create with new job id each time * retrying for rate limit exceeded on create Job regardless of the status code 200 * removing flattened pom * refactoring, removing unnecesary variables, adding comments * linting code * 🦉 Updates from OwlBot See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * 🦉 Updates from OwlBot See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * refactoring, removing comments Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
1 parent faf5897 commit 97a61dc

File tree

3 files changed

+67
-19
lines changed

3 files changed

+67
-19
lines changed

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

Lines changed: 23 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -347,37 +347,46 @@ public JobId get() {
347347

348348
@InternalApi("visible for testing")
349349
Job create(JobInfo jobInfo, Supplier<JobId> idProvider, JobOption... options) {
350-
boolean idRandom = false;
351-
if (jobInfo.getJobId() == null) {
352-
jobInfo = jobInfo.toBuilder().setJobId(idProvider.get()).build();
353-
idRandom = true;
354-
}
355-
final com.google.api.services.bigquery.model.Job jobPb =
356-
jobInfo.setProjectId(getOptions().getProjectId()).toPb();
357-
final Map<BigQueryRpc.Option, ?> optionsMap = optionMap(options);
350+
final boolean idRandom = (jobInfo.getJobId() == null);
358351

352+
final Map<BigQueryRpc.Option, ?> optionsMap = optionMap(options);
359353
BigQueryException createException;
360354
// NOTE(pongad): This double-try structure is admittedly odd.
361355
// translateAndThrow itself throws, and pretends to return an exception only
362356
// so users can pretend to throw.
363357
// This makes it difficult to translate without throwing.
364358
// Fixing this entails some work on BaseServiceException.translate.
365359
// Since that affects a bunch of APIs, we should fix this as a separate change.
360+
final JobId[] finalJobId = new JobId[1];
366361
try {
367362
try {
368363
return Job.fromPb(
369364
this,
370-
runWithRetries(
365+
BigQueryRetryHelper.runWithRetries(
371366
new Callable<com.google.api.services.bigquery.model.Job>() {
372367
@Override
373368
public com.google.api.services.bigquery.model.Job call() {
374-
return bigQueryRpc.create(jobPb, optionsMap);
369+
if (idRandom) {
370+
// re-generate a new random job with the same jobInfo when jobId is not
371+
// provided by the user
372+
JobInfo recreatedJobInfo =
373+
jobInfo.toBuilder().setJobId(idProvider.get()).build();
374+
com.google.api.services.bigquery.model.Job newJobPb =
375+
recreatedJobInfo.setProjectId(getOptions().getProjectId()).toPb();
376+
finalJobId[0] = recreatedJobInfo.getJobId();
377+
return bigQueryRpc.create(newJobPb, optionsMap);
378+
} else {
379+
com.google.api.services.bigquery.model.Job jobPb =
380+
jobInfo.setProjectId(getOptions().getProjectId()).toPb();
381+
return bigQueryRpc.create(jobPb, optionsMap);
382+
}
375383
}
376384
},
377385
getOptions().getRetrySettings(),
378-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
379-
getOptions().getClock()));
380-
} catch (RetryHelper.RetryHelperException e) {
386+
EXCEPTION_HANDLER,
387+
getOptions().getClock(),
388+
DEFAULT_RETRY_CONFIG));
389+
} catch (BigQueryRetryHelper.BigQueryRetryHelperException e) {
381390
throw BigQueryException.translateAndThrow(e);
382391
}
383392
} catch (BigQueryException e) {
@@ -396,7 +405,7 @@ public com.google.api.services.bigquery.model.Job call() {
396405
// fetch a job created by someone else.
397406
Job job;
398407
try {
399-
job = getJob(jobInfo.getJobId());
408+
job = getJob(finalJobId[0]);
400409
} catch (BigQueryException e) {
401410
throw createException;
402411
}

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

Lines changed: 20 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,8 @@ public boolean shouldRetry(
7272
// the exception messages
7373
boolean shouldRetry =
7474
(shouldRetryBasedOnResult(context, previousThrowable, previousResponse)
75-
|| shouldRetryBasedOnBigQueryRetryConfig(previousThrowable, bigQueryRetryConfig))
75+
|| shouldRetryBasedOnBigQueryRetryConfig(
76+
previousThrowable, bigQueryRetryConfig, previousResponse))
7677
&& shouldRetryBasedOnTiming(context, nextAttemptSettings);
7778

7879
if (LOG.isLoggable(Level.FINEST)) {
@@ -92,13 +93,26 @@ public boolean shouldRetry(
9293
}
9394

9495
private boolean shouldRetryBasedOnBigQueryRetryConfig(
95-
Throwable previousThrowable, BigQueryRetryConfig bigQueryRetryConfig) {
96+
Throwable previousThrowable,
97+
BigQueryRetryConfig bigQueryRetryConfig,
98+
ResponseT previousResponse) {
9699
/*
97100
We are deciding if a given error should be retried on the basis of error message.
98101
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
99102
*/
100-
String errorDesc;
101-
if (previousThrowable != null && (errorDesc = previousThrowable.getMessage()) != null) {
103+
String errorDesc = null;
104+
if (previousThrowable != null) {
105+
errorDesc = previousThrowable.getMessage();
106+
} else if (previousResponse != null) {
107+
/*
108+
In some cases error messages may come without an exception
109+
e.g. status code 200 with a rate limit exceeded for job create
110+
in these cases there is now previousThrowable so we need to check previousResponse
111+
*/
112+
errorDesc = previousResponse.toString();
113+
}
114+
115+
if (errorDesc != null) {
102116
errorDesc = errorDesc.toLowerCase(); // for case insensitive comparison
103117
for (Iterator<String> retriableMessages =
104118
bigQueryRetryConfig.getRetriableErrorMessages().iterator();
@@ -161,7 +175,8 @@ public TimedAttemptSettings createNextAttempt(
161175
if (!((shouldRetryBasedOnResult(context, previousThrowable, previousResponse)
162176
|| shouldRetryBasedOnBigQueryRetryConfig(
163177
previousThrowable,
164-
bigQueryRetryConfig)))) { // Calling shouldRetryBasedOnBigQueryRetryConfig to check if
178+
bigQueryRetryConfig,
179+
previousResponse)))) { // Calling shouldRetryBasedOnBigQueryRetryConfig to check if
165180
// the error message could be retried
166181
return null;
167182
}

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

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1529,6 +1529,30 @@ public void testCreateJobSuccess() {
15291529
verify(bigqueryRpcMock).create(jobCapture.capture(), eq(EMPTY_RPC_OPTIONS));
15301530
}
15311531

1532+
@Test
1533+
public void testCreateJobFailureShouldRetry() {
1534+
when(bigqueryRpcMock.create(jobCapture.capture(), eq(EMPTY_RPC_OPTIONS)))
1535+
.thenThrow(new BigQueryException(500, "InternalError"))
1536+
.thenThrow(new BigQueryException(502, "Bad Gateway"))
1537+
.thenThrow(new BigQueryException(503, "Service Unavailable"))
1538+
.thenThrow(
1539+
new BigQueryException(
1540+
400, RATE_LIMIT_ERROR_MSG)) // retrial on based on RATE_LIMIT_EXCEEDED_MSG
1541+
.thenThrow(new BigQueryException(200, RATE_LIMIT_ERROR_MSG))
1542+
.thenReturn(newJobPb());
1543+
1544+
bigquery = options.getService();
1545+
bigquery =
1546+
options
1547+
.toBuilder()
1548+
.setRetrySettings(ServiceOptions.getDefaultRetrySettings())
1549+
.build()
1550+
.getService();
1551+
1552+
((BigQueryImpl) bigquery).create(JobInfo.of(QUERY_JOB_CONFIGURATION_FOR_DMLQUERY));
1553+
verify(bigqueryRpcMock, times(6)).create(jobCapture.capture(), eq(EMPTY_RPC_OPTIONS));
1554+
}
1555+
15321556
@Test
15331557
public void testCreateJobWithSelectedFields() {
15341558
when(bigqueryRpcMock.create(

0 commit comments

Comments
 (0)