Skip to content

Retry transient errors in BigqueryClient.queryUnflattened #40793

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mxtymoshyk wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Retry transient errors in BigqueryClient.queryUnflattened #40793

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ca6a705
Retry transient errors in BigqueryClient.queryUnflattened
mxtymoshyk Oct 25, 2026
8e4e1b7
Retry the whole query job in BigqueryClient.queryUnflattened
mxtymoshyk Oct 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
114 changes: 78 additions & 36 deletions ...oogle-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/BigqueryClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;

import com.google.api.client.googleapis.json.GoogleJsonResponseException;
import com.google.api.client.http.HttpResponseException;
import com.google.api.client.http.HttpTransport;
import com.google.api.client.json.JsonFactory;
import com.google.api.client.util.BackOff;
Expand Down Expand Up @@ -49,6 +50,7 @@
import com.google.auth.Credentials;
import com.google.auth.http.HttpCredentialsAdapter;
import com.google.auth.oauth2.GoogleCredentials;
import com.google.cloud.hadoop.util.ApiErrorExtractor;
import java.io.IOException;
import java.util.Collection;
import java.util.Collections;
Expand Down Expand Up @@ -317,46 +319,86 @@ public List<TableRow> queryUnflattened(
.setTableId(temporaryTableId);

createNewDataset(projectId, temporaryDatasetId, null, location);
createNewTable(
projectId,
temporaryDatasetId,
new Table().setTableReference(tempTableReference).setLocation(location));

JobConfigurationQuery jcQuery =
new JobConfigurationQuery()
.setFlattenResults(false)
.setAllowLargeResults(true)
.setDestinationTable(tempTableReference)
.setUseLegacySql(!useStandardSql)
.setQuery(query);
JobConfiguration jc = new JobConfiguration().setQuery(jcQuery);
try {
createNewTable(
projectId,
temporaryDatasetId,
new Table().setTableReference(tempTableReference).setLocation(location));

JobConfigurationQuery jcQuery =
new JobConfigurationQuery()
.setFlattenResults(false)
.setAllowLargeResults(true)
.setDestinationTable(tempTableReference)
.setWriteDisposition("WRITE_TRUNCATE")
.setUseLegacySql(!useStandardSql)
.setQuery(query);
GetQueryResultsResponse qResponse = runQueryJobWithRetries(projectId, jcQuery, location);

final TableSchema schema = qResponse.getSchema();
final List<TableRow> rows = qResponse.getRows();
if (rows == null) {
return Collections.EMPTY_LIST;
}
return !typed
? rows
: rows.stream()
.map(r -> getTypedTableRow(schema.getFields(), r))
.collect(Collectors.toList());
} finally {
deleteDataset(projectId, temporaryDatasetId);
}
}

Job job = new Job().setConfiguration(jc);
/**
* Starts a query job and waits for its results. Polling again can't revive a job that failed with
* a transient error such as a 503 backendError, so each retry starts a new job.
*/
private GetQueryResultsResponse runQueryJobWithRetries(
String projectId, JobConfigurationQuery jcQuery, @Nullable String location)
throws IOException, InterruptedException {
Sleeper sleeper = Sleeper.DEFAULT;
BackOff backoff = BackOffAdapter.toGcpBackOff(BACKOFF_FACTORY.backoff());
IOException lastException = null;
do {
if (lastException != null) {
LOG.warn("Retrying query ({}) after exception", jcQuery.getQuery(), lastException);
}
try {
Job job = new Job().setConfiguration(new JobConfiguration().setQuery(jcQuery));
String jobId =
bqClient.jobs().insert(projectId, job).execute().getJobReference().getJobId();
GetQueryResultsResponse qResponse;
do {
qResponse =
bqClient.jobs().getQueryResults(projectId, jobId).setLocation(location).execute();
} while (!qResponse.getJobComplete());
return qResponse;
} catch (IOException e) {
if (!isTransient(e)) {
throw e;
}
lastException = e;
}
} while (BackOffUtils.next(sleeper, backoff));

Job insertedJob = bqClient.jobs().insert(projectId, job).execute();
throw new RuntimeException(
String.format(
"Unable to get BigQuery query results after retrying %d times using query (%s)",
MAX_QUERY_RETRIES, jcQuery.getQuery()),
lastException);
}

