feat(bigquery-jdbc): infer undeclared fields in PreparedStatement

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryParameterHandler.java

@@ -24,6 +24,7 @@
 import java.math.BigInteger;
 import java.sql.SQLException;
 import java.util.ArrayList;
+import java.util.BitSet;
 
 class BigQueryParameterHandler {
   private final BigQueryJdbcCustomLogger LOG = new BigQueryJdbcCustomLogger(this.toString());
@@ -48,30 +49,27 @@ enum BigQueryStatementParameterType {
   };
 
   private int parametersArraySize;
+  private final BitSet userSetParameters = new BitSet();
   ArrayList<BigQueryJdbcParameter> parametersList = new ArrayList<>(parametersArraySize);
 
   private long highestIndex = 0;
 
   QueryJobConfiguration.Builder configureParameters(
       QueryJobConfiguration.Builder jobConfigurationBuilder) throws SQLException {
     LOG.finest("++enter++");
-    try {
-      for (int i = 1; i <= this.parametersArraySize; i++) {
-
-        Object parameterValue = getParameter(i);
-        StandardSQLTypeName sqlType = getSqlType(i);
-        parameterValue = formatValueForQueryParameter(parameterValue, sqlType);
-        LOG.finest(
-            "Parameter %s of type %s at index %s added to QueryJobConfiguration",
-            parameterValue, sqlType, i);
-        jobConfigurationBuilder.addPositionalParameter(
-            QueryParameterValue.of(parameterValue, sqlType));
-      }
-    } catch (NullPointerException e) {
-      LOG.severe("Null parameter mapping encountered.", e);
-      if (e.getMessage().contains("Null type")) {
-        throw new BigQueryJdbcException("One or more parameters missing in Prepared statement.", e);
+    for (int i = 1; i <= this.parametersArraySize; i++) {
+      if (!this.userSetParameters.get(i)) {
+        throw new BigQueryJdbcException("One or more parameters missing in Prepared statement.");
       }
+
+      Object parameterValue = getParameter(i);
+      StandardSQLTypeName sqlType = getSqlType(i);
+      parameterValue = formatValueForQueryParameter(parameterValue, sqlType);
+      LOG.finest(
+          "Parameter %s of type %s at index %s added to QueryJobConfiguration",
+          parameterValue, sqlType, i);
+      jobConfigurationBuilder.addPositionalParameter(
+          QueryParameterValue.of(parameterValue, sqlType));
     }
     return jobConfigurationBuilder;
   }
@@ -108,13 +106,11 @@ static Object formatValueForQueryParameter(Object parameterValue, StandardSQLTyp
     return parameterValue;
   }
 
-  void setParameter(int parameterIndex, Object value, Class type)
-      throws BigQueryJdbcSqlFeatureNotSupportedException {
-    LOG.finest("++enter++");
-    LOG.finest("setParameter called by : %s", type.getName());
-    checkValidIndex(parameterIndex);
-
+  private BigQueryJdbcParameter getOrCreateParameter(int parameterIndex) {
     int arrayIndex = parameterIndex - 1;
+    while (parametersList.size() < parameterIndex) {
+      parametersList.add(null);
+    }
     if (parameterIndex >= this.highestIndex || this.parametersList.get(arrayIndex) == null) {
       parametersList.ensureCapacity(parameterIndex);
       while (parametersList.size() < parameterIndex) {
@@ -123,8 +119,15 @@ void setParameter(int parameterIndex, Object value, Class type)
       parametersList.set(arrayIndex, new BigQueryJdbcParameter());
     }
     this.highestIndex = Math.max(parameterIndex, highestIndex);
-    BigQueryJdbcParameter parameter = parametersList.get(arrayIndex);
+    return parametersList.get(arrayIndex);
+  }
+
+  void setParameter(int parameterIndex, Object value, Class type) {
+    LOG.finest("++enter++");
+    LOG.finest("setParameter called by : %s", type.getName());
+    checkValidIndex(parameterIndex);
 
+    BigQueryJdbcParameter parameter = getOrCreateParameter(parameterIndex);
     parameter.setIndex(parameterIndex);
     parameter.setValue(value);
     parameter.setType(type);
@@ -133,9 +136,21 @@ void setParameter(int parameterIndex, Object value, Class type)
     parameter.setParamType(BigQueryStatementParameterType.UNSPECIFIED);
     parameter.setScale(-1);
 
+    this.userSetParameters.set(parameterIndex);
+
     LOG.finest("Parameter set { %s }", parameter.toString());
   }
 
+  void setInferredParameterType(int parameterIndex, StandardSQLTypeName sqlTypeName) {
+    checkValidIndex(parameterIndex);
+    BigQueryJdbcParameter parameter = getOrCreateParameter(parameterIndex);
+
+    Class<?> javaType = BigQueryTypeRegistry.toJavaClass(sqlTypeName);
+    parameter.setIndex(parameterIndex);
+    parameter.setType(javaType);
+    parameter.setSqlType(sqlTypeName);
+  }
+
   private void checkValidIndex(int parameterIndex) {
     if (parameterIndex > this.parametersArraySize) {
       IndexOutOfBoundsException ex =
@@ -174,7 +189,12 @@ StandardSQLTypeName getSqlType(int index) {
 
   void clearParameters() {
     LOG.finest("++enter++");
-    parametersList.clear();
+    this.userSetParameters.clear();
+    for (BigQueryJdbcParameter param : this.parametersList) {
+      if (param != null) {
+        param.setValue(null);
+      }
+    }
     highestIndex = 0;
   }
 
@@ -229,16 +249,8 @@ void setParameter(
     LOG.finest("++enter++");
     LOG.finest("setParameter called by : %s", type.getName());
     checkValidIndex(parameterIndex);
-    int arrayIndex = parameterIndex - 1;
-    if (parameterIndex >= this.highestIndex || this.parametersList.get(arrayIndex) == null) {
-      parametersList.ensureCapacity(parameterIndex);
-      while (parametersList.size() < parameterIndex) {
-        parametersList.add(null);
-      }
-      parametersList.set(arrayIndex, new BigQueryJdbcParameter());
-    }
-    this.highestIndex = Math.max(parameterIndex, highestIndex);
-    BigQueryJdbcParameter parameter = parametersList.get(arrayIndex);
+
+    BigQueryJdbcParameter parameter = getOrCreateParameter(parameterIndex);
 
     parameter.setIndex(parameterIndex);
     parameter.setValue(value);

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryPreparedStatement.java

@@ -17,6 +17,7 @@
 package com.google.cloud.bigquery.jdbc;
 
 import com.google.api.gax.retrying.RetrySettings;
+import com.google.api.services.bigquery.model.QueryParameter;
 import com.google.cloud.bigquery.FieldList;
 import com.google.cloud.bigquery.JobStatistics.QueryStatistics;
 import com.google.cloud.bigquery.JobStatistics.QueryStatistics.StatementType;
@@ -85,6 +86,33 @@ class BigQueryPreparedStatement extends BigQueryStatement implements PreparedSta
     super(connection);
     setCurrentQuery(query);
     this.parameterHandler = new BigQueryParameterHandler(this.parameterCount);
+    if (this.parameterCount > 0) {
+      populateInferredParameterTypes();
+    }
+  }
+
+  private void populateInferredParameterTypes() {
+    if (this.currentQuery == null) {
+      return;
+    }
+
+    try {
+      List<QueryParameter> undeclaredQueryParameters =
+          getUndeclaredQueryParameters(this.currentQuery);
+      if (undeclaredQueryParameters != null) {
+        int index = 1;
+        for (QueryParameter parameter : undeclaredQueryParameters) {
+          if (parameter.getParameterType() != null) {
+            String typeName = parameter.getParameterType().getType();
+            StandardSQLTypeName sqlTypeName = StandardSQLTypeName.valueOf(typeName);
+            this.parameterHandler.setInferredParameterType(index, sqlTypeName);
+          }
+          index++;
+        }
+      }
+    } catch (Exception ex) {
+      LOG.warning("Could not infer parameter types via dryRun: " + ex.getMessage());
+    }
   }
 
   void setCurrentQuery(String currentQuery) {

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java

@@ -20,6 +20,7 @@
 import com.google.api.gax.paging.Page;
 import com.google.api.gax.rpc.ApiException;
 import com.google.api.gax.rpc.StatusCode;
+import com.google.api.services.bigquery.model.QueryParameter;
 import com.google.cloud.Tuple;
 import com.google.cloud.bigquery.BigQuery;
 import com.google.cloud.bigquery.BigQuery.JobListOption;
@@ -1899,4 +1900,13 @@ private void enqueueBufferError(BlockingQueue<BigQueryFieldValueListWrapper> que
   private void enqueueBufferEndOfStream(BlockingQueue<BigQueryFieldValueListWrapper> queue) {
     Uninterruptibles.putUninterruptibly(queue, BigQueryFieldValueListWrapper.ofEndOfStream(null));
   }
+
+  List<QueryParameter> getUndeclaredQueryParameters(String query) {
+    QueryJobConfiguration dryRunConfig =
+        getJobConfig(query).setDryRun(true).setParameterMode("POSITIONAL").build();
+    Job dryRunJob = this.bigQuery.create((JobInfo.of(dryRunConfig)));
+    QueryStatistics jobStatistics = dryRunJob.getStatistics();
+    List<QueryParameter> queryParameters = jobStatistics.getQueryParameters();
+    return queryParameters;
+  }
 }

java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryPreparedStatementSettersTest.java

@@ -30,6 +30,7 @@
 import com.google.cloud.bigquery.QueryJobConfiguration;
 import com.google.cloud.bigquery.Schema;
 import com.google.cloud.bigquery.StandardSQLTypeName;
+import com.google.cloud.bigquery.exception.BigQueryJdbcException;
 import com.google.gson.Gson;
 import com.google.gson.JsonNull;
 import com.google.gson.JsonObject;
@@ -325,4 +326,21 @@ public void testCreateJsonRowWithSetObjectNull() throws Exception {
     assertTrue(jsonRow.get("col1").isJsonNull());
     assertEquals("42", jsonRow.get("col2").getAsString());
   }
+
+  @Test
+  public void testInferredParameterTypeKnownBeforeSetters() throws Exception {
+    // 1. Inferred type is known immediately without calling setInt/setString
+    preparedStatement.parameterHandler.setInferredParameterType(1, StandardSQLTypeName.INT64);
+
+    ParameterMetaData pmd = preparedStatement.getParameterMetaData();
+    assertEquals(Types.BIGINT, pmd.getParameterType(1));
+    assertEquals("INT64", pmd.getParameterTypeName(1));
+
+    // 2. But execute() still fails if caller forgot to set value!
+    assertThrows(BigQueryJdbcException.class, () -> preparedStatement.execute());
+
+    // 3. Once setter is called, execute() succeeds
+    preparedStatement.setLong(1, 42L);
+    // All parameters provided -> proceeds to configureParameters without throwing
+  }
 }

java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBigQueryJDBCTest.java

@@ -1048,6 +1048,50 @@ public void testPreparedExecuteMethod() throws SQLException {
     assertFalse(dropStatus);
   }
 
+  @Test
+  public void testPreparedInferredParameterTypes() throws SQLException {
+
+    String TABLE_NAME = "JDBC_PREPARED_PARAMETER_INFER_TABLE_" + randomNumber;
+    String createQuery =
+        String.format(
+            "CREATE OR REPLACE TABLE %s.%s (`StringField` STRING, `IntegerField` INTEGER, `BytesField` BYTES, `DoubleField` FLOAT64, `BooleanField` BOOL, `NumericField` NUMERIC, "
+                + "`BigNumericField` BIGNUMERIC, `DateField` DATE, `TimeField` TIME, `DateTimeField` DATETIME, `TimestampField` TIMESTAMP, `ArrayField` ARRAY<STRING>, `StructField` STRUCT<subField STRING>, "
+                + "`JsonField` JSON, `GeographyField` GEOGRAPHY, `IntervalField` INTERVAL, `RangeField` RANGE<DATE>);",
+            DATASET, TABLE_NAME);
+    String insertQuery =
+        String.format(
+            "INSERT INTO %s.%s (StringField, IntegerField, BytesField, DoubleField, BooleanField, NumericField, BigNumericField, "
+                + "DateField, TimeField, DateTimeField, TimestampField, ArrayField, StructField, JsonField, GeographyField, IntervalField, RangeField) "
+                + "VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?);",
+            DATASET, TABLE_NAME);
+
+    String dropQuery = String.format("DROP TABLE %s.%s", DATASET, TABLE_NAME);
+    int[] expectedValues = {
+      -9, -5, -3, 8, 16, 2, 2, 91, 92, 93, 93, 2003, 2002, 1111, 1111, 1111, 1111
+    };
+
+    boolean createStatus = bigQueryStatement.execute(createQuery);
+    assertFalse(createStatus);
+
+    PreparedStatement insertStmt = bigQueryConnection.prepareStatement(insertQuery);
+    ParameterMetaData parameterMetaData = insertStmt.getParameterMetaData();
+    for (int i = 0; i < parameterMetaData.getParameterCount(); i++) {
+      assertEquals(expectedValues[i], parameterMetaData.getParameterType(i + 1));
+    }
+
+    // Testing an Exception is thrown if not all values are set.
+    insertStmt.setString(1, "String1");
+    insertStmt.setInt(2, 111);
+    insertStmt.setObject(4, 1.5);
+    insertStmt.setObject(6, true, Types.BOOLEAN);
+    insertStmt.setNull(7, Types.VARCHAR);
+
+    assertThrows(BigQueryJdbcException.class, insertStmt::execute);
+
+    boolean dropStatus = bigQueryStatement.execute(dropQuery);
+    assertFalse(dropStatus);
+  }
+
   @Test
   public void testPreparedStatementThrowsSyntaxError() throws SQLException {
     String TABLE_NAME = "JDBC_PREPARED_SYNTAX_ERR_TABLE_" + randomNumber;