Fix timestamp handling in JDBC connectors

presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/BaseJdbcClient.java

@@ -229,7 +229,7 @@ public List<JdbcColumnHandle> getColumns(ConnectorSession session, JdbcTableHand
     @Override
     public Optional<ColumnMapping> toPrestoType(ConnectorSession session, JdbcTypeHandle typeHandle)
     {
-        return jdbcTypeToPrestoType(typeHandle);
+        return jdbcTypeToPrestoType(session, typeHandle);
     }
 
     @Override

presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/StandardColumnMappings.java

@@ -16,6 +16,7 @@
 import com.google.common.base.CharMatcher;
 import com.google.common.primitives.Shorts;
 import com.google.common.primitives.SignedBytes;
+import io.prestosql.spi.connector.ConnectorSession;
 import io.prestosql.spi.type.CharType;
 import io.prestosql.spi.type.DecimalType;
 import io.prestosql.spi.type.Decimals;
@@ -32,6 +33,9 @@
 import java.sql.Time;
 import java.sql.Timestamp;
 import java.sql.Types;
+import java.time.Instant;
+import java.time.LocalDateTime;
+import java.time.ZoneId;
 import java.util.Optional;
 
 import static com.google.common.base.Preconditions.checkArgument;
@@ -61,10 +65,10 @@
 import static java.lang.Math.max;
 import static java.lang.Math.min;
 import static java.lang.Math.toIntExact;
+import static java.time.ZoneOffset.UTC;
 import static java.util.Objects.requireNonNull;
 import static java.util.concurrent.TimeUnit.DAYS;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
-import static org.joda.time.DateTimeZone.UTC;
 
 public final class StandardColumnMappings
 {
@@ -235,7 +239,7 @@ public static ColumnMapping dateColumnMapping()
                      */
                     long localMillis = resultSet.getDate(columnIndex).getTime();
                     // Convert it to a ~midnight in UTC.
-                    long utcMillis = ISOChronology.getInstance().getZone().getMillisKeepLocal(UTC, localMillis);
+                    long utcMillis = ISOChronology.getInstance().getZone().getMillisKeepLocal(DateTimeZone.UTC, localMillis);
                     // convert to days
                     return MILLISECONDS.toDays(utcMillis);
                 },
@@ -247,7 +251,7 @@ public static LongWriteFunction dateWriteFunction()
         return (statement, index, value) -> {
             // convert to midnight in default time zone
             long millis = DAYS.toMillis(value);
-            statement.setDate(index, new Date(UTC.getMillisKeepLocal(DateTimeZone.getDefault(), millis)));
+            statement.setDate(index, new Date(DateTimeZone.UTC.getMillisKeepLocal(DateTimeZone.getDefault(), millis)));
         };
     }
 
@@ -276,33 +280,107 @@ public static LongWriteFunction timeWriteFunction()
         };
     }
 
