feat: add dmlStatistics support (#1431)

Provides detailed statistics for DML statements Present only for DML statements INSERT, UPDATE, DELETE or TRUNCATE.

Towards b/186432630

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();
+  }
+}

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<TableId> referencedTables;
     private final StatementType statementType;
     private final Long totalBytesBilled;
@@ -406,6 +407,7 @@ static final class Builder extends JobStatistics.Builder<QueryStatistics, Builde
       private RoutineId ddlTargetRoutine;
       private Long estimatedBytesProcessed;
       private Long numDmlAffectedRows;
+      private DmlStats dmlStats;
       private List<TableId> 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<TableId> 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));

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());
+  }
+}

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());

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 =