Skip to content

Commit

Permalink
feat: add support connection properties for job and query (#320)
Browse files Browse the repository at this point in the history
  • Loading branch information
Praful Makani committed May 5, 2020
1 parent 3dade31 commit dcd3663
Show file tree
Hide file tree
Showing 5 changed files with 283 additions and 2 deletions.
@@ -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();
}
}
Expand Up @@ -68,6 +68,7 @@ public final class QueryJobConfiguration extends JobConfiguration {
private final Long jobTimeoutMs;
private final Map<String, String> labels;
private final RangePartitioning rangePartitioning;
private final List<ConnectionProperty> connectionProperties;

/**
* Priority levels for a query. If not specified the priority is assumed to be {@link
Expand Down Expand Up @@ -116,6 +117,7 @@ public static final class Builder
private Long jobTimeoutMs;
private Map<String, String> labels;
private RangePartitioning rangePartitioning;
private List<ConnectionProperty> connectionProperties;

private Builder() {
super(Type.QUERY);
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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. */
Expand Down Expand Up @@ -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<ConnectionProperty> connectionProperties) {
this.connectionProperties = ImmutableList.copyOf(connectionProperties);
return this;
}

public QueryJobConfiguration build() {
return new QueryJobConfiguration(this);
}
Expand Down Expand Up @@ -619,6 +643,7 @@ private QueryJobConfiguration(Builder builder) {
this.jobTimeoutMs = builder.jobTimeoutMs;
this.labels = builder.labels;
this.rangePartitioning = builder.rangePartitioning;
this.connectionProperties = builder.connectionProperties;
}

/**
Expand Down Expand Up @@ -803,6 +828,11 @@ public RangePartitioning getRangePartitioning() {
return rangePartitioning;
}

/** Returns the connection properties for connection string with this job */
public List<ConnectionProperty> getConnectionProperties() {
return connectionProperties;
}

@Override
public Builder toBuilder() {
return new Builder(this);
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -869,7 +900,8 @@ public int hashCode() {
clustering,
jobTimeoutMs,
labels,
rangePartitioning);
rangePartitioning,
connectionProperties);
}

@Override
Expand Down Expand Up @@ -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;
}
Expand Down
@@ -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());
}
}
Expand Up @@ -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<String> 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<ConnectionProperty> CONNECTION_PROPERTIES =
ImmutableList.of(CONNECTION_PROPERTY);
private static final Field FIELD_SCHEMA1 =
Field.newBuilder("StringField", LegacySQLTypeName.STRING)
.setMode(Field.Mode.NULLABLE)
Expand Down Expand Up @@ -112,6 +118,7 @@ public class QueryJobConfigurationTest {
.setJobTimeoutMs(TIMEOUT)
.setLabels(LABELS)
.setRangePartitioning(RANGE_PARTITIONING)
.setConnectionProperties(CONNECTION_PROPERTIES)
.build();

@Test
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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());
}
}

0 comments on commit dcd3663

Please sign in to comment.