CR suggestions

snowflakedb · Mar 28, 2024 · 2659cde · 2659cde
1 parent 155aaf7
commit 2659cde
Show file tree

Hide file tree

Showing 5 changed files with 67 additions and 56 deletions.
diff --git a/src/main/java/net/snowflake/client/core/ArrowSqlInput.java b/src/main/java/net/snowflake/client/core/ArrowSqlInput.java
@@ -25,12 +25,12 @@
 @SnowflakeJdbcInternalApi
 public class ArrowSqlInput extends BaseSqlInput {
 
-  private final JsonStringHashMap<String, Object> input;
+  private final Map<String, Object> input;
   private final Iterator<Object> structuredTypeFields;
   private int currentIndex = 0;
 
   public ArrowSqlInput(
-      JsonStringHashMap<String, Object> input,
+      Map<String, Object> input,
       SFBaseSession session,
       Converters converters,
       List<FieldMetadata> fields) {
@@ -214,7 +214,7 @@ public <T> T readObject(Class<T> type) throws SQLException {
           SQLData instance = (SQLData) SQLDataCreationHelper.create(type);
           instance.readSQL(
               new ArrowSqlInput(
-                  (JsonStringHashMap<String, Object>) value,
+                  (Map<String, Object>) value,
                   session,
                   converters,
                   fieldMetadata.getFields()),

diff --git a/src/main/java/net/snowflake/client/core/SFArrowResultSet.java b/src/main/java/net/snowflake/client/core/SFArrowResultSet.java
@@ -19,6 +19,7 @@
 import java.sql.Time;
 import java.sql.Timestamp;
 import java.sql.Types;
+import java.util.Map;
 import java.util.TimeZone;
 import net.snowflake.client.core.arrow.ArrowVectorConverter;
 import net.snowflake.client.core.arrow.StructConverter;
@@ -362,6 +363,27 @@ public Converters getConverters() {
     return converters;
   }
 
+  @Override
+  public Date convertToDate(Object object, TimeZone tz) throws SFException {
+    return converters
+            .getStructuredTypeDateTimeConverter()
+            .getDate((int) object, tz);
+  }
+
+  @Override
+  public Time convertToTime(Object object, int scale) throws SFException {
+    return      converters.getStructuredTypeDateTimeConverter()
+            .getTime((long) object, scale);
+  }
+
+  @Override
+  public Timestamp convertToTimestamp(Object object, int columnType, int columnSubType, TimeZone tz, int scale) throws SFException {
+     return converters
+            .getStructuredTypeDateTimeConverter()
+            .getTimestamp(
+                    (JsonStringHashMap<String, Object>) object, columnType, columnSubType, tz, scale);
+  }
+
   /**
    * Advance to next row
    *
@@ -516,7 +538,7 @@ public Object getObject(int columnIndex) throws SFException {
       if (converter instanceof VarCharConverter) {
         return createJsonSqlInput(columnIndex, obj);
       } else if (converter instanceof StructConverter) {
-        return createArrowSqlInput(columnIndex, (JsonStringHashMap<String, Object>) obj);
+        return createArrowSqlInput(columnIndex, (Map<String, Object>) obj);
       }
     }
     return obj;
@@ -536,7 +558,7 @@ private Object createJsonSqlInput(int columnIndex, Object obj) throws SFExceptio
     }
   }
 
-  private Object createArrowSqlInput(int columnIndex, JsonStringHashMap<String, Object> input) {
+  private Object createArrowSqlInput(int columnIndex, Map<String, Object> input) {
     return new ArrowSqlInput(
         input,
         session,

diff --git a/src/main/java/net/snowflake/client/core/SFBaseResultSet.java b/src/main/java/net/snowflake/client/core/SFBaseResultSet.java
@@ -191,15 +191,15 @@ public boolean isArrayBindSupported() {
    * data size.
    *
    * @param maxSizeInBytes The expected max data size wrapped in the ResultSetSerializables object.
-   *     NOTE: this parameter is intended to make the data size in each serializable object to be
-   *     less than it. But if user specifies a small value which may be smaller than the data size
-   *     of one result chunk. So the definition can't be guaranteed completely. For this special
-   *     case, one serializable object is used to wrap the data chunk.
+   *                       NOTE: this parameter is intended to make the data size in each serializable object to be
+   *                       less than it. But if user specifies a small value which may be smaller than the data size
+   *                       of one result chunk. So the definition can't be guaranteed completely. For this special
+   *                       case, one serializable object is used to wrap the data chunk.
    * @return a list of SnowflakeResultSetSerializable
    * @throws SQLException if fails to split objects.
    */
   public List<SnowflakeResultSetSerializable> getResultSetSerializables(long maxSizeInBytes)
-      throws SQLException {
+          throws SQLException {
     return this.resultSetSerializable.splitBySize(maxSizeInBytes);
   }
 
@@ -213,4 +213,15 @@ public Converters getConverters() {
   public TimeZone getSessionTimeZone() {
     return resultSetSerializable.getTimeZone();
   }
+
+  @SnowflakeJdbcInternalApi
+  public abstract Date convertToDate(Object object, TimeZone tz) throws SFException;
+
+  @SnowflakeJdbcInternalApi
+  public abstract Time convertToTime(Object object, int scale) throws SFException;
+
+  @SnowflakeJdbcInternalApi
+  public abstract Timestamp convertToTimestamp(
+          Object object, int columnType, int columnSubType, TimeZone tz, int scale) throws SFException;
 }
+
diff --git a/src/main/java/net/snowflake/client/core/SFJsonResultSet.java b/src/main/java/net/snowflake/client/core/SFJsonResultSet.java
@@ -426,4 +426,21 @@ private static Object convert(JsonStringToTypeConverter converter, JsonNode node
       return converter.convert(node.toString());
     }
   }
+
+  @Override
+  public Date convertToDate(Object object, TimeZone tz) throws SFException {
+    return (Date) converters.dateConverter(session).convert((String) object);
+  }
+
+  @Override
+  public Time convertToTime(Object object, int scale) throws SFException {
+    return (Time) converters.timeConverter(session).convert((String) object);
+  }
+
+  @Override
+  public Timestamp convertToTimestamp(Object object, int columnType, int columnSubType, TimeZone tz, int scale) throws SFException {
+    return (Timestamp) converters
+            .timestampConverter(columnSubType, columnType, scale, session, null, tz)
+            .convert((String) object);
+  }
 }
diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeBaseResultSet.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeBaseResultSet.java
@@ -1577,17 +1577,19 @@ public <T> Map<String, T> getMap(int columnIndex, Class<T> type) throws SQLExcep
                       .get(columnIndex - 1)
                       .getFields(),
                   sfBaseResultSet.getSessionTimezone());
-        } else {
+        } else if (object instanceof ArrowSqlInput || object instanceof Map) {
           sqlInput =
               new ArrowSqlInput(
-                  (JsonStringHashMap<String, Object>) entry.getValue(),
+                  (Map<String, Object>) entry.getValue(),
                   session,
                   sfBaseResultSet.getConverters(),
                   sfBaseResultSet
                       .getMetaData()
                       .getColumnMetadata()
                       .get(columnIndex - 1)
                       .getFields());
+        } else {
+          throw new SQLException("SqlInput type " + object.getClass() + " is not supported when mapping to SQLData class");
         }
         instance.readSQL(sqlInput, null);
         resultMap.put(entry.getKey(), (T) instance);
@@ -1689,19 +1691,17 @@ public <T> Map<String, T> getMap(int columnIndex, Class<T> type) throws SQLExcep
       } else if (Date.class.isAssignableFrom(type)) {
         resultMap.put(
             entry.getKey(),
-            mapSFExceptionToSQLException(() -> (T) convertToDate(entry.getValue(), tz)));
+            mapSFExceptionToSQLException(() -> (T) sfBaseResultSet.convertToDate(entry.getValue(), tz)));
       } else if (Time.class.isAssignableFrom(type)) {
         resultMap.put(
             entry.getKey(),
-            mapSFExceptionToSQLException(() -> (T) convertToTime(entry.getValue(), tz, scale)));
+            mapSFExceptionToSQLException(() -> (T) sfBaseResultSet.convertToTime(entry.getValue(), scale)));
 
       } else if (Timestamp.class.isAssignableFrom(type)) {
         resultMap.put(
             entry.getKey(),
             mapSFExceptionToSQLException(
-                () ->
-                    (T)
-                        convertToTimestamp(
+                () -> (T) sfBaseResultSet.convertToTimestamp(
                             entry.getValue(), columnType, columnSubType, tz, scale)));
 
       } else {
@@ -1741,43 +1741,4 @@ public boolean isWrapperFor(Class<?> iface) throws SQLException {
 
     return iface.isInstance(this);
   }
-
-  private Date convertToDate(Object object, TimeZone tz) throws SFException {
-    if (sfBaseResultSet instanceof SFArrowResultSet) {
-      return sfBaseResultSet
-          .getConverters()
-          .getStructuredTypeDateTimeConverter()
-          .getDate((int) object, tz);
-    } else {
-      return (Date) sfBaseResultSet.getConverters().dateConverter(session).convert((String) object);
-    }
-  }
-
-  private Time convertToTime(Object object, TimeZone tz, int scale) throws SFException {
-    if (sfBaseResultSet instanceof SFArrowResultSet) {
-      return sfBaseResultSet
-          .getConverters()
-          .getStructuredTypeDateTimeConverter()
-          .getTime((long) object, scale);
-    } else {
-      return (Time) sfBaseResultSet.getConverters().timeConverter(session).convert((String) object);
-    }
-  }
-
-  private Timestamp convertToTimestamp(
-      Object object, int columnType, int columnSubType, TimeZone tz, int scale) throws SFException {
-    if (sfBaseResultSet instanceof SFArrowResultSet) {
-      return sfBaseResultSet
-          .getConverters()
-          .getStructuredTypeDateTimeConverter()
-          .getTimestamp(
-              (JsonStringHashMap<String, Object>) object, columnType, columnSubType, tz, scale);
-    } else {
-      return (Timestamp)
-          sfBaseResultSet
-              .getConverters()
-              .timestampConverter(columnSubType, columnType, scale, session, null, tz)
-              .convert((String) object);
-    }
-  }
 }