feat: add reservation usage in job statistics (#1018)

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java

@@ -43,6 +43,7 @@ public abstract class JobStatistics implements Serializable {
   private final Long numChildJobs;
   private final String parentJobId;
   private final ScriptStatistics scriptStatistics;
+  private final List<ReservationUsage> reservationUsage;
 
   /** A Google BigQuery Copy Job statistics. */
   public static class CopyStatistics extends JobStatistics {
@@ -1047,6 +1048,118 @@ static ScriptStatistics fromPb(
     }
   }
 
+  /** ReservationUsage contains information about a job's usage of a single reservation. */
+  public static class ReservationUsage {
+
+    static final Function<
+            com.google.api.services.bigquery.model.JobStatistics.ReservationUsage, ReservationUsage>
+        FROM_PB_FUNCTION =
+            new Function<
+                com.google.api.services.bigquery.model.JobStatistics.ReservationUsage,
+                ReservationUsage>() {
+              @Override
+              public ReservationUsage apply(
+                  com.google.api.services.bigquery.model.JobStatistics.ReservationUsage usage) {
+                return ReservationUsage.fromPb(usage);
+              }
+            };
+
+    static final Function<
+            ReservationUsage, com.google.api.services.bigquery.model.JobStatistics.ReservationUsage>
+        TO_PB_FUNCTION =
+            new Function<
+                ReservationUsage,
+                com.google.api.services.bigquery.model.JobStatistics.ReservationUsage>() {
+              @Override
+              public com.google.api.services.bigquery.model.JobStatistics.ReservationUsage apply(
+                  ReservationUsage usage) {
+                return usage.toPb();
+              }
+            };
+
+    private final String name;
+    private final Long slotMs;
+
+    public static class Builder {
+
+      private String name;
+      private Long slotMs;
+
+      private Builder() {};
+
+      Builder setName(String name) {
+        this.name = name;
+        return this;
+      }
+
+      Builder setSlotMs(Long slotMs) {
+        this.slotMs = slotMs;
+        return this;
+      }
+
+      ReservationUsage build() {
+        return new ReservationUsage(this);
+      }
+    }
+
+    private ReservationUsage(Builder builder) {
+      this.name = builder.name;
+      this.slotMs = builder.slotMs;
+    }
+
+    // Return mame indicates the utilized reservation name, or "unreserved" for ondemand usage.
+    public String getName() {
+      return name;
+    }
+
+    // Returns slotMs reports the slot milliseconds utilized within in the given reservation.
+    public Long getSlotMs() {
+      return slotMs;
+    }
+
+    static Builder newBuilder() {
+      return new Builder();
+    }
+
+    ToStringHelper toStringHelper() {
+      return MoreObjects.toStringHelper(this).add("name", name).add("slotMs", slotMs);
+    }
+
+    @Override
+    public String toString() {
+      return toStringHelper().toString();
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+      return obj == this
+          || obj != null
+              && obj.getClass().equals(ReservationUsage.class)
+              && Objects.equals(toPb(), ((ReservationUsage) obj).toPb());
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(name, slotMs);
+    }
+
+    com.google.api.services.bigquery.model.JobStatistics.ReservationUsage toPb() {
+      com.google.api.services.bigquery.model.JobStatistics.ReservationUsage usage =
+          new com.google.api.services.bigquery.model.JobStatistics.ReservationUsage();
+      usage.setName(name);
+      usage.setSlotMs(slotMs);
+      return usage;
+    }
+
+    static ReservationUsage fromPb(
+        com.google.api.services.bigquery.model.JobStatistics.ReservationUsage usage) {
+      Builder builder = newBuilder();
+      builder.setName(usage.getName());
+      builder.setSlotMs(usage.getSlotMs());
+      return builder.build();
+    }
+  }
+
   abstract static class Builder<T extends JobStatistics, B extends Builder<T, B>> {
 
     private Long creationTime;
@@ -1055,6 +1168,7 @@ static ScriptStatistics fromPb(
     private Long numChildJobs;
     private String parentJobId;
     private ScriptStatistics scriptStatistics;
+    private List<ReservationUsage> reservationUsage;
 
     protected Builder() {}
 
@@ -1067,6 +1181,10 @@ protected Builder(com.google.api.services.bigquery.model.JobStatistics statistic
       if (statisticsPb.getScriptStatistics() != null) {
         this.scriptStatistics = ScriptStatistics.fromPb(statisticsPb.getScriptStatistics());
       }
+      if (reservationUsage != null) {
+        this.reservationUsage =
+            Lists.transform(statisticsPb.getReservationUsage(), ReservationUsage.FROM_PB_FUNCTION);
+      }
     }
 
     @SuppressWarnings("unchecked")
@@ -1099,6 +1217,7 @@ protected JobStatistics(Builder builder) {
     this.numChildJobs = builder.numChildJobs;
     this.parentJobId = builder.parentJobId;
     this.scriptStatistics = builder.scriptStatistics;
+    this.reservationUsage = builder.reservationUsage;
   }
 
   /** Returns the creation time of the job in milliseconds since epoch. */
@@ -1137,14 +1256,20 @@ public ScriptStatistics getScriptStatistics() {
     return scriptStatistics;
   }
 
+  /** ReservationUsage contains information about a job's usage of a single reservation. */
+  public List<ReservationUsage> getReservationUsage() {
+    return reservationUsage;
+  }
+
   ToStringHelper toStringHelper() {
     return MoreObjects.toStringHelper(this)
         .add("creationTime", creationTime)
         .add("endTime", endTime)
         .add("startTime", startTime)
         .add("numChildJobs", numChildJobs)
         .add("parentJobId", parentJobId)
-        .add("scriptStatistics", scriptStatistics);
+        .add("scriptStatistics", scriptStatistics)
+        .add("reservationUsage", reservationUsage);
   }
 
   @Override
@@ -1154,7 +1279,13 @@ public String toString() {
 
   final int baseHashCode() {
     return Objects.hash(
-        creationTime, endTime, startTime, numChildJobs, parentJobId, scriptStatistics);
+        creationTime,
+        endTime,
+        startTime,
+        numChildJobs,
+        parentJobId,
+        scriptStatistics,
+        reservationUsage);
   }
 
   final boolean baseEquals(JobStatistics jobStatistics) {
@@ -1172,6 +1303,10 @@ com.google.api.services.bigquery.model.JobStatistics toPb() {
     if (scriptStatistics != null) {
       statistics.setScriptStatistics(scriptStatistics.toPb());
     }
+    if (reservationUsage != null) {
+      statistics.setReservationUsage(
+          Lists.transform(reservationUsage, ReservationUsage.TO_PB_FUNCTION));
+    }
     return statistics;
   }
 

google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java

@@ -23,6 +23,7 @@
 import com.google.cloud.bigquery.JobStatistics.ExtractStatistics;
 import com.google.cloud.bigquery.JobStatistics.LoadStatistics;
 import com.google.cloud.bigquery.JobStatistics.QueryStatistics;
+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.QueryStage.QueryStep;
@@ -58,6 +59,8 @@ public class JobStatisticsTest {
   private static final Long CREATION_TIME = 10L;
   private static final Long END_TIME = 20L;
   private static final Long START_TIME = 15L;
+  private static final String NAME = "reservation-name";
+  private static final Long SLOTMS = 12545L;
   private static final CopyStatistics COPY_STATISTICS =
       CopyStatistics.newBuilder()
           .setCreationTimestamp(CREATION_TIME)
@@ -200,6 +203,8 @@ public class JobStatisticsTest {
           .setEvaluationKind(EVALUATIONKIND_TYPE_EXPRESSION)
           .setStackFrames(ImmutableList.of(EXPRESSION_STACK_FRAME))
           .build();
+  private static final ReservationUsage RESERVATION_USAGE =
+      ReservationUsage.newBuilder().setName(NAME).setSlotMs(SLOTMS).build();
 
   @Test
   public void testBuilder() {
@@ -268,6 +273,8 @@ public void testBuilder() {
     assertEquals(EVALUATIONKIND_TYPE_EXPRESSION, EXPRESSION_SCRIPT_STATISTICS.getEvaluationKind());
     assertEquals(
         ImmutableList.of(EXPRESSION_STACK_FRAME), EXPRESSION_SCRIPT_STATISTICS.getStackFrames());
+    assertEquals(NAME, RESERVATION_USAGE.getName());
+    assertEquals(SLOTMS, RESERVATION_USAGE.getSlotMs());
   }
 
   @Test
@@ -292,6 +299,7 @@ public void testToPbAndFromPb() {
     for (ScriptStackFrame stackFrame : EXPRESSION_SCRIPT_STATISTICS.getStackFrames()) {
       compareStackFrames(stackFrame, ScriptStackFrame.fromPb(stackFrame.toPb()));
     }
+    compareReservation(RESERVATION_USAGE, ReservationUsage.fromPb(RESERVATION_USAGE.toPb()));
   }
 
   @Test
@@ -392,4 +400,13 @@ private void compareStackFrames(
     assertEquals(expected.getStartLine(), value.getStartLine());
     assertEquals(expected.getText(), value.getText());
   }
+
+  private void compareReservation(ReservationUsage expected, ReservationUsage value) {
+    assertEquals(expected, value);
+    assertEquals(expected.hashCode(), value.hashCode());
+    assertEquals(expected.toString(), value.toString());
+    assertEquals(expected.toPb(), value.toPb());
+    assertEquals(expected.getName(), value.getName());
+    assertEquals(expected.getSlotMs(), value.getSlotMs());
+  }
 }