return toStringResult for getObject and getBytes called on structured…

… type
snowflakedb · Oct 29, 2024 · 93e98ea · 93e98ea
1 parent bb4593e
commit 93e98ea
Show file tree

Hide file tree

Showing 7 changed files with 59 additions and 18 deletions.
diff --git a/src/main/java/net/snowflake/client/core/SFArrowResultSet.java b/src/main/java/net/snowflake/client/core/SFArrowResultSet.java
@@ -573,16 +573,8 @@ public Object getObject(int columnIndex) throws SFException {
     converter.setTreatNTZAsUTC(treatNTZAsUTC);
     converter.setUseSessionTimezone(useSessionTimezone);
     converter.setSessionTimeZone(sessionTimeZone);
-    Object obj = converter.toObject(index);
-    boolean isStructuredType = resultSetMetaData.isStructuredTypeColumn(columnIndex);
-    if (type == Types.STRUCT && isStructuredType) {
-      if (converter instanceof VarCharConverter) {
-        return createJsonSqlInput(columnIndex, obj);
-      } else if (converter instanceof StructConverter) {
-        return createArrowSqlInput(columnIndex, (Map<String, Object>) obj);
-      }
-    }
-    return obj;
+
+    return converter.toString(index);
   }
 
   private Object createJsonSqlInput(int columnIndex, Object obj) throws SFException {

diff --git a/src/main/java/net/snowflake/client/core/arrow/ArrayConverter.java b/src/main/java/net/snowflake/client/core/arrow/ArrayConverter.java
@@ -19,7 +19,12 @@ public ArrayConverter(ListVector valueVector, int vectorIndex, DataConversionCon
 
   @Override
   public Object toObject(int index) throws SFException {
-    return vector.getObject(index);
+    return toString(index);
+  }
+
+  @Override
+  public byte[] toBytes(int index) throws SFException {
+    return toString(index).getBytes();
   }
 
   @Override

diff --git a/src/main/java/net/snowflake/client/core/arrow/MapConverter.java b/src/main/java/net/snowflake/client/core/arrow/MapConverter.java
@@ -22,11 +22,12 @@ public MapConverter(MapVector valueVector, int columnIndex, DataConversionContex
 
   @Override
   public Object toObject(int index) throws SFException {
-    List<JsonStringHashMap<String, Object>> entriesList =
-        (List<JsonStringHashMap<String, Object>>) vector.getObject(index);
-    return entriesList.stream()
-        .collect(
-            Collectors.toMap(entry -> entry.get("key").toString(), entry -> entry.get("value")));
+    return toString(index);
+  }
+
+  @Override
+  public byte[] toBytes(int index) throws SFException {
+    return toString(index).getBytes();
   }
 
   @Override

diff --git a/src/main/java/net/snowflake/client/core/arrow/StructConverter.java b/src/main/java/net/snowflake/client/core/arrow/StructConverter.java
@@ -21,7 +21,12 @@ public StructConverter(StructVector vector, int columnIndex, DataConversionConte
 
   @Override
   public Object toObject(int index) throws SFException {
-    return structVector.getObject(index);
+    return toString(index);
+  }
+
+  @Override
+  public byte[] toBytes(int index) throws SFException {
+    return toString(index).getBytes();
   }
 
   @Override

diff --git a/src/main/java/net/snowflake/client/core/arrow/VectorTypeConverter.java b/src/main/java/net/snowflake/client/core/arrow/VectorTypeConverter.java
@@ -21,6 +21,11 @@ public Object toObject(int index) throws SFException {
     return vector.getObject(index);
   }
 
+  @Override
+  public byte[] toBytes(int index) throws SFException {
+    return toString(index).getBytes();
+  }
+
   @Override
   public String toString(int index) throws SFException {
     List<?> object = vector.getObject(index);

diff --git a/...owflake/client/jdbc/structuredtypes/StructuredTypesGetStringArrowJsonCompatibilityIT.java b/...owflake/client/jdbc/structuredtypes/StructuredTypesGetStringArrowJsonCompatibilityIT.java
@@ -62,6 +62,24 @@ public void testRunAsGetString() throws SQLException {
         (resultSet) -> assertGetStringIsCompatible(resultSet, expectedStructureTypeRepresentation));
   }
 
+  @Test
+  @ConditionalIgnoreRule.ConditionalIgnore(condition = RunningOnGithubAction.class)
+  public void testRunAsGetObject() throws SQLException {
+    withFirstRow(
+            connections.get(queryResultFormat),
+            selectSql,
+            (resultSet) -> assertGetObjectIsCompatible(resultSet, expectedStructureTypeRepresentation));
+  }
+
+  @Test
+  @ConditionalIgnoreRule.ConditionalIgnore(condition = RunningOnGithubAction.class)
+  public void testRunAsGetBytes() throws SQLException {
+    withFirstRow(
+            connections.get(queryResultFormat),
+            selectSql,
+            (resultSet) -> assertGetBytesIsCompatible(resultSet, expectedStructureTypeRepresentation));
+  }
+
   @Parameterized.Parameters(name = "format={0},sql={1}")
   public static Collection<Object[]> data() {
     Map<String, String> samples = new LinkedHashMap<>();

diff --git a/src/test/java/net/snowflake/client/jdbc/structuredtypes/StructuredTypesGetStringBaseIT.java b/src/test/java/net/snowflake/client/jdbc/structuredtypes/StructuredTypesGetStringBaseIT.java
@@ -2,6 +2,7 @@
 
 import static org.junit.Assert.assertTrue;
 
+import java.nio.charset.StandardCharsets;
 import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.SQLException;
@@ -50,12 +51,26 @@ protected static Connection initConnection(ResultSetFormatType queryResultFormat
     return conn;
   }
 
+  protected void assertGetBytesIsCompatible(ResultSet resultSet, String expected)
+          throws SQLException {
+    String result = new String(resultSet.getBytes(1), StandardCharsets.UTF_8);
+    TestUtil.assertEqualsIgnoringWhitespace(expected, result);
+  }
+
   protected void assertGetStringIsCompatible(ResultSet resultSet, String expected)
-      throws SQLException {
+          throws SQLException {
     String result = resultSet.getString(1);
     TestUtil.assertEqualsIgnoringWhitespace(expected, result);
   }
 
+  protected void assertGetObjectIsCompatible(ResultSet resultSet, String expected)
+      throws SQLException {
+    String result = resultSet.getObject(1, String.class);
+    String resultCasted = (String) resultSet.getObject(1);
+    TestUtil.assertEqualsIgnoringWhitespace(expected, result);
+    TestUtil.assertEqualsIgnoringWhitespace(expected, resultCasted);
+  }
+
   protected void withFirstRow(
       Connection connection, String sqlText, ThrowingConsumer<ResultSet, SQLException> consumer)
       throws SQLException {