Skip to content

Commit 7d16de8

Browse files
keshavdandevakeshavdandeva
authored
feat(bigquery): expose StatementType and query execution stats on TableResult (#14145)
b/549680449 This PR exposes `StatementType` and jobless query execution metrics on `TableResult` by plumbing them from `QueryResponse` and `JobStatistics.QueryStatistics`. This enables downstream consumers (such as the BigQuery JDBC driver) to inspect query statement types and execution statistics directly from `TableResult` without needing to issue secondary jobs or dry-run queries. ## Changes - **Dependency Update**: Bumped `google-api-services-bigquery` to `v2-rev20260731-2.0.0` in `google-cloud-jar-parent/pom.xml` and `java-bigquery/pom.xml`. - **TableResult**: Added getters, builder setters, pagination propagation (`getNextPage()`), `toString()`, `hashCode()`, and `equals()` for: - `getStatementType()` (`StatementType`) - `getTotalBytesBilled()` (`Long`) - `getTotalBytesProcessed()` (`Long`) - `getTotalSlotMs()` (`Long`) - `getNumDmlAffectedRows()` (`Long`) - **BigQueryImpl**: Plumbed these fields from `QueryResponse` when building `TableResult` in `queryRpc()`. - **Job**: Plumbed `StatementType` and query metrics from `QueryStatistics` when creating `TableResult` in `getQueryResults()`. - **Tests**: Added unit tests in `TableResultTest` and `BigQueryImplTest` verifying field retrieval, pagination propagation, builder modifications, and serialization.
1 parent 0d5fac0 commit 7d16de8

6 files changed

Lines changed: 291 additions & 15 deletions

File tree

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

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,8 @@
4141
import com.google.cloud.Tuple;
4242
import com.google.cloud.bigquery.BigQueryRetryHelper.BigQueryRetryHelperException;
4343
import com.google.cloud.bigquery.InsertAllRequest.RowToInsert;
44+
import com.google.cloud.bigquery.JobStatistics.QueryStatistics.StatementType;
45+
import com.google.cloud.bigquery.JobStatistics.SessionInfo;
4446
import com.google.cloud.bigquery.spi.v2.BigQueryRpc;
4547
import com.google.cloud.bigquery.spi.v2.HttpBigQueryRpc;
4648
import com.google.common.annotations.VisibleForTesting;
@@ -2095,6 +2097,17 @@ public com.google.api.services.bigquery.model.QueryResponse call()
20952097
return job;
20962098
}
20972099

2100+
StatementType statementType =
2101+
results.getStatementType() != null
2102+
? StatementType.valueOf(results.getStatementType())
2103+
: null;
2104+
Long totalBytesBilled = results.getTotalBytesBilled();
2105+
Long totalBytesProcessed = results.getTotalBytesProcessed();
2106+
Long totalSlotMs = results.getTotalSlotMs();
2107+
Long numDmlAffectedRows = results.getNumDmlAffectedRows();
2108+
SessionInfo sessionInfo =
2109+
results.getSessionInfo() != null ? SessionInfo.fromPb(results.getSessionInfo()) : null;
2110+
20982111
if (results.getPageToken() != null) {
20992112
JobId jobId = JobId.fromPb(results.getJobReference());
21002113
String cursor = results.getPageToken();
@@ -2114,6 +2127,12 @@ public com.google.api.services.bigquery.model.QueryResponse call()
21142127
.setQueryId(results.getQueryId())
21152128
.setJobCreationReason(JobCreationReason.fromPb(results.getJobCreationReason()))
21162129
.setRowsInPage(results.getRows() != null ? (long) results.getRows().size() : 0L)
2130+
.setStatementType(statementType)
2131+
.setTotalBytesBilled(totalBytesBilled)
2132+
.setTotalBytesProcessed(totalBytesProcessed)
2133+
.setTotalSlotMs(totalSlotMs)
2134+
.setNumDmlAffectedRows(numDmlAffectedRows)
2135+
.setSessionInfo(sessionInfo)
21172136
.build();
21182137
}
21192138
// only 1 page of result
@@ -2134,6 +2153,12 @@ public com.google.api.services.bigquery.model.QueryResponse call()
21342153
.setQueryId(results.getQueryId())
21352154
.setJobCreationReason(JobCreationReason.fromPb(results.getJobCreationReason()))
21362155
.setRowsInPage(results.getRows() != null ? (long) results.getRows().size() : 0L)
2156+
.setStatementType(statementType)
2157+
.setTotalBytesBilled(totalBytesBilled)
2158+
.setTotalBytesProcessed(totalBytesProcessed)
2159+
.setTotalSlotMs(totalSlotMs)
2160+
.setNumDmlAffectedRows(numDmlAffectedRows)
2161+
.setSessionInfo(sessionInfo)
21372162
.build();
21382163
}
21392164

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

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,9 @@
2828
import com.google.cloud.bigquery.BigQuery.QueryResultsOption;
2929
import com.google.cloud.bigquery.BigQuery.TableDataListOption;
3030
import com.google.cloud.bigquery.JobConfiguration.Type;
31+
import com.google.cloud.bigquery.JobStatistics.QueryStatistics;
32+
import com.google.cloud.bigquery.JobStatistics.QueryStatistics.StatementType;
33+
import com.google.cloud.bigquery.JobStatistics.SessionInfo;
3134
import com.google.common.collect.ImmutableList;
3235
import io.opentelemetry.api.common.Attributes;
3336
import io.opentelemetry.api.trace.Span;
@@ -414,6 +417,17 @@ public TableResult getQueryResults(QueryResultsOption... options)
414417
: ImmutableList.copyOf(job.getStatus().getExecutionErrors()));
415418
}
416419

