From 3174c46954de1642fa8bfda77e49ea8ec1dcb468 Mon Sep 17 00:00:00 2001 From: Liam Huffman Date: Wed, 30 Apr 2025 13:19:06 -0700 Subject: [PATCH] feat(bigquery): Add support for reservation field in jobs. --- .../cloud/bigquery/CopyJobConfiguration.java | 34 +++++++++++++++++-- .../bigquery/ExtractJobConfiguration.java | 34 +++++++++++++++++-- .../cloud/bigquery/LoadJobConfiguration.java | 33 ++++++++++++++++-- .../cloud/bigquery/QueryJobConfiguration.java | 34 +++++++++++++++++-- .../cloud/bigquery/QueryRequestInfo.java | 9 ++++- .../bigquery/CopyJobConfigurationTest.java | 4 +++ .../bigquery/ExtractJobConfigurationTest.java | 7 ++++ .../bigquery/LoadJobConfigurationTest.java | 5 +++ .../bigquery/QueryJobConfigurationTest.java | 4 +++ .../cloud/bigquery/QueryRequestInfoTest.java | 4 +++ 10 files changed, 158 insertions(+), 10 deletions(-) diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/CopyJobConfiguration.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/CopyJobConfiguration.java index 37955fec0..54e612271 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/CopyJobConfiguration.java +++ b/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 labels; private final Long jobTimeoutMs; + private final String reservation; public static final class Builder extends JobConfiguration.Builder { @@ -58,6 +59,7 @@ public static final class Builder private EncryptionConfiguration destinationEncryptionConfiguration; private Map 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; } diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ExtractJobConfiguration.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ExtractJobConfiguration.java index 29a256e9e..d79959ee0 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ExtractJobConfiguration.java +++ b/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 labels; private final Long jobTimeoutMs; + private final String reservation; public static final class Builder extends JobConfiguration.Builder { @@ -61,6 +62,7 @@ public static final class Builder private Boolean useAvroLogicalTypes; private Map 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; } diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/LoadJobConfiguration.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/LoadJobConfiguration.java index 2d348f0d4..5d1755459 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/LoadJobConfiguration.java +++ b/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 connectionProperties; - private final Boolean createSession; + private final String reservation; public static final class Builder extends JobConfiguration.Builder implements LoadConfiguration.Builder { @@ -95,6 +94,7 @@ public static final class Builder extends JobConfiguration.Builder 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; diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryJobConfiguration.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryJobConfiguration.java index 0ad85137b..b460b6121 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryJobConfiguration.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryJobConfiguration.java @@ -74,6 +74,7 @@ public final class QueryJobConfiguration extends JobConfiguration { // maxResults is only used for fast query path private final Long maxResults; private final JobCreationMode jobCreationMode; + private final String reservation; /** * Priority levels for a query. If not specified the priority is assumed to be {@link @@ -142,6 +143,7 @@ public static final class Builder private List connectionProperties; private Long maxResults; private JobCreationMode jobCreationMode; + private String reservation; private Builder() { super(Type.QUERY); @@ -178,6 +180,7 @@ private Builder(QueryJobConfiguration jobConfiguration) { this.connectionProperties = jobConfiguration.connectionProperties; this.maxResults = jobConfiguration.maxResults; this.jobCreationMode = jobConfiguration.jobCreationMode; + this.reservation = jobConfiguration.reservation; } private Builder(com.google.api.services.bigquery.model.JobConfiguration configurationPb) { @@ -285,6 +288,9 @@ private Builder(com.google.api.services.bigquery.model.JobConfiguration configur queryConfigurationPb.getConnectionProperties(), ConnectionProperty.FROM_PB_FUNCTION); } + if (configurationPb.getReservation() != null) { + this.reservation = configurationPb.getReservation(); + } } /** Sets the BigQuery SQL query to execute. */ @@ -682,6 +688,19 @@ Builder setJobCreationMode(JobCreationMode jobCreationMode) { 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 QueryJobConfiguration build() { return new QueryJobConfiguration(this); } @@ -727,6 +746,7 @@ private QueryJobConfiguration(Builder builder) { this.connectionProperties = builder.connectionProperties; this.maxResults = builder.maxResults; this.jobCreationMode = builder.jobCreationMode; + this.reservation = builder.reservation; } /** @@ -943,6 +963,11 @@ JobCreationMode getJobCreationMode() { return jobCreationMode; } + /** Returns the reservation associated with this job */ + public String getReservation() { + return reservation; + } + @Override public Builder toBuilder() { return new Builder(this); @@ -978,7 +1003,8 @@ ToStringHelper toStringHelper() { .add("labels", labels) .add("rangePartitioning", rangePartitioning) .add("connectionProperties", connectionProperties) - .add("jobCreationMode", jobCreationMode); + .add("jobCreationMode", jobCreationMode) + .add("reservation", reservation); } @Override @@ -1016,7 +1042,8 @@ public int hashCode() { jobTimeoutMs, labels, rangePartitioning, - connectionProperties); + connectionProperties, + reservation); } @Override @@ -1125,6 +1152,9 @@ com.google.api.services.bigquery.model.JobConfiguration toPb() { queryConfigurationPb.setConnectionProperties( Lists.transform(connectionProperties, ConnectionProperty.TO_PB_FUNCTION)); } + if (reservation != null) { + configurationPb.setReservation(reservation); + } configurationPb.setQuery(queryConfigurationPb); return configurationPb; } diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java index 1a56872a1..7eebfea96 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java @@ -44,6 +44,7 @@ final class QueryRequestInfo { private final Boolean useLegacySql; private final JobCreationMode jobCreationMode; private final DataFormatOptions formatOptions; + private final String reservation; QueryRequestInfo(QueryJobConfiguration config, Boolean useInt64Timestamps) { this.config = config; @@ -61,6 +62,7 @@ final class QueryRequestInfo { this.useQueryCache = config.useQueryCache(); this.jobCreationMode = config.getJobCreationMode(); this.formatOptions = new DataFormatOptions().setUseInt64Timestamp(useInt64Timestamps); + this.reservation = config.getReservation(); } boolean isFastQuerySupported(JobId jobId) { @@ -128,6 +130,9 @@ QueryRequest toPb() { if (formatOptions != null) { request.setFormatOptions(formatOptions); } + if (reservation != null) { + request.setReservation(reservation); + } return request; } @@ -148,6 +153,7 @@ public String toString() { .add("useLegacySql", useLegacySql) .add("jobCreationMode", jobCreationMode) .add("formatOptions", formatOptions.getUseInt64Timestamp()) + .add("reservation", reservation) .toString(); } @@ -167,7 +173,8 @@ public int hashCode() { useQueryCache, useLegacySql, jobCreationMode, - formatOptions); + formatOptions, + reservation); } @Override diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/CopyJobConfigurationTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/CopyJobConfigurationTest.java index 70a4983df..3f21bf1c0 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/CopyJobConfigurationTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/CopyJobConfigurationTest.java @@ -44,6 +44,7 @@ public class CopyJobConfigurationTest { EncryptionConfiguration.newBuilder().setKmsKeyName("KMS_KEY_1").build(); private static final Map LABELS = ImmutableMap.of("job-name", "copy"); private static final Long TIMEOUT = 10L; + private static final String RESERVATION = "reservation"; private static final CopyJobConfiguration COPY_JOB_CONFIGURATION = CopyJobConfiguration.newBuilder(DESTINATION_TABLE, SOURCE_TABLE) .setCreateDisposition(CREATE_DISPOSITION) @@ -51,6 +52,7 @@ public class CopyJobConfigurationTest { .setDestinationEncryptionConfiguration(COPY_JOB_ENCRYPTION_CONFIGURATION) .setLabels(LABELS) .setJobTimeoutMs(TIMEOUT) + .setReservation(RESERVATION) .build(); private static final CopyJobConfiguration COPY_JOB_CONFIGURATION_MULTIPLE_TABLES = CopyJobConfiguration.newBuilder(DESTINATION_TABLE, SOURCE_TABLES) @@ -58,6 +60,7 @@ public class CopyJobConfigurationTest { .setWriteDisposition(WRITE_DISPOSITION) .setLabels(LABELS) .setJobTimeoutMs(TIMEOUT) + .setReservation(RESERVATION) .build(); @Test @@ -181,5 +184,6 @@ private void compareCopyJobConfiguration( value.getDestinationEncryptionConfiguration()); assertEquals(expected.getLabels(), value.getLabels()); assertEquals(expected.getJobTimeoutMs(), value.getJobTimeoutMs()); + assertEquals(expected.getReservation(), value.getReservation()); } } diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ExtractJobConfigurationTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ExtractJobConfigurationTest.java index ce9e35ddd..2bf1e80a2 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ExtractJobConfigurationTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ExtractJobConfigurationTest.java @@ -43,6 +43,7 @@ public class ExtractJobConfigurationTest { private static final Map LABELS = ImmutableMap.of("test-job-name", "test-extract-job"); private static final Long TIMEOUT = 10L; + private static final String RESERVATION = "reservation"; private static final ExtractJobConfiguration EXTRACT_CONFIGURATION = ExtractJobConfiguration.newBuilder(TABLE_ID, DESTINATION_URIS) .setPrintHeader(PRINT_HEADER) @@ -51,6 +52,7 @@ public class ExtractJobConfigurationTest { .setFormat(FORMAT) .setLabels(LABELS) .setJobTimeoutMs(TIMEOUT) + .setReservation(RESERVATION) .build(); private static final ExtractJobConfiguration EXTRACT_CONFIGURATION_ONE_URI = ExtractJobConfiguration.newBuilder(TABLE_ID, DESTINATION_URI) @@ -60,6 +62,7 @@ public class ExtractJobConfigurationTest { .setFormat(FORMAT) .setLabels(LABELS) .setJobTimeoutMs(TIMEOUT) + .setReservation(RESERVATION) .build(); private static final ExtractJobConfiguration EXTRACT_CONFIGURATION_AVRO = ExtractJobConfiguration.newBuilder(TABLE_ID, DESTINATION_URI) @@ -70,6 +73,7 @@ public class ExtractJobConfigurationTest { .setUseAvroLogicalTypes(USEAVROLOGICALTYPES) .setLabels(LABELS) .setJobTimeoutMs(TIMEOUT) + .setReservation(RESERVATION) .build(); private static final ExtractJobConfiguration EXTRACT_CONFIGURATION_MODEL = ExtractJobConfiguration.newBuilder(MODEL_ID, DESTINATION_URIS) @@ -80,6 +84,7 @@ public class ExtractJobConfigurationTest { .setUseAvroLogicalTypes(USEAVROLOGICALTYPES) .setLabels(LABELS) .setJobTimeoutMs(TIMEOUT) + .setReservation(RESERVATION) .build(); @Test @@ -185,6 +190,7 @@ public void testBuilder() { assertEquals(FORMAT, EXTRACT_CONFIGURATION_MODEL.getFormat()); assertEquals(LABELS, EXTRACT_CONFIGURATION_MODEL.getLabels()); assertEquals(TIMEOUT, EXTRACT_CONFIGURATION_MODEL.getJobTimeoutMs()); + assertEquals(RESERVATION, EXTRACT_CONFIGURATION_MODEL.getReservation()); } @Test @@ -256,5 +262,6 @@ private void compareExtractJobConfiguration( assertEquals(expected.getFormat(), value.getFormat()); assertEquals(expected.getLabels(), value.getLabels()); assertEquals(expected.getJobTimeoutMs(), value.getJobTimeoutMs()); + assertEquals(expected.getReservation(), value.getReservation()); } } diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/LoadJobConfigurationTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/LoadJobConfigurationTest.java index 92a617187..6d181a328 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/LoadJobConfigurationTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/LoadJobConfigurationTest.java @@ -67,6 +67,7 @@ public class LoadJobConfigurationTest { private static final Map LABELS = ImmutableMap.of("test-job-name", "test-load-job"); private static final Long TIMEOUT = 10L; + private static final String RESERVATION = "reservation"; private static final RangePartitioning.Range RANGE = RangePartitioning.Range.newBuilder().setStart(1L).setInterval(2L).setEnd(10L).build(); private static final RangePartitioning RANGE_PARTITIONING = @@ -108,6 +109,7 @@ public class LoadJobConfigurationTest { .setHivePartitioningOptions(HIVE_PARTITIONING_OPTIONS) .setConnectionProperties(CONNECTION_PROPERTIES) .setCreateSession(CREATE_SESSION) + .setReservation(RESERVATION) .build(); private static final DatastoreBackupOptions BACKUP_OPTIONS = @@ -127,6 +129,7 @@ public class LoadJobConfigurationTest { .setLabels(LABELS) .setJobTimeoutMs(TIMEOUT) .setRangePartitioning(RANGE_PARTITIONING) + .setReservation(RESERVATION) .build(); private static final LoadJobConfiguration LOAD_CONFIGURATION_AVRO = LoadJobConfiguration.newBuilder(TABLE_ID, SOURCE_URIS) @@ -145,6 +148,7 @@ public class LoadJobConfigurationTest { .setLabels(LABELS) .setJobTimeoutMs(TIMEOUT) .setRangePartitioning(RANGE_PARTITIONING) + .setReservation(RESERVATION) .build(); @Test @@ -266,5 +270,6 @@ private void compareLoadJobConfiguration( assertEquals(expected.getHivePartitioningOptions(), value.getHivePartitioningOptions()); assertEquals(expected.getConnectionProperties(), value.getConnectionProperties()); assertEquals(expected.getCreateSession(), value.getCreateSession()); + assertEquals(expected.getReservation(), value.getReservation()); } } diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryJobConfigurationTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryJobConfigurationTest.java index 9dd393004..f25aa47ed 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryJobConfigurationTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryJobConfigurationTest.java @@ -112,6 +112,7 @@ public class QueryJobConfigurationTest { ImmutableMap.of("string", STRING_PARAMETER, "timestamp", TIMESTAMP_PARAMETER); private static final String PARAMETER_MODE = "POSITIONAL"; private static final JobCreationMode JOB_CREATION_MODE = JobCreationMode.JOB_CREATION_OPTIONAL; + private static final String RESERVATION = "reservation"; private static final QueryJobConfiguration QUERY_JOB_CONFIGURATION = QueryJobConfiguration.newBuilder(QUERY) .setUseQueryCache(USE_QUERY_CACHE) @@ -139,6 +140,7 @@ public class QueryJobConfigurationTest { .setConnectionProperties(CONNECTION_PROPERTIES) .setPositionalParameters(POSITIONAL_PARAMETER) .setParameterMode(PARAMETER_MODE) + .setReservation(RESERVATION) .build(); private static final QueryJobConfiguration QUERY_JOB_CONFIGURATION_ADD_POSITIONAL_PARAMETER = QUERY_JOB_CONFIGURATION.toBuilder() @@ -188,6 +190,7 @@ public void testToPbAndFromPb() { assertNotNull(QUERY_JOB_CONFIGURATION.getConnectionProperties()); assertNotNull(QUERY_JOB_CONFIGURATION.getPositionalParameters()); assertNotNull(QUERY_JOB_CONFIGURATION.getNamedParameters()); + assertNotNull(QUERY_JOB_CONFIGURATION.getReservation()); compareQueryJobConfiguration( QUERY_JOB_CONFIGURATION, QueryJobConfiguration.fromPb(QUERY_JOB_CONFIGURATION.toPb())); QueryJobConfiguration job = QueryJobConfiguration.of(QUERY); @@ -271,5 +274,6 @@ private void compareQueryJobConfiguration( assertEquals(expected.getConnectionProperties(), value.getConnectionProperties()); assertEquals(expected.getPositionalParameters(), value.getPositionalParameters()); assertEquals(expected.getNamedParameters(), value.getNamedParameters()); + assertEquals(expected.getReservation(), value.getReservation()); } } diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryRequestInfoTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryRequestInfoTest.java index 420c01484..ed9effe0b 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryRequestInfoTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryRequestInfoTest.java @@ -110,6 +110,7 @@ public class QueryRequestInfoTest { ImmutableMap.of("string", STRING_PARAMETER, "timestamp", TIMESTAMP_PARAMETER); private static final JobCreationMode jobCreationModeRequired = JobCreationMode.JOB_CREATION_REQUIRED; + private static final String RESERVATION = "reservation"; private static final QueryJobConfiguration QUERY_JOB_CONFIGURATION = QueryJobConfiguration.newBuilder(QUERY) .setUseQueryCache(USE_QUERY_CACHE) @@ -137,6 +138,7 @@ public class QueryRequestInfoTest { .setPositionalParameters(POSITIONAL_PARAMETER) .setMaxResults(100L) .setJobCreationMode(jobCreationModeRequired) + .setReservation(RESERVATION) .build(); QueryRequestInfo REQUEST_INFO = new QueryRequestInfo(QUERY_JOB_CONFIGURATION, false); private static final QueryJobConfiguration QUERY_JOB_CONFIGURATION_SUPPORTED = @@ -151,6 +153,7 @@ public class QueryRequestInfoTest { .setPositionalParameters(POSITIONAL_PARAMETER) .setCreateSession(CREATE_SESSION) .setMaxResults(100L) + .setReservation(RESERVATION) .build(); QueryRequestInfo REQUEST_INFO_SUPPORTED = new QueryRequestInfo(QUERY_JOB_CONFIGURATION_SUPPORTED, false); @@ -214,5 +217,6 @@ private void compareQueryRequestInfo(QueryRequestInfo expected, QueryRequestInfo assertEquals(expectedQueryReq.getUseLegacySql(), actualQueryReq.getUseLegacySql()); assertEquals(expectedQueryReq.get("jobCreationMode"), actualQueryReq.get("jobCreationMode")); assertEquals(expectedQueryReq.getFormatOptions(), actualQueryReq.getFormatOptions()); + assertEquals(expectedQueryReq.getReservation(), actualQueryReq.getReservation()); } }