add RestClientConfig params: builderCallback, connectTimeout and read…

…Timeout
apache · Mar 12, 2024 · ae67f8b · ae67f8b
1 parent 44f99e9
commit ae67f8b
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 1 deletion.
diff --git a/hugegraph-common/src/main/java/org/apache/hugegraph/rest/AbstractRestClient.java b/hugegraph-common/src/main/java/org/apache/hugegraph/rest/AbstractRestClient.java
@@ -183,6 +183,12 @@ private OkHttpClient buildOkHttpClient(RestClientConfig config) {
             builder.connectTimeout(config.getTimeout(), TimeUnit.MILLISECONDS)
                    .readTimeout(config.getTimeout(), TimeUnit.MILLISECONDS);
         }
+        if(config.getConnectTimeout() != null) {
+            builder.connectTimeout(config.getConnectTimeout(), TimeUnit.MILLISECONDS);
+        }
+        if(config.getReadTimeout() != null) {
+            builder.readTimeout(config.getReadTimeout(), TimeUnit.MILLISECONDS);
+        }
 
         if (config.getMaxIdleConns() != null || config.getIdleTime() != null) {
             ConnectionPool connectionPool = new ConnectionPool(config.getMaxIdleConns(),
@@ -205,6 +211,11 @@ private OkHttpClient buildOkHttpClient(RestClientConfig config) {
         configSsl(builder, this.baseUrl, config.getTrustStoreFile(),
                   config.getTrustStorePassword());
 
+        //execute builder callback before builder.build()
+        if (config.getBuilderCallback() != null) {
+            config.getBuilderCallback().accept(builder);
+        }
+
         OkHttpClient okHttpClient = builder.build();
 
         if (config.getMaxConns() != null) {

diff --git a/hugegraph-common/src/main/java/org/apache/hugegraph/rest/RestClientConfig.java b/hugegraph-common/src/main/java/org/apache/hugegraph/rest/RestClientConfig.java
@@ -17,9 +17,12 @@
 
 package org.apache.hugegraph.rest;
 
+import java.util.function.Consumer;
+
 import lombok.Builder;
 import lombok.Getter;
 import lombok.Setter;
+import okhttp3.OkHttpClient;
 
 @Builder
 @Getter
@@ -30,13 +33,21 @@ public class RestClientConfig {
     private String user;
     private String password;
     private String token;
-    // unit in milliseconds
+    /**
+     * @deprecated  use connectTimeout and readTimeout instead
+     */
+    @Deprecated
     private Integer timeout;
+    // unit in milliseconds
+    private Integer connectTimeout;
+    // unit in milliseconds
+    private Integer readTimeout;
     private Integer maxConns;
     private Integer maxConnsPerRoute;
     // unit in seconds
     private Integer idleTime = 30;
     private Integer maxIdleConns = 5;
     private String trustStoreFile;
     private String trustStorePassword;
+    private Consumer<OkHttpClient.Builder> builderCallback;
 }