feat: add support for parameterMode in QueryJobConfiguration to unblock JDBC migration to the Java client library (#1223)

stephaniewang526 · web-flow · commit 3ce49334478d · 2021-04-09T17:48:05.000-04:00
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
@@ -47,6 +47,7 @@ public final class QueryJobConfiguration extends JobConfiguration {
   private final String query;
   private final ImmutableList<QueryParameterValue> positionalParameters;
   private final ImmutableMap<String, QueryParameterValue> namedParameters;
+  private final String parameterMode;
   private final TableId destinationTable;
   private final Map<String, ExternalTableDefinition> tableDefinitions;
   private final List<UserDefinedFunction> userDefinedFunctions;
@@ -98,6 +99,7 @@ public static final class Builder
     private String query;
     private List<QueryParameterValue> positionalParameters = Lists.newArrayList();
     private Map<String, QueryParameterValue> namedParameters = Maps.newHashMap();
+    private String parameterMode;
     private TableId destinationTable;
     private Map<String, ExternalTableDefinition> tableDefinitions;
     private List<UserDefinedFunction> userDefinedFunctions;
@@ -131,6 +133,7 @@ private Builder(QueryJobConfiguration jobConfiguration) {
       this.query = jobConfiguration.query;
       this.namedParameters = jobConfiguration.namedParameters;
       this.positionalParameters = jobConfiguration.positionalParameters;
+      this.parameterMode = jobConfiguration.parameterMode;
       this.destinationTable = jobConfiguration.destinationTable;
       this.tableDefinitions = jobConfiguration.tableDefinitions;
       this.userDefinedFunctions = jobConfiguration.userDefinedFunctions;
@@ -163,11 +166,13 @@ private Builder(com.google.api.services.bigquery.model.JobConfiguration configur
       if (queryConfigurationPb.getQueryParameters() != null
           && !queryConfigurationPb.getQueryParameters().isEmpty()) {
         if (queryConfigurationPb.getQueryParameters().get(0).getName() == null) {
+          parameterMode = "POSITIONAL";
           setPositionalParameters(
               Lists.transform(
                   queryConfigurationPb.getQueryParameters(),
                   POSITIONAL_PARAMETER_FROM_PB_FUNCTION));
         } else {
+          parameterMode = "NAMED";
           Map<String, QueryParameterValue> values = Maps.newHashMap();
           for (QueryParameter queryParameterPb : queryConfigurationPb.getQueryParameters()) {
             checkNotNull(queryParameterPb.getName());
@@ -277,6 +282,16 @@ public Builder addPositionalParameter(QueryParameterValue value) {
       return this;
     }
 
+    /**
+     * Standard SQL only. Set to POSITIONAL to use positional (?) query parameters or to NAMED to
+     * use named (@myparam) query parameters in this query.
+     */
+    public Builder setParameterMode(String parameterMode) {
+      checkNotNull(parameterMode);
+      this.parameterMode = parameterMode;
+      return this;
+    }
+
     /**
      * Sets the query parameters to a list of positional query parameters to use in the query.
      *
@@ -639,6 +654,7 @@ private QueryJobConfiguration(Builder builder) {
     }
     positionalParameters = ImmutableList.copyOf(builder.positionalParameters);
     namedParameters = ImmutableMap.copyOf(builder.namedParameters);
+    this.parameterMode = builder.parameterMode;
     this.allowLargeResults = builder.allowLargeResults;
     this.createDisposition = builder.createDisposition;
     this.defaultDataset = builder.defaultDataset;
@@ -876,6 +892,7 @@ ToStringHelper toStringHelper() {
         .add("query", query)
         .add("positionalParameters", positionalParameters)
         .add("namedParameters", namedParameters)
+        .add("parameterMode", parameterMode)
         .add("destinationTable", destinationTable)
         .add("destinationEncryptionConfiguration", destinationEncryptionConfiguration)
         .add("defaultDataset", defaultDataset)
@@ -919,6 +936,7 @@ public int hashCode() {
         query,
         positionalParameters,
         namedParameters,
+        parameterMode,
         tableDefinitions,
         useQueryCache,
         userDefinedFunctions,
@@ -963,6 +981,9 @@ com.google.api.services.bigquery.model.JobConfiguration toPb() {
           Lists.transform(namedParameters.entrySet().asList(), NAMED_PARAMETER_TO_PB_FUNCTION);
       queryConfigurationPb.setQueryParameters(queryParametersPb);
     }
+    if (parameterMode != null) {
+      queryConfigurationPb.setParameterMode(parameterMode);
+    }
     configurationPb.setDryRun(dryRun());
     if (allowLargeResults != null) {
       queryConfigurationPb.setAllowLargeResults(allowLargeResults);
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
@@ -108,6 +108,7 @@ public class QueryJobConfigurationTest {
       ImmutableList.of(STRING_PARAMETER, TIMESTAMP_PARAMETER, BIGNUMERIC_PARAMETER);
   private static final Map<String, QueryParameterValue> NAME_PARAMETER =
       ImmutableMap.of("string", STRING_PARAMETER, "timestamp", TIMESTAMP_PARAMETER);
+  private static final String PARAMETER_MODE = "POSITIONAL";
   private static final QueryJobConfiguration QUERY_JOB_CONFIGURATION =
       QueryJobConfiguration.newBuilder(QUERY)
           .setUseQueryCache(USE_QUERY_CACHE)
@@ -133,6 +134,7 @@ public class QueryJobConfigurationTest {
           .setRangePartitioning(RANGE_PARTITIONING)
           .setConnectionProperties(CONNECTION_PROPERTIES)
           .setPositionalParameters(POSITIONAL_PARAMETER)
+          .setParameterMode(PARAMETER_MODE)
           .build();
   private static final QueryJobConfiguration QUERY_JOB_CONFIGURATION_ADD_POSITIONAL_PARAMETER =
       QUERY_JOB_CONFIGURATION
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
@@ -2099,6 +2099,33 @@ public void testScriptStatistics() throws InterruptedException {
     }
   }
 
+  @Test
+  public void testQueryParameterModeWithDryRun() {
+    String query =
+        "SELECT TimestampField, StringField, BooleanField, BigNumericField, BigNumericField1, BigNumericField2, BigNumericField3, BigNumericField4 FROM "
+            + TABLE_ID.getTable()
+            + " WHERE StringField = ?"
+            + " AND TimestampField > ?"
+            + " AND IntegerField IN UNNEST(?)"
+            + " AND IntegerField < ?"
+            + " AND FloatField > ?"
+            + " AND NumericField < ?"
+            + " AND BigNumericField = ?";
+
+    QueryJobConfiguration queryConfig =
+        QueryJobConfiguration.newBuilder(query)
+            .setDefaultDataset(DatasetId.of(DATASET))
+            .setParameterMode("POSITIONAL")
+            .setUseLegacySql(false)
+            .setDryRun(true)
+            .build();
+
+    Job job = bigquery.create(JobInfo.of(queryConfig));
+    JobStatistics.QueryStatistics statistics = job.getStatistics();
+
+    assertNotNull(statistics.getTotalBytesProcessed());
+  }
+
   @Test
   public void testPositionalQueryParameters() throws InterruptedException {
     String query =
