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..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; @@ -317,46 +319,86 @@ 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); + 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 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 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 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..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 @@ -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; @@ -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 mockStatic; @Before @@ -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 rows) { + return new GetQueryResultsResponse().setJobComplete(true).setRows(rows); + } + + @Test + public void testQueryUnflattenedStartsNewJobAfterServerError() throws Exception { + stubQueryUnflattened(); + List 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 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, "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()); + } }