From 4c3b2de16584b9079dd4afd5a33aa8c00bc75f20 Mon Sep 17 00:00:00 2001 From: Stephanie Wang Date: Wed, 11 Aug 2021 12:20:36 -0400 Subject: [PATCH] feat: add support for transactioninfo in query statistics (#1497) Fixes #1467 --- .../google/cloud/bigquery/JobStatistics.java | 92 ++++++++++++++++++- .../cloud/bigquery/JobStatisticsTest.java | 16 ++++ .../cloud/bigquery/it/ITBigQueryTest.java | 22 +++++ 3 files changed, 128 insertions(+), 2 deletions(-) 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 ea5555b9b..a3d04c017 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 @@ -44,6 +44,7 @@ public abstract class JobStatistics implements Serializable { private final String parentJobId; private final ScriptStatistics scriptStatistics; private final List reservationUsage; + private final TransactionInfo transactionInfo; /** A Google BigQuery Copy Job statistics. */ public static class CopyStatistics extends JobStatistics { @@ -1178,6 +1179,78 @@ static ReservationUsage fromPb( } } + // TransactionInfo contains information about a multi-statement transaction that may have + // associated with a job. + public static class TransactionInfo { + + // TransactionID is the system-generated identifier for the transaction. + private final String transactionId; + + public static class Builder { + + private String transactionId; + + private Builder() {}; + + Builder setTransactionId(String transactionId) { + this.transactionId = transactionId; + return this; + } + + TransactionInfo build() { + return new TransactionInfo(this); + } + } + + private TransactionInfo(Builder builder) { + this.transactionId = builder.transactionId; + } + + public String getTransactionId() { + return transactionId; + } + + static Builder newbuilder() { + return new Builder(); + } + + ToStringHelper toStringHelper() { + return MoreObjects.toStringHelper(this).add("transactionId", transactionId); + } + + @Override + public String toString() { + return toStringHelper().toString(); + } + + @Override + public boolean equals(Object obj) { + return obj == this + || obj != null + && obj.getClass().equals(TransactionInfo.class) + && Objects.equals(toPb(), ((TransactionInfo) obj).toPb()); + } + + @Override + public int hashCode() { + return Objects.hash(transactionId); + } + + com.google.api.services.bigquery.model.TransactionInfo toPb() { + com.google.api.services.bigquery.model.TransactionInfo transactionInfo = + new com.google.api.services.bigquery.model.TransactionInfo(); + transactionInfo.setTransactionId(transactionId); + return transactionInfo; + } + + static TransactionInfo fromPb( + com.google.api.services.bigquery.model.TransactionInfo transactionInfo) { + Builder builder = newbuilder(); + builder.setTransactionId(transactionInfo.getTransactionId()); + return builder.build(); + } + } + abstract static class Builder> { private Long creationTime; @@ -1187,6 +1260,7 @@ abstract static class Builder> private String parentJobId; private ScriptStatistics scriptStatistics; private List reservationUsage; + private TransactionInfo transactionInfo; protected Builder() {} @@ -1203,6 +1277,9 @@ protected Builder(com.google.api.services.bigquery.model.JobStatistics statistic this.reservationUsage = Lists.transform(statisticsPb.getReservationUsage(), ReservationUsage.FROM_PB_FUNCTION); } + if (statisticsPb.getTransactionInfo() != null) { + this.transactionInfo = TransactionInfo.fromPb(statisticsPb.getTransactionInfo()); + } } @SuppressWarnings("unchecked") @@ -1236,6 +1313,7 @@ protected JobStatistics(Builder builder) { this.parentJobId = builder.parentJobId; this.scriptStatistics = builder.scriptStatistics; this.reservationUsage = builder.reservationUsage; + this.transactionInfo = builder.transactionInfo; } /** Returns the creation time of the job in milliseconds since epoch. */ @@ -1279,6 +1357,11 @@ public List getReservationUsage() { return reservationUsage; } + /** Info indicates the transaction ID associated with the job, if any. */ + public TransactionInfo getTransactionInfo() { + return transactionInfo; + } + ToStringHelper toStringHelper() { return MoreObjects.toStringHelper(this) .add("creationTime", creationTime) @@ -1287,7 +1370,8 @@ ToStringHelper toStringHelper() { .add("numChildJobs", numChildJobs) .add("parentJobId", parentJobId) .add("scriptStatistics", scriptStatistics) - .add("reservationUsage", reservationUsage); + .add("reservationUsage", reservationUsage) + .add("transactionInfo", transactionInfo); } @Override @@ -1303,7 +1387,8 @@ final int baseHashCode() { numChildJobs, parentJobId, scriptStatistics, - reservationUsage); + reservationUsage, + transactionInfo); } final boolean baseEquals(JobStatistics jobStatistics) { @@ -1325,6 +1410,9 @@ com.google.api.services.bigquery.model.JobStatistics toPb() { statistics.setReservationUsage( Lists.transform(reservationUsage, ReservationUsage.TO_PB_FUNCTION)); } + if (transactionInfo != null) { + statistics.setTransactionInfo(transactionInfo.toPb()); + } return statistics; } 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 edd76f821..c421c0370 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 @@ -26,9 +26,11 @@ import com.google.cloud.bigquery.JobStatistics.ReservationUsage; import com.google.cloud.bigquery.JobStatistics.ScriptStatistics; import com.google.cloud.bigquery.JobStatistics.ScriptStatistics.ScriptStackFrame; +import com.google.cloud.bigquery.JobStatistics.TransactionInfo; import com.google.cloud.bigquery.QueryStage.QueryStep; import com.google.common.collect.ImmutableList; import java.util.List; +import java.util.UUID; import org.junit.Test; public class JobStatisticsTest { @@ -70,6 +72,7 @@ public class JobStatisticsTest { private static final Long START_TIME = 15L; private static final String NAME = "reservation-name"; private static final Long SLOTMS = 12545L; + private static final String TRANSACTION_ID = UUID.randomUUID().toString().substring(0, 8); private static final CopyStatistics COPY_STATISTICS = CopyStatistics.newBuilder() .setCreationTimestamp(CREATION_TIME) @@ -216,6 +219,9 @@ public class JobStatisticsTest { private static final ReservationUsage RESERVATION_USAGE = ReservationUsage.newBuilder().setName(NAME).setSlotMs(SLOTMS).build(); + private static final TransactionInfo TRANSACTION_INFO = + TransactionInfo.newbuilder().setTransactionId(TRANSACTION_ID).build(); + @Test public void testBuilder() { assertEquals(CREATION_TIME, EXTRACT_STATISTICS.getCreationTime()); @@ -286,6 +292,7 @@ public void testBuilder() { ImmutableList.of(EXPRESSION_STACK_FRAME), EXPRESSION_SCRIPT_STATISTICS.getStackFrames()); assertEquals(NAME, RESERVATION_USAGE.getName()); assertEquals(SLOTMS, RESERVATION_USAGE.getSlotMs()); + assertEquals(TRANSACTION_ID, TRANSACTION_INFO.getTransactionId()); } @Test @@ -311,6 +318,7 @@ public void testToPbAndFromPb() { compareStackFrames(stackFrame, ScriptStackFrame.fromPb(stackFrame.toPb())); } compareReservation(RESERVATION_USAGE, ReservationUsage.fromPb(RESERVATION_USAGE.toPb())); + compareTransactionInfo(TRANSACTION_INFO, TransactionInfo.fromPb(TRANSACTION_INFO.toPb())); } @Test @@ -425,4 +433,12 @@ private void compareReservation(ReservationUsage expected, ReservationUsage valu assertEquals(expected.getName(), value.getName()); assertEquals(expected.getSlotMs(), value.getSlotMs()); } + + private void compareTransactionInfo(TransactionInfo expected, TransactionInfo value) { + assertEquals(expected, value); + assertEquals(expected.hashCode(), value.hashCode()); + assertEquals(expected.toString(), value.toString()); + assertEquals(expected.toPb(), value.toPb()); + assertEquals(expected.getTransactionId(), value.getTransactionId()); + } } 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 ff4068940..484e4c0f5 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 @@ -71,6 +71,7 @@ import com.google.cloud.bigquery.JobInfo; import com.google.cloud.bigquery.JobStatistics; import com.google.cloud.bigquery.JobStatistics.LoadStatistics; +import com.google.cloud.bigquery.JobStatistics.TransactionInfo; import com.google.cloud.bigquery.LegacySQLTypeName; import com.google.cloud.bigquery.LoadJobConfiguration; import com.google.cloud.bigquery.MaterializedViewDefinition; @@ -2219,6 +2220,27 @@ public void testDmlStatistics() throws InterruptedException { assertEquals(2L, statistics.getDmlStats().getUpdatedRowCount().longValue()); } + @Test + public void testTransactionInfo() throws InterruptedException { + String tableName = TABLE_ID_FASTQUERY.getTable(); + String transaction = + String.format( + "BEGIN TRANSACTION;\n" + + " UPDATE %s.%s SET StringField = 'hello' WHERE TRUE;\n" + + " COMMIT TRANSACTION;\n", + DATASET, tableName); + QueryJobConfiguration config = QueryJobConfiguration.of(transaction); + Job remoteJob = bigquery.create(JobInfo.of(config)); + JobInfo parentJobInfo = remoteJob.waitFor(); + String parentJobId = parentJobInfo.getJobId().getJob(); + Page childJobs = bigquery.listJobs(JobListOption.parentJobId(parentJobId)); + for (Job job : childJobs.iterateAll()) { + // only those child jobs inside the transaction would have transactionInfo populated + TransactionInfo transactionInfo = job.getStatistics().getTransactionInfo(); + assertNotNull(transactionInfo.getTransactionId()); + } + } + @Test public void testScriptStatistics() throws InterruptedException { String script =