diff --git a/beeline/src/java/org/apache/hive/beeline/BeeLine.java b/beeline/src/java/org/apache/hive/beeline/BeeLine.java index 4a82635..c4eed9a 100644 --- a/beeline/src/java/org/apache/hive/beeline/BeeLine.java +++ b/beeline/src/java/org/apache/hive/beeline/BeeLine.java @@ -824,12 +824,13 @@ public int begin(String[] args, InputStream inputStream) throws IOException { } } else { int code = initArgsFromCliVars(args); + defaultConnect(false); if (code != 0) return code; } if (getOpts().getScriptFile() != null) { - return executeFile(getOpts().getScriptFile()); + return executeFile(getOpts().getScriptFile(), false); } try { info(getApplicationTitle()); @@ -837,6 +838,7 @@ public int begin(String[] args, InputStream inputStream) throws IOException { // ignore } ConsoleReader reader = getConsoleReader(inputStream); + return execute(reader, false); } finally { close(); @@ -848,7 +850,7 @@ int runInit() { if (initFile != null) { info("Running init script " + initFile); try { - return executeFile(initFile); + return executeFile(initFile, false); } finally { exit = false; } @@ -883,7 +885,7 @@ public int defaultConnect(boolean exitOnError) { return ERRNO_OK; } - private int executeFile(String fileName) { + private int executeFile(String fileName, boolean isSourceCMD) { FileInputStream initStream = null; try { initStream = new FileInputStream(fileName); @@ -893,28 +895,57 @@ private int executeFile(String fileName) { return ERRNO_OTHER; } finally { IOUtils.closeStream(initStream); - consoleReader = null; - output(""); // dummy new line + if(!isSourceCMD) { + consoleReader = null; + output(""); // dummy new line + } + } + } + + private boolean isSourceCMD(String cmd) { + if (cmd == null || cmd.isEmpty()) + return false; + String[] tokens = tokenizeCmd(cmd); + return tokens[0].equalsIgnoreCase("source"); + } + + private boolean sourceFile(String cmd) { + String trimedCmd = cmd.trim(); + String[] tokens = tokenizeCmd(trimedCmd); + String cmd_1 = getFirstCmd(trimedCmd, tokens[0].length()); + File sourceFile = new File(cmd_1); + if (!sourceFile.isFile()) { + return false; + } else { + boolean ret = (executeFile(cmd_1, true) == ERRNO_OK); + // For source command, we should not exit even when meeting some empty line. + setExit(false); + return ret; } } private int execute(ConsoleReader reader, boolean exitOnError) { String line; - if (!isBeeLine) { - if (defaultConnect(exitOnError) != ERRNO_OK && exitOnError) { - return ERRNO_OTHER; - } - } + boolean ret; while (!exit) { try { // Execute one instruction; terminate on executing a script if there is an error // in silent mode, prevent the query and prompt being echoed back to terminal - line = (getOpts().isSilent() && getOpts().getScriptFile() != null) ? - reader.readLine(null, ConsoleReader.NULL_MASK) : reader.readLine(getPrompt()); + line = (getOpts().isSilent() && getOpts().getScriptFile() != null) ? reader + .readLine(null, ConsoleReader.NULL_MASK) : reader.readLine(getPrompt()); + + // trim line + line = (line == null) ? null : line.trim(); if (!isBeeLine) { line = cliToBeelineCmd(line); } - if (!dispatch(line) && exitOnError) { + + if (isSourceCMD(line)) { + ret = sourceFile(line); + } else { + ret = dispatch(line); + } + if (!ret && exitOnError) { return ERRNO_OTHER; } } catch (Throwable t) { @@ -1020,19 +1051,20 @@ void usage() { return cmd.split("\\s+"); } - public String cliToBeelineCmd(String cmd) { + /** + * Extract and clean up the first command in the input. + */ + private String getFirstCmd(String cmd, int length) { + return cmd.substring(length).trim(); + } + + private String cliToBeelineCmd(String cmd) { if (cmd == null) return null; - String cmd_trimmed = cmd.trim(); - String[] tokens = tokenizeCmd(cmd_trimmed); - - if (cmd_trimmed.equalsIgnoreCase("quit") || cmd_trimmed.equalsIgnoreCase("exit")) { + if (cmd.equalsIgnoreCase("quit") || cmd.equalsIgnoreCase("exit")) { return null; - } else if (tokens[0].equalsIgnoreCase("source")) { - //TODO - return cmd; - } else if (cmd_trimmed.startsWith("!")) { - String shell_cmd = cmd_trimmed.substring(1); + } else if (cmd.startsWith("!")) { + String shell_cmd = cmd.substring(1); return "!sh " + shell_cmd; } else { // local mode // command like dfs @@ -1040,7 +1072,6 @@ public String cliToBeelineCmd(String cmd) { } } - /** * Dispatch the specified line to the appropriate {@link CommandHandler}. * diff --git a/beeline/src/test/org/apache/hive/beeline/cli/TestHiveCli.java b/beeline/src/test/org/apache/hive/beeline/cli/TestHiveCli.java index cc0b598..6cbb030 100644 --- a/beeline/src/test/org/apache/hive/beeline/cli/TestHiveCli.java +++ b/beeline/src/test/org/apache/hive/beeline/cli/TestHiveCli.java @@ -36,10 +36,15 @@ public class TestHiveCli { private static final Log LOG = LogFactory.getLog(TestHiveCli.class.getName()); - - final static String CMD = "create database if not exists test;\ncreate table if not exists test" + - ".testTbl(a " + - "" + "string, b string);\n"; + private static final int ERRNO_OK = 0; + private static final int ERRNO_ARGS = 1; + private static final int ERRNO_OTHER = 2; + + private final static String SOURCE_CONTEXT = + "create table if not exists test.testSrcTbl(a string, b string);"; + final static String CMD = + "create database if not exists test;\ncreate table if not exists test.testTbl(a string, b " + + "string);\n"; private HiveCli cli; private OutputStream os; private PrintStream ps; @@ -65,53 +70,53 @@ private void executeCMD(String[] args, String input, int retCode) { } } - private void verifyCMD(String CMD, String keywords, OutputStream os, String[] options, int - retCode) { + private void verifyCMD(String CMD, String keywords, OutputStream os, String[] options, + int retCode) { executeCMD(options, CMD, retCode); String output = os.toString(); Assert.assertTrue(output.contains(keywords)); } - @Test - public void testInValidCmd() { - verifyCMD("!lss\n", "Failed to execute lss", errS, null, 0); + @Test public void testInValidCmd() { + verifyCMD("!lss\n", "Failed to execute lss", errS, null, ERRNO_OK); + } + + @Test public void testHelp() { + verifyCMD(null, "usage: hive", os, new String[] { "-H" }, ERRNO_ARGS); } - @Test - public void testHelp() { - verifyCMD(null, "usage: hive", os, new String[]{"-H"}, 1); + @Test public void testInvalidDatabaseOptions() { + verifyCMD("\nshow tables\nquit\n", "Database does not exist: invalidDB", errS, + new String[] { "--database", "invalidDB" }, ERRNO_OK); } - @Test - public void testInvalidDatabaseOptions() { - verifyCMD("\nshow tables\nquit\n", "Database does not exist: invalidDB", errS, new - String[]{"--database", "invalidDB"}, 0); + @Test public void testDatabaseOptions() { + verifyCMD("\nshow tables;\nquit;", "testTbl", os, new String[] { "--database", "test" }, + ERRNO_OK); } - @Test - public void testDatabaseOptions() { - verifyCMD("\nshow tables;\nquit;", "testTbl", os, new String[]{"--database", "test"}, 0); + @Test public void testSourceCmd() { + File f = generateTmpFile(SOURCE_CONTEXT); + verifyCMD("source " + f.getPath() + "\n" + "desc testSrcTbl\n" + "quit\n", "col_name", os, + new String[] { "--database", "test" }, ERRNO_OK); } - @Test - public void testSqlFromCmd() { - verifyCMD(null, "", os, new String[]{"-e", "show databases;"}, 0); + @Test public void testSqlFromCmd() { + verifyCMD(null, "", os, new String[] { "-e", "show databases;" }, ERRNO_OK); } - @Test - public void testSqlFromCmdWithDBName() { - verifyCMD(null, "testTbl", os, new String[]{"-e", "show tables;", "--database", "test"}, 0); + @Test public void testSqlFromCmdWithDBName() { + verifyCMD(null, "testTbl", os, new String[] { "-e", "show tables;", "--database", "test" }, + ERRNO_OK); } - @Test - public void testInvalidOptions() { - verifyCMD(null, "The '-e' and '-f' options cannot be specified simultaneously", errS, new - String[]{"-e", "show tables;", "-f", "path/to/file"}, 1); + @Test public void testInvalidOptions() { + verifyCMD(null, "The '-e' and '-f' options cannot be specified simultaneously", errS, + new String[] { "-e", "show tables;", "-f", "path/to/file" }, ERRNO_ARGS); } - @Test - public void testInvalidOptions2() { - verifyCMD(null, "Unrecognized option: -k", errS, new String[]{"-k"}, 1); + @Test public void testInvalidOptions2() { + verifyCMD(null, "Unrecognized option: -k", errS, new String[] { "-k" }, ERRNO_ARGS); } private void redirectOutputStream() { @@ -124,30 +129,36 @@ private void redirectOutputStream() { System.setErr(errPs); } - private void initFileFromFile() { + private void initFromFile() { + tmp = generateTmpFile(CMD); + if (tmp == null) { + Assert.fail("Fail to create the initial file"); + } + executeCMD(new String[] { "-f", "\"" + tmp.getAbsolutePath() + "\"" }, null, 0); + } + + private File generateTmpFile(String context) { + File file = null; BufferedWriter bw = null; try { - // create a tmp file - tmp = File.createTempFile("test", ".sql"); - bw = new BufferedWriter(new FileWriter(tmp)); - bw.write(CMD); + file = File.createTempFile("test", ".sql"); + bw = new BufferedWriter(new FileWriter(file)); + bw.write(context); } catch (IOException e) { LOG.error("Failed to write tmp file due to the exception: " + e); } finally { IOUtils.closeQuietly(bw); } - executeCMD(new String[]{"-f", "\"" + tmp.getAbsolutePath() + "\""}, null, 0); + return file; } - @Before - public void setup() { + @Before public void setup() { cli = new HiveCli(); redirectOutputStream(); - initFileFromFile(); + initFromFile(); } - @After - public void tearDown() { + @After public void tearDown() { tmp.delete(); } }