From 9d67e05d2a6fab0c9e6017dec33b0d74ad821825 Mon Sep 17 00:00:00 2001 From: Stephanie Wang Date: Thu, 8 Jul 2021 14:32:07 -0400 Subject: [PATCH] feat: add dmlStatistics support (#1431) Provides detailed statistics for DML statements Present only for DML statements INSERT, UPDATE, DELETE or TRUNCATE. Towards b/186432630 --- .../com/google/cloud/bigquery/DmlStats.java | 112 ++++++++++++++++++ .../google/cloud/bigquery/JobStatistics.java | 20 +++- .../google/cloud/bigquery/DmlStatsTest.java | 55 +++++++++ .../cloud/bigquery/JobStatisticsTest.java | 11 ++ .../cloud/bigquery/it/ITBigQueryTest.java | 20 ++++ 5 files changed, 217 insertions(+), 1 deletion(-) create mode 100644 google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/DmlStats.java create mode 100644 google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/DmlStatsTest.java diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/DmlStats.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/DmlStats.java new file mode 100644 index 000000000..0ce97dc0c --- /dev/null +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/DmlStats.java @@ -0,0 +1,112 @@ +/* + * Copyright 2021 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery; + +import com.google.api.services.bigquery.model.DmlStatistics; +import com.google.auto.value.AutoValue; +import java.io.Serializable; +import javax.annotation.Nullable; + +/** Represents DML statistics information. */ +@AutoValue +public abstract class DmlStats implements Serializable { + + @AutoValue.Builder + public abstract static class Builder { + /** + * Number of deleted Rows. populated by DML DELETE, MERGE and TRUNCATE statements. + * + * @param deletedRowCount deletedRowCount or {@code null} for none + */ + public abstract Builder setDeletedRowCount(Long deletedRowCount); + + /** + * Number of inserted Rows. Populated by DML INSERT and MERGE statements. + * + * @param insertedRowCount insertedRowCount or {@code null} for none + */ + public abstract Builder setInsertedRowCount(Long insertedRowCount); + + /** + * Number of updated Rows. Populated by DML UPDATE and MERGE statements. + * + * @param updatedRowCount updatedRowCount or {@code null} for none + */ + public abstract Builder setUpdatedRowCount(Long updatedRowCount); + + /** Creates a {@code DmlStats} object. */ + public abstract DmlStats build(); + } + + /** + * Returns number of deleted Rows. populated by DML DELETE, MERGE and TRUNCATE statements. + * + * @return value or {@code null} for none + */ + @Nullable + public abstract Long getDeletedRowCount(); + + /** + * Returns number of inserted Rows. Populated by DML INSERT and MERGE statements. + * + * @return value or {@code null} for none + */ + @Nullable + public abstract Long getInsertedRowCount(); + + /** + * Returns number of updated Rows. Populated by DML UPDATE and MERGE statements. + * + * @return value or {@code null} for none + */ + @Nullable + public abstract Long getUpdatedRowCount(); + + public abstract Builder toBuilder(); + + public static Builder newBuilder() { + return new AutoValue_DmlStats.Builder(); + } + + DmlStatistics toPb() { + DmlStatistics dmlStatisticsPb = new DmlStatistics(); + if (getDeletedRowCount() != null) { + dmlStatisticsPb.setDeletedRowCount(getDeletedRowCount()); + } + if (getInsertedRowCount() != null) { + dmlStatisticsPb.setInsertedRowCount(getInsertedRowCount()); + } + if (getUpdatedRowCount() != null) { + dmlStatisticsPb.setUpdatedRowCount(getUpdatedRowCount()); + } + return dmlStatisticsPb; + } + + static DmlStats fromPb(DmlStatistics dmlStatisticsPb) { + Builder builder = newBuilder(); + if (dmlStatisticsPb.getDeletedRowCount() != null) { + builder.setDeletedRowCount(dmlStatisticsPb.getDeletedRowCount()); + } + if (dmlStatisticsPb.getInsertedRowCount() != null) { + builder.setInsertedRowCount(dmlStatisticsPb.getInsertedRowCount()); + } + if (dmlStatisticsPb.getUpdatedRowCount() != null) { + builder.setUpdatedRowCount(dmlStatisticsPb.getUpdatedRowCount()); + } + return builder.build(); + } +} diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java index d971a6fb2..ea5555b9b 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java @@ -326,6 +326,7 @@ public static class QueryStatistics extends JobStatistics { private final RoutineId ddlTargetRoutine; private final Long estimatedBytesProcessed; private final Long numDmlAffectedRows; + private final DmlStats dmlStats; private final List referencedTables; private final StatementType statementType; private final Long totalBytesBilled; @@ -406,6 +407,7 @@ static final class Builder extends JobStatistics.Builder referencedTables; private StatementType statementType; private Long totalBytesBilled; @@ -458,6 +460,9 @@ private Builder(com.google.api.services.bigquery.model.JobStatistics statisticsP if (statisticsPb.getQuery().getSchema() != null) { this.schema = Schema.fromPb(statisticsPb.getQuery().getSchema()); } + if (statisticsPb.getQuery().getDmlStats() != null) { + this.dmlStats = DmlStats.fromPb(statisticsPb.getQuery().getDmlStats()); + } } } @@ -496,6 +501,11 @@ Builder setNumDmlAffectedRows(Long numDmlAffectedRows) { return self(); } + Builder setDmlStats(DmlStats dmlStats) { + this.dmlStats = dmlStats; + return self(); + } + Builder setReferenceTables(List referencedTables) { this.referencedTables = referencedTables; return self(); @@ -561,6 +571,7 @@ private QueryStatistics(Builder builder) { this.ddlTargetRoutine = builder.ddlTargetRoutine; this.estimatedBytesProcessed = builder.estimatedBytesProcessed; this.numDmlAffectedRows = builder.numDmlAffectedRows; + this.dmlStats = builder.dmlStats; this.referencedTables = builder.referencedTables; this.statementType = builder.statementType; this.totalBytesBilled = builder.totalBytesBilled; @@ -614,6 +625,11 @@ public Long getNumDmlAffectedRows() { return numDmlAffectedRows; } + /** Detailed statistics for DML statements. */ + public DmlStats getDmlStats() { + return dmlStats; + } + /** * Referenced tables for the job. Queries that reference more than 50 tables will not have a * complete list. @@ -729,7 +745,9 @@ com.google.api.services.bigquery.model.JobStatistics toPb() { if (ddlTargetRoutine != null) { queryStatisticsPb.setDdlTargetRoutine(ddlTargetRoutine.toPb()); } - + if (dmlStats != null) { + queryStatisticsPb.setDmlStats(dmlStats.toPb()); + } if (referencedTables != null) { queryStatisticsPb.setReferencedTables( Lists.transform(referencedTables, TableId.TO_PB_FUNCTION)); diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/DmlStatsTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/DmlStatsTest.java new file mode 100644 index 000000000..48950831a --- /dev/null +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/DmlStatsTest.java @@ -0,0 +1,55 @@ +/* + * Copyright 2021 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery; + +import static org.junit.Assert.assertEquals; + +import org.junit.Test; + +public class DmlStatsTest { + + private static final Long DELETED_ROW_COUNT = 10L; + private static final Long INSERTED_ROW_COUNT = 20L; + private static final Long UPDATED_ROW_COUNT = 30L; + private static final DmlStats DML_STATS = + DmlStats.newBuilder() + .setDeletedRowCount(DELETED_ROW_COUNT) + .setInsertedRowCount(INSERTED_ROW_COUNT) + .setUpdatedRowCount(UPDATED_ROW_COUNT) + .build(); + + @Test + public void testBuilder() { + assertEquals(DELETED_ROW_COUNT, DML_STATS.getDeletedRowCount()); + assertEquals(UPDATED_ROW_COUNT, DML_STATS.getUpdatedRowCount()); + assertEquals(INSERTED_ROW_COUNT, DML_STATS.getInsertedRowCount()); + } + + @Test + public void testToPbAndFromPb() { + compareDmlStats(DML_STATS, DmlStats.fromPb(DML_STATS.toPb())); + } + + private void compareDmlStats(DmlStats expected, DmlStats actual) { + assertEquals(expected, actual); + assertEquals(expected.hashCode(), actual.hashCode()); + assertEquals(expected.toString(), actual.toString()); + assertEquals(expected.getDeletedRowCount(), actual.getDeletedRowCount()); + assertEquals(expected.getInsertedRowCount(), actual.getInsertedRowCount()); + assertEquals(expected.getUpdatedRowCount(), actual.getUpdatedRowCount()); + } +} diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java index 96bfa3f08..a28eae26b 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java @@ -42,6 +42,15 @@ public class JobStatisticsTest { private static final RoutineId DDL_TARGET_ROUTINE = RoutineId.of("alpha", "beta", "gamma"); private static final Long ESTIMATE_BYTES_PROCESSED = 101L; private static final Long NUM_DML_AFFECTED_ROWS = 88L; + private static final Long DELETED_ROW_COUNT = 10L; + private static final Long INSERTED_ROW_COUNT = 20L; + private static final Long UPDATED_ROW_COUNT = 30L; + private static final DmlStats DML_STATS = + DmlStats.newBuilder() + .setDeletedRowCount(DELETED_ROW_COUNT) + .setInsertedRowCount(INSERTED_ROW_COUNT) + .setUpdatedRowCount(UPDATED_ROW_COUNT) + .build(); private static final QueryStatistics.StatementType STATEMENT_TYPE = QueryStatistics.StatementType.SELECT; private static final Long TOTAL_BYTES_BILLED = 24L; @@ -147,6 +156,7 @@ public class JobStatisticsTest { .setDDLTargetRoutine(DDL_TARGET_ROUTINE) .setEstimatedBytesProcessed(ESTIMATE_BYTES_PROCESSED) .setNumDmlAffectedRows(NUM_DML_AFFECTED_ROWS) + .setDmlStats(DML_STATS) .setReferenceTables(REFERENCED_TABLES) .setStatementType(STATEMENT_TYPE) .setTotalBytesBilled(TOTAL_BYTES_BILLED) @@ -232,6 +242,7 @@ public void testBuilder() { assertEquals(DDL_TARGET_ROUTINE, QUERY_STATISTICS.getDdlTargetRoutine()); assertEquals(ESTIMATE_BYTES_PROCESSED, QUERY_STATISTICS.getEstimatedBytesProcessed()); assertEquals(NUM_DML_AFFECTED_ROWS, QUERY_STATISTICS.getNumDmlAffectedRows()); + assertEquals(DML_STATS, QUERY_STATISTICS.getDmlStats()); assertEquals(REFERENCED_TABLES, QUERY_STATISTICS.getReferencedTables()); assertEquals(STATEMENT_TYPE, QUERY_STATISTICS.getStatementType()); assertEquals(TOTAL_BYTES_BILLED, QUERY_STATISTICS.getTotalBytesBilled()); diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java index a0135fb9e..1669343fb 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java @@ -2205,6 +2205,26 @@ public void testFastQueryHTTPException() throws InterruptedException { } } + @Test + public void testDmlStatistics() throws InterruptedException { + String tableName = TABLE_ID_FASTQUERY.getTable(); + // Run a DML statement to UPDATE 2 rows of data + String dmlQuery = + String.format("UPDATE %s.%s SET StringField = 'hello' WHERE TRUE", DATASET, tableName); + QueryJobConfiguration dmlConfig = QueryJobConfiguration.newBuilder(dmlQuery).build(); + Job remoteJob = bigquery.create(JobInfo.of(dmlConfig)); + remoteJob = remoteJob.waitFor(); + assertNull(remoteJob.getStatus().getError()); + + TableResult result = remoteJob.getQueryResults(); + assertEquals(TABLE_SCHEMA, result.getSchema()); + + Job queryJob = bigquery.getJob(remoteJob.getJobId()); + JobStatistics.QueryStatistics statistics = queryJob.getStatistics(); + assertEquals(2L, statistics.getNumDmlAffectedRows().longValue()); + assertEquals(2L, statistics.getDmlStats().getUpdatedRowCount().longValue()); + } + @Test public void testScriptStatistics() throws InterruptedException { String script =