diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ConnectionProperty.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ConnectionProperty.java new file mode 100644 index 000000000..a6206712e --- /dev/null +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ConnectionProperty.java @@ -0,0 +1,146 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery; + +import static com.google.common.base.Preconditions.checkNotNull; + +import com.google.common.base.Function; +import com.google.common.base.MoreObjects; +import com.google.common.base.Objects; + +public final class ConnectionProperty { + + static final Function< + com.google.api.services.bigquery.model.ConnectionProperty, ConnectionProperty> + FROM_PB_FUNCTION = + new Function< + com.google.api.services.bigquery.model.ConnectionProperty, ConnectionProperty>() { + @Override + public ConnectionProperty apply( + com.google.api.services.bigquery.model.ConnectionProperty connectionProperty) { + return ConnectionProperty.fromPb(connectionProperty); + } + }; + static final Function< + ConnectionProperty, com.google.api.services.bigquery.model.ConnectionProperty> + TO_PB_FUNCTION = + new Function< + ConnectionProperty, com.google.api.services.bigquery.model.ConnectionProperty>() { + @Override + public com.google.api.services.bigquery.model.ConnectionProperty apply( + ConnectionProperty connectionProperty) { + return connectionProperty.toPb(); + } + }; + + private final String key; + private final String value; + + /** A builder for {@code ConnectionProperty} objects. */ + public static final class Builder { + private String key; + private String value; + + private Builder() {}; + + private Builder(ConnectionProperty properties) { + this.key = properties.key; + this.value = properties.value; + } + + /** [Required] Name of the connection property to set. */ + public Builder setKey(String key) { + this.key = key; + return this; + } + + /** [Required] Value of the connection property. */ + public Builder setValue(String value) { + this.value = value; + return this; + } + + /** Creates a {@code ConnectionProperty} object. */ + public ConnectionProperty build() { + return new ConnectionProperty(this); + } + } + + private ConnectionProperty(Builder builder) { + this.key = checkNotNull(builder.key, "Required key is null or empty"); + this.value = checkNotNull(builder.value, "Required value is null or empty"); + } + + /** Return the key of property. */ + public String getKey() { + return key; + } + + /** Return the value of property. */ + public String getValue() { + return value; + } + + /** Return a connection property for the given key and value. */ + public static ConnectionProperty of(String key, String value) { + return newBuilder().setKey(key).setValue(value).build(); + } + + /** Returns a builder for the {@code ConnectionProperty} object. */ + public static Builder newBuilder() { + return new Builder(); + } + + /** Returns a builder for the {@code ConnectionProperty} object. */ + public Builder toBuilder() { + return new Builder(this); + } + + @Override + public String toString() { + return MoreObjects.toStringHelper(this).add("key", key).add("value", value).toString(); + } + + @Override + public int hashCode() { + return Objects.hashCode(key, value); + } + + @Override + public boolean equals(Object obj) { + return obj == this + || obj != null + && obj.getClass().equals(ConnectionProperty.class) + && java.util.Objects.equals(toPb(), ((ConnectionProperty) obj).toPb()); + } + + com.google.api.services.bigquery.model.ConnectionProperty toPb() { + com.google.api.services.bigquery.model.ConnectionProperty properties = + new com.google.api.services.bigquery.model.ConnectionProperty(); + properties.setKey(key); + properties.setValue(value); + return properties; + } + + static ConnectionProperty fromPb( + com.google.api.services.bigquery.model.ConnectionProperty properties) { + Builder builder = newBuilder(); + builder.setKey(properties.getKey()); + builder.setValue(properties.getValue()); + return builder.build(); + } +} 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 2c1a9a43b..c26d5cf0c 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 @@ -68,6 +68,7 @@ public final class QueryJobConfiguration extends JobConfiguration { private final Long jobTimeoutMs; private final Map labels; private final RangePartitioning rangePartitioning; + private final List connectionProperties; /** * Priority levels for a query. If not specified the priority is assumed to be {@link @@ -116,6 +117,7 @@ public static final class Builder private Long jobTimeoutMs; private Map labels; private RangePartitioning rangePartitioning; + private List connectionProperties; private Builder() { super(Type.QUERY); @@ -147,6 +149,7 @@ private Builder(QueryJobConfiguration jobConfiguration) { this.jobTimeoutMs = jobConfiguration.jobTimeoutMs; this.labels = jobConfiguration.labels; this.rangePartitioning = jobConfiguration.rangePartitioning; + this.connectionProperties = jobConfiguration.connectionProperties; } private Builder(com.google.api.services.bigquery.model.JobConfiguration configurationPb) { @@ -240,6 +243,12 @@ private Builder(com.google.api.services.bigquery.model.JobConfiguration configur this.rangePartitioning = RangePartitioning.fromPb(queryConfigurationPb.getRangePartitioning()); } + if (queryConfigurationPb.getConnectionProperties() != null) { + this.connectionProperties = + Lists.transform( + queryConfigurationPb.getConnectionProperties(), + ConnectionProperty.FROM_PB_FUNCTION); + } } /** Sets the BigQuery SQL query to execute. */ @@ -579,6 +588,21 @@ public Builder setRangePartitioning(RangePartitioning rangePartitioning) { return this; } + /** + * A connection-level property to customize query behavior. Under JDBC, these correspond + * directly to connection properties passed to the DriverManager. Under ODBC, these correspond + * to properties in the connection string. Currently, the only supported connection property is + * "time_zone", whose value represents the default timezone used to run the query. Additional + * properties are allowed, but ignored. Specifying multiple connection properties with the same + * key is an error. + * + * @param connectionProperties connectionProperties or {@code null} for none + */ + public Builder setConnectionProperties(List connectionProperties) { + this.connectionProperties = ImmutableList.copyOf(connectionProperties); + return this; + } + public QueryJobConfiguration build() { return new QueryJobConfiguration(this); } @@ -619,6 +643,7 @@ private QueryJobConfiguration(Builder builder) { this.jobTimeoutMs = builder.jobTimeoutMs; this.labels = builder.labels; this.rangePartitioning = builder.rangePartitioning; + this.connectionProperties = builder.connectionProperties; } /** @@ -803,6 +828,11 @@ public RangePartitioning getRangePartitioning() { return rangePartitioning; } + /** Returns the connection properties for connection string with this job */ + public List getConnectionProperties() { + return connectionProperties; + } + @Override public Builder toBuilder() { return new Builder(this); @@ -834,7 +864,8 @@ ToStringHelper toStringHelper() { .add("clustering", clustering) .add("jobTimeoutMs", jobTimeoutMs) .add("labels", labels) - .add("rangePartitioning", rangePartitioning); + .add("rangePartitioning", rangePartitioning) + .add("connectionProperties", connectionProperties); } @Override @@ -869,7 +900,8 @@ public int hashCode() { clustering, jobTimeoutMs, labels, - rangePartitioning); + rangePartitioning, + connectionProperties); } @Override @@ -968,6 +1000,10 @@ com.google.api.services.bigquery.model.JobConfiguration toPb() { if (rangePartitioning != null) { queryConfigurationPb.setRangePartitioning(rangePartitioning.toPb()); } + if (connectionProperties != null) { + queryConfigurationPb.setConnectionProperties( + Lists.transform(connectionProperties, ConnectionProperty.TO_PB_FUNCTION)); + } configurationPb.setQuery(queryConfigurationPb); return configurationPb; } diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ConnectionPropertyTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ConnectionPropertyTest.java new file mode 100644 index 000000000..9177720e8 --- /dev/null +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ConnectionPropertyTest.java @@ -0,0 +1,64 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery; + +import static com.google.common.truth.Truth.assertThat; + +import org.junit.Test; + +public class ConnectionPropertyTest { + + private static final String KEY = "time_zone"; + private static final String VALUE = "US/Eastern"; + private static final ConnectionProperty CONNECTION_PROPERTY = + ConnectionProperty.newBuilder().setKey(KEY).setValue(VALUE).build(); + + @Test + public void testToBuilder() { + compareConnectionProperty(CONNECTION_PROPERTY, CONNECTION_PROPERTY.toBuilder().build()); + ConnectionProperty property = CONNECTION_PROPERTY.toBuilder().setKey("time-zone").build(); + assertThat(property.getKey()).isEqualTo("time-zone"); + property = CONNECTION_PROPERTY.toBuilder().setKey(KEY).build(); + compareConnectionProperty(CONNECTION_PROPERTY, property); + } + + @Test + public void testToBuilderIncomplete() { + ConnectionProperty connectionProperty = ConnectionProperty.of(KEY, VALUE); + compareConnectionProperty(connectionProperty, connectionProperty.toBuilder().build()); + } + + @Test + public void testBuilder() { + assertThat(CONNECTION_PROPERTY.getKey()).isEqualTo(KEY); + assertThat(CONNECTION_PROPERTY.getValue()).isEqualTo(VALUE); + } + + @Test + public void testToAndFromPb() { + compareConnectionProperty( + CONNECTION_PROPERTY, ConnectionProperty.fromPb(CONNECTION_PROPERTY.toPb())); + } + + private void compareConnectionProperty(ConnectionProperty expected, ConnectionProperty value) { + assertThat(value).isEqualTo(expected); + assertThat(value.getKey()).isEqualTo(expected.getKey()); + assertThat(value.getValue()).isEqualTo(expected.getValue()); + assertThat(value.toString()).isEqualTo(expected.toString()); + assertThat(value.hashCode()).isEqualTo(expected.hashCode()); + } +} 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 b31312a34..406d7edfa 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 @@ -38,6 +38,12 @@ public class QueryJobConfigurationTest { private static final DatasetId DATASET_ID = DatasetId.of("dataset"); private static final TableId TABLE_ID = TableId.of("dataset", "table"); private static final List SOURCE_URIS = ImmutableList.of("uri1", "uri2"); + private static final String KEY = "time_zone"; + private static final String VALUE = "US/Eastern"; + private static final ConnectionProperty CONNECTION_PROPERTY = + ConnectionProperty.newBuilder().setKey(KEY).setValue(VALUE).build(); + private static final List CONNECTION_PROPERTIES = + ImmutableList.of(CONNECTION_PROPERTY); private static final Field FIELD_SCHEMA1 = Field.newBuilder("StringField", LegacySQLTypeName.STRING) .setMode(Field.Mode.NULLABLE) @@ -112,6 +118,7 @@ public class QueryJobConfigurationTest { .setJobTimeoutMs(TIMEOUT) .setLabels(LABELS) .setRangePartitioning(RANGE_PARTITIONING) + .setConnectionProperties(CONNECTION_PROPERTIES) .build(); @Test @@ -146,6 +153,7 @@ public void testToPbAndFromPb() { assertNotNull(QUERY_JOB_CONFIGURATION.getJobTimeoutMs()); assertNotNull(QUERY_JOB_CONFIGURATION.getLabels()); assertNotNull(QUERY_JOB_CONFIGURATION.getRangePartitioning()); + assertNotNull(QUERY_JOB_CONFIGURATION.getConnectionProperties()); compareQueryJobConfiguration( QUERY_JOB_CONFIGURATION, QueryJobConfiguration.fromPb(QUERY_JOB_CONFIGURATION.toPb())); QueryJobConfiguration job = QueryJobConfiguration.of(QUERY); @@ -204,5 +212,6 @@ private void compareQueryJobConfiguration( assertEquals(expected.getJobTimeoutMs(), value.getJobTimeoutMs()); assertEquals(expected.getLabels(), value.getLabels()); assertEquals(expected.getRangePartitioning(), value.getRangePartitioning()); + assertEquals(expected.getConnectionProperties(), value.getConnectionProperties()); } } diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java index 472d8ecf8..165122d99 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java @@ -47,6 +47,7 @@ import com.google.cloud.bigquery.BigQueryError; import com.google.cloud.bigquery.BigQueryException; import com.google.cloud.bigquery.Clustering; +import com.google.cloud.bigquery.ConnectionProperty; import com.google.cloud.bigquery.CopyJobConfiguration; import com.google.cloud.bigquery.Dataset; import com.google.cloud.bigquery.DatasetId; @@ -301,6 +302,12 @@ public class ITBigQueryTest { + " \"GeographyField\": \"POINT(-122.35022 47.649154)\"," + " \"NumericField\": \"123456.789012345\"" + "}"; + private static final String KEY = "time_zone"; + private static final String VALUE = "US/Eastern"; + private static final ConnectionProperty CONNECTION_PROPERTY = + ConnectionProperty.newBuilder().setKey(KEY).setValue(VALUE).build(); + private static final List CONNECTION_PROPERTIES = + ImmutableList.of(CONNECTION_PROPERTY); private static final Set PUBLIC_DATASETS = ImmutableSet.of("github_repos", "hacker_news", "noaa_gsod", "samples", "usa_names"); @@ -1799,6 +1806,25 @@ public void testQueryJob() throws InterruptedException, TimeoutException { assertNotNull(statistics.getQueryPlan()); } + @Test + public void testQueryJobWithConnectionProperties() throws InterruptedException { + String tableName = "test_query_job_table_connection_properties"; + String query = "SELECT TimestampField, StringField, BooleanField FROM " + TABLE_ID.getTable(); + TableId destinationTable = TableId.of(DATASET, tableName); + QueryJobConfiguration configuration = + QueryJobConfiguration.newBuilder(query) + .setDefaultDataset(DatasetId.of(DATASET)) + .setDestinationTable(destinationTable) + .setConnectionProperties(CONNECTION_PROPERTIES) + .build(); + Job remoteJob = bigquery.create(JobInfo.of(configuration)); + remoteJob = remoteJob.waitFor(); + assertNull(remoteJob.getStatus().getError()); + QueryJobConfiguration jobConfiguration = remoteJob.getConfiguration(); + assertEquals(CONNECTION_PROPERTIES, jobConfiguration.getConnectionProperties()); + assertTrue(bigquery.delete(destinationTable)); + } + @Test public void testQueryJobWithLabels() throws InterruptedException, TimeoutException { String tableName = "test_query_job_table";