Add support for TIMESTAMP types in Exasol connector

Author: skyglass

docs/src/main/sphinx/connector/exasol.md

@@ -101,6 +101,13 @@ Trino data type mapping:
   * - ``DATE``
     - ``DATE``
     -
+  * - ``TIMESTAMP(n)``
+    - ``TIMESTAMP(n)``
+    -  See :ref:`exasol-timestamp-mapping`
+  * - ``TIMESTAMP(n) WITH LOCAL TIME ZONE``
+    - ``TIMESTAMP(n) WITH TIME ZONE``
+    -  See :ref:`exasol-timestamp-with-local-time-zone-mapping`
+    
 ```
 
 No other types are supported.
@@ -127,6 +134,45 @@ If no length is specified, the connector uses 2000000.
 Trino's `CHAR(n)` maps to `CHAR(n)` and vice versa if `n` is no greater than 2000.
 Exasol does not support longer values.
 
+(exasol-timestamp-mapping)=
+### Mapping `TIMESTAMP` Types
+
+Exasol `TIMESTAMP(n)` columns are mapped to Trino's `TIMESTAMP(n)` type 
+and vice versa, with the following exceptions:
+
+- **No precision specified**:  
+  If the precision is omitted (i.e., the column is defined as `TIMESTAMP` 
+  without `(n)`), Exasol defaults to a precision of 3. 
+  This maps to Trino's `TIMESTAMP(3)`.
+
+- **Precision greater than 9**:  
+  Exasol supports up to 9 digits of fractional seconds. 
+  If the precision in Exasol exceeds 9, an exception is thrown
+
+- **Negative precision**:  
+  Negative values for precision are invalid. If encountered, 
+  an exception is thrown.
+
+(exasol-timestamp-with-local-time-zone-mapping)=
+### Mapping `TIMESTAMP WITH LOCAL TIME ZONE` Types
+
+Exasol `TIMESTAMP(n) WITH LOCAL TIME ZONE` columns are mapped 
+to Trino's `TIMESTAMP(n) WITH TIME ZONE` type and vice versa, 
+with the following exceptions:
+
+- **No precision specified**:  
+  If the precision is omitted (i.e., the column is defined 
+  as `TIMESTAMP WITH LOCAL TIME ZONE` without `(n)`), Exasol defaults 
+  to a precision of 3. This maps to Trino's `TIMESTAMP(3) WITH TIME ZONE`.
+
+- **Precision greater than 9**:  
+  Exasol supports up to 9 digits of fractional seconds. 
+  If the precision in Exasol exceeds 9, an exception is thrown.
+
+- **Negative precision**:  
+  Negative values for precision are invalid. 
+  If encountered, an exception is thrown.
+ 
 ```{include} jdbc-type-mapping.fragment
 ```
 

plugin/trino-exasol/src/main/java/io/trino/plugin/exasol/ExasolClient.java

@@ -28,6 +28,8 @@
 import io.trino.plugin.jdbc.JdbcTypeHandle;
 import io.trino.plugin.jdbc.LongReadFunction;
 import io.trino.plugin.jdbc.LongWriteFunction;
+import io.trino.plugin.jdbc.ObjectReadFunction;
+import io.trino.plugin.jdbc.ObjectWriteFunction;
 import io.trino.plugin.jdbc.QueryBuilder;
 import io.trino.plugin.jdbc.WriteFunction;
 import io.trino.plugin.jdbc.WriteMapping;
@@ -39,42 +41,68 @@
 import io.trino.spi.connector.ColumnPosition;
 import io.trino.spi.connector.ConnectorSession;
 import io.trino.spi.connector.ConnectorTableMetadata;
+import io.trino.spi.type.LongTimestamp;
+import io.trino.spi.type.LongTimestampWithTimeZone;
+import io.trino.spi.type.TimestampType;
+import io.trino.spi.type.TimestampWithTimeZoneType;
 import io.trino.spi.type.Type;
 
 import java.sql.Connection;
 import java.sql.Date;
+import java.sql.PreparedStatement;
+import java.sql.SQLException;
+import java.sql.Timestamp;
 import java.sql.Types;
+import java.time.Instant;
 import java.time.LocalDate;
+import java.time.LocalDateTime;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.OptionalLong;
 import java.util.Set;
 
+import static com.google.common.base.Preconditions.checkArgument;
+import static io.trino.plugin.jdbc.PredicatePushdownController.FULL_PUSHDOWN;
 import static io.trino.plugin.jdbc.StandardColumnMappings.bigintColumnMapping;
 import static io.trino.plugin.jdbc.StandardColumnMappings.booleanColumnMapping;
 import static io.trino.plugin.jdbc.StandardColumnMappings.decimalColumnMapping;
 import static io.trino.plugin.jdbc.StandardColumnMappings.defaultCharColumnMapping;
 import static io.trino.plugin.jdbc.StandardColumnMappings.defaultVarcharColumnMapping;
 import static io.trino.plugin.jdbc.StandardColumnMappings.doubleColumnMapping;
+import static io.trino.plugin.jdbc.StandardColumnMappings.fromLongTrinoTimestamp;
+import static io.trino.plugin.jdbc.StandardColumnMappings.fromTrinoTimestamp;
 import static io.trino.plugin.jdbc.StandardColumnMappings.integerColumnMapping;
 import static io.trino.plugin.jdbc.StandardColumnMappings.smallintColumnMapping;
+import static io.trino.plugin.jdbc.StandardColumnMappings.toLongTrinoTimestamp;
+import static io.trino.plugin.jdbc.StandardColumnMappings.toTrinoTimestamp;
 import static io.trino.plugin.jdbc.TypeHandlingJdbcSessionProperties.getUnsupportedTypeHandling;
 import static io.trino.plugin.jdbc.UnsupportedTypeHandling.CONVERT_TO_VARCHAR;
 import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED;
 import static io.trino.spi.connector.ConnectorMetadata.MODIFYING_ROWS_MESSAGE;
+import static io.trino.spi.type.DateTimeEncoding.packDateTimeWithZone;
 import static io.trino.spi.type.DateType.DATE;
 import static io.trino.spi.type.DecimalType.createDecimalType;
+import static io.trino.spi.type.TimeZoneKey.UTC_KEY;
+import static io.trino.spi.type.TimestampType.createTimestampType;
+import static io.trino.spi.type.TimestampWithTimeZoneType.createTimestampWithTimeZoneType;
+import static io.trino.spi.type.Timestamps.NANOSECONDS_PER_MILLISECOND;
+import static io.trino.spi.type.Timestamps.PICOSECONDS_PER_NANOSECOND;
+import static java.lang.String.format;
 import static java.util.Locale.ENGLISH;
 
 public class ExasolClient
         extends BaseJdbcClient
 {
+    private static final int EXASOL_TIMESTAMP_WITH_TIMEZONE = 124;
+
     private static final Set<String> INTERNAL_SCHEMAS = ImmutableSet.<String>builder()
             .add("EXA_STATISTICS")
             .add("SYS")
             .build();
 
+    private static final int MAX_EXASOL_TIMESTAMP_PRECISION = 9;
+
     @Inject
     public ExasolClient(
             BaseJdbcConfig config,
@@ -229,6 +257,10 @@ public Optional<ColumnMapping> toColumnMapping(ConnectorSession session, Connect
                 return Optional.of(defaultVarcharColumnMapping(typeHandle.requiredColumnSize(), true));
             case Types.DATE:
                 return Optional.of(dateColumnMapping());
+            case Types.TIMESTAMP:
+                return Optional.of(timestampColumnMapping(typeHandle));
+            case EXASOL_TIMESTAMP_WITH_TIMEZONE:
+                return Optional.of(timestampWithTimeZoneColumnMapping(typeHandle));
         }
 
         if (getUnsupportedTypeHandling(session) == CONVERT_TO_VARCHAR) {
@@ -237,6 +269,247 @@ public Optional<ColumnMapping> toColumnMapping(ConnectorSession session, Connect
         return Optional.empty();
     }
 
+    private static ColumnMapping timestampColumnMapping(JdbcTypeHandle typeHandle)
+    {
+        int timestampPrecision = typeHandle.requiredDecimalDigits();
+        TimestampType timestampType = createTimestampType(timestampPrecision);
+        if (timestampType.isShort()) {
+            return ColumnMapping.longMapping(
+                    timestampType,
+                    longTimestampReadFunction(timestampType),
+                    longTimestampWriteFunction(timestampType),
+                    FULL_PUSHDOWN);
+        }
+        return ColumnMapping.objectMapping(
+                timestampType,
+                objectTimestampReadFunction(timestampType),
+                objectTimestampWriteFunction(timestampType),
+                FULL_PUSHDOWN);
+    }
+
+    private static LongReadFunction longTimestampReadFunction(TimestampType timestampType)
+    {
+        return (resultSet, columnIndex) -> {
+            Timestamp timestamp = resultSet.getObject(columnIndex, Timestamp.class);
+            return toTrinoTimestamp(timestampType, timestamp.toLocalDateTime());
+        };
+    }
+
+    private static ObjectReadFunction objectTimestampReadFunction(TimestampType timestampType)
+    {
+        verifyObjectTimestampPrecision(timestampType);
+        return ObjectReadFunction.of(
+                LongTimestamp.class,
+                (resultSet, columnIndex) -> {
+                    Timestamp timestamp = resultSet.getObject(columnIndex, Timestamp.class);
+                    return toLongTrinoTimestamp(timestampType, timestamp.toLocalDateTime());
+                });
+    }
+
+    private static void verifyObjectTimestampPrecision(TimestampType timestampType)
+    {
+        int precision = timestampType.getPrecision();
+        checkArgument(precision > TimestampType.MAX_SHORT_PRECISION && precision <= MAX_EXASOL_TIMESTAMP_PRECISION,
+                "Precision is out of range: %s", precision);
+    }
+
+    private static ObjectWriteFunction objectTimestampWriteFunction(TimestampType timestampType)
+    {
+        int precision = timestampType.getPrecision();
+        verifyObjectTimestampPrecision(timestampType);
+
+        return new ObjectWriteFunction() {
+            @Override
+            public Class<?> getJavaType()
+            {
+                return LongTimestamp.class;
+            }
+
+            @Override
+            public void set(PreparedStatement statement, int index, Object value)
+                    throws SQLException
+            {
+                LocalDateTime localDateTime = fromLongTrinoTimestamp((LongTimestamp) value, precision);
+                Timestamp timestamp = Timestamp.valueOf(localDateTime);
+                statement.setObject(index, timestamp);
+            }
+
+            @Override
+            public String getBindExpression()
+            {
+                return getTimestampBindExpression(precision);
+            }
+
+            @Override
+            public void setNull(PreparedStatement statement, int index)
+                    throws SQLException
+            {
+                statement.setNull(index, Types.VARCHAR);
+            }
+        };
+    }
+
+    private static LongWriteFunction longTimestampWriteFunction(TimestampType timestampType)
+    {
+        return new LongWriteFunction()
+        {
+            @Override
+            public String getBindExpression()
+            {
+                return getTimestampBindExpression(timestampType.getPrecision());
+            }
+
+            @Override
+            public void set(PreparedStatement statement, int index, long epochMicros)
+                    throws SQLException
+            {
+                LocalDateTime localDateTime = fromTrinoTimestamp(epochMicros);
+                Timestamp timestampValue = Timestamp.valueOf(localDateTime);
+                statement.setObject(index, timestampValue);
+            }
+
+            @Override
+            public void setNull(PreparedStatement statement, int index)
+                    throws SQLException
+            {
+                statement.setNull(index, Types.VARCHAR);
+            }
+        };
+    }
+
+    private static ColumnMapping timestampWithTimeZoneColumnMapping(JdbcTypeHandle typeHandle)
+    {
+        int timestampPrecision = typeHandle.requiredDecimalDigits();
+        TimestampWithTimeZoneType timestampWithTimeZoneType = createTimestampWithTimeZoneType(timestampPrecision);
+
+        if (timestampWithTimeZoneType.isShort()) {
+            return ColumnMapping.longMapping(
+                    timestampWithTimeZoneType,
+                    longTimestampWithTimeZoneReadFunction(),
+                    longTimestampWithTimeZoneWriteFunction(timestampWithTimeZoneType),
+                    FULL_PUSHDOWN);
+        }
+        return ColumnMapping.objectMapping(
+                timestampWithTimeZoneType,
+                objectTimestampWithTimeZoneReadFunction(timestampWithTimeZoneType),
+                objectTimestampWithTimeZoneWriteFunction(timestampWithTimeZoneType),
+                FULL_PUSHDOWN);
+    }
+
+    private static LongReadFunction longTimestampWithTimeZoneReadFunction()
+    {
+        return (resultSet, columnIndex) -> {
+            Timestamp timestamp = resultSet.getObject(columnIndex, Timestamp.class);
+            return packDateTimeWithZone(timestamp.getTime(), UTC_KEY);
+        };
+    }
+
+    private static ObjectReadFunction objectTimestampWithTimeZoneReadFunction(
+            TimestampWithTimeZoneType timestampType)
+    {
+        verifyObjectTimestampWithTimeZonePrecision(timestampType);
+        return ObjectReadFunction.of(
+                LongTimestampWithTimeZone.class,
+                (resultSet, columnIndex) -> {
+                    Timestamp timestamp = resultSet.getObject(columnIndex, Timestamp.class);
+
+                    long millisUtc = timestamp.getTime();
+                    long nanosUtc = millisUtc * NANOSECONDS_PER_MILLISECOND + timestamp.getNanos() % NANOSECONDS_PER_MILLISECOND;
+                    int picosOfMilli = (int) ((nanosUtc - millisUtc * NANOSECONDS_PER_MILLISECOND) * PICOSECONDS_PER_NANOSECOND);
+
+                    return LongTimestampWithTimeZone.fromEpochMillisAndFraction(
+                            millisUtc,
+                            picosOfMilli,
+                            UTC_KEY);
+                });
+    }
+
+    private static void verifyObjectTimestampWithTimeZonePrecision(TimestampWithTimeZoneType timestampType)
+    {
+        int precision = timestampType.getPrecision();
+        checkArgument(precision > TimestampWithTimeZoneType.MAX_SHORT_PRECISION && precision <= MAX_EXASOL_TIMESTAMP_PRECISION,
+                "Precision is out of range: %s", precision);
+    }
+
+    private static ObjectWriteFunction objectTimestampWithTimeZoneWriteFunction(TimestampWithTimeZoneType timestampType)
+    {
+        int precision = timestampType.getPrecision();
+        verifyObjectTimestampWithTimeZonePrecision(timestampType);
+
+        return new ObjectWriteFunction() {
+            @Override
+            public Class<?> getJavaType()
+            {
+                return LongTimestampWithTimeZone.class;
+            }
+
+            @Override
+            public void set(PreparedStatement statement, int index, Object value)
+                    throws SQLException
+            {
+                if (value == null) {
+                    statement.setNull(index, Types.VARCHAR);
+                    return;
+                }
+
+                LongTimestampWithTimeZone timestampValue = (LongTimestampWithTimeZone) value;
+                Instant instant = Instant.ofEpochMilli(timestampValue.getEpochMillis())
+                        .plusNanos(timestampValue.getPicosOfMilli() / PICOSECONDS_PER_NANOSECOND);
+                Timestamp timestamp = Timestamp.from(instant);
+                statement.setObject(index, timestamp);
+            }
+
+            @Override
+            public String getBindExpression()
+            {
+                return getTimestampBindExpression(precision);
+            }
+
+            @Override
+            public void setNull(PreparedStatement statement, int index)
+                    throws SQLException
+            {
+                statement.setNull(index, Types.VARCHAR);
+            }
+        };
+    }
+
+    private static LongWriteFunction longTimestampWithTimeZoneWriteFunction(TimestampWithTimeZoneType timestampType)
+    {
+        return new LongWriteFunction()
+        {
+            @Override
+            public String getBindExpression()
+            {
+                return getTimestampBindExpression(timestampType.getPrecision());
+            }
+
+            @Override
+            public void set(PreparedStatement statement, int index, long epochMicros)
+                    throws SQLException
+            {
+                LocalDateTime localDateTime = fromTrinoTimestamp(epochMicros);
+                Timestamp timestampValue = Timestamp.valueOf(localDateTime);
+                statement.setObject(index, timestampValue);
+            }
+
+            @Override
+            public void setNull(PreparedStatement statement, int index)
+                    throws SQLException
+            {
+                statement.setNull(index, Types.VARCHAR);
+            }
+        };
+    }
+
+    private static String getTimestampBindExpression(int precision)
+    {
+        if (precision <= 0) {
+            return "TO_TIMESTAMP(?, 'YYYY-MM-DD HH24:MI:SS')";
+        }
+        return format("TO_TIMESTAMP(?, 'YYYY-MM-DD HH24:MI:SS.FF%d')", precision);
+    }
+
     private static ColumnMapping dateColumnMapping()
     {
         // Exasol driver does not support LocalDate