Skip to content

Commit 10d9583

Browse files
aiguoferaiguofer
committed
Add test
1 parent b0b497a commit 10d9583

File tree

2 files changed

+181
-0
lines changed

2 files changed

+181
-0
lines changed

‎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: 175 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,175 @@
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.TimeStampMilliTZVector;
38+
import org.apache.arrow.vector.TimeStampMilliVector;
39+
import org.apache.arrow.vector.VectorSchemaRoot;
40+
import org.apache.arrow.vector.types.TimeUnit;
41+
import org.apache.arrow.vector.types.pojo.ArrowType;
42+
import org.apache.arrow.vector.types.pojo.Field;
43+
import org.apache.arrow.vector.types.pojo.Schema;
44+
import org.junit.jupiter.api.BeforeAll;
45+
import org.junit.jupiter.api.Test;
46+
import org.junit.jupiter.api.extension.RegisterExtension;
47+
48+
/**
49+
* Timestamps have a lot of nuances in JDBC. This class is here to test that timestamp behavior is
50+
* correct for different types of Timestamp vectors as well as different methods of retrieving the
51+
* timestamps in JDBC.
52+
*/
53+
public class TimestampResultSetTest {
54+
private static final MockFlightSqlProducer FLIGHT_SQL_PRODUCER = new MockFlightSqlProducer();
55+
56+
@RegisterExtension public static FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION;
57+
58+
static {
59+
FLIGHT_SERVER_TEST_EXTENSION =
60+
FlightServerTestExtension.createStandardTestExtension(FLIGHT_SQL_PRODUCER);
61+
}
62+
63+
private static final String QUERY_STRING = "SELECT * FROM TIMESTAMPS";
64+
private static final Schema QUERY_SCHEMA =
65+
new Schema(
66+
ImmutableList.of(
67+
Field.nullable("no_tz", new ArrowType.Timestamp(TimeUnit.MILLISECOND, null)),
68+
Field.nullable("utc", new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")),
69+
Field.nullable("utc+1", new ArrowType.Timestamp(TimeUnit.MILLISECOND, "GMT+1:00")),
70+
Field.nullable("utc-1", new ArrowType.Timestamp(TimeUnit.MILLISECOND, "GMT-1:00"))));
71+
72+
@BeforeAll
73+
public static void setup() throws SQLException {
74+
Instant firstDay2025 = OffsetDateTime.of(2025, 1, 1, 0, 0, 0, 0, ZoneOffset.UTC).toInstant();
75+
76+
FLIGHT_SQL_PRODUCER.addSelectQuery(
77+
QUERY_STRING,
78+
QUERY_SCHEMA,
79+
Collections.singletonList(
80+
listener -> {
81+
try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
82+
final VectorSchemaRoot root = VectorSchemaRoot.create(QUERY_SCHEMA, allocator)) {
83+
listener.start(root);
84+
root.getFieldVectors()
85+
.forEach(
86+
v -> {
87+
if (v instanceof TimeStampMilliVector) {
88+
((TimeStampMilliVector) v).setSafe(0, firstDay2025.toEpochMilli());
89+
} else if (v instanceof TimeStampMilliTZVector) {
90+
String vecTz =
91+
((ArrowType.Timestamp) v.getField().getType()).getTimezone();
92+
long millis = firstDay2025.toEpochMilli();
93+
((TimeStampMilliTZVector) v).setSafe(0, millis);
94+
}
95+
});
96+
root.setRowCount(1);
97+
listener.putNext();
98+
} catch (final Throwable throwable) {
99+
listener.error(throwable);
100+
} finally {
101+
listener.completed();
102+
}
103+
}));
104+
}
105+
106+
/**
107+
* This test doesn't yet test anything other than ensuring all ResultSet methods to retrieve a
108+
* timestamp succeed.
109+
*
110+
* <p>This is a good starting point to add more tests to ensure the values are correct when we
111+
* change the "local calendar" either through changing the JVM default or through the connection
112+
* property.
113+
*/
114+
@Test
115+
public void test() {
116+
TimeZone.setDefault(TimeZone.getTimeZone("UTC"));
117+
try (Connection connection = FLIGHT_SERVER_TEST_EXTENSION.getConnection("UTC")) {
118+
try (PreparedStatement s = connection.prepareStatement(QUERY_STRING)) {
119+
try (ResultSet rs = s.executeQuery()) {
120+
int numCols = rs.getMetaData().getColumnCount();
121+
try {
122+
rs.next();
123+
for (int i = 1; i <= numCols; i++) {
124+
int type = rs.getMetaData().getColumnType(i);
125+
String name = rs.getMetaData().getColumnName(i);
126+
System.out.println(name);
127+
System.out.print("- getDate:\t\t\t\t\t\t\t");
128+
System.out.print(rs.getDate(i));
129+
System.out.println();
130+
System.out.print("- getTimestamp:\t\t\t\t\t\t");
131+
System.out.print(rs.getTimestamp(i));
132+
System.out.println();
133+
System.out.print("- getString:\t\t\t\t\t\t");
134+
System.out.print(rs.getString(i));
135+
System.out.println();
136+
System.out.print("- getObject:\t\t\t\t\t\t");
137+
System.out.print(rs.getObject(i));
138+
System.out.println();
139+
System.out.print("- getObject(Timestamp.class):\t\t");
140+
System.out.print(rs.getObject(i, Timestamp.class));
141+
System.out.println();
142+
System.out.print("- getTimestamp(default Calendar):\t");
143+
System.out.print(rs.getTimestamp(i, Calendar.getInstance()));
144+
System.out.println();
145+
System.out.print("- getTimestamp(UTC Calendar):\t\t");
146+
System.out.print(
147+
rs.getTimestamp(i, Calendar.getInstance(TimeZone.getTimeZone("UTC"))));
148+
System.out.println();
149+
System.out.print("- getObject(LocalDateTime.class):\t");
150+
System.out.print(rs.getObject(i, LocalDateTime.class));
151+
System.out.println();
152+
if (type == Types.TIMESTAMP_WITH_TIMEZONE) {
153+
System.out.print("- getObject(Instant.class):\t\t\t");
154+
System.out.print(rs.getObject(i, Instant.class));
155+
System.out.println();
156+
System.out.print("- getObject(OffsetDateTime.class):\t");
157+
System.out.print(rs.getObject(i, OffsetDateTime.class));
158+
System.out.println();
159+
System.out.print("- getObject(ZonedDateTime.class):\t");
160+
System.out.print(rs.getObject(i, ZonedDateTime.class));
161+
System.out.println();
162+
}
163+
System.out.println();
164+
}
165+
System.out.println();
166+
} catch (SQLException e) {
167+
throw new RuntimeException(e);
168+
}
169+
}
170+
}
171+
} catch (SQLException e) {
172+
throw new RuntimeException(e);
173+
}
174+
}
175+
}

0 commit comments

Comments
 (0)