diff --git a/beeline/src/java/org/apache/hive/beeline/BeeLine.java b/beeline/src/java/org/apache/hive/beeline/BeeLine.java index 4a82635..4a42138 100644 --- a/beeline/src/java/org/apache/hive/beeline/BeeLine.java +++ b/beeline/src/java/org/apache/hive/beeline/BeeLine.java @@ -141,9 +141,9 @@ private static final String SCRIPT_OUTPUT_PREFIX = ">>>"; private static final int SCRIPT_OUTPUT_PAD_SIZE = 5; - private static final int ERRNO_OK = 0; - private static final int ERRNO_ARGS = 1; - private static final int ERRNO_OTHER = 2; + public static final int ERRNO_OK = 0; + public static final int ERRNO_ARGS = 1; + public static final int ERRNO_OTHER = 2; private static final String HIVE_VAR_PREFIX = "--hivevar"; private static final String HIVE_CONF_PREFIX = "--hiveconf"; @@ -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 isSource) { FileInputStream initStream = null; try { initStream = new FileInputStream(fileName); @@ -893,28 +895,54 @@ private int executeFile(String fileName) { return ERRNO_OTHER; } finally { IOUtils.closeStream(initStream); - consoleReader = null; - output(""); // dummy new line + if(!isSource) { + consoleReader = null; + output(""); // dummy new line + }else{ + // For source command, we should not exit when meeting some empty line. + setExit(false); + } + } + } + + private boolean isSourceCMD(String cmd) { + if(cmd == null || cmd.isEmpty()) return false; + String trimmedCmd = cmd.trim(); + String[] tokens = tokenizeCmd(trimmedCmd); + 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 { + return executeFile(cmd_1, true) == ERRNO_OK; } } private int execute(ConsoleReader reader, boolean exitOnError) { String line; - if (!isBeeLine) { - if (defaultConnect(exitOnError) != ERRNO_OK && exitOnError) { - return ERRNO_OTHER; - } - } 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()); if (!isBeeLine) { line = cliToBeelineCmd(line); } - if (!dispatch(line) && exitOnError) { + + boolean ret; + if (isSourceCMD(line)) { + ret = sourceFile(line); + } else { + ret = dispatch(line); + } + if (!ret && exitOnError) { return ERRNO_OTHER; } } catch (Throwable t) { @@ -1020,23 +1048,25 @@ void usage() { return cmd.split("\\s+"); } + /** + * Extract and clean up the first command in the input. + */ + private String getFirstCmd(String cmd, int length) { + return cmd.substring(length).trim(); + } + public 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")) { return null; - } else if (tokens[0].equalsIgnoreCase("source")) { - //TODO - return cmd; } else if (cmd_trimmed.startsWith("!")) { String shell_cmd = cmd_trimmed.substring(1); return "!sh " + shell_cmd; } else { // local mode // command like dfs - return cmd; + return cmd_trimmed; } } 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..0ee8e8d 100644 --- a/beeline/src/test/org/apache/hive/beeline/cli/TestHiveCli.java +++ b/beeline/src/test/org/apache/hive/beeline/cli/TestHiveCli.java @@ -21,6 +21,7 @@ import org.apache.commons.io.IOUtils; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; +import org.apache.hive.beeline.BeeLine; import org.junit.After; import org.junit.Before; import org.junit.Test; @@ -37,9 +38,10 @@ 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 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; @@ -74,44 +76,50 @@ private void verifyCMD(String CMD, String keywords, OutputStream os, String[] op @Test public void testInValidCmd() { - verifyCMD("!lss\n", "Failed to execute lss", errS, null, 0); + verifyCMD("!lss\n", "Failed to execute lss", errS, null, BeeLine.ERRNO_OK); } @Test public void testHelp() { - verifyCMD(null, "usage: hive", os, new String[]{"-H"}, 1); + verifyCMD(null, "usage: hive", os, new String[]{"-H"}, BeeLine.ERRNO_ARGS); } @Test public void testInvalidDatabaseOptions() { verifyCMD("\nshow tables\nquit\n", "Database does not exist: invalidDB", errS, new - String[]{"--database", "invalidDB"}, 0); + String[]{"--database", "invalidDB"}, BeeLine.ERRNO_OK); } - @Test - public void testDatabaseOptions() { - verifyCMD("\nshow tables;\nquit;", "testTbl", os, new String[]{"--database", "test"}, 0); + @Test public void testDatabaseOptions() { + verifyCMD("\nshow tables;\nquit;", "testTbl", os, new String[] { "--database", "test" }, + BeeLine.ERRNO_OK); + } + + @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" }, BeeLine.ERRNO_OK); } @Test public void testSqlFromCmd() { - verifyCMD(null, "", os, new String[]{"-e", "show databases;"}, 0); + verifyCMD(null, "", os, new String[]{"-e", "show databases;"}, BeeLine.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" }, + BeeLine.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); + String[]{"-e", "show tables;", "-f", "path/to/file"}, BeeLine.ERRNO_ARGS); } @Test public void testInvalidOptions2() { - verifyCMD(null, "Unrecognized option: -k", errS, new String[]{"-k"}, 1); + verifyCMD(null, "Unrecognized option: -k", errS, new String[]{"-k"}, BeeLine.ERRNO_ARGS); } private void redirectOutputStream() { @@ -124,26 +132,34 @@ 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() { cli = new HiveCli(); redirectOutputStream(); - initFileFromFile(); + initFromFile(); } @After