-    public static ColumnMapping timestampColumnMapping()
+    /**
+     * @deprecated This method uses {@link java.sql.Timestamp} and the class cannot represent date-time value when JVM zone had
+     * forward offset change (a 'gap'). This includes regular DST changes (e.g. Europe/Warsaw) and one-time policy changes
+     * (Asia/Kathmandu's shift by 15 minutes on January 1, 1986, 00:00:00). If driver only supports {@link LocalDateTime}, use
+     * {@link #timestampColumnMapping} instead.
+     */
+    @Deprecated
+    public static ColumnMapping timestampColumnMappingUsingSqlTimestamp(ConnectorSession session)
     {
+        if (session.isLegacyTimestamp()) {
+            ZoneId sessionZone = ZoneId.of(session.getTimeZoneKey().getId());
+            return ColumnMapping.longMapping(
+                    TIMESTAMP,
+                    (resultSet, columnIndex) -> {
+                        Timestamp timestamp = resultSet.getTimestamp(columnIndex);
+                        return toPrestoLegacyTimestamp(timestamp.toLocalDateTime(), sessionZone);
+                    },
+                    timestampWriteFunctionUsingSqlTimestamp(session));
+        }
+
         return ColumnMapping.longMapping(
                 TIMESTAMP,
                 (resultSet, columnIndex) -> {
-                    /*
-                     * TODO `resultSet.getTimestamp(columnIndex)` returns wrong value if JVM's zone had forward offset change and the local time
-                     * corresponding to timestamp value being retrieved was not present (a 'gap'), this includes regular DST changes (e.g. Europe/Warsaw)
-                     * and one-time policy changes (Asia/Kathmandu's shift by 15 minutes on January 1, 1986, 00:00:00).
-                     * The problem can be averted by using `resultSet.getObject(columnIndex, LocalDateTime.class)` -- but this is not universally supported by JDBC drivers.
-                     */
                     Timestamp timestamp = resultSet.getTimestamp(columnIndex);
-                    return timestamp.getTime();
+                    return toPrestoTimestamp(timestamp.toLocalDateTime());
                 },
-                timestampWriteFunction());
+                timestampWriteFunctionUsingSqlTimestamp(session));
+    }
+
+    public static ColumnMapping timestampColumnMapping(ConnectorSession session)
+    {
+        if (session.isLegacyTimestamp()) {
+            ZoneId sessionZone = ZoneId.of(session.getTimeZoneKey().getId());
+            return ColumnMapping.longMapping(
+                    TIMESTAMP,
+                    (resultSet, columnIndex) -> toPrestoLegacyTimestamp(resultSet.getObject(columnIndex, LocalDateTime.class), sessionZone),
+                    timestampWriteFunction(session));
+        }
+
+        return ColumnMapping.longMapping(
+                TIMESTAMP,
+                (resultSet, columnIndex) -> toPrestoTimestamp(resultSet.getObject(columnIndex, LocalDateTime.class)),
+                timestampWriteFunction(session));
+    }
+
+    /**
+     * @deprecated This method uses {@link java.sql.Timestamp} and the class cannot represent date-time value when JVM zone had
+     * forward offset change (a 'gap'). This includes regular DST changes (e.g. Europe/Warsaw) and one-time policy changes
+     * (Asia/Kathmandu's shift by 15 minutes on January 1, 1986, 00:00:00). If driver only supports {@link LocalDateTime}, use
+     * {@link #timestampWriteFunction} instead.
+     */
+    @Deprecated
+    public static LongWriteFunction timestampWriteFunctionUsingSqlTimestamp(ConnectorSession connectorSession)
+    {
+        if (connectorSession.isLegacyTimestamp()) {
+            ZoneId sessionZone = ZoneId.of(connectorSession.getTimeZoneKey().getId());
+            return (statement, index, value) -> statement.setTimestamp(index, Timestamp.valueOf(fromPrestoLegacyTimestamp(value, sessionZone)));
+        }
+        return (statement, index, value) -> statement.setTimestamp(index, Timestamp.valueOf(fromPrestoTimestamp(value)));
     }
 
-    public static LongWriteFunction timestampWriteFunction()
+    public static LongWriteFunction timestampWriteFunction(ConnectorSession session)
     {
+        if (session.isLegacyTimestamp()) {
+            ZoneId sessionZone = ZoneId.of(session.getTimeZoneKey().getId());
+            return (statement, index, value) -> statement.setObject(index, fromPrestoLegacyTimestamp(value, sessionZone));
+        }
         return (statement, index, value) -> {
-            // Copied from `QueryBuilder.buildSql`
-            // TODO verify correctness, add tests and support non-legacy timestamp
-            statement.setTimestamp(index, new Timestamp(value));
+            statement.setObject(index, fromPrestoTimestamp(value));
         };
     }
 
-    public static Optional<ColumnMapping> jdbcTypeToPrestoType(JdbcTypeHandle type)
+    /**
+     * @deprecated applicable in legacy timestamp semantics only
+     */
+    @Deprecated
+    private static long toPrestoLegacyTimestamp(LocalDateTime localDateTime, ZoneId sessionZone)
+    {
+        return localDateTime.atZone(sessionZone).toInstant().toEpochMilli();
+    }
+
+    private static long toPrestoTimestamp(LocalDateTime localDateTime)
+    {
+        return localDateTime.atZone(UTC).toInstant().toEpochMilli();
+    }
+
+    /**
+     * @deprecated applicable in legacy timestamp semantics only
+     */
+    @Deprecated
+    private static LocalDateTime fromPrestoLegacyTimestamp(long value, ZoneId sessionZone)
+    {
+        return Instant.ofEpochMilli(value).atZone(sessionZone).toLocalDateTime();
+    }
+
+    private static LocalDateTime fromPrestoTimestamp(long value)
+    {
+        return Instant.ofEpochMilli(value).atZone(UTC).toLocalDateTime();
+    }
+
+    public static Optional<ColumnMapping> jdbcTypeToPrestoType(ConnectorSession session, JdbcTypeHandle type)
     {
         int columnSize = type.getColumnSize();
         switch (type.getJdbcType()) {
@@ -365,7 +443,8 @@ public static Optional<ColumnMapping> jdbcTypeToPrestoType(JdbcTypeHandle type)
                 return Optional.of(timeColumnMapping());
 
             case Types.TIMESTAMP:
-                return Optional.of(timestampColumnMapping());
+                // TODO default to `timestampColumnMapping`
+                return Optional.of(timestampColumnMappingUsingSqlTimestamp(session));
         }
         return Optional.empty();
     }

