[GH-16351] Do not call System.exit from water.tools

h2oai · Aug 14, 2024 · 6113424 · 6113424
1 parent 9e9fe40
commit 6113424
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 16 deletions.
diff --git a/h2o-algos/src/main/java/water/tools/MojoConvertTool.java b/h2o-algos/src/main/java/water/tools/MojoConvertTool.java
@@ -33,25 +33,28 @@ void convert() throws IOException {
  Files.write(pojoPath, pojo.getBytes(StandardCharsets.UTF_8));
  }
 
- private static void usage() {
- System.err.println("java -cp h2o.jar " + MojoConvertTool.class.getName() + " source_mojo.zip target_pojo.java");
+ public static void main(String[] args) throws IOException {
+ try {
+ mainInternal(args);
+ }
+ catch (IllegalArgumentException e) {
+ System.err.println(e.getMessage());
+ System.exit(1);
+ }
  }
 
- public static void main(String[] args) throws IOException {
+ public static void mainInternal(String[] args) throws IOException {
  if (args.length < 2) {
- usage();
- System.exit(1);
+ throw new IllegalArgumentException("java -cp h2o.jar " + MojoConvertTool.class.getName() + " source_mojo.zip target_pojo.java");
  }
 
  File mojoFile = new File(args[0]);
  if (!mojoFile.isFile()) {
- System.err.println("Specified MOJO file (" + mojoFile.getAbsolutePath() + ") doesn't exist!");
- System.exit(2);
+ throw new IllegalArgumentException("Specified MOJO file (" + mojoFile.getAbsolutePath() + ") doesn't exist!");
  }
  File pojoFile = new File(args[1]);
  if (pojoFile.isDirectory() || (pojoFile.getParentFile() != null && !pojoFile.getParentFile().isDirectory())) {
- System.err.println("Invalid target POJO file (" + pojoFile.getAbsolutePath() + ")! Please specify a file in an existing directory.");
- System.exit(3);
+ throw new IllegalArgumentException("Invalid target POJO file (" + pojoFile.getAbsolutePath() + ")! Please specify a file in an existing directory.");
  }
 
  System.out.println();

diff --git a/h2o-core/src/main/java/water/rapids/ast/prims/internal/AstRunTool.java b/h2o-core/src/main/java/water/rapids/ast/prims/internal/AstRunTool.java
@@ -33,7 +33,7 @@ public ValStr apply(Env env, Env.StackHelp stk, AstRoot[] asts) {
  try {
  // only allow to run approved tools (from our package), not just anything on classpath
  Class<?> clazz = Class.forName(TOOLS_PACKAGE + toolClassName);
- Method mainMethod = clazz.getDeclaredMethod("main", String[].class);
+ Method mainMethod = clazz.getDeclaredMethod("mainInternal", String[].class);
  mainMethod.invoke(null, new Object[]{args});
  } catch (Exception e) {
  throw new RuntimeException(e);

diff --git a/h2o-extensions/xgboost/src/main/java/water/tools/XGBoostLibExtractTool.java b/h2o-extensions/xgboost/src/main/java/water/tools/XGBoostLibExtractTool.java
@@ -10,19 +10,25 @@
 public class XGBoostLibExtractTool {
 
  public static void main(String[] args) throws IOException {
+ try {
+ mainInternal(args);
+ } catch (IllegalArgumentException e) {
+ System.err.println((e.getMessage());
+ System.exit(1);
+ }
+ }
+
+ public static void mainInternal(String[] args) throws IOException {
  if (args.length != 1) {
- System.err.println("XGBoostLibExtractTool: Specify target directory where to extract XGBoost native libraries.");
- System.exit(-1);
+ throw new IllegalArgumentException("XGBoostLibExtractTool: Specify target directory where to extract XGBoost native libraries.");
  }
  File dir = new File(args[0]);
  if (!dir.exists()) {
- System.err.println("XGBoostLibExtractTool: Directory '" + dir.getAbsolutePath() + "' doesn't exist.");
- System.exit(-1);
+ throw new IllegalArgumentException("XGBoostLibExtractTool: Directory '" + dir.getAbsolutePath() + "' doesn't exist.");
  }
  NativeLibraryLoaderChain loader = XGBoostExtension.getLoader();
  if (loader == null) {
- System.err.println("XGBoostLibExtractTool: Failed to locate native libraries.");
- System.exit(-1);
+ throw new IllegalArgumentException("XGBoostLibExtractTool: Failed to locate native libraries.");
  }
  for (NativeLibrary lib : loader.getNativeLibs()) {
  if (!lib.isBundled())