Support internal sessions through inherited client RPC services

iotdb-client/isession/src/main/java/org/apache/iotdb/isession/IInternalSessionBuilder.java

@@ -0,0 +1,32 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.isession;
+
+import org.apache.iotdb.rpc.IoTDBConnectionException;
+
+/** Capability implemented by builders that can construct sessions on internal endpoints. */
+public interface IInternalSessionBuilder<T> {
+
+  /**
+   * Builds a session using the caller-provided internal node supplier. Whether the returned session
+   * is already open follows the builder's ordinary build contract.
+   */
+  T buildInternal(INodeSupplier nodeSupplier) throws IoTDBConnectionException;
+}

iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java

@@ -20,6 +20,11 @@
 package org.apache.iotdb.session.i18n;
 
 public final class SessionMessages {
+  public static final String EXCEPTION_NO_DATANODE_INTERNAL_ENDPOINT_IS_AVAILABLE_249D1081 =
+      "No DataNode internal endpoint is available";
+
+  public static final String EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392 =
+      "Internal node supplier must not be null";
 
   // Session
   public static final String NODE_URLS_EMPTY = "nodeUrls shouldn't be empty.";

iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java

@@ -20,6 +20,11 @@
 package org.apache.iotdb.session.i18n;
 
 public final class SessionMessages {
+  public static final String EXCEPTION_NO_DATANODE_INTERNAL_ENDPOINT_IS_AVAILABLE_249D1081 =
+      "没有可用的 DataNode 内部端点";
+
+  public static final String EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392 =
+      "内部节点 supplier 不能为空";
 
   // Session
   public static final String NODE_URLS_EMPTY = "nodeUrls 不能为空。";

iotdb-client/session/src/main/java/org/apache/iotdb/session/InternalSession.java

@@ -0,0 +1,77 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.session;
+
+import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+import org.apache.iotdb.isession.INodeSupplier;
+import org.apache.iotdb.rpc.IoTDBConnectionException;
+import org.apache.iotdb.session.i18n.SessionMessages;
+
+import java.time.ZoneId;
+import java.util.List;
+import java.util.Objects;
+
+/** Session using caller-supplied internal endpoints and the caller's username. */
+public class InternalSession extends Session {
+
+  private final INodeSupplier nodeSupplier;
+
+  public InternalSession(AbstractSessionBuilder builder, INodeSupplier nodeSupplier) {
+    super(builder);
+    this.nodeSupplier =
+        Objects.requireNonNull(
+            nodeSupplier,
+            SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392);
+    password = null;
+    useEncryptedPassword = false;
+    enableAutoFetch = false;
+    enableRedirection = false;
+    enableQueryRedirection = false;
+  }
+
+  @Override
+  public synchronized void open() throws IoTDBConnectionException {
+    open(enableThriftRpcCompaction, connectionTimeoutInMs);
+  }
+
+  @Override
+  protected INodeSupplier createNodeSupplier(boolean enableRPCCompaction, int timeoutInMs) {
+    return nodeSupplier;
+  }
+
+  @Override
+  public SessionConnection constructSessionConnection(
+      Session session, TEndPoint endpoint, ZoneId zoneId) throws IoTDBConnectionException {
+    final List<TEndPoint> candidates = availableNodes.get();
+    if (candidates.isEmpty()) {
+      throw new IoTDBConnectionException(
+          SessionMessages.EXCEPTION_NO_DATANODE_INTERNAL_ENDPOINT_IS_AVAILABLE_249D1081);
+    }
+    defaultEndPoint =
+        candidates.contains(endpoint)
+            ? endpoint.deepCopy()
+            : nodeSupplier
+                .getQueryEndPoint()
+                .filter(candidates::contains)
+                .orElse(candidates.get(0))
+                .deepCopy();
+    return super.constructSessionConnection(session, defaultEndPoint, zoneId);
+  }
+}

iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java

@@ -21,6 +21,7 @@
 
 import org.apache.iotdb.common.rpc.thrift.TAggregationType;
 import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+import org.apache.iotdb.isession.IInternalSessionBuilder;
 import org.apache.iotdb.isession.INodeSupplier;
 import org.apache.iotdb.isession.ISession;
 import org.apache.iotdb.isession.SessionConfig;
@@ -549,33 +550,7 @@
       this.availableNodes = null;
     }
 
-    if (enableAutoFetch) {
-      initThreadPool();
-      this.availableNodes =
-          NodesSupplier.createNodeSupplier(
-              getNodeUrls(),
-              SessionUtils.parseSeedNodeUrls(initialNodeUrls),
-              executorService,
-              username,
-              password,
-              zoneId,
-              thriftDefaultBufferSize,
-              thriftMaxFrameSize,
-              connectionTimeoutInMs,
-              useSSL,
-              trustStore,
-              trustStorePwd,
-              keyStore,
-              keyStorePwd,
-              sslProtocol,
-              enableRPCCompaction,
-              version.toString());
-    } else {
-      // A static supplier has no separate seed tier. Keep every configured address after reopen.
-      Set<TEndPoint> staticNodes = new LinkedHashSet<>(getNodeUrls());
-      staticNodes.addAll(SessionUtils.parseSeedNodeUrls(initialNodeUrls));
-      this.availableNodes = new DummyNodesSupplier(new ArrayList<>(staticNodes));
-    }
+    this.availableNodes = createNodeSupplier(enableRPCCompaction, connectionTimeoutInMs);
 
     this.connectionTimeoutInMs = connectionTimeoutInMs;
     setDefaultSessionConnection(constructSessionConnection(this, defaultEndPoint, zoneId));
@@ -589,6 +564,37 @@
     }
   }
 
