feat(bigquery): expose StatementType and query execution stats on TableResult

java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java

@@ -41,6 +41,8 @@
 import com.google.cloud.Tuple;
 import com.google.cloud.bigquery.BigQueryRetryHelper.BigQueryRetryHelperException;
 import com.google.cloud.bigquery.InsertAllRequest.RowToInsert;
+import com.google.cloud.bigquery.JobStatistics.QueryStatistics.StatementType;
+import com.google.cloud.bigquery.JobStatistics.SessionInfo;
 import com.google.cloud.bigquery.spi.v2.BigQueryRpc;
 import com.google.cloud.bigquery.spi.v2.HttpBigQueryRpc;
 import com.google.common.annotations.VisibleForTesting;
@@ -2095,6 +2097,17 @@ public com.google.api.services.bigquery.model.QueryResponse call()
       return job;
     }
 
+    StatementType statementType =
+        results.getStatementType() != null
+            ? StatementType.valueOf(results.getStatementType())
+            : null;
+    Long totalBytesBilled = results.getTotalBytesBilled();
+    Long totalBytesProcessed = results.getTotalBytesProcessed();
+    Long totalSlotMs = results.getTotalSlotMs();
+    Long numDmlAffectedRows = results.getNumDmlAffectedRows();
+    SessionInfo sessionInfo =
+        results.getSessionInfo() != null ? SessionInfo.fromPb(results.getSessionInfo()) : null;
+
     if (results.getPageToken() != null) {
       JobId jobId = JobId.fromPb(results.getJobReference());
       String cursor = results.getPageToken();
@@ -2114,6 +2127,12 @@ public com.google.api.services.bigquery.model.QueryResponse call()
           .setQueryId(results.getQueryId())
           .setJobCreationReason(JobCreationReason.fromPb(results.getJobCreationReason()))
           .setRowsInPage(results.getRows() != null ? (long) results.getRows().size() : 0L)
+          .setStatementType(statementType)
+          .setTotalBytesBilled(totalBytesBilled)
+          .setTotalBytesProcessed(totalBytesProcessed)
+          .setTotalSlotMs(totalSlotMs)
+          .setNumDmlAffectedRows(numDmlAffectedRows)
+          .setSessionInfo(sessionInfo)
           .build();
     }
     // only 1 page of result
@@ -2134,6 +2153,12 @@ public com.google.api.services.bigquery.model.QueryResponse call()
         .setQueryId(results.getQueryId())
         .setJobCreationReason(JobCreationReason.fromPb(results.getJobCreationReason()))
         .setRowsInPage(results.getRows() != null ? (long) results.getRows().size() : 0L)
+        .setStatementType(statementType)
+        .setTotalBytesBilled(totalBytesBilled)
+        .setTotalBytesProcessed(totalBytesProcessed)
+        .setTotalSlotMs(totalSlotMs)
+        .setNumDmlAffectedRows(numDmlAffectedRows)
+        .setSessionInfo(sessionInfo)
         .build();
   }
 

java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Job.java

@@ -28,6 +28,9 @@
 import com.google.cloud.bigquery.BigQuery.QueryResultsOption;
 import com.google.cloud.bigquery.BigQuery.TableDataListOption;
 import com.google.cloud.bigquery.JobConfiguration.Type;
+import com.google.cloud.bigquery.JobStatistics.QueryStatistics;
+import com.google.cloud.bigquery.JobStatistics.QueryStatistics.StatementType;
+import com.google.cloud.bigquery.JobStatistics.SessionInfo;
 import com.google.common.collect.ImmutableList;
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.trace.Span;
@@ -414,6 +417,17 @@ public TableResult getQueryResults(QueryResultsOption... options)
                 : ImmutableList.copyOf(job.getStatus().getExecutionErrors()));
       }
 
