diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java index 0dc2b07c97..622e5fe7f6 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java @@ -19,6 +19,7 @@ import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; +import java.sql.Types; import java.util.HashSet; import java.util.List; import java.util.Objects; @@ -28,14 +29,17 @@ import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.calcite.avatica.AvaticaConnection; import org.apache.calcite.avatica.AvaticaResultSet; import org.apache.calcite.avatica.AvaticaResultSetMetaData; +import org.apache.calcite.avatica.AvaticaSite; import org.apache.calcite.avatica.AvaticaStatement; import org.apache.calcite.avatica.ColumnMetaData; import org.apache.calcite.avatica.Meta; import org.apache.calcite.avatica.Meta.Frame; import org.apache.calcite.avatica.Meta.Signature; import org.apache.calcite.avatica.QueryState; +import org.apache.calcite.avatica.util.Cursor; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -102,6 +106,33 @@ void populateData(final VectorSchemaRoot vectorSchemaRoot, final Schema schema) execute2(new ArrowFlightJdbcCursor(vectorSchemaRoot), this.signature.columns); } + /** + * The default method in AvaticaResultSet does not properly handle TIMESTASMP_WITH_TIMEZONE, so we + * override here to add support. + * + * @param columnIndex the first column is 1, the second is 2, ... + * @return Object + * @throws SQLException if there is an underlying exception + */ + @Override + public Object getObject(int columnIndex) throws SQLException { + this.checkOpen(); + + Cursor.Accessor accessor; + try { + accessor = accessorList.get(columnIndex - 1); + } catch (IndexOutOfBoundsException e) { + throw AvaticaConnection.HELPER.createException("invalid column ordinal: " + columnIndex); + } + + ColumnMetaData metaData = columnMetaDataList.get(columnIndex - 1); + if (metaData.type.id == Types.TIMESTAMP_WITH_TIMEZONE) { + return accessor.getTimestamp(localCalendar); + } else { + return AvaticaSite.get(accessor, metaData.type.id, localCalendar); + } + } + @Override protected void cancel() { signature.columns.clear(); diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestExtension.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestExtension.java index db0438059f..f71114e1b5 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestExtension.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestExtension.java @@ -130,6 +130,12 @@ public Connection getConnection(boolean useEncryption) throws SQLException { return this.createDataSource().getConnection(); } + public Connection getConnection(String timezone) throws SQLException { + setUseEncryption(false); + properties.put("timezone", timezone); + return this.createDataSource().getConnection(); + } + private void setUseEncryption(boolean useEncryption) { properties.put("useEncryption", useEncryption); } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/TimestampResultSetTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/TimestampResultSetTest.java new file mode 100644 index 0000000000..0921ae2d38 --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/TimestampResultSetTest.java @@ -0,0 +1,164 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.driver.jdbc; + +import com.google.common.collect.ImmutableList; +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.LocalDateTime; +import java.time.OffsetDateTime; +import java.time.ZoneOffset; +import java.time.ZonedDateTime; +import java.util.Calendar; +import java.util.Collections; +import java.util.TimeZone; +import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.TimeStampVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.TimeUnit; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +/** + * Timestamps have a lot of nuances in JDBC. This class is here to test that timestamp behavior is + * correct for different types of Timestamp vectors as well as different methods of retrieving the + * timestamps in JDBC. + */ +public class TimestampResultSetTest { + private static final MockFlightSqlProducer FLIGHT_SQL_PRODUCER = new MockFlightSqlProducer(); + + @RegisterExtension public static FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION; + + static { + FLIGHT_SERVER_TEST_EXTENSION = + FlightServerTestExtension.createStandardTestExtension(FLIGHT_SQL_PRODUCER); + } + + private static final String QUERY_STRING = "SELECT * FROM TIMESTAMPS"; + private static final Schema QUERY_SCHEMA = + new Schema( + ImmutableList.of( + Field.nullable("no_tz", new ArrowType.Timestamp(TimeUnit.MILLISECOND, null)), + Field.nullable("utc", new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")), + Field.nullable("utc+1", new ArrowType.Timestamp(TimeUnit.MILLISECOND, "GMT+1")), + Field.nullable("utc-1", new ArrowType.Timestamp(TimeUnit.MILLISECOND, "GMT-1")))); + + @BeforeAll + public static void setup() throws SQLException { + Instant firstDay2025 = OffsetDateTime.of(2025, 1, 1, 0, 0, 0, 0, ZoneOffset.UTC).toInstant(); + + FLIGHT_SQL_PRODUCER.addSelectQuery( + QUERY_STRING, + QUERY_SCHEMA, + Collections.singletonList( + listener -> { + try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + final VectorSchemaRoot root = VectorSchemaRoot.create(QUERY_SCHEMA, allocator)) { + listener.start(root); + root.getFieldVectors() + .forEach(v -> ((TimeStampVector) v).setSafe(0, firstDay2025.toEpochMilli())); + root.setRowCount(1); + listener.putNext(); + } catch (final Throwable throwable) { + listener.error(throwable); + } finally { + listener.completed(); + } + })); + } + + /** + * This test doesn't yet test anything other than ensuring all ResultSet methods to retrieve a + * timestamp succeed. + * + *

This is a good starting point to add more tests to ensure the values are correct when we + * change the "local calendar" either through changing the JVM default or through the connection + * property. + */ + @Test + public void test() { + TimeZone.setDefault(TimeZone.getTimeZone("UTC")); + try (Connection connection = FLIGHT_SERVER_TEST_EXTENSION.getConnection("UTC")) { + try (PreparedStatement s = connection.prepareStatement(QUERY_STRING)) { + try (ResultSet rs = s.executeQuery()) { + int numCols = rs.getMetaData().getColumnCount(); + try { + rs.next(); + for (int i = 1; i <= numCols; i++) { + int type = rs.getMetaData().getColumnType(i); + String name = rs.getMetaData().getColumnName(i); + System.out.println(name); + System.out.print("- getDate:\t\t\t\t\t\t\t"); + System.out.print(rs.getDate(i)); + System.out.println(); + System.out.print("- getTimestamp:\t\t\t\t\t\t"); + System.out.print(rs.getTimestamp(i)); + System.out.println(); + System.out.print("- getString:\t\t\t\t\t\t"); + System.out.print(rs.getString(i)); + System.out.println(); + System.out.print("- getObject:\t\t\t\t\t\t"); + System.out.print(rs.getObject(i)); + System.out.println(); + System.out.print("- getObject(Timestamp.class):\t\t"); + System.out.print(rs.getObject(i, Timestamp.class)); + System.out.println(); + System.out.print("- getTimestamp(default Calendar):\t"); + System.out.print(rs.getTimestamp(i, Calendar.getInstance())); + System.out.println(); + System.out.print("- getTimestamp(UTC Calendar):\t\t"); + System.out.print( + rs.getTimestamp(i, Calendar.getInstance(TimeZone.getTimeZone("UTC")))); + System.out.println(); + System.out.print("- getObject(LocalDateTime.class):\t"); + System.out.print(rs.getObject(i, LocalDateTime.class)); + System.out.println(); + if (type == Types.TIMESTAMP_WITH_TIMEZONE) { + System.out.print("- getObject(Instant.class):\t\t\t"); + System.out.print(rs.getObject(i, Instant.class)); + System.out.println(); + System.out.print("- getObject(OffsetDateTime.class):\t"); + System.out.print(rs.getObject(i, OffsetDateTime.class)); + System.out.println(); + System.out.print("- getObject(ZonedDateTime.class):\t"); + System.out.print(rs.getObject(i, ZonedDateTime.class)); + System.out.println(); + } + System.out.println(); + } + System.out.println(); + } catch (SQLException e) { + throw new RuntimeException(e); + } + } + } + } catch (SQLException e) { + throw new RuntimeException(e); + } + } +}