Skip to content

Commit

Permalink
Add support for LocalDateTime and Instant in getObject and setObject
Browse files Browse the repository at this point in the history
  • Loading branch information
@huw0
huw0 committed Nov 11, 2024
1 parent 79df77e commit 201b6c3
Show file tree
Hide file tree
Showing 3 changed files with 159 additions and 9 deletions.
43 changes: 34 additions & 9 deletions client/trino-jdbc/src/main/java/io/trino/jdbc/AbstractTrinoResultSet.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@
import java.sql.Time;
import java.sql.Timestamp;
import java.sql.Types;
import java.time.Instant;
import java.time.LocalDateTime;
import java.time.ZoneId;
import java.time.ZonedDateTime;
Expand Down Expand Up @@ -156,9 +157,11 @@ abstract class AbstractTrinoResultSet
.add("date", String.class, java.time.LocalDate.class, string -> parseDate(string, CURRENT_TIME_ZONE, CURRENT_JAVA_TIME_ZONE).toLocalDate())
.add("time", String.class, Time.class, string -> parseTime(string, SYSTEM_DEFAULT_ZONE_ID))
.add("time with time zone", String.class, Time.class, AbstractTrinoResultSet::parseTimeWithTimeZone)
.add("timestamp", String.class, LocalDateTime.class, AbstractTrinoResultSet::parseTimestampAsLocalDateTime)
.add("timestamp", String.class, Timestamp.class, string -> parseTimestampAsSqlTimestamp(string, SYSTEM_DEFAULT_ZONE_ID))
.add("timestamp with time zone", String.class, Instant.class, AbstractTrinoResultSet::parseTimestampWithTimeZoneAsInstant)
.add("timestamp with time zone", String.class, Timestamp.class, AbstractTrinoResultSet::parseTimestampWithTimeZoneAsSqlTimestamp)
.add("timestamp with time zone", String.class, ZonedDateTime.class, AbstractTrinoResultSet::parseTimestampWithTimeZone)
.add("timestamp with time zone", String.class, ZonedDateTime.class, AbstractTrinoResultSet::parseTimestampWithTimeZoneAsZonedDateTime)
.add("interval year to month", String.class, TrinoIntervalYearMonth.class, AbstractTrinoResultSet::parseIntervalYearMonth)
.add("interval day to second", String.class, TrinoIntervalDayTime.class, AbstractTrinoResultSet::parseIntervalDayTime)
.add("array", List.class, List.class, (type, list) -> (List<?>) convertFromClientRepresentation(type, list))
Expand Down Expand Up @@ -447,10 +450,16 @@ private Timestamp getTimestamp(int columnIndex, DateTimeZone localTimeZone)
throw new IllegalArgumentException("Expected column to be a timestamp type but is " + columnInfo.getColumnTypeName());
}

private static ZonedDateTime parseTimestampWithTimeZone(String value)
private static Instant parseTimestampWithTimeZoneAsInstant(String value)
{
ParsedTimestamp parsed = parseTimestamp(value);
return toZonedDateTime(parsed, timezone -> ZoneId.of(timezone.orElseThrow(() -> new IllegalArgumentException("Time zone missing: " + value))));
return toInstant(parsed, timezone -> timezone.map(ZoneId::of).orElseThrow(() -> new IllegalArgumentException("Time zone missing: " + value)));
}

private static ZonedDateTime parseTimestampWithTimeZoneAsZonedDateTime(String value)
{
ParsedTimestamp parsed = parseTimestamp(value);
return toZonedDateTime(parsed, timezone -> timezone.map(ZoneId::of).orElseThrow(() -> new IllegalArgumentException("Time zone missing: " + value)));
}

@Override
Expand Down Expand Up @@ -1975,6 +1984,11 @@ private static Timestamp parseTimestampWithTimeZoneAsSqlTimestamp(String value)
ZoneId.of(timezone.orElseThrow(() -> new IllegalArgumentException("Time zone missing: " + value))));
}

private static LocalDateTime parseTimestampAsLocalDateTime(String value)
{
return toLocalDateTime(parseTimestamp(value));
}