+      QueryStatistics stats =
+          job.getStatistics() instanceof QueryStatistics
+              ? (QueryStatistics) job.getStatistics()
+              : null;
+      StatementType statementType = stats != null ? stats.getStatementType() : null;
+      Long totalBytesBilled = stats != null ? stats.getTotalBytesBilled() : null;
+      Long totalBytesProcessed = stats != null ? stats.getTotalBytesProcessed() : null;
+      Long totalSlotMs = stats != null ? stats.getTotalSlotMs() : null;
+      Long numDmlAffectedRows = stats != null ? stats.getNumDmlAffectedRows() : null;
+      SessionInfo sessionInfo = stats != null ? stats.getSessionInfo() : null;
+
       // If there are no rows in the result, this may have been a DDL query.
       // Listing table data might fail, such as with CREATE VIEW queries.
       // Avoid a tabledata.list API request by returning an empty TableResult.
@@ -425,6 +439,12 @@ public TableResult getQueryResults(QueryResultsOption... options)
                 .setTotalRows(0L)
                 .setPageNoSchema(new PageImpl<FieldValueList>(null, "", null))
                 .setRowsInPage(0L)
+                .setStatementType(statementType)
+                .setTotalBytesBilled(totalBytesBilled)
+                .setTotalBytesProcessed(totalBytesProcessed)
+                .setTotalSlotMs(totalSlotMs)
+                .setNumDmlAffectedRows(numDmlAffectedRows)
+                .setSessionInfo(sessionInfo)
                 .build();
         return emptyTableResult;
       }
@@ -436,7 +456,16 @@ public TableResult getQueryResults(QueryResultsOption... options)
       TableResult tableResult =
           bigquery.listTableData(
               table, response.getSchema(), listOptions.toArray(new TableDataListOption[0]));
