[GH-16351] Do not call System.exit from water.tools [nocheck] (#16366)

* [GH-16351] Do not call System.exit from water.tools * add mainInternal to EncryptionTool * make the err msg short * make the err msg short - part 2 * improve error handling
h2oai · Oct 23, 2024 · d0899f8 · d0899f8
1 parent eaccd47
commit d0899f8
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 19 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 {
-        if (args.length < 2) {
-            usage();
+        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 < 2 || args[0] == null || args[1] == null) {
+            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);
+        if (!mojoFile.exists() || !mojoFile.isFile()) {
+            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,10 +33,12 @@ 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);
+            RuntimeException shorterException = new RuntimeException(e.getCause().getMessage());
+            shorterException.setStackTrace(new StackTraceElement[0]);
+            throw shorterException;
         }
         return new ValStr("OK");
     }

diff --git a/h2o-core/src/main/java/water/tools/EncryptionTool.java b/h2o-core/src/main/java/water/tools/EncryptionTool.java
@@ -47,6 +47,9 @@ public void encrypt(File input, File output) throws IOException, GeneralSecurity
     }
 
     public static void main(String[] args) throws GeneralSecurityException, IOException {
+        mainInternal(args);
+    }
+    public static void mainInternal(String[] args) throws GeneralSecurityException, IOException {
         EncryptionTool et = new EncryptionTool();
         et._keystore_file = new File(args[0]);
         et._keystore_type = args[1];

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())