diff --git a/kyuubi-hive-beeline/src/main/java/org/apache/hive/beeline/BeeLine.java b/kyuubi-hive-beeline/src/main/java/org/apache/hive/beeline/BeeLine.java index e3b29206c05..b66dca26c1a 100644 --- a/kyuubi-hive-beeline/src/main/java/org/apache/hive/beeline/BeeLine.java +++ b/kyuubi-hive-beeline/src/main/java/org/apache/hive/beeline/BeeLine.java @@ -134,7 +134,7 @@ public class BeeLine implements Closeable { private final BeeLineOpts opts = new BeeLineOpts(this, System.getProperties()); private String lastProgress = null; private final Map seenWarnings = new HashMap(); - private final Commands commands = new Commands(this); + private Commands commands = new Commands(this); private OutputFile scriptOutputFile = null; private OutputFile recordOutputFile = null; private PrintStream outputStream = new PrintStream(System.out, true); @@ -2382,4 +2382,8 @@ void setIsTestMode(boolean isTestMode) { boolean isTestMode() { return isTestMode; } + + protected void setCommands(Commands commands) { + this.commands = commands; + } } 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 index c786da35f24..878c95d86dc 100644 --- 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 @@ -48,7 +48,6 @@ public class KyuubiBeeLine extends BeeLine { public static final String KYUUBI_BEELINE_DEFAULT_JDBC_DRIVER = "org.apache.kyuubi.jdbc.KyuubiHiveDriver"; - protected KyuubiCommands commands = new KyuubiCommands(this); private Driver defaultDriver; // copied from org.apache.hive.beeline.BeeLine @@ -66,11 +65,7 @@ public KyuubiBeeLine() { @SuppressWarnings("deprecation") public KyuubiBeeLine(boolean isBeeLine) { super(isBeeLine); - try { - DynFields.builder().hiddenImpl(BeeLine.class, "commands").buildChecked(this).set(commands); - } catch (Throwable t) { - throw new ExceptionInInitializerError("Failed to inject kyuubi commands"); - } + setCommands(new KyuubiCommands(this)); try { defaultDriver = DynConstructors.builder() diff --git a/kyuubi-hive-beeline/src/test/java/org/apache/hive/beeline/KyuubiBeeLineTest.java b/kyuubi-hive-beeline/src/test/java/org/apache/hive/beeline/KyuubiBeeLineTest.java index 9c7aec35a42..8b87b13a06c 100644 --- a/kyuubi-hive-beeline/src/test/java/org/apache/hive/beeline/KyuubiBeeLineTest.java +++ b/kyuubi-hive-beeline/src/test/java/org/apache/hive/beeline/KyuubiBeeLineTest.java @@ -26,8 +26,13 @@ import java.io.PrintStream; import org.apache.kyuubi.util.reflect.DynFields; import org.junit.Test; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; public class KyuubiBeeLineTest { + + private static final Logger LOG = LoggerFactory.getLogger(KyuubiBeeLineTest.class); + @Test public void testKyuubiBeelineWithoutArgs() { KyuubiBeeLine kyuubiBeeLine = new KyuubiBeeLine(); @@ -90,15 +95,46 @@ public void testKyuubiBeeLinePythonMode() { @Test public void testKyuubiBeelineComment() { - KyuubiBeeLine kyuubiBeeLine = new KyuubiBeeLine(); - int result = kyuubiBeeLine.initArgsFromCliVars(new String[] {"-e", "--comment show database;"}); - assertEquals(0, result); - result = kyuubiBeeLine.initArgsFromCliVars(new String[] {"-e", "--comment\n show database;"}); - assertEquals(1, result); - result = - kyuubiBeeLine.initArgsFromCliVars( - new String[] {"-e", "--comment line 1 \n --comment line 2 \n show database;"}); - assertEquals(1, result); + KyuubiBeeLine interceptedKyuubiBeeLine = + new KyuubiBeeLine() { + @Override + boolean dispatch(String line) { + if (line != null && line.startsWith("!connect")) { + LOG.info("Return true for command: {} to pretend connection is established.", line); + return true; + } + return super.dispatch(line); + } + }; + + String[] cmd = new String[] {""}; + KyuubiCommands interceptedCommands = + new KyuubiCommands(interceptedKyuubiBeeLine) { + @Override + public boolean sql(String line, boolean entireLineAsCommand) { + LOG.info("Return true for sql {} to pretend sql is executed successfully.", line); + cmd[0] = line; + return true; + } + }; + interceptedKyuubiBeeLine.setCommands(interceptedCommands); + + interceptedKyuubiBeeLine.initArgs( + new String[] {"-u", "dummy_url", "-e", "--comment show database;"}); + assertEquals(0, cmd[0].length()); + + // Beeline#exit must be false to execute sql + interceptedKyuubiBeeLine.setExit(false); + interceptedKyuubiBeeLine.initArgs( + new String[] {"-u", "dummy_url", "-e", "--comment\n show database;"}); + assertEquals("show database;", cmd[0]); + + interceptedKyuubiBeeLine.setExit(false); + interceptedKyuubiBeeLine.initArgs( + new String[] { + "-u", "dummy_url", "-e", "--comment line 1 \n --comment line 2 \n show database;" + }); + assertEquals("show database;", cmd[0]); } static class BufferPrintStream extends PrintStream {