-      TableResult tableResultWithJobId = tableResult.toBuilder().setJobId(job.getJobId()).build();
+      TableResult tableResultWithJobId =
+          tableResult.toBuilder()
+              .setJobId(job.getJobId())
+              .setStatementType(statementType)
+              .setTotalBytesBilled(totalBytesBilled)
+              .setTotalBytesProcessed(totalBytesProcessed)
+              .setTotalSlotMs(totalSlotMs)
+              .setNumDmlAffectedRows(numDmlAffectedRows)
+              .setSessionInfo(sessionInfo)
+              .build();
       return tableResultWithJobId;
     } finally {
       if (getQueryResults != null) {

java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableResult.java

@@ -18,13 +18,15 @@
 
 import com.google.api.gax.paging.Page;
 import com.google.auto.value.AutoValue;
+import com.google.cloud.bigquery.JobStatistics.QueryStatistics.StatementType;
+import com.google.cloud.bigquery.JobStatistics.SessionInfo;
 import com.google.common.base.Function;
 import com.google.common.base.MoreObjects;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Iterators;
 import java.io.Serializable;
 import java.util.Objects;
-import javax.annotation.Nullable;
+import org.jspecify.annotations.Nullable;
 
 @AutoValue
 public abstract class TableResult implements Page<FieldValueList>, Serializable {
@@ -51,6 +53,18 @@ public abstract static class Builder {
 
     abstract TableResult.Builder setRowsInPage(Long rowsInPage);
 
+    abstract TableResult.Builder setStatementType(@Nullable StatementType statementType);
+
+    abstract TableResult.Builder setTotalBytesBilled(@Nullable Long totalBytesBilled);
+
+    abstract TableResult.Builder setTotalBytesProcessed(@Nullable Long totalBytesProcessed);
+
+    abstract TableResult.Builder setTotalSlotMs(@Nullable Long totalSlotMs);
+
+    abstract TableResult.Builder setNumDmlAffectedRows(@Nullable Long numDmlAffectedRows);
+
+    abstract TableResult.Builder setSessionInfo(@Nullable SessionInfo sessionInfo);
+
     /** Creates a @code TableResult} object. */
     public abstract TableResult build();
   }
@@ -62,8 +76,7 @@ public static Builder newBuilder() {
   }
 
   /** Returns the schema of the results. Null if the schema is not supplied. */
-  @Nullable
-  public abstract Schema getSchema();
+  public abstract @Nullable Schema getSchema();
 
   /**
    * Returns the total number of rows in the complete result set, which can be more than the number
@@ -74,18 +87,58 @@ public static Builder newBuilder() {
 
   public abstract Page<FieldValueList> getPageNoSchema();
 
-  @Nullable
-  public abstract JobId getJobId();
+  public abstract @Nullable JobId getJobId();
 
-  @Nullable
-  public abstract String getQueryId();
+  public abstract @Nullable String getQueryId();
 
-  @Nullable
-  public abstract JobCreationReason getJobCreationReason();
+  public abstract @Nullable JobCreationReason getJobCreationReason();
 
   /** Returns the number of rows in the current page of results. */
-  @Nullable
-  public abstract Long getRowsInPage();
+  public abstract @Nullable Long getRowsInPage();
+
+  /**
+   * Returns the statement type of the query (e.g. SELECT, INSERT, UPDATE, DDL, SCRIPT).
+   *
+   * @return statement type, or {@code null} if not populated by the service
+   */
+  public abstract @Nullable StatementType getStatementType();
+
+  /**
+   * Returns the total number of bytes billed for the query.
+   *
+   * @return total bytes billed, or {@code null} if not populated by the service
+   */
+  public abstract @Nullable Long getTotalBytesBilled();
+
+  /**
+   * Returns the total number of bytes processed by the query.
+   *
+   * @return total bytes processed, or {@code null} if not populated by the service
+   */
+  public abstract @Nullable Long getTotalBytesProcessed();
+
+  /**
+   * Returns the total slot milliseconds consumed by the query.
+   *
+   * @return total slot milliseconds, or {@code null} if not populated by the service
+   */
+  public abstract @Nullable Long getTotalSlotMs();
+
+  /**
+   * Returns the number of rows affected by a DML statement (INSERT, UPDATE, DELETE, MERGE).
+   *
+   * @return number of affected rows for DML queries, or {@code null} if not populated by the
+   *     service
+   */
+  public abstract @Nullable Long getNumDmlAffectedRows();
+
+  /**
+   * Returns information about the BigQuery session if this query was executed within or created a
+   * session.
+   *
+   * @return session information, or {@code null} if not populated by the service
+   */
+  public abstract @Nullable SessionInfo getSessionInfo();
 
   @Override
   public boolean hasNextPage() {
@@ -109,6 +162,12 @@ public TableResult getNextPage() {
           .setQueryId(getQueryId())
           .setJobCreationReason(getJobCreationReason())
           .setRowsInPage(nextRows)
+          .setStatementType(getStatementType())
+          .setTotalBytesBilled(getTotalBytesBilled())
+          .setTotalBytesProcessed(getTotalBytesProcessed())
+          .setTotalSlotMs(getTotalSlotMs())
+          .setNumDmlAffectedRows(getNumDmlAffectedRows())
+          .setSessionInfo(getSessionInfo())
           .build();
     }
     return null;
@@ -147,13 +206,29 @@ public String toString() {
         .add("cursor", getNextPageToken())
         .add("queryId", getQueryId())
         .add("rowsInPage", getRowsInPage())
+        .add("statementType", getStatementType())
+        .add("totalBytesBilled", getTotalBytesBilled())
+        .add("totalBytesProcessed", getTotalBytesProcessed())
+        .add("totalSlotMs", getTotalSlotMs())
+        .add("numDmlAffectedRows", getNumDmlAffectedRows())
+        .add("sessionInfo", getSessionInfo())
         .toString();
   }
 
   @Override
   public final int hashCode() {
     return Objects.hash(
-        getPageNoSchema(), getSchema(), getTotalRows(), getQueryId(), getRowsInPage());
+        getPageNoSchema(),
+        getSchema(),
+        getTotalRows(),
+        getQueryId(),
+        getRowsInPage(),
+        getStatementType(),
+        getTotalBytesBilled(),
+        getTotalBytesProcessed(),
+        getTotalSlotMs(),
+        getNumDmlAffectedRows(),
+        getSessionInfo());
   }
 
   @Override
@@ -170,6 +245,12 @@ public final boolean equals(Object obj) {
         && Objects.equals(getSchema(), response.getSchema())
         && getTotalRows() == response.getTotalRows()
         && Objects.equals(getQueryId(), response.getQueryId())
-        && Objects.equals(getRowsInPage(), response.getRowsInPage());
+        && Objects.equals(getRowsInPage(), response.getRowsInPage())
+        && Objects.equals(getStatementType(), response.getStatementType())
+        && Objects.equals(getTotalBytesBilled(), response.getTotalBytesBilled())
+        && Objects.equals(getTotalBytesProcessed(), response.getTotalBytesProcessed())
+        && Objects.equals(getTotalSlotMs(), response.getTotalSlotMs())
+        && Objects.equals(getNumDmlAffectedRows(), response.getNumDmlAffectedRows())
+        && Objects.equals(getSessionInfo(), response.getSessionInfo());
   }
 }

java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java

@@ -52,6 +52,7 @@
 import com.google.api.services.bigquery.model.ProjectList;
 import com.google.api.services.bigquery.model.ProjectReference;
 import com.google.api.services.bigquery.model.QueryRequest;
+import com.google.api.services.bigquery.model.SessionInfo;
 import com.google.api.services.bigquery.model.TableCell;
 import com.google.api.services.bigquery.model.TableDataInsertAllRequest;
 import com.google.api.services.bigquery.model.TableDataInsertAllResponse;
@@ -65,6 +66,7 @@
 import com.google.cloud.bigquery.BigQuery.JobOption;
 import com.google.cloud.bigquery.BigQuery.QueryResultsOption;
 import com.google.cloud.bigquery.InsertAllRequest.RowToInsert;
+import com.google.cloud.bigquery.JobStatistics.QueryStatistics.StatementType;
 import com.google.cloud.bigquery.spi.BigQueryRpcFactory;
 import com.google.cloud.bigquery.spi.v2.BigQueryRpc;
 import com.google.cloud.bigquery.spi.v2.HttpBigQueryRpc;
@@ -169,6 +171,8 @@ public class BigQueryImplTest {
           .setField("timestampField");
   private static final TimePartitioning TIME_PARTITIONING_NULL_TYPE =
       TimePartitioning.fromPb(PB_TIMEPARTITIONING);
+  private static final String SESSION_ID = "test-session-id";
+  private static final SessionInfo PB_SESSION_INFO = new SessionInfo().setSessionId(SESSION_ID);
   private static final ImmutableMap<String, String> LABELS = ImmutableMap.of("key", "value");
   private static final StandardTableDefinition TABLE_DEFINITION_WITH_PARTITIONING =
       StandardTableDefinition.newBuilder()
@@ -2874,7 +2878,12 @@ void testQueryWithTimeoutSetsTimeout() throws InterruptedException, IOException
             .setPageToken(null)
             .setRows(ImmutableList.of(TABLE_ROW))
             .setSchema(TABLE_SCHEMA.toPb())
+            .setStatementType("SELECT")
+            .setTotalBytesBilled(100L)
             .setTotalBytesProcessed(42L)
+            .setTotalSlotMs(50L)
+            .setNumDmlAffectedRows(0L)
+            .setSessionInfo(PB_SESSION_INFO)
             .setTotalRows(BigInteger.valueOf(1L));
 
     when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture()))
@@ -2883,6 +2892,14 @@ void testQueryWithTimeoutSetsTimeout() throws InterruptedException, IOException
     bigquery = options.getService();
     Object result = bigquery.queryWithTimeout(QUERY_JOB_CONFIGURATION_FOR_QUERY, null, 1000L);
     assertTrue(result instanceof TableResult);
+    TableResult tableResult = (TableResult) result;
+    assertEquals(StatementType.SELECT, tableResult.getStatementType());
+    assertEquals((Long) 100L, tableResult.getTotalBytesBilled());
+    assertEquals((Long) 42L, tableResult.getTotalBytesProcessed());
+    assertEquals((Long) 50L, tableResult.getTotalSlotMs());
+    assertEquals((Long) 0L, tableResult.getNumDmlAffectedRows());
+    assertNotNull(tableResult.getSessionInfo());
+    assertEquals(SESSION_ID, tableResult.getSessionInfo().getSessionId());
     QueryRequest requestPb = requestPbCapture.getValue();
     assertEquals((Long) 1000L, requestPb.getTimeoutMs());
   }