presto-base-jdbc/src/test/java/io/prestosql/plugin/jdbc/TestJdbcQueryBuilder.java

@@ -22,6 +22,8 @@
 import io.prestosql.spi.predicate.SortedRangeSet;
 import io.prestosql.spi.predicate.TupleDomain;
 import io.prestosql.spi.type.CharType;
+import io.prestosql.spi.type.SqlTimestamp;
+import io.prestosql.testing.DateTimeTestingUtils;
 import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
@@ -61,13 +63,15 @@
 import static io.prestosql.spi.type.RealType.REAL;
 import static io.prestosql.spi.type.SmallintType.SMALLINT;
 import static io.prestosql.spi.type.TimeType.TIME;
+import static io.prestosql.spi.type.TimeZoneKey.UTC_KEY;
 import static io.prestosql.spi.type.TimestampType.TIMESTAMP;
 import static io.prestosql.spi.type.TinyintType.TINYINT;
 import static io.prestosql.spi.type.VarcharType.VARCHAR;
 import static io.prestosql.testing.TestingConnectorSession.SESSION;
 import static java.lang.Float.floatToRawIntBits;
 import static java.lang.String.format;
 import static java.time.temporal.ChronoUnit.DAYS;
+import static org.joda.time.DateTimeZone.UTC;
 import static org.testng.Assert.assertEquals;
 
 @Test(singleThreaded = true)
@@ -335,9 +339,9 @@ public void testBuildSqlWithTimestamp()
         TupleDomain<ColumnHandle> tupleDomain = TupleDomain.withColumnDomains(ImmutableMap.of(
                 columns.get(6), Domain.create(SortedRangeSet.copyOf(TIMESTAMP,
                         ImmutableList.of(
-                                Range.equal(TIMESTAMP, toTimestamp(2016, 6, 3, 0, 23, 37).getTime()),
-                                Range.equal(TIMESTAMP, toTimestamp(2016, 10, 19, 16, 23, 37).getTime()),
-                                Range.range(TIMESTAMP, toTimestamp(2016, 6, 7, 8, 23, 37).getTime(), false, toTimestamp(2016, 6, 9, 12, 23, 37).getTime(), true))),
+                                Range.equal(TIMESTAMP, toPrestoTimestamp(2016, 6, 3, 0, 23, 37)),
+                                Range.equal(TIMESTAMP, toPrestoTimestamp(2016, 10, 19, 16, 23, 37)),
+                                Range.range(TIMESTAMP, toPrestoTimestamp(2016, 6, 7, 8, 23, 37), false, toPrestoTimestamp(2016, 6, 9, 12, 23, 37), true))),
                         false)));
 
         Connection connection = database.getConnection();
@@ -374,6 +378,15 @@ public void testEmptyBuildSql()
         }
     }
 
