feat: add support for transactioninfo in query statistics (#1497)

Fixes #1467

Author: stephaniewang526

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> 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<T extends JobStatistics, B extends Builder<T, B>> {
 
     private Long creationTime;
@@ -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() {}
 
@@ -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<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)
@@ -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;
   }
 

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

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<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 =