SNOW-1465374: Return consistent timestamps_ltz between JSON and ARROW…

… result sets (#1792)
snowflakedb · Jul 2, 2024 · 8e916ae · 8e916ae
1 parent 3fee6f9
commit 8e916ae
Show file tree

Hide file tree

Showing 4 changed files with 134 additions and 3 deletions.
diff --git a/src/main/java/net/snowflake/client/core/arrow/ArrowResultUtil.java b/src/main/java/net/snowflake/client/core/arrow/ArrowResultUtil.java
@@ -11,6 +11,7 @@
 import java.util.TimeZone;
 import net.snowflake.client.core.ResultUtil;
 import net.snowflake.client.core.SFException;
+import net.snowflake.client.core.SnowflakeJdbcInternalApi;
 import net.snowflake.client.jdbc.ErrorCode;
 import net.snowflake.client.jdbc.SnowflakeTimestampWithTimezone;
 import net.snowflake.client.log.ArgSupplier;
@@ -151,14 +152,27 @@ public static Timestamp moveToTimeZone(Timestamp ts, TimeZone oldTZ, TimeZone ne
    * @return
    */
   public static Timestamp toJavaTimestamp(long epoch, int scale) {
+    return toJavaTimestamp(epoch, scale, TimeZone.getDefault(), false);
+  }
+
+  /**
+   * generate Java Timestamp object
+   *
+   * @param epoch the value since epoch time
+   * @param scale the scale of the value
+   * @return
+   */
+  @SnowflakeJdbcInternalApi
+  public static Timestamp toJavaTimestamp(
+      long epoch, int scale, TimeZone sessionTimezone, boolean useSessionTimezone) {
     long seconds = epoch / powerOfTen(scale);
     int fraction = (int) ((epoch % powerOfTen(scale)) * powerOfTen(9 - scale));
     if (fraction < 0) {
       // handle negative case here
       seconds--;
       fraction += 1000000000;
     }
-    return createTimestamp(seconds, fraction, TimeZone.getDefault(), false);
+    return createTimestamp(seconds, fraction, sessionTimezone, useSessionTimezone);
   }
 
   /**

diff --git a/src/main/java/net/snowflake/client/core/arrow/BigIntToTimestampLTZConverter.java b/src/main/java/net/snowflake/client/core/arrow/BigIntToTimestampLTZConverter.java
@@ -11,6 +11,7 @@
 import net.snowflake.client.core.DataConversionContext;
 import net.snowflake.client.core.ResultUtil;
 import net.snowflake.client.core.SFException;
+import net.snowflake.client.core.SnowflakeJdbcInternalApi;
 import net.snowflake.client.jdbc.ErrorCode;
 import net.snowflake.client.jdbc.SnowflakeType;
 import net.snowflake.client.jdbc.SnowflakeUtil;
@@ -65,7 +66,7 @@ public Timestamp toTimestamp(int index, TimeZone tz) throws SFException {
   private Timestamp getTimestamp(int index, TimeZone tz) throws SFException {
     long val = bigIntVector.getDataBuffer().getLong(index * BigIntVector.TYPE_WIDTH);
     int scale = context.getScale(columnIndex);
-    return getTimestamp(val, scale);
+    return getTimestamp(val, scale, sessionTimeZone, useSessionTimezone);
   }
 
   @Override
@@ -90,8 +91,25 @@ public boolean toBoolean(int index) throws SFException {
         SnowflakeUtil.BOOLEAN_STR, val);
   }
 
+  /**
+   * Use {@link #getTimestamp(long, int, TimeZone, boolean)}
+   *
+   * @param val epoch
+   * @param scale scale
+   * @return Timestamp value without timezone take into account
+   * @throws SFException
+   */
+  @Deprecated
   public static Timestamp getTimestamp(long val, int scale) throws SFException {
     Timestamp ts = ArrowResultUtil.toJavaTimestamp(val, scale);
     return ResultUtil.adjustTimestamp(ts);
   }
+
+  @SnowflakeJdbcInternalApi
+  public static Timestamp getTimestamp(
+      long epoch, int scale, TimeZone sessionTimeZone, boolean useSessionTimezone)
+      throws SFException {
+    return ResultUtil.adjustTimestamp(
+        ArrowResultUtil.toJavaTimestamp(epoch, scale, sessionTimeZone, useSessionTimezone));
+  }
 }
diff --git a/src/main/java/net/snowflake/client/core/arrow/StructuredTypeDateTimeConverter.java b/src/main/java/net/snowflake/client/core/arrow/StructuredTypeDateTimeConverter.java
@@ -92,7 +92,8 @@ private Timestamp convertTimestampLtz(Object obj, int scale) throws SFException
             false);
       }
     } else if (obj instanceof Long) {
-      return BigIntToTimestampLTZConverter.getTimestamp((long) obj, scale);
+      return BigIntToTimestampLTZConverter.getTimestamp(
+          (long) obj, scale, sessionTimeZone, useSessionTimezone);
     }
     throw new SFException(
         ErrorCode.INVALID_VALUE_CONVERT,

diff --git a/src/test/java/net/snowflake/client/jdbc/ConnectionLatestIT.java b/src/test/java/net/snowflake/client/jdbc/ConnectionLatestIT.java
@@ -40,6 +40,7 @@
 import java.sql.Statement;
 import java.time.Duration;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.Enumeration;
 import java.util.List;
@@ -59,6 +60,8 @@
 import net.snowflake.client.core.SecurityUtil;
 import net.snowflake.client.core.SessionUtil;
 import net.snowflake.client.jdbc.telemetryOOB.TelemetryService;
+import net.snowflake.client.log.SFLogger;
+import net.snowflake.client.log.SFLoggerFactory;
 import net.snowflake.common.core.ClientAuthnDTO;
 import net.snowflake.common.core.ClientAuthnParameter;
 import net.snowflake.common.core.SqlState;
@@ -82,6 +85,7 @@
 @Category(TestCategoryConnection.class)
 public class ConnectionLatestIT extends BaseJDBCTest {
   @Rule public TemporaryFolder tmpFolder = new TemporaryFolder();
+  private static final SFLogger logger = SFLoggerFactory.getLogger(ConnectionLatestIT.class);
 
   private boolean defaultState;
 
@@ -1374,4 +1378,98 @@ public void testDataSourceOktaGenerates429StatusCode() throws Exception {
       thread.join();
     }
   }
+
+  /**
+   * SNOW-1465374: For TIMESTAMP_LTZ we were returning timestamps without timezone when scale was
+   * set e.g. to 6 in Arrow format The problem wasn't visible when calling getString, but was
+   * visible when we called toString on passed getTimestamp since we returned {@link
+   * java.sql.Timestamp}, not {@link SnowflakeTimestampWithTimezone}
+   *
+   * <p>Timestamps before 1582-10-05 are always returned as {@link java.sql.Timestamp}, not {@link
+   * SnowflakeTimestampWithTimezone} {SnowflakeTimestampWithTimezone}
+   *
+   * <p>Added in > 3.16.1
+   */
+  @Test
+  public void shouldGetDifferentTimestampLtzConsistentBetweenFormats() throws Exception {
+    try (Connection connection = getConnection();
+        Statement statement = connection.createStatement()) {
+      statement.executeUpdate(
+          "create or replace table DATETIMETZ_TYPE(timestamp_tzcol timestamp_ltz, timestamp_tzpcol timestamp_ltz(6), timestamptzcol timestampltz, timestampwtzcol timestamp with local time zone);");
+      Arrays.asList(
+              "insert into DATETIMETZ_TYPE values('9999-12-31 23:59:59.999999999','9999-12-31 23:59:59.999999','9999-12-31 23:59:59.999999999','9999-12-31 23:59:59.999999999');",
+              "insert into DATETIMETZ_TYPE values('1582-01-01 00:00:00.000000001','1582-01-01 00:00:00.000001','1582-01-01 00:00:00.000000001','1582-01-01 00:00:00.000000001');",
+              "insert into DATETIMETZ_TYPE values('2000-06-18 18:29:30.123456789 +0100','2000-06-18 18:29:30.123456 +0100','2000-06-18 18:29:30.123456789 +0100','2000-06-18 18:29:30.123456789 +0100');",
+              "insert into DATETIMETZ_TYPE values(current_timestamp(),current_timestamp(),current_timestamp(),current_timestamp());",
+              "insert into DATETIMETZ_TYPE values('2000-06-18 18:29:30.12345 -0530','2000-06-18 18:29:30.123 -0530','2000-06-18 18:29:30.123456 -0530','2000-06-18 18:29:30.123 -0530');",
+              "insert into DATETIMETZ_TYPE values('2000-06-18 18:29:30','2000-06-18 18:29:30','2000-06-18 18:29:30','2000-06-18 18:29:30');",
+              "insert into DATETIMETZ_TYPE values('1582-10-04 00:00:00.000000001','1582-10-04 00:00:00.000001','1582-10-04 00:00:00.000000001','1582-10-04 00:00:00.000000001');",
+              "insert into DATETIMETZ_TYPE values('1582-10-05 00:00:00.000000001','1582-10-05 00:00:00.000001','1582-10-05 00:00:00.000000001','1582-10-05 00:00:00.000000001');",
+              "insert into DATETIMETZ_TYPE values('1583-10-05 00:00:00.000000001','1583-10-05 00:00:00.000001','1583-10-05 00:00:00.000000001','1583-10-05 00:00:00.000000001');")
+          .forEach(
+              insert -> {
+                try {
+                  statement.executeUpdate(insert);
+                } catch (SQLException e) {
+                  throw new RuntimeException(e);
+                }
+              });
+      try (ResultSet arrowResultSet = statement.executeQuery("select * from DATETIMETZ_TYPE")) {
+        try (Connection jsonConnection = getConnection();
+            Statement jsonStatement = jsonConnection.createStatement()) {
+          jsonStatement.execute("alter session set JDBC_QUERY_RESULT_FORMAT=JSON");
+          try (ResultSet jsonResultSet =
+              jsonStatement.executeQuery("select * from DATETIMETZ_TYPE")) {
+            int rowIdx = 0;
+            while (arrowResultSet.next()) {
+              logger.debug("Checking row " + rowIdx);
+              assertTrue(jsonResultSet.next());
+              for (int column = 1; column <= 4; ++column) {
+                logger.trace(
+                    "JSON row[{}],column[{}] as string '{}', timestamp string '{}', as timestamp numeric '{}', tz offset={}, timestamp class {}",
+                    rowIdx,
+                    column,
+                    jsonResultSet.getString(column),
+                    jsonResultSet.getTimestamp(column),
+                    jsonResultSet.getTimestamp(column).getTime(),
+                    jsonResultSet.getTimestamp(column).getTimezoneOffset(),
+                    jsonResultSet.getTimestamp(column).getClass());
+                logger.trace(
+                    "ARROW row[{}],column[{}] as string '{}', timestamp string '{}', as timestamp numeric '{}', tz offset={}, timestamp class {}",
+                    rowIdx,
+                    column,
+                    arrowResultSet.getString(column),
+                    arrowResultSet.getTimestamp(column),
+                    arrowResultSet.getTimestamp(column).getTime(),
+                    arrowResultSet.getTimestamp(column).getTimezoneOffset(),
+                    arrowResultSet.getTimestamp(column).getClass());
+                assertEquals(
+                    "Expecting that string representation are the same for row "
+                        + rowIdx
+                        + " and column "
+                        + column,
+                    jsonResultSet.getString(column),
+                    arrowResultSet.getString(column));
+                assertEquals(
+                    "Expecting that string representation (via toString) are the same for row "
+                        + rowIdx
+                        + " and column "
+                        + column,
+                    jsonResultSet.getTimestamp(column).toString(),
+                    arrowResultSet.getTimestamp(column).toString());
+                assertEquals(
+                    "Expecting that timestamps are the same for row "
+                        + rowIdx
+                        + " and column "
+                        + column,
+                    jsonResultSet.getTimestamp(column),
+                    arrowResultSet.getTimestamp(column));
+              }
+              rowIdx++;
+            }
+          }
+        }
+      }
+    }
+  }
 }