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
googleapis · Oct 19, 2020 · f2799dd · f2799dd
1 parent 137eeaa
commit f2799dd
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 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
@@ -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
@@ -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