Skip to content

Commit d31d9b0

Browse files
aiguoferaiguofer
authored
GH-839: Fix support for ResultSet.getObject for TIMESTAMP_WITH_TIMEZONE (#840)
## What's Changed Turns out AvaticaSite.get does not account for TIMESTAMP_WITH_TIMEZONE types so we add support on an override function. Closes #818. Closes #839.
1 parent d529557 commit d31d9b0

File tree

3 files changed

+201
-0
lines changed

3 files changed

+201
-0
lines changed

‎flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import java.sql.ResultSet;
2020
import java.sql.ResultSetMetaData;
2121
import java.sql.SQLException;
22+
import java.sql.Types;
2223
import java.util.HashSet;
2324
import java.util.List;
2425
import java.util.Objects;
@@ -28,14 +29,17 @@
2829
import org.apache.arrow.util.AutoCloseables;
2930
import org.apache.arrow.vector.VectorSchemaRoot;
3031
import org.apache.arrow.vector.types.pojo.Schema;
32+
import org.apache.calcite.avatica.AvaticaConnection;
3133
import org.apache.calcite.avatica.AvaticaResultSet;
3234
import org.apache.calcite.avatica.AvaticaResultSetMetaData;
35+
import org.apache.calcite.avatica.AvaticaSite;
3336
import org.apache.calcite.avatica.AvaticaStatement;
3437
import org.apache.calcite.avatica.ColumnMetaData;
3538
import org.apache.calcite.avatica.Meta;
3639
import org.apache.calcite.avatica.Meta.Frame;
3740
import org.apache.calcite.avatica.Meta.Signature;
3841
import org.apache.calcite.avatica.QueryState;
42+
import org.apache.calcite.avatica.util.Cursor;
3943
import org.slf4j.Logger;
4044
import org.slf4j.LoggerFactory;
4145

@@ -102,6 +106,33 @@ void populateData(final VectorSchemaRoot vectorSchemaRoot, final Schema schema)
102106
execute2(new ArrowFlightJdbcCursor(vectorSchemaRoot), this.signature.columns);
103107
}
104108

