Skip to content

Commit fd7b80c

Browse files
DaZuiZuiDaZuiZui
authored
fix(jdbc): forward column mapping in table-model PreparedStatement
1 parent 2cd3b07 commit fd7b80c

4 files changed

Lines changed: 328 additions & 1 deletion

File tree

‎integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBTablePreparedStatementIT.java‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -382,4 +382,39 @@ public void testPrepareAndExecuteWithStringParameter() {
382382
fail("testPrepareAndExecuteWithStringParameter failed: " + e.getMessage());
383383
}
384384
}
385+
386+
@Test
387+
public void testPrepareAndExecuteWithReorderedProjection() {
388+
String[] expectedHeader = new String[] {"value", "name"};
389+
String[] retArray = new String[] {"200.3,Bob,"};
390+
try (Connection connection = EnvFactory.getEnv().getTableConnection();
391+
Statement statement = connection.createStatement()) {
392+
statement.execute("USE " + DATABASE_NAME);
393+
// Projection order differs from the table's column order
394+
statement.execute("PREPARE stmt9 FROM SELECT value, name FROM test_table WHERE id = ?");
395+
executePreparedStatementAndVerify(
396+
connection, statement, "EXECUTE stmt9 USING 2", expectedHeader, retArray);
397+
statement.execute("DEALLOCATE PREPARE stmt9");
398+
} catch (SQLException e) {
399+
fail("testPrepareAndExecuteWithReorderedProjection failed: " + e.getMessage());
400+
}
401+
}
402+
403+
@Test
404+
public void testPrepareAndExecuteWithRepeatedProjection() {
405+
String[] expectedHeader = new String[] {"x", "y"};
406+
String[] retArray = new String[] {"Bob,Bob,"};
407+
try (Connection connection = EnvFactory.getEnv().getTableConnection();
408+
Statement statement = connection.createStatement()) {
409+
statement.execute("USE " + DATABASE_NAME);
410+
// The same physical column is projected twice under different aliases
411+
statement.execute(
412+
"PREPARE stmt10 FROM SELECT name AS x, name AS y FROM test_table WHERE id = ?");
413+
executePreparedStatementAndVerify(
414+
connection, statement, "EXECUTE stmt10 USING 2", expectedHeader, retArray);
415+
statement.execute("DEALLOCATE PREPARE stmt10");
416+
} catch (SQLException e) {
417+
fail("testPrepareAndExecuteWithRepeatedProjection failed: " + e.getMessage());
418+
}
419+
}
385420
}

