Skip to content

Commit

Permalink
[apache#768] feat(cli): Cli method for blacklist update
Browse files Browse the repository at this point in the history
  • Loading branch information
wangjunbo committed Jun 14, 2023
1 parent c274b37 commit d26cca0
Show file tree
Hide file tree
Showing 6 changed files with 156 additions and 15 deletions.
2 changes: 1 addition & 1 deletion bin/uniffle
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ function uniffle_cmd_case
UNIFFLE_CLASSNAME=org.apache.uniffle.cli.UniffleCLI
;;
admin-cli)
UNIFFLE_CLASSNAME=org.apache.uniffle.cli.UniffleCLI
UNIFFLE_CLASSNAME=org.apache.uniffle.cli.UniffleAdminCLI
;;
*)
UNIFFLE_CLASSNAME="${subcmd}"
Expand Down
131 changes: 131 additions & 0 deletions cli/src/main/java/org/apache/uniffle/cli/UniffleAdminCLI.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,131 @@
/*
* 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.uniffle.cli;

import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.Option;
import org.apache.commons.cli.Options;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.uniffle.AbstractCustomCommandLine;
import org.apache.uniffle.UniffleCliArgsException;
import org.apache.uniffle.api.AdminRestApi;
import org.apache.uniffle.client.UniffleRestClient;

public class UniffleAdminCLI extends AbstractCustomCommandLine {

private static final Logger LOG = LoggerFactory.getLogger(UniffleAdminCLI.class);

private final Options allOptions;
private final Option checkerClass;
private final Option coordinatorHost;
private final Option coordPort;
private final Option refreshAccessCli;
private final Option help;
protected UniffleRestClient client;

public UniffleAdminCLI(String shortPrefix, String longPrefix) {
allOptions = new Options();
refreshAccessCli = new Option(shortPrefix + "r", longPrefix + "refreshChecker",
false, "This is an admin command that will refresh access checker.");
checkerClass = new Option(shortPrefix + "c", longPrefix + "checkerClass",
true, "This is the checker class that will be refreshed.");
help = new Option(shortPrefix + "h", longPrefix + "help",
false, "Help for the Uniffle Admin CLI.");
coordinatorHost = new Option(shortPrefix + "s", longPrefix + "coordinatorHost",
true, "This is coordinator server host.");
coordPort = new Option(shortPrefix + "p", longPrefix + "port",
true, "This is coordinator server port.");
allOptions.addOption(coordinatorHost);
allOptions.addOption(coordPort);
allOptions.addOption(refreshAccessCli);
allOptions.addOption(help);
}

public UniffleAdminCLI(String shortPrefix, String longPrefix, UniffleRestClient client) {
this(shortPrefix, longPrefix);
this.client = client;
}

public int run(String[] args) throws UniffleCliArgsException {
final CommandLine cmd = parseCommandLineOptions(args, true);

if (cmd.hasOption(help.getOpt())) {
printUsage();
return 0;
}

if (cmd.hasOption(coordinatorHost.getOpt()) && cmd.hasOption(coordPort.getOpt())) {
String host = cmd.getOptionValue(coordinatorHost.getOpt()).trim();
int port = Integer.parseInt(cmd.getOptionValue(coordPort.getOpt()).trim());
String hostUrl = String.format("https://%s:%d", host, port);
LOG.info("connected to coordinator: https://{}:{}.", host, port);
client = UniffleRestClient.builder(hostUrl).build();
}

if (cmd.hasOption(refreshAccessCli.getOpt())) {
if (!cmd.hasOption(checkerClass.getOpt())) {
throw new UniffleCliArgsException("");
}
String checker = cmd.getOptionValue(checkerClass.getOpt());
if (checker == null) {
throw new UniffleCliArgsException("");
}
LOG.info("uniffle-admin-cli : refresh access checker {}!", checker);
refreshAccessChecker(checker);
return 0;
}

return 1;
}

private String refreshAccessChecker(String checker) throws UniffleCliArgsException {
if (client == null) {
throw new UniffleCliArgsException("Missing Coordinator host address and grpc port parameters.");
}
AdminRestApi adminRestApi = new AdminRestApi(client);
return adminRestApi.refreshAccessChecker(checker);
}

@Override
public void addRunOptions(Options baseOptions) {
baseOptions.addOption(refreshAccessCli);
baseOptions.addOption(checkerClass);
baseOptions.addOption(coordinatorHost);
baseOptions.addOption(coordPort);
}

@Override
public void addGeneralOptions(Options baseOptions) {
baseOptions.addOption(help);
}

public static void main(String[] args) {
int retCode;
try {
final UniffleAdminCLI cli = new UniffleAdminCLI("", "");
retCode = cli.run(args);
} catch (UniffleCliArgsException e) {
retCode = AbstractCustomCommandLine.handleCliArgsException(e, LOG);
} catch (Exception e) {
retCode = AbstractCustomCommandLine.handleError(e, LOG);
}
System.exit(retCode);
}
}
22 changes: 14 additions & 8 deletions cli/src/main/java/org/apache/uniffle/cli/UniffleCLI.java
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ public class UniffleCLI extends AbstractCustomCommandLine {
private final Option uniffleClientCli;
private final Option uniffleAdminCli;
private final Option checkerClass;
private final Option coordServer;
private final Option coordinatorHost;
private final Option coordPort;
private final Option refreshAccessCli;
private final Option help;
Expand All @@ -54,13 +54,13 @@ public UniffleCLI(String shortPrefix, String longPrefix) {
true, "This is the checker class that will be refreshed.");
help = new Option(shortPrefix + "h", longPrefix + "help",
false, "Help for the Uniffle CLI.");
coordServer = new Option(shortPrefix + "s", longPrefix + "host",
coordinatorHost = new Option(shortPrefix + "ch", longPrefix + "coordinatorHost",
true, "This is coordinator server host.");
coordPort = new Option(shortPrefix + "p", longPrefix + "port",
true, "This is coordinator server port.");
allOptions.addOption(uniffleClientCli);
allOptions.addOption(uniffleAdminCli);
allOptions.addOption(coordServer);
allOptions.addOption(coordinatorHost);
allOptions.addOption(coordPort);
allOptions.addOption(refreshAccessCli);
allOptions.addOption(help);
Expand Down Expand Up @@ -90,16 +90,22 @@ public int run(String[] args) throws UniffleCliArgsException {
System.out.println("uniffle-admin-cli : " + cliArgs);
return 0;
}
if (cmd.hasOption(coordServer.getOpt()) && cmd.hasOption(coordPort.getOpt())) {
String host = cmd.getOptionValue(coordServer.getOpt()).trim();
if (cmd.hasOption(coordinatorHost.getOpt()) && cmd.hasOption(coordPort.getOpt())) {
String host = cmd.getOptionValue(coordinatorHost.getOpt()).trim();
int port = Integer.parseInt(cmd.getOptionValue(coordPort.getOpt()).trim());
String hostUrl = String.format("http://%s:%d", host, port);
LOG.info("uniffle-admin-cli : coordinator server host {}, port {}.", host, port);
String hostUrl = String.format("https://%s:%d", host, port);
LOG.info("connected to coordinator: https://{}:{}.", host, port);
client = UniffleRestClient.builder(hostUrl).build();
}

if (cmd.hasOption(refreshAccessCli.getOpt())) {
if (!cmd.hasOption(checkerClass.getOpt())) {
throw new UniffleCliArgsException("");
}
String checker = cmd.getOptionValue(checkerClass.getOpt());
if (checker == null) {
throw new UniffleCliArgsException("");
}
LOG.info("uniffle-admin-cli : refresh access checker {}!", checker);
refreshAccessChecker(checker);
return 0;
Expand All @@ -122,7 +128,7 @@ public void addRunOptions(Options baseOptions) {
baseOptions.addOption(uniffleAdminCli);
baseOptions.addOption(refreshAccessCli);
baseOptions.addOption(checkerClass);
baseOptions.addOption(coordServer);
baseOptions.addOption(coordinatorHost);
baseOptions.addOption(coordPort);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@
public class UniffleRestClient implements AutoCloseable, Cloneable {

private RestClient restClient;

private RestClientConf conf;

private String hostUrl;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -38,9 +38,9 @@ public void testAdminRefreshCLI() throws UniffleCliArgsException, IOException {
UniffleRestClient client = Mockito.mock(UniffleRestClient.class, Answers.RETURNS_DEEP_STUBS);
Mockito.when(client.getHttpClient().get(Mockito.anyString(), Mockito.anyMap(), Mockito.anyString()))
.thenReturn("OK");
UniffleCLI uniffleCLI = new UniffleCLI("", "", client);
String[] args = {"-rc", "-cc", "org.apache.uniffle.test.AccessClusterTest$MockedAccessChecker"};
assertEquals(0, uniffleCLI.run(args));
UniffleAdminCLI uniffleAdminCLI = new UniffleAdminCLI("", "", client);
String[] args = {"-r", "-c", "org.apache.uniffle.test.AccessClusterTest$MockedAccessChecker"};
assertEquals(0, uniffleAdminCLI.run(args));
Mockito.verify(client.getHttpClient(),
Mockito.times(1)).get("/api/server/admin/refreshChecker",
Collections.singletonMap("class", "org.apache.uniffle.test.AccessClusterTest$MockedAccessChecker"), null);
Expand All @@ -49,9 +49,9 @@ public void testAdminRefreshCLI() throws UniffleCliArgsException, IOException {

@Test
public void testMissingClientCLI() throws UniffleCliArgsException, IOException {
UniffleCLI uniffleCLI = new UniffleCLI("", "");
UniffleAdminCLI uniffleAdminCLI = new UniffleAdminCLI("", "");
String[] args = {"-rc", "-cc", "org.apache.uniffle.test.AccessClusterTest$MockedAccessChecker"};
assertThrows(UniffleCliArgsException.class, () -> uniffleCLI.run(args));
assertThrows(UniffleCliArgsException.class, () -> uniffleAdminCLI.run(args));
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,11 @@ public AccessQuotaChecker(AccessManager accessManager) throws Exception {
hostIp = RssUtils.getHostIp();
}

@Override
public void refreshAccessChecker() {

}

@Override
public AccessCheckResult check(AccessInfo accessInfo) {
COUNTER.increment();
Expand Down

0 comments on commit d26cca0

Please sign in to comment.