refact(core): support auth config for computer task (#265)

* fix: hugegraph client authentication configuration * chore: improve description --------- Co-authored-by: diaohancai <[email protected]>
apache · Oct 8, 2023 · a21a191 · a21a191
1 parent 15fc9bb
commit a21a191
Show file tree

Hide file tree

Showing 6 changed files with 35 additions and 9 deletions.
diff --git a/computer-api/src/main/java/org/apache/hugegraph/computer/core/config/ComputerOptions.java b/computer-api/src/main/java/org/apache/hugegraph/computer/core/config/ComputerOptions.java
@@ -645,6 +645,22 @@ public static synchronized ComputerOptions instance() {
                     "hugegraph"
             );
 
+    public static final ConfigOption<String> HUGEGRAPH_USERNAME =
+            new ConfigOption<>(
+                    "hugegraph.username",
+                    "The username of graph for authentication.",
+                    null,
+                    ""
+            );
+
+    public static final ConfigOption<String> HUGEGRAPH_PASSWORD =
+            new ConfigOption<>(
+                    "hugegraph.password",
+                    "The password of graph for authentication.",
+                    null,
+                    ""
+            );
+
     public static final ConfigOption<String> TRANSPORT_SERVER_HOST =
             new ConfigOption<>(
                     "transport.server_host",

diff --git a/...uter-api/src/main/java/org/apache/hugegraph/computer/core/output/hg/task/TaskManager.java b/...uter-api/src/main/java/org/apache/hugegraph/computer/core/output/hg/task/TaskManager.java
@@ -56,7 +56,9 @@ public TaskManager(Config config) {
         this.config = config;
         String url = config.get(ComputerOptions.HUGEGRAPH_URL);
         String graph = config.get(ComputerOptions.HUGEGRAPH_GRAPH_NAME);
-        this.client = new HugeClientBuilder(url, graph).build();
+        String username = config.get(ComputerOptions.HUGEGRAPH_USERNAME);
+        String password = config.get(ComputerOptions.HUGEGRAPH_PASSWORD);
+        this.client = new HugeClientBuilder(url, graph).configUser(username, password).build();
         // Try to make all batch threads running and don't wait for producer
         this.batchSemaphore = new Semaphore(this.batchSemaphoreNum());
         /*

diff --git a/...uter-core/src/main/java/org/apache/hugegraph/computer/core/input/hg/HugeGraphFetcher.java b/...uter-core/src/main/java/org/apache/hugegraph/computer/core/input/hg/HugeGraphFetcher.java
@@ -37,7 +37,9 @@ public class HugeGraphFetcher implements GraphFetcher {
     public HugeGraphFetcher(Config config, InputSplitRpcService rpcService) {
         String url = config.get(ComputerOptions.HUGEGRAPH_URL);
         String graph = config.get(ComputerOptions.HUGEGRAPH_GRAPH_NAME);
-        this.client = new HugeClientBuilder(url, graph).build();
+        String username = config.get(ComputerOptions.HUGEGRAPH_USERNAME);
+        String password = config.get(ComputerOptions.HUGEGRAPH_PASSWORD);
+        this.client = new HugeClientBuilder(url, graph).configUser(username, password).build();
         this.vertexFetcher = new HugeVertexFetcher(config, this.client);
         this.edgeFetcher = new HugeEdgeFetcher(config, this.client);
         this.rpcService = rpcService;

diff --git a/...core/src/main/java/org/apache/hugegraph/computer/core/input/hg/HugeInputSplitFetcher.java b/...core/src/main/java/org/apache/hugegraph/computer/core/input/hg/HugeInputSplitFetcher.java
@@ -38,8 +38,11 @@ public HugeInputSplitFetcher(Config config) {
         this.config = config;
         String url = config.get(ComputerOptions.HUGEGRAPH_URL);
         String graph = config.get(ComputerOptions.HUGEGRAPH_GRAPH_NAME);
+        String username = config.get(ComputerOptions.HUGEGRAPH_USERNAME);
+        String password = config.get(ComputerOptions.HUGEGRAPH_PASSWORD);
         int timeout = config.get(ComputerOptions.INPUT_SPLIT_FETCH_TIMEOUT);
-        this.client = new HugeClientBuilder(url, graph).configTimeout(timeout)
+        this.client = new HugeClientBuilder(url, graph).configUser(username, password)
+                                                       .configTimeout(timeout)
                                                        .build();
     }
 

diff --git a/computer-dist/src/assembly/static/conf/computer.properties b/computer-dist/src/assembly/static/conf/computer.properties
@@ -27,6 +27,8 @@ bsp.etcd_endpoints=http://127.0.0.1:2379
 
 hugegraph.url=http://127.0.0.1:8080
 hugegraph.name=hugegraph
+hugegraph.username=
+hugegraph.password=
 
 algorithm.params_class=org.apache.hugegraph.computer.algorithm.centrality.pagerank.PageRankParams
 

diff --git a/computer-test/src/main/java/org/apache/hugegraph/computer/suite/unit/UnitTestBase.java b/computer-test/src/main/java/org/apache/hugegraph/computer/suite/unit/UnitTestBase.java
@@ -69,6 +69,8 @@ public class UnitTestBase {
                                         "abcdefghijklmnopqrstuvxyz";
     private static String URL;
     private static String GRAPH;
+    private static String USERNAME;
+    private static String PASSWORD;
     private static HugeClient CLIENT = null;
 
     protected static void clearAll() {
@@ -121,11 +123,10 @@ public static void init() throws ClassNotFoundException {
                 "defaultValue",
                 "src/main/resources/hdfs_input_test/struct.json");
 
-        URL = ComputerOptions.HUGEGRAPH_URL
-                .defaultValue();
-
-        GRAPH = ComputerOptions.HUGEGRAPH_GRAPH_NAME
-                .defaultValue();
+        URL = ComputerOptions.HUGEGRAPH_URL.defaultValue();
+        GRAPH = ComputerOptions.HUGEGRAPH_GRAPH_NAME.defaultValue();
+        USERNAME = ComputerOptions.HUGEGRAPH_USERNAME.defaultValue();
+        PASSWORD = ComputerOptions.HUGEGRAPH_PASSWORD.defaultValue();
 
         Class.forName(IdType.class.getName());
         // Don't forget to register options
@@ -286,7 +287,7 @@ protected static StreamGraphOutput newStreamGraphOutput(
 
     protected static synchronized HugeClient client() {
         if (CLIENT == null) {
-            CLIENT = HugeClient.builder(URL, GRAPH).build();
+            CLIENT = HugeClient.builder(URL, GRAPH).configUser(USERNAME, PASSWORD).build();
         }
         return CLIENT;
     }