+    private static long toPrestoTimestamp(int year, int month, int day, int hour, int minute, int second)
+    {
+        SqlTimestamp sqlTimestamp = DateTimeTestingUtils.sqlTimestampOf(year, month, day, hour, minute, second, 0, UTC, UTC_KEY, SESSION);
+        if (SESSION.isLegacyTimestamp()) {
+            return sqlTimestamp.getMillisUtc();
+        }
+        return sqlTimestamp.getMillis();
+    }
+
     private static Timestamp toTimestamp(int year, int month, int day, int hour, int minute, int second)
     {
         return Timestamp.valueOf(LocalDateTime.of(year, month, day, hour, minute, second));

presto-mysql/src/main/java/io/prestosql/plugin/mysql/MySqlClient.java

@@ -48,7 +48,7 @@
 import static io.prestosql.plugin.jdbc.DriverConnectionFactory.basicConnectionProperties;
 import static io.prestosql.plugin.jdbc.JdbcErrorCode.JDBC_ERROR;
 import static io.prestosql.plugin.jdbc.StandardColumnMappings.realWriteFunction;
-import static io.prestosql.plugin.jdbc.StandardColumnMappings.timestampWriteFunction;
+import static io.prestosql.plugin.jdbc.StandardColumnMappings.timestampWriteFunctionUsingSqlTimestamp;
 import static io.prestosql.plugin.jdbc.StandardColumnMappings.varbinaryWriteFunction;
 import static io.prestosql.plugin.jdbc.StandardColumnMappings.varcharWriteFunction;
 import static io.prestosql.spi.StandardErrorCode.ALREADY_EXISTS;
@@ -167,7 +167,8 @@ public WriteMapping toWriteMapping(ConnectorSession session, Type type)
             throw new PrestoException(NOT_SUPPORTED, "Unsupported column type: " + type.getDisplayName());
         }
         if (TIMESTAMP.equals(type)) {
-            return WriteMapping.longMapping("datetime", timestampWriteFunction());
+            // TODO use `timestampWriteFunction`
+            return WriteMapping.longMapping("datetime", timestampWriteFunctionUsingSqlTimestamp(session));
         }
         if (VARBINARY.equals(type)) {
             return WriteMapping.sliceMapping("mediumblob", varbinaryWriteFunction());

presto-mysql/src/test/java/io/prestosql/plugin/mysql/TestMySqlTypeMapping.java

@@ -251,14 +251,14 @@ public void testDate()
     @Test
     public void testDatetime()
     {
-        // TODO MySQL datetime is not correctly read (see comment in StandardColumnMappings.timestampColumnMapping)
+        // TODO MySQL datetime is not correctly read (see comment in StandardColumnMappings.timestampColumnMappingUsingSqlTimestamp)
         // testing this is hard because of https://github.com/prestodb/presto/issues/7122
     }
 
     @Test
     public void testTimestamp()
     {
-        // TODO MySQL timestamp is not correctly read (see comment in StandardColumnMappings.timestampColumnMapping)
+        // TODO MySQL timestamp is not correctly read (see comment in StandardColumnMappings.timestampColumnMappingUsingSqlTimestamp)
         // testing this is hard because of https://github.com/prestodb/presto/issues/7122
     }
 

presto-postgresql/src/main/java/io/prestosql/plugin/postgresql/PostgreSqlClient.java

@@ -58,10 +58,13 @@
 import static io.airlift.slice.Slices.utf8Slice;
 import static io.prestosql.plugin.jdbc.ColumnMapping.DISABLE_PUSHDOWN;
 import static io.prestosql.plugin.jdbc.JdbcErrorCode.JDBC_ERROR;
+import static io.prestosql.plugin.jdbc.StandardColumnMappings.timestampColumnMapping;
+import static io.prestosql.plugin.jdbc.StandardColumnMappings.timestampWriteFunction;
 import static io.prestosql.plugin.jdbc.StandardColumnMappings.varbinaryWriteFunction;
 import static io.prestosql.spi.StandardErrorCode.ALREADY_EXISTS;
 import static io.prestosql.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT;
 import static io.prestosql.spi.StandardErrorCode.NOT_SUPPORTED;
+import static io.prestosql.spi.type.TimestampType.TIMESTAMP;
 import static io.prestosql.spi.type.VarbinaryType.VARBINARY;
 import static java.lang.String.format;
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -147,6 +150,9 @@ public Optional<ColumnMapping> toPrestoType(ConnectorSession session, JdbcTypeHa
             // This can be e.g. an ENUM
             return Optional.of(typedVarcharColumnMapping(typeHandle.getJdbcTypeName()));
         }
+        if (typeHandle.getJdbcType() == Types.TIMESTAMP) {
+            return Optional.of(timestampColumnMapping(session));
+        }
         // TODO support PostgreSQL's TIMESTAMP WITH TIME ZONE and TIME WITH TIME ZONE explicitly, otherwise predicate pushdown for these types may be incorrect
         return super.toPrestoType(session, typeHandle);
     }
@@ -157,6 +163,9 @@ public WriteMapping toWriteMapping(ConnectorSession session, Type type)
         if (VARBINARY.equals(type)) {
             return WriteMapping.sliceMapping("bytea", varbinaryWriteFunction());
         }
+        if (TIMESTAMP.equals(type)) {
+            return WriteMapping.longMapping("timestamp", timestampWriteFunction(session));
+        }
         if (type.getTypeSignature().getBase().equals(StandardTypes.JSON)) {
             return WriteMapping.sliceMapping("jsonb", typedVarcharWriteFunction("json"));
         }