From ca6a705c47191bea623c5af3e52bf5f09ac55f3b Mon Sep 17 00:00:00 2001 From: Maksym Tymoshyk Date: Sun, 25 Oct 2026 19:45:54 +0200 Subject: [PATCH 1/2] Retry transient errors in BigqueryClient.queryUnflattened queryUnflattened polled jobs.getQueryResults with no retry, so one 503 backendError failed an integration test after its pipeline had already succeeded. Poll through getQueryResultsWithRetries, which uses the same backoff as the other BigqueryClient methods. It retries 5xx and 429 and throws other 4xx errors right away. Also delete the temporary dataset in a finally block. Before, the method skipped the delete when the query returned no rows or threw. --- .../sdk/io/gcp/testing/BigqueryClient.java | 109 ++++++++++++------ .../io/gcp/testing/BigqueryClientTest.java | 79 +++++++++++++ 2 files changed, 150 insertions(+), 38 deletions(-) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/BigqueryClient.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/BigqueryClient.java index a83c49404912..238183675bee 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/BigqueryClient.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/BigqueryClient.java @@ -317,46 +317,79 @@ public List 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 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 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 diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/testing/BigqueryClientTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/testing/BigqueryClientTest.java index b1260e6fac7e..1d4bed8963cc 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/testing/BigqueryClientTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/testing/BigqueryClientTest.java @@ -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; @@ -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 mockStatic; @Before @@ -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 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()); + } } From 8e4e1b7efaef0a6948d54aa5acebbd71e76f4c47 Mon Sep 17 00:00:00 2001 From: Maksym Tymoshyk Date: Sun, 25 Oct 2026 20:36:48 +0200 Subject: [PATCH 2/2] Retry the whole query job in BigqueryClient.queryUnflattened Polling getQueryResults again can't fix a job that itself failed with a transient error such as a 503 backendError, because every poll returns the same error. Each retry now starts a new query job, which also covers a 503 from jobs.insert. The temp table uses WRITE_TRUNCATE so a retried job overwrites anything a failed job left behind. The retry check also treats 403 rateLimitExceeded as transient, applies the status-code rules to non-JSON HTTP errors, and keeps the original jobComplete check so a missing value fails instead of looping. --- .../sdk/io/gcp/testing/BigqueryClient.java | 59 +++++++++------- .../io/gcp/testing/BigqueryClientTest.java | 67 +++++++++++++++++-- 2 files changed, 95 insertions(+), 31 deletions(-) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/BigqueryClient.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/BigqueryClient.java index 238183675bee..46a2829e6f78 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/BigqueryClient.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/BigqueryClient.java @@ -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; @@ -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; @@ -328,20 +330,10 @@ public List queryUnflattened( .setFlattenResults(false) .setAllowLargeResults(true) .setDestinationTable(tempTableReference) + .setWriteDisposition("WRITE_TRUNCATE") .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())); + GetQueryResultsResponse qResponse = runQueryJobWithRetries(projectId, jcQuery, location); final TableSchema schema = qResponse.getSchema(); final List rows = qResponse.getRows(); @@ -359,39 +351,56 @@ public List queryUnflattened( } /** - * 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. + * 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 getQueryResultsWithRetries( - String projectId, String jobId, @Nullable String location) + 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 getQueryResults for job ({}) after exception", jobId, lastException); + LOG.warn("Retrying query ({}) after exception", jcQuery.getQuery(), lastException); } try { - return bqClient.jobs().getQueryResults(projectId, jobId).setLocation(location).execute(); - } catch (GoogleJsonResponseException e) { - if (e.getStatusCode() < 500 && e.getStatusCode() != 429) { + 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; - } 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), + "Unable to get BigQuery query results after retrying %d times using query (%s)", + MAX_QUERY_RETRIES, jcQuery.getQuery()), lastException); } + /** + * 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; + } + int statusCode = ((HttpResponseException) e).getStatusCode(); + return statusCode >= 500 || statusCode == 429 || ApiErrorExtractor.INSTANCE.rateLimited(e); + } + @Nonnull public QueryResponse queryWithRetries(String query, String projectId, boolean typed) throws IOException, InterruptedException { diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/testing/BigqueryClientTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/testing/BigqueryClientTest.java index 1d4bed8963cc..55a4f3d4409a 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/testing/BigqueryClientTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/testing/BigqueryClientTest.java @@ -28,6 +28,8 @@ 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; @@ -118,8 +120,18 @@ public void testQueryWithRetriesWhenQueryResponseNull() throws Exception { } 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, "error", new HttpHeaders()), null); + new HttpResponseException.Builder(statusCode, reason, new HttpHeaders()), details); } private void stubQueryUnflattened() throws IOException { @@ -138,23 +150,52 @@ private void stubQueryUnflattened() throws IOException { when(mockGetQueryResults.setLocation(any())).thenReturn(mockGetQueryResults); } + private static GetQueryResultsResponse completeResponse(List rows) { + return new GetQueryResultsResponse().setJobComplete(true).setRows(rows); + } + @Test - public void testQueryUnflattenedRetriesServerError() throws Exception { + public void testQueryUnflattenedStartsNewJobAfterServerError() throws Exception { stubQueryUnflattened(); List rows = Collections.singletonList(new TableRow().set("f0_", "1")); when(mockGetQueryResults.execute()) - .thenThrow(jsonError(503)) - .thenReturn(new GetQueryResultsResponse().setJobComplete(true).setRows(rows)); + .thenThrow(jsonError(503, "backendError")) + .thenReturn(completeResponse(rows)); assertEquals(rows, bqClient.queryUnflattened(query, projectId, false, true)); - verify(mockGetQueryResults, times(2)).execute(); + verify(mockJobs, times(2)).insert(eq(projectId), any(Job.class)); verify(mockDatasets).delete(eq(projectId), anyString()); } + @Test + public void testQueryUnflattenedRetriesJobInsertError() throws Exception { + stubQueryUnflattened(); + List 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 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)); + when(mockGetQueryResults.execute()).thenThrow(jsonError(400, "invalidQuery")); thrown.expect(GoogleJsonResponseException.class); try { @@ -165,6 +206,20 @@ public void testQueryUnflattenedDoesNotRetryClientError() throws Exception { } } + @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();