fix: NPE for executeSelect nonFast path with empty result (#3445)

PhongChuong · web-flow · commit d0d758a6e5e9 · 2024-08-22T12:48:48.000-04:00
* fix: NPE for executeSelect nonFast path with empty result

* remove commented out extraneous code
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
@@ -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<ExecuteSelectResponse> 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
@@ -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()
@@ -375,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"));
@@ -401,6 +409,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 {
