Updated tests

snowflakedb · Mar 13, 2024 · 5428951 · 5428951
1 parent 0e4e3df
commit 5428951
Show file tree

Hide file tree

Showing 7 changed files with 39 additions and 87 deletions.
diff --git a/src/main/java/net/snowflake/client/core/arrow/AbstractArrowVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/AbstractArrowVectorConverter.java
@@ -146,6 +146,14 @@ public BigDecimal toBigDecimal(int index) throws SFException {
         ErrorCode.INVALID_VALUE_CONVERT, logicalTypeStr, SnowflakeUtil.BIG_DECIMAL_STR, "");
   }
 
+  public boolean shouldTreatDecimalAsInt(){
+    if (!context.getSession().isJdbcArrowTreatDecimalAsInt()
+            && !context.getSession().isJdbcTreatDecimalAsInt()) {
+      return false;
+    }
+    return true;
+  }
+
   @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,8 +137,7 @@ public BigDecimal toBigDecimal(int index) {
   public Object toObject(int index) throws SFException {
     if (bigIntVector.isNull(index)) {
       return null;
-    } else if (!context.getSession().isJdbcArrowTreatDecimalAsInt()
-        && !context.getSession().isJdbcTreatDecimalAsInt()) {
+    } else if (!shouldTreatDecimalAsInt()) {
       return BigDecimal.valueOf(getLong(index), sfScale);
     }
     return getLong(index);

diff --git a/src/main/java/net/snowflake/client/core/arrow/IntToFixedConverter.java b/src/main/java/net/snowflake/client/core/arrow/IntToFixedConverter.java
@@ -107,8 +107,7 @@ public BigDecimal toBigDecimal(int index) throws SFException {
   public Object toObject(int index) throws SFException {
     if (isNull(index)) {
       return null;
-    } else if (!context.getSession().isJdbcArrowTreatDecimalAsInt()
-        && !context.getSession().isJdbcTreatDecimalAsInt()) {
+    } else if (!shouldTreatDecimalAsInt()) {
       return BigDecimal.valueOf((long) getInt(index), sfScale);
     }
     return (long) getInt(index);

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,8 +99,7 @@ public double toDouble(int index) throws SFException {
   public Object toObject(int index) throws SFException {
     if (isNull(index)) {
       return null;
-    } else if (!context.getSession().isJdbcArrowTreatDecimalAsInt()
-        && !context.getSession().isJdbcTreatDecimalAsInt()) {
+    } 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
@@ -92,8 +92,7 @@ public BigDecimal toBigDecimal(int index) throws SFException {
   public Object toObject(int index) throws SFException {
     if (isNull(index)) {
       return null;
-    } else if (!context.getSession().isJdbcArrowTreatDecimalAsInt()
-        && !context.getSession().isJdbcTreatDecimalAsInt()) {
+    } else if (!shouldTreatDecimalAsInt()) {
       return BigDecimal.valueOf((long) getByte(index), sfScale);
     }
     return (long) toByte(index);

diff --git a/src/test/java/net/snowflake/client/jdbc/ResultSetArrowLatestIT.java b/src/test/java/net/snowflake/client/jdbc/ResultSetArrowLatestIT.java
@@ -4,17 +4,8 @@
 package net.snowflake.client.jdbc;
 
 import net.snowflake.client.category.TestCategoryArrow;
-import org.junit.Test;
 import org.junit.experimental.categories.Category;
 
-import java.sql.Connection;
-import java.sql.ResultSet;
-import java.sql.SQLException;
-import java.sql.Statement;
-import java.util.Properties;
-
-import static org.junit.Assert.assertEquals;
-
 /**
  * ResultSet integration tests for the latest JDBC driver. This doesn't work for the oldest
  * supported driver. Drop this file when ResultSetLatestIT is dropped.

diff --git a/src/test/java/net/snowflake/client/jdbc/ResultSetLatestIT.java b/src/test/java/net/snowflake/client/jdbc/ResultSetLatestIT.java
@@ -980,6 +980,24 @@ public void testLargeStringRetrieval() throws SQLException {
     }
   }
 
+  private static void assertAllColumnsAreLongButBigIntIsBigDecimal(ResultSet rs) throws SQLException {
+    while(rs.next()) {
+      assertEquals("class java.lang.Long", rs.getObject(1).getClass().toString());
+      assertEquals("class java.math.BigDecimal", rs.getObject(2).getClass().toString());
+      assertEquals("class java.lang.Long", rs.getObject(3).getClass().toString());
+      assertEquals("class java.lang.Long", rs.getObject(4).getClass().toString());
+    }
+  }
+
+  private static void assertAllColumnsAreBigDecimal(ResultSet rs) throws SQLException {
+    while (rs.next()) {
+      assertEquals("class java.math.BigDecimal", rs.getObject(1).getClass().toString());
+      assertEquals("class java.math.BigDecimal", rs.getObject(2).getClass().toString());
+      assertEquals("class java.math.BigDecimal", rs.getObject(3).getClass().toString());
+      assertEquals("class java.math.BigDecimal", rs.getObject(4).getClass().toString());
+    }
+  }
+
   // Test setting new connection property jdbc_arrow_treat_decimal_as_int=false. Connection property introduced after version 3.15.0.
   @Test
   public void testGetObjectForArrowResultFormatJDBCArrowDecimalAsIntFalse() throws SQLException {
@@ -993,23 +1011,13 @@ public void testGetObjectForArrowResultFormatJDBCArrowDecimalAsIntFalse() throws
 
       // Test with jdbc_arrow_treat_decimal_as_int=false and JDBC_TREAT_DECIMAL_AS_INT=true
       try (ResultSet rs = stmt.executeQuery(selectQuery)) {
-        while(rs.next()) {
-          assertEquals("class java.lang.Long", rs.getObject(1).getClass().toString());
-          assertEquals("class java.math.BigDecimal", rs.getObject(2).getClass().toString());
-          assertEquals("class java.lang.Long", rs.getObject(3).getClass().toString());
-          assertEquals("class java.lang.Long", rs.getObject(4).getClass().toString());
-        }
+        assertAllColumnsAreLongButBigIntIsBigDecimal(rs);
       }
 
       // Test with jdbc_arrow_treat_decimal_as_int=false and JDBC_TREAT_DECIMAL_AS_INT=false
       stmt.execute(setJdbcTreatDecimalAsIntFalse);
       try (ResultSet rs = stmt.executeQuery(selectQuery)) {
-        while (rs.next()) {
-          assertEquals("class java.math.BigDecimal", rs.getObject(1).getClass().toString());
-          assertEquals("class java.math.BigDecimal", rs.getObject(2).getClass().toString());
-          assertEquals("class java.math.BigDecimal", rs.getObject(3).getClass().toString());
-          assertEquals("class java.math.BigDecimal", rs.getObject(4).getClass().toString());
-        }
+        assertAllColumnsAreBigDecimal(rs);
       }
     }
   }
@@ -1036,80 +1044,29 @@ public void testGetObjectForArrowResultFormatJDBCArrowDecimalAsIntTrue() throws
       // Test with jdbc_arrow_treat_decimal_as_int=true and JDBC_TREAT_DECIMAL_AS_INT=false
       stmt.execute(setJdbcTreatDecimalAsIntFalse);
       try (ResultSet rs = stmt.executeQuery(selectQuery)) {
-        while (rs.next()) {
-          assertEquals("class java.lang.Long", rs.getObject(1).getClass().toString());
-          assertEquals("class java.math.BigDecimal", rs.getObject(2).getClass().toString());
-          assertEquals("class java.lang.Long", rs.getObject(3).getClass().toString());
-          assertEquals("class java.lang.Long", rs.getObject(4).getClass().toString());
-        }
+        assertAllColumnsAreLongButBigIntIsBigDecimal(rs);
       }
     }
   }
 
-  // Test setting new connection property jdbc_arrow_treat_decimal_as_int=false. Connection property
-  // introduced after version 3.15.0.
+  // Test getObject for numeric types when JDBC_TREAT_DECIMAL_AS_INT is set and using JSON result format.
   @Test
-  public void testGetObjectForJSONResultFormatJDBCArrowDecimalAsIntFalse() throws SQLException {
-    Properties properties = new Properties();
-    properties.put("jdbc_arrow_treat_decimal_as_int", false);
-    try (Connection con = getConnection(properties);
-        Statement stmt = con.createStatement()) {
-      stmt.execute("alter session set jdbc_query_result_format = 'JSON'");
-      stmt.execute(createTableSql);
-      stmt.execute(insertStmt);
-
-      // Test with jdbc_arrow_treat_decimal_as_int=false and JDBC_TREAT_DECIMAL_AS_INT=true
-      try (ResultSet rs = stmt.executeQuery(selectQuery)) {
-        while (rs.next()) {
-          assertEquals("class java.lang.Long", rs.getObject(1).getClass().toString());
-          assertEquals("class java.math.BigDecimal", rs.getObject(2).getClass().toString());
-          assertEquals("class java.lang.Long", rs.getObject(3).getClass().toString());
-          assertEquals("class java.lang.Long", rs.getObject(4).getClass().toString());
-        }
-      }
-
-      // Test with jdbc_arrow_treat_decimal_as_int=false and JDBC_TREAT_DECIMAL_AS_INT=false
-      stmt.execute(setJdbcTreatDecimalAsIntFalse);
-      try (ResultSet rs = stmt.executeQuery(selectQuery)) {
-        while (rs.next()) {
-          assertEquals("class java.math.BigDecimal", rs.getObject(1).getClass().toString());
-          assertEquals("class java.math.BigDecimal", rs.getObject(2).getClass().toString());
-          assertEquals("class java.math.BigDecimal", rs.getObject(3).getClass().toString());
-          assertEquals("class java.math.BigDecimal", rs.getObject(4).getClass().toString());
-        }
-      }
-    }
-  }
-
-  // Test default setting of new connection property jdbc_arrow_treat_decimal_as_int=true.
-  // Connection property introduced after version 3.15.0.
-  @Test
-  public void testGetObjectForJSONResultFormatJDBCArrowDecimalAsIntTrue() throws SQLException {
+  public void testGetObjectForJSONResultFormatJDBCDecimalAsIntFalse() throws SQLException {
     try (Connection con = BaseJDBCTest.getConnection();
         Statement stmt = con.createStatement()) {
       stmt.execute("alter session set jdbc_query_result_format = 'JSON'");
       stmt.execute(createTableSql);
       stmt.execute(insertStmt);
 
-      // Test with jdbc_arrow_treat_decimal_as_int=true and JDBC_TREAT_DECIMAL_AS_INT=true
+      // Test with JDBC_TREAT_DECIMAL_AS_INT=true (default value)
       try (ResultSet rs = stmt.executeQuery(selectQuery)) {
-        while (rs.next()) {
-          assertEquals("class java.lang.Long", rs.getObject(1).getClass().toString());
-          assertEquals("class java.math.BigDecimal", rs.getObject(2).getClass().toString());
-          assertEquals("class java.lang.Long", rs.getObject(3).getClass().toString());
-          assertEquals("class java.lang.Long", rs.getObject(4).getClass().toString());
-        }
+        assertAllColumnsAreLongButBigIntIsBigDecimal(rs);
       }
 
-      // Test with jdbc_arrow_treat_decimal_as_int=true and JDBC_TREAT_DECIMAL_AS_INT=false
+      // Test with JDBC_TREAT_DECIMAL_AS_INT=false
       stmt.execute(setJdbcTreatDecimalAsIntFalse);
       try (ResultSet rs = stmt.executeQuery(selectQuery)) {
-        while (rs.next()) {
-          assertEquals("class java.math.BigDecimal", rs.getObject(1).getClass().toString());
-          assertEquals("class java.math.BigDecimal", rs.getObject(2).getClass().toString());
-          assertEquals("class java.math.BigDecimal", rs.getObject(3).getClass().toString());
-          assertEquals("class java.math.BigDecimal", rs.getObject(4).getClass().toString());
-        }
+        assertAllColumnsAreBigDecimal(rs);
       }
     }
   }