feat(bigquery-jdbc): expose QueryStatistics for jobless queries (#14596)

Author: logachev

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java

@@ -115,6 +115,10 @@ public QueryStatistics getQueryStatistics() {
     return queryStatistics;
   }
 
+  protected void setQueryStatistics(QueryStatistics queryStatistics) {
+    this.queryStatistics = queryStatistics;
+  }
+
   public void setJobId(JobId jobId) {
     this.jobId = jobId;
   }

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryResultSet.java

@@ -38,7 +38,7 @@ public interface BigQueryResultSet {
 
   /*
    * Returns com.google.cloud.bigquery.JobStatistics.QueryStatistics object with statistics for the
-   * completed Job for non-low latency queries.
+   * completed query.
    *
    * @return QueryStatistics object or null.
    */

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java

@@ -918,6 +918,9 @@ ResultSet processArrowResultSet(TableResult results, Job job) throws SQLExceptio
               arrowResultSet, referenceQueueArrowRs, populateBufferWorker));
       arrowResultSet.setJobId(currentJobId);
       arrowResultSet.setQueryId(results.getQueryId());
+      if (job == null && currentJobId == null) {
+        arrowResultSet.setQueryStatistics(results.extractQueryStatistics());
+      }
       return arrowResultSet;
 
     } catch (Exception | OutOfMemoryError ex) {
@@ -1208,6 +1211,9 @@ BigQueryJsonResultSet processJsonResultSet(TableResult results, Job job) throws
             job);
     jsonResultSet.setJobId(jobId);
     jsonResultSet.setQueryId(results.getQueryId());
+    if (job == null && jobId == null) {
+      jsonResultSet.setQueryStatistics(results.extractQueryStatistics());
+    }
     jsonResultSetFinalizers.add(
         new BigQueryResultSetFinalizers.JsonResultSetFinalizer(
             jsonResultSet, referenceQueueJsonRs, jsonWorkers));

java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSetTest.java

@@ -322,4 +322,23 @@ public void testSetJobUpdatesJobId() {
 
     assertThat(resultSet.getJobId()).isEqualTo(newJobId);
   }
+
+  @Test
+  public void testQueryStatisticsWithPrePopulatedStatisticsForStatelessQuery() {
+    QueryStatistics statelessStats = mock(QueryStatistics.class);
+    doReturn(10485760L).when(statelessStats).getTotalBytesProcessed();
+    doReturn(false).when(statelessStats).getCacheHit();
+    doReturn(250L).when(statelessStats).getTotalSlotMs();
+    resultSet.setQueryId("stateless-query-id");
+    resultSet.setQueryStatistics(statelessStats);
+
+    assertThat(resultSet.getJobId()).isNull();
+    assertThat(resultSet.getJob()).isNull();
+    QueryStatistics stats = resultSet.getQueryStatistics();
+    assertThat(stats).isSameInstanceAs(statelessStats);
+    assertThat(stats.getTotalBytesProcessed()).isEqualTo(10485760L);
+    assertThat(stats.getCacheHit()).isFalse();
+    assertThat(stats.getTotalSlotMs()).isEqualTo(250L);
+    verify(bigQuery, never()).getJob(any(JobId.class));
+  }
 }

java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java

@@ -1270,4 +1270,52 @@ public void testEnableTimestampPicosPropagation() {
     BigQueryStatement statement = new BigQueryStatement(bigQueryConnection);
     assertTrue(statement.isEnableTimestampPicos());
   }
+
+  @Test
+  public void testJoblessQueryPopulatesQueryStatistics() throws Exception {
+    doReturn(true).when(bigQueryConnection).getUseStatelessQueryMode();
+    BigQueryStatement joblessStatement = new BigQueryStatement(bigQueryConnection);
+
+    Schema expectedSchema = Schema.of(fieldList);
+    SessionInfo expectedSessionInfo = mock(SessionInfo.class);
+    doReturn("session_abc").when(expectedSessionInfo).getSessionId();
+    TableResult tableResultMock = mock(TableResult.class);
+    doReturn("stateless-query-id").when(tableResultMock).getQueryId();
+    doReturn(null).when(tableResultMock).getJobId();
+    doReturn(expectedSchema).when(tableResultMock).getSchema();
+    doReturn(1L).when(tableResultMock).getTotalRows();
+    doReturn(ImmutableList.<FieldValueList>of()).when(tableResultMock).getValues();
+    doReturn(StatementType.SELECT).when(tableResultMock).getStatementType();
+    QueryStatistics expectedStats = mock(QueryStatistics.class);
+    doReturn(10485760L).when(expectedStats).getTotalBytesProcessed();
+    doReturn(10485760L).when(expectedStats).getTotalBytesBilled();
+    doReturn(320L).when(expectedStats).getTotalSlotMs();
+    doReturn(false).when(expectedStats).getCacheHit();
+    doReturn(StatementType.SELECT).when(expectedStats).getStatementType();
+    doReturn(expectedSessionInfo).when(expectedStats).getSessionInfo();
+    doReturn(expectedSchema).when(expectedStats).getSchema();
+    doReturn(expectedStats).when(tableResultMock).extractQueryStatistics();
+
+    doReturn(tableResultMock)
+        .when(bigquery)
+        .queryWithTimeout(any(QueryJobConfiguration.class), any(), any());
+
+    ResultSet rs = joblessStatement.executeQuery("SELECT * FROM test");
+    assertNotNull(rs);
+    BigQueryResultSet bqRs = rs.unwrap(BigQueryResultSet.class);
+    assertNull(bqRs.getJobId());
+    assertEquals("stateless-query-id", bqRs.getQueryId());
+
+    QueryStatistics stats = bqRs.getQueryStatistics();
+    assertNotNull(stats);
+    assertEquals(10485760L, stats.getTotalBytesProcessed().longValue());
+    assertEquals(10485760L, stats.getTotalBytesBilled().longValue());
+    assertEquals(320L, stats.getTotalSlotMs().longValue());
+    assertFalse(stats.getCacheHit());
+    assertEquals(StatementType.SELECT, stats.getStatementType());
+    assertEquals(expectedSessionInfo, stats.getSessionInfo());
+    assertEquals(expectedSchema, stats.getSchema());
+
+    verify(bigquery, Mockito.never()).getJob(any(JobId.class));
+  }
 }