GetQueryResultsResponse qResponse;
do {
qResponse =
bqClient
.jobs()
.getQueryResults(projectId, insertedJob.getJobReference().getJobId())
.setLocation(location)
.execute();

} while (!qResponse.getJobComplete());

final TableSchema schema = qResponse.getSchema();
final List<TableRow> rows = qResponse.getRows();
if (rows == null) {
return Collections.EMPTY_LIST;
/**
* Returns true for errors that may succeed on retry: 5xx responses, rate limiting (429, or 403
* rateLimitExceeded) and failures without an HTTP response, such as timeouts.
*/
private boolean isTransient(IOException e) {
if (!(e instanceof HttpResponseException)) {
return true;
}
deleteDataset(projectId, temporaryDatasetId);
return !typed
? rows
: rows.stream()
.map(r -> getTypedTableRow(schema.getFields(), r))
.collect(Collectors.toList());
int statusCode = ((HttpResponseException) e).getStatusCode();
return statusCode >= 500 || statusCode == 429 || ApiErrorExtractor.INSTANCE.rateLimited(e);
}

@Nonnull
Expand Down
134 changes: 134 additions & 0 deletions ...e-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/testing/BigqueryClientTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,17 +17,33 @@
*/
package org.apache.beam.sdk.io.gcp.testing;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.atLeast;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import com.google.api.client.googleapis.json.GoogleJsonError;
import com.google.api.client.googleapis.json.GoogleJsonError.ErrorInfo;
import com.google.api.client.googleapis.json.GoogleJsonResponseException;
import com.google.api.client.http.HttpHeaders;
import com.google.api.client.http.HttpResponseException;
import com.google.api.services.bigquery.Bigquery;
import com.google.api.services.bigquery.model.Dataset;
import com.google.api.services.bigquery.model.GetQueryResultsResponse;
import com.google.api.services.bigquery.model.Job;
import com.google.api.services.bigquery.model.JobReference;
import com.google.api.services.bigquery.model.QueryRequest;
import com.google.api.services.bigquery.model.Table;
import com.google.api.services.bigquery.model.TableRow;
import java.io.IOException;
import java.util.Collections;
import java.util.List;
import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
Expand All @@ -51,6 +67,13 @@ public class BigqueryClientTest {
@Mock private Bigquery mockBigqueryClient;
@Mock private Bigquery.Jobs mockJobs;
@Mock private Bigquery.Jobs.Query mockQuery;
@Mock private Bigquery.Jobs.Insert mockJobInsert;
@Mock private Bigquery.Jobs.GetQueryResults mockGetQueryResults;
@Mock private Bigquery.Datasets mockDatasets;
@Mock private Bigquery.Datasets.Insert mockDatasetInsert;
@Mock private Bigquery.Datasets.Delete mockDatasetDelete;
@Mock private Bigquery.Tables mockTables;
@Mock private Bigquery.Tables.Insert mockTableInsert;
private MockedStatic<BigqueryClient> mockStatic;

@Before
Expand Down Expand Up @@ -95,4 +118,115 @@ public void testQueryWithRetriesWhenQueryResponseNull() throws Exception {
.query(eq(projectId), any(QueryRequest.class));
}
}

private static GoogleJsonResponseException jsonError(int statusCode) {
return jsonError(statusCode, "error");
}

private static GoogleJsonResponseException jsonError(int statusCode, String reason) {
ErrorInfo errorInfo = new ErrorInfo();
errorInfo.setDomain("global");
errorInfo.setReason(reason);
GoogleJsonError details = new GoogleJsonError();
details.setCode(statusCode);
details.setErrors(Collections.singletonList(errorInfo));
return new GoogleJsonResponseException(
new HttpResponseException.Builder(statusCode, reason, new HttpHeaders()), details);
}

private void stubQueryUnflattened() throws IOException {
when(mockBigqueryClient.datasets()).thenReturn(mockDatasets);
when(mockDatasets.insert(anyString(), any(Dataset.class))).thenReturn(mockDatasetInsert);
when(mockDatasetInsert.execute()).thenReturn(new Dataset());
when(mockDatasets.delete(anyString(), anyString())).thenReturn(mockDatasetDelete);
when(mockDatasetDelete.setDeleteContents(true)).thenReturn(mockDatasetDelete);
when(mockBigqueryClient.tables()).thenReturn(mockTables);
when(mockTables.insert(anyString(), anyString(), any(Table.class))).thenReturn(mockTableInsert);
when(mockTableInsert.execute()).thenReturn(new Table());
when(mockJobs.insert(anyString(), any(Job.class))).thenReturn(mockJobInsert);
when(mockJobInsert.execute())
.thenReturn(new Job().setJobReference(new JobReference().setJobId("job-id")));
when(mockJobs.getQueryResults(anyString(), anyString())).thenReturn(mockGetQueryResults);
when(mockGetQueryResults.setLocation(any())).thenReturn(mockGetQueryResults);
}

private static GetQueryResultsResponse completeResponse(List<TableRow> rows) {
return new GetQueryResultsResponse().setJobComplete(true).setRows(rows);
}

@Test
public void testQueryUnflattenedStartsNewJobAfterServerError() throws Exception {
stubQueryUnflattened();
List<TableRow> rows = Collections.singletonList(new TableRow().set("f0_", "1"));
when(mockGetQueryResults.execute())
.thenThrow(jsonError(503, "backendError"))
.thenReturn(completeResponse(rows));

assertEquals(rows, bqClient.queryUnflattened(query, projectId, false, true));
verify(mockJobs, times(2)).insert(eq(projectId), any(Job.class));
verify(mockDatasets).delete(eq(projectId), anyString());
}

@Test
public void testQueryUnflattenedRetriesJobInsertError() throws Exception {
stubQueryUnflattened();
List<TableRow> rows = Collections.singletonList(new TableRow().set("f0_", "1"));
when(mockJobInsert.execute())
.thenThrow(jsonError(503, "backendError"))
.thenReturn(new Job().setJobReference(new JobReference().setJobId("job-id")));
when(mockGetQueryResults.execute()).thenReturn(completeResponse(rows));

assertEquals(rows, bqClient.queryUnflattened(query, projectId, false, true));
verify(mockJobInsert, times(2)).execute();
}

@Test
public void testQueryUnflattenedRetriesRateLimitExceeded() throws Exception {
stubQueryUnflattened();
List<TableRow> rows = Collections.singletonList(new TableRow().set("f0_", "1"));
when(mockGetQueryResults.execute())
.thenThrow(jsonError(403, "rateLimitExceeded"))
.thenReturn(completeResponse(rows));

assertEquals(rows, bqClient.queryUnflattened(query, projectId, false, true));
verify(mockGetQueryResults, times(2)).execute();
}

@Test
public void testQueryUnflattenedDoesNotRetryClientError() throws Exception {
stubQueryUnflattened();
when(mockGetQueryResults.execute()).thenThrow(jsonError(400, "invalidQuery"));

thrown.expect(GoogleJsonResponseException.class);
try {
bqClient.queryUnflattened(query, projectId, false, true);
} finally {
verify(mockGetQueryResults, times(1)).execute();
verify(mockDatasets).delete(eq(projectId), anyString());
}
}

@Test
public void testQueryUnflattenedDoesNotRetryNonJsonClientError() throws Exception {
stubQueryUnflattened();
when(mockGetQueryResults.execute())
.thenThrow(new HttpResponseException.Builder(404, "Not Found", new HttpHeaders()).build());

thrown.expect(HttpResponseException.class);
try {
bqClient.queryUnflattened(query, projectId, false, true);
} finally {
verify(mockGetQueryResults, times(1)).execute();
}
}

@Test
public void testQueryUnflattenedDeletesDatasetWhenNoRows() throws Exception {
stubQueryUnflattened();
when(mockGetQueryResults.execute())
.thenReturn(new GetQueryResultsResponse().setJobComplete(true));

assertTrue(bqClient.queryUnflattened(query, projectId, false, true).isEmpty());
verify(mockDatasets).delete(eq(projectId), anyString());
}
}
Loading