From e5a271504520e04172e0b6395ad860f9ad47dfbf Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Thu, 12 Nov 2026 12:07:23 +0800 Subject: [PATCH 1/6] Add database-scoped tablet insertion for table sessions --- .../it/session/IoTDBSessionRelationalIT.java | 57 +++++++++++++++++++ .../apache/iotdb/isession/ITableSession.java | 15 +++++ .../org/apache/iotdb/session/Session.java | 34 +++++++++-- .../apache/iotdb/session/TableSession.java | 6 ++ .../session/pool/TableSessionWrapper.java | 8 ++- .../plan/parser/StatementGenerator.java | 17 ++++++ .../src/main/thrift/client.thrift | 3 +- 7 files changed, 132 insertions(+), 8 deletions(-) diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBSessionRelationalIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBSessionRelationalIT.java index 73e9881dfcaa2..6cecfe79182a5 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBSessionRelationalIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBSessionRelationalIT.java @@ -177,6 +177,63 @@ public static void main(String[] args) } } + @Test + public void testInsertTabletIntoSpecifiedDatabase() + throws IoTDBConnectionException, StatementExecutionException { + try (ITableSession session = EnvFactory.getEnv().getTableSessionConnection()) { + session.executeNonQueryStatement( + "CREATE TABLE db1.tablet_database (tag1 STRING TAG, m1 DOUBLE FIELD)"); + session.executeNonQueryStatement( + "CREATE TABLE db2.tablet_database (tag1 STRING TAG, m1 DOUBLE FIELD)"); + + List columnNames = Arrays.asList("tag1", "m1"); + List dataTypes = Arrays.asList(TSDataType.STRING, TSDataType.DOUBLE); + List columnCategories = + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD); + Tablet tablet = new Tablet("tablet_database", columnNames, dataTypes, columnCategories, 2); + for (int row = 0; row < 2; row++) { + tablet.addTimestamp(row, row); + tablet.addValue("tag1", row, "tag:" + row); + tablet.addValue("m1", row, row * 1.0); + } + + session.insert(tablet, "db1"); + session.executeNonQueryStatement("USE db1"); + session.insert(tablet, "db2"); + Tablet currentDatabaseTablet = + new Tablet("tablet_database", columnNames, dataTypes, columnCategories, 1); + currentDatabaseTablet.addTimestamp(0, 2); + currentDatabaseTablet.addValue("tag1", 0, "tag:2"); + currentDatabaseTablet.addValue("m1", 0, 2.0); + session.insert(currentDatabaseTablet, null); + + try (SessionDataSet dataSet = + session.executeQueryStatement("SELECT * FROM db1.tablet_database")) { + assertTrue(dataSet.hasNext()); + dataSet.next(); + assertTrue(dataSet.hasNext()); + dataSet.next(); + assertTrue(dataSet.hasNext()); + dataSet.next(); + assertFalse(dataSet.hasNext()); + } + try (SessionDataSet dataSet = + session.executeQueryStatement("SELECT * FROM db2.tablet_database")) { + assertTrue(dataSet.hasNext()); + dataSet.next(); + assertTrue(dataSet.hasNext()); + dataSet.next(); + assertFalse(dataSet.hasNext()); + } + try (SessionDataSet dataSet = session.executeQueryStatement("SHOW CURRENT_DATABASE")) { + assertTrue(dataSet.hasNext()); + assertEquals("db1", dataSet.next().getField(0).getStringValue()); + } + + assertThrows(StatementExecutionException.class, () -> session.insert(tablet, "")); + } + } + private static void insertRelationalTabletPerformanceTest() throws IoTDBConnectionException, StatementExecutionException { try (ITableSession session = EnvFactory.getEnv().getTableSessionConnection()) { diff --git a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ITableSession.java b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ITableSession.java index 75a37f8384719..dfb7372f3ac34 100644 --- a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ITableSession.java +++ b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ITableSession.java @@ -46,6 +46,21 @@ public interface ITableSession extends AutoCloseable { */ void insert(Tablet tablet) throws StatementExecutionException, IoTDBConnectionException; + /** + * Inserts a {@link Tablet} into the specified database for this invocation only. + * + *

The database does not become the current database of this session and subsequent operations + * are unaffected. + * + * @param tablet the tablet containing time-series data to be inserted. + * @param database the target database for this insertion; if null, the current database of this + * session is used. + * @throws StatementExecutionException if an error occurs while executing the statement. + * @throws IoTDBConnectionException if there is an issue with the IoTDB connection. + */ + void insert(Tablet tablet, String database) + throws StatementExecutionException, IoTDBConnectionException; + /** * Executes a non-query SQL statement, such as a DDL or DML command. * diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java index d8bcce5e13def..59bdcf2b7af26 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java @@ -2907,11 +2907,22 @@ private void insertTabletInternal(Tablet tablet, TSInsertTabletReq request) */ public void insertRelationalTablet(Tablet tablet) throws IoTDBConnectionException, StatementExecutionException { + insertRelationalTablet(tablet, null); + } + + /** + * Inserts a relational {@link Tablet} into the specified database for this invocation only. + * + * @param tablet data batch + * @param database target database; if null, use the current database of this session + */ + public void insertRelationalTablet(Tablet tablet, String database) + throws IoTDBConnectionException, StatementExecutionException { if (tablet.getRowSize() == 0) { return; } if (enableRedirection) { - insertRelationalTabletWithLeaderCache(tablet); + insertRelationalTabletWithLeaderCache(tablet, database); } else { TSInsertTabletReq request = genTSInsertTabletReq(tablet, false, false); if (request == null) { @@ -2923,6 +2934,7 @@ public void insertRelationalTablet(Tablet tablet) return; } request.setWriteToTable(true); + setRelationalRequestDatabase(request, database); try { getDefaultSessionConnection().insertTablet(request); } catch (RedirectException ignored) { @@ -2930,7 +2942,7 @@ public void insertRelationalTablet(Tablet tablet) } } - private void insertRelationalTabletWithLeaderCache(Tablet tablet) + private void insertRelationalTabletWithLeaderCache(Tablet tablet, String database) throws IoTDBConnectionException, StatementExecutionException { Map relationalTabletGroup = new HashMap<>(); if (tableModelDeviceIdToEndpoint.isEmpty()) { @@ -2975,13 +2987,14 @@ private void insertRelationalTabletWithLeaderCache(Tablet tablet) } } if (relationalTabletGroup.size() == 1) { - insertRelationalTabletOnce(relationalTabletGroup); + insertRelationalTabletOnce(relationalTabletGroup, database); } else { - insertRelationalTabletByGroup(relationalTabletGroup); + insertRelationalTabletByGroup(relationalTabletGroup, database); } } - private void insertRelationalTabletOnce(Map relationalTabletGroup) + private void insertRelationalTabletOnce( + Map relationalTabletGroup, String database) throws IoTDBConnectionException, StatementExecutionException { Map.Entry entry = relationalTabletGroup.entrySet().iterator().next(); SessionConnection connection = entry.getKey(); @@ -2996,6 +3009,7 @@ private void insertRelationalTabletOnce(Map relationa return; } request.setWriteToTable(true); + setRelationalRequestDatabase(request, database); try { connection.insertTablet(request); } catch (RedirectException e) { @@ -3024,7 +3038,8 @@ private void insertRelationalTabletOnce(Map relationa @SuppressWarnings({ "squid:S3776" }) // ignore Cognitive Complexity of methods should not be too high - private void insertRelationalTabletByGroup(Map relationalTabletGroup) + private void insertRelationalTabletByGroup( + Map relationalTabletGroup, String database) throws IoTDBConnectionException, StatementExecutionException { List> completableFutures = relationalTabletGroup.entrySet().stream() @@ -3044,6 +3059,7 @@ private void insertRelationalTabletByGroup(Map relati return; } request.setWriteToTable(true); + setRelationalRequestDatabase(request, database); InsertConsumer insertConsumer = SessionConnection::insertTablet; try { @@ -3101,6 +3117,12 @@ private void insertRelationalTabletByGroup(Map relati } } + private static void setRelationalRequestDatabase(TSInsertTabletReq request, String database) { + if (database != null) { + request.setDatabase(database); + } + } + /** * insert the aligned timeseries data of a device. For each timestamp, the number of measurements * is the same. diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSession.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSession.java index 059fdea346948..b51a3a59831d8 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSession.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSession.java @@ -39,6 +39,12 @@ public void insert(Tablet tablet) throws StatementExecutionException, IoTDBConne session.insertRelationalTablet(tablet); } + @Override + public void insert(Tablet tablet, String database) + throws StatementExecutionException, IoTDBConnectionException { + session.insertRelationalTablet(tablet, database); + } + @Override public void executeNonQueryStatement(String sql) throws IoTDBConnectionException, StatementExecutionException { diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionWrapper.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionWrapper.java index 4cf7ca980f347..5cb7a89fc0eea 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionWrapper.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionWrapper.java @@ -58,8 +58,14 @@ protected TableSessionWrapper(Session session, SessionPool sessionPool) { @Override public void insert(Tablet tablet) throws StatementExecutionException, IoTDBConnectionException { + insert(tablet, null); + } + + @Override + public void insert(Tablet tablet, String database) + throws StatementExecutionException, IoTDBConnectionException { try { - session.insertRelationalTablet(tablet); + session.insertRelationalTablet(tablet, database); } catch (IoTDBConnectionException e) { sessionPool.cleanSessionAndMayThrowConnectionException(session); closed.set(true); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGenerator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGenerator.java index 046a94c1fa243..8fc1826994e1a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGenerator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGenerator.java @@ -24,11 +24,14 @@ import org.apache.iotdb.common.rpc.thrift.TAggregationType; import org.apache.iotdb.commons.exception.IllegalPathException; import org.apache.iotdb.commons.exception.MetadataException; +import org.apache.iotdb.commons.exception.SemanticException; import org.apache.iotdb.commons.path.MeasurementPath; import org.apache.iotdb.commons.path.PartialPath; import org.apache.iotdb.commons.queryengine.utils.TimestampPrecisionUtils; import org.apache.iotdb.commons.schema.table.column.TsTableColumnCategory; import org.apache.iotdb.commons.service.metric.PerformanceOverviewMetrics; +import org.apache.iotdb.commons.utils.PathUtils; +import org.apache.iotdb.db.i18n.DataNodeMiscMessages; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.qp.sql.IoTDBSqlParser; import org.apache.iotdb.db.qp.sql.SqlLexer; @@ -116,6 +119,7 @@ import java.util.HashMap; import java.util.LinkedHashMap; import java.util.List; +import java.util.Locale; import java.util.Map; import java.util.stream.Collectors; @@ -385,6 +389,19 @@ public static InsertTabletStatement createStatement(TSInsertTabletReq insertTabl ColumnCategory.values()[insertTabletReq.columnCategories.get(i).intValue()]); } insertStatement.setColumnCategories(columnCategories); + if (insertTabletReq.isSetDatabase()) { + if (insertTabletReq.getDatabase().isEmpty()) { + throw new IllegalArgumentException( + DataNodeQueryMessages.DATABASE_IS_NOT_SPECIFIED_FOR_INSERT); + } + if (!PathUtils.isTableModelDatabase(insertTabletReq.getDatabase())) { + throw new SemanticException( + String.format( + DataNodeMiscMessages.DATABASE_NAME_IS_TREE_MODEL_DATABASE, + insertTabletReq.getDatabase())); + } + insertStatement.setDatabaseName(insertTabletReq.getDatabase().toLowerCase(Locale.ENGLISH)); + } } PERFORMANCE_OVERVIEW_METRICS.recordParseCost(System.nanoTime() - startTime); return insertStatement; diff --git a/iotdb-protocol/thrift-datanode/src/main/thrift/client.thrift b/iotdb-protocol/thrift-datanode/src/main/thrift/client.thrift index 7ff12ceb02154..4ac00b7e270b1 100644 --- a/iotdb-protocol/thrift-datanode/src/main/thrift/client.thrift +++ b/iotdb-protocol/thrift-datanode/src/main/thrift/client.thrift @@ -274,6 +274,7 @@ struct TSInsertTabletReq { 11: optional bool isCompressed 12: optional list encodingTypes 13: optional byte compressType + 14: optional string database } struct TSInsertTabletsReq { @@ -705,4 +706,4 @@ service IClientRPCService { /** For other node's call */ common.TSStatus testConnectionEmptyRPC(); -} \ No newline at end of file +} From 7ab0db7be40a02355539a9630dd39b1ff1e1960e Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Thu, 12 Nov 2026 15:19:18 +0800 Subject: [PATCH 2/6] Fix table model redirect cache across databases --- .../InternalTableSessionRedirectIT.java | 5 +- .../InternalSubscriptionTableIT.java | 4 +- .../org/apache/iotdb/isession/ISession.java | 2 +- .../org/apache/iotdb/session/Session.java | 81 ++++++++++++------- .../iotdb/session/pool/SessionPool.java | 3 +- .../iotdb/session/SessionCacheLeaderTest.java | 73 +++++++++++++++-- 6 files changed, 130 insertions(+), 38 deletions(-) diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/InternalTableSessionRedirectIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/InternalTableSessionRedirectIT.java index 9bd60069ccb61..a9b0b162d13cd 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/InternalTableSessionRedirectIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/InternalTableSessionRedirectIT.java @@ -114,8 +114,11 @@ public void testTabletInsertRedirectsToInternalLeaderEndpoint() throws Exception session.insert(createTablet(2L)); final Session underlyingSession = InternalSessionITUtils.getUnderlyingSession(session); - final Map tableModelDeviceIdToEndpoint = + final Map> tableModelDeviceIdToEndpointByDatabase = InternalSessionITUtils.getSessionField(underlyingSession, "tableModelDeviceIdToEndpoint"); + assertEquals(1, tableModelDeviceIdToEndpointByDatabase.size()); + final Map tableModelDeviceIdToEndpoint = + tableModelDeviceIdToEndpointByDatabase.values().iterator().next(); assertEquals(1, tableModelDeviceIdToEndpoint.size()); assertTrue(leaderEndPoints.contains(tableModelDeviceIdToEndpoint.values().iterator().next())); } diff --git a/integration-test/src/test/java/org/apache/iotdb/subscription/it/consensus/local/tablemodel/InternalSubscriptionTableIT.java b/integration-test/src/test/java/org/apache/iotdb/subscription/it/consensus/local/tablemodel/InternalSubscriptionTableIT.java index 57468e4bb2567..a37b3bf8d048c 100644 --- a/integration-test/src/test/java/org/apache/iotdb/subscription/it/consensus/local/tablemodel/InternalSubscriptionTableIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/subscription/it/consensus/local/tablemodel/InternalSubscriptionTableIT.java @@ -163,10 +163,10 @@ public void testInternalPullConsumerPollsAcrossDataNodes() throws Exception { session.executeNonQueryStatement("flush"); final Session underlyingSession = InternalSessionITUtils.getUnderlyingSession(session); - final Map deviceToEndpoint = + final Map> deviceToEndpointByDatabase = InternalSessionITUtils.getSessionField( underlyingSession, "tableModelDeviceIdToEndpoint"); - leaders.addAll(deviceToEndpoint.values()); + deviceToEndpointByDatabase.values().forEach(leaders::addAll); } Assert.assertTrue( diff --git a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ISession.java b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ISession.java index 70d4266d6c889..653753272878d 100644 --- a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ISession.java +++ b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ISession.java @@ -70,7 +70,7 @@ void open( boolean enableRPCCompression, int connectionTimeoutInMs, Map deviceIdToEndpoint, - Map tabletModelDeviceIdToEndpoint, + Map> tableModelDeviceIdToEndpoint, INodeSupplier nodeSupplier) throws IoTDBConnectionException; diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java index 59bdcf2b7af26..166e74abfd268 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java @@ -173,12 +173,13 @@ public class Session implements ISession { private static final double CONVERT_THRESHOLD = 0.5; private static final double SAMPLE_PROPORTION = 0.05; private static final int MIN_RECORDS_SIZE = 40; + private static final String UNSPECIFIED_DATABASE_KEY = ""; @SuppressWarnings("squid:S3077") // Non-primitive fields should not be "volatile" protected volatile Map deviceIdToEndpoint; @SuppressWarnings("squid:S3077") // Non-primitive fields should not be "volatile" - protected volatile Map tableModelDeviceIdToEndpoint; + protected volatile Map> tableModelDeviceIdToEndpoint; @SuppressWarnings("squid:S3077") // Non-primitive fields should not be "volatile" protected volatile Map endPointToSessionConnection; @@ -663,7 +664,7 @@ public synchronized void open( boolean enableThriftRpcCompaction, int connectionTimeoutInMs, Map deviceIdToEndpoint, - Map tableModelDeviceIdToEndpoint, + Map> tableModelDeviceIdToEndpoint, INodeSupplier nodesSupplier) throws IoTDBConnectionException { if (!isClosed) { @@ -1420,12 +1421,20 @@ private SessionConnection getSessionConnection(String deviceId) throws IoTDBConn } } - private SessionConnection getSessionConnection(IDeviceID deviceId) + private Map getTableModelDatabaseRoutes(String routeDatabase) { + Map databaseToEndpoint = tableModelDeviceIdToEndpoint.get(routeDatabase); + return databaseToEndpoint == null + ? tableModelDeviceIdToEndpoint.computeIfAbsent( + routeDatabase, ignored -> new ConcurrentHashMap<>()) + : databaseToEndpoint; + } + + private SessionConnection getSessionConnection( + Map databaseToEndpoint, IDeviceID deviceId) throws IoTDBConnectionException { TEndPoint endPoint; if (enableRedirection - && tableModelDeviceIdToEndpoint != null - && (endPoint = tableModelDeviceIdToEndpoint.get(deviceId)) != null + && (endPoint = databaseToEndpoint.get(deviceId)) != null && endPointToSessionConnection.containsKey(endPoint)) { return endPointToSessionConnection.get(endPoint); } else { @@ -1465,12 +1474,13 @@ public void removeBrokenSessionConnection(SessionConnection sessionConnection) { } } if (tableModelDeviceIdToEndpoint != null && !tableModelDeviceIdToEndpoint.isEmpty()) { - for (Iterator> it = - tableModelDeviceIdToEndpoint.entrySet().iterator(); - it.hasNext(); ) { - Entry entry = it.next(); - if (entry.getValue().equals(endPoint)) { - it.remove(); + for (Map deviceIdToEndpoint : tableModelDeviceIdToEndpoint.values()) { + for (Iterator> it = deviceIdToEndpoint.entrySet().iterator(); + it.hasNext(); ) { + Entry entry = it.next(); + if (Objects.equals(entry.getValue(), endPoint)) { + it.remove(); + } } } } @@ -1503,15 +1513,15 @@ private void handleRedirection(String deviceId, TEndPoint endpoint) { } } - private void handleRedirection(IDeviceID deviceId, TEndPoint endpoint) { + private void handleRedirection( + Map databaseToEndpoint, IDeviceID deviceId, TEndPoint endpoint) { if (enableRedirection) { // no need to redirection if (UrlUtils.isWildcardAddress(endpoint.ip)) { return; } - if (!tableModelDeviceIdToEndpoint.containsKey(deviceId) - || !tableModelDeviceIdToEndpoint.get(deviceId).equals(endpoint)) { - tableModelDeviceIdToEndpoint.put(deviceId, endpoint); + if (!endpoint.equals(databaseToEndpoint.get(deviceId))) { + databaseToEndpoint.put(deviceId, endpoint); } SessionConnection connection = endPointToSessionConnection.computeIfAbsent( @@ -1524,11 +1534,16 @@ private void handleRedirection(IDeviceID deviceId, TEndPoint endpoint) { } }); if (connection == null) { - tableModelDeviceIdToEndpoint.remove(deviceId); + databaseToEndpoint.remove(deviceId, endpoint); } } } + private String resolveTableModelRouteDatabase(String database) { + String routeDatabase = database == null ? this.database : database; + return routeDatabase == null ? UNSPECIFIED_DATABASE_KEY : routeDatabase; + } + private void handleQueryRedirection(TEndPoint endPoint) throws IoTDBConnectionException { if (enableQueryRedirection) { AtomicReference exceptionReference = new AtomicReference<>(); @@ -2921,8 +2936,10 @@ public void insertRelationalTablet(Tablet tablet, String database) if (tablet.getRowSize() == 0) { return; } + // Capture the effective database before grouping so a later USE cannot change cached routes. + String routeDatabase = resolveTableModelRouteDatabase(database); if (enableRedirection) { - insertRelationalTabletWithLeaderCache(tablet, database); + insertRelationalTabletWithLeaderCache(tablet, database, routeDatabase); } else { TSInsertTabletReq request = genTSInsertTabletReq(tablet, false, false); if (request == null) { @@ -2942,17 +2959,20 @@ public void insertRelationalTablet(Tablet tablet, String database) } } - private void insertRelationalTabletWithLeaderCache(Tablet tablet, String database) + private void insertRelationalTabletWithLeaderCache( + Tablet tablet, String database, String routeDatabase) throws IoTDBConnectionException, StatementExecutionException { Map relationalTabletGroup = new HashMap<>(); - if (tableModelDeviceIdToEndpoint.isEmpty()) { + Map databaseToEndpoint = getTableModelDatabaseRoutes(routeDatabase); + if (databaseToEndpoint.isEmpty()) { relationalTabletGroup.put(getDefaultSessionConnection(), tablet); } else if (SessionUtils.isTabletContainsSingleDevice(tablet)) { - relationalTabletGroup.put(getSessionConnection(tablet.getDeviceID(0)), tablet); + relationalTabletGroup.put( + getSessionConnection(databaseToEndpoint, tablet.getDeviceID(0)), tablet); } else { for (int i = 0; i < tablet.getRowSize(); i++) { IDeviceID iDeviceID = tablet.getDeviceID(i); - final SessionConnection connection = getSessionConnection(iDeviceID); + final SessionConnection connection = getSessionConnection(databaseToEndpoint, iDeviceID); int finalI = i; relationalTabletGroup.compute( connection, @@ -2987,14 +3007,16 @@ private void insertRelationalTabletWithLeaderCache(Tablet tablet, String databas } } if (relationalTabletGroup.size() == 1) { - insertRelationalTabletOnce(relationalTabletGroup, database); + insertRelationalTabletOnce(relationalTabletGroup, database, databaseToEndpoint); } else { - insertRelationalTabletByGroup(relationalTabletGroup, database); + insertRelationalTabletByGroup(relationalTabletGroup, database, databaseToEndpoint); } } private void insertRelationalTabletOnce( - Map relationalTabletGroup, String database) + Map relationalTabletGroup, + String database, + Map databaseToEndpoint) throws IoTDBConnectionException, StatementExecutionException { Map.Entry entry = relationalTabletGroup.entrySet().iterator().next(); SessionConnection connection = entry.getKey(); @@ -3020,7 +3042,8 @@ private void insertRelationalTabletOnce( endPointMap.put(tablet.getDeviceID(i), endPointList.get(i)); } } - endPointMap.forEach(this::handleRedirection); + endPointMap.forEach( + (deviceId, endpoint) -> handleRedirection(databaseToEndpoint, deviceId, endpoint)); } catch (IoTDBConnectionException e) { if (endPointToSessionConnection != null && endPointToSessionConnection.size() > 1) { // remove the broken session @@ -3039,7 +3062,9 @@ private void insertRelationalTabletOnce( "squid:S3776" }) // ignore Cognitive Complexity of methods should not be too high private void insertRelationalTabletByGroup( - Map relationalTabletGroup, String database) + Map relationalTabletGroup, + String database, + Map databaseToEndpoint) throws IoTDBConnectionException, StatementExecutionException { List> completableFutures = relationalTabletGroup.entrySet().stream() @@ -3072,7 +3097,9 @@ private void insertRelationalTabletByGroup( endPointMap.put(subTablet.getDeviceID(i), endPointList.get(i)); } } - endPointMap.forEach(this::handleRedirection); + endPointMap.forEach( + (deviceId, endpoint) -> + handleRedirection(databaseToEndpoint, deviceId, endpoint)); } catch (StatementExecutionException e) { throw new CompletionException(e); } catch (IoTDBConnectionException e) { diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java index 9585d65c3902e..2df5866f38882 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java @@ -43,7 +43,6 @@ import org.apache.thrift.TException; import org.apache.tsfile.enums.TSDataType; -import org.apache.tsfile.file.metadata.IDeviceID; import org.apache.tsfile.file.metadata.enums.CompressionType; import org.apache.tsfile.file.metadata.enums.TSEncoding; import org.apache.tsfile.write.record.Tablet; @@ -139,7 +138,7 @@ public class SessionPool implements ISessionPool { private boolean enableQueryRedirection = false; private Map deviceIdToEndpoint; - private Map tableModelDeviceIdToEndpoint; + private Map> tableModelDeviceIdToEndpoint; private int thriftDefaultBufferSize; private int thriftMaxFrameSize; diff --git a/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionCacheLeaderTest.java b/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionCacheLeaderTest.java index 46ec3ee1f5d21..bd140c1e633cb 100644 --- a/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionCacheLeaderTest.java +++ b/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionCacheLeaderTest.java @@ -52,6 +52,7 @@ import java.util.Map; import java.util.Random; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNull; @@ -704,11 +705,70 @@ public void testInsertRelationalTablet() tablet.reset(); } - assertEquals(4, session.tableModelDeviceIdToEndpoint.size()); + assertEquals(1, session.tableModelDeviceIdToEndpoint.size()); + assertEquals(4, session.tableModelDeviceIdToEndpoint.values().iterator().next().size()); assertEquals(4, session.endPointToSessionConnection.size()); session.close(); } + @Test + public void testRelationalTabletRedirectCacheIsScopedByDatabase() + throws IoTDBConnectionException, StatementExecutionException { + final TEndPoint db1Endpoint = endpoints.get(1); + final TEndPoint db2Endpoint = endpoints.get(2); + final AtomicInteger db1AcceptedWrites = new AtomicInteger(); + final AtomicInteger db2AcceptedWrites = new AtomicInteger(); + session = + new MockSession("127.0.0.1", 55560, true, "table") { + @Override + public SessionConnection constructSessionConnection( + Session owner, TEndPoint endpoint, ZoneId zoneId) { + final TEndPoint connectionEndpoint = endpoint; + return new MockSessionConnection(owner, endpoint, zoneId) { + @Override + protected void insertTablet(TSInsertTabletReq request) throws RedirectException { + TEndPoint leader = "db1".equals(request.getDatabase()) ? db1Endpoint : db2Endpoint; + if (connectionEndpoint.equals(leader)) { + if (leader.equals(db1Endpoint)) { + db1AcceptedWrites.incrementAndGet(); + } else { + db2AcceptedWrites.incrementAndGet(); + } + return; + } + throw new RedirectException(leader); + } + }; + } + }; + session.open(); + + Tablet tablet = + new Tablet( + "table1", + Arrays.asList("id", "field"), + Arrays.asList(TSDataType.STRING, TSDataType.INT64), + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD), + 1); + tablet.addTimestamp(0, 1); + tablet.addValue("id", 0, "id0"); + tablet.addValue("field", 0, 1L); + + session.insertRelationalTablet(tablet, "db1"); + session.insertRelationalTablet(tablet, "db2"); + session.insertRelationalTablet(tablet, "db1"); + session.insertRelationalTablet(tablet, "db2"); + + assertEquals(1, db1AcceptedWrites.get()); + assertEquals(1, db2AcceptedWrites.get()); + assertEquals(2, session.tableModelDeviceIdToEndpoint.size()); + assertEquals( + db1Endpoint, session.tableModelDeviceIdToEndpoint.get("db1").get(tablet.getDeviceID(0))); + assertEquals( + db2Endpoint, session.tableModelDeviceIdToEndpoint.get("db2").get(tablet.getDeviceID(0))); + session.close(); + } + @Test public void testGroupedRelationalTabletPreservesStatementFailures() throws Exception { AtomicBoolean failWrites = new AtomicBoolean(); @@ -750,7 +810,8 @@ protected void insertTablet(TSInsertTabletReq request) tablet.addValue("field", row, 1L); } session.insertRelationalTablet(tablet); - assertEquals(4, session.tableModelDeviceIdToEndpoint.size()); + assertEquals(1, session.tableModelDeviceIdToEndpoint.size()); + assertEquals(4, session.tableModelDeviceIdToEndpoint.values().iterator().next().size()); failWrites.set(true); tablet.reset(); @@ -1217,9 +1278,11 @@ public void testInsertRelationalTabletWithSessionBroken() throws StatementExecut timestamp++; } - assertEquals(3, session.tableModelDeviceIdToEndpoint.size()); - for (Map.Entry endPointEntry : - session.tableModelDeviceIdToEndpoint.entrySet()) { + assertEquals(1, session.tableModelDeviceIdToEndpoint.size()); + Map tableModelDeviceRoutes = + session.tableModelDeviceIdToEndpoint.values().iterator().next(); + assertEquals(3, tableModelDeviceRoutes.size()); + for (Map.Entry endPointEntry : tableModelDeviceRoutes.entrySet()) { assertEquals(getDeviceIdBelongedEndpoint(endPointEntry.getKey()), endPointEntry.getValue()); } assertEquals(3, session.endPointToSessionConnection.size()); From e8a15c64233973fe617cd65ea136b2807f5d4a33 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Thu, 12 Nov 2026 15:24:10 +0800 Subject: [PATCH 3/6] Fix missing table model device ID import --- .../src/main/java/org/apache/iotdb/session/pool/SessionPool.java | 1 + 1 file changed, 1 insertion(+) diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java index 2df5866f38882..96aa834009384 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java @@ -43,6 +43,7 @@ import org.apache.thrift.TException; import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.metadata.IDeviceID; import org.apache.tsfile.file.metadata.enums.CompressionType; import org.apache.tsfile.file.metadata.enums.TSEncoding; import org.apache.tsfile.write.record.Tablet; From 2f49f14c786eb8788de71a38eb7b3bdb1b33d33f Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Thu, 12 Nov 2026 15:35:42 +0800 Subject: [PATCH 4/6] Preserve ISession flat route cache API --- .../org/apache/iotdb/isession/ISession.java | 2 +- .../org/apache/iotdb/session/Session.java | 21 +++++++++++++++++++ .../iotdb/session/pool/SessionPool.java | 4 ++-- .../iotdb/session/SessionCacheLeaderTest.java | 3 ++- 4 files changed, 26 insertions(+), 4 deletions(-) diff --git a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ISession.java b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ISession.java index 653753272878d..48dc055746563 100644 --- a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ISession.java +++ b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ISession.java @@ -70,7 +70,7 @@ void open( boolean enableRPCCompression, int connectionTimeoutInMs, Map deviceIdToEndpoint, - Map> tableModelDeviceIdToEndpoint, + Map tableModelDeviceIdToEndpoint, INodeSupplier nodeSupplier) throws IoTDBConnectionException; diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java index 166e74abfd268..237b0bbc6fed4 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java @@ -661,6 +661,27 @@ public synchronized void open( @Override public synchronized void open( + boolean enableThriftRpcCompaction, + int connectionTimeoutInMs, + Map deviceIdToEndpoint, + Map tableModelDeviceIdToEndpoint, + INodeSupplier nodesSupplier) + throws IoTDBConnectionException { + Map> databaseToDeviceIdToEndpoint = new ConcurrentHashMap<>(); + if (tableModelDeviceIdToEndpoint != null) { + databaseToDeviceIdToEndpoint.put( + resolveTableModelRouteDatabase(null), tableModelDeviceIdToEndpoint); + } + openWithTableModelDatabaseCache( + enableThriftRpcCompaction, + connectionTimeoutInMs, + deviceIdToEndpoint, + databaseToDeviceIdToEndpoint, + nodesSupplier); + } + + /** Opens this Session with a database-aware table-model route cache for internal callers. */ + public synchronized void openWithTableModelDatabaseCache( boolean enableThriftRpcCompaction, int connectionTimeoutInMs, Map deviceIdToEndpoint, diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java index 96aa834009384..9ab73fdd8f20a 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java @@ -774,7 +774,7 @@ private ISession getSession() throws IoTDBConnectionException { session = constructNewSession(); try { - session.open( + session.openWithTableModelDatabaseCache( enableThriftCompression, connectionTimeoutInMs, deviceIdToEndpoint, @@ -890,7 +890,7 @@ public void closeResultSet(SessionDataSetWrapper wrapper) { private void tryConstructNewSession() { Session session = constructNewSession(); try { - session.open( + session.openWithTableModelDatabaseCache( enableThriftCompression, connectionTimeoutInMs, deviceIdToEndpoint, diff --git a/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionCacheLeaderTest.java b/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionCacheLeaderTest.java index bd140c1e633cb..772e5528267ca 100644 --- a/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionCacheLeaderTest.java +++ b/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionCacheLeaderTest.java @@ -47,6 +47,7 @@ import java.time.ZoneId; import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -736,7 +737,7 @@ protected void insertTablet(TSInsertTabletReq request) throws RedirectException } return; } - throw new RedirectException(leader); + throw new RedirectException(Collections.singletonList(leader)); } }; } From b092a12a21bfe926805dbbb91ffdf491135b9b9d Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Thu, 12 Nov 2026 15:42:41 +0800 Subject: [PATCH 5/6] Expose database-aware session route cache open API --- .../src/main/java/org/apache/iotdb/isession/ISession.java | 8 ++++++++ .../src/main/java/org/apache/iotdb/session/Session.java | 1 + 2 files changed, 9 insertions(+) diff --git a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ISession.java b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ISession.java index 48dc055746563..299dff20f4542 100644 --- a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ISession.java +++ b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ISession.java @@ -74,6 +74,14 @@ void open( INodeSupplier nodeSupplier) throws IoTDBConnectionException; + void openWithTableModelDatabaseCache( + boolean enableRPCCompression, + int connectionTimeoutInMs, + Map deviceIdToEndpoint, + Map> tableModelDeviceIdToEndpoint, + INodeSupplier nodeSupplier) + throws IoTDBConnectionException; + void close() throws IoTDBConnectionException; String getTimeZone() throws IoTDBConnectionException; diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java index 237b0bbc6fed4..347de6fa82ce0 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java @@ -681,6 +681,7 @@ public synchronized void open( } /** Opens this Session with a database-aware table-model route cache for internal callers. */ + @Override public synchronized void openWithTableModelDatabaseCache( boolean enableThriftRpcCompaction, int connectionTimeoutInMs, From f53ab9aa42feb3b1876c6b29fd71f775b5d99921 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Thu, 12 Nov 2026 16:06:46 +0800 Subject: [PATCH 6/6] Fix nested route cache test compilation --- .../local/tablemodel/InternalSubscriptionTableIT.java | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/integration-test/src/test/java/org/apache/iotdb/subscription/it/consensus/local/tablemodel/InternalSubscriptionTableIT.java b/integration-test/src/test/java/org/apache/iotdb/subscription/it/consensus/local/tablemodel/InternalSubscriptionTableIT.java index a37b3bf8d048c..260bbcac187c4 100644 --- a/integration-test/src/test/java/org/apache/iotdb/subscription/it/consensus/local/tablemodel/InternalSubscriptionTableIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/subscription/it/consensus/local/tablemodel/InternalSubscriptionTableIT.java @@ -166,7 +166,9 @@ public void testInternalPullConsumerPollsAcrossDataNodes() throws Exception { final Map> deviceToEndpointByDatabase = InternalSessionITUtils.getSessionField( underlyingSession, "tableModelDeviceIdToEndpoint"); - deviceToEndpointByDatabase.values().forEach(leaders::addAll); + deviceToEndpointByDatabase + .values() + .forEach(deviceToEndpoint -> leaders.addAll(deviceToEndpoint.values())); } Assert.assertTrue(