SNOW-1234216 Add checks for structured types for getMap, getObject an…

…d getArray … (#1694)
snowflakedb · Apr 3, 2024 · ebd0f06 · ebd0f06
1 parent 846c7f9
commit ebd0f06
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 14 deletions.
diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeBaseResultSet.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeBaseResultSet.java
@@ -46,6 +46,7 @@
 import net.snowflake.client.core.SFBaseSession;
 import net.snowflake.client.core.SFException;
 import net.snowflake.client.core.structs.SQLDataCreationHelper;
+import net.snowflake.client.core.structs.StructureTypeHelper;
 import net.snowflake.client.log.SFLogger;
 import net.snowflake.client.log.SFLoggerFactory;
 import net.snowflake.common.core.SqlState;
@@ -1352,21 +1353,24 @@ public void updateNClob(String columnLabel, Reader reader) throws SQLException {
   @Override
   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)) {
-      SQLData instance = (SQLData) SQLDataCreationHelper.create(type);
-      SQLInput sqlInput = (SQLInput) getObject(columnIndex);
-      instance.readSQL(sqlInput, null);
-      return (T) instance;
-    } else if (Map.class.isAssignableFrom(type)) {
-      Object object = getObject(columnIndex);
-      if (object instanceof JsonSqlInput) {
-        JsonNode jsonNode = ((JsonSqlInput) object).getInput();
-        return (T)
-            OBJECT_MAPPER.convertValue(jsonNode, new TypeReference<Map<String, Object>>() {});
-      } else {
-        return (T) ((ArrowSqlInput) object).getInput();
+    if (StructureTypeHelper.isStructureTypeEnabled()) {
+      if (SQLData.class.isAssignableFrom(type)) {
+        SQLData instance = (SQLData) SQLDataCreationHelper.create(type);
+        SQLInput sqlInput = (SQLInput) getObject(columnIndex);
+        instance.readSQL(sqlInput, null);
+        return (T) instance;
+      } else if (Map.class.isAssignableFrom(type)) {
+        Object object = getObject(columnIndex);
+        if (object instanceof JsonSqlInput) {
+          JsonNode jsonNode = ((JsonSqlInput) object).getInput();
+          return (T)
+              OBJECT_MAPPER.convertValue(jsonNode, new TypeReference<Map<String, Object>>() {});
+        } else {
+          return (T) ((ArrowSqlInput) object).getInput();
+        }
       }
-    } else if (String.class.isAssignableFrom(type)) {
+    }
+    if (String.class.isAssignableFrom(type)) {
       return (T) getString(columnIndex);
     } else if (Boolean.class.isAssignableFrom(type)) {
       return (T) (Boolean) getBoolean(columnIndex);
@@ -1405,6 +1409,10 @@ public <T> List<T> getList(int columnIndex, Class<T> type) throws SQLException {
   }
 
   public <T> T[] getArray(int columnIndex, Class<T> type) throws SQLException {
+    logger.debug("public <T> T[] getArray(int columnIndex, Class<T> type)", false);
+    if (!StructureTypeHelper.isStructureTypeEnabled()) {
+      throw new SnowflakeLoggedFeatureNotSupportedException(session);
+    }
     FieldMetadata fieldMetadata =
         sfBaseResultSet.getMetaData().getColumnMetadata().get(columnIndex - 1).getFields().get(0);
     int columnSubType = fieldMetadata.getType();
@@ -1545,6 +1553,10 @@ public <T> T[] getArray(int columnIndex, Class<T> type) throws SQLException {
   }
 
   public <T> Map<String, T> getMap(int columnIndex, Class<T> type) throws SQLException {
+    logger.debug("public <T> Map<String, T> getMap(int columnIndex, Class<T> type)", false);
+    if (!StructureTypeHelper.isStructureTypeEnabled()) {
+      throw new SnowflakeLoggedFeatureNotSupportedException(session);
+    }
     FieldMetadata valueFieldMetadata =
         sfBaseResultSet.getMetaData().getColumnMetadata().get(columnIndex - 1).getFields().get(1);
     int columnSubType = valueFieldMetadata.getType();

diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeResultSetV1.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeResultSetV1.java
@@ -31,6 +31,7 @@
 import net.snowflake.client.core.QueryStatus;
 import net.snowflake.client.core.SFBaseResultSet;
 import net.snowflake.client.core.SFException;
+import net.snowflake.client.core.structs.StructureTypeHelper;
 
 /** Snowflake ResultSet implementation */
 public class SnowflakeResultSetV1 extends SnowflakeBaseResultSet
@@ -272,6 +273,9 @@ public Object getObject(int columnIndex) throws SQLException {
   }
 
   public Array getArray(int columnIndex) throws SQLException {
+    if (!StructureTypeHelper.isStructureTypeEnabled()) {
+      throw new SnowflakeLoggedFeatureNotSupportedException(session);
+    }
     raiseSQLExceptionIfResultSetIsClosed();
     try {
       return sfBaseResultSet.getArray(columnIndex);