Skip to content

Commit

Permalink
Add the handle for timestamptz back to fix testDataMappingSmokeTest
Browse files Browse the repository at this point in the history
  • Loading branch information
@yuuteng
yuuteng committed Mar 6, 2024
1 parent 943b18f commit 19accce
Show file tree
Hide file tree
Showing 2 changed files with 56 additions and 58 deletions.
87 changes: 55 additions & 32 deletions plugin/trino-snowflake/src/main/java/io/trino/plugin/snowflake/SnowflakeClient.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -181,6 +181,7 @@ public Optional<ColumnMapping> toColumnMapping(ConnectorSession session, Connect
final Map<String, ColumnMappingFunction> snowflakeColumnMappings = ImmutableMap.<String, ColumnMappingFunction>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 -> {
Expand Down Expand Up @@ -229,38 +230,14 @@ public WriteMapping toWriteMapping(ConnectorSession session, Type type)

final Map<String, WriteMappingFunction> snowflakeWriteMappings = ImmutableMap.<String, WriteMappingFunction>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);
Expand Down Expand Up @@ -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
Expand Down
27 changes: 1 addition & 26 deletions ...n/trino-snowflake/src/test/java/io/trino/plugin/snowflake/TestSnowflakeConnectorTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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()
Expand Down

0 comments on commit 19accce

Please sign in to comment.