From 943b18f873dd8dc38935adcb0d8dde405028a8ae Mon Sep 17 00:00:00 2001 From: Teng YU Date: Tue, 5 Mar 2024 17:52:56 +0100 Subject: [PATCH] Fix tests --- .../plugin/snowflake/SnowflakeClient.java | 46 +++++++++++++++++++ .../snowflake/TestSnowflakeConnectorTest.java | 26 +++++++++++ 2 files changed, 72 insertions(+) diff --git a/plugin/trino-snowflake/src/main/java/io/trino/plugin/snowflake/SnowflakeClient.java b/plugin/trino-snowflake/src/main/java/io/trino/plugin/snowflake/SnowflakeClient.java index 2b592c0a4f513..e4e2f87543a96 100644 --- a/plugin/trino-snowflake/src/main/java/io/trino/plugin/snowflake/SnowflakeClient.java +++ b/plugin/trino-snowflake/src/main/java/io/trino/plugin/snowflake/SnowflakeClient.java @@ -29,6 +29,7 @@ import io.trino.plugin.jdbc.JdbcTableHandle; import io.trino.plugin.jdbc.JdbcTypeHandle; import io.trino.plugin.jdbc.LongWriteFunction; +import io.trino.plugin.jdbc.ObjectReadFunction; import io.trino.plugin.jdbc.ObjectWriteFunction; import io.trino.plugin.jdbc.PredicatePushdownController; import io.trino.plugin.jdbc.QueryBuilder; @@ -66,17 +67,25 @@ import java.math.RoundingMode; import java.sql.Connection; import java.sql.PreparedStatement; +import java.sql.ResultSet; 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.time.LocalTime; import java.time.ZoneId; +import java.time.ZoneOffset; import java.time.ZonedDateTime; import java.time.format.DateTimeFormatter; +import java.util.Calendar; +import java.util.Date; +import java.util.GregorianCalendar; import java.util.Locale; import java.util.Map; import java.util.Optional; +import java.util.TimeZone; import java.util.function.BiFunction; import static com.google.common.base.Preconditions.checkArgument; @@ -101,6 +110,7 @@ public class SnowflakeClient private static final DateTimeFormatter SNOWFLAKE_DATE_FORMATTER = DateTimeFormatter.ofPattern("uuuu-MM-dd"); private static final DateTimeFormatter SNOWFLAKE_TIMESTAMP_FORMATTER = DateTimeFormatter.ofPattern("u-MM-dd'T'HH:mm:ss.SSSSSSSSS"); private static final DateTimeFormatter SNOWFLAKE_TIME_FORMATTER = DateTimeFormatter.ofPattern("HH:mm:ss.SSSSSSSSS"); + private static final TimeZone UTC_TZ = TimeZone.getTimeZone(ZoneId.of("UTC")); private final AggregateFunctionRewriter aggregateFunctionRewriter; private interface WriteMappingFunction @@ -170,6 +180,7 @@ public Optional toColumnMapping(ConnectorSession session, Connect final Map snowflakeColumnMappings = ImmutableMap.builder() .put("time", handle -> Optional.of(timeColumnMapping(handle.getRequiredDecimalDigits()))) + .put("timestampntz", handle -> Optional.of(timestampColumnMapping(handle.getRequiredDecimalDigits()))) .put("date", handle -> Optional.of(ColumnMapping.longMapping(DateType.DATE, (resultSet, columnIndex) -> LocalDate.ofEpochDay(resultSet.getLong(columnIndex)).toEpochDay(), snowFlakeDateWriter()))) .put("varchar", handle -> Optional.of(varcharColumnMapping(handle.getRequiredColumnSize()))) .put("number", handle -> { @@ -312,6 +323,41 @@ private static ColumnMapping timeColumnMapping(int precision) PredicatePushdownController.FULL_PUSHDOWN); } + private static ColumnMapping timestampColumnMapping(int precision) + { + // <= 6 fits into a long + if (precision <= 6) { + return ColumnMapping.longMapping(TimestampType.createTimestampType(precision), (resultSet, columnIndex) -> StandardColumnMappings.toTrinoTimestamp(TimestampType.createTimestampType(precision), toLocalDateTime(resultSet, columnIndex)), timestampWriteFunction()); + } + + // Too big. Put it in an object + return ColumnMapping.objectMapping(TimestampType.createTimestampType(precision), longTimestampReader(), longTimestampWriter(precision)); + } + + private static LocalDateTime toLocalDateTime(ResultSet resultSet, int columnIndex) + throws SQLException + { + Calendar calendar = new GregorianCalendar(UTC_TZ, Locale.ENGLISH); + calendar.setTime(new Date(0)); + Timestamp ts = resultSet.getTimestamp(columnIndex, calendar); + return LocalDateTime.ofInstant(Instant.ofEpochMilli(ts.getTime()), ZoneOffset.UTC); + } + + private static ObjectReadFunction longTimestampReader() + { + return ObjectReadFunction.of(LongTimestamp.class, (resultSet, columnIndex) -> { + Calendar calendar = new GregorianCalendar(UTC_TZ, Locale.ENGLISH); + calendar.setTime(new Date(0)); + Timestamp ts = resultSet.getTimestamp(columnIndex, calendar); + long epochMillis = ts.getTime(); + int nanosInTheSecond = ts.getNanos(); + int nanosInTheMilli = nanosInTheSecond % Timestamps.NANOSECONDS_PER_MILLISECOND; + long micro = epochMillis * Timestamps.MICROSECONDS_PER_MILLISECOND + (nanosInTheMilli / Timestamps.NANOSECONDS_PER_MICROSECOND); + int picosOfMicro = nanosInTheMilli % 1000 * 1000; + return new LongTimestamp(micro, picosOfMicro); + }); + } + private static LongWriteFunction timeWriteFunction(int precision) { checkArgument(precision <= MAX_SUPPORTED_TEMPORAL_PRECISION, "Unsupported precision: %s", precision); diff --git a/plugin/trino-snowflake/src/test/java/io/trino/plugin/snowflake/TestSnowflakeConnectorTest.java b/plugin/trino-snowflake/src/test/java/io/trino/plugin/snowflake/TestSnowflakeConnectorTest.java index ee7d624cd6aa3..05fba3555e1ef 100644 --- a/plugin/trino-snowflake/src/test/java/io/trino/plugin/snowflake/TestSnowflakeConnectorTest.java +++ b/plugin/trino-snowflake/src/test/java/io/trino/plugin/snowflake/TestSnowflakeConnectorTest.java @@ -21,6 +21,7 @@ import io.trino.testing.TestingConnectorBehavior; import io.trino.testing.sql.SqlExecutor; import io.trino.testing.sql.TestTable; +import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.TestInstance; import org.junit.jupiter.api.parallel.Execution; @@ -247,6 +248,31 @@ public void testLimitPushdown() abort("TODO"); } + @Test + @Disabled + @Override + public void testCommentTable() + { + // This connector does not support setting table comments + } + + @Test + @Override + public void testCommentTableSpecialCharacter() + { + assertThatThrownBy(super::testCommentTableSpecialCharacter) + .hasStackTraceContaining("This connector does not support setting table comments"); + } + + @Test + @Override + public void testDataMappingSmokeTest() + { + // TODO: Connector does not support type timestamptz for the first commit + assertThatThrownBy(super::testDataMappingSmokeTest) + .hasStackTraceContaining("Unsupported column type(2014):timestamptz"); + } + @Test @Override public void testInsertIntoNotNullColumn()