Skip to content
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.

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

feat: add reservation usage in job statistics #1018

Merged
merged 1 commit into from Feb 5, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
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
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());
}
}