Support TIMESTAMPS with precisions higher than 3 in MySQL

Author: guyco33

plugin/trino-mysql/src/main/java/io/trino/plugin/mysql/MySqlClient.java

@@ -25,7 +25,6 @@
 import io.trino.plugin.jdbc.JdbcSortItem;
 import io.trino.plugin.jdbc.JdbcTableHandle;
 import io.trino.plugin.jdbc.JdbcTypeHandle;
-import io.trino.plugin.jdbc.LongReadFunction;
 import io.trino.plugin.jdbc.PreparedQuery;
 import io.trino.plugin.jdbc.WriteMapping;
 import io.trino.plugin.jdbc.expression.AggregateFunctionRewriter;
@@ -54,6 +53,7 @@
 import io.trino.spi.type.DecimalType;
 import io.trino.spi.type.Decimals;
 import io.trino.spi.type.StandardTypes;
+import io.trino.spi.type.TimestampType;
 import io.trino.spi.type.Type;
 import io.trino.spi.type.TypeManager;
 import io.trino.spi.type.TypeSignature;
@@ -67,7 +67,6 @@
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Types;
-import java.time.LocalDateTime;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
@@ -101,11 +100,13 @@
 import static io.trino.plugin.jdbc.StandardColumnMappings.integerColumnMapping;
 import static io.trino.plugin.jdbc.StandardColumnMappings.integerWriteFunction;
 import static io.trino.plugin.jdbc.StandardColumnMappings.longDecimalWriteFunction;
+import static io.trino.plugin.jdbc.StandardColumnMappings.longTimestampWriteFunction;
 import static io.trino.plugin.jdbc.StandardColumnMappings.realColumnMapping;
 import static io.trino.plugin.jdbc.StandardColumnMappings.realWriteFunction;
 import static io.trino.plugin.jdbc.StandardColumnMappings.shortDecimalWriteFunction;
 import static io.trino.plugin.jdbc.StandardColumnMappings.smallintColumnMapping;
 import static io.trino.plugin.jdbc.StandardColumnMappings.smallintWriteFunction;
+import static io.trino.plugin.jdbc.StandardColumnMappings.timestampColumnMapping;
 import static io.trino.plugin.jdbc.StandardColumnMappings.timestampWriteFunction;
 import static io.trino.plugin.jdbc.StandardColumnMappings.tinyintColumnMapping;
 import static io.trino.plugin.jdbc.StandardColumnMappings.tinyintWriteFunction;
@@ -122,22 +123,25 @@
 import static io.trino.spi.type.RealType.REAL;
 import static io.trino.spi.type.SmallintType.SMALLINT;
 import static io.trino.spi.type.TimeWithTimeZoneType.TIME_WITH_TIME_ZONE;
-import static io.trino.spi.type.TimestampType.TIMESTAMP_MILLIS;
+import static io.trino.spi.type.TimestampType.createTimestampType;
 import static io.trino.spi.type.TimestampWithTimeZoneType.TIMESTAMP_TZ_MILLIS;
-import static io.trino.spi.type.Timestamps.MICROSECONDS_PER_MILLISECOND;
-import static io.trino.spi.type.Timestamps.MICROSECONDS_PER_SECOND;
-import static io.trino.spi.type.Timestamps.NANOSECONDS_PER_MILLISECOND;
 import static io.trino.spi.type.TinyintType.TINYINT;
 import static io.trino.spi.type.VarbinaryType.VARBINARY;
 import static java.lang.Math.max;
 import static java.lang.Math.min;
 import static java.lang.String.format;
