From 6a209b0d778e8f155c4e5099d62e5dd8111af2ab Mon Sep 17 00:00:00 2001
From: Przemyslaw Motacki <przemyslaw.motacki@snowflake.com>
Date: Tue, 2 Apr 2024 09:07:52 +0200
Subject: [PATCH] SNOW-1234214 Fix for json Maps (#1688)

* Mapping to json node
---
 .../client/core/SFJsonResultSet.java          |  9 +++--
 .../client/core/json/Converters.java          |  2 +-
 .../client/jdbc/SnowflakeBaseResultSet.java   | 35 +++++++++++++------
 .../snowflake/client/jdbc/SnowflakeUtil.java  |  4 +++
 4 files changed, 34 insertions(+), 16 deletions(-)

diff --git a/src/main/java/net/snowflake/client/core/SFJsonResultSet.java b/src/main/java/net/snowflake/client/core/SFJsonResultSet.java
index 5a91b7591..93946a0a3 100644
--- a/src/main/java/net/snowflake/client/core/SFJsonResultSet.java
+++ b/src/main/java/net/snowflake/client/core/SFJsonResultSet.java
@@ -4,6 +4,8 @@
 
 package net.snowflake.client.core;
 
+import static net.snowflake.client.jdbc.SnowflakeUtil.getJsonNodeStringValue;
+
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -420,11 +422,8 @@ private Stream getStream(Iterator nodeElements, JsonStringToTypeConverter conver
 
   private static Object convert(JsonStringToTypeConverter converter, JsonNode node)
       throws SFException {
-    if (node.isValueNode()) {
-      return converter.convert(node.asText());
-    } else {
-      return converter.convert(node.toString());
-    }
+    String nodeValue = getJsonNodeStringValue(node);
+    return converter.convert(nodeValue);
   }
 
   @Override
diff --git a/src/main/java/net/snowflake/client/core/json/Converters.java b/src/main/java/net/snowflake/client/core/json/Converters.java
index ab8bf60bb..cfdc532d4 100644
--- a/src/main/java/net/snowflake/client/core/json/Converters.java
+++ b/src/main/java/net/snowflake/client/core/json/Converters.java
@@ -199,7 +199,7 @@ public JsonStringToTypeConverter timeConverter(SFBaseSession session) {
       SnowflakeDateTimeFormat formatter =
           SnowflakeDateTimeFormat.fromSqlFormat(
               (String) session.getCommonParameters().get("TIME_OUTPUT_FORMAT"));
-      SFTimestamp timestamp = formatter.parse((String) value);
+      SFTimestamp timestamp = formatter.parse(value);
       return Time.valueOf(
           Instant.ofEpochMilli(timestamp.getTime()).atZone(ZoneOffset.UTC).toLocalTime());
     };
diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeBaseResultSet.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeBaseResultSet.java
index 2bad0a4e5..0cc0c225e 100644
--- a/src/main/java/net/snowflake/client/jdbc/SnowflakeBaseResultSet.java
+++ b/src/main/java/net/snowflake/client/jdbc/SnowflakeBaseResultSet.java
@@ -34,6 +34,7 @@
 import java.util.Arrays;
 import java.util.Calendar;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.TimeZone;
@@ -43,6 +44,7 @@
 import net.snowflake.client.core.ObjectMapperFactory;
 import net.snowflake.client.core.SFBaseResultSet;
 import net.snowflake.client.core.SFBaseSession;
+import net.snowflake.client.core.SFException;
 import net.snowflake.client.core.structs.SQLDataCreationHelper;
 import net.snowflake.client.log.SFLogger;
 import net.snowflake.client.log.SFLoggerFactory;
@@ -1550,16 +1552,8 @@ public <T> Map<String, T> getMap(int columnIndex, Class<T> type) throws SQLExcep
     int scale = valueFieldMetadata.getScale();
     TimeZone tz = sfBaseResultSet.getSessionTimeZone();
     Object object = getObject(columnIndex);
-    Map<String, Object> map;
-    if (object instanceof JsonSqlInput) {
-      map = new HashMap<>();
-      JsonNode jsonNode = ((JsonSqlInput) object).getInput();
-      jsonNode
-          .fieldNames()
-          .forEachRemaining(node -> map.put(node.toString(), jsonNode.get(node.toString())));
-    } else {
-      map = (Map<String, Object>) object;
-    }
+    Map<String, Object> map =
+        mapSFExceptionToSQLException(() -> prepareMapWithValues(object, type));
     Map<String, T> resultMap = new HashMap<>();
     for (Map.Entry<String, Object> entry : map.entrySet()) {
       if (SQLData.class.isAssignableFrom(type)) {
@@ -1748,4 +1742,25 @@ public boolean isWrapperFor(Class<?> iface) throws SQLException {
 
     return iface.isInstance(this);
   }
+
+  private <T> Map<String, Object> prepareMapWithValues(Object object, Class<T> type)
+      throws SFException {
+    if (object instanceof JsonSqlInput) {
+      Map map = new HashMap<>();
+      JsonNode jsonNode = ((JsonSqlInput) object).getInput();
+      for (Iterator<String> it = jsonNode.fieldNames(); it.hasNext(); ) {
+        String name = it.next();
+        map.put(
+            name,
+            SQLData.class.isAssignableFrom(type)
+                ? jsonNode.get(name)
+                : SnowflakeUtil.getJsonNodeStringValue(jsonNode.get(name)));
+      }
+      return map;
+    } else if (object instanceof Map) {
+      return (Map<String, Object>) object;
+    } else {
+      throw new SFException(ErrorCode.INVALID_STRUCT_DATA, "Object couldn't be converted to map");
+    }
+  }
 }
diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeUtil.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeUtil.java
index 00f698cfd..82dbca171 100644
--- a/src/main/java/net/snowflake/client/jdbc/SnowflakeUtil.java
+++ b/src/main/java/net/snowflake/client/jdbc/SnowflakeUtil.java
@@ -800,4 +800,8 @@ public static <T> T mapSFExceptionToSQLException(ThrowingCallable<T, SFException
       throw new SQLException(e);
     }
   }
+
+  public static String getJsonNodeStringValue(JsonNode node) throws SFException {
+    return node.isValueNode() ? node.asText() : node.toString();
+  }
 }