From a2dd07bedf3442d1f4bd5b0787f88c2663266da5 Mon Sep 17 00:00:00 2001 From: PhongChuong Date: Wed, 21 Aug 2024 15:12:27 -0400 Subject: [PATCH 1/2] fix: NPE for executeSelect nonFast path with empty result --- .../google/cloud/bigquery/ConnectionImpl.java | 16 ++++++---- .../cloud/bigquery/ConnectionImplTest.java | 32 +++++++++++++++++++ 2 files changed, 42 insertions(+), 6 deletions(-) diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ConnectionImpl.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ConnectionImpl.java index 2d0367790..413243396 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ConnectionImpl.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ConnectionImpl.java @@ -53,6 +53,7 @@ import com.google.common.util.concurrent.ListeningExecutorService; import com.google.common.util.concurrent.MoreExecutors; import java.io.IOException; +import java.math.BigInteger; import java.util.AbstractList; import java.util.ArrayList; import java.util.Collections; @@ -418,12 +419,15 @@ public ListenableFuture executeSelectAsync( @VisibleForTesting BigQueryResult getResultSet( GetQueryResultsResponse firstPage, JobId jobId, String sql, Boolean hasQueryParameters) { - return getSubsequentQueryResultsWithJob( - firstPage.getTotalRows().longValue(), - (long) firstPage.getRows().size(), - jobId, - firstPage, - hasQueryParameters); + if (firstPage.getTotalRows().compareTo(BigInteger.ZERO) > 0) { + return getSubsequentQueryResultsWithJob( + firstPage.getTotalRows().longValue(), + (long) firstPage.getRows().size(), + jobId, + firstPage, + hasQueryParameters); + } + return new BigQueryResultImpl(Schema.fromPb(firstPage.getSchema()), 0, null, null); } static class EndOfFieldValueList diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ConnectionImplTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ConnectionImplTest.java index dff73d6bd..fc488e4f4 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ConnectionImplTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ConnectionImplTest.java @@ -100,6 +100,15 @@ public class ConnectionImplTest { .setTotalBytesProcessed(42L) .setTotalRows(BigInteger.valueOf(1L)) .setSchema(FAST_QUERY_TABLESCHEMA); + private static final GetQueryResultsResponse GET_QUERY_RESULTS_RESPONSE_EMPTY = + new GetQueryResultsResponse() + .setJobReference(QUERY_JOB.toPb()) + .setJobComplete(true) + .setCacheHit(false) + .setPageToken(PAGE_TOKEN) + .setTotalBytesProcessed(0L) + .setTotalRows(BigInteger.valueOf(0L)) + .setSchema(FAST_QUERY_TABLESCHEMA); private static final GetQueryResultsResponse GET_QUERY_RESULTS_RESPONSE_NULL_SCHEMA = new GetQueryResultsResponse() @@ -401,6 +410,29 @@ public void testLegacyQuerySinglePage() throws BigQuerySQLException { .createJobForQuery(any(com.google.api.services.bigquery.model.Job.class)); } + // calls executeSelect with a nonFast query where the query returns an empty result. + @Test + public void testLegacyQuerySinglePageEmptyResults() throws BigQuerySQLException { + ConnectionImpl connectionSpy = Mockito.spy(connection); + com.google.api.services.bigquery.model.Job jobResponseMock = + new com.google.api.services.bigquery.model.Job() + .setJobReference(QUERY_JOB.toPb()) + .setId(JOB) + .setStatus(new com.google.api.services.bigquery.model.JobStatus().setState("DONE")); + // emulating a legacy query + doReturn(false).when(connectionSpy).isFastQuerySupported(); + doReturn(GET_QUERY_RESULTS_RESPONSE_EMPTY) + .when(connectionSpy) + .getQueryResultsFirstPage(any(JobId.class)); + when(bigqueryRpcMock.createJobForQuery(any(com.google.api.services.bigquery.model.Job.class))) + .thenReturn(jobResponseMock); // RPC call in createQueryJob + BigQueryResult res = connectionSpy.executeSelect(SQL_QUERY); + assertEquals(res.getTotalRows(), 0); + assertEquals(QUERY_SCHEMA, res.getSchema()); + verify(bigqueryRpcMock, times(1)) + .createJobForQuery(any(com.google.api.services.bigquery.model.Job.class)); + } + // exercises getSubsequentQueryResultsWithJob for fast running queries @Test public void testFastQueryLongRunning() throws SQLException { From 45a18ca7992c8dbd60ecd1fac813fe80f3ec31c1 Mon Sep 17 00:00:00 2001 From: PhongChuong Date: Wed, 21 Aug 2024 15:18:40 -0400 Subject: [PATCH 2/2] remove commented out extraneous code --- .../test/java/com/google/cloud/bigquery/ConnectionImplTest.java | 1 - 1 file changed, 1 deletion(-) diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ConnectionImplTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ConnectionImplTest.java index fc488e4f4..58cb69ba7 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ConnectionImplTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ConnectionImplTest.java @@ -384,7 +384,6 @@ public void testLegacyQuerySinglePage() throws BigQuerySQLException { ConnectionImpl connectionSpy = Mockito.spy(connection); com.google.api.services.bigquery.model.Job jobResponseMock = new com.google.api.services.bigquery.model.Job() - // .setConfiguration(QUERY_JOB.g) .setJobReference(QUERY_JOB.toPb()) .setId(JOB) .setStatus(new com.google.api.services.bigquery.model.JobStatus().setState("DONE"));