‎integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBTablePreparedStatementJDBCIT.java‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -495,4 +495,109 @@ public void testPreparedStatementWithAllDataTypes() throws SQLException {
495495
}
496496
}
497497
}
498+
499+
@Test
500+
public void testPreparedStatementWithTimestampProjection() throws SQLException {
501+
try (Connection connection = getConnection();
502+
Statement stmt = connection.createStatement()) {
503+
stmt.execute("USE " + DATABASE_NAME);
504+
// Self-contained fixture: the all-types table is created by another test method
505+
stmt.execute(
506+
"CREATE TABLE IF NOT EXISTS timestamp_projection_table("
507+
+ "int32_col INT32 FIELD, timestamp_col TIMESTAMP FIELD)");
508+
stmt.execute(
509+
"INSERT INTO timestamp_projection_table VALUES "
510+
+ "(2025-06-01T00:00:00, 123, 2025-06-15T10:30:00)");
511+
512+
// TIMESTAMP is a valid table-model field type and must keep its mapping as well
513+
try (PreparedStatement ps =
514+
connection.prepareStatement(
515+
"SELECT timestamp_col, int32_col FROM timestamp_projection_table "
516+
+ "WHERE int32_col = ?")) {
517+
ps.setInt(1, 123);
518+
try (ResultSet rs = ps.executeQuery()) {
519+
assertTrue(rs.next());
520+
assertEquals(
521+
Timestamp.valueOf("2025-06-15 10:30:00").getTime(), rs.getTimestamp(1).getTime());
522+
assertEquals(123, rs.getInt(2));
523+
assertEquals(
524+
Timestamp.valueOf("2025-06-15 10:30:00").getTime(),
525+
rs.getTimestamp("timestamp_col").getTime());
526+
assertEquals(123, rs.getInt("int32_col"));
527+
assertFalse(rs.next());
528+
}
529+
}
530+
}
531+
}
532+
533+
@Test
534+
public void testPreparedStatementWithReorderedProjection() throws SQLException {
535+
try (Connection connection = getConnection();
536+
Statement stmt = connection.createStatement()) {
537+
stmt.execute("USE " + DATABASE_NAME);
538+
539+
// Projection order ("value", "name") differs from the materialized order
540+
try (PreparedStatement ps =
541+
connection.prepareStatement("SELECT value, name FROM test_table WHERE int_value = ?")) {
542+
ps.setInt(1, 20);
543+
try (ResultSet rs = ps.executeQuery()) {
544+
assertTrue(rs.next());
545+
ResultSetMetaData metaData = rs.getMetaData();
546+
assertEquals(2, metaData.getColumnCount());
547+
assertEquals("value", metaData.getColumnName(1));
548+
assertEquals("name", metaData.getColumnName(2));
549+
assertEquals(200.3, rs.getDouble(1), 0.001);
550+
assertEquals("Bob", rs.getString(2));
551+
assertEquals("Bob", rs.getString("name"));
552+
assertEquals(200.3, rs.getDouble("value"), 0.001);
553+
assertFalse(rs.next());
554+
}
555+
}
556+
}
557+
}
558+
559+
@Test
560+
public void testPreparedStatementWithRepeatedProjection() throws SQLException {
561+
try (Connection connection = getConnection();
562+
Statement stmt = connection.createStatement()) {
563+
stmt.execute("USE " + DATABASE_NAME);
564+
565+
// The same physical column is projected twice under different aliases
566+
try (PreparedStatement ps =
567+
connection.prepareStatement(
568+
"SELECT name AS x, name AS y FROM test_table WHERE int_value = ?")) {
569+
ps.setInt(1, 20);
570+
try (ResultSet rs = ps.executeQuery()) {
571+
assertTrue(rs.next());
572+
assertEquals("Bob", rs.getString(1));
573+
assertEquals("Bob", rs.getString(2));
574+
assertEquals("Bob", rs.getString("x"));
575+
assertEquals("Bob", rs.getString("y"));
576+
assertFalse(rs.next());
577+
}
578+
}
579+
}
580+
}
581+
582+
@Test
583+
public void testPreparedStatementWithIdentityProjection() throws SQLException {
584+
try (Connection connection = getConnection();
585+
Statement stmt = connection.createStatement()) {
586+
stmt.execute("USE " + DATABASE_NAME);
587+
588+
// Identity mapping: a regression control for the two cases above
589+
try (PreparedStatement ps =
590+
connection.prepareStatement("SELECT name, value FROM test_table WHERE int_value = ?")) {
591+
ps.setInt(1, 20);
592+
try (ResultSet rs = ps.executeQuery()) {
593+
assertTrue(rs.next());
594+
assertEquals("Bob", rs.getString(1));
595+
assertEquals(200.3, rs.getDouble(2), 0.001);
596+
assertEquals("Bob", rs.getString("name"));
597+
assertEquals(200.3, rs.getDouble("value"), 0.001);
598+
assertFalse(rs.next());
599+
}
600+
}
601+
}
602+
}
498603
}

‎iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/IoTDBTablePreparedStatement.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -233,7 +233,12 @@ private ResultSet processQueryResult(TSExecuteStatementResp resp) throws SQLExce
233233
resp.tracingInfo,
234234
(long) queryTimeout * 1000,
235235
resp.isSetMoreData() && resp.isMoreData(),
236-
zoneId);
236+
zoneId,
237+
charset,
238+
resp.isSetTableModel() && resp.isTableModel(),
239+
resp.getColumnIndex2TsBlockColumnIndexList() == null
240+
? null
241+
: new ArrayList<>(resp.getColumnIndex2TsBlockColumnIndexList()));
237242
return resultSet;
238243
}
239244
return null;

