feat(bigquery-jdbc): support picosecond in REST JSON path and nested types

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

@@ -35,11 +35,15 @@ class BigQueryJsonArray extends BigQueryBaseArray {
   private List<FieldValue> values;
 
   BigQueryJsonArray(Field schema, FieldValue values) {
-    this(schema, values, BigQueryJdbcResultSetLogger.getLogger(BigQueryJsonArray.class));
+    this(schema, values, false, BigQueryJdbcResultSetLogger.getLogger(BigQueryJsonArray.class));
   }
 
-  BigQueryJsonArray(Field schema, FieldValue values, BigQueryJdbcResultSetLogger log) {
-    super(schema, log);
+  BigQueryJsonArray(
+      Field schema,
+      FieldValue values,
+      boolean enableTimestampPicos,
+      BigQueryJdbcResultSetLogger log) {
+    super(schema, enableTimestampPicos, log);
     this.values = (values == null || values.isNull()) ? null : values.getRepeatedValue();
   }
 
@@ -74,7 +78,11 @@ public ResultSet getResultSet() {
     BigQueryFieldValueListWrapper bigQueryFieldValueListWrapper =
         getNestedFieldValueListWrapper(FieldList.of(singleElementSchema()), this.values);
     return BigQueryJsonResultSet.getNestedResultSet(
-        Schema.of(this.schema), bigQueryFieldValueListWrapper, 0, this.values.size());
+        Schema.of(this.schema),
+        bigQueryFieldValueListWrapper,
+        0,
+        this.values.size(),
+        this.enableTimestampPicos);
   }
 
   @Override
@@ -88,7 +96,11 @@ public ResultSet getResultSet(long index, int count) {
     BigQueryFieldValueListWrapper bigQueryFieldValueListWrapper =
         getNestedFieldValueListWrapper(FieldList.of(singleElementSchema()), this.values);
     return BigQueryJsonResultSet.getNestedResultSet(
-        Schema.of(this.schema), bigQueryFieldValueListWrapper, range.x(), range.y());
+        Schema.of(this.schema),
+        bigQueryFieldValueListWrapper,
+        range.x(),
+        range.y(),
+        this.enableTimestampPicos);
   }
 
   @Override
@@ -99,10 +111,24 @@ public void free() {
 
   @Override
   Object getCoercedValue(int index) throws SQLException {
+    LOG.finestTrace("getCoercedValue");
     FieldValue fieldValue = this.values.get(index);
-    return this.arrayOfStruct
-        ? new BigQueryJsonStruct(
-            this.schema.getSubFields(), fieldValue, this.LOG.getJsonStructLogger())
-        : BigQueryTypeRegistry.convert(fieldValue, this.schema.getType().getStandardType(), null);
+    if (fieldValue == null || fieldValue.isNull()) {
+      return null;
+    }
+    if (this.arrayOfStruct) {
+      return new BigQueryJsonStruct(
+          this.schema.getSubFields(),
+          fieldValue,
+          this.enableTimestampPicos,
+          this.LOG.getJsonStructLogger());
+    }
+    if (this.enableTimestampPicos && BigQueryTemporalUtility.isPicosecondTimestamp(this.schema)) {
+      return BigQueryTemporalUtility.formatTimestampValue(fieldValue.getStringValue(), true);
+    }
+    if (this.enableTimestampPicos && BigQueryJsonResultSet.isRangeTimestamp(this.schema)) {
+      return BigQueryJsonResultSet.formatRangeTimestamp(fieldValue);
+    }
+    return BigQueryTypeRegistry.convert(fieldValue, this.schema.getType().getStandardType(), null);
   }
 }

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

@@ -24,7 +24,9 @@
 import com.google.cloud.bigquery.FieldValue;
 import com.google.cloud.bigquery.FieldValue.Attribute;
 import com.google.cloud.bigquery.Job;
+import com.google.cloud.bigquery.Range;
 import com.google.cloud.bigquery.Schema;
+import com.google.cloud.bigquery.StandardSQLTypeName;
 import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException;
 import java.sql.ResultSet;
 import java.sql.SQLException;