private static Timestamp parseTimestampAsSqlTimestamp(String value, ZoneId localTimeZone)
{
requireNonNull(localTimeZone, "localTimeZone is null");
Expand Down Expand Up @@ -2049,7 +2063,13 @@ private static Timestamp toTimestamp(String originalValue, ParsedTimestamp parse
return timestamp;
}

private static ZonedDateTime toZonedDateTime(ParsedTimestamp parsed, Function<Optional<String>, ZoneId> timeZoneParser)
private static Instant toInstant(ParsedTimestamp parsed, Function<Optional<String>, ZoneId> timeZoneParser)
{
return toZonedDateTime(parsed, timeZoneParser)
.toInstant();
}

private static LocalDateTime toLocalDateTime(ParsedTimestamp parsed)
{
int year = parsed.year;
int month = parsed.month;
Expand All @@ -2058,14 +2078,19 @@ private static ZonedDateTime toZonedDateTime(ParsedTimestamp parsed, Function<Op
int minute = parsed.minute;
int second = parsed.second;
long picosOfSecond = parsed.picosOfSecond;
ZoneId zoneId = timeZoneParser.apply(parsed.timezone);

ZonedDateTime zonedDateTime = LocalDateTime.of(year, month, day, hour, minute, second, 0)
.atZone(zoneId);
LocalDateTime localDateTime = LocalDateTime.of(year, month, day, hour, minute, second, 0);

int nanoOfSecond = (int) rescale(picosOfSecond, 12, 9);
zonedDateTime = zonedDateTime.plusNanos(nanoOfSecond);
return zonedDateTime;
return localDateTime.plusNanos(nanoOfSecond);
}

private static ZonedDateTime toZonedDateTime(ParsedTimestamp parsed, Function<Optional<String>, ZoneId> timeZoneParser)
{
ZoneId zoneId = timeZoneParser.apply(parsed.timezone);

return toLocalDateTime(parsed)
.atZone(zoneId);
}

private static Time parseTime(String value, ZoneId localTimeZone)
Expand Down
19 changes: 19 additions & 0 deletions client/trino-jdbc/src/main/java/io/trino/jdbc/TrinoPreparedStatement.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,10 +45,12 @@
import java.sql.Time;
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.OffsetTime;
import java.time.ZoneOffset;
import java.time.format.DateTimeFormatter;
import java.time.format.DateTimeFormatterBuilder;
import java.util.ArrayList;
Expand Down Expand Up @@ -94,6 +96,14 @@ public class TrinoPreparedStatement
.append(ISO_LOCAL_TIME)
.toFormatter();

private static final DateTimeFormatter OFFSET_DATE_TIME_FORMATTER =
new DateTimeFormatterBuilder()
.append(ISO_LOCAL_DATE)
.appendLiteral(' ')
.append(ISO_LOCAL_TIME)
.appendOffset("+HH:mm", "+00:00")
.toFormatter();

private static final DateTimeFormatter OFFSET_TIME_FORMATTER =
new DateTimeFormatterBuilder()
.append(ISO_LOCAL_TIME)
Expand Down Expand Up @@ -432,6 +442,9 @@ private String toTimestampWithTimeZoneLiteral(Object value)
// TODO validate proper format
return (String) value;
}
else if (value instanceof Instant) {
return OFFSET_DATE_TIME_FORMATTER.format(((Instant) value).atOffset(ZoneOffset.UTC));
}
throw invalidConversion(value, "timestamp with time zone");
}

Expand Down Expand Up @@ -598,6 +611,12 @@ else if (x instanceof Date) {
else if (x instanceof LocalDate) {
setAsDate(parameterIndex, x);
}
else if (x instanceof LocalDateTime) {
setAsTimestamp(parameterIndex, x);
}
else if (x instanceof Instant) {
setAsTimestampWithTimeZone(parameterIndex, x);
}
else if (x instanceof Time) {
setTime(parameterIndex, (Time) x);
}
Expand Down
106 changes: 106 additions & 0 deletions client/trino-jdbc/src/test/java/io/trino/jdbc/TestJdbcPreparedStatement.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,12 +44,14 @@
import java.sql.Time;
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.OffsetTime;
import java.time.ZoneId;
import java.time.ZoneOffset;
import java.time.temporal.ChronoUnit;
import java.util.Calendar;
import java.util.TimeZone;
import java.util.stream.LongStream;
Expand Down Expand Up @@ -1102,6 +1104,91 @@ private void testConvertLocalDate(boolean explicitPrepare)
.roundTripsAs(Types.DATE, Date.valueOf(jvmGapDate));
}

@Test
public void testConvertInstant()
throws SQLException
{
testConvertInstant(true);
testConvertInstant(false);
}

private void testConvertInstant(boolean explicitPrepare)
throws SQLException
{
LocalDateTime dateTime = LocalDateTime.of(2001, 5, 6, 12, 34, 56);
Instant instant = dateTime.toInstant(ZoneOffset.UTC);

assertBind((ps, i) -> ps.setObject(i, instant, Types.TIMESTAMP_WITH_TIMEZONE), explicitPrepare)
.resultsIn("timestamp(0) with time zone", "TIMESTAMP '2001-05-06 12:34:56 +00:00'")
.roundTripsAs(Types.TIMESTAMP_WITH_TIMEZONE, Instant.class, instant);

Instant instantWithDeciSecond = instant.plus(100, ChronoUnit.MILLIS);

assertBind((ps, i) -> ps.setObject(i, instantWithDeciSecond), explicitPrepare)
.resultsIn("timestamp(1) with time zone", "TIMESTAMP '2001-05-06 12:34:56.1 +00:00'")
.roundTripsAs(Types.TIMESTAMP_WITH_TIMEZONE, Instant.class, instantWithDeciSecond);

assertBind((ps, i) -> ps.setObject(i, instantWithDeciSecond, Types.TIMESTAMP_WITH_TIMEZONE), explicitPrepare)
.resultsIn("timestamp(1) with time zone", "TIMESTAMP '2001-05-06 12:34:56.1 +00:00'")
.roundTripsAs(Types.TIMESTAMP_WITH_TIMEZONE, Instant.class, instantWithDeciSecond);

Instant instantWithMilliSecond = instant.plus(123, ChronoUnit.MILLIS);

assertBind((ps, i) -> ps.setObject(i, instantWithMilliSecond), explicitPrepare)
.resultsIn("timestamp(3) with time zone", "TIMESTAMP '2001-05-06 12:34:56.123 +00:00'")
.roundTripsAs(Types.TIMESTAMP_WITH_TIMEZONE, Instant.class, instantWithMilliSecond);

assertBind((ps, i) -> ps.setObject(i, instantWithMilliSecond, Types.TIMESTAMP_WITH_TIMEZONE), explicitPrepare)
.resultsIn("timestamp(3) with time zone", "TIMESTAMP '2001-05-06 12:34:56.123 +00:00'")
.roundTripsAs(Types.TIMESTAMP_WITH_TIMEZONE, Instant.class, instantWithMilliSecond);
}