‎iotdb-client/jdbc/src/test/java/org/apache/iotdb/jdbc/IoTDBTablePreparedStatementTest.java‎

Lines changed: 182 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,12 +25,17 @@
2525
import org.apache.iotdb.rpc.stmt.PreparedParameterSerde;
2626
import org.apache.iotdb.service.rpc.thrift.IClientRPCService.Iface;
2727
import org.apache.iotdb.service.rpc.thrift.TSCancelOperationReq;
28+
import org.apache.iotdb.service.rpc.thrift.TSCloseOperationReq;
2829
import org.apache.iotdb.service.rpc.thrift.TSExecutePreparedReq;
2930
import org.apache.iotdb.service.rpc.thrift.TSExecuteStatementReq;
3031
import org.apache.iotdb.service.rpc.thrift.TSExecuteStatementResp;
3132
import org.apache.iotdb.service.rpc.thrift.TSPrepareReq;
3233
import org.apache.iotdb.service.rpc.thrift.TSPrepareResp;
3334

35+
import org.apache.tsfile.enums.TSDataType;
36+
import org.apache.tsfile.read.common.block.TsBlockBuilder;
37+
import org.apache.tsfile.read.common.block.column.TsBlockSerde;
38+
import org.apache.tsfile.utils.Binary;
3439
import org.junit.Before;
3540
import org.junit.Test;
3641
import org.mockito.ArgumentCaptor;
@@ -39,17 +44,23 @@
3944
import org.mockito.invocation.InvocationOnMock;
4045
import org.mockito.stubbing.Answer;
4146

47+
import java.io.IOException;
4248
import java.nio.ByteBuffer;
49+
import java.nio.charset.Charset;
50+
import java.sql.ResultSet;
4351
import java.sql.SQLException;
4452
import java.sql.Timestamp;
4553
import java.sql.Types;
4654
import java.time.Instant;
4755
import java.time.LocalDate;
4856
import java.time.ZoneId;
57+
import java.util.ArrayList;
4958
import java.util.Arrays;
59+
import java.util.Collections;
5060
import java.util.List;
5161
import java.util.stream.Collectors;
5262

