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

Author: Praful Makani

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 @@ abstract static class Builder<T extends JobStatistics, B extends Builder<T, B>>
  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());
+ }
 }