Java tutorial
/** * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.apache.tajo.cli.tsql; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; import com.google.common.collect.Maps; import com.google.protobuf.ServiceException; import jline.UnsupportedTerminal; import jline.console.ConsoleReader; import jline.console.completer.AggregateCompleter; import jline.console.completer.ArgumentCompleter; import jline.console.completer.Completer; import jline.console.completer.StringsCompleter; import org.apache.commons.cli.*; import org.apache.tajo.*; import org.apache.tajo.TajoProtos.QueryState; import org.apache.tajo.annotation.Nullable; import org.apache.tajo.catalog.TableDesc; import org.apache.tajo.cli.tsql.ParsedResult.StatementType; import org.apache.tajo.cli.tsql.SimpleParser.ParsingState; import org.apache.tajo.cli.tsql.commands.*; import org.apache.tajo.client.*; import org.apache.tajo.conf.TajoConf; import org.apache.tajo.conf.TajoConf.ConfVars; import org.apache.tajo.exception.ExceptionUtil; import org.apache.tajo.exception.ReturnStateUtil; import org.apache.tajo.exception.TajoException; import org.apache.tajo.ipc.ClientProtos; import org.apache.tajo.parser.sql.SQLLexer; import org.apache.tajo.service.ServiceTrackerFactory; import org.apache.tajo.util.*; import java.io.*; import java.lang.reflect.Constructor; import java.net.InetSocketAddress; import java.sql.ResultSet; import java.sql.SQLException; import java.util.*; import java.util.stream.Collectors; import java.util.stream.Stream; public class TajoCli implements Closeable { public static final int SHUTDOWN_HOOK_PRIORITY = 50; public static final String ERROR_PREFIX = "ERROR: "; public static final String KILL_PREFIX = "KILL: "; private final TajoConf conf; private TajoClient client; private final TajoCliContext context; // Jline and Console related things private final ConsoleReader reader; private final InputStream sin; private final PrintWriter sout; private final PrintWriter serr; private TajoFileHistory history; private final boolean reconnect; // reconnect on invalid session // Current States private String currentDatabase; private TajoCliOutputFormatter displayFormatter; private boolean wasError = false; private static final Class[] registeredCommands = { DescTableCommand.class, DescFunctionCommand.class, HelpCommand.class, ExitCommand.class, CopyrightCommand.class, VersionCommand.class, ConnectDatabaseCommand.class, ListDatabaseCommand.class, SetCommand.class, UnsetCommand.class, ExecExternalShellCommand.class, HdfsCommand.class, TajoAdminCommand.class, TajoGetConfCommand.class, TajoHAAdminCommand.class }; private AggregateCompleter cliCompleter; private ArgumentCompleter sqlCompleter; private final Map<String, TajoShellCommand> commands = new TreeMap<>(); protected static final Options options; private static final String HOME_DIR = System.getProperty("user.home"); private static final String HISTORY_FILE = ".tajo_history"; static { options = new Options(); options.addOption("c", "command", true, "execute only single command, then exit"); options.addOption("f", "file", true, "execute commands from file, then exit"); options.addOption("h", "host", true, "Tajo server host"); options.addOption("p", "port", true, "Tajo server port"); options.addOption("B", "background", false, "execute as background process"); options.addOption("conf", "conf", true, "configuration value"); options.addOption("param", "param", true, "parameter value in SQL file"); options.addOption("reconnect", "reconnect", false, "reconnect on invalid session"); options.addOption("help", "help", false, "help"); } public class TajoCliContext extends OverridableConf { public TajoCliContext(TajoConf conf) { super(conf, ConfigKey.ConfigType.SESSION); } public TajoClient getTajoClient() { return client; } public void setCurrentDatabase(String databasae) { currentDatabase = databasae; } public String getCurrentDatabase() { return currentDatabase; } public PrintWriter getOutput() { return sout; } public PrintWriter getError() { return serr; } public TajoConf getConf() { return conf; } @VisibleForTesting public String getCliSideVar(String key) { if (SessionVars.exists(key)) { ConfigKey configKey = SessionVars.get(key); return get(configKey); } else { return get(key); } } public void setCliSideVar(String key, String value) { Preconditions.checkNotNull(key); Preconditions.checkNotNull(value); boolean shouldReloadFormatter = false; if (SessionVars.exists(key)) { SessionVars configKey = SessionVars.get(key); put(configKey, value); shouldReloadFormatter = configKey.getMode() == SessionVars.VariableMode.CLI_SIDE_VAR; } else { set(key, value); // It is hard to recognize it is a client side variable. So, we always reload formatter. shouldReloadFormatter = true; } if (shouldReloadFormatter) { try { initFormatter(); } catch (Exception e) { System.err.println(ERROR_PREFIX + e.getMessage()); } } } public Map<String, TajoShellCommand> getCommands() { return commands; } } public TajoCli(TajoConf c, String[] args, @Nullable Properties clientParams, InputStream in, OutputStream out, OutputStream err) throws Exception { CommandLineParser parser = new PosixParser(); CommandLine cmd = parser.parse(options, args); this.conf = new TajoConf(c); context = new TajoCliContext(conf); this.sin = in; if (cmd.hasOption("B")) { this.reader = new ConsoleReader(sin, out, new UnsupportedTerminal()); } else { this.reader = new ConsoleReader(sin, out); } this.reader.setExpandEvents(false); this.sout = new PrintWriter(reader.getOutput()); this.serr = new PrintWriter(new OutputStreamWriter(err, "UTF-8")); initFormatter(); if (cmd.hasOption("help")) { printUsage(); System.exit(0); } String hostName = null; Integer port = null; if (cmd.hasOption("h")) { hostName = cmd.getOptionValue("h"); } if (cmd.hasOption("p")) { port = Integer.parseInt(cmd.getOptionValue("p")); } String baseDatabase = null; if (cmd.getArgList().size() > 0) { baseDatabase = (String) cmd.getArgList().get(0); } if (cmd.getOptionValues("conf") != null) { processConfVarCommand(cmd.getOptionValues("conf")); } this.reconnect = cmd.hasOption("reconnect"); // if there is no "-h" option, InetSocketAddress address = conf.getSocketAddrVar(TajoConf.ConfVars.TAJO_MASTER_CLIENT_RPC_ADDRESS, TajoConf.ConfVars.TAJO_MASTER_UMBILICAL_RPC_ADDRESS); if (hostName == null) { hostName = address.getHostName(); } if (port == null) { port = address.getPort(); } // Get connection parameters Properties defaultConnParams = CliClientParamsFactory.get(clientParams); final KeyValueSet actualConnParams = new KeyValueSet(Maps.fromProperties(defaultConnParams)); if ((hostName == null) ^ (port == null)) { System.err.println(ERROR_PREFIX + "cannot find valid Tajo server address"); throw new RuntimeException("cannot find valid Tajo server address"); } else if (hostName != null && port != null) { conf.setVar(ConfVars.TAJO_MASTER_CLIENT_RPC_ADDRESS, NetUtils.getHostPortString(hostName, port)); client = new TajoClientImpl(ServiceTrackerFactory.get(conf), baseDatabase, actualConnParams); } else if (hostName == null && port == null) { client = new TajoClientImpl(ServiceTrackerFactory.get(conf), baseDatabase, actualConnParams); } try { context.setCurrentDatabase(client.getCurrentDatabase()); initHistory(); initCommands(); reader.addCompleter(cliCompleter); reader.addCompleter(sqlCompleter); if (cmd.getOptionValues("conf") != null) { processSessionVarCommand(cmd.getOptionValues("conf")); } if (cmd.hasOption("c")) { displayFormatter.setScriptMode(); int exitCode = executeScript(cmd.getOptionValue("c")); sout.flush(); serr.flush(); System.exit(exitCode); } if (cmd.hasOption("f")) { displayFormatter.setScriptMode(); cmd.getOptionValues(""); File sqlFile = new File(cmd.getOptionValue("f")); if (sqlFile.exists()) { String script = FileUtil.readTextFile(new File(cmd.getOptionValue("f"))); script = replaceParam(script, cmd.getOptionValues("param")); int exitCode = executeScript(script); sout.flush(); serr.flush(); System.exit(exitCode); } else { System.err.println(ERROR_PREFIX + "No such a file \"" + cmd.getOptionValue("f") + "\""); System.exit(-1); } } } catch (Exception e) { System.err.println(ERROR_PREFIX + "Exception was thrown. Caused by " + e.getMessage()); if (client != null) { client.close(); } throw e; } addShutdownHook(); } private void processConfVarCommand(String[] confCommands) throws SQLException { for (String eachParam : confCommands) { String[] tokens = eachParam.split("="); if (tokens.length != 2) { continue; } if (!SessionVars.exists(tokens[0])) { conf.set(tokens[0], tokens[1]); } } } private void processSessionVarCommand(String[] confCommands) throws TajoException { for (String eachParam : confCommands) { String[] tokens = eachParam.split("="); if (tokens.length != 2) { continue; } if (SessionVars.exists(tokens[0])) { ((SetCommand) commands.get("\\set")).set(tokens[0], tokens[1]); } } } private void initFormatter() throws Exception { Class formatterClass = context.getClass(SessionVars.CLI_FORMATTER_CLASS); if (displayFormatter == null || !displayFormatter.getClass().equals(formatterClass)) { displayFormatter = (TajoCliOutputFormatter) formatterClass.newInstance(); } displayFormatter.init(context); } public TajoCliContext getContext() { return context; } protected static String replaceParam(String script, String[] params) { if (params == null || params.length == 0) { return script; } for (String eachParam : params) { String[] tokens = eachParam.split("="); if (tokens.length != 2) { continue; } script = script.replace("${" + tokens[0] + "}", tokens[1]); } return script; } private void initHistory() { try { String historyPath = HOME_DIR + File.separator + HISTORY_FILE; if ((new File(HOME_DIR)).exists()) { history = new TajoFileHistory(new File(historyPath)); history.setAutoTrim(false); history.setIgnoreDuplicates(false); reader.setHistory(history); } else { System.err.println(ERROR_PREFIX + "home directory : '" + HOME_DIR + "' does not exist."); } } catch (Exception e) { System.err.println(ERROR_PREFIX + e.getMessage()); } } private void initCommands() { List<Completer> compList = new ArrayList<>(); for (Class clazz : registeredCommands) { TajoShellCommand cmd = null; try { Constructor cons = clazz.getConstructor(new Class[] { TajoCliContext.class }); cmd = (TajoShellCommand) cons.newInstance(context); } catch (Exception e) { System.err.println(e.getMessage()); throw new RuntimeException(e.getMessage()); } // make completers for console auto-completion compList.add(cmd.getArgumentCompleter()); commands.put(cmd.getCommand(), cmd); for (String alias : cmd.getAliases()) { commands.put(alias, cmd); } } cliCompleter = new AggregateCompleter(compList); sqlCompleter = new ArgumentCompleter(new ArgumentCompleter.AbstractArgumentDelimiter() { @Override public boolean isDelimiterChar(CharSequence buf, int pos) { char c = buf.charAt(pos); return Character.isWhitespace(c) || !(Character.isLetterOrDigit(c)) && c != '_'; } }, new StringsCompleter(getKeywords())); } private Collection<String> getKeywords() { // SQL reserved keywords Stream<String> tokens = Arrays.stream(SQLLexer.tokenNames); Stream<String> rules = Arrays.stream(SQLLexer.ruleNames); List<String> keywords = Stream.concat(tokens, rules) .filter((str) -> str.matches("[A-Z_]+") && str.length() > 1).distinct().map(String::toLowerCase) .collect(Collectors.toList()); // DB and table names for (String db : client.getAllDatabaseNames()) { keywords.add(db); keywords.addAll(client.getTableList(db)); } tokens.close(); rules.close(); return keywords; } private void addShutdownHook() { ShutdownHookManager.get().addShutdownHook(new Runnable() { @Override public void run() { try { history.flush(); } catch (IOException e) { } client.close(); } }, SHUTDOWN_HOOK_PRIORITY); } private String updatePrompt(ParsingState state) throws ServiceException { if (state == ParsingState.WITHIN_QUOTE) { return "'"; } else if (state == ParsingState.TOK_START) { return context.getCurrentDatabase(); } else { return ""; } } public int runShell() throws Exception { String line; String currentPrompt = context.getCurrentDatabase(); int exitCode; ParsingState latestState = SimpleParser.START_STATE; sout.write("welcome to \n" + " _____ ___ _____ ___\n" + " /_ _/ _ |/_ _/ /\n" + " / // /_| |_/ // / /\n" + " /_//_/ /_/___/ \\__/ " + VersionInfo.getVersion() + "\n\n"); sout.write("Try \\? for help.\n"); SimpleParser parser = new SimpleParser(); try { while ((line = reader.readLine(currentPrompt + "> ")) != null) { if (line.equals("")) { continue; } wasError = false; if (line.startsWith("{")) { executeJsonQuery(line); } else { List<ParsedResult> parsedResults = parser.parseLines(line); if (latestState != ParsingState.TOK_START && parsedResults.size() > 0) { // Add multi-line statements to history in addition to individual lines. ParsedResult parsed = parsedResults.get(0); history.add(parsed.getHistoryStatement() + (parsed.getType() == StatementType.STATEMENT ? ";" : "")); } exitCode = executeParsedResults(parsedResults); latestState = parser.getState(); currentPrompt = updatePrompt(latestState); // If the ON_ERROR_STOP flag is set, Cli should stop on query failure. if (exitCode != 0 && context.getBool(SessionVars.ON_ERROR_STOP)) { return exitCode; } } } } catch (Exception e) { System.err.println(ERROR_PREFIX + "Exception was thrown. Caused by " + e.getMessage()); if (client != null) { client.close(); } throw e; } return 0; } private int executeParsedResults(Collection<ParsedResult> parsedResults) throws Exception { int exitCode = 0; for (ParsedResult parsedResult : parsedResults) { if (parsedResult.getType() == StatementType.META) { exitCode = executeMetaCommand(parsedResult.getStatement()); } else { exitCode = executeQuery(parsedResult.getStatement()); } if (exitCode != 0 && context.getBool(SessionVars.ON_ERROR_STOP)) { return exitCode; } } return exitCode; } public int executeMetaCommand(String line) { String[] metaCommands = line.split(";"); for (String metaCommand : metaCommands) { String arguments[] = metaCommand.split(" "); TajoShellCommand invoked = commands.get(arguments[0]); if (invoked == null) { printInvalidCommand(arguments[0]); wasError = true; return -1; } try { invoked.invoke(arguments); } catch (Throwable t) { onError(t); return -1; } finally { context.getError().flush(); context.getOutput().flush(); } if (wasError && context.getBool(SessionVars.ON_ERROR_STOP)) { break; } } return 0; } private void executeJsonQuery(String json) throws TajoException { long startTime = System.currentTimeMillis(); ClientProtos.SubmitQueryResponse response = client.executeQueryWithJson(json); if (ReturnStateUtil.isSuccess(response.getState())) { switch (response.getResultType()) { case FETCH: QueryId queryId = new QueryId(response.getQueryId()); waitForQueryCompleted(queryId); break; case ENCLOSED: localQueryCompleted(response, startTime); break; default: displayFormatter.printMessage(sout, "OK"); } } else { onError((Throwable) ExceptionUtil.toTajoExceptionCommon(response.getState())); } } private int executeQuery(String statement) throws SQLException { long startTime = System.currentTimeMillis(); ClientProtos.SubmitQueryResponse response = null; try { response = client.executeQuery(statement); } catch (Throwable t) { onError(t); } if (ReturnStateUtil.isSuccess(response.getState())) { switch (response.getResultType()) { case FETCH: QueryId queryId = new QueryId(response.getQueryId()); waitForQueryCompleted(queryId); break; case ENCLOSED: localQueryCompleted(response, startTime); break; default: displayFormatter.printMessage(sout, "OK"); } } else { onError((Throwable) ExceptionUtil.toTajoExceptionCommon(response.getState())); } return wasError ? -1 : 0; } private void localQueryCompleted(ClientProtos.SubmitQueryResponse response, long startTime) { ResultSet res = null; try { QueryId queryId = new QueryId(response.getQueryId()); float responseTime = ((float) (System.currentTimeMillis() - startTime) / 1000.0f); TableDesc desc = new TableDesc(response.getTableDesc()); // non-forwarded INSERT INTO query does not have any query id. // In this case, it just returns succeeded query information without printing the query results. if (response.getMaxRowNum() < 0 && queryId.equals(QueryIdFactory.NULL_QUERY_ID)) { displayFormatter.printResult(sout, sin, desc, responseTime, res); } else { int fetchRows = conf.getIntVar(TajoConf.ConfVars.$RESULT_SET_FETCH_ROWNUM); res = TajoClientUtil.createResultSet(client, response, fetchRows); displayFormatter.printResult(sout, sin, desc, responseTime, res); } } catch (Throwable t) { onError(t); } finally { if (res != null) { try { res.close(); } catch (SQLException e) { // ignore } } } } private void waitForQueryCompleted(QueryId queryId) { // if query is empty string if (queryId.equals(QueryIdFactory.NULL_QUERY_ID)) { return; } // query execute ResultSet res = null; QueryStatus status = null; try { int initRetries = 0; int progressRetries = 0; while (true) { status = client.getQueryStatus(queryId); if (TajoClientUtil.isQueryWaitingForSchedule(status.getState())) { Thread.sleep(Math.min(20 * initRetries, 1000)); initRetries++; continue; } if (TajoClientUtil.isQueryRunning(status.getState())) { displayFormatter.printProgress(serr, status); } if (TajoClientUtil.isQueryComplete(status.getState()) && status.getState() != QueryState.QUERY_KILL_WAIT) { displayFormatter.printProgress(serr, status); break; } else { Thread.sleep(Math.min(200 * progressRetries, 1000)); progressRetries += 2; } } if (status.getState() == QueryState.QUERY_ERROR || status.getState() == QueryState.QUERY_FAILED) { displayFormatter.printErrorMessage(serr, status); wasError = true; } else if (status.getState() == QueryState.QUERY_KILLED) { displayFormatter.printKilledMessage(serr, queryId); wasError = true; } else { if (status.getState() == QueryState.QUERY_SUCCEEDED) { float responseTime = ((float) (status.getFinishTime() - status.getSubmitTime()) / 1000.0f); ClientProtos.GetQueryResultResponse response = client.getResultResponse(queryId); if (status.hasResult()) { int fetchRows = conf.getIntVar(TajoConf.ConfVars.$RESULT_SET_FETCH_ROWNUM); res = TajoClientUtil.createResultSet(client, queryId, response, fetchRows); TableDesc desc = new TableDesc(response.getTableDesc()); displayFormatter.printResult(sout, sin, desc, responseTime, res); } else { TableDesc desc = new TableDesc(response.getTableDesc()); displayFormatter.printResult(sout, sin, desc, responseTime, res); } } } } catch (Throwable t) { onError(t); } finally { if (res != null) { try { res.close(); } catch (SQLException e) { } } else { if (status != null && status.getQueryId() != null) { client.closeQuery(status.getQueryId()); } } } } public int executeScript(String script) throws Exception { wasError = false; List<ParsedResult> results = SimpleParser.parseScript(script); return executeParsedResults(results); } private void printUsage() { HelpFormatter formatter = new HelpFormatter(); formatter.printUsage(this.serr, 80, "tsql [options] [database]", options); } private void printInvalidCommand(String command) { serr.println("Invalid command " + command + ". Try \\? for help."); } private void onError(Throwable t) { Preconditions.checkNotNull(t); wasError = true; displayFormatter.printErrorMessage(serr, t.getMessage()); if (reconnect && (t instanceof InvalidClientSessionException)) { try { ((SessionConnection) client).reconnect(); } catch (Exception e) { // ignore } } } @VisibleForTesting public void close() { //for testcase if (client != null) { client.close(); } if (reader != null) { reader.shutdown(); } } public static void main(String[] args) throws Exception { TajoConf conf = new TajoConf(); TajoCli shell = new TajoCli(conf, args, new Properties(), System.in, System.out, System.err); System.out.println(); System.exit(shell.runShell()); } }