Skip to content

Commit 3401348

Browse files
dmitry-chirkov-dremiodmitry-chirkov-dremio
committed
Address review feedback: defensive copy, setObject override, negative epoch tests
1 parent 015a689 commit 3401348

5 files changed

Lines changed: 182 additions & 6 deletions

File tree

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

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,7 @@ public synchronized void close() throws SQLException {
8383
@Override
8484
public void setTimestamp(int parameterIndex, Timestamp x) throws SQLException {
8585
if (x != null) {
86-
rawTimestamps.put(parameterIndex, x);
86+
rawTimestamps.put(parameterIndex, (Timestamp) x.clone());
8787
} else {
8888
rawTimestamps.remove(parameterIndex);
8989
}
@@ -93,13 +93,32 @@ public void setTimestamp(int parameterIndex, Timestamp x) throws SQLException {
9393
@Override
9494
public void setTimestamp(int parameterIndex, Timestamp x, Calendar cal) throws SQLException {
9595
if (x != null) {
96-
rawTimestamps.put(parameterIndex, x);
96+
rawTimestamps.put(parameterIndex, (Timestamp) x.clone());
9797
} else {
9898
rawTimestamps.remove(parameterIndex);
9999
}
100100
super.setTimestamp(parameterIndex, x, cal);
101101
}
102102

103+
@Override
104+
public void setObject(int parameterIndex, Object x, int targetSqlType) throws SQLException {
105+
rawTimestamps.remove(parameterIndex);
106+
super.setObject(parameterIndex, x, targetSqlType);
107+
}
108+
109+
@Override
110+
public void setObject(int parameterIndex, Object x) throws SQLException {
111+
rawTimestamps.remove(parameterIndex);
112+
super.setObject(parameterIndex, x);
113+
}
114+
115+
@Override
116+
public void setObject(int parameterIndex, Object x, int targetSqlType, int scaleOrLength)
117+
throws SQLException {
118+
rawTimestamps.remove(parameterIndex);
119+
super.setObject(parameterIndex, x, targetSqlType, scaleOrLength);
120+
}
121+
103122
@Override
104123
public void clearParameters() throws SQLException {
105124
rawTimestamps.clear();

‎flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/TimestampAvaticaParameterConverter.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717
package org.apache.arrow.driver.jdbc.converter.impl;
1818

1919
import java.sql.Timestamp;
20-
import org.checkerframework.checker.nullness.qual.Nullable;
2120
import org.apache.arrow.vector.FieldVector;
2221
import org.apache.arrow.vector.TimeStampMicroTZVector;
2322
import org.apache.arrow.vector.TimeStampMicroVector;
@@ -31,6 +30,7 @@
3130
import org.apache.arrow.vector.types.pojo.Field;
3231
import org.apache.calcite.avatica.AvaticaParameter;
3332
import org.apache.calcite.avatica.remote.TypedValue;
33+
import org.checkerframework.checker.nullness.qual.Nullable;
3434

3535
/** AvaticaParameterConverter for Timestamp Arrow types. */
3636
public class TimestampAvaticaParameterConverter extends BaseAvaticaParameterConverter {
@@ -69,7 +69,7 @@ private long convertFromTimestamp(Timestamp ts) {
6969
private long convertFromMillis(long epochMillis) {
7070
switch (type.getUnit()) {
7171
case SECOND:
72-
return epochMillis / 1_000L;
72+
return Math.floorDiv(epochMillis, 1_000L);
7373
case MILLISECOND:
7474
return epochMillis;
7575
case MICROSECOND:

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -105,8 +105,7 @@ public void bind(List<TypedValue> typedValues, int index) {
105105
* @param index index for parameter.
106106
* @param rawTimestamps Raw java.sql.Timestamp objects keyed by 1-based parameter index.
107107
*/
108-
public void bind(
109-
List<TypedValue> typedValues, int index, Map<Integer, Timestamp> rawTimestamps) {
108+
public void bind(List<TypedValue> typedValues, int index, Map<Integer, Timestamp> rawTimestamps) {
110109
if (preparedStatement.getParameterSchema().getFields().size() != typedValues.size()) {
111110
throw new IllegalStateException(
112111
String.format(

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

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -304,4 +304,117 @@ public void testTimestampParameterWithMillisecondPrecision() throws SQLException
304304
assertEquals(1, updated);
305305
}
306306
}
307+
308+
@Test
309+
public void testTimestampMutationAfterSetDoesNotAffectBoundValue() throws SQLException {
310+
String query = "Fake timestamp mutation test";
311+
Schema parameterSchema =
312+
new Schema(
313+
Collections.singletonList(
314+
Field.nullable("ts", new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC"))));
315+
316+
// Original: epochSeconds=1730637909, nanos=869885001 → micros = 1730637909869885
317+
List<List<Object>> expected =
318+
Collections.singletonList(Collections.singletonList(1730637909869885L));
319+
320+
PRODUCER.addUpdateQuery(query, 1);
321+
PRODUCER.addExpectedParameters(query, parameterSchema, expected);
322+
323+
try (PreparedStatement stmt = connection.prepareStatement(query)) {
324+
Timestamp ts = new Timestamp(1730637909869L);
325+
ts.setNanos(869885001);
326+
stmt.setTimestamp(1, ts);
327+
328+
// Mutate the Timestamp after setTimestamp — should not affect the stored value
329+
ts.setNanos(0);
330+
ts.setTime(0);
331+
332+
int updated = stmt.executeUpdate();
333+
assertEquals(1, updated);
334+
}
335+
}
336+
337+
@Test
338+
public void testTimestampSetObjectFallsBackToMillis() throws SQLException {
339+
String query = "Fake timestamp setObject test";
340+
Schema parameterSchema =
341+
new Schema(
342+
Collections.singletonList(
343+
Field.nullable("ts", new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC"))));
344+
345+
// setObject goes through Avatica's millis-only path: 1730637909869 * 1000 = 1730637909869000
346+
List<List<Object>> expected =
347+
Collections.singletonList(Collections.singletonList(1730637909869000L));
348+
349+
PRODUCER.addUpdateQuery(query, 1);
350+
PRODUCER.addExpectedParameters(query, parameterSchema, expected);
351+
352+
try (PreparedStatement stmt = connection.prepareStatement(query)) {
353+
Timestamp ts = new Timestamp(1730637909869L);
354+
ts.setNanos(869885001); // sub-ms nanos present but will be lost via setObject
355+
stmt.setObject(1, ts);
356+
int updated = stmt.executeUpdate();
357+
assertEquals(1, updated);
358+
}
359+
}
360+
361+
@Test
362+
public void testSetObjectAfterSetTimestampClearsRawTimestamp() throws SQLException {
363+
String query = "Fake timestamp setObject after setTimestamp";
364+
Schema parameterSchema =
365+
new Schema(
366+
Collections.singletonList(
367+
Field.nullable("ts", new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC"))));
368+
369+
// After setObject replaces setTimestamp, millis-only path is used:
370+
// 1000L (epoch millis) * 1000 = 1000000 (epoch micros)
371+
List<List<Object>> expected = Collections.singletonList(Collections.singletonList(1000000L));
372+
373+
PRODUCER.addUpdateQuery(query, 1);
374+
PRODUCER.addExpectedParameters(query, parameterSchema, expected);
375+
376+
try (PreparedStatement stmt = connection.prepareStatement(query)) {
377+
// First set with setTimestamp (populates rawTimestamps)
378+
Timestamp ts1 = new Timestamp(1730637909869L);
379+
ts1.setNanos(869885001);
380+
stmt.setTimestamp(1, ts1);
381+
382+
// Then replace with setObject (should clear rawTimestamps for this index)
383+
Timestamp ts2 = new Timestamp(1000L);
384+
stmt.setObject(1, ts2);
385+
386+
int updated = stmt.executeUpdate();
387+
assertEquals(1, updated);
388+
}
389+
}
390+
391+
@Test
392+
public void testSetTimestampAfterSetObjectPreservesSubMillis() throws SQLException {
393+
String query = "Fake timestamp setTimestamp after setObject";
394+
Schema parameterSchema =
395+
new Schema(
396+
Collections.singletonList(
397+
Field.nullable("ts", new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC"))));
398+
399+
// setTimestamp called last, so rawTimestamp is used: epochSeconds=1730637909, nanos=869885001
400+
List<List<Object>> expected =
401+
Collections.singletonList(Collections.singletonList(1730637909869885L));
402+
403+
PRODUCER.addUpdateQuery(query, 1);
404+
PRODUCER.addExpectedParameters(query, parameterSchema, expected);
405+
406+
try (PreparedStatement stmt = connection.prepareStatement(query)) {
407+
// First set with setObject (millis only)
408+
Timestamp ts1 = new Timestamp(1000L);
409+
stmt.setObject(1, ts1);
410+
411+
// Then replace with setTimestamp (populates rawTimestamps with sub-ms precision)
412+
Timestamp ts2 = new Timestamp(1730637909869L);
413+
ts2.setNanos(869885001);
414+
stmt.setTimestamp(1, ts2);
415+
416+
int updated = stmt.executeUpdate();
417+
assertEquals(1, updated);
418+
}
419+
}
307420
}

‎flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/converter/impl/TimestampAvaticaParameterConverterTest.java‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -201,6 +201,51 @@ public void testSecVectorTruncatesSubSecond() {
201201
}
202202
}
203203

204+
@Test
205+
public void testSecVectorNegativeEpochFloorDivision() {
206+
BufferAllocator allocator = rootAllocatorTestExtension.getRootAllocator();
207+
ArrowType.Timestamp type = new ArrowType.Timestamp(TimeUnit.SECOND, null);
208+
TimestampAvaticaParameterConverter converter = new TimestampAvaticaParameterConverter(type);
209+
210+
try (TimeStampSecVector vector = new TimeStampSecVector("ts", allocator)) {
211+
vector.allocateNew(3);
212+
// -999 millis: regular division gives 0 (wrong), floorDiv gives -1 (correct)
213+
TypedValue tv1 = TypedValue.ofLocal(ColumnMetaData.Rep.JAVA_SQL_TIMESTAMP, -999L);
214+
assertTrue(converter.bindParameter(vector, tv1, 0));
215+
assertEquals(-1L, vector.get(0));
216+
217+
// -1000 millis: exactly -1 second
218+
TypedValue tv2 = TypedValue.ofLocal(ColumnMetaData.Rep.JAVA_SQL_TIMESTAMP, -1000L);
219+
assertTrue(converter.bindParameter(vector, tv2, 1));
220+
assertEquals(-1L, vector.get(1));
221+
222+
// -1001 millis: floorDiv gives -2 (correct), regular division gives -1 (wrong)
223+
TypedValue tv3 = TypedValue.ofLocal(ColumnMetaData.Rep.JAVA_SQL_TIMESTAMP, -1001L);
224+
assertTrue(converter.bindParameter(vector, tv3, 2));
225+
assertEquals(-2L, vector.get(2));
226+
}
227+
}
228+
229+
@Test
230+
public void testNegativeEpochRawTimestampPreservesSubMillis() {
231+
BufferAllocator allocator = rootAllocatorTestExtension.getRootAllocator();
232+
ArrowType.Timestamp type = new ArrowType.Timestamp(TimeUnit.MICROSECOND, null);
233+
TimestampAvaticaParameterConverter converter = new TimestampAvaticaParameterConverter(type);
234+
235+
// 1969-12-31 23:59:59.000123456 UTC → epochMillis=-1000, nanos=123456
236+
// epochSeconds = floorDiv(-1000, 1000) = -1
237+
// micros = -1 * 1_000_000 + 123456/1000 = -1_000_000 + 123 = -999877
238+
Timestamp ts = new Timestamp(-1000L);
239+
ts.setNanos(123456);
240+
241+
try (TimeStampMicroVector vector = new TimeStampMicroVector("ts", allocator)) {
242+
vector.allocateNew(1);
243+
TypedValue typedValue = TypedValue.ofLocal(ColumnMetaData.Rep.JAVA_SQL_TIMESTAMP, -1000L);
244+
assertTrue(converter.bindParameter(vector, typedValue, 0, ts));
245+
assertEquals(-999877L, vector.get(0));
246+
}
247+
}
248+
204249
private void assertBindConvertsMillis(TimeUnit unit, String tz, long expectedValue) {
205250
BufferAllocator allocator = rootAllocatorTestExtension.getRootAllocator();
206251
ArrowType.Timestamp type = new ArrowType.Timestamp(unit, tz);

0 commit comments

Comments
 (0)