Skip to content

Commit

Permalink
feat: add support for transactioninfo in query statistics (#1497)
Browse files Browse the repository at this point in the history
Fixes #1467
  • Loading branch information
stephaniewang526 authored Aug 11, 2021
1 parent ec68c11 commit 4c3b2de
Show file tree
Hide file tree
Showing 3 changed files with 128 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ public abstract class JobStatistics implements Serializable {
private final String parentJobId;
private final ScriptStatistics scriptStatistics;
private final List<ReservationUsage> reservationUsage;
private final TransactionInfo transactionInfo;

/** A Google BigQuery Copy Job statistics. */
public static class CopyStatistics extends JobStatistics {
Expand Down Expand Up @@ -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<T extends JobStatistics, B extends Builder<T, B>> {

private Long creationTime;
Expand All @@ -1187,6 +1260,7 @@ abstract static class Builder<T extends JobStatistics, B extends Builder<T, B>>
private String parentJobId;
private ScriptStatistics scriptStatistics;
private List<ReservationUsage> reservationUsage;
private TransactionInfo transactionInfo;

protected Builder() {}

Expand All @@ -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")
Expand Down Expand Up @@ -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. */
Expand Down Expand Up @@ -1279,6 +1357,11 @@ public List<ReservationUsage> 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)
Expand All @@ -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
Expand All @@ -1303,7 +1387,8 @@ final int baseHashCode() {
numChildJobs,
parentJobId,
scriptStatistics,
reservationUsage);
reservationUsage,
transactionInfo);
}

final boolean baseEquals(JobStatistics jobStatistics) {
Expand All @@ -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;
}

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

0 comments on commit 4c3b2de

Please sign in to comment.