SNOW-1234378 Return correct column type for structs depending on the …

…run flag
snowflakedb · Mar 21, 2024 · 88786c1 · 88786c1
1 parent 066e25a
commit 88786c1
Show file tree

Hide file tree

Showing 8 changed files with 103 additions and 14 deletions.
diff --git a/src/main/java/net/snowflake/client/core/SFArrowResultSet.java b/src/main/java/net/snowflake/client/core/SFArrowResultSet.java
@@ -22,6 +22,7 @@
 import java.util.TimeZone;
 import net.snowflake.client.core.arrow.ArrowVectorConverter;
 import net.snowflake.client.core.json.Converters;
+import net.snowflake.client.core.structs.StructureTypeHelper;
 import net.snowflake.client.jdbc.ArrowResultChunk;
 import net.snowflake.client.jdbc.ArrowResultChunk.ArrowChunkIterator;
 import net.snowflake.client.jdbc.ErrorCode;
@@ -517,8 +518,7 @@ public Array getArray(int columnIndex) throws SFException {
 
   private Object handleObjectType(int columnIndex, Object obj) throws SFException {
     int columnType = resultSetMetaData.getColumnType(columnIndex);
-    if (columnType == Types.STRUCT
-        && Boolean.valueOf(System.getProperty(STRUCTURED_TYPE_ENABLED_PROPERTY_NAME))) {
+    if (columnType == Types.STRUCT && StructureTypeHelper.isStructureTypeEnabled()) {
       try {
         JsonNode jsonNode = OBJECT_MAPPER.readTree((String) obj);
         return new JsonSqlInput(

diff --git a/src/main/java/net/snowflake/client/core/SFBaseResultSet.java b/src/main/java/net/snowflake/client/core/SFBaseResultSet.java
@@ -28,7 +28,6 @@
 /** Base class for query result set and metadata result set */
 public abstract class SFBaseResultSet {
   private static final SFLogger logger = SFLoggerFactory.getLogger(SFBaseResultSet.class);
-  static final String STRUCTURED_TYPE_ENABLED_PROPERTY_NAME = "STRUCTURED_TYPE_ENABLED";
 
   boolean wasNull = false;
 

diff --git a/src/main/java/net/snowflake/client/core/SFJsonResultSet.java b/src/main/java/net/snowflake/client/core/SFJsonResultSet.java
@@ -23,6 +23,7 @@
 import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
 import net.snowflake.client.core.json.Converters;
+import net.snowflake.client.core.structs.StructureTypeHelper;
 import net.snowflake.client.jdbc.ErrorCode;
 import net.snowflake.client.jdbc.FieldMetadata;
 import net.snowflake.client.jdbc.SnowflakeColumnMetadata;
@@ -95,13 +96,13 @@ public Object getObject(int columnIndex) throws SFException {
         return getBoolean(columnIndex);
 
       case Types.STRUCT:
-        if (Boolean.parseBoolean(System.getProperty(STRUCTURED_TYPE_ENABLED_PROPERTY_NAME))) {
+        if (StructureTypeHelper.isStructureTypeEnabled()) {
           return getSqlInput((String) obj, columnIndex);
         } else {
           throw new SFException(ErrorCode.FEATURE_UNSUPPORTED, "data type: " + type);
         }
       case Types.ARRAY:
-        if (Boolean.parseBoolean(System.getProperty(STRUCTURED_TYPE_ENABLED_PROPERTY_NAME))) {
+        if (StructureTypeHelper.isStructureTypeEnabled()) {
           return getArray(columnIndex);
         } else {
           throw new SFException(ErrorCode.FEATURE_UNSUPPORTED, "data type: " + type);

diff --git a/src/main/java/net/snowflake/client/core/structs/StructureTypeHelper.java b/src/main/java/net/snowflake/client/core/structs/StructureTypeHelper.java
@@ -0,0 +1,22 @@
+package net.snowflake.client.core.structs;
+
+import net.snowflake.client.core.SnowflakeJdbcInternalApi;
+
+@SnowflakeJdbcInternalApi
+public class StructureTypeHelper {
+  private static final String STRUCTURED_TYPE_ENABLED_PROPERTY_NAME = "STRUCTURED_TYPE_ENABLED";
+  private static boolean structuredTypeEnabled =
+      Boolean.valueOf(System.getProperty(STRUCTURED_TYPE_ENABLED_PROPERTY_NAME));
+
+  public static boolean isStructureTypeEnabled() {
+    return structuredTypeEnabled;
+  }
+
+  public static void enableStructuredType() {
+    structuredTypeEnabled = true;
+  }
+
+  public static void disableStructuredType() {
+    structuredTypeEnabled = false;
+  }
+}
diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeType.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeType.java
@@ -15,6 +15,7 @@
 import java.util.HashSet;
 import java.util.Set;
 import net.snowflake.client.core.SFBaseSession;
+import net.snowflake.client.core.structs.StructureTypeHelper;
 import net.snowflake.common.core.SFBinary;
 import net.snowflake.common.core.SqlState;
 
@@ -79,8 +80,13 @@ public static JavaDataType getJavaType(SnowflakeType type) {
       case BINARY:
         return JavaDataType.JAVA_BYTES;
       case ANY:
-      case OBJECT:
         return JavaDataType.JAVA_OBJECT;
+      case OBJECT:
+        if (StructureTypeHelper.isStructureTypeEnabled()) {
+          return JavaDataType.JAVA_OBJECT;
+        } else {
+          return JavaDataType.JAVA_STRING;
+        }
       default:
         // Those are not supported, but no reason to panic
         return JavaDataType.JAVA_STRING;

diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeUtil.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeUtil.java
@@ -38,6 +38,7 @@
 import net.snowflake.client.core.SFBaseSession;
 import net.snowflake.client.core.SFSessionProperty;
 import net.snowflake.client.core.SnowflakeJdbcInternalApi;
+import net.snowflake.client.core.structs.StructureTypeHelper;
 import net.snowflake.client.log.SFLogger;
 import net.snowflake.client.log.SFLoggerFactory;
 import net.snowflake.common.core.SnowflakeDateTimeFormat;
@@ -278,14 +279,22 @@ static ColumnTypeInfo getSnowflakeType(
             new ColumnTypeInfo(Types.ARRAY, defaultIfNull(extColTypeName, "ARRAY"), baseType);
         break;
 
-      case OBJECT:
       case MAP:
-        int targetType =
-            "GEOGRAPHY".equals(extColTypeName) || "GEOMETRY".equals(extColTypeName)
-                ? Types.VARCHAR
-                : Types.STRUCT;
         columnTypeInfo =
-            new ColumnTypeInfo(targetType, defaultIfNull(extColTypeName, "OBJECT"), baseType);
+            new ColumnTypeInfo(Types.STRUCT, defaultIfNull(extColTypeName, "OBJECT"), baseType);
+        break;
+
+      case OBJECT:
+        if (StructureTypeHelper.isStructureTypeEnabled()) {
+          boolean isGeoType =
+              "GEOMETRY".equals(extColTypeName) || "GEOGRAPHY".equals(extColTypeName);
+          int type = isGeoType ? Types.VARCHAR : Types.STRUCT;
+          columnTypeInfo =
+              new ColumnTypeInfo(type, defaultIfNull(extColTypeName, "OBJECT"), baseType);
+        } else {
+          columnTypeInfo =
+              new ColumnTypeInfo(Types.VARCHAR, defaultIfNull(extColTypeName, "OBJECT"), baseType);
+        }
         break;
 
       case VARIANT:

diff --git a/src/test/java/net/snowflake/client/jdbc/ResultSetStructuredTypesLatestIT.java b/src/test/java/net/snowflake/client/jdbc/ResultSetStructuredTypesLatestIT.java
@@ -15,6 +15,7 @@
 import java.sql.Statement;
 import java.sql.Time;
 import java.sql.Timestamp;
+import java.sql.Types;
 import java.time.LocalDate;
 import java.time.LocalDateTime;
 import java.time.LocalTime;
@@ -23,8 +24,10 @@
 import net.snowflake.client.ConditionalIgnoreRule;
 import net.snowflake.client.RunningOnGithubAction;
 import net.snowflake.client.ThrowingConsumer;
+import net.snowflake.client.ThrowingRunnable;
 import net.snowflake.client.category.TestCategoryStructuredType;
 import net.snowflake.client.core.structs.SnowflakeObjectTypeFactories;
+import net.snowflake.client.core.structs.StructureTypeHelper;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
@@ -404,6 +407,32 @@ public void testMapArrayOfArrays() throws SQLException {
         });
   }
 
+  @Test
+  @ConditionalIgnoreRule.ConditionalIgnore(condition = RunningOnGithubAction.class)
+  public void testColumnTypeWhenStructureTypeIsDisabled() throws Exception {
+    withStructureTypeTemporaryDisabled(
+        () -> {
+          withFirstRow(
+              "SELECT {'string':'a'}::OBJECT(string VARCHAR)",
+              resultSet -> {
+                assertEquals(Types.VARCHAR, resultSet.getMetaData().getColumnType(1));
+              });
+        });
+  }
+
+  @Test
+  @ConditionalIgnoreRule.ConditionalIgnore(condition = RunningOnGithubAction.class)
+  public void testColumnTypeWhenStructureTypeIsEnabled() throws Exception {
+    withStructureTypeTemporaryEnabled(
+        () -> {
+          withFirstRow(
+              "SELECT {'string':'a'}::OBJECT(string VARCHAR)",
+              resultSet -> {
+                assertEquals(Types.STRUCT, resultSet.getMetaData().getColumnType(1));
+              });
+        });
+  }
+
   private void withFirstRow(String sqlText, ThrowingConsumer<ResultSet, SQLException> consumer)
       throws SQLException {
     try (Connection connection = init();
@@ -413,5 +442,28 @@ private void withFirstRow(String sqlText, ThrowingConsumer<ResultSet, SQLExcepti
       consumer.accept(rs);
     }
   }
-  ;
+
+  private void withStructureTypeTemporaryEnabled(ThrowingRunnable action) throws Exception {
+    boolean isStructureTypeEnabled = StructureTypeHelper.isStructureTypeEnabled();
+    try {
+      StructureTypeHelper.enableStructuredType();
+      action.run();
+    } finally {
+      if (!isStructureTypeEnabled) {
+        StructureTypeHelper.disableStructuredType();
+      }
+    }
+  }
+
+  private void withStructureTypeTemporaryDisabled(ThrowingRunnable action) throws Exception {
+    boolean isStructureTypeEnabled = StructureTypeHelper.isStructureTypeEnabled();
+    try {
+      StructureTypeHelper.disableStructuredType();
+      action.run();
+    } finally {
+      if (isStructureTypeEnabled) {
+        StructureTypeHelper.enableStructuredType();
+      }
+    }
+  }
 }
diff --git a/src/test/java/net/snowflake/client/jdbc/SnowflakeDriverLatestIT.java b/src/test/java/net/snowflake/client/jdbc/SnowflakeDriverLatestIT.java
@@ -838,7 +838,7 @@ public void testGeoOutputTypes() throws Throwable {
       regularStatement.execute("insert into t_geo values ('POINT(0 0)'), ('LINESTRING(1 1, 2 2)')");
 
       testGeoOutputTypeSingle(
-          regularStatement, false, "geoJson", "OBJECT", "java.lang.String", Types.STRUCT);
+          regularStatement, false, "geoJson", "OBJECT", "java.lang.String", Types.VARCHAR);
 
       testGeoOutputTypeSingle(
           regularStatement, true, "geoJson", "GEOGRAPHY", "java.lang.String", Types.VARCHAR);