From f2799dd9a4323460843aa6df7ac7c2f092c7f1b3 Mon Sep 17 00:00:00 2001
From: Praful Makani <praful@qlogic.io>
Date: Mon, 19 Oct 2020 22:44:57 +0530
Subject: [PATCH] fix: schema and totalRows for duplicate request ids (#804)

* fix: schema and totalRows for duplicate request ids

* fix: modify code and test case

* fix: address comment
---
 .../google/cloud/bigquery/BigQueryImpl.java   | 25 +++++++----
 .../cloud/bigquery/it/ITBigQueryTest.java     | 44 +++++++++++++++++++
 2 files changed, 61 insertions(+), 8 deletions(-)

diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java
index 76067be0d..a19c65e63 100644
--- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java
+++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java
@@ -1239,7 +1239,8 @@ public TableResult query(QueryJobConfiguration configuration, JobOption... optio
   }
 
   private TableResult queryRpc(
-      final String projectId, final QueryRequest content, JobOption... options) {
+      final String projectId, final QueryRequest content, JobOption... options)
+      throws InterruptedException {
     com.google.api.services.bigquery.model.QueryResponse results;
     try {
       results =
@@ -1265,14 +1266,22 @@ public com.google.api.services.bigquery.model.QueryResponse call() {
       throw new BigQueryException(bigQueryErrors);
     }
 
-    Schema schema = results.getSchema() == null ? null : Schema.fromPb(results.getSchema());
-    Long numRows;
-    if (results.getNumDmlAffectedRows() == null && results.getTotalRows() == null) {
-      numRows = 0L;
-    } else if (results.getNumDmlAffectedRows() != null) {
-      numRows = results.getNumDmlAffectedRows();
+    long numRows;
+    Schema schema;
+    if (results.getSchema() == null && results.getJobComplete()) {
+      JobId jobId = JobId.fromPb(results.getJobReference());
+      Job job = getJob(jobId, options);
+      TableResult tableResult = job.getQueryResults();
+      return tableResult;
     } else {
-      numRows = results.getTotalRows().longValue();
+      schema = results.getSchema() == null ? null : Schema.fromPb(results.getSchema());
+      if (results.getNumDmlAffectedRows() == null && results.getTotalRows() == null) {
+        numRows = 0L;
+      } else if (results.getNumDmlAffectedRows() != null) {
+        numRows = results.getNumDmlAffectedRows();
+      } else {
+        numRows = results.getTotalRows().longValue();
+      }
     }
 
     if (results.getPageToken() != null) {
diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java
index 05332ba6a..36ad12d68 100644
--- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java
+++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java
@@ -1608,6 +1608,35 @@ public void testFastQueryMultipleRuns() throws InterruptedException {
     assertFalse(result2.hasNextPage());
   }
 
+  @Test
+  public void testFastQuerySinglePageDuplicateRequestIds() throws InterruptedException {
+    String query =
+        "SELECT TimestampField, StringField, BooleanField FROM " + TABLE_ID_FASTQUERY.getTable();
+    QueryJobConfiguration config =
+        QueryJobConfiguration.newBuilder(query).setDefaultDataset(DatasetId.of(DATASET)).build();
+    TableResult result = bigquery.query(config);
+    assertEquals(QUERY_RESULT_SCHEMA, result.getSchema());
+    assertEquals(2, result.getTotalRows());
+    assertNull(result.getNextPage());
+    assertNull(result.getNextPageToken());
+    assertFalse(result.hasNextPage());
+
+    TableResult result1 = bigquery.query(config);
+    assertEquals(QUERY_RESULT_SCHEMA, result1.getSchema());
+    assertEquals(2, result1.getTotalRows());
+    assertNull(result1.getNextPage());
+    assertNull(result1.getNextPageToken());
+    assertFalse(result1.hasNextPage());
+
+    config.toBuilder().setQuery(query).build();
+    TableResult result2 = bigquery.query(config);
+    assertEquals(QUERY_RESULT_SCHEMA, result2.getSchema());
+    assertEquals(2, result2.getTotalRows());
+    assertNull(result2.getNextPage());
+    assertNull(result2.getNextPageToken());
+    assertFalse(result2.hasNextPage());
+  }
+
   @Test
   public void testFastSQLQuery() throws InterruptedException {
     String query =
@@ -1650,6 +1679,21 @@ public void testFastSQLQueryMultiPage() throws InterruptedException {
     assertNotNull(result.getNextPage());
     assertNotNull(result.getNextPageToken());
     assertTrue(result.hasNextPage());
+
+    TableResult result1 = bigquery.query(config);
+    assertEquals(LARGE_TABLE_SCHEMA, result.getSchema());
+    assertEquals(313348, result.getTotalRows());
+    assertNotNull(result1.getNextPage());
+    assertNotNull(result1.getNextPageToken());
+    assertTrue(result1.hasNextPage());
+
+    config.toBuilder().setQuery(query).build();
+    TableResult result2 = bigquery.query(config);
+    assertEquals(LARGE_TABLE_SCHEMA, result2.getSchema());
+    assertEquals(313348, result2.getTotalRows());
+    assertNotNull(result2.getNextPage());
+    assertNotNull(result2.getNextPageToken());
+    assertTrue(result2.hasNextPage());
   }
 
   @Test