feat: add support connection properties for job and query (#320)

Author: Praful Makani

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();
+  }
+}

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<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
@@ -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);
@@ -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<ConnectionProperty> 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<ConnectionProperty> 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;
   }

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());
+  }
+}

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<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)
@@ -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());
   }
 }