feat(bigquery-jdbc): support picosecond in PreparedStatement parameters and batching

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

@@ -23,18 +23,32 @@
 import com.google.cloud.bigquery.exception.BigQueryJdbcSqlFeatureNotSupportedException;
 import java.math.BigInteger;
 import java.sql.SQLException;
+import java.sql.Time;
+import java.sql.Timestamp;
 import java.util.ArrayList;
 
 class BigQueryParameterHandler {
   private final BigQueryJdbcCustomLogger LOG = new BigQueryJdbcCustomLogger(this.toString());
+  private final int parametersArraySize;
+  private final boolean enableTimestampPicos;
+  final ArrayList<BigQueryJdbcParameter> parametersList;
+  private long highestIndex = 0;
 
-  public BigQueryParameterHandler(int parameterCount) {
-    this.parametersArraySize = parameterCount;
+  BigQueryParameterHandler(int parameterCount, boolean enableTimestampPicos) {
+    this(parameterCount, new ArrayList<>(parameterCount), enableTimestampPicos);
+  }
+
+  BigQueryParameterHandler(int parameterCount) {
+    this(parameterCount, false);
   }
 
-  BigQueryParameterHandler(int parameterCount, ArrayList<BigQueryJdbcParameter> parametersList) {
+  BigQueryParameterHandler(
+      int parameterCount,
+      ArrayList<BigQueryJdbcParameter> parametersList,
+      boolean enableTimestampPicos) {
     this.parametersArraySize = parameterCount;
     this.parametersList = parametersList;
+    this.enableTimestampPicos = enableTimestampPicos;
   }
 
   // Indicates whether the parameter is input, output or both
@@ -47,11 +61,6 @@ enum BigQueryStatementParameterType {
     INOUT
   };
 
-  private int parametersArraySize;
-  ArrayList<BigQueryJdbcParameter> parametersList = new ArrayList<>(parametersArraySize);
-
-  private long highestIndex = 0;
-
   QueryJobConfiguration.Builder configureParameters(
       QueryJobConfiguration.Builder jobConfigurationBuilder) throws SQLException {
     LOG.finest("++enter++");
@@ -60,7 +69,8 @@ QueryJobConfiguration.Builder configureParameters(
 
         Object parameterValue = getParameter(i);
         StandardSQLTypeName sqlType = getSqlType(i);
-        parameterValue = formatValueForQueryParameter(parameterValue, sqlType);
+        parameterValue =
+            formatValueForQueryParameter(parameterValue, sqlType, this.enableTimestampPicos);
         LOG.finest(
             "Parameter %s of type %s at index %s added to QueryJobConfiguration",
             parameterValue, sqlType, i);
@@ -76,7 +86,8 @@ QueryJobConfiguration.Builder configureParameters(
     return jobConfigurationBuilder;
   }
 
-  static Object formatValueForQueryParameter(Object parameterValue, StandardSQLTypeName sqlType) {
+  static Object formatValueForQueryParameter(
+      Object parameterValue, StandardSQLTypeName sqlType, boolean enableTimestampPicos) {
     if (parameterValue == null) {
       return null;
     }
@@ -89,13 +100,17 @@ static Object formatValueForQueryParameter(Object parameterValue, StandardSQLTyp
     if (sqlType == StandardSQLTypeName.FLOAT64 && parameterValue instanceof Float) {
       return ((Number) parameterValue).doubleValue();
     }
-    if (parameterValue instanceof java.sql.Timestamp) {
-      java.sql.Timestamp ts = (java.sql.Timestamp) parameterValue;
-      java.sql.Timestamp copy = new java.sql.Timestamp(ts.getTime());
-      copy.setNanos((ts.getNanos() / 1000) * 1000);
-      return copy.toString();
+    if (parameterValue instanceof Timestamp) {
+      return formatTimestampParameter((Timestamp) parameterValue, enableTimestampPicos);
+    }
+    if (sqlType == StandardSQLTypeName.TIMESTAMP && parameterValue instanceof String) {
+      String str = ((String) parameterValue).trim();
+      if (str.length() > 10 && str.charAt(10) == 'T') {
+        str = str.substring(0, 10) + ' ' + str.substring(11);
+      }
+      return BigQueryTemporalUtility.truncateFractionalSeconds(str, enableTimestampPicos ? 12 : 6);
     }
-    if (parameterValue instanceof java.sql.Time) {
+    if (parameterValue instanceof Time) {
       String timeStr = parameterValue.toString();
       if (timeStr.length() == 8) {
         return timeStr + ".000000";
@@ -108,6 +123,15 @@ static Object formatValueForQueryParameter(Object parameterValue, StandardSQLTyp
     return parameterValue;
   }
 
+  private static String formatTimestampParameter(Timestamp ts, boolean enableTimestampPicos) {
+    if (enableTimestampPicos) {
+      return ts.toString();
+    }
+    Timestamp copy = new Timestamp(ts.getTime());
+    copy.setNanos((ts.getNanos() / 1000) * 1000);
+    return copy.toString();
+  }
+
   void setParameter(int parameterIndex, Object value, Class type)
       throws BigQueryJdbcSqlFeatureNotSupportedException {
     LOG.finest("++enter++");

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

@@ -84,7 +84,8 @@ class BigQueryPreparedStatement extends BigQueryStatement implements PreparedSta
   BigQueryPreparedStatement(BigQueryConnection connection, String query) {
     super(connection);
     setCurrentQuery(query);
-    this.parameterHandler = new BigQueryParameterHandler(this.parameterCount);
+    this.parameterHandler =
+        new BigQueryParameterHandler(this.parameterCount, this.isEnableTimestampPicos());
   }
 
   void setCurrentQuery(String currentQuery) {
@@ -99,6 +100,7 @@ private int getParameterCount(String query) {
 
   @Override
   public ResultSet executeQuery() throws SQLException {
+    validateExecution();
     return BigQueryJdbcOpenTelemetry.withTracing(
         "BigQueryPreparedStatement.executeQuery",
         this.connection,
@@ -108,6 +110,7 @@ public ResultSet executeQuery() throws SQLException {
 
   @Override
   public long executeLargeUpdate() throws SQLException {
+    validateExecution();
     return BigQueryJdbcOpenTelemetry.withTracing(
         "BigQueryPreparedStatement.executeLargeUpdate",
         this.connection,
@@ -122,6 +125,7 @@ public int executeUpdate() throws SQLException {
 
   @Override
   public boolean execute() throws SQLException {
+    validateExecution();
     return BigQueryJdbcOpenTelemetry.withTracing(
         "BigQueryPreparedStatement.execute",
         this.connection,
@@ -314,6 +318,7 @@ private ArrayList<BigQueryJdbcParameter> deepCopyParameterList(
 
   @Override
   public int[] executeBatch() throws SQLException {
+    validateExecution();
     int[] result = new int[this.batchParameters.size()];
     if (this.batchParameters.isEmpty()) {
       return result;
@@ -463,7 +468,8 @@ QueryJobConfiguration getWriteBatchJobConfiguration(
       ArrayList<BigQueryJdbcParameter> currentParameterList) throws SQLException {
     LOG.finer("++enter++");
     BigQueryParameterHandler batchHandler =
-        new BigQueryParameterHandler(this.parameterCount, currentParameterList);
+        new BigQueryParameterHandler(
+            this.parameterCount, currentParameterList, this.isEnableTimestampPicos());
     QueryJobConfiguration.Builder jobConfiguration = getJobConfig(this.currentQuery);
     jobConfiguration.setParameterMode("POSITIONAL");
     jobConfiguration = batchHandler.configureParameters(jobConfiguration);
@@ -482,7 +488,7 @@ QueryJobConfiguration getStandardBatchJobConfiguration(String query) throws SQLE
       for (BigQueryJdbcParameter parameter : parameterList) {
         Object parameterValue =
             BigQueryParameterHandler.formatValueForQueryParameter(
-                parameter.getValue(), parameter.getSqlType());
+                parameter.getValue(), parameter.getSqlType(), this.isEnableTimestampPicos());
         StandardSQLTypeName sqlType = parameter.getSqlType();
         LOG.finer(
             "Parameter %s of type %s at index %s added to QueryJobConfiguration",

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

@@ -242,7 +242,7 @@ private BigQuerySettings generateBigQuerySettings() {
    */
   @Override
   public ResultSet executeQuery(String sql) throws SQLException {
-    checkClosed();
+    validateExecution();
     return BigQueryJdbcOpenTelemetry.withTracing(
         "BigQueryStatement.executeQuery", this.connection, sql, () -> executeQueryImpl(sql));
   }
@@ -267,7 +267,7 @@ private ResultSet executeQueryImpl(String sql) throws SQLException {
 
   @Override
   public long executeLargeUpdate(String sql) throws SQLException {
-    checkClosed();
+    validateExecution();
     return BigQueryJdbcOpenTelemetry.withTracing(
         "BigQueryStatement.executeLargeUpdate",
         this.connection,
@@ -308,7 +308,7 @@ int checkUpdateCount(long updateCount) {
 
   @Override
   public boolean execute(String sql) throws SQLException {
-    checkClosed();
+    validateExecution();
     return BigQueryJdbcOpenTelemetry.withTracing(
         "BigQueryStatement.execute", this.connection, sql, () -> executeImpl(sql));
   }
@@ -1705,6 +1705,7 @@ public void clearBatch() {
   @Override
   public int[] executeBatch() throws SQLException {
     LOG.finest("++enter++");
+    validateExecution();
     return BigQueryJdbcOpenTelemetry.withTracing(
         "BigQueryStatement.executeBatch",
         this.connection,
@@ -1866,6 +1867,21 @@ void checkClosed() throws SQLException {
     }
   }
 
+  /**
+   * Validates that the statement is open and that execution configuration settings are compatible.
+   *
+   * @throws SQLException if the statement is closed or settings are incompatible
+   */
+  void validateExecution() throws SQLException {
+    checkClosed();
+    if (isEnableTimestampPicos() && getUseLegacySql()) {
+      throw new BigQueryJdbcException(
+          "Picosecond data is incompatible with Legacy SQL. "
+              + "To query your Picosecond data, please set QueryDialect to SQL "
+              + "and restructure your query as a Standard SQL query.");
+    }
+  }
+
   enum SqlType {
     SELECT,
     DML,

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

@@ -59,7 +59,7 @@ private BigQueryTemporalUtility() {}
    */
   public static Timestamp boxDateTime(String val, ZoneId zoneId) {
     ZoneId targetZone = zoneId != null ? zoneId : ZoneId.systemDefault();
-    String isoString = truncateIsoFractionToNanos(val.replace(' ', 'T'));
+    String isoString = truncateFractionalSeconds(val.replace(' ', 'T'), 9);
     return Timestamp.from(LocalDateTime.parse(isoString).atZone(targetZone).toInstant());
   }
 
@@ -80,7 +80,7 @@ public static Date boxDate(String val, ZoneId zoneId) {
    * perfectly accurate modern conversions.
    */
   public static Time boxTime(String val, ZoneId zoneId) {
-    LocalTime localTime = LocalTime.parse(truncateIsoFractionToNanos(val));
+    LocalTime localTime = LocalTime.parse(truncateFractionalSeconds(val, 9));
 
     if (zoneId == null) {
       // JDBC 4.2 Modern API (no Calendar provided):
@@ -139,7 +139,7 @@ public static Timestamp boxTimestamp(String val) {
     }
 
     // Truncate sub-nanosecond fraction (> 9 digits) to prevent Instant.parse failure
-    iso = truncateIsoFractionToNanos(iso);
+    iso = truncateFractionalSeconds(iso, 9);
 
     try {
       return Timestamp.from(Instant.parse(iso));
@@ -344,30 +344,35 @@ private static StringBuilder formatDateTimeBase(LocalDateTime dt, int scale) {
   }
 
   /**
-   * Truncates sub-second fractional digits to at most 9 digits (nanoseconds) so that standard JDK
-   * temporal parsers (which cap at nanosecond precision) can parse the string without throwing
-   * {@link java.time.format.DateTimeParseException}. Any trailing timezone offset or suffix is
-   * preserved intact.
+   * Truncates sub-second fractional digits in a timestamp string to at most {@code maxDigits}. Any
+   * trailing timezone offset or suffix is preserved intact.
    */
-  private static String truncateIsoFractionToNanos(String iso) {
-    int dotIdx = iso.indexOf('.');
+  static String truncateFractionalSeconds(String str, int maxDigits) {
+    if (str == null) {
+      return null;
+    }
+    int dotIdx = str.indexOf('.');
     if (dotIdx < 0) {
-      return iso;
+      return str;
     }
 
     int fractionStart = dotIdx + 1;
     int fractionEnd = fractionStart;
-    while (fractionEnd < iso.length() && Character.isDigit(iso.charAt(fractionEnd))) {
+    int len = str.length();
+    while (fractionEnd < len) {
+      char c = str.charAt(fractionEnd);
+      if (c < '0' || c > '9') {
+        break;
+      }
       fractionEnd++;
     }
 
     int fractionDigits = fractionEnd - fractionStart;
-    if (fractionDigits <= 9) {
-      return iso;
+    if (fractionDigits <= maxDigits) {
+      return str;
     }
 
-    // Retain the first 9 fractional digits and append any trailing suffix (e.g., 'Z' or offset)
-    return iso.substring(0, fractionStart + 9) + iso.substring(fractionEnd);
+    return str.substring(0, fractionStart + maxDigits) + str.substring(fractionEnd);
   }
 
   /**

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

@@ -22,6 +22,7 @@
 import com.google.cloud.bigquery.QueryJobConfiguration;
 import com.google.cloud.bigquery.StandardSQLTypeName;
 import com.google.cloud.bigquery.jdbc.BigQueryParameterHandler.BigQueryStatementParameterType;
+import java.sql.Timestamp;
 import org.junit.jupiter.api.Test;
 
 public class BigQueryParameterHandlerTest {
@@ -156,4 +157,87 @@ public void testConfigureParametersWidenNumericTypes() throws Exception {
     assertEquals("5", config.getPositionalParameters().get(0).getValue());
     assertEquals("10", config.getPositionalParameters().get(1).getValue());
   }
+
+  @Test
+  public void testFormatValueForQueryParameter_timestampPrecision() {
+    Timestamp ts = Timestamp.valueOf("2024-01-01 12:34:56.123456789");
+
+    // When enableTimestampPicos is true, full nanoseconds are preserved
+    Object picosFormatted =
+        BigQueryParameterHandler.formatValueForQueryParameter(
+            ts, StandardSQLTypeName.TIMESTAMP, true);
+    assertEquals("2024-01-01 12:34:56.123456789", picosFormatted);
+
+    // When enableTimestampPicos is false, truncated to microseconds
+    Object microsFormatted =
+        BigQueryParameterHandler.formatValueForQueryParameter(
+            ts, StandardSQLTypeName.TIMESTAMP, false);
+    assertEquals("2024-01-01 12:34:56.123456", microsFormatted);
+  }
+
+  @Test
+  public void testFormatValueForQueryParameter_stringTimestampPrecision() {
+    String picosString = "2024-01-01 12:34:56.123456789012";
+
+    // When enableTimestampPicos is true, up to 12 digits are preserved
+    Object picosFormatted =
+        BigQueryParameterHandler.formatValueForQueryParameter(
+            picosString, StandardSQLTypeName.TIMESTAMP, true);
+    assertEquals("2024-01-01 12:34:56.123456789012", picosFormatted);
+
+    // 15 digits truncated to 12 digits
+    String excessString = "2024-01-01 12:34:56.123456789012345";
+    Object truncatedPicos =
+        BigQueryParameterHandler.formatValueForQueryParameter(
+            excessString, StandardSQLTypeName.TIMESTAMP, true);
+    assertEquals("2024-01-01 12:34:56.123456789012", truncatedPicos);
+
+    // 'T' delimiter at index 10 is replaced with space and whitespace is trimmed
+    String isoString = "  2024-01-01T12:34:56.123456789012  ";
+    Object replacedT =
+        BigQueryParameterHandler.formatValueForQueryParameter(
+            isoString, StandardSQLTypeName.TIMESTAMP, true);
+    assertEquals("2024-01-01 12:34:56.123456789012", replacedT);
+
+    // Characters matching 'T' in timezone names are not replaced
+    String timezoneWithT = "  2024-01-01T12:34:56.123456 America/Toronto  ";
+    Object preservedTimezone =
+        BigQueryParameterHandler.formatValueForQueryParameter(
+            timezoneWithT, StandardSQLTypeName.TIMESTAMP, false);
+    assertEquals("2024-01-01 12:34:56.123456 America/Toronto", preservedTimezone);
+
+    String utcTimezone = "2024-01-01 12:34:56.123456 UTC";
+    Object preservedUtc =
+        BigQueryParameterHandler.formatValueForQueryParameter(
+            utcTimezone, StandardSQLTypeName.TIMESTAMP, false);
+    assertEquals("2024-01-01 12:34:56.123456 UTC", preservedUtc);
+
+    // When enableTimestampPicos is false, truncated to 6 digits
+    Object microsFormatted =
+        BigQueryParameterHandler.formatValueForQueryParameter(
+            picosString, StandardSQLTypeName.TIMESTAMP, false);
+    assertEquals("2024-01-01 12:34:56.123456", microsFormatted);
+  }
+
+  @Test
+  public void testConfigureParameters_withEnableTimestampPicos() throws Exception {
+    BigQueryParameterHandler paramHandler = new BigQueryParameterHandler(2, true);
+    Timestamp ts = Timestamp.valueOf("2024-01-01 12:34:56.123456789");
+    String picosStr = "2024-01-01 12:34:56.123456789012";
+
+    paramHandler.setParameter(1, ts, Timestamp.class);
+    paramHandler.setParameter(2, picosStr, Timestamp.class);
+
+    QueryJobConfiguration.Builder builder = QueryJobConfiguration.newBuilder("SELECT ?, ?");
+    paramHandler.configureParameters(builder);
+
+    QueryJobConfiguration config = builder.build();
+    assertEquals(2, config.getPositionalParameters().size());
+    assertEquals(
+        "2024-01-01 12:34:56.123456789", config.getPositionalParameters().get(0).getValue());
+    assertEquals(StandardSQLTypeName.TIMESTAMP, config.getPositionalParameters().get(0).getType());
+    assertEquals(
+        "2024-01-01 12:34:56.123456789012", config.getPositionalParameters().get(1).getValue());
+    assertEquals(StandardSQLTypeName.TIMESTAMP, config.getPositionalParameters().get(1).getType());
+  }
 }