63+
import static org.junit.Assert.assertArrayEquals;
5364
import static org.junit.Assert.assertEquals;
5465
import static org.junit.Assert.assertThrows;
5566
import static org.junit.Assert.assertTrue;
@@ -75,6 +86,7 @@ public class IoTDBTablePreparedStatementTest {
7586
public void before() throws Exception {
7687
MockitoAnnotations.initMocks(this);
7788
when(connection.getSqlDialect()).thenReturn("table");
89+
when(connection.getTimeFactor()).thenReturn(1_000);
7890
when(execStatementResp.getStatus()).thenReturn(Status_SUCCESS);
7991
when(execStatementResp.getQueryId()).thenReturn(queryId);
8092

@@ -98,6 +110,8 @@ public TSPrepareResp answer(InvocationOnMock invocation) throws Throwable {
98110
// Mock for executePreparedStatement
99111
when(client.executePreparedStatement(any(TSExecutePreparedReq.class)))
100112
.thenReturn(execStatementResp);
113+
114+
when(client.closeOperation(any(TSCloseOperationReq.class))).thenReturn(Status_SUCCESS);
101115
}
102116

103117
/** Count the number of '?' placeholders in a SQL string, ignoring those inside quotes */
@@ -255,6 +269,174 @@ public void testTableModelStringWithNull() throws Exception {
255269
assertTrue(argument.getValue().getParameters() != null);
256270
}
257271

272+
// ========== Table Model result column mapping tests ==========
273+
274+
private static TsBlockBuilder int32BlockBuilder(int columnCount) {
275+
TsBlockBuilder builder = new TsBlockBuilder(Collections.nCopies(columnCount, TSDataType.INT32));
276+
builder.getTimeColumnBuilder().writeLong(1000L);
277+
builder.declarePosition();
278+
return builder;
279+
}
280+
281+
private TSExecuteStatementResp buildQueryResp(
282+
List<String> columns, List<Integer> mapping, TsBlockBuilder blockBuilder) throws IOException {
283+
return buildQueryResp(
284+
columns, Collections.nCopies(columns.size(), TSDataType.INT32), mapping, blockBuilder);
285+
}
286+
287+
private TSExecuteStatementResp buildQueryResp(
288+
List<String> columns,
289+
List<TSDataType> dataTypes,
290+
List<Integer> mapping,
291+
TsBlockBuilder blockBuilder)
292+
throws IOException {
293+
TSExecuteStatementResp resp = new TSExecuteStatementResp();
294+
resp.setStatus(Status_SUCCESS);
295+
resp.setQueryId(1L);
296+
resp.setColumns(columns);
297+
List<String> dataTypeList = new ArrayList<>();
298+
for (TSDataType dataType : dataTypes) {
299+
dataTypeList.add(dataType.name());
300+
}
301+
resp.setDataTypeList(dataTypeList);
302+
resp.setIgnoreTimeStamp(true);
303+
resp.setTableModel(true);
304+
if (mapping != null) {
305+
resp.setColumnIndex2TsBlockColumnIndexList(mapping);
306+
}
307+
resp.setQueryResult(
308+
Collections.singletonList(new TsBlockSerde().serialize(blockBuilder.build())));
309+
resp.setMoreData(false);
310+
return resp;
311+
}
312+
313+
@SuppressWarnings("resource")
314+
@Test
315+
public void testTableModelPreparedStatementReorderedColumns() throws Exception {
316+
TsBlockBuilder blockBuilder = int32BlockBuilder(2);
317+
blockBuilder.getColumnBuilder(0).writeInt(10); // physical column 0
318+
blockBuilder.getColumnBuilder(1).writeInt(20); // physical column 1
319+
320+
when(client.executePreparedStatement(any(TSExecutePreparedReq.class)))
321+
.thenReturn(buildQueryResp(Arrays.asList("b", "a"), Arrays.asList(1, 0), blockBuilder));
322+
323+
IoTDBTablePreparedStatement ps =
324+
new IoTDBTablePreparedStatement(
325+
connection, client, sessionId, "SELECT b, a FROM t WHERE a > ?", zoneId);
326+
ps.setInt(1, 0);
327+
try (ResultSet rs = ps.executeQuery()) {
328+
assertTrue(rs.next());
329+
assertEquals(20, rs.getInt(1));
330+
assertEquals(10, rs.getInt(2));
331+
assertEquals(20, rs.getInt("b"));
332+
assertEquals(10, rs.getInt("a"));
333+
}
334+
}
335+
336+
@SuppressWarnings("resource")
337+
@Test
338+
public void testTableModelPreparedStatementRepeatedProjection() throws Exception {
339+
TsBlockBuilder blockBuilder = int32BlockBuilder(1);
340+
blockBuilder.getColumnBuilder(0).writeInt(10); // physical column 0, referenced twice
341+
342+
when(client.executePreparedStatement(any(TSExecutePreparedReq.class)))
343+
.thenReturn(buildQueryResp(Arrays.asList("x", "y"), Arrays.asList(0, 0), blockBuilder));
344+
345+
IoTDBTablePreparedStatement ps =
346+
new IoTDBTablePreparedStatement(
347+
connection, client, sessionId, "SELECT a AS x, a AS y FROM t WHERE a > ?", zoneId);
348+
ps.setInt(1, 0);
349+
try (ResultSet rs = ps.executeQuery()) {
350+
assertTrue(rs.next());
351+
assertEquals(10, rs.getInt(1));
352+
assertEquals(10, rs.getInt(2));
353+
assertEquals(10, rs.getInt("x"));
354+
assertEquals(10, rs.getInt("y"));
355+
}
356+
}
357+
358+
@SuppressWarnings("resource")
359+
@Test
360+
public void testTableModelPreparedStatementIdentityMapping() throws Exception {
361+
TsBlockBuilder blockBuilder = int32BlockBuilder(2);
362+
blockBuilder.getColumnBuilder(0).writeInt(10);
363+
blockBuilder.getColumnBuilder(1).writeInt(20);
364+
365+
when(client.executePreparedStatement(any(TSExecutePreparedReq.class)))
366+
.thenReturn(buildQueryResp(Arrays.asList("a", "b"), Arrays.asList(0, 1), blockBuilder));
367+
368+
IoTDBTablePreparedStatement ps =
369+
new IoTDBTablePreparedStatement(
370+
connection, client, sessionId, "SELECT a, b FROM t WHERE a > ?", zoneId);
371+
ps.setInt(1, 0);
372+
try (ResultSet rs = ps.executeQuery()) {
373+
assertTrue(rs.next());
374+
assertEquals(10, rs.getInt(1));
375+
assertEquals(20, rs.getInt(2));
376+
assertEquals(10, rs.getInt("a"));
377+
assertEquals(20, rs.getInt("b"));
378+
}
379+
}
380+
381+
@SuppressWarnings("resource")
382+
@Test
383+
public void testTableModelPreparedStatementMissingMappingFallsBack() throws Exception {
384+
TsBlockBuilder blockBuilder = int32BlockBuilder(2);
385+
blockBuilder.getColumnBuilder(0).writeInt(10);
386+
blockBuilder.getColumnBuilder(1).writeInt(20);
387+
388+
// No columnIndex2TsBlockColumnIndexList: the client must keep its own derivation
389+
when(client.executePreparedStatement(any(TSExecutePreparedReq.class)))
390+
.thenReturn(buildQueryResp(Arrays.asList("a", "b"), null, blockBuilder));
391+
392+
IoTDBTablePreparedStatement ps =
393+
new IoTDBTablePreparedStatement(
394+
connection, client, sessionId, "SELECT a, b FROM t WHERE a > ?", zoneId);
395+
ps.setInt(1, 0);
396+
try (ResultSet rs = ps.executeQuery()) {
397+
assertTrue(rs.next());
398+
assertEquals(10, rs.getInt(1));
399+
assertEquals(20, rs.getInt(2));
400+
assertEquals(10, rs.getInt("a"));
401+
assertEquals(20, rs.getInt("b"));
402+
}
403+
}
404+
405+
// ========== Charset forwarding ==========
406+
407+
@SuppressWarnings("resource")
408+
@Test
409+
public void testTableModelPreparedStatementGetBytesUsesConnectionCharset() throws Exception {
410+
// A STRING column holding 中文; the value travels as UTF-8 in the TsBlock, but getBytes must
411+
// encode the decoded String with the charset the connection was opened with.
412+
TsBlockBuilder blockBuilder = new TsBlockBuilder(Collections.singletonList(TSDataType.STRING));
413+
blockBuilder.getTimeColumnBuilder().writeLong(1000L);
414+
blockBuilder.declarePosition();
415+
blockBuilder.getColumnBuilder(0).writeBinary(new Binary("中文", Charset.forName("UTF-8")));
416+
when(client.executePreparedStatement(any(TSExecutePreparedReq.class)))
417+
.thenReturn(
418+
buildQueryResp(
419+
Collections.singletonList("name"),
420+
Collections.singletonList(TSDataType.STRING),
421+
null,
422+
blockBuilder));
423+
424+
IoTDBTablePreparedStatement ps =
425+
new IoTDBTablePreparedStatement(
426+
connection,
427+
client,
428+
sessionId,
429+
"SELECT name FROM t WHERE name = ?",
430+
zoneId,
431+
Charset.forName("GBK"));
432+
ps.setString(1, "中文");
433+
try (ResultSet rs = ps.executeQuery()) {
434+
assertTrue(rs.next());
435+
assertArrayEquals("中文".getBytes(Charset.forName("GBK")), rs.getBytes(1));
436+
assertArrayEquals("中文".getBytes(Charset.forName("GBK")), rs.getBytes("name"));
437+
}
438+
}
439+
258440
// ========== Cancel, fetch size and timestamp precision ==========
259441

260442
@SuppressWarnings("resource")

0 commit comments

Comments
 (0)