Skip to content
This repository was archived by the owner on Mar 23, 2026. It is now read-only.

Commit 9d67e05

Browse files
stephaniewang526stephaniewang526
authored
feat: add dmlStatistics support (#1431)
Provides detailed statistics for DML statements Present only for DML statements INSERT, UPDATE, DELETE or TRUNCATE. Towards b/186432630
1 parent 94ce14f commit 9d67e05

5 files changed

Lines changed: 217 additions & 1 deletion

File tree

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

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
/*
2+
* Copyright 2021 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.bigquery;
18+
19+
import com.google.api.services.bigquery.model.DmlStatistics;
20+
import com.google.auto.value.AutoValue;
21+
import java.io.Serializable;
22+
import javax.annotation.Nullable;
23+
24+
/** Represents DML statistics information. */
25+
@AutoValue
26+
public abstract class DmlStats implements Serializable {
27+
28+
@AutoValue.Builder
29+
public abstract static class Builder {
30+
/**
31+
* Number of deleted Rows. populated by DML DELETE, MERGE and TRUNCATE statements.
32+
*
33+
* @param deletedRowCount deletedRowCount or {@code null} for none
34+
*/
35+
public abstract Builder setDeletedRowCount(Long deletedRowCount);
36+
37+
/**
38+
* Number of inserted Rows. Populated by DML INSERT and MERGE statements.
39+
*
40+
* @param insertedRowCount insertedRowCount or {@code null} for none
41+
*/
42+
public abstract Builder setInsertedRowCount(Long insertedRowCount);
43+
44+
/**
45+
* Number of updated Rows. Populated by DML UPDATE and MERGE statements.
46+
*
47+
* @param updatedRowCount updatedRowCount or {@code null} for none
48+
*/
49+
public abstract Builder setUpdatedRowCount(Long updatedRowCount);
50+
51+
/** Creates a {@code DmlStats} object. */
52+
public abstract DmlStats build();
53+
}
54+
55+
/**
56+
* Returns number of deleted Rows. populated by DML DELETE, MERGE and TRUNCATE statements.
57+
*
58+
* @return value or {@code null} for none
59+
*/
60+
@Nullable
61+
public abstract Long getDeletedRowCount();
62+
63+
/**
64+
* Returns number of inserted Rows. Populated by DML INSERT and MERGE statements.
65+
*
66+
* @return value or {@code null} for none
67+
*/
68+
@Nullable
69+
public abstract Long getInsertedRowCount();
70+
71+
/**
72+
* Returns number of updated Rows. Populated by DML UPDATE and MERGE statements.
73+
*
74+
* @return value or {@code null} for none
75+
*/
76+
@Nullable
77+
public abstract Long getUpdatedRowCount();
78+
79+
public abstract Builder toBuilder();
80+
81+
public static Builder newBuilder() {
82+
return new AutoValue_DmlStats.Builder();
83+
}
84+
85+
DmlStatistics toPb() {
86+
DmlStatistics dmlStatisticsPb = new DmlStatistics();
87+
if (getDeletedRowCount() != null) {
88+
dmlStatisticsPb.setDeletedRowCount(getDeletedRowCount());
89+
}
90+
if (getInsertedRowCount() != null) {
91+
dmlStatisticsPb.setInsertedRowCount(getInsertedRowCount());
92+
}
93+
if (getUpdatedRowCount() != null) {
94+
dmlStatisticsPb.setUpdatedRowCount(getUpdatedRowCount());
95+
}
96+
return dmlStatisticsPb;
97+
}
98+
99+
static DmlStats fromPb(DmlStatistics dmlStatisticsPb) {
100+
Builder builder = newBuilder();
101+
if (dmlStatisticsPb.getDeletedRowCount() != null) {
102+
builder.setDeletedRowCount(dmlStatisticsPb.getDeletedRowCount());
103+
}
104+
if (dmlStatisticsPb.getInsertedRowCount() != null) {
105+
builder.setInsertedRowCount(dmlStatisticsPb.getInsertedRowCount());
106+
}
107+
if (dmlStatisticsPb.getUpdatedRowCount() != null) {
108+
builder.setUpdatedRowCount(dmlStatisticsPb.getUpdatedRowCount());
109+
}
110+
return builder.build();
111+
}
112+
}

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

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -326,6 +326,7 @@ public static class QueryStatistics extends JobStatistics {
326326
private final RoutineId ddlTargetRoutine;
327327
private final Long estimatedBytesProcessed;
328328
private final Long numDmlAffectedRows;
329+
private final DmlStats dmlStats;
329330
private final List<TableId> referencedTables;
330331
private final StatementType statementType;
331332
private final Long totalBytesBilled;
@@ -406,6 +407,7 @@ static final class Builder extends JobStatistics.Builder<QueryStatistics, Builde
406407
private RoutineId ddlTargetRoutine;
407408
private Long estimatedBytesProcessed;
408409
private Long numDmlAffectedRows;
410+
private DmlStats dmlStats;
409411
private List<TableId> referencedTables;
410412
private StatementType statementType;
411413
private Long totalBytesBilled;
@@ -458,6 +460,9 @@ private Builder(com.google.api.services.bigquery.model.JobStatistics statisticsP
458460
if (statisticsPb.getQuery().getSchema() != null) {
459461
this.schema = Schema.fromPb(statisticsPb.getQuery().getSchema());
460462
}
463+
if (statisticsPb.getQuery().getDmlStats() != null) {
464+
this.dmlStats = DmlStats.fromPb(statisticsPb.getQuery().getDmlStats());
465+
}
461466
}
462467
}
463468

@@ -496,6 +501,11 @@ Builder setNumDmlAffectedRows(Long numDmlAffectedRows) {
496501
return self();
497502
}
498503

504+
Builder setDmlStats(DmlStats dmlStats) {
505+
this.dmlStats = dmlStats;
506+
return self();
507+
}
508+
499509
Builder setReferenceTables(List<TableId> referencedTables) {
500510
this.referencedTables = referencedTables;
501511
return self();
@@ -561,6 +571,7 @@ private QueryStatistics(Builder builder) {
561571
this.ddlTargetRoutine = builder.ddlTargetRoutine;
562572
this.estimatedBytesProcessed = builder.estimatedBytesProcessed;
563573
this.numDmlAffectedRows = builder.numDmlAffectedRows;
574+
this.dmlStats = builder.dmlStats;
564575
this.referencedTables = builder.referencedTables;
565576
this.statementType = builder.statementType;
566577
this.totalBytesBilled = builder.totalBytesBilled;
@@ -614,6 +625,11 @@ public Long getNumDmlAffectedRows() {
614625
return numDmlAffectedRows;
615626
}
616627

628+
/** Detailed statistics for DML statements. */
629+
public DmlStats getDmlStats() {
630+
return dmlStats;
631+
}
632+
617633
/**
618634
* Referenced tables for the job. Queries that reference more than 50 tables will not have a
619635
* complete list.
@@ -729,7 +745,9 @@ com.google.api.services.bigquery.model.JobStatistics toPb() {
729745
if (ddlTargetRoutine != null) {
730746
queryStatisticsPb.setDdlTargetRoutine(ddlTargetRoutine.toPb());
731747
}
732-
748+
if (dmlStats != null) {
749+
queryStatisticsPb.setDmlStats(dmlStats.toPb());
750+
}
733751
if (referencedTables != null) {
734752
queryStatisticsPb.setReferencedTables(
735753
Lists.transform(referencedTables, TableId.TO_PB_FUNCTION));

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

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
/*
2+
* Copyright 2021 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.bigquery;
18+
19+
import static org.junit.Assert.assertEquals;
20+
21+
import org.junit.Test;
22+
23+
public class DmlStatsTest {
24+
25+
private static final Long DELETED_ROW_COUNT = 10L;
26+
private static final Long INSERTED_ROW_COUNT = 20L;
27+
private static final Long UPDATED_ROW_COUNT = 30L;
28+
private static final DmlStats DML_STATS =
29+
DmlStats.newBuilder()
30+
.setDeletedRowCount(DELETED_ROW_COUNT)
31+
.setInsertedRowCount(INSERTED_ROW_COUNT)
32+
.setUpdatedRowCount(UPDATED_ROW_COUNT)
33+
.build();
34+
35+
@Test
36+
public void testBuilder() {
37+
assertEquals(DELETED_ROW_COUNT, DML_STATS.getDeletedRowCount());
38+
assertEquals(UPDATED_ROW_COUNT, DML_STATS.getUpdatedRowCount());
39+
assertEquals(INSERTED_ROW_COUNT, DML_STATS.getInsertedRowCount());
40+
}
41+
42+
@Test
43+
public void testToPbAndFromPb() {
44+
compareDmlStats(DML_STATS, DmlStats.fromPb(DML_STATS.toPb()));
45+
}
46+
47+
private void compareDmlStats(DmlStats expected, DmlStats actual) {
48+
assertEquals(expected, actual);
49+
assertEquals(expected.hashCode(), actual.hashCode());
50+
assertEquals(expected.toString(), actual.toString());
51+
assertEquals(expected.getDeletedRowCount(), actual.getDeletedRowCount());
52+
assertEquals(expected.getInsertedRowCount(), actual.getInsertedRowCount());
53+
assertEquals(expected.getUpdatedRowCount(), actual.getUpdatedRowCount());
54+
}
55+
}

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,15 @@ public class JobStatisticsTest {
4242
private static final RoutineId DDL_TARGET_ROUTINE = RoutineId.of("alpha", "beta", "gamma");
4343
private static final Long ESTIMATE_BYTES_PROCESSED = 101L;
4444
private static final Long NUM_DML_AFFECTED_ROWS = 88L;
45+
private static final Long DELETED_ROW_COUNT = 10L;
46+
private static final Long INSERTED_ROW_COUNT = 20L;
47+
private static final Long UPDATED_ROW_COUNT = 30L;
48+
private static final DmlStats DML_STATS =
49+
DmlStats.newBuilder()
50+
.setDeletedRowCount(DELETED_ROW_COUNT)
51+
.setInsertedRowCount(INSERTED_ROW_COUNT)
52+
.setUpdatedRowCount(UPDATED_ROW_COUNT)
53+
.build();
4554
private static final QueryStatistics.StatementType STATEMENT_TYPE =
4655
QueryStatistics.StatementType.SELECT;
4756
private static final Long TOTAL_BYTES_BILLED = 24L;
@@ -147,6 +156,7 @@ public class JobStatisticsTest {
147156
.setDDLTargetRoutine(DDL_TARGET_ROUTINE)
148157
.setEstimatedBytesProcessed(ESTIMATE_BYTES_PROCESSED)
149158
.setNumDmlAffectedRows(NUM_DML_AFFECTED_ROWS)
159+
.setDmlStats(DML_STATS)
150160
.setReferenceTables(REFERENCED_TABLES)
151161
.setStatementType(STATEMENT_TYPE)
152162
.setTotalBytesBilled(TOTAL_BYTES_BILLED)
@@ -232,6 +242,7 @@ public void testBuilder() {
232242
assertEquals(DDL_TARGET_ROUTINE, QUERY_STATISTICS.getDdlTargetRoutine());
233243
assertEquals(ESTIMATE_BYTES_PROCESSED, QUERY_STATISTICS.getEstimatedBytesProcessed());
234244
assertEquals(NUM_DML_AFFECTED_ROWS, QUERY_STATISTICS.getNumDmlAffectedRows());
245+
assertEquals(DML_STATS, QUERY_STATISTICS.getDmlStats());
235246
assertEquals(REFERENCED_TABLES, QUERY_STATISTICS.getReferencedTables());
236247
assertEquals(STATEMENT_TYPE, QUERY_STATISTICS.getStatementType());
237248
assertEquals(TOTAL_BYTES_BILLED, QUERY_STATISTICS.getTotalBytesBilled());

‎google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2205,6 +2205,26 @@ public void testFastQueryHTTPException() throws InterruptedException {
22052205
}
22062206
}
22072207

2208+
@Test
2209+
public void testDmlStatistics() throws InterruptedException {
2210+
String tableName = TABLE_ID_FASTQUERY.getTable();
2211+
// Run a DML statement to UPDATE 2 rows of data
2212+
String dmlQuery =
2213+
String.format("UPDATE %s.%s SET StringField = 'hello' WHERE TRUE", DATASET, tableName);
2214+
QueryJobConfiguration dmlConfig = QueryJobConfiguration.newBuilder(dmlQuery).build();
2215+
Job remoteJob = bigquery.create(JobInfo.of(dmlConfig));
2216+
remoteJob = remoteJob.waitFor();
2217+
assertNull(remoteJob.getStatus().getError());
2218+
2219+
TableResult result = remoteJob.getQueryResults();
2220+
assertEquals(TABLE_SCHEMA, result.getSchema());
2221+
2222+
Job queryJob = bigquery.getJob(remoteJob.getJobId());
2223+
JobStatistics.QueryStatistics statistics = queryJob.getStatistics();
2224+
assertEquals(2L, statistics.getNumDmlAffectedRows().longValue());
2225+
assertEquals(2L, statistics.getDmlStats().getUpdatedRowCount().longValue());
2226+
}
2227+
22082228
@Test
22092229
public void testScriptStatistics() throws InterruptedException {
22102230
String script =

0 commit comments

Comments
 (0)