Skip to content

feat(bigquery): Add support for reservation field in jobs. #3768

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.

Sign up for GitHub

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

Merged
whuffman36 merged 1 commit into from
Apr 30, 2025
Merged

feat(bigquery): Add support for reservation field in jobs. #3768

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 32 additions & 2 deletions google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/CopyJobConfiguration.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ public final class CopyJobConfiguration extends JobConfiguration {
private final EncryptionConfiguration destinationEncryptionConfiguration;
private final Map<String, String> labels;
private final Long jobTimeoutMs;
private final String reservation;

public static final class Builder
extends JobConfiguration.Builder<CopyJobConfiguration, Builder> {
Expand All @@ -58,6 +59,7 @@ public static final class Builder
private EncryptionConfiguration destinationEncryptionConfiguration;
private Map<String, String> labels;
private Long jobTimeoutMs;
private String reservation;

private Builder() {
super(Type.COPY);
Expand All @@ -74,6 +76,7 @@ private Builder(CopyJobConfiguration jobConfiguration) {
this.destinationEncryptionConfiguration = jobConfiguration.destinationEncryptionConfiguration;
this.labels = jobConfiguration.labels;
this.jobTimeoutMs = jobConfiguration.jobTimeoutMs;
this.reservation = jobConfiguration.reservation;
}

private Builder(com.google.api.services.bigquery.model.JobConfiguration configurationPb) {
Expand Down Expand Up @@ -113,6 +116,9 @@ private Builder(com.google.api.services.bigquery.model.JobConfiguration configur
if (configurationPb.getJobTimeoutMs() != null) {
this.jobTimeoutMs = configurationPb.getJobTimeoutMs();
}
if (configurationPb.getReservation() != null) {
this.reservation = configurationPb.getReservation();
}
}

/** Sets the source tables to copy. */
Expand Down Expand Up @@ -201,6 +207,19 @@ public Builder setJobTimeoutMs(Long jobTimeoutMs) {
return this;
}

/**
* [Optional] The reservation that job would use. User can specify a reservation to execute the
* job. If reservation is not set, reservation is determined based on the rules defined by the
* reservation assignments. The expected format is
* `projects/{project}/locations/{location}/reservations/{reservation}`.
*
* @param reservation reservation or {@code null} for none
*/
public Builder setReservation(String reservation) {
this.reservation = reservation;
return this;
}

public CopyJobConfiguration build() {
return new CopyJobConfiguration(this);
}
Expand All @@ -217,6 +236,7 @@ private CopyJobConfiguration(Builder builder) {
this.destinationEncryptionConfiguration = builder.destinationEncryptionConfiguration;
this.labels = builder.labels;
this.jobTimeoutMs = builder.jobTimeoutMs;
this.reservation = builder.reservation;
}

/** Returns the source tables to copy. */
Expand Down Expand Up @@ -275,6 +295,11 @@ public Long getJobTimeoutMs() {
return jobTimeoutMs;
}

/** Returns the reservation associated with this job */
public String getReservation() {
return reservation;
}

@Override
public Builder toBuilder() {
return new Builder(this);
Expand All @@ -291,7 +316,8 @@ ToStringHelper toStringHelper() {
.add("createDisposition", createDisposition)
.add("writeDisposition", writeDisposition)
.add("labels", labels)
.add("jobTimeoutMs", jobTimeoutMs);
.add("jobTimeoutMs", jobTimeoutMs)
.add("reservation", reservation);
}

@Override
Expand All @@ -311,7 +337,8 @@ public int hashCode() {
createDisposition,
writeDisposition,
labels,
jobTimeoutMs);
jobTimeoutMs,
reservation);
}

@Override
Expand Down Expand Up @@ -366,6 +393,9 @@ com.google.api.services.bigquery.model.JobConfiguration toPb() {
if (jobTimeoutMs != null) {
jobConfiguration.setJobTimeoutMs(jobTimeoutMs);
}
if (reservation != null) {
jobConfiguration.setReservation(reservation);
}
jobConfiguration.setCopy(configurationPb);
return jobConfiguration;
}
Expand Down
34 changes: 32 additions & 2 deletions google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ExtractJobConfiguration.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ public final class ExtractJobConfiguration extends JobConfiguration {
private final Boolean useAvroLogicalTypes;
private final Map<String, String> labels;
private final Long jobTimeoutMs;
private final String reservation;

public static final class Builder
extends JobConfiguration.Builder<ExtractJobConfiguration, Builder> {
Expand All @@ -61,6 +62,7 @@ public static final class Builder
private Boolean useAvroLogicalTypes;
private Map<String, String> labels;
private Long jobTimeoutMs;
private String reservation;

private Builder() {
super(Type.EXTRACT);
Expand All @@ -78,6 +80,7 @@ private Builder(ExtractJobConfiguration jobInfo) {
this.useAvroLogicalTypes = jobInfo.useAvroLogicalTypes;
this.labels = jobInfo.labels;
this.jobTimeoutMs = jobInfo.jobTimeoutMs;
this.reservation = jobInfo.reservation;
}

private Builder(com.google.api.services.bigquery.model.JobConfiguration configurationPb) {
Expand All @@ -101,6 +104,9 @@ private Builder(com.google.api.services.bigquery.model.JobConfiguration configur
if (configurationPb.getJobTimeoutMs() != null) {
this.jobTimeoutMs = configurationPb.getJobTimeoutMs();
}
if (configurationPb.getReservation() != null) {
this.reservation = configurationPb.getReservation();
}
}

/** Sets the table to export. */
Expand Down Expand Up @@ -198,6 +204,19 @@ public Builder setJobTimeoutMs(Long jobTimeoutMs) {
return this;
}

/**
* [Optional] The reservation that job would use. User can specify a reservation to execute the
* job. If reservation is not set, reservation is determined based on the rules defined by the
* reservation assignments. The expected format is
* `projects/{project}/locations/{location}/reservations/{reservation}`.
*
* @param reservation reservation or {@code null} for none
*/
public Builder setReservation(String reservation) {
this.reservation = reservation;
return this;
}

public ExtractJobConfiguration build() {
return new ExtractJobConfiguration(this);
}
Expand All @@ -215,6 +234,7 @@ private ExtractJobConfiguration(Builder builder) {
this.useAvroLogicalTypes = builder.useAvroLogicalTypes;
this.labels = builder.labels;
this.jobTimeoutMs = builder.jobTimeoutMs;
this.reservation = builder.reservation;
}

/** Returns the table to export. */
Expand Down Expand Up @@ -274,6 +294,11 @@ public Long getJobTimeoutMs() {
return jobTimeoutMs;
}

/** Returns the reservation associated with this job */
public String getReservation() {
return reservation;
}

@Override
public Builder toBuilder() {
return new Builder(this);
Expand All @@ -291,7 +316,8 @@ ToStringHelper toStringHelper() {
.add("compression", compression)
.add("useAvroLogicalTypes", useAvroLogicalTypes)
.add("labels", labels)
.add("jobTimeoutMs", jobTimeoutMs);
.add("jobTimeoutMs", jobTimeoutMs)
.add("reservation", reservation);
}

@Override
Expand All @@ -313,7 +339,8 @@ public int hashCode() {
compression,
useAvroLogicalTypes,
labels,
jobTimeoutMs);
jobTimeoutMs,
reservation);
}

@Override
Expand Down Expand Up @@ -350,6 +377,9 @@ com.google.api.services.bigquery.model.JobConfiguration toPb() {
if (jobTimeoutMs != null) {
jobConfiguration.setJobTimeoutMs(jobTimeoutMs);
}
if (reservation != null) {
jobConfiguration.setReservation(reservation);
}
jobConfiguration.setExtract(extractConfigurationPb);
return jobConfiguration;
}
Expand Down
33 changes: 30 additions & 3 deletions google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/LoadJobConfiguration.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,10 +60,9 @@ public final class LoadJobConfiguration extends JobConfiguration implements Load
private final RangePartitioning rangePartitioning;
private final HivePartitioningOptions hivePartitioningOptions;
private final String referenceFileSchemaUri;

private final List<ConnectionProperty> connectionProperties;

private final Boolean createSession;
private final String reservation;

public static final class Builder extends JobConfiguration.Builder<LoadJobConfiguration, Builder>
implements LoadConfiguration.Builder {
Expand Down Expand Up @@ -95,6 +94,7 @@ public static final class Builder extends JobConfiguration.Builder<LoadJobConfig
private String referenceFileSchemaUri;
private List<ConnectionProperty> connectionProperties;
private Boolean createSession;
private String reservation;

private Builder() {
super(Type.LOAD);
Expand Down Expand Up @@ -128,6 +128,7 @@ private Builder(LoadJobConfiguration loadConfiguration) {
this.referenceFileSchemaUri = loadConfiguration.referenceFileSchemaUri;
this.connectionProperties = loadConfiguration.connectionProperties;
this.createSession = loadConfiguration.createSession;
this.reservation = loadConfiguration.reservation;
}

private Builder(com.google.api.services.bigquery.model.JobConfiguration configurationPb) {
Expand Down Expand Up @@ -234,6 +235,9 @@ private Builder(com.google.api.services.bigquery.model.JobConfiguration configur
loadConfigurationPb.getConnectionProperties(), ConnectionProperty.FROM_PB_FUNCTION);
}
createSession = loadConfigurationPb.getCreateSession();
if (configurationPb.getReservation() != null) {
this.reservation = configurationPb.getReservation();
}
}

@Override
Expand Down Expand Up @@ -432,6 +436,19 @@ public Builder setCreateSession(Boolean createSession) {
return this;
}

/**
* [Optional] The reservation that job would use. User can specify a reservation to execute the
* job. If reservation is not set, reservation is determined based on the rules defined by the
* reservation assignments. The expected format is
* `projects/{project}/locations/{location}/reservations/{reservation}`.
*
* @param reservation reservation or {@code null} for none
*/
public Builder setReservation(String reservation) {
this.reservation = reservation;
return this;
}

@Override
public LoadJobConfiguration build() {
return new LoadJobConfiguration(this);
Expand Down Expand Up @@ -465,6 +482,7 @@ private LoadJobConfiguration(Builder builder) {
this.referenceFileSchemaUri = builder.referenceFileSchemaUri;
this.connectionProperties = builder.connectionProperties;
this.createSession = builder.createSession;
this.reservation = builder.reservation;
}

@Override
Expand Down Expand Up @@ -611,6 +629,11 @@ public Boolean getCreateSession() {
return createSession;
}

/** Returns the reservation associated with this job */
public String getReservation() {
return reservation;
}

@Override
public Builder toBuilder() {
return new Builder(this);
Expand Down Expand Up @@ -643,7 +666,8 @@ ToStringHelper toStringHelper() {
.add("hivePartitioningOptions", hivePartitioningOptions)
.add("referenceFileSchemaUri", referenceFileSchemaUri)
.add("connectionProperties", connectionProperties)
.add("createSession", createSession);
.add("createSession", createSession)
.add("reservation", reservation);
}

@Override
Expand Down Expand Up @@ -762,6 +786,9 @@ com.google.api.services.bigquery.model.JobConfiguration toPb() {
if (createSession != null) {
loadConfigurationPb.setCreateSession(createSession);
}
if (reservation != null) {
jobConfiguration.setReservation(reservation);
}

jobConfiguration.setLoad(loadConfigurationPb);
return jobConfiguration;
Expand Down
Loading