diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java index 2ad09c33d7cb..da4b11e676dd 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java +++ b/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(); } diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Job.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Job.java index 43832f3410b8..ca68b87db97c 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Job.java +++ b/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(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) { diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableResult.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableResult.java index a791628c5b16..b49a253e0c67 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableResult.java +++ b/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, 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 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()); } } diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java index 9f2320ab3c35..9a398e74a67d 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java +++ b/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 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()); } diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TableResultTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TableResultTest.java index 90ae2692f00c..f697049dc2f9 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TableResultTest.java +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TableResultTest.java @@ -21,6 +21,8 @@ import com.google.api.gax.paging.Page; import com.google.cloud.PageImpl; +import com.google.cloud.bigquery.JobStatistics.QueryStatistics.StatementType; +import com.google.cloud.bigquery.JobStatistics.SessionInfo; import com.google.common.collect.ImmutableList; import org.junit.jupiter.api.Test; @@ -46,6 +48,15 @@ public Page getNextPage() { null, ImmutableList.of(newFieldValueList("2"))); private static final Schema SCHEMA = Schema.of(Field.of("field", LegacySQLTypeName.INTEGER)); + private static final String SESSION_ID = "session_123"; + private static final SessionInfo SESSION_INFO = + SessionInfo.newBuilder().setSessionId(SESSION_ID).build(); + private static final String SESSION_ID_1 = "session_1"; + private static final SessionInfo SESSION_INFO_1 = + SessionInfo.newBuilder().setSessionId(SESSION_ID_1).build(); + private static final String SESSION_ID_2 = "session_2"; + private static final SessionInfo SESSION_INFO_2 = + SessionInfo.newBuilder().setSessionId(SESSION_ID_2).build(); private static FieldValueList newFieldValueList(String s) { return FieldValueList.of(ImmutableList.of(FieldValue.of(PRIMITIVE, s))); @@ -115,4 +126,117 @@ void testSchema() { newFieldValueList("2").withSchema(SCHEMA.getFields())) .inOrder(); } + + @Test + void testStatementTypeAndExecutionStats() { + TableResult result = + TableResult.newBuilder() + .setSchema(SCHEMA) + .setTotalRows(3L) + .setPageNoSchema(INNER_PAGE_0) + .setRowsInPage(2L) + .setStatementType(StatementType.SELECT) + .setTotalBytesBilled(1024L) + .setTotalBytesProcessed(2048L) + .setTotalSlotMs(500L) + .setNumDmlAffectedRows(0L) + .setSessionInfo(SESSION_INFO) + .build(); + + assertThat(result.getStatementType()).isEqualTo(StatementType.SELECT); + assertThat(result.getTotalBytesBilled()).isEqualTo(1024L); + assertThat(result.getTotalBytesProcessed()).isEqualTo(2048L); + assertThat(result.getTotalSlotMs()).isEqualTo(500L); + assertThat(result.getNumDmlAffectedRows()).isEqualTo(0L); + assertThat(result.getSessionInfo()).isEqualTo(SESSION_INFO); + assertThat(result.getSessionInfo().getSessionId()).isEqualTo(SESSION_ID); + + TableResult next = result.getNextPage(); + assertThat(next.getStatementType()).isEqualTo(StatementType.SELECT); + assertThat(next.getTotalBytesBilled()).isEqualTo(1024L); + assertThat(next.getTotalBytesProcessed()).isEqualTo(2048L); + assertThat(next.getTotalSlotMs()).isEqualTo(500L); + assertThat(next.getNumDmlAffectedRows()).isEqualTo(0L); + assertThat(next.getSessionInfo()).isEqualTo(SESSION_INFO); + } + + @Test + void testToBuilder() { + TableResult result = + TableResult.newBuilder() + .setSchema(SCHEMA) + .setTotalRows(3L) + .setPageNoSchema(INNER_PAGE_0) + .setRowsInPage(2L) + .setStatementType(StatementType.INSERT) + .setTotalBytesBilled(500L) + .setTotalBytesProcessed(1000L) + .setTotalSlotMs(250L) + .setNumDmlAffectedRows(5L) + .setSessionInfo(SESSION_INFO) + .build(); + + TableResult modified = + result.toBuilder() + .setStatementType(StatementType.UPDATE) + .setNumDmlAffectedRows(10L) + .build(); + + assertThat(modified.getStatementType()).isEqualTo(StatementType.UPDATE); + assertThat(modified.getNumDmlAffectedRows()).isEqualTo(10L); + assertThat(modified.getTotalBytesBilled()).isEqualTo(500L); + assertThat(modified.getSessionInfo()).isEqualTo(SESSION_INFO); + } + + @Test + void testEqualsAndHashCode() { + TableResult result1 = + TableResult.newBuilder() + .setSchema(SCHEMA) + .setTotalRows(3L) + .setPageNoSchema(INNER_PAGE_0) + .setRowsInPage(2L) + .setStatementType(StatementType.SELECT) + .setTotalBytesBilled(100L) + .setTotalBytesProcessed(200L) + .setTotalSlotMs(50L) + .setNumDmlAffectedRows(0L) + .setSessionInfo(SESSION_INFO_1) + .build(); + + TableResult result2 = + TableResult.newBuilder() + .setSchema(SCHEMA) + .setTotalRows(3L) + .setPageNoSchema(INNER_PAGE_0) + .setRowsInPage(2L) + .setStatementType(StatementType.SELECT) + .setTotalBytesBilled(100L) + .setTotalBytesProcessed(200L) + .setTotalSlotMs(50L) + .setNumDmlAffectedRows(0L) + .setSessionInfo(SESSION_INFO_1) + .build(); + + TableResult result3 = + TableResult.newBuilder() + .setSchema(SCHEMA) + .setTotalRows(3L) + .setPageNoSchema(INNER_PAGE_0) + .setRowsInPage(2L) + .setStatementType(StatementType.DELETE) + .setTotalBytesBilled(100L) + .setTotalBytesProcessed(200L) + .setTotalSlotMs(50L) + .setNumDmlAffectedRows(1L) + .setSessionInfo(SESSION_INFO_2) + .build(); + + assertThat(result1).isEqualTo(result2); + assertThat(result1.hashCode()).isEqualTo(result2.hashCode()); + assertThat(result1).isNotEqualTo(result3); + assertThat(result1.toString()).contains("statementType=SELECT"); + assertThat(result1.toString()).contains("totalBytesBilled=100"); + assertThat(result1.toString()).contains("sessionId=" + SESSION_ID_1); + } } diff --git a/java-bigquery/pom.xml b/java-bigquery/pom.xml index ee1fcd22dfc8..b480390dd2e7 100644 --- a/java-bigquery/pom.xml +++ b/java-bigquery/pom.xml @@ -55,7 +55,7 @@ UTF-8 github google-cloud-bigquery-parent - v2-rev20251012-2.0.0 + v2-rev20260731-2.0.0