diff --git a/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/rest/StoreAPI.java b/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/rest/StoreAPI.java index b3c30ca9e1..4fcf3660f5 100644 --- a/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/rest/StoreAPI.java +++ b/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/rest/StoreAPI.java @@ -150,9 +150,14 @@ public Map> shardLeaders() throws PDException { return leaders; } - @GetMapping(value = "/balanceLeaders") - public Map balanceLeaders() throws PDException { - return pdRestService.balancePartitionLeader(); + @GetMapping(value = "/balanceLeaders", produces = MediaType.APPLICATION_JSON_VALUE) + @ResponseBody + public String balanceLeaders() { + try { + return toJSON(pdRestService.balancePartitionLeader()); + } catch (PDException e) { + return toJSON(e); + } } @DeleteMapping(value = "/store/{storeId}") diff --git a/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/rest/TaskAPI.java b/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/rest/TaskAPI.java index a1876141b9..60aad2d554 100644 --- a/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/rest/TaskAPI.java +++ b/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/rest/TaskAPI.java @@ -89,9 +89,16 @@ public String splitPartitions() { } } - @GetMapping(value = "/balanceLeaders") - public Map balanceLeaders() throws PDException { - return pdRestService.balancePartitionLeader(); + @GetMapping(value = "/balanceLeaders", produces = MediaType.APPLICATION_JSON_VALUE) + @ResponseBody + public String balanceLeaders() { + try { + Map leaders = pdRestService.balancePartitionLeader(); + return toJSON(leaders); + } catch (PDException e) { + e.printStackTrace(); + return toJSON(e); + } } @GetMapping(value = "/compact") diff --git a/hugegraph-pd/hg-pd-service/src/test/java/org/apache/hugegraph/pd/rest/BalanceLeadersAPITest.java b/hugegraph-pd/hg-pd-service/src/test/java/org/apache/hugegraph/pd/rest/BalanceLeadersAPITest.java new file mode 100644 index 0000000000..e8410de215 --- /dev/null +++ b/hugegraph-pd/hg-pd-service/src/test/java/org/apache/hugegraph/pd/rest/BalanceLeadersAPITest.java @@ -0,0 +1,101 @@ +/* + * 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.hugegraph.pd.rest; + +import java.util.Map; + +import org.apache.hugegraph.pd.common.PDException; +import org.apache.hugegraph.pd.service.PDRestService; +import org.junit.Assert; +import org.junit.Test; + +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.ObjectMapper; + +/** + * Pins the body of the two routes to the same leader balance, {@code GET /v1/task/balanceLeaders} + * and {@code GET /v1/balanceLeaders}, when the balance is refused, for example inside the window + * after {@code balancePartitions} sets the balance-shard key. Both must answer the status and + * reason like their sibling endpoints do, not let the {@link PDException} escape as a bare + * HTTP 500. + */ +public class BalanceLeadersAPITest { + + private static final String REASON = "balance shard is processing, please try later!"; + + private static final LeaderBalance REFUSED = () -> { + throw new PDException(1001, REASON); + }; + + private static final LeaderBalance BALANCED = () -> Map.of(1, 2L); + + @Test + public void testTaskRouteRefusedBalanceReturnsErrorBody() throws Exception { + TaskAPI api = new TaskAPI(); + api.pdRestService = service(REFUSED); + + assertErrorBody(api.balanceLeaders()); + } + + @Test + public void testTaskRouteSuccessfulBalanceKeepsBody() { + TaskAPI api = new TaskAPI(); + api.pdRestService = service(BALANCED); + + Assert.assertEquals("{\"1\":2}", api.balanceLeaders()); + } + + @Test + public void testStoreRouteRefusedBalanceReturnsErrorBody() throws Exception { + StoreAPI api = new StoreAPI(); + api.pdRestService = service(REFUSED); + + assertErrorBody(api.balanceLeaders()); + } + + @Test + public void testStoreRouteSuccessfulBalanceKeepsBody() { + StoreAPI api = new StoreAPI(); + api.pdRestService = service(BALANCED); + + Assert.assertEquals("{\"1\":2}", api.balanceLeaders()); + } + + private static PDRestService service(LeaderBalance balance) { + return new PDRestService() { + @Override + public Map balancePartitionLeader() throws PDException { + return balance.run(); + } + }; + } + + private static void assertErrorBody(String json) throws Exception { + Map body = new ObjectMapper().readValue( + json, new TypeReference>() { + }); + + Assert.assertEquals(1001, body.get("status")); + Assert.assertEquals(REASON, body.get("error")); + } + + private interface LeaderBalance { + + Map run() throws PDException; + } +}