@@ -45,6 +47,7 @@ class BigQueryJsonResultSet extends BigQueryBaseResultSet {
   private final int fromIndex;
   private final int toIndexExclusive;
   private final Future<?>[] ownedTasks;
+  private final boolean enableTimestampPicos;
 
   private BigQueryJsonResultSet(
       Schema schema,
@@ -57,7 +60,8 @@ private BigQueryJsonResultSet(
       int toIndexExclusive,
       Future<?>[] ownedTasks,
       BigQuery bigQuery,
-      Job job) {
+      Job job,
+      boolean enableTimestampPicos) {
     super(bigQuery, statement, schema, isNested, job);
     this.totalRows = totalRows;
     this.buffer = buffer;
@@ -66,6 +70,7 @@ private BigQueryJsonResultSet(
     this.toIndexExclusive = toIndexExclusive;
     this.nestedRowIndex = fromIndex - 1;
     this.ownedTasks = ownedTasks;
+    this.enableTimestampPicos = enableTimestampPicos;
   }
 
   /**
@@ -95,7 +100,18 @@ static BigQueryJsonResultSet of(
       Job job) {
 
     return new BigQueryJsonResultSet(
-        schema, totalRows, buffer, statement, false, null, -1, -1, ownedTasks, bigQuery, job);
+        schema,
+        totalRows,
+        buffer,
+        statement,
+        false,
+        null,
+        -1,
+        -1,
+        ownedTasks,
+        bigQuery,
+        job,
+        statement != null && statement.isEnableTimestampPicos());
   }
 
   static BigQueryJsonResultSet of(
@@ -106,7 +122,18 @@ static BigQueryJsonResultSet of(
       Future<?>[] ownedTasks) {
 
     return new BigQueryJsonResultSet(
-        schema, totalRows, buffer, statement, false, null, -1, -1, ownedTasks, null, null);
+        schema,
+        totalRows,
+        buffer,
+        statement,
+        false,
+        null,
+        -1,
+        -1,
+        ownedTasks,
+        null,
+        null,
+        statement != null && statement.isEnableTimestampPicos());
   }
 
   static BigQueryJsonResultSet of(
@@ -133,6 +160,7 @@ static BigQueryJsonResultSet of(
     fromIndex = 0;
     ownedTasks = new Future<?>[0];
     toIndexExclusive = 0;
+    this.enableTimestampPicos = false;
   }
 
   //
@@ -145,10 +173,15 @@ static BigQueryJsonResultSet of(
    * @param cursor Points to the current record
    * @param fromIndex starting index under consideration
    * @param toIndexExclusive last index under consideration
+   * @param enableTimestampPicos whether picosecond timestamp precision is enabled
    * @return The BigQueryJsonResultSet
    */
   static BigQueryJsonResultSet getNestedResultSet(
-      Schema schema, BigQueryFieldValueListWrapper cursor, int fromIndex, int toIndexExclusive) {
+      Schema schema,
+      BigQueryFieldValueListWrapper cursor,
+      int fromIndex,
+      int toIndexExclusive,
+      boolean enableTimestampPicos) {
     return new BigQueryJsonResultSet(
         schema,
         -1,
@@ -160,7 +193,8 @@ static BigQueryJsonResultSet getNestedResultSet(
         toIndexExclusive,
         null,
         null,
-        null);
+        null,
+        enableTimestampPicos);
   }
 
   /* Advances the result set to the next row, returning false if no such row exists. Potentially blocking operation */
@@ -232,21 +266,62 @@ public Object getObject(int columnIndex) throws SQLException {
       Field arrayField = this.schema.getFields().get(0);
       if (isStruct(arrayField)) {
         return new BigQueryJsonStruct(
-            arrayField.getSubFields(), value, this.LOG.getJsonStructLogger());
+            arrayField.getSubFields(),
+            value,
+            this.enableTimestampPicos,
+            this.LOG.getJsonStructLogger());
+      }
+      if (this.enableTimestampPicos && BigQueryTemporalUtility.isPicosecondTimestamp(arrayField)) {
+        return BigQueryTemporalUtility.formatTimestampValue(value.getStringValue(), true);
+      }
+      if (this.enableTimestampPicos && isRangeTimestamp(arrayField)) {
+        return formatRangeTimestamp(value);
       }
       return BigQueryTypeRegistry.convert(value, arrayField.getType().getStandardType(), null);
     }
 
-    int extraIndex = this.isNested ? 2 : 1;
-    Field fieldSchema = this.schemaFieldList.get(columnIndex - extraIndex);
+    Field fieldSchema = this.schemaFieldList.get(columnIndex - 1);
     if (isArray(fieldSchema)) {
-      return new BigQueryJsonArray(fieldSchema, value, this.LOG.getJsonArrayLogger());
-    } else if (isStruct(fieldSchema)) {
+      return new BigQueryJsonArray(
+          fieldSchema, value, this.enableTimestampPicos, this.LOG.getJsonArrayLogger());
+    }
+    if (isStruct(fieldSchema)) {
       return new BigQueryJsonStruct(
-          fieldSchema.getSubFields(), value, this.LOG.getJsonStructLogger());
-    } else {
-      return BigQueryTypeRegistry.convert(value, fieldSchema.getType().getStandardType(), null);
+          fieldSchema.getSubFields(),
+          value,
+          this.enableTimestampPicos,
+          this.LOG.getJsonStructLogger());
     }
+    if (this.enableTimestampPicos && isRangeTimestamp(fieldSchema)) {
+      return formatRangeTimestamp(value);
+    }
+    if (this.enableTimestampPicos && BigQueryTemporalUtility.isPicosecondTimestamp(fieldSchema)) {
+      return BigQueryTemporalUtility.formatTimestampValue(value.getStringValue(), true);
+    }
+    return BigQueryTypeRegistry.convert(value, fieldSchema.getType().getStandardType(), null);
+  }
+
+  static String formatRangeTimestamp(FieldValue value) throws SQLException {
+    Range range = value.getRangeValue();
+    String start =
+        range.getStart().isNull()
+            ? "UNBOUNDED"
+            : BigQueryTemporalUtility.formatTimestampValue(range.getStart().getStringValue(), true);
+    String end =
+        range.getEnd().isNull()
+            ? "UNBOUNDED"
+            : BigQueryTemporalUtility.formatTimestampValue(range.getEnd().getStringValue(), true);
+    return String.format("[%s, %s)", start, end);
+  }
+
+  static boolean isRangeTimestamp(Field field) {
+    return field != null
+        && field.getRangeElementType() != null
+        && field.getRangeElementType().getType() != null
+        && field
+            .getRangeElementType()
+            .getType()
+            .equalsIgnoreCase(StandardSQLTypeName.TIMESTAMP.name());
   }
 
   /**

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

@@ -36,11 +36,15 @@ class BigQueryJsonStruct extends BigQueryBaseStruct {
   private final List<FieldValue> values;
 
   public BigQueryJsonStruct(FieldList schema, FieldValue values) {
-    this(schema, values, BigQueryJdbcResultSetLogger.getLogger(BigQueryJsonStruct.class));
+    this(schema, values, false, BigQueryJdbcResultSetLogger.getLogger(BigQueryJsonStruct.class));
   }
 
-  public BigQueryJsonStruct(FieldList schema, FieldValue values, BigQueryJdbcResultSetLogger log) {
-    super(log);
+  public BigQueryJsonStruct(
+      FieldList schema,
+      FieldValue values,
+      boolean enableTimestampPicos,
+      BigQueryJdbcResultSetLogger log) {
+    super(enableTimestampPicos, log);
     this.schema = schema;
     this.values = (values == null || values.isNull()) ? null : values.getRecordValue();
   }
@@ -67,14 +71,27 @@ public Object[] getAttributes() throws SQLException {
 
   private Object getValue(Field currentSchema, FieldValue currentValue) throws SQLException {
     LOG.finestTrace("getValue");
+    if (currentValue == null || currentValue.isNull()) {
+      return null;
+    }
     if (isArray(currentSchema)) {
-      return new BigQueryJsonArray(currentSchema, currentValue, this.LOG.getJsonArrayLogger());
-    } else if (isStruct(currentSchema)) {
+      return new BigQueryJsonArray(
+          currentSchema, currentValue, this.enableTimestampPicos, this.LOG.getJsonArrayLogger());
+    }
+    if (isStruct(currentSchema)) {
       return new BigQueryJsonStruct(
-          currentSchema.getSubFields(), currentValue, this.LOG.getJsonStructLogger());
-    } else {
-      return BigQueryTypeRegistry.convert(
-          currentValue, currentSchema.getType().getStandardType(), null);
+          currentSchema.getSubFields(),
+          currentValue,
+          this.enableTimestampPicos,
+          this.LOG.getJsonStructLogger());
+    }
+    if (this.enableTimestampPicos && BigQueryTemporalUtility.isPicosecondTimestamp(currentSchema)) {
+      return BigQueryTemporalUtility.formatTimestampValue(currentValue.getStringValue(), true);
+    }
+    if (this.enableTimestampPicos && BigQueryJsonResultSet.isRangeTimestamp(currentSchema)) {
+      return BigQueryJsonResultSet.formatRangeTimestamp(currentValue);
     }
+    return BigQueryTypeRegistry.convert(
+        currentValue, currentSchema.getType().getStandardType(), null);
   }
 }

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

@@ -486,7 +486,7 @@ public void testJsonArrayConnectionIdPropagation() {
             Attribute.REPEATED,
             FieldValueList.of(
                 Collections.singletonList(FieldValue.of(Attribute.PRIMITIVE, "123"))));
-    BigQueryJsonArray array = new BigQueryJsonArray(arraySchema, arrayValue, logger);
+    BigQueryJsonArray array = new BigQueryJsonArray(arraySchema, arrayValue, false, logger);
 
     assertConnectionIdPropagated(
         logger, connectionId, "Log from JSON Array", () -> array.LOG.fine("Log from JSON Array"));
@@ -504,7 +504,8 @@ public void testJsonStructConnectionIdPropagation() {
             Attribute.RECORD,
             FieldValueList.of(
                 Collections.singletonList(FieldValue.of(Attribute.PRIMITIVE, "456"))));
-    BigQueryJsonStruct struct = new BigQueryJsonStruct(structSchema, structValue, structLogger);
+    BigQueryJsonStruct struct =
+        new BigQueryJsonStruct(structSchema, structValue, false, structLogger);
 
     assertConnectionIdPropagated(
         structLogger,
@@ -531,7 +532,8 @@ public void testNestedStructInJsonArrayConnectionIdPropagation() throws Exceptio
                 Collections.singletonList(FieldValue.of(Attribute.PRIMITIVE, "789"))));
     FieldValue listVal =
         FieldValue.of(Attribute.REPEATED, FieldValueList.of(Collections.singletonList(recordVal)));
-    BigQueryJsonArray arrayWithNested = new BigQueryJsonArray(arrayNestedSchema, listVal, logger);
+    BigQueryJsonArray arrayWithNested =
+        new BigQueryJsonArray(arrayNestedSchema, listVal, false, logger);
 
     Object[] result = (Object[]) arrayWithNested.getArray();
     BigQueryJsonStruct nestedStruct = (BigQueryJsonStruct) result[0];
@@ -565,7 +567,7 @@ public void testNestedArrayInJsonStructConnectionIdPropagation() throws Exceptio
         FieldValue.of(Attribute.RECORD, FieldValueList.of(Collections.singletonList(arrayVal)));
 
     BigQueryJsonStruct structWithNestedArray =
-        new BigQueryJsonStruct(nestedSchema, rootVal, structLogger);
+        new BigQueryJsonStruct(nestedSchema, rootVal, false, structLogger);
     Object[] attributes = structWithNestedArray.getAttributes();
     BigQueryJsonArray nestedArray = (BigQueryJsonArray) attributes[0];