From ebd0f0627523f67f84a56666a23a4c2c33fab5f2 Mon Sep 17 00:00:00 2001
From: Dawid Heyman <dawid.heyman@snowflake.com>
Date: Wed, 3 Apr 2024 15:26:36 +0200
Subject: [PATCH] =?UTF-8?q?SNOW-1234216=20Add=20checks=20for=20structured?=
 =?UTF-8?q?=20types=20for=20getMap,=20getObject=20and=20getArray=20?=
 =?UTF-8?q?=E2=80=A6=20(#1694)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../client/jdbc/SnowflakeBaseResultSet.java   | 40 ++++++++++++-------
 .../client/jdbc/SnowflakeResultSetV1.java     |  4 ++
 2 files changed, 30 insertions(+), 14 deletions(-)

diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeBaseResultSet.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeBaseResultSet.java
index be8a2783d..bd330e68e 100644
--- 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
index 02909f4ff..bdd8ffcee 100644
--- 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);