Fix TODOs, introduce factories

snowflakedb · Nov 16, 2023 · 750cee2 · 750cee2
1 parent 1888ade
commit 750cee2
Show file tree

Hide file tree

Showing 14 changed files with 98 additions and 61 deletions.
diff --git a/src/main/java/net/snowflake/client/core/SFArrowResultSet.java b/src/main/java/net/snowflake/client/core/SFArrowResultSet.java
@@ -479,7 +479,10 @@ public Object getObject(int columnIndex) throws SFException {
     converter.setUseSessionTimezone(useSessionTimezone);
     converter.setSessionTimeZone(timeZone);
     Object obj = converter.toObject(index);
-    // TODO structuredType clean this up
+    return handleObjectType(columnIndex, obj);
+  }
+
+  private Object handleObjectType(int columnIndex, Object obj) throws SFException {
     if (resultSetMetaData.getColumnType(columnIndex) == Types.STRUCT) {
       try {
         JsonNode jsonNode = OBJECT_MAPPER.readTree((String) obj);

diff --git a/src/main/java/net/snowflake/client/core/SFJsonResultSet.java b/src/main/java/net/snowflake/client/core/SFJsonResultSet.java
@@ -10,7 +10,10 @@
 import java.math.BigDecimal;
 import java.math.RoundingMode;
 import java.nio.ByteBuffer;
-import java.sql.*;
+import java.sql.Date;
+import java.sql.Time;
+import java.sql.Timestamp;
+import java.sql.Types;
 import java.util.TimeZone;
 import net.snowflake.client.core.arrow.ArrowResultUtil;
 import net.snowflake.client.jdbc.*;
@@ -26,8 +29,7 @@
 /** Abstract class used to represent snowflake result set in json format */
 public abstract class SFJsonResultSet extends SFBaseResultSet {
   private static final SFLogger logger = SFLoggerFactory.getLogger(SFJsonResultSet.class);
-  private static final ObjectMapper OBJECT_MAPPER =
-      new ObjectMapper(); // TODO structuredType is it proper init?
+  private static final ObjectMapper OBJECT_MAPPER = ObjectMapperFactory.getObjectMapper();
 
   TimeZone sessionTimeZone;
 
@@ -46,23 +48,11 @@ public abstract class SFJsonResultSet extends SFBaseResultSet {
    */
   protected abstract Object getObjectInternal(int columnIndex) throws SFException;
 
-  //  public <T extends SQLData> getObject(int columnIndex) throws SFException {
-  //    // check metaData
-  //    int type = resultSetMetaData.getColumnType(columnIndex);
-  //
-  ////    if (typeCanBeConvertToT)
-  //
-  //    T t = new T();
-  //    t.readSQL(new SQLInput());
-  //
-  //
-  //
-  //  }
   public Object getObject(int columnIndex) throws SFException {
 
     int type = resultSetMetaData.getColumnType(columnIndex);
 
-    Object obj = getObjectInternal(columnIndex); // JsonNode
+    Object obj = getObjectInternal(columnIndex);
     if (obj == null) {
       return null;
     }

diff --git a/src/main/java/net/snowflake/client/core/SFResultSet.java b/src/main/java/net/snowflake/client/core/SFResultSet.java
@@ -306,21 +306,16 @@ protected Object getObjectInternal(int columnIndex) throws SFException {
     final int internalColumnIndex = columnIndex - 1;
     Object retValue;
     if (sortResult) {
-      // StructuredType should return JsonNode too
       retValue = firstChunkSortedRowSet[currentChunkRowIndex][internalColumnIndex];
     } else if (firstChunkRowset != null) {
       retValue =
-          JsonResultChunk.extractCell(
-              firstChunkRowset, currentChunkRowIndex, internalColumnIndex, resultSetMetaData);
+          JsonResultChunk.extractCell(firstChunkRowset, currentChunkRowIndex, internalColumnIndex);
     } else if (currentChunk != null) {
-      // StructuredType should return JsonNode too
       retValue = currentChunk.getCell(currentChunkRowIndex, internalColumnIndex);
     } else {
       throw new SFException(ErrorCode.COLUMN_DOES_NOT_EXIST, columnIndex);
     }
     wasNull = retValue == null;
-
-    // TODO structuredType must be JsonNode instance
     return retValue;
   }
 
@@ -332,7 +327,7 @@ private void sortResultSet() {
       firstChunkSortedRowSet[rowIdx] = new Object[columnCount];
       for (int colIdx = 0; colIdx < columnCount; colIdx++) {
         firstChunkSortedRowSet[rowIdx][colIdx] =
-            JsonResultChunk.extractCell(firstChunkRowset, rowIdx, colIdx, resultSetMetaData);
+            JsonResultChunk.extractCell(firstChunkRowset, rowIdx, colIdx);
       }
     }
 

diff --git a/src/main/java/net/snowflake/client/core/structs/SnowflakeObjectTypeFactories.java b/src/main/java/net/snowflake/client/core/structs/SnowflakeObjectTypeFactories.java
@@ -0,0 +1,29 @@
+package net.snowflake.client.core.structs;
+
+import static net.snowflake.client.util.Predicates.notNull;
+
+import java.sql.SQLData;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.function.Supplier;
+
+public class SnowflakeObjectTypeFactories {
+  private static final Map<Class<?>, Supplier<SQLData>> factories = new HashMap<>();
+
+  public static synchronized void register(Class<?> type, Supplier<SQLData> factory) {
+    notNull(type, "type cannot be null");
+    notNull(factory, "factory cannot be null");
+    factories.put(type, factory);
+  }
+
+  public static synchronized void unregister(Class<?> type) {
+    notNull(type, "type cannot be null");
+    factories.remove(type);
+  }
+
+  public static synchronized Optional<Supplier<SQLData>> get(Class<?> type) {
+    notNull(type, "type cannot be null");
+    return Optional.ofNullable(factories.get(type));
+  }
+}
diff --git a/src/main/java/net/snowflake/client/jdbc/JsonResultChunk.java b/src/main/java/net/snowflake/client/jdbc/JsonResultChunk.java
@@ -12,7 +12,6 @@
 import java.util.LinkedList;
 import java.util.List;
 import net.snowflake.client.core.SFBaseSession;
-import net.snowflake.client.core.SFResultSetMetaData;
 import net.snowflake.client.log.SFLogger;
 import net.snowflake.client.log.SFLoggerFactory;
 import net.snowflake.common.core.SqlState;
@@ -33,24 +32,19 @@ public JsonResultChunk(
     this.session = session;
   }
 
-  public static Object extractCell(
-      JsonNode resultData, int rowIdx, int colIdx, SFResultSetMetaData resultSetMetaData) {
+  public static Object extractCell(JsonNode resultData, int rowIdx, int colIdx) {
     JsonNode currentRow = resultData.get(rowIdx);
 
     JsonNode colNode = currentRow.get(colIdx);
 
     if (colNode.isTextual()) {
-      return colNode.textValue();
+      return colNode.asText();
     } else if (colNode.isNumber()) {
       return colNode.numberValue();
-    }
-    // TODO: structuredType - need response with json but not as String
-    else if (colNode.isObject()) {
-      return colNode;
     } else if (colNode.isNull()) {
       return null;
     }
-    throw new RuntimeException("Unknown json type");
+    throw new RuntimeException("Unknow json type");
   }
 
   public void tryReuse(ResultChunkDataCache cache) {

diff --git a/src/main/java/net/snowflake/client/jdbc/SFAsyncResultSet.java b/src/main/java/net/snowflake/client/jdbc/SFAsyncResultSet.java
@@ -286,8 +286,7 @@ public ResultSetMetaData getMetaData() throws SQLException {
 
   public Object getObject(int columnIndex) throws SQLException {
     raiseSQLExceptionIfResultSetIsClosed();
-    return resultSetForNext.getObject(
-        columnIndex); // TODO structuredType does it work out of the box?
+    return resultSetForNext.getObject(columnIndex);
   }
 
   public BigDecimal getBigDecimal(int columnIndex) throws SQLException {

diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeBaseResultSet.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeBaseResultSet.java
@@ -10,11 +10,11 @@
 import java.math.BigDecimal;
 import java.net.URL;
 import java.sql.*;
-import java.util.Calendar;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.TimeZone;
+import java.sql.Date;
+import java.util.*;
+import java.util.function.Supplier;
 import net.snowflake.client.core.SFBaseSession;
+import net.snowflake.client.core.structs.SnowflakeObjectTypeFactories;
 import net.snowflake.client.log.SFLogger;
 import net.snowflake.client.log.SFLoggerFactory;
 import net.snowflake.common.core.SqlState;
@@ -1320,25 +1320,31 @@ public void updateNClob(String columnLabel, Reader reader) throws SQLException {
   public <T> T getObject(int columnIndex, Class<T> type) throws SQLException {
     logger.debug("public <T> T getObject(int columnIndex,Class<T> type)", false);
     if (SQLData.class.isAssignableFrom(type)) {
-      try {
-        SQLData instance = (SQLData) type.newInstance();
-        SQLInput sqlInput = (SQLInput) getObject(columnIndex);
-        instance.readSQL(sqlInput, null);
-        return (T) instance;
-      } // TODO structuredType clean exceptions
-      catch (Exception e) {
-        throw new RuntimeException(e);
-      }
+      Optional<Supplier<SQLData>> typeFactory = SnowflakeObjectTypeFactories.get(type);
+      SQLData instance =
+          typeFactory
+              .map(Supplier::get)
+              .orElseGet(() -> createUsingReflection((Class<SQLData>) type));
+      SQLInput sqlInput = (SQLInput) getObject(columnIndex);
+      instance.readSQL(sqlInput, null);
+      return (T) instance;
     } else {
       return (T) getObject(columnIndex);
     }
   }
 
-  // @Override
+  private SQLData createUsingReflection(Class<? extends SQLData> type) {
+    try {
+      return type.newInstance();
+    } catch (InstantiationException | IllegalAccessException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  @Override
   public <T> T getObject(String columnLabel, Class<T> type) throws SQLException {
     logger.debug("public <T> T getObject(String columnLabel,Class<T> type)", false);
     return getObject(findColumn(columnLabel), type);
-    //    throw new SnowflakeLoggedFeatureNotSupportedException(session);
   }
 
   @SuppressWarnings("unchecked")

diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeResultSetMetaDataV1.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeResultSetMetaDataV1.java
@@ -106,7 +106,7 @@ public boolean isCaseSensitive(int column) throws SQLException {
     int colType = getColumnType(column);
 
     switch (colType) {
-        // TODO structuredType how about other types?
+        // TODO structuredType fill for Array and Map
         // Note: SF types ARRAY, OBJECT, GEOMETRY are also represented as VARCHAR.
       case Types.VARCHAR:
       case Types.CHAR:

diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeType.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeType.java
@@ -48,7 +48,7 @@ public static SnowflakeType fromString(String name) {
   }
 
   public static JavaDataType getJavaType(SnowflakeType type) {
-    // TODO structuredType
+    // TODO structuredType fill for Array and Map
     switch (type) {
       case TEXT:
         return JavaDataType.JAVA_STRING;
@@ -72,11 +72,10 @@ public static JavaDataType getJavaType(SnowflakeType type) {
       case ARRAY:
       case VARIANT:
         return JavaDataType.JAVA_STRING;
-      case OBJECT:
-        return JavaDataType.JAVA_STRING;
       case BINARY:
         return JavaDataType.JAVA_BYTES;
       case ANY:
+      case OBJECT:
         return JavaDataType.JAVA_OBJECT;
       default:
         // Those are not supported, but no reason to panic

diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeUtil.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeUtil.java
@@ -217,15 +217,15 @@ public static SnowflakeColumnMetadata extractColumnMetadata(
         extColTypeName = "BOOLEAN";
         break;
 
+        // TODO structuredType fill for Array and Map
       case ARRAY:
         colType = Types.VARCHAR;
         extColTypeName = "ARRAY";
         break;
 
       case OBJECT:
         colType = Types.STRUCT;
-        // TODO : structuredType
-        extColTypeName = "STRUCT";
+        extColTypeName = "OBJECT";
         break;
 
       case VARIANT:

diff --git a/src/main/java/net/snowflake/client/util/Predicates.java b/src/main/java/net/snowflake/client/util/Predicates.java
@@ -0,0 +1,9 @@
+package net.snowflake.client.util;
+
+public class Predicates {
+  public static void notNull(Object o, String msg) {
+    if (o == null) {
+      throw new NullPointerException(msg);
+    }
+  }
+}
diff --git a/src/test/java/net/snowflake/client/jdbc/MockConnectionTest.java b/src/test/java/net/snowflake/client/jdbc/MockConnectionTest.java
@@ -619,8 +619,7 @@ public boolean next() {
 
     @Override
     protected Object getObjectInternal(int columnIndex) {
-      return JsonResultChunk.extractCell(
-          resultJson, currentRowIdx, columnIndex - 1, resultSetMetaData);
+      return JsonResultChunk.extractCell(resultJson, currentRowIdx, columnIndex - 1);
     }
 
     @Override

diff --git a/src/test/java/net/snowflake/client/jdbc/ResultSetAsyncIT.java b/src/test/java/net/snowflake/client/jdbc/ResultSetAsyncIT.java
@@ -224,6 +224,7 @@ public void testGetMethods() throws Throwable {
     Clob clob = connection.createClob();
     clob.setString(1, "hello world");
     Statement statement = connection.createStatement();
+    // TODO structureType - add to test when WRITE is ready
     statement.execute(
         "create or replace table test_get(colA integer, colB number, colC number, colD string, colE double, colF float, colG boolean, colH text, colI binary(3), colJ number(38,9), colK int, colL date, colM time, colN timestamp_ltz)");
 

diff --git a/src/test/java/net/snowflake/client/jdbc/ResultSetIT.java b/src/test/java/net/snowflake/client/jdbc/ResultSetIT.java
@@ -17,6 +17,7 @@
 import net.snowflake.client.ConditionalIgnoreRule;
 import net.snowflake.client.RunningOnGithubAction;
 import net.snowflake.client.category.TestCategoryResultSet;
+import net.snowflake.client.core.structs.SnowflakeObjectTypeFactories;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
 
@@ -68,8 +69,21 @@ public void writeSQL(SQLOutput stream) throws SQLException {}
   }
 
   @Test
-  public void testMapJson() throws SQLException {
+  public void testMapStructToObjectWithFactory() throws SQLException {
+    testMapJson(true);
+  }
 
+  @Test
+  public void testMapStructToObjectWithReflection() throws SQLException {
+    testMapJson(false);
+  }
+
+  private void testMapJson(boolean registerFactory) throws SQLException {
+    if (registerFactory) {
+      SnowflakeObjectTypeFactories.register(TestClass.class, TestClass::new);
+    } else {
+      SnowflakeObjectTypeFactories.unregister(TestClass.class);
+    }
     Connection connection = init();
     Statement statement = connection.createStatement();
     ResultSet resultSet = statement.executeQuery("select {'x':'a'}::OBJECT(x VARCHAR)");
@@ -81,7 +95,7 @@ public void testMapJson() throws SQLException {
   }
 
   @Test
-  public void testMapJsonFromChunks() throws SQLException {
+  public void testMapStructsFromChunks() throws SQLException {
 
     Connection connection = init();
     Statement statement = connection.createStatement();
@@ -90,7 +104,6 @@ public void testMapJsonFromChunks() throws SQLException {
             "select {'x':'a'}::OBJECT(x VARCHAR) FROM TABLE(GENERATOR(ROWCOUNT=>30000))");
     int i = 0;
     while (resultSet.next()) {
-      System.out.println(i++);
       TestClass object = resultSet.getObject(1, TestClass.class);
       assertEquals("a", object.x);
     }