From 19acccec78c00fd129400a87974218f077c8eda0 Mon Sep 17 00:00:00 2001 From: Teng YU Date: Wed, 6 Mar 2024 15:34:18 +0100 Subject: [PATCH] Add the handle for timestamptz back to fix testDataMappingSmokeTest --- .../plugin/snowflake/SnowflakeClient.java | 87 ++++++++++++------- .../snowflake/TestSnowflakeConnectorTest.java | 27 +----- 2 files changed, 56 insertions(+), 58 deletions(-) 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 e4e2f87543a96..7f0d7c1bd7957 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 @@ -181,6 +181,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("timestamptz", handle -> Optional.of(timestampTZColumnMapping(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 -> { @@ -229,38 +230,14 @@ public WriteMapping toWriteMapping(ConnectorSession session, Type type) final Map snowflakeWriteMappings = ImmutableMap.builder() .put("TimeType", writeType -> WriteMapping.longMapping("time", timeWriteFunction(((TimeType) writeType).getPrecision()))) - .put("ShortTimestampType", writeType -> { - WriteMapping myMap = SnowflakeClient.snowFlakeTimestampWriter(writeType); - return myMap; - }) - .put("ShortTimestampWithTimeZoneType", writeType -> { - WriteMapping myMap = SnowflakeClient.snowFlakeTimestampWithTZWriter(writeType); - return myMap; - }) - .put("LongTimestampType", writeType -> { - WriteMapping myMap = SnowflakeClient.snowFlakeTimestampWithTZWriter(writeType); - return myMap; - }) - .put("LongTimestampWithTimeZoneType", writeType -> { - WriteMapping myMap = SnowflakeClient.snowFlakeTimestampWithTZWriter(writeType); - return myMap; - }) - .put("VarcharType", writeType -> { - WriteMapping myMap = SnowflakeClient.snowFlakeVarCharWriter(writeType); - return myMap; - }) - .put("CharType", writeType -> { - WriteMapping myMap = SnowflakeClient.snowFlakeCharWriter(writeType); - return myMap; - }) - .put("LongDecimalType", writeType -> { - WriteMapping myMap = SnowflakeClient.snowFlakeDecimalWriter(writeType); - return myMap; - }) - .put("ShortDecimalType", writeType -> { - WriteMapping myMap = SnowflakeClient.snowFlakeDecimalWriter(writeType); - return myMap; - }) + .put("ShortTimestampType", writeType -> SnowflakeClient.snowFlakeTimestampWriter(writeType)) + .put("ShortTimestampWithTimeZoneType", writeType -> SnowflakeClient.snowFlakeTimestampWithTZWriter(writeType)) + .put("LongTimestampType", writeType -> SnowflakeClient.snowFlakeTimestampWithTZWriter(writeType)) + .put("LongTimestampWithTimeZoneType", writeType -> SnowflakeClient.snowFlakeTimestampWithTZWriter(writeType)) + .put("VarcharType", writeType -> SnowflakeClient.snowFlakeVarCharWriter(writeType)) + .put("CharType", writeType -> SnowflakeClient.snowFlakeCharWriter(writeType)) + .put("LongDecimalType", writeType -> SnowflakeClient.snowFlakeDecimalWriter(writeType)) + .put("ShortDecimalType", writeType -> SnowflakeClient.snowFlakeDecimalWriter(writeType)) .buildOrThrow(); WriteMappingFunction writeMappingFunction = snowflakeWriteMappings.get(simple); @@ -323,6 +300,52 @@ private static ColumnMapping timeColumnMapping(int precision) PredicatePushdownController.FULL_PUSHDOWN); } + private static ColumnMapping timestampTZColumnMapping(int precision) + { + if (precision <= 3) { + return ColumnMapping.longMapping(TimestampWithTimeZoneType.createTimestampWithTimeZoneType(precision), + (resultSet, columnIndex) -> { + ZonedDateTime timestamp = SNOWFLAKE_DATETIME_FORMATTER.parse(resultSet.getString(columnIndex), ZonedDateTime::from); + return DateTimeEncoding.packDateTimeWithZone(timestamp.toInstant().toEpochMilli(), timestamp.getZone().getId()); + }, + timestampWithTZWriter(), PredicatePushdownController.FULL_PUSHDOWN); + } + else { + return ColumnMapping.objectMapping(TimestampWithTimeZoneType.createTimestampWithTimeZoneType(precision), longTimestampWithTimezoneReadFunction(), longTimestampWithTZWriteFunction()); + } + } + + private static LongWriteFunction timestampWithTZWriter() + { + return (statement, index, encodedTimeWithZone) -> { + Instant timeI = Instant.ofEpochMilli(DateTimeEncoding.unpackMillisUtc(encodedTimeWithZone)); + ZoneId zone = ZoneId.of(DateTimeEncoding.unpackZoneKey(encodedTimeWithZone).getId()); + statement.setString(index, SNOWFLAKE_DATETIME_FORMATTER.format(timeI.atZone(zone))); + }; + } + + private static ObjectReadFunction longTimestampWithTimezoneReadFunction() + { + return ObjectReadFunction.of(LongTimestampWithTimeZone.class, (resultSet, columnIndex) -> { + ZonedDateTime timestamp = SNOWFLAKE_DATETIME_FORMATTER.parse(resultSet.getString(columnIndex), ZonedDateTime::from); + return LongTimestampWithTimeZone.fromEpochSecondsAndFraction(timestamp.toEpochSecond(), + (long) timestamp.getNano() * Timestamps.PICOSECONDS_PER_NANOSECOND, + TimeZoneKey.getTimeZoneKey(timestamp.getZone().getId())); + }); + } + + private static ObjectWriteFunction longTimestampWithTZWriteFunction() + { + return ObjectWriteFunction.of(LongTimestampWithTimeZone.class, (statement, index, value) -> { + long epoMilli = value.getEpochMillis(); + long epoSeconds = Math.floorDiv(epoMilli, Timestamps.MILLISECONDS_PER_SECOND); + long adjNano = (long) Math.floorMod(epoMilli, Timestamps.MILLISECONDS_PER_SECOND) * Timestamps.NANOSECONDS_PER_MILLISECOND + value.getPicosOfMilli() / Timestamps.PICOSECONDS_PER_NANOSECOND; + ZoneId zone = TimeZoneKey.getTimeZoneKey(value.getTimeZoneKey()).getZoneId(); + Instant timeI = Instant.ofEpochSecond(epoSeconds, adjNano); + statement.setString(index, SNOWFLAKE_DATETIME_FORMATTER.format(ZonedDateTime.ofInstant(timeI, zone))); + }); + } + private static ColumnMapping timestampColumnMapping(int precision) { // <= 6 fits into a long 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 05fba3555e1ef..84bbb18c4809d 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,7 +21,6 @@ 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; @@ -73,6 +72,7 @@ protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior) SUPPORTS_AGGREGATION_PUSHDOWN, SUPPORTS_ARRAY, SUPPORTS_COMMENT_ON_COLUMN, + SUPPORTS_COMMENT_ON_TABLE, SUPPORTS_CREATE_TABLE_WITH_COLUMN_COMMENT, SUPPORTS_CREATE_TABLE_WITH_TABLE_COMMENT, SUPPORTS_LIMIT_PUSHDOWN, @@ -248,31 +248,6 @@ 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()