109+
/**
110+
* The default method in AvaticaResultSet does not properly handle TIMESTASMP_WITH_TIMEZONE, so we
111+
* override here to add support.
112+
*
113+
* @param columnIndex the first column is 1, the second is 2, ...
114+
* @return Object
115+
* @throws SQLException if there is an underlying exception
116+
*/
117+
@Override
118+
public Object getObject(int columnIndex) throws SQLException {
119+
this.checkOpen();
120+
121+
Cursor.Accessor accessor;
122+
try {
123+
accessor = accessorList.get(columnIndex - 1);
124+
} catch (IndexOutOfBoundsException e) {
125+
throw AvaticaConnection.HELPER.createException("invalid column ordinal: " + columnIndex);
126+
}
127+
128+
ColumnMetaData metaData = columnMetaDataList.get(columnIndex - 1);
129+
if (metaData.type.id == Types.TIMESTAMP_WITH_TIMEZONE) {
130+
return accessor.getTimestamp(localCalendar);
131+
} else {
132+
return AvaticaSite.get(accessor, metaData.type.id, localCalendar);
133+
}
134+
}
135+
105136
@Override
106137
protected void cancel() {
107138
signature.columns.clear();

‎flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestExtension.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,12 @@ public Connection getConnection(boolean useEncryption) throws SQLException {
130130
return this.createDataSource().getConnection();
131131
}
132132

133+
public Connection getConnection(String timezone) throws SQLException {
134+
setUseEncryption(false);
135+
properties.put("timezone", timezone);
136+
return this.createDataSource().getConnection();
137+
}
138+
133139
private void setUseEncryption(boolean useEncryption) {
134140
properties.put("useEncryption", useEncryption);
135141
}

‎flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/TimestampResultSetTest.java‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package org.apache.arrow.driver.jdbc;
18+
19+
import com.google.common.collect.ImmutableList;
20+
import java.sql.Connection;
21+
import java.sql.PreparedStatement;
22+
import java.sql.ResultSet;
23+
import java.sql.SQLException;
24+
import java.sql.Timestamp;
25+
import java.sql.Types;
26+
import java.time.Instant;
27+
import java.time.LocalDateTime;
28+
import java.time.OffsetDateTime;
29+
import java.time.ZoneOffset;
30+
import java.time.ZonedDateTime;
31+
import java.util.Calendar;
32+
import java.util.Collections;
33+
import java.util.TimeZone;
34+
import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer;
35+
import org.apache.arrow.memory.BufferAllocator;
36+
import org.apache.arrow.memory.RootAllocator;
37+
import org.apache.arrow.vector.TimeStampVector;
38+
import org.apache.arrow.vector.VectorSchemaRoot;
39+
import org.apache.arrow.vector.types.TimeUnit;
40+
import org.apache.arrow.vector.types.pojo.ArrowType;
41+
import org.apache.arrow.vector.types.pojo.Field;
42+
import org.apache.arrow.vector.types.pojo.Schema;
43+
import org.junit.jupiter.api.BeforeAll;
44+
import org.junit.jupiter.api.Test;
45+
import org.junit.jupiter.api.extension.RegisterExtension;
46+
47+
/**
48+
* Timestamps have a lot of nuances in JDBC. This class is here to test that timestamp behavior is
49+
* correct for different types of Timestamp vectors as well as different methods of retrieving the
50+
* timestamps in JDBC.
51+
*/
52+
public class TimestampResultSetTest {
53+
private static final MockFlightSqlProducer FLIGHT_SQL_PRODUCER = new MockFlightSqlProducer();
54+
55+
@RegisterExtension public static FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION;
56+
57+
static {
58+
FLIGHT_SERVER_TEST_EXTENSION =
59+
FlightServerTestExtension.createStandardTestExtension(FLIGHT_SQL_PRODUCER);
60+
}
61+
62+
private static final String QUERY_STRING = "SELECT * FROM TIMESTAMPS";
63+
private static final Schema QUERY_SCHEMA =
64+
new Schema(
65+
ImmutableList.of(
66+
Field.nullable("no_tz", new ArrowType.Timestamp(TimeUnit.MILLISECOND, null)),
67+
Field.nullable("utc", new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")),
68+
Field.nullable("utc+1", new ArrowType.Timestamp(TimeUnit.MILLISECOND, "GMT+1")),
69+
Field.nullable("utc-1", new ArrowType.Timestamp(TimeUnit.MILLISECOND, "GMT-1"))));
70+
71+
@BeforeAll
72+
public static void setup() throws SQLException {
73+
Instant firstDay2025 = OffsetDateTime.of(2025, 1, 1, 0, 0, 0, 0, ZoneOffset.UTC).toInstant();
74+
75+
FLIGHT_SQL_PRODUCER.addSelectQuery(
76+
QUERY_STRING,
77+
QUERY_SCHEMA,
78+
Collections.singletonList(
79+
listener -> {
80+
try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
81+
final VectorSchemaRoot root = VectorSchemaRoot.create(QUERY_SCHEMA, allocator)) {
82+
listener.start(root);
83+
root.getFieldVectors()
84+
.forEach(v -> ((TimeStampVector) v).setSafe(0, firstDay2025.toEpochMilli()));
85+
root.setRowCount(1);
86+
listener.putNext();
87+
} catch (final Throwable throwable) {
88+
listener.error(throwable);
89+
} finally {
90+
listener.completed();
91+
}
92+
}));
93+
}
94+
95+
/**
96+
* This test doesn't yet test anything other than ensuring all ResultSet methods to retrieve a
97+
* timestamp succeed.
98+
*
99+
* <p>This is a good starting point to add more tests to ensure the values are correct when we
100+
* change the "local calendar" either through changing the JVM default or through the connection
101+
* property.
102+
*/
103+
@Test
104+
public void test() {
105+
TimeZone.setDefault(TimeZone.getTimeZone("UTC"));
106+
try (Connection connection = FLIGHT_SERVER_TEST_EXTENSION.getConnection("UTC")) {
107+
try (PreparedStatement s = connection.prepareStatement(QUERY_STRING)) {
108+
try (ResultSet rs = s.executeQuery()) {
109+
int numCols = rs.getMetaData().getColumnCount();
110+
try {
111+
rs.next();
112+
for (int i = 1; i <= numCols; i++) {
113+
int type = rs.getMetaData().getColumnType(i);
114+
String name = rs.getMetaData().getColumnName(i);
115+
System.out.println(name);
116+
System.out.print("- getDate:\t\t\t\t\t\t\t");
117+
System.out.print(rs.getDate(i));
118+
System.out.println();
119+
System.out.print("- getTimestamp:\t\t\t\t\t\t");
120+
System.out.print(rs.getTimestamp(i));
121+
System.out.println();
122+
System.out.print("- getString:\t\t\t\t\t\t");
123+
System.out.print(rs.getString(i));
124+
System.out.println();
125+
System.out.print("- getObject:\t\t\t\t\t\t");
126+
System.out.print(rs.getObject(i));
127+
System.out.println();
128+
System.out.print("- getObject(Timestamp.class):\t\t");
129+
System.out.print(rs.getObject(i, Timestamp.class));
130+
System.out.println();
131+
System.out.print("- getTimestamp(default Calendar):\t");
132+
System.out.print(rs.getTimestamp(i, Calendar.getInstance()));
133+
System.out.println();
134+
System.out.print("- getTimestamp(UTC Calendar):\t\t");
135+
System.out.print(
136+
rs.getTimestamp(i, Calendar.getInstance(TimeZone.getTimeZone("UTC"))));
137+
System.out.println();
138+
System.out.print("- getObject(LocalDateTime.class):\t");
139+
System.out.print(rs.getObject(i, LocalDateTime.class));
140+
System.out.println();
141+
if (type == Types.TIMESTAMP_WITH_TIMEZONE) {
142+
System.out.print("- getObject(Instant.class):\t\t\t");
143+
System.out.print(rs.getObject(i, Instant.class));
144+
System.out.println();
145+
System.out.print("- getObject(OffsetDateTime.class):\t");
146+
System.out.print(rs.getObject(i, OffsetDateTime.class));
147+
System.out.println();
148+
System.out.print("- getObject(ZonedDateTime.class):\t");
149+
System.out.print(rs.getObject(i, ZonedDateTime.class));
150+
System.out.println();
151+
}
152+
System.out.println();
153+
}
154+
System.out.println();
155+
} catch (SQLException e) {
156+
throw new RuntimeException(e);
157+
}
158+
}
159+
}
160+
} catch (SQLException e) {
161+
throw new RuntimeException(e);
162+
}
163+
}
164+
}

0 commit comments

Comments
 (0)