options

python options save save remove comments for non python mode save ut
apache · Jun 15, 2023 · 3337ca8 · 3337ca8
1 parent 01320c4
commit 3337ca8
Show file tree

Hide file tree

Showing 3 changed files with 95 additions and 4 deletions.
diff --git a/kyuubi-hive-beeline/src/main/java/org/apache/hive/beeline/KyuubiBeeLine.java b/kyuubi-hive-beeline/src/main/java/org/apache/hive/beeline/KyuubiBeeLine.java
@@ -19,10 +19,10 @@
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.lang.reflect.Field;
+import java.lang.reflect.Modifier;
 import java.sql.Driver;
-import java.util.Arrays;
-import java.util.Collections;
-import java.util.List;
+import java.util.*;
 import org.apache.commons.cli.CommandLine;
 import org.apache.commons.cli.Options;
 import org.apache.commons.cli.ParseException;
@@ -41,6 +41,12 @@ public class KyuubiBeeLine extends BeeLine {
   private static final int ERRNO_ARGS = 1;
   private static final int ERRNO_OTHER = 2;
 
+  private static final ResourceBundle beelineResourceBundle =
+      ResourceBundle.getBundle(BeeLine.class.getSimpleName());
+  private static final ResourceBundle kyuubiResourceBundle = new KyuubiBeelineResourceBundle();
+  private static final String PYTHON_MODE_PREFIX = "--python-mode";
+  private boolean pythonMode = false;
+
   public KyuubiBeeLine() {
     this(true);
   }
@@ -50,6 +56,13 @@ public KyuubiBeeLine(boolean isBeeLine) {
     super(isBeeLine);
     try {
       DynFields.builder().hiddenImpl(BeeLine.class, "commands").buildChecked(this).set(commands);
+
+      Field resourceBundleField = BeeLine.class.getDeclaredField("resourceBundle");
+      resourceBundleField.setAccessible(true);
+      Field modifiers = Field.class.getDeclaredField("modifiers");
+      modifiers.setAccessible(true);
+      modifiers.setInt(resourceBundleField, resourceBundleField.getModifiers() & ~Modifier.FINAL);
+      resourceBundleField.set(null, kyuubiResourceBundle);
     } catch (Throwable t) {
       throw new ExceptionInInitializerError("Failed to inject kyuubi commands");
     }
@@ -64,6 +77,15 @@ public KyuubiBeeLine(boolean isBeeLine) {
     }
   }
 
+  public boolean isPythonMode() {
+    return pythonMode;
+  }
+
+  // Visible for testing
+  public void setPythonMode(boolean pythonMode) {
+    this.pythonMode = pythonMode;
+  }
+
   /** Starts the program. */
   public static void main(String[] args) throws IOException {
     mainWithInputRedirection(args, null);
@@ -125,7 +147,20 @@ int initArgs(String[] args) {
               .<Options>buildStaticChecked()
               .get();
 
-      beelineParser = new BeelineParser();
+      beelineParser =
+          new BeelineParser() {
+            @Override
+            protected void processOption(String arg, ListIterator iter) throws ParseException {
+              if (PYTHON_MODE_PREFIX.equals(arg)) {
+                String stripped = arg.substring(2, arg.length());
+                String[] parts = split(stripped, "=");
+                String value = parts.length >= 2 ? parts[1] : "true";
+                pythonMode = Boolean.parseBoolean(value);
+              } else {
+                super.processOption(arg, iter);
+              }
+            }
+          };
       cl = beelineParser.parse(options, args);
 
       connSuccessful =
@@ -248,4 +283,34 @@ int runInit() {
     }
     return executionResult;
   }
+
+  static class KyuubiBeelineResourceBundle extends ListResourceBundle {
+    static String CMD_USAGE = "cmd-usage";
+
+    private Object[][] contents = new Object[beelineResourceBundle.keySet().size()][];
+
+    public KyuubiBeelineResourceBundle() {
+      int i = 0;
+      for (String key : beelineResourceBundle.keySet()) {
+        String value = beelineResourceBundle.getString(key);
+        if (key.equals(CMD_USAGE)) {
+          StringBuilder stringBuilder = new StringBuilder();
+          stringBuilder.append(value).append("\n");
+          stringBuilder
+              .append("Usage: java " + KyuubiBeeLine.class.getCanonicalName())
+              .append("\n");
+          stringBuilder.append(
+              "   --python-mode=[true/false]                      Execute python code/script.");
+          value = stringBuilder.toString();
+        }
+        contents[i] = new Object[] {key, value};
+        i++;
+      }
+    }
+
+    @Override
+    protected Object[][] getContents() {
+      return contents;
+    }
+  }
 }
diff --git a/kyuubi-hive-beeline/src/main/java/org/apache/hive/beeline/KyuubiCommands.java b/kyuubi-hive-beeline/src/main/java/org/apache/hive/beeline/KyuubiCommands.java
@@ -24,6 +24,7 @@
 import java.sql.*;
 import java.util.*;
 import org.apache.hive.beeline.logs.KyuubiBeelineInPlaceUpdateStream;
+import org.apache.hive.common.util.HiveStringUtils;
 import org.apache.kyuubi.jdbc.hive.KyuubiStatement;
 import org.apache.kyuubi.jdbc.hive.Utils;
 import org.apache.kyuubi.jdbc.hive.logs.InPlaceUpdateStream;
@@ -522,6 +523,10 @@ public String handleMultiLineCmd(String line) throws IOException {
             ? null
             : jline.console.ConsoleReader.NULL_MASK;
 
+    int[] startQuote = {-1};
+    if (!beeLine.isPythonMode()) {
+      line = HiveStringUtils.removeComments(line, startQuote);
+    }
     while (isMultiLine(line) && beeLine.getOpts().isAllowMultiLineCommand()) {
       StringBuilder prompt = new StringBuilder(beeLine.getPrompt());
       if (!beeLine.getOpts().isSilent()) {
@@ -547,6 +552,9 @@ public String handleMultiLineCmd(String line) throws IOException {
       if (extra == null) { // it happens when using -f and the line of cmds does not end with ;
         break;
       }
+      if (!beeLine.isPythonMode()) {
+        extra = HiveStringUtils.removeComments(extra, startQuote);
+      }
       if (!extra.isEmpty()) {
         line += "\n" + extra;
       }

diff --git a/kyuubi-hive-beeline/src/test/java/org/apache/hive/beeline/KyuubiCommandsTest.java b/kyuubi-hive-beeline/src/test/java/org/apache/hive/beeline/KyuubiCommandsTest.java
@@ -34,6 +34,7 @@ public void testParsePythonSnippets() throws IOException {
     Mockito.when(reader.readLine()).thenReturn(pythonSnippets);
 
     KyuubiBeeLine beeline = new KyuubiBeeLine();
+    beeline.setPythonMode(true);
     beeline.setConsoleReader(reader);
     KyuubiCommands commands = new KyuubiCommands(beeline);
     String line = commands.handleMultiLineCmd(pythonSnippets);
@@ -42,4 +43,21 @@ public void testParsePythonSnippets() throws IOException {
     assertEquals(cmdList.size(), 1);
     assertEquals(cmdList.get(0), pythonSnippets);
   }
+
+  @Test
+  public void testHandleMultiLineCmd() throws IOException {
+    ConsoleReader reader = Mockito.mock(ConsoleReader.class);
+    String snippets = "select 1;--comments1\nselect 2;--comments2";
+    Mockito.when(reader.readLine()).thenReturn(snippets);
+
+    KyuubiBeeLine beeline = new KyuubiBeeLine();
+    beeline.setConsoleReader(reader);
+    beeline.setPythonMode(false);
+    KyuubiCommands commands = new KyuubiCommands(beeline);
+    String line = commands.handleMultiLineCmd(snippets);
+    List<String> cmdList = commands.getCmdList(line, false);
+    assertEquals(cmdList.size(), 2);
+    assertEquals(cmdList.get(0), "select 1");
+    assertEquals(cmdList.get(1), "\nselect 2");
+  }
 }