Skip to content

Commit

Permalink
feat: add reservation usage in job statistics (#1018)
Browse files Browse the repository at this point in the history
  • Loading branch information
Praful Makani committed Feb 5, 2021
1 parent 24a6a3b commit 49cbb0f
Show file tree
Hide file tree
Showing 2 changed files with 154 additions and 2 deletions.
139 changes: 137 additions & 2 deletions google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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;
Expand All @@ -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() {}

Expand All @@ -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")
Expand Down Expand Up @@ -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. */
Expand Down Expand Up @@ -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
Expand All @@ -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) {
Expand All @@ -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;
}

Expand Down
17 changes: 17 additions & 0 deletions google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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() {
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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());
}
}

0 comments on commit 49cbb0f

Please sign in to comment.