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

Author: whuffman36

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

@@ -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> {
@@ -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);
@@ -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) {
@@ -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. */
@@ -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);
     }
@@ -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. */
@@ -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);
@@ -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
@@ -311,7 +337,8 @@ public int hashCode() {
         createDisposition,
         writeDisposition,
         labels,
-        jobTimeoutMs);
+        jobTimeoutMs,
+        reservation);
   }
 
   @Override
@@ -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;
   }

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

@@ -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> {
@@ -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);
@@ -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) {
@@ -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. */
@@ -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);
     }
@@ -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. */
@@ -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);
@@ -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
@@ -313,7 +339,8 @@ public int hashCode() {
         compression,
         useAvroLogicalTypes,
         labels,
-        jobTimeoutMs);
+        jobTimeoutMs,
+        reservation);
   }
 
   @Override
@@ -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;
   }

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

@@ -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 {
@@ -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);
@@ -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) {
@@ -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
@@ -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);
@@ -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
@@ -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);
@@ -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
@@ -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;