Conflicts resolved

snowflakedb · Mar 22, 2024 · 0afe368 · 0afe368
2 parents 348cfee + a36390a
commit 0afe368
Show file tree

Hide file tree

Showing 18 changed files with 366 additions and 19 deletions.
diff --git a/src/main/java/net/snowflake/client/core/SFArrowResultSet.java b/src/main/java/net/snowflake/client/core/SFArrowResultSet.java
@@ -24,6 +24,7 @@
 import net.snowflake.client.core.arrow.StructConverter;
 import net.snowflake.client.core.arrow.VarCharConverter;
 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;
@@ -505,8 +506,7 @@ public Object getObject(int columnIndex) throws SFException {
     converter.setSessionTimeZone(sessionTimezone);
     Object obj = converter.toObject(index);
     int type = resultSetMetaData.getColumnType(columnIndex);
-    if (type == Types.STRUCT
-        && Boolean.parseBoolean(System.getProperty(STRUCTURED_TYPE_ENABLED_PROPERTY_NAME))) {
+    if (type == Types.STRUCT && StructureTypeHelper.isStructureTypeEnabled()) {
       if (converter instanceof VarCharConverter) {
         return createJsonSqlInput(columnIndex, obj);
       } else if (converter instanceof StructConverter) {

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/SFBaseSession.java b/src/main/java/net/snowflake/client/core/SFBaseSession.java
@@ -142,6 +142,8 @@ public abstract class SFBaseSession {
 
   private Map<String, Object> commonParameters;
 
+  private boolean isJdbcArrowTreatDecimalAsInt = true;
+
   protected SFBaseSession(SFConnectionHandler sfConnectionHandler) {
     this.sfConnectionHandler = sfConnectionHandler;
   }
@@ -270,6 +272,14 @@ public void setJdbcTreatDecimalAsInt(boolean jdbcTreatDecimalAsInt) {
     isJdbcTreatDecimalAsInt = jdbcTreatDecimalAsInt;
   }
 
+  public boolean isJdbcArrowTreatDecimalAsInt() {
+    return isJdbcArrowTreatDecimalAsInt;
+  }
+
+  public void setJdbcArrowTreatDecimalAsInt(boolean jdbcArrowTreatDecimalAsInt) {
+    isJdbcArrowTreatDecimalAsInt = jdbcArrowTreatDecimalAsInt;
+  }
+
   public String getServerUrl() {
     if (connectionPropertiesMap.containsKey(SFSessionProperty.SERVER_URL)) {
       return (String) connectionPropertiesMap.get(SFSessionProperty.SERVER_URL);

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/SFSession.java b/src/main/java/net/snowflake/client/core/SFSession.java
@@ -475,6 +475,12 @@ public void addSFSessionProperty(String propertyName, Object propertyValue) thro
           }
           break;
 
+        case JDBC_ARROW_TREAT_DECIMAL_AS_INT:
+          if (propertyValue != null) {
+            setJdbcArrowTreatDecimalAsInt(getBooleanValue(propertyValue));
+          }
+          break;
+
         default:
           break;
       }

diff --git a/src/main/java/net/snowflake/client/core/SFSessionProperty.java b/src/main/java/net/snowflake/client/core/SFSessionProperty.java
@@ -80,7 +80,9 @@ public enum SFSessionProperty {
 
   ENABLE_PATTERN_SEARCH("enablePatternSearch", false, Boolean.class),
 
-  DISABLE_GCS_DEFAULT_CREDENTIALS("disableGcsDefaultCredentials", false, Boolean.class);
+  DISABLE_GCS_DEFAULT_CREDENTIALS("disableGcsDefaultCredentials", false, Boolean.class),
+
+  JDBC_ARROW_TREAT_DECIMAL_AS_INT("JDBC_ARROW_TREAT_DECIMAL_AS_INT", false, Boolean.class);
 
   // property key in string
   private String propertyKey;

diff --git a/src/main/java/net/snowflake/client/core/arrow/AbstractArrowVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/AbstractArrowVectorConverter.java
@@ -38,6 +38,8 @@ abstract class AbstractArrowVectorConverter implements ArrowVectorConverter {
 
   protected TimeZone sessionTimeZone;
 
+  private boolean shouldTreatDecimalAsInt;
+
   /** Field names of the struct vectors used by timestamp */
   public static final String FIELD_NAME_EPOCH = "epoch"; // seconds since epoch
 
@@ -53,6 +55,11 @@ abstract class AbstractArrowVectorConverter implements ArrowVectorConverter {
     this.valueVector = valueVector;
     this.columnIndex = vectorIndex + 1;
     this.context = context;
+    this.shouldTreatDecimalAsInt =
+        context == null
+            || context.getSession() == null
+            || context.getSession().isJdbcArrowTreatDecimalAsInt()
+            || context.getSession().isJdbcTreatDecimalAsInt();
   }
 
   @Override
@@ -146,6 +153,10 @@ public BigDecimal toBigDecimal(int index) throws SFException {
         ErrorCode.INVALID_VALUE_CONVERT, logicalTypeStr, SnowflakeUtil.BIG_DECIMAL_STR, "");
   }
 
+  boolean shouldTreatDecimalAsInt() {
+    return shouldTreatDecimalAsInt;
+  }
+
   @Override
   public void setTreatNTZAsUTC(boolean isUTC) {
     this.treatNTZasUTC = isUTC;

diff --git a/src/main/java/net/snowflake/client/core/arrow/BigIntToFixedConverter.java b/src/main/java/net/snowflake/client/core/arrow/BigIntToFixedConverter.java
@@ -137,9 +137,10 @@ public BigDecimal toBigDecimal(int index) {
   public Object toObject(int index) throws SFException {
     if (bigIntVector.isNull(index)) {
       return null;
-    } else {
-      return getLong(index);
+    } else if (!shouldTreatDecimalAsInt()) {
+      return BigDecimal.valueOf(getLong(index), sfScale);
     }
+    return getLong(index);
   }
 
   @Override

diff --git a/src/main/java/net/snowflake/client/core/arrow/IntToFixedConverter.java b/src/main/java/net/snowflake/client/core/arrow/IntToFixedConverter.java
@@ -105,7 +105,12 @@ public BigDecimal toBigDecimal(int index) throws SFException {
 
   @Override
   public Object toObject(int index) throws SFException {
-    return isNull(index) ? null : (long) getInt(index);
+    if (isNull(index)) {
+      return null;
+    } else if (!shouldTreatDecimalAsInt()) {
+      return BigDecimal.valueOf((long) getInt(index), sfScale);
+    }
+    return (long) getInt(index);
   }
 
   @Override

diff --git a/src/main/java/net/snowflake/client/core/arrow/SmallIntToFixedConverter.java b/src/main/java/net/snowflake/client/core/arrow/SmallIntToFixedConverter.java
@@ -99,6 +99,8 @@ public double toDouble(int index) throws SFException {
   public Object toObject(int index) throws SFException {
     if (isNull(index)) {
       return null;
+    } else if (!shouldTreatDecimalAsInt()) {
+      return BigDecimal.valueOf((long) getShort(index), sfScale);
     }
     return (long) getShort(index);
   }

diff --git a/src/main/java/net/snowflake/client/core/arrow/TinyIntToFixedConverter.java b/src/main/java/net/snowflake/client/core/arrow/TinyIntToFixedConverter.java
@@ -90,7 +90,12 @@ public BigDecimal toBigDecimal(int index) throws SFException {
 
   @Override
   public Object toObject(int index) throws SFException {
-    return isNull(index) ? null : (long) toByte(index);
+    if (isNull(index)) {
+      return null;
+    } else if (!shouldTreatDecimalAsInt()) {
+      return BigDecimal.valueOf((long) getByte(index), sfScale);
+    }
+    return (long) toByte(index);
   }
 
   @Override

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/SnowflakeBaseResultSet.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeBaseResultSet.java
@@ -1354,8 +1354,38 @@ public <T> T getObject(int columnIndex, Class<T> type) throws SQLException {
     } else if (Map.class.isAssignableFrom(type)) {
       JsonNode jsonNode = ((JsonSqlInput) getObject(columnIndex)).getInput();
       return (T) OBJECT_MAPPER.convertValue(jsonNode, new TypeReference<Map<String, Object>>() {});
+    } else if (String.class.isAssignableFrom(type)) {
+      return (T) getString(columnIndex);
+    } else if (Boolean.class.isAssignableFrom(type)) {
+      return (T) (Boolean) getBoolean(columnIndex);
+    } else if (Byte.class.isAssignableFrom(type)) {
+      return (T) (Byte) getByte(columnIndex);
+    } else if (Short.class.isAssignableFrom(type)) {
+      return (T) (Short) getShort(columnIndex);
+    } else if (Integer.class.isAssignableFrom(type)) {
+      return (T) (Integer) getInt(columnIndex);
+    } else if (Long.class.isAssignableFrom(type)) {
+      return (T) (Long) getLong(columnIndex);
+    } else if (Float.class.isAssignableFrom(type)) {
+      return (T) (Float) getFloat(columnIndex);
+    } else if (Double.class.isAssignableFrom(type)) {
+      return (T) (Double) getDouble(columnIndex);
+    } else if (Date.class.isAssignableFrom(type)) {
+      return (T) getDate(columnIndex);
+    } else if (Date.class.isAssignableFrom(type)) {
+      return (T) getDate(columnIndex);
+    } else if (Time.class.isAssignableFrom(type)) {
+      return (T) getTime(columnIndex);
+    } else if (Timestamp.class.isAssignableFrom(type)) {
+      return (T) getTimestamp(columnIndex);
+    } else if (BigDecimal.class.isAssignableFrom(type)) {
+      return (T) getBigDecimal(columnIndex);
     } else {
-      return (T) getObject(columnIndex);
+      logger.debug(
+          "Unsupported type passed to getObject(int columnIndex,Class<T> type): " + type.getName());
+      throw new SQLException(
+          "Type passed to 'getObject(int columnIndex,Class<T> type)' is unsupported. Type: "
+              + type.getName());
     }
   }
 

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;
 
@@ -80,8 +81,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
@@ -39,6 +39,7 @@
 import net.snowflake.client.core.SFException;
 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.client.util.ThrowingCallable;
@@ -279,14 +280,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: