feat(client): support user defined OKHTTPClient configs (#590)

Follow apache/incubator-hugegraph-commons#140, add a callback builder for user defined configs in (OK)HTTPClient After: ```java HugeClient.builder(url, graph) .configUser(username, password) .configTimeout(timeout) .configHttpBuilder(x -> x.changeHTTPClientConfigs(....)) .build(); ``` --------- Co-authored-by: imbajin <[email protected]>
apache · Mar 15, 2024 · 7aa9cf1 · 7aa9cf1
1 parent d70e83f
commit 7aa9cf1
Show file tree

Hide file tree

Showing 7 changed files with 73 additions and 12 deletions.
diff --git a/hugegraph-client/src/main/java/org/apache/hugegraph/client/RestClient.java b/hugegraph-client/src/main/java/org/apache/hugegraph/client/RestClient.java
@@ -20,6 +20,7 @@
 import org.apache.hugegraph.exception.ServerException;
 import org.apache.hugegraph.rest.AbstractRestClient;
 import org.apache.hugegraph.rest.ClientException;
+import org.apache.hugegraph.rest.RestClientConfig;
 import org.apache.hugegraph.rest.RestResult;
 import org.apache.hugegraph.serializer.PathDeserializer;
 import org.apache.hugegraph.structure.graph.Path;
@@ -45,13 +46,18 @@ public RestClient(String url, String username, String password, int timeout) {
         super(url, username, password, timeout * SECOND);
     }
 
+    @Deprecated
     public RestClient(String url, String username, String password, int timeout,
                       int maxConns, int maxConnsPerRoute,
                       String trustStoreFile, String trustStorePassword) {
         super(url, username, password, timeout * SECOND, maxConns,
               maxConnsPerRoute, trustStoreFile, trustStorePassword);
     }
 
+    public RestClient(String url, RestClientConfig config) {
+        super(url, config);
+    }
+
     public void apiVersion(Version version) {
         E.checkNotNull(version, "api version");
         this.apiVersion = version;

diff --git a/hugegraph-client/src/main/java/org/apache/hugegraph/driver/HugeClient.java b/hugegraph-client/src/main/java/org/apache/hugegraph/driver/HugeClient.java
@@ -21,6 +21,7 @@
 
 import org.apache.hugegraph.client.RestClient;
 import org.apache.hugegraph.rest.ClientException;
+import org.apache.hugegraph.rest.RestClientConfig;
 import org.apache.hugegraph.util.VersionUtil;
 import org.apache.hugegraph.version.ClientVersion;
 import org.slf4j.Logger;
@@ -52,19 +53,26 @@ public class HugeClient implements Closeable {
 
     public HugeClient(HugeClientBuilder builder) {
         this.borrowedClient = false;
+        RestClientConfig config;
         try {
-            this.client = new RestClient(builder.url(),
-                                         builder.username(),
-                                         builder.password(),
-                                         builder.timeout(),
-                                         builder.maxConns(),
-                                         builder.maxConnsPerRoute(),
-                                         builder.trustStoreFile(),
-                                         builder.trustStorePassword());
+            config = RestClientConfig.builder()
+                                     .user(builder.username())
+                                     .password(builder.password())
+                                     .timeout(builder.timeout())
+                                     .connectTimeout(builder.connectTimeout())
+                                     .readTimeout(builder.readTimeout())
+                                     .maxConns(builder.maxConns())
+                                     .maxConnsPerRoute(builder.maxConnsPerRoute())
+                                     .trustStoreFile(builder.trustStoreFile())
+                                     .trustStorePassword(builder.trustStorePassword())
+                                     .builderCallback(builder.httpBuilderConsumer())
+                                     .build();
+            this.client = new RestClient(builder.url(), config);
         } catch (Exception e) {
             LOG.warn("Failed to create RestClient instance", e);
             throw new ClientException("Failed to connect url '%s'", builder.url());
         }
+
         try {
             this.initManagers(this.client, builder.graph());
         } catch (Throwable e) {

diff --git a/hugegraph-client/src/main/java/org/apache/hugegraph/driver/HugeClientBuilder.java b/hugegraph-client/src/main/java/org/apache/hugegraph/driver/HugeClientBuilder.java
@@ -17,15 +17,20 @@
 
 package org.apache.hugegraph.driver;
 
+import java.util.function.Consumer;
+
 import org.apache.hugegraph.util.E;
 
+import okhttp3.OkHttpClient;
+
 public class HugeClientBuilder {
 
     private static final int CPUS = Runtime.getRuntime().availableProcessors();
     private static final int DEFAULT_TIMEOUT = 20;
     private static final int DEFAULT_MAX_CONNS = 4 * CPUS;
     private static final int DEFAULT_MAX_CONNS_PER_ROUTE = 2 * CPUS;
     private static final int DEFAULT_IDLE_TIME = 30;
+    private static final int SECOND = 1000;
 
     private String url;
     private String graph;
@@ -37,6 +42,10 @@ public class HugeClientBuilder {
     private int idleTime;
     private String trustStoreFile;
     private String trustStorePassword;
+    private Consumer<OkHttpClient.Builder> httpBuilderConsumer;
+    /** Set them null by default to keep compatibility with 'timeout' */
+    private Integer connectTimeout;
+    private Integer readTimeout;
 
     public HugeClientBuilder(String url, String graph) {
         E.checkArgument(url != null && !url.isEmpty(),
@@ -48,12 +57,16 @@ public HugeClientBuilder(String url, String graph) {
         this.graph = graph;
         this.username = "";
         this.password = "";
-        this.timeout = DEFAULT_TIMEOUT;
+        this.timeout = DEFAULT_TIMEOUT * SECOND;
+
         this.maxConns = DEFAULT_MAX_CONNS;
         this.maxConnsPerRoute = DEFAULT_MAX_CONNS_PER_ROUTE;
         this.trustStoreFile = "";
         this.trustStorePassword = "";
         this.idleTime = DEFAULT_IDLE_TIME;
+        this.httpBuilderConsumer = null;
+        this.connectTimeout = null;
+        this.readTimeout = null;
     }
 
     public HugeClient build() {
@@ -96,7 +109,21 @@ public HugeClientBuilder configTimeout(int timeout) {
         if (timeout == 0) {
             timeout = DEFAULT_TIMEOUT;
         }
-        this.timeout = timeout;
+        this.timeout = timeout * SECOND;
+        return this;
+    }
+
+    public HugeClientBuilder configConnectTimeout(Integer connectTimeout) {
+        if (connectTimeout != null) {
+            this.connectTimeout = connectTimeout * SECOND;
+        }
+        return this;
+    }
+
+    public HugeClientBuilder configReadTimeout(Integer readTimeout) {
+        if (readTimeout != null) {
+            this.readTimeout = readTimeout * SECOND;
+        }
         return this;
     }
 
@@ -114,7 +141,11 @@ public HugeClientBuilder configUser(String username, String password) {
         }
         this.username = username;
         this.password = password;
+        return this;
+    }
 
+    public HugeClientBuilder configHttpBuilder(Consumer<OkHttpClient.Builder> builderConsumer) {
+        this.httpBuilderConsumer = builderConsumer;
         return this;
     }
 
@@ -138,6 +169,14 @@ public int timeout() {
         return this.timeout;
     }
 
+    public Integer connectTimeout() {
+        return this.connectTimeout;
+    }
+
+    public Integer readTimeout() {
+        return this.readTimeout;
+    }
+
     public int maxConns() {
         return maxConns;
     }
@@ -157,4 +196,8 @@ public String trustStoreFile() {
     public String trustStorePassword() {
         return this.trustStorePassword;
     }
+
+    public Consumer<OkHttpClient.Builder> httpBuilderConsumer() {
+        return httpBuilderConsumer;
+    }
 }
diff --git a/hugegraph-client/src/test/java/org/apache/hugegraph/functional/HugeClientHttpsTest.java b/hugegraph-client/src/test/java/org/apache/hugegraph/functional/HugeClientHttpsTest.java
@@ -71,7 +71,8 @@ public void testHttpsClientBuilderWithConnection() {
     @Test
     public void testHttpsClientWithConnectionPoolNoUserParam() {
         client = HugeClient.builder(BASE_URL, GRAPH)
-                           .configTimeout(TIMEOUT)
+                           .configConnectTimeout(3)
+                           .configReadTimeout(10)
                            .configPool(MAX_CONNS, MAX_CONNS_PER_ROUTE)
                            .configSSL(TRUST_STORE_PATH, TRUST_STORE_PASSWORD)
                            .build();

diff --git a/hugegraph-client/src/test/java/org/apache/hugegraph/functional/HugeClientTest.java b/hugegraph-client/src/test/java/org/apache/hugegraph/functional/HugeClientTest.java
@@ -32,6 +32,7 @@ public class HugeClientTest {
     public void testContext() {
         HugeClient client = HugeClient.builder(BASE_URL, GRAPH)
                                       .configUser(USERNAME, PASSWORD)
+                                      .configHttpBuilder(builder -> builder.followRedirects(false))
                                       .build();
 
         String token = "Bearer token";

diff --git a/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/util/HugeClientUtil.java b/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/util/HugeClientUtil.java
@@ -65,8 +65,10 @@ public static HugeClient tryConnect(GraphConnection connection) {
         try {
             client = HugeClient.builder(url, graph)
                                .configUser(username, password)
+                               // TODO: change it to connTimeout & readTimeout
                                .configTimeout(timeout)
                                .configSSL(trustStoreFile, trustStorePassword)
+                               .configHttpBuilder(http -> http.followRedirects(false))
                                .build();
         } catch (IllegalStateException e) {
             String message = e.getMessage();

diff --git a/pom.xml b/pom.xml
@@ -99,7 +99,7 @@
 
     <properties>
         <revision>1.2.0</revision>
-        <hugegraph.common.version>1.2.0</hugegraph.common.version>
+        <hugegraph.common.version>1.3.0</hugegraph.common.version>
         <release.name>${project.artifactId}</release.name>
         <final.name>apache-${release.name}-incubating-${project.version}</final.name>
         <assembly.dir>${project.basedir}/assembly</assembly.dir>