Skip to content

feat: add support for transactioninfo in query statistics #1497

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
gcf-merge-on-green merged 2 commits into from
Aug 11, 2021
Merged

feat: add support for transactioninfo in query statistics #1497

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
55fcb36
feat: add support for transactioninfo in query statistics
stephaniewang526 Aug 10, 2021
72aad98
simplify test case
stephaniewang526 Aug 10, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
92 changes: 90 additions & 2 deletions google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java
View file
Open in desktop
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
16 changes: 16 additions & 0 deletions google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java
View file
Open in desktop
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());
}
}
22 changes: 22 additions & 0 deletions google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java
View file
Open in desktop
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();
Copy link
Contributor

@shollyman shollyman Aug 10, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you need to worry about test ordering, or is this created in a setup?

Copy link
Contributor Author

@stephaniewang526 stephaniewang526 Aug 10, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for pointing this out. Yes, this is fine since the later test case will still update the table given the WHERE TRUE condition.

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