+  /** Creates the endpoint source used by initial connections and reconnects. */
+  protected INodeSupplier createNodeSupplier(
+      boolean enableRPCCompaction, int connectionTimeoutInMs) {
+    if (enableAutoFetch) {
+      initThreadPool();
+      return NodesSupplier.createNodeSupplier(
+          getNodeUrls(),
+          SessionUtils.parseSeedNodeUrls(initialNodeUrls),
+          executorService,
+          username,
+          password,
+          zoneId,
+          thriftDefaultBufferSize,
+          thriftMaxFrameSize,
+          connectionTimeoutInMs,
+          useSSL,
+          trustStore,
+          trustStorePwd,
+          keyStore,
+          keyStorePwd,
+          sslProtocol,
+          enableRPCCompaction,
+          version.toString());
+    } else {
+      // A static supplier has no separate seed tier. Keep every configured address after reopen.
+      Set<TEndPoint> staticNodes = new LinkedHashSet<>(getNodeUrls());
+      staticNodes.addAll(SessionUtils.parseSeedNodeUrls(initialNodeUrls));
+      return new DummyNodesSupplier(new ArrayList<>(staticNodes));
+    }
+  }
+
   private void initThreadPool() {
     this.executorService =
         Executors.newSingleThreadScheduledExecutor(
@@ -699,7 +705,7 @@
       isClosed = true;
     }
   }
 
   public SessionConnection constructSessionConnection(
       Session session, TEndPoint endpoint, ZoneId zoneId) throws IoTDBConnectionException {
     if (endpoint == null) {
@@ -4445,7 +4451,8 @@
     this.enableThriftRpcCompaction = enableThriftRpcCompaction;
   }
 
-  public static class Builder extends AbstractSessionBuilder {
+  public static class Builder extends AbstractSessionBuilder
+      implements IInternalSessionBuilder<InternalSession> {
 
     public Builder host(String host) {
       this.host = host;
@@ -4595,5 +4602,11 @@
       }
       return new Session(this);
     }
+
+    /** Builds an unopened internal session using caller-supplied internal endpoints. */
+    @Override
+    public InternalSession buildInternal(INodeSupplier nodeSupplier) {
+      return new InternalSession(this, nodeSupplier);
+    }
   }
 }

iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java

@@ -19,6 +19,8 @@
 
 package org.apache.iotdb.session;
 
+import org.apache.iotdb.isession.IInternalSessionBuilder;
+import org.apache.iotdb.isession.INodeSupplier;
 import org.apache.iotdb.isession.ITableSession;
 import org.apache.iotdb.isession.SessionConfig;
 import org.apache.iotdb.rpc.IoTDBConnectionException;
@@ -41,7 +43,8 @@
  *
  * <p>All configurations have reasonable default values, which can be overridden as needed.
  */
-public class TableSessionBuilder extends AbstractSessionBuilder {
+public class TableSessionBuilder extends AbstractSessionBuilder
+    implements IInternalSessionBuilder<ITableSession> {
 
   private boolean enableCompression = false;
   private int connectionTimeoutInMs = SessionConfig.DEFAULT_CONNECTION_TIMEOUT_MS;
@@ -437,7 +440,17 @@ public ITableSession build() throws IoTDBConnectionException {
           Collections.singletonList(SessionConfig.DEFAULT_HOST + ":" + SessionConfig.DEFAULT_PORT);
     }
     this.sqlDialect = TABLE;
-    Session newSession = new Session(this);
+    return buildSession(new Session(this));
+  }
+
+  /** Builds and opens a table session using caller-supplied internal endpoints. */
+  @Override
+  public ITableSession buildInternal(INodeSupplier nodeSupplier) throws IoTDBConnectionException {
+    this.sqlDialect = TABLE;
+    return buildSession(new InternalSession(this, nodeSupplier));
+  }
+
+  private ITableSession buildSession(Session newSession) throws IoTDBConnectionException {
     newSession.enableIoTDBRpcCompression = isIoTDBRpcCompressionEnabled;
     newSession.tabletCompressionMinRowSize = tabletCompressionMinRowSize;
 

iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/InternalTableSessionPool.java

@@ -0,0 +1,38 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.session.pool;
+
+import org.apache.iotdb.isession.INodeSupplier;
+import org.apache.iotdb.session.i18n.SessionMessages;
+
+import java.util.Objects;
+
+/** Table session pool using caller-supplied internal endpoints and the caller's username. */
+public class InternalTableSessionPool extends TableSessionPool {
+
+  InternalTableSessionPool(AbstractSessionPoolBuilder builder, INodeSupplier nodeSupplier) {
+    super(
+        new SessionPool(
+            builder,
+            Objects.requireNonNull(
+                nodeSupplier,
+                SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392)));
+  }
+}