diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java index 586becb710b2..9248818c9972 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java @@ -248,13 +248,13 @@ public Page getNextPage() { } } - private class QueryPageFetcher extends Thread implements NextPageFetcher { + private class QueryPageFetcher implements NextPageFetcher { private static final long serialVersionUID = -8501991114794410114L; private final Map requestOptions; private final BigQueryOptions serviceOptions; - private Job job; - private final TableId table; + private final JobId jobId; + private TableId table; private final Schema schema; QueryPageFetcher( @@ -266,20 +266,26 @@ private class QueryPageFetcher extends Thread implements NextPageFetcher getNextPage() { - while (!JobStatus.State.DONE.equals(job.getStatus().getState())) { - try { - sleep(5000); - } catch (InterruptedException ex) { - throw new RuntimeException(ex.getMessage()); - } - job = job.reload(); + // Lazily resolve the destination table on the first getNextPage() call rather than in the + // constructor. All first-page rows and metadata on TableResult are already populated from + // QueryResponse (jobs.query), and the destination table is only ever used here to call + // listTableData() for subsequent pages. Deferring getJob(jobId) is therefore behaviorally + // identical when iterating multiple pages while avoiding an unused blocking jobs.get RPC + // when only the first page is consumed. + // + // Historically, getNextPage() also polled job.reload() with a 5s sleep until State.DONE + // (from before queryRpc() checked results.getJobComplete()). Because QueryPageFetcher is + // only created when results.getJobComplete() is true, the job is already complete and does + // not need status polling. + if (table == null) { + Job job = getJob(jobId); + table = ((QueryJobConfiguration) job.getConfiguration()).getDestinationTable(); } return listTableData(table, schema, serviceOptions, requestOptions).x(); } @@ -2456,7 +2462,7 @@ public com.google.api.services.bigquery.model.QueryResponse call() long numRows; Schema schema; - if (results.getJobComplete() && results.getSchema() != null) { + if (Boolean.TRUE.equals(results.getJobComplete()) && results.getSchema() != null) { schema = Schema.fromPb(results.getSchema()); if (results.getNumDmlAffectedRows() == null && results.getTotalRows() == null) { numRows = 0L; @@ -2470,8 +2476,7 @@ public com.google.api.services.bigquery.model.QueryResponse call() // return the schema, fallback. Some operations don't return the schema and can be optimized // here, but this is left as future work. JobId jobId = JobId.fromPb(results.getJobReference()); - Job job = getJob(jobId, options); - return job; + return getJob(jobId, options); } if (results.getPageToken() != null) { diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java index fe16314c4360..6aab5de04b81 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java @@ -2687,6 +2687,8 @@ PROJECT, DATASET, TABLE, optionMap(BigQuery.TableDataListOption.pageToken(CURSOR TableResult result = bigquery.query(QUERY_JOB_CONFIGURATION_FOR_QUERY); assertTrue(result.hasNextPage()); assertNotNull(result.getNextPageToken()); + verify(bigqueryRpcMock, never()) + .getJobSkipExceptionTranslation(PROJECT, JOB, null, EMPTY_RPC_OPTIONS); assertNotNull(result.getNextPage()); QueryRequest requestPb = requestPbCapture.getValue(); @@ -2704,6 +2706,85 @@ PROJECT, DATASET, TABLE, optionMap(BigQuery.TableDataListOption.pageToken(CURSOR .queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture()); } + @Test + void testFastQueryMultiplePagesFirstPageOnlySkipsGetJob() + throws InterruptedException, IOException { + JobId queryJob = JobId.of(PROJECT, JOB); + com.google.api.services.bigquery.model.QueryResponse queryResponsePb = + new com.google.api.services.bigquery.model.QueryResponse() + .setCacheHit(false) + .setJobReference(queryJob.toPb()) + .setJobComplete(true) + .setKind("bigquery#queryResponse") + .setPageToken(CURSOR) + .setRows(ImmutableList.of(TABLE_ROW)) + .setSchema(TABLE_SCHEMA.toPb()) + .setTotalBytesProcessed(42L) + .setTotalRows(BigInteger.valueOf(2L)); + + when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture())) + .thenReturn(queryResponsePb); + + bigquery = options.getService(); + TableResult result = bigquery.query(QUERY_JOB_CONFIGURATION_FOR_QUERY); + assertTrue(result.hasNextPage()); + assertEquals(CURSOR, result.getNextPageToken()); + assertThat(result.getSchema()).isEqualTo(TABLE_SCHEMA); + assertThat(result.getTotalRows()).isEqualTo(2L); + for (FieldValueList row : result.getValues()) { + assertThat(row.get(0).getBooleanValue()).isFalse(); + assertThat(row.get(1).getLongValue()).isEqualTo(1); + } + + verify(bigqueryRpcMock) + .queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture()); + verify(bigqueryRpcMock, never()).getJobSkipExceptionTranslation(any(), any(), any(), any()); + verify(bigqueryRpcMock, never()) + .listTableDataSkipExceptionTranslation(any(), any(), any(), any()); + } + + @Test + void testFastQueryMultiplePagesCachesDestinationTable() throws InterruptedException, IOException { + JobId queryJob = JobId.of(PROJECT, JOB); + when(bigqueryRpcMock.getJobSkipExceptionTranslation(PROJECT, JOB, null, EMPTY_RPC_OPTIONS)) + .thenReturn(COMPLETE_QUERY_JOB.toPb()); + when(bigqueryRpcMock.listTableDataSkipExceptionTranslation( + PROJECT, DATASET, TABLE, optionMap(BigQuery.TableDataListOption.pageToken(CURSOR)))) + .thenReturn( + new TableDataList() + .setPageToken(null) + .setRows(ImmutableList.of(TABLE_ROW)) + .setTotalRows(2L)); + + com.google.api.services.bigquery.model.QueryResponse queryResponsePb = + new com.google.api.services.bigquery.model.QueryResponse() + .setCacheHit(false) + .setJobReference(queryJob.toPb()) + .setJobComplete(true) + .setKind("bigquery#queryResponse") + .setPageToken(CURSOR) + .setRows(ImmutableList.of(TABLE_ROW)) + .setSchema(TABLE_SCHEMA.toPb()) + .setTotalBytesProcessed(42L) + .setTotalRows(BigInteger.valueOf(2L)); + + when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture())) + .thenReturn(queryResponsePb); + + bigquery = options.getService(); + TableResult firstPage = bigquery.query(QUERY_JOB_CONFIGURATION_FOR_QUERY); + assertNotNull(firstPage.getNextPage()); + assertNotNull(firstPage.getNextPage()); + + // Destination table is cached after the first getNextPage() call, so jobs.get is only called + // once even if getNextPage() is invoked multiple times on the first page. + verify(bigqueryRpcMock, times(1)) + .getJobSkipExceptionTranslation(PROJECT, JOB, null, EMPTY_RPC_OPTIONS); + verify(bigqueryRpcMock, times(2)) + .listTableDataSkipExceptionTranslation( + PROJECT, DATASET, TABLE, optionMap(BigQuery.TableDataListOption.pageToken(CURSOR))); + } + @Test void testFastQuerySlowDdl() throws InterruptedException, IOException { // mock new fast query path response when running a query that takes more than 10s