-import static java.time.ZoneOffset.UTC;
 import static java.util.stream.Collectors.joining;
 
 public class MySqlClient
         extends BaseJdbcClient
 {
+    private static final int MAX_SUPPORTED_TIMESTAMP_PRECISION = 6;
+    // MySQL driver returns width of timestamp types instead of precision.
+    // 19 characters are used for zero-precision timestamps while others
+    // require 19 + precision + 1 characters with the additional character
+    // required for the decimal separator.
+    private static final int ZERO_PRECISION_TIMESTAMP_COLUMN_SIZE = 19;
+
     private final Type jsonType;
     private final AggregateFunctionRewriter aggregateFunctionRewriter;
 
@@ -319,29 +323,22 @@ public Optional<ColumnMapping> toColumnMapping(ConnectorSession session, Connect
                 return Optional.of(dateColumnMapping());
 
             case Types.TIMESTAMP:
-                // TODO support higher precisions (https://github.com/trinodb/trino/issues/6910)
-                return Optional.of(ColumnMapping.longMapping(
-                        TIMESTAMP_MILLIS,
-                        timestampMillisReadFunction(),
-                        timestampWriteFunction(TIMESTAMP_MILLIS),
-                        DISABLE_PUSHDOWN));
+                TimestampType timestampType = createTimestampType(getTimestampPrecision(typeHandle.getRequiredColumnSize()));
+                return Optional.of(timestampColumnMapping(timestampType));
         }
 
         // TODO add explicit mappings
         return legacyColumnMapping(session, connection, typeHandle);
     }
 
-    // TODO support higher precisions (https://github.com/trinodb/trino/issues/6910)
-    private static LongReadFunction timestampMillisReadFunction()
-    {
-        return (resultSet, columnIndex) -> toTrinoTimestampMillis(resultSet.getObject(columnIndex, LocalDateTime.class));
-    }
-
-    private static long toTrinoTimestampMillis(LocalDateTime localDateTime)
+    private static int getTimestampPrecision(int timestampColumnSize)
     {
-        return localDateTime.toEpochSecond(UTC) * MICROSECONDS_PER_SECOND
-                // rounding to millis, TODO support higher precisions (https://github.com/trinodb/trino/issues/6910)
-                + Math.round(localDateTime.getNano() * 1.0 / NANOSECONDS_PER_MILLISECOND) * MICROSECONDS_PER_MILLISECOND;
+        if (timestampColumnSize == ZERO_PRECISION_TIMESTAMP_COLUMN_SIZE) {
+            return 0;
+        }
+        int timestampPrecision = timestampColumnSize - ZERO_PRECISION_TIMESTAMP_COLUMN_SIZE - 1;
+        verify(1 <= timestampPrecision && timestampPrecision <= MAX_SUPPORTED_TIMESTAMP_PRECISION, "Unexpected timestamp precision %s calculated from timestamp column size %s", timestampPrecision, timestampColumnSize);
+        return timestampPrecision;
     }
 
     @Override
@@ -382,10 +379,16 @@ public WriteMapping toWriteMapping(ConnectorSession session, Type type)
         if (TIME_WITH_TIME_ZONE.equals(type) || TIMESTAMP_TZ_MILLIS.equals(type)) {
             throw new TrinoException(NOT_SUPPORTED, "Unsupported column type: " + type.getDisplayName());
         }
-        if (TIMESTAMP_MILLIS.equals(type)) {
-            // TODO use `timestampWriteFunction` (https://github.com/trinodb/trino/issues/6910)
-            return WriteMapping.longMapping("datetime(3)", timestampWriteFunction(TIMESTAMP_MILLIS));
+
+        if (type instanceof TimestampType) {
+            TimestampType timestampType = (TimestampType) type;
+            if (timestampType.getPrecision() <= MAX_SUPPORTED_TIMESTAMP_PRECISION) {
+                verify(timestampType.getPrecision() <= TimestampType.MAX_SHORT_PRECISION);
+                return WriteMapping.longMapping(format("datetime(%s)", timestampType.getPrecision()), timestampWriteFunction(timestampType));
+            }
+            return WriteMapping.objectMapping(format("datetime(%s)", MAX_SUPPORTED_TIMESTAMP_PRECISION), longTimestampWriteFunction(timestampType));
         }
+
         if (VARBINARY.equals(type)) {
             return WriteMapping.sliceMapping("mediumblob", varbinaryWriteFunction());
         }