Skip to content

Commit

Permalink
feat: add dmlStatistics support (#1431)
Browse files Browse the repository at this point in the history
Provides detailed statistics for DML statements Present only for DML statements INSERT, UPDATE, DELETE or TRUNCATE.

Towards b/186432630
  • Loading branch information
stephaniewang526 authored Jul 8, 2021
1 parent 94ce14f commit 9d67e05
Show file tree
Hide file tree
Showing 5 changed files with 217 additions and 1 deletion.
Original file line number Diff line number Diff line change
@@ -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();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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());
}
}
}

Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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));
Expand Down
Original file line number Diff line number Diff line change
@@ -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());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand Down

0 comments on commit 9d67e05

Please sign in to comment.