refactor user defined http-client usage

apache · Mar 14, 2024 · 5cacaed · 5cacaed
1 parent 3484629
commit 5cacaed
Show file tree

Hide file tree

Showing 5 changed files with 36 additions and 22 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
@@ -46,6 +46,7 @@ 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) {

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
@@ -30,6 +30,7 @@
 public class HugeClient implements Closeable {
 
     private static final Logger LOG = LoggerFactory.getLogger(RestClient.class);
+    private static final int SECOND = 1000;
 
     static {
         ClientVersion.check();
@@ -53,28 +54,26 @@ public class HugeClient implements Closeable {
 
     public HugeClient(HugeClientBuilder builder) {
         this.borrowedClient = false;
+        RestClientConfig config;
         try {
-            RestClientConfig 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(okHttpClientBuilder -> {
-                                                          okHttpClientBuilder.followRedirects(
-                                                                  builder.followRedirects());
-                                                      })
-                                                      .build();
+            config = RestClientConfig.builder()
+                                     .user(builder.username())
+                                     .password(builder.password())
+                                     .timeout(builder.timeout() * SECOND)
+                                     .connectTimeout(builder.connectTimeout() * SECOND)
+                                     .readTimeout(builder.readTimeout() * SECOND)
+                                     .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,8 +17,12 @@
 
 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();
@@ -32,15 +36,15 @@ public class HugeClientBuilder {
     private String username;
     private String password;
     private int timeout;
-    private Integer connectTimeout;
-    private Integer readTimeout;
     private int maxConns;
     private int maxConnsPerRoute;
     private int idleTime;
     private String trustStoreFile;
     private String trustStorePassword;
-
-    private Boolean followRedirects;
+    private Consumer<OkHttpClient.Builder> httpBuilderConsumer;
+    /** Set them null by default to keep compatibility with 'timeout' */
+    private final Integer connectTimeout;
+    private final Integer readTimeout;
 
     public HugeClientBuilder(String url, String graph) {
         E.checkArgument(url != null && !url.isEmpty(),
@@ -53,11 +57,15 @@ public HugeClientBuilder(String url, String graph) {
         this.username = "";
         this.password = "";
         this.timeout = DEFAULT_TIMEOUT;
+
         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() {
@@ -118,7 +126,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;
     }
 
@@ -170,7 +182,7 @@ public String trustStorePassword() {
         return this.trustStorePassword;
     }
 
-    public Boolean followRedirects() {
-        return this.followRedirects;
+    public Consumer<OkHttpClient.Builder> httpBuilderConsumer() {
+        return httpBuilderConsumer;
     }
 }
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
@@ -67,6 +67,7 @@ public static HugeClient tryConnect(GraphConnection connection) {
                                .configUser(username, password)
                                .configTimeout(timeout)
                                .configSSL(trustStoreFile, trustStorePassword)
+                               .configHttpBuilder(http -> http.followRedirects(false))
                                .build();
         } catch (IllegalStateException e) {
             String message = e.getMessage();