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

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

@@ -28,6 +28,7 @@ class BigQueryJdbcParameter {
   private String paramName;
   private BigQueryStatementParameterType paramType;
   private int scale;
+  private boolean isUserSet = false;
 
   BigQueryJdbcParameter() {}
 
@@ -36,6 +37,7 @@ class BigQueryJdbcParameter {
     this.value = parameter.value;
     this.type = parameter.type;
     this.sqlType = parameter.sqlType;
+    this.isUserSet = parameter.isUserSet;
   }
 
   int getIndex() {
@@ -94,6 +96,14 @@ void setScale(int scale) {
     this.scale = scale;
   }
 
+  boolean isUserSet() {
+    return isUserSet;
+  }
+
+  void setUserSet(boolean userSet) {
+    isUserSet = userSet;
+  }
+
   @Override
   public String toString() {
     return "BigQueryJdbcParameter{"
@@ -112,6 +122,8 @@ public String toString() {
         + paramType.name()
         + ", scale="
         + scale
+        + ", isUserSet="
+        + isUserSet
         + '}';
   }
 }

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

@@ -64,24 +64,24 @@ enum BigQueryStatementParameterType {
   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, this.enableTimestampPicos);
-        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++) {
+
+      int arrayIndex = i - 1;
+      if (this.parametersList.size() <= arrayIndex
+          || this.parametersList.get(arrayIndex) == null
+          || !this.parametersList.get(arrayIndex).isUserSet()) {
+        throw new BigQueryJdbcException("One or more parameters missing in Prepared statement.");
       }
+
+      Object parameterValue = getParameter(i);
+      StandardSQLTypeName sqlType = getSqlType(i);
+      parameterValue =
+          formatValueForQueryParameter(parameterValue, sqlType, this.enableTimestampPicos);
+      LOG.finest(
+          "Parameter %s of type %s at index %s added to QueryJobConfiguration",
+          parameterValue, sqlType, i);
+      jobConfigurationBuilder.addPositionalParameter(
+          QueryParameterValue.of(parameterValue, sqlType));
     }
     return jobConfigurationBuilder;
   }
@@ -132,13 +132,11 @@ private static String formatTimestampParameter(Timestamp ts, boolean enableTimes
     return copy.toString();
   }
 
-  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) {
@@ -147,19 +145,37 @@ 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);
     parameter.setSqlType(BigQueryTypeRegistry.toBigQueryType(type));
     parameter.setParamName("");
     parameter.setParamType(BigQueryStatementParameterType.UNSPECIFIED);
     parameter.setScale(-1);
+    parameter.setUserSet(true);
 
     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 =
@@ -198,7 +214,12 @@ StandardSQLTypeName getSqlType(int index) {
 
   void clearParameters() {
     LOG.finest("++enter++");
-    parametersList.clear();
+    for (BigQueryJdbcParameter param : this.parametersList) {
+      if (param != null) {
+        param.setValue(null);
+        param.setUserSet(false);
+      }
+    }
     highestIndex = 0;
   }
 
@@ -236,6 +257,8 @@ void setParameter(
     parameter.setParamName(paramName);
     parameter.setParamType(paramType);
     parameter.setScale(scale);
+    parameter.setUserSet(true);
+
     if (parameter.getIndex() == -1) {
       parametersList.add(parameter);
     }
@@ -253,16 +276,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);
@@ -271,6 +286,7 @@ void setParameter(
     parameter.setParamName("");
     parameter.setParamType(paramType);
     parameter.setScale(scale);
+    parameter.setUserSet(true);
 
     LOG.finest("Parameter set { %s }", parameter.toString());
   }

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;
@@ -86,6 +87,33 @@ class BigQueryPreparedStatement extends BigQueryStatement implements PreparedSta
     setCurrentQuery(query);
     this.parameterHandler =
         new BigQueryParameterHandler(this.parameterCount, this.isEnableTimestampPicos());
+    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;
@@ -1915,4 +1916,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

@@ -17,6 +17,7 @@
 package com.google.cloud.bigquery.jdbc;
 
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
@@ -31,6 +32,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;
@@ -61,6 +63,7 @@ public class BigQueryPreparedStatementSettersTest {
   public void setUp() throws Exception {
     connection = mock(BigQueryConnection.class);
     when(connection.getQueryDialect()).thenReturn("SQL");
+    when(connection.getConnectionId()).thenReturn("test-connection-id");
     preparedStatement = new BigQueryPreparedStatement(connection, "SELECT ?, ?, ?, ?, ?");
   }
 
@@ -327,6 +330,35 @@ public void testCreateJsonRowWithSetObjectNull() throws Exception {
     assertEquals("42", jsonRow.get("col2").getAsString());
   }
 
+  @Test
+  public void testInferredParameterTypeKnownBeforeSetters() throws Exception {
+    preparedStatement = new BigQueryPreparedStatement(connection, "SELECT ?");
+
+    // 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. configureParameters fails before value is supplied
+    QueryJobConfiguration.Builder configBuilder = QueryJobConfiguration.newBuilder("SELECT ?");
+    BigQueryJdbcException ex =
+        assertThrows(
+            BigQueryJdbcException.class,
+            () -> preparedStatement.parameterHandler.configureParameters(configBuilder));
+    assertTrue(ex.getMessage().contains("One or more parameters missing"));
+    // 3. Once setter is called, configureParameters succeeds and populates QueryJobConfiguration
+
+    preparedStatement.setLong(1, 42L);
+    assertDoesNotThrow(() -> preparedStatement.parameterHandler.configureParameters(configBuilder));
+
+    QueryJobConfiguration config = configBuilder.build();
+    assertEquals(1, config.getPositionalParameters().size());
+    assertEquals("42", config.getPositionalParameters().get(0).getValue());
+    assertEquals(StandardSQLTypeName.INT64, config.getPositionalParameters().get(0).getType());
+  }
+
   @Test
   public void testSetObjectWithTimestampStringAndTypesTimestamp_picosEnabled() throws Exception {
     BigQueryConnection picosConnection = mock(BigQueryConnection.class);

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;