Add support for TIMESTAMP types in exasol connector

Author: skyglass

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

@@ -103,6 +103,12 @@ Trino data type mapping:
 * - `HASHTYPE`
   - `VARBINARY`
   -
+* - `TIMESTAMP(n)`
+  - `TIMESTAMP(n)`
+  -
+* - `TIMESTAMP(n) WITH LOCAL TIME ZONE`
+  - `TIMESTAMP(n)`
+  -
 :::
 
 No other types are supported.

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

@@ -30,6 +30,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.SliceReadFunction;
 import io.trino.plugin.jdbc.SliceWriteFunction;
@@ -43,12 +45,18 @@
 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.TimestampType;
 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.LocalDate;
+import java.time.LocalDateTime;
 import java.util.HexFormat;
 import java.util.List;
 import java.util.Map;
@@ -57,20 +65,27 @@
 import java.util.Set;
 import java.util.function.BiFunction;
 
+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.DateType.DATE;
 import static io.trino.spi.type.DecimalType.createDecimalType;
+import static io.trino.spi.type.TimestampType.createTimestampType;
 import static io.trino.spi.type.VarbinaryType.VARBINARY;
 import static java.lang.String.format;
 import static java.util.Locale.ENGLISH;
@@ -79,11 +94,15 @@
 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,
@@ -239,8 +258,13 @@ public Optional<ColumnMapping> toColumnMapping(ConnectorSession session, Connect
                 // String data is sorted by its binary representation.
                 // https://docs.exasol.com/db/latest/sql/select.htm#UsageNotes
                 return Optional.of(defaultVarcharColumnMapping(typeHandle.requiredColumnSize(), true));
+            // DATE, TIMESTAMP and TIMESTAMP WITH LOCAL TIME ZONE Types are described here in more details:
+            // https://docs.exasol.com/db/latest/sql_references/data_types/datatypedetails.htm
             case Types.DATE:
                 return Optional.of(dateColumnMapping());
+            case Types.TIMESTAMP:
+            case EXASOL_TIMESTAMP_WITH_TIMEZONE:
+                return Optional.of(timestampColumnMapping(typeHandle));
         }
 
         if (getUnsupportedTypeHandling(session) == CONVERT_TO_VARCHAR) {
@@ -256,6 +280,128 @@ private boolean isHashType(JdbcTypeHandle typeHandle)
                 && typeHandle.jdbcTypeName().get().equalsIgnoreCase("HASHTYPE");
     }
 
+    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(timestampType.getPrecision());
+            }
+
+            @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);
+            }
+        };
+    }
+
+    /**
+     * Returns a {@code TO_TIMESTAMP} bind expression using the appropriate format model
+     * based on the given fractional seconds precision.
+     * See for more details: https://docs.exasol.com/db/latest/sql_references/formatmodels.htm
+     */
+    private static String getTimestampBindExpression(int precision)
+    {
+        // Negative precisions are prevented by TimestampType
+        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