From 6003cedc009ea6e4da4e53ee359df76f5e5d8332 Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Thu, 20 Aug 2026 14:40:49 +0000 Subject: [PATCH 1/2] feat(bigquery): expose StatementType and query execution stats on TableResult --- .../google/cloud/bigquery/BigQueryImpl.java | 20 ++++ .../java/com/google/cloud/bigquery/Job.java | 27 ++++- .../google/cloud/bigquery/TableResult.java | 65 ++++++++++- .../cloud/bigquery/BigQueryImplTest.java | 11 ++ .../cloud/bigquery/TableResultTest.java | 106 ++++++++++++++++++ java-bigquery/pom.xml | 2 +- 6 files changed, 227 insertions(+), 4 deletions(-) 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..40d78bb60c6d 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,7 @@ 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.spi.v2.BigQueryRpc; import com.google.cloud.bigquery.spi.v2.HttpBigQueryRpc; import com.google.common.annotations.VisibleForTesting; @@ -2095,6 +2096,15 @@ 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(); + if (results.getPageToken() != null) { JobId jobId = JobId.fromPb(results.getJobReference()); String cursor = results.getPageToken(); @@ -2114,6 +2124,11 @@ 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) .build(); } // only 1 page of result @@ -2134,6 +2149,11 @@ 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) .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..614a436a5744 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,8 @@ 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.common.collect.ImmutableList; import io.opentelemetry.api.common.Attributes; import io.opentelemetry.api.trace.Span; @@ -414,6 +416,16 @@ 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; + // 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 +437,11 @@ public TableResult getQueryResults(QueryResultsOption... options) .setTotalRows(0L) .setPageNoSchema(new PageImpl(null, "", null)) .setRowsInPage(0L) + .setStatementType(statementType) + .setTotalBytesBilled(totalBytesBilled) + .setTotalBytesProcessed(totalBytesProcessed) + .setTotalSlotMs(totalSlotMs) + .setNumDmlAffectedRows(numDmlAffectedRows) .build(); return emptyTableResult; } @@ -436,7 +453,15 @@ 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) + .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..df62176af99f 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,6 +18,7 @@ import com.google.api.gax.paging.Page; import com.google.auto.value.AutoValue; +import com.google.cloud.bigquery.JobStatistics.QueryStatistics.StatementType; import com.google.common.base.Function; import com.google.common.base.MoreObjects; import com.google.common.collect.Iterables; @@ -51,6 +52,16 @@ public abstract static class Builder { abstract TableResult.Builder setRowsInPage(Long rowsInPage); + public abstract TableResult.Builder setStatementType(StatementType statementType); + + public abstract TableResult.Builder setTotalBytesBilled(Long totalBytesBilled); + + public abstract TableResult.Builder setTotalBytesProcessed(Long totalBytesProcessed); + + public abstract TableResult.Builder setTotalSlotMs(Long totalSlotMs); + + public abstract TableResult.Builder setNumDmlAffectedRows(Long numDmlAffectedRows); + /** Creates a @code TableResult} object. */ public abstract TableResult build(); } @@ -87,6 +98,32 @@ public static Builder newBuilder() { @Nullable public abstract Long getRowsInPage(); + /** + * Returns the statement type of the query (e.g. SELECT, INSERT, UPDATE, DDL, SCRIPT), if + * available. + */ + @Nullable + public abstract StatementType getStatementType(); + + /** Returns the total number of bytes billed for the query, if available. */ + @Nullable + public abstract Long getTotalBytesBilled(); + + /** Returns the total number of bytes processed by the query, if available. */ + @Nullable + public abstract Long getTotalBytesProcessed(); + + /** Returns the total slot milliseconds for the query, if available. */ + @Nullable + public abstract Long getTotalSlotMs(); + + /** + * Returns the number of rows affected by a DML statement (INSERT, UPDATE, DELETE, MERGE), if + * available. + */ + @Nullable + public abstract Long getNumDmlAffectedRows(); + @Override public boolean hasNextPage() { return getPageNoSchema().hasNextPage(); @@ -109,6 +146,11 @@ public TableResult getNextPage() { .setQueryId(getQueryId()) .setJobCreationReason(getJobCreationReason()) .setRowsInPage(nextRows) + .setStatementType(getStatementType()) + .setTotalBytesBilled(getTotalBytesBilled()) + .setTotalBytesProcessed(getTotalBytesProcessed()) + .setTotalSlotMs(getTotalSlotMs()) + .setNumDmlAffectedRows(getNumDmlAffectedRows()) .build(); } return null; @@ -147,13 +189,27 @@ 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()) .toString(); } @Override public final int hashCode() { return Objects.hash( - getPageNoSchema(), getSchema(), getTotalRows(), getQueryId(), getRowsInPage()); + getPageNoSchema(), + getSchema(), + getTotalRows(), + getQueryId(), + getRowsInPage(), + getStatementType(), + getTotalBytesBilled(), + getTotalBytesProcessed(), + getTotalSlotMs(), + getNumDmlAffectedRows()); } @Override @@ -170,6 +226,11 @@ 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()); } } 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..434d317c420f 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 @@ -65,6 +65,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; @@ -2874,7 +2875,11 @@ 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) .setTotalRows(BigInteger.valueOf(1L)); when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture())) @@ -2883,6 +2888,12 @@ 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()); 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..fe9d33d519ec 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 @@ -115,4 +115,110 @@ 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(JobStatistics.QueryStatistics.StatementType.SELECT) + .setTotalBytesBilled(1024L) + .setTotalBytesProcessed(2048L) + .setTotalSlotMs(500L) + .setNumDmlAffectedRows(0L) + .build(); + + assertThat(result.getStatementType()) + .isEqualTo(JobStatistics.QueryStatistics.StatementType.SELECT); + assertThat(result.getTotalBytesBilled()).isEqualTo(1024L); + assertThat(result.getTotalBytesProcessed()).isEqualTo(2048L); + assertThat(result.getTotalSlotMs()).isEqualTo(500L); + assertThat(result.getNumDmlAffectedRows()).isEqualTo(0L); + + TableResult next = result.getNextPage(); + assertThat(next.getStatementType()) + .isEqualTo(JobStatistics.QueryStatistics.StatementType.SELECT); + assertThat(next.getTotalBytesBilled()).isEqualTo(1024L); + assertThat(next.getTotalBytesProcessed()).isEqualTo(2048L); + assertThat(next.getTotalSlotMs()).isEqualTo(500L); + assertThat(next.getNumDmlAffectedRows()).isEqualTo(0L); + } + + @Test + void testToBuilder() { + TableResult result = + TableResult.newBuilder() + .setSchema(SCHEMA) + .setTotalRows(3L) + .setPageNoSchema(INNER_PAGE_0) + .setRowsInPage(2L) + .setStatementType(JobStatistics.QueryStatistics.StatementType.INSERT) + .setTotalBytesBilled(500L) + .setTotalBytesProcessed(1000L) + .setTotalSlotMs(250L) + .setNumDmlAffectedRows(5L) + .build(); + + TableResult modified = + result.toBuilder() + .setStatementType(JobStatistics.QueryStatistics.StatementType.UPDATE) + .setNumDmlAffectedRows(10L) + .build(); + + assertThat(modified.getStatementType()) + .isEqualTo(JobStatistics.QueryStatistics.StatementType.UPDATE); + assertThat(modified.getNumDmlAffectedRows()).isEqualTo(10L); + assertThat(modified.getTotalBytesBilled()).isEqualTo(500L); + } + + @Test + void testEqualsAndHashCode() { + TableResult result1 = + TableResult.newBuilder() + .setSchema(SCHEMA) + .setTotalRows(3L) + .setPageNoSchema(INNER_PAGE_0) + .setRowsInPage(2L) + .setStatementType(JobStatistics.QueryStatistics.StatementType.SELECT) + .setTotalBytesBilled(100L) + .setTotalBytesProcessed(200L) + .setTotalSlotMs(50L) + .setNumDmlAffectedRows(0L) + .build(); + + TableResult result2 = + TableResult.newBuilder() + .setSchema(SCHEMA) + .setTotalRows(3L) + .setPageNoSchema(INNER_PAGE_0) + .setRowsInPage(2L) + .setStatementType(JobStatistics.QueryStatistics.StatementType.SELECT) + .setTotalBytesBilled(100L) + .setTotalBytesProcessed(200L) + .setTotalSlotMs(50L) + .setNumDmlAffectedRows(0L) + .build(); + + TableResult result3 = + TableResult.newBuilder() + .setSchema(SCHEMA) + .setTotalRows(3L) + .setPageNoSchema(INNER_PAGE_0) + .setRowsInPage(2L) + .setStatementType(JobStatistics.QueryStatistics.StatementType.DELETE) + .setTotalBytesBilled(100L) + .setTotalBytesProcessed(200L) + .setTotalSlotMs(50L) + .setNumDmlAffectedRows(1L) + .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"); + } } 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 From 94fbc074d41f427f78b6e3a2fe0e88576043eb93 Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Thu, 20 Aug 2026 22:16:58 +0000 Subject: [PATCH 2/2] refactor(bigquery): consolidate query execution metrics into JobStatistics on TableResult --- .../google/cloud/bigquery/BigQueryImpl.java | 22 +--- .../java/com/google/cloud/bigquery/Job.java | 27 +---- .../google/cloud/bigquery/JobStatistics.java | 28 +++++ .../google/cloud/bigquery/TableResult.java | 61 ++--------- .../cloud/bigquery/BigQueryImplTest.java | 14 ++- .../cloud/bigquery/TableResultTest.java | 100 ++++++++++-------- 6 files changed, 105 insertions(+), 147 deletions(-) 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 40d78bb60c6d..96bb33e43bc3 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,7 +41,6 @@ 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.spi.v2.BigQueryRpc; import com.google.cloud.bigquery.spi.v2.HttpBigQueryRpc; import com.google.common.annotations.VisibleForTesting; @@ -2096,14 +2095,7 @@ 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(); + JobStatistics jobStatistics = JobStatistics.QueryStatistics.fromPb(results, schema); if (results.getPageToken() != null) { JobId jobId = JobId.fromPb(results.getJobReference()); @@ -2124,11 +2116,7 @@ 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) + .setJobStatistics(jobStatistics) .build(); } // only 1 page of result @@ -2149,11 +2137,7 @@ 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) + .setJobStatistics(jobStatistics) .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 614a436a5744..fa411ca88ad7 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,8 +28,6 @@ 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.common.collect.ImmutableList; import io.opentelemetry.api.common.Attributes; import io.opentelemetry.api.trace.Span; @@ -416,15 +414,7 @@ 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; + JobStatistics stats = job.getStatistics(); // 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. @@ -437,11 +427,7 @@ public TableResult getQueryResults(QueryResultsOption... options) .setTotalRows(0L) .setPageNoSchema(new PageImpl(null, "", null)) .setRowsInPage(0L) - .setStatementType(statementType) - .setTotalBytesBilled(totalBytesBilled) - .setTotalBytesProcessed(totalBytesProcessed) - .setTotalSlotMs(totalSlotMs) - .setNumDmlAffectedRows(numDmlAffectedRows) + .setJobStatistics(stats) .build(); return emptyTableResult; } @@ -454,14 +440,7 @@ public TableResult getQueryResults(QueryResultsOption... options) bigquery.listTableData( table, response.getSchema(), listOptions.toArray(new TableDataListOption[0])); TableResult tableResultWithJobId = - tableResult.toBuilder() - .setJobId(job.getJobId()) - .setStatementType(statementType) - .setTotalBytesBilled(totalBytesBilled) - .setTotalBytesProcessed(totalBytesProcessed) - .setTotalSlotMs(totalSlotMs) - .setNumDmlAffectedRows(numDmlAffectedRows) - .build(); + tableResult.toBuilder().setJobId(job.getJobId()).setJobStatistics(stats).build(); return tableResultWithJobId; } finally { if (getQueryResults != null) { diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java index c78378ab53f7..77153be12ae5 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java @@ -24,6 +24,7 @@ import com.google.api.services.bigquery.model.JobStatistics4; import com.google.api.services.bigquery.model.JobStatistics5; import com.google.api.services.bigquery.model.QueryParameter; +import com.google.api.services.bigquery.model.QueryResponse; import com.google.auto.value.AutoValue; import com.google.cloud.StringEnumType; import com.google.cloud.StringEnumValue; @@ -1073,6 +1074,33 @@ static QueryStatistics fromPb( com.google.api.services.bigquery.model.JobStatistics statisticPb) { return new Builder(statisticPb).build(); } + + /** Creates a {@code QueryStatistics} object from a {@link QueryResponse} and schema. */ + static QueryStatistics fromPb(QueryResponse queryResponse, Schema schema) { + if (queryResponse == null) { + return null; + } + Builder builder = newBuilder(); + if (queryResponse.getStatementType() != null) { + builder.setStatementType(queryResponse.getStatementType()); + } + if (queryResponse.getTotalBytesBilled() != null) { + builder.setTotalBytesBilled(queryResponse.getTotalBytesBilled()); + } + if (queryResponse.getTotalBytesProcessed() != null) { + builder.setTotalBytesProcessed(queryResponse.getTotalBytesProcessed()); + } + if (queryResponse.getTotalSlotMs() != null) { + builder.setTotalSlotMs(queryResponse.getTotalSlotMs()); + } + if (queryResponse.getNumDmlAffectedRows() != null) { + builder.setNumDmlAffectedRows(queryResponse.getNumDmlAffectedRows()); + } + if (schema != null) { + builder.setSchema(schema); + } + return builder.build(); + } } /** A Google BigQuery Script statistics. */ 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 df62176af99f..a142cc4764d6 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,7 +18,6 @@ import com.google.api.gax.paging.Page; import com.google.auto.value.AutoValue; -import com.google.cloud.bigquery.JobStatistics.QueryStatistics.StatementType; import com.google.common.base.Function; import com.google.common.base.MoreObjects; import com.google.common.collect.Iterables; @@ -52,15 +51,7 @@ public abstract static class Builder { abstract TableResult.Builder setRowsInPage(Long rowsInPage); - public abstract TableResult.Builder setStatementType(StatementType statementType); - - public abstract TableResult.Builder setTotalBytesBilled(Long totalBytesBilled); - - public abstract TableResult.Builder setTotalBytesProcessed(Long totalBytesProcessed); - - public abstract TableResult.Builder setTotalSlotMs(Long totalSlotMs); - - public abstract TableResult.Builder setNumDmlAffectedRows(Long numDmlAffectedRows); + public abstract TableResult.Builder setJobStatistics(JobStatistics jobStatistics); /** Creates a @code TableResult} object. */ public abstract TableResult build(); @@ -98,31 +89,9 @@ public static Builder newBuilder() { @Nullable public abstract Long getRowsInPage(); - /** - * Returns the statement type of the query (e.g. SELECT, INSERT, UPDATE, DDL, SCRIPT), if - * available. - */ - @Nullable - public abstract StatementType getStatementType(); - - /** Returns the total number of bytes billed for the query, if available. */ - @Nullable - public abstract Long getTotalBytesBilled(); - - /** Returns the total number of bytes processed by the query, if available. */ - @Nullable - public abstract Long getTotalBytesProcessed(); - - /** Returns the total slot milliseconds for the query, if available. */ - @Nullable - public abstract Long getTotalSlotMs(); - - /** - * Returns the number of rows affected by a DML statement (INSERT, UPDATE, DELETE, MERGE), if - * available. - */ + /** Returns the statistics of the job/query that produced these results, if available. */ @Nullable - public abstract Long getNumDmlAffectedRows(); + public abstract JobStatistics getJobStatistics(); @Override public boolean hasNextPage() { @@ -146,11 +115,7 @@ public TableResult getNextPage() { .setQueryId(getQueryId()) .setJobCreationReason(getJobCreationReason()) .setRowsInPage(nextRows) - .setStatementType(getStatementType()) - .setTotalBytesBilled(getTotalBytesBilled()) - .setTotalBytesProcessed(getTotalBytesProcessed()) - .setTotalSlotMs(getTotalSlotMs()) - .setNumDmlAffectedRows(getNumDmlAffectedRows()) + .setJobStatistics(getJobStatistics()) .build(); } return null; @@ -189,11 +154,7 @@ 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("jobStatistics", getJobStatistics()) .toString(); } @@ -205,11 +166,7 @@ public final int hashCode() { getTotalRows(), getQueryId(), getRowsInPage(), - getStatementType(), - getTotalBytesBilled(), - getTotalBytesProcessed(), - getTotalSlotMs(), - getNumDmlAffectedRows()); + getJobStatistics()); } @Override @@ -227,10 +184,6 @@ public final boolean equals(Object obj) { && getTotalRows() == response.getTotalRows() && Objects.equals(getQueryId(), response.getQueryId()) && 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(getJobStatistics(), response.getJobStatistics()); } } 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 434d317c420f..ef97a4721a5e 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 @@ -65,6 +65,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; import com.google.cloud.bigquery.JobStatistics.QueryStatistics.StatementType; import com.google.cloud.bigquery.spi.BigQueryRpcFactory; import com.google.cloud.bigquery.spi.v2.BigQueryRpc; @@ -2889,11 +2890,14 @@ void testQueryWithTimeoutSetsTimeout() throws InterruptedException, IOException 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.getJobStatistics()); + assertTrue(tableResult.getJobStatistics() instanceof QueryStatistics); + QueryStatistics stats = (QueryStatistics) tableResult.getJobStatistics(); + assertEquals(StatementType.SELECT, stats.getStatementType()); + assertEquals((Long) 100L, stats.getTotalBytesBilled()); + assertEquals((Long) 42L, stats.getTotalBytesProcessed()); + assertEquals((Long) 50L, stats.getTotalSlotMs()); + assertEquals((Long) 0L, stats.getNumDmlAffectedRows()); 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 fe9d33d519ec..a615fa470c1f 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 @@ -117,76 +117,95 @@ void testSchema() { } @Test - void testStatementTypeAndExecutionStats() { + void testJobStatistics() { + JobStatistics.QueryStatistics stats = + JobStatistics.QueryStatistics.newBuilder() + .setStatementType(JobStatistics.QueryStatistics.StatementType.SELECT) + .setTotalBytesBilled(1024L) + .setTotalBytesProcessed(2048L) + .setTotalSlotMs(500L) + .setNumDmlAffectedRows(0L) + .build(); + TableResult result = TableResult.newBuilder() .setSchema(SCHEMA) .setTotalRows(3L) .setPageNoSchema(INNER_PAGE_0) .setRowsInPage(2L) - .setStatementType(JobStatistics.QueryStatistics.StatementType.SELECT) - .setTotalBytesBilled(1024L) - .setTotalBytesProcessed(2048L) - .setTotalSlotMs(500L) - .setNumDmlAffectedRows(0L) + .setJobStatistics(stats) .build(); - assertThat(result.getStatementType()) + assertThat(result.getJobStatistics()).isEqualTo(stats); + assertThat(result.getJobStatistics()).isInstanceOf(JobStatistics.QueryStatistics.class); + JobStatistics.QueryStatistics queryStats = + (JobStatistics.QueryStatistics) result.getJobStatistics(); + assertThat(queryStats.getStatementType()) .isEqualTo(JobStatistics.QueryStatistics.StatementType.SELECT); - assertThat(result.getTotalBytesBilled()).isEqualTo(1024L); - assertThat(result.getTotalBytesProcessed()).isEqualTo(2048L); - assertThat(result.getTotalSlotMs()).isEqualTo(500L); - assertThat(result.getNumDmlAffectedRows()).isEqualTo(0L); + assertThat(queryStats.getTotalBytesBilled()).isEqualTo(1024L); + assertThat(queryStats.getTotalBytesProcessed()).isEqualTo(2048L); + assertThat(queryStats.getTotalSlotMs()).isEqualTo(500L); + assertThat(queryStats.getNumDmlAffectedRows()).isEqualTo(0L); TableResult next = result.getNextPage(); - assertThat(next.getStatementType()) - .isEqualTo(JobStatistics.QueryStatistics.StatementType.SELECT); - assertThat(next.getTotalBytesBilled()).isEqualTo(1024L); - assertThat(next.getTotalBytesProcessed()).isEqualTo(2048L); - assertThat(next.getTotalSlotMs()).isEqualTo(500L); - assertThat(next.getNumDmlAffectedRows()).isEqualTo(0L); + assertThat(next.getJobStatistics()).isEqualTo(stats); } @Test void testToBuilder() { + JobStatistics.QueryStatistics stats1 = + JobStatistics.QueryStatistics.newBuilder() + .setStatementType(JobStatistics.QueryStatistics.StatementType.INSERT) + .setNumDmlAffectedRows(5L) + .build(); + TableResult result = TableResult.newBuilder() .setSchema(SCHEMA) .setTotalRows(3L) .setPageNoSchema(INNER_PAGE_0) .setRowsInPage(2L) - .setStatementType(JobStatistics.QueryStatistics.StatementType.INSERT) - .setTotalBytesBilled(500L) - .setTotalBytesProcessed(1000L) - .setTotalSlotMs(250L) - .setNumDmlAffectedRows(5L) + .setJobStatistics(stats1) .build(); - TableResult modified = - result.toBuilder() + JobStatistics.QueryStatistics stats2 = + JobStatistics.QueryStatistics.newBuilder() .setStatementType(JobStatistics.QueryStatistics.StatementType.UPDATE) .setNumDmlAffectedRows(10L) .build(); - assertThat(modified.getStatementType()) - .isEqualTo(JobStatistics.QueryStatistics.StatementType.UPDATE); - assertThat(modified.getNumDmlAffectedRows()).isEqualTo(10L); - assertThat(modified.getTotalBytesBilled()).isEqualTo(500L); + TableResult modified = result.toBuilder().setJobStatistics(stats2).build(); + + assertThat(modified.getJobStatistics()).isEqualTo(stats2); } @Test void testEqualsAndHashCode() { + JobStatistics.QueryStatistics stats1 = + JobStatistics.QueryStatistics.newBuilder() + .setStatementType(JobStatistics.QueryStatistics.StatementType.SELECT) + .setTotalBytesBilled(100L) + .build(); + + JobStatistics.QueryStatistics stats2 = + JobStatistics.QueryStatistics.newBuilder() + .setStatementType(JobStatistics.QueryStatistics.StatementType.SELECT) + .setTotalBytesBilled(100L) + .build(); + + JobStatistics.QueryStatistics stats3 = + JobStatistics.QueryStatistics.newBuilder() + .setStatementType(JobStatistics.QueryStatistics.StatementType.DELETE) + .setTotalBytesBilled(100L) + .build(); + TableResult result1 = TableResult.newBuilder() .setSchema(SCHEMA) .setTotalRows(3L) .setPageNoSchema(INNER_PAGE_0) .setRowsInPage(2L) - .setStatementType(JobStatistics.QueryStatistics.StatementType.SELECT) - .setTotalBytesBilled(100L) - .setTotalBytesProcessed(200L) - .setTotalSlotMs(50L) - .setNumDmlAffectedRows(0L) + .setJobStatistics(stats1) .build(); TableResult result2 = @@ -195,11 +214,7 @@ void testEqualsAndHashCode() { .setTotalRows(3L) .setPageNoSchema(INNER_PAGE_0) .setRowsInPage(2L) - .setStatementType(JobStatistics.QueryStatistics.StatementType.SELECT) - .setTotalBytesBilled(100L) - .setTotalBytesProcessed(200L) - .setTotalSlotMs(50L) - .setNumDmlAffectedRows(0L) + .setJobStatistics(stats2) .build(); TableResult result3 = @@ -208,17 +223,12 @@ void testEqualsAndHashCode() { .setTotalRows(3L) .setPageNoSchema(INNER_PAGE_0) .setRowsInPage(2L) - .setStatementType(JobStatistics.QueryStatistics.StatementType.DELETE) - .setTotalBytesBilled(100L) - .setTotalBytesProcessed(200L) - .setTotalSlotMs(50L) - .setNumDmlAffectedRows(1L) + .setJobStatistics(stats3) .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("jobStatistics"); } }