@Test
public void testConvertLocalDateTime()
throws SQLException
{
testConvertLocalDateTime(true);
testConvertLocalDateTime(false);
}

private void testConvertLocalDateTime(boolean explicitPrepare)
throws SQLException
{
LocalDateTime dateTime = LocalDateTime.of(2001, 5, 6, 12, 34, 56);
Timestamp sqlTimestamp = Timestamp.valueOf(dateTime);

assertBind((ps, i) -> ps.setObject(i, dateTime, Types.TIMESTAMP), explicitPrepare)
.resultsIn("timestamp(0)", "TIMESTAMP '2001-05-06 12:34:56'")
.roundTripsAs(Types.TIMESTAMP, sqlTimestamp)
.roundTripsAs(Types.TIMESTAMP, LocalDateTime.class, dateTime);

LocalDateTime dateTimeWithDeciSecond = dateTime.plus(100, ChronoUnit.MILLIS);
Timestamp timestampWithWithDecisecond = new Timestamp(sqlTimestamp.getTime() + 100);

assertBind((ps, i) -> ps.setObject(i, dateTimeWithDeciSecond), explicitPrepare)
.resultsIn("timestamp(1)", "TIMESTAMP '2001-05-06 12:34:56.1'")
.roundTripsAs(Types.TIMESTAMP, timestampWithWithDecisecond)
.roundTripsAs(Types.TIMESTAMP, LocalDateTime.class, dateTimeWithDeciSecond);

assertBind((ps, i) -> ps.setObject(i, dateTimeWithDeciSecond, Types.TIMESTAMP), explicitPrepare)
.resultsIn("timestamp(1)", "TIMESTAMP '2001-05-06 12:34:56.1'")
.roundTripsAs(Types.TIMESTAMP, timestampWithWithDecisecond)
.roundTripsAs(Types.TIMESTAMP, LocalDateTime.class, dateTimeWithDeciSecond);

LocalDateTime dateTimeWithMilliSecond = dateTime.plus(123, ChronoUnit.MILLIS);
Timestamp timestampWithMillisecond = new Timestamp(sqlTimestamp.getTime() + 123);

assertBind((ps, i) -> ps.setObject(i, dateTimeWithMilliSecond), explicitPrepare)
.resultsIn("timestamp(3)", "TIMESTAMP '2001-05-06 12:34:56.123'")
.roundTripsAs(Types.TIMESTAMP, timestampWithMillisecond)
.roundTripsAs(Types.TIMESTAMP, LocalDateTime.class, dateTimeWithMilliSecond);

assertBind((ps, i) -> ps.setObject(i, dateTimeWithMilliSecond, Types.TIMESTAMP), explicitPrepare)
.resultsIn("timestamp(3)", "TIMESTAMP '2001-05-06 12:34:56.123'")
.roundTripsAs(Types.TIMESTAMP, timestampWithMillisecond)
.roundTripsAs(Types.TIMESTAMP, LocalDateTime.class, dateTimeWithMilliSecond);
}

@Test
public void testConvertTime()
throws SQLException
Expand Down Expand Up @@ -1548,6 +1635,25 @@ public BindAssertion roundTripsAs(int expectedSqlType, Object expectedValue)
return this;
}

public BindAssertion roundTripsAs(int expectedSqlType, Class expectedClass, Object expectedValue)
throws SQLException
{
try (Connection connection = connectionFactory.createConnection();
PreparedStatement statement = connection.prepareStatement("SELECT ?")) {
binder.bind(statement, 1);

try (ResultSet rs = statement.executeQuery()) {
verify(rs.next(), "no row returned");
assertThat(rs.getObject(1, expectedClass)).isEqualTo(expectedValue);
verify(!rs.next(), "unexpected second row");

assertThat(rs.getMetaData().getColumnType(1)).isEqualTo(expectedSqlType);
}
}

return this;
}

public BindAssertion resultsIn(String type, String expectedValueLiteral)
throws SQLException
{
Expand Down

0 comments on commit 201b6c3

Please sign in to comment.