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 1 commit
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
109 changes: 71 additions & 38 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 @@ -317,46 +317,79 @@ 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);

Job job = new Job().setConfiguration(jc);

Job insertedJob = bqClient.jobs().insert(projectId, job).execute();
try {
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);

Job job = new Job().setConfiguration(jc);

Job insertedJob = bqClient.jobs().insert(projectId, job).execute();

GetQueryResultsResponse qResponse;
do {
qResponse =
getQueryResultsWithRetries(
projectId, insertedJob.getJobReference().getJobId(), location);
} while (!Boolean.TRUE.equals(qResponse.getJobComplete()));

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);
}
}

GetQueryResultsResponse qResponse;
/**
* Calls jobs.getQueryResults and retries 5xx errors, such as a 503 backendError, and 429 rate
* limiting. Throws other 4xx errors, such as a failed query, right away.
*/
private GetQueryResultsResponse getQueryResultsWithRetries(
String projectId, String jobId, @Nullable String location)
throws IOException, InterruptedException {
Sleeper sleeper = Sleeper.DEFAULT;
BackOff backoff = BackOffAdapter.toGcpBackOff(BACKOFF_FACTORY.backoff());
IOException lastException = null;
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;
}
deleteDataset(projectId, temporaryDatasetId);
return !typed
? rows
: rows.stream()
.map(r -> getTypedTableRow(schema.getFields(), r))
.collect(Collectors.toList());
if (lastException != null) {
LOG.warn("Retrying getQueryResults for job ({}) after exception", jobId, lastException);
}
try {
return bqClient.jobs().getQueryResults(projectId, jobId).setLocation(location).execute();
} catch (GoogleJsonResponseException e) {
if (e.getStatusCode() < 500 && e.getStatusCode() != 429) {
throw e;
}
lastException = e;
} catch (IOException e) {
// ignore and retry
lastException = e;
}
} while (BackOffUtils.next(sleeper, backoff));

throw new RuntimeException(
String.format(
"Unable to get BigQuery query results after retrying %d times for job (%s)",
MAX_QUERY_RETRIES, jobId),
lastException);
}

@Nonnull
Expand Down
79 changes: 79 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,31 @@
*/
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.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 +65,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 +116,62 @@ public void testQueryWithRetriesWhenQueryResponseNull() throws Exception {
.query(eq(projectId), any(QueryRequest.class));
}
}

private static GoogleJsonResponseException jsonError(int statusCode) {
return new GoogleJsonResponseException(
new HttpResponseException.Builder(statusCode, "error", new HttpHeaders()), null);
}

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);
}

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

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

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

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 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