420+
QueryStatistics stats =
421+
job.getStatistics() instanceof QueryStatistics
422+
? (QueryStatistics) job.getStatistics()
423+
: null;
424+
StatementType statementType = stats != null ? stats.getStatementType() : null;
425+
Long totalBytesBilled = stats != null ? stats.getTotalBytesBilled() : null;
426+
Long totalBytesProcessed = stats != null ? stats.getTotalBytesProcessed() : null;
427+
Long totalSlotMs = stats != null ? stats.getTotalSlotMs() : null;
428+
Long numDmlAffectedRows = stats != null ? stats.getNumDmlAffectedRows() : null;
429+
SessionInfo sessionInfo = stats != null ? stats.getSessionInfo() : null;
430+
417431
// If there are no rows in the result, this may have been a DDL query.
418432
// Listing table data might fail, such as with CREATE VIEW queries.
419433
// Avoid a tabledata.list API request by returning an empty TableResult.
@@ -425,6 +439,12 @@ public TableResult getQueryResults(QueryResultsOption... options)
425439
.setTotalRows(0L)
426440
.setPageNoSchema(new PageImpl<FieldValueList>(null, "", null))
427441
.setRowsInPage(0L)
442+
.setStatementType(statementType)
443+
.setTotalBytesBilled(totalBytesBilled)
444+
.setTotalBytesProcessed(totalBytesProcessed)
445+
.setTotalSlotMs(totalSlotMs)
446+
.setNumDmlAffectedRows(numDmlAffectedRows)
447+
.setSessionInfo(sessionInfo)
428448
.build();
429449
return emptyTableResult;
430450
}
@@ -436,7 +456,16 @@ public TableResult getQueryResults(QueryResultsOption... options)
436456
TableResult tableResult =
437457
bigquery.listTableData(
438458
table, response.getSchema(), listOptions.toArray(new TableDataListOption[0]));
439-
TableResult tableResultWithJobId = tableResult.toBuilder().setJobId(job.getJobId()).build();
459+
TableResult tableResultWithJobId =
460+
tableResult.toBuilder()
461+
.setJobId(job.getJobId())
462+
.setStatementType(statementType)
463+
.setTotalBytesBilled(totalBytesBilled)
464+
.setTotalBytesProcessed(totalBytesProcessed)
465+
.setTotalSlotMs(totalSlotMs)
466+
.setNumDmlAffectedRows(numDmlAffectedRows)
467+
.setSessionInfo(sessionInfo)
468+
.build();
440469
return tableResultWithJobId;
441470
} finally {
442471
if (getQueryResults != null) {

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

Lines changed: 94 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,15 @@
1818

1919
import com.google.api.gax.paging.Page;
2020
import com.google.auto.value.AutoValue;
21+
import com.google.cloud.bigquery.JobStatistics.QueryStatistics.StatementType;
22+
import com.google.cloud.bigquery.JobStatistics.SessionInfo;
2123
import com.google.common.base.Function;
2224
import com.google.common.base.MoreObjects;
2325
import com.google.common.collect.Iterables;
2426
import com.google.common.collect.Iterators;
2527
import java.io.Serializable;
2628
import java.util.Objects;
27-
import javax.annotation.Nullable;
29+
import org.jspecify.annotations.Nullable;
2830

2931
@AutoValue
3032
public abstract class TableResult implements Page<FieldValueList>, Serializable {
@@ -51,6 +53,18 @@ public abstract static class Builder {
5153

5254
abstract TableResult.Builder setRowsInPage(Long rowsInPage);
5355

56+
abstract TableResult.Builder setStatementType(@Nullable StatementType statementType);
57+
58+
abstract TableResult.Builder setTotalBytesBilled(@Nullable Long totalBytesBilled);
59+
60+
abstract TableResult.Builder setTotalBytesProcessed(@Nullable Long totalBytesProcessed);
61+
62+
abstract TableResult.Builder setTotalSlotMs(@Nullable Long totalSlotMs);
63+
64+
abstract TableResult.Builder setNumDmlAffectedRows(@Nullable Long numDmlAffectedRows);
65+
66+
abstract TableResult.Builder setSessionInfo(@Nullable SessionInfo sessionInfo);
67+
5468
/** Creates a @code TableResult} object. */
5569
public abstract TableResult build();
5670
}
@@ -62,8 +76,7 @@ public static Builder newBuilder() {
6276
}
6377

6478
/** Returns the schema of the results. Null if the schema is not supplied. */
65-
@Nullable
66-
public abstract Schema getSchema();
79+
public abstract @Nullable Schema getSchema();
6780

6881
/**
6982
* 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() {
7487

7588
public abstract Page<FieldValueList> getPageNoSchema();
7689

77-
@Nullable
78-
public abstract JobId getJobId();
90+
public abstract @Nullable JobId getJobId();
7991

80-
@Nullable
81-
public abstract String getQueryId();
92+
public abstract @Nullable String getQueryId();
8293

83-
@Nullable
84-
public abstract JobCreationReason getJobCreationReason();
94+
public abstract @Nullable JobCreationReason getJobCreationReason();
8595

8696
/** Returns the number of rows in the current page of results. */
87-
@Nullable
88-
public abstract Long getRowsInPage();
97+
public abstract @Nullable Long getRowsInPage();
98+
99+
/**
100+
* Returns the statement type of the query (e.g. SELECT, INSERT, UPDATE, DDL, SCRIPT).
101+
*
102+
* @return statement type, or {@code null} if not populated by the service
103+
*/
104+
public abstract @Nullable StatementType getStatementType();
105+
106+
/**
107+
* Returns the total number of bytes billed for the query.
108+
*
109+
* @return total bytes billed, or {@code null} if not populated by the service
110+
*/
111+
public abstract @Nullable Long getTotalBytesBilled();
112+
113+
/**
114+
* Returns the total number of bytes processed by the query.
115+
*
116+
* @return total bytes processed, or {@code null} if not populated by the service
117+
*/
118+
public abstract @Nullable Long getTotalBytesProcessed();
119+
120+
/**
121+
* Returns the total slot milliseconds consumed by the query.
122+
*
123+
* @return total slot milliseconds, or {@code null} if not populated by the service
124+
*/
125+
public abstract @Nullable Long getTotalSlotMs();
126+
127+
/**
128+
* Returns the number of rows affected by a DML statement (INSERT, UPDATE, DELETE, MERGE).
129+
*
130+
* @return number of affected rows for DML queries, or {@code null} if not populated by the
131+
* service
132+
*/
133+
public abstract @Nullable Long getNumDmlAffectedRows();
134+
135+
/**
136+
* Returns information about the BigQuery session if this query was executed within or created a
137+
* session.
138+
*
139+
* @return session information, or {@code null} if not populated by the service
140+
*/
141+
public abstract @Nullable SessionInfo getSessionInfo();
89142

90143
@Override
91144
public boolean hasNextPage() {
@@ -109,6 +162,12 @@ public TableResult getNextPage() {
109162
.setQueryId(getQueryId())
110163
.setJobCreationReason(getJobCreationReason())
111164
.setRowsInPage(nextRows)
165+
.setStatementType(getStatementType())
166+
.setTotalBytesBilled(getTotalBytesBilled())
167+
.setTotalBytesProcessed(getTotalBytesProcessed())
168+
.setTotalSlotMs(getTotalSlotMs())
169+
.setNumDmlAffectedRows(getNumDmlAffectedRows())
170+
.setSessionInfo(getSessionInfo())
112171
.build();
113172
}
114173
return null;
@@ -147,13 +206,29 @@ public String toString() {
147206
.add("cursor", getNextPageToken())
148207
.add("queryId", getQueryId())
149208
.add("rowsInPage", getRowsInPage())
209+
.add("statementType", getStatementType())
210+
.add("totalBytesBilled", getTotalBytesBilled())
211+
.add("totalBytesProcessed", getTotalBytesProcessed())
212+
.add("totalSlotMs", getTotalSlotMs())
213+
.add("numDmlAffectedRows", getNumDmlAffectedRows())
214+
.add("sessionInfo", getSessionInfo())
150215
.toString();
151216
}
152217

153218
@Override
154219
public final int hashCode() {
155220
return Objects.hash(
156-
getPageNoSchema(), getSchema(), getTotalRows(), getQueryId(), getRowsInPage());
221+
getPageNoSchema(),
222+
getSchema(),
223+
getTotalRows(),
224+
getQueryId(),
225+
getRowsInPage(),
226+
getStatementType(),
227+
getTotalBytesBilled(),
228+
getTotalBytesProcessed(),
229+
getTotalSlotMs(),
230+
getNumDmlAffectedRows(),
231+
getSessionInfo());
157232
}
158233

159234
@Override
@@ -170,6 +245,12 @@ public final boolean equals(Object obj) {
170245
&& Objects.equals(getSchema(), response.getSchema())
171246
&& getTotalRows() == response.getTotalRows()
172247
&& Objects.equals(getQueryId(), response.getQueryId())
173-
&& Objects.equals(getRowsInPage(), response.getRowsInPage());
248+
&& Objects.equals(getRowsInPage(), response.getRowsInPage())
249+
&& Objects.equals(getStatementType(), response.getStatementType())
250+
&& Objects.equals(getTotalBytesBilled(), response.getTotalBytesBilled())
251+
&& Objects.equals(getTotalBytesProcessed(), response.getTotalBytesProcessed())
252+
&& Objects.equals(getTotalSlotMs(), response.getTotalSlotMs())
253+
&& Objects.equals(getNumDmlAffectedRows(), response.getNumDmlAffectedRows())
254+
&& Objects.equals(getSessionInfo(), response.getSessionInfo());
174255
}
175256
}

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@
5252
import com.google.api.services.bigquery.model.ProjectList;
5353
import com.google.api.services.bigquery.model.ProjectReference;
5454
import com.google.api.services.bigquery.model.QueryRequest;
55+
import com.google.api.services.bigquery.model.SessionInfo;
5556
import com.google.api.services.bigquery.model.TableCell;
5657
import com.google.api.services.bigquery.model.TableDataInsertAllRequest;
5758
import com.google.api.services.bigquery.model.TableDataInsertAllResponse;
@@ -65,6 +66,7 @@
6566
import com.google.cloud.bigquery.BigQuery.JobOption;
6667
import com.google.cloud.bigquery.BigQuery.QueryResultsOption;
6768
import com.google.cloud.bigquery.InsertAllRequest.RowToInsert;
69+
import com.google.cloud.bigquery.JobStatistics.QueryStatistics.StatementType;
6870
import com.google.cloud.bigquery.spi.BigQueryRpcFactory;
6971
import com.google.cloud.bigquery.spi.v2.BigQueryRpc;
7072
import com.google.cloud.bigquery.spi.v2.HttpBigQueryRpc;
@@ -169,6 +171,8 @@ public class BigQueryImplTest {
169171
.setField("timestampField");
170172
private static final TimePartitioning TIME_PARTITIONING_NULL_TYPE =
171173
TimePartitioning.fromPb(PB_TIMEPARTITIONING);
174+
private static final String SESSION_ID = "test-session-id";
175+
private static final SessionInfo PB_SESSION_INFO = new SessionInfo().setSessionId(SESSION_ID);
172176
private static final ImmutableMap<String, String> LABELS = ImmutableMap.of("key", "value");
173177
private static final StandardTableDefinition TABLE_DEFINITION_WITH_PARTITIONING =
174178
StandardTableDefinition.newBuilder()
@@ -2874,7 +2878,12 @@ void testQueryWithTimeoutSetsTimeout() throws InterruptedException, IOException
28742878
.setPageToken(null)
28752879
.setRows(ImmutableList.of(TABLE_ROW))
28762880
.setSchema(TABLE_SCHEMA.toPb())
2881+
.setStatementType("SELECT")
2882+
.setTotalBytesBilled(100L)
28772883
.setTotalBytesProcessed(42L)
2884+
.setTotalSlotMs(50L)
2885+
.setNumDmlAffectedRows(0L)
2886+
.setSessionInfo(PB_SESSION_INFO)
28782887
.setTotalRows(BigInteger.valueOf(1L));
28792888

28802889
when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture()))
@@ -2883,6 +2892,14 @@ void testQueryWithTimeoutSetsTimeout() throws InterruptedException, IOException
28832892
bigquery = options.getService();
28842893
Object result = bigquery.queryWithTimeout(QUERY_JOB_CONFIGURATION_FOR_QUERY, null, 1000L);
28852894
assertTrue(result instanceof TableResult);
2895+
TableResult tableResult = (TableResult) result;
2896+
assertEquals(StatementType.SELECT, tableResult.getStatementType());
2897+
assertEquals((Long) 100L, tableResult.getTotalBytesBilled());
2898+
assertEquals((Long) 42L, tableResult.getTotalBytesProcessed());
2899+
assertEquals((Long) 50L, tableResult.getTotalSlotMs());
2900+
assertEquals((Long) 0L, tableResult.getNumDmlAffectedRows());
2901+
assertNotNull(tableResult.getSessionInfo());
2902+
assertEquals(SESSION_ID, tableResult.getSessionInfo().getSessionId());
28862903
QueryRequest requestPb = requestPbCapture.getValue();
28872904
assertEquals((Long) 1000L, requestPb.getTimeoutMs());
28882905
}

0 commit comments

Comments
 (0)