exec/java-exec/src/main/java/org/apache/drill/exec/client/DrillSqlLineApplication.java

/*
 * 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.drill.exec.client;

import com.typesafe.config.Config;
import com.typesafe.config.ConfigFactory;
import org.apache.drill.common.scanner.ClassPathScanner;
import org.apache.drill.common.util.DrillVersionInfo;
import com.google.common.annotations.VisibleForTesting;
import org.jline.reader.impl.completer.StringsCompleter;
import sqlline.Application;
import sqlline.CommandHandler;
import sqlline.ConnectionMetadata;
import sqlline.OutputFormat;
import sqlline.PromptHandler;
import sqlline.ReflectiveCommandHandler;
import sqlline.SqlLine;
import sqlline.SqlLineOpts;

import java.net.URL;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * <p>Drill SqlLine application configuration.</p>
 *
 * <p>Customizes SqlLine for Drill, i.e. overrides application info message,
 * known drivers, connection url examples, removes non applicable commands, sets SqlLine properties.</p>
 *
 * <p>Uses {@link #DRILL_SQLLINE_CONF} as base configuration, allows to override it using {@link #DRILL_SQLLINE_OVERRIDE_CONF}.
 * If configuration files are missing in the classpath, issues warning and proceeds with default SqlLine configuration.</p>
 */
public class DrillSqlLineApplication extends Application {

  private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(DrillSqlLineApplication.class);

  private static final String DRILL_SQLLINE_CONF = "drill-sqlline.conf";
  private static final String DRILL_SQLLINE_OVERRIDE_CONF = "drill-sqlline-override.conf";

  private static final String INFO_MESSAGE_TEMPLATE_CONF = "drill.sqlline.info_message_template";
  private static final String QUOTES_CONF = "drill.sqlline.quotes";
  private static final String DRIVERS_CONF = "drill.sqlline.drivers";
  private static final String CONNECTION_URL_EXAMPLES_CONF = "drill.sqlline.connection_url_examples";
  private static final String COMMANDS_TO_EXCLUDE_CONF = "drill.sqlline.commands.exclude";
  private static final String OPTS_CONF = "drill.sqlline.opts";
  private static final String PROMPT_WITH_SCHEMA = "drill.sqlline.prompt.with_schema";

  private final Config config;

  public DrillSqlLineApplication() {
    this(DRILL_SQLLINE_CONF, DRILL_SQLLINE_OVERRIDE_CONF);
  }

  @VisibleForTesting
  public DrillSqlLineApplication(String configName, String overrideConfigName) {
    this.config = overrideConfig(overrideConfigName, loadConfig(configName));
    if (config.isEmpty()) {
      logger.warn("Was unable to find / load [{}]. Will use default SqlLine configuration.", configName);
    }
  }

  public Config getConfig() {
    return config;
  }

  @Override
  public String getInfoMessage() {
    if (config.hasPath(INFO_MESSAGE_TEMPLATE_CONF)) {
      String quote = "";
      if (config.hasPath(QUOTES_CONF)) {
        List<String> quotes = config.getStringList(QUOTES_CONF);
        quote = quotes.get(new Random().nextInt(quotes.size()));
      }
      return String.format(config.getString(INFO_MESSAGE_TEMPLATE_CONF), getVersion(), quote);
    }

    return super.getInfoMessage();
  }

  @Override
  public String getVersion() {
    return DrillVersionInfo.getVersion();
  }

  @Override
  public List<String> allowedDrivers() {
    if (config.hasPath(DRIVERS_CONF)) {
      return config.getStringList(DRIVERS_CONF);
    }
    return super.allowedDrivers();
  }

  @Override
  public Map<String, OutputFormat> getOutputFormats(SqlLine sqlLine) {
    return sqlLine.getOutputFormats();
  }

  @Override
  public Collection<String> getConnectionUrlExamples() {
    if (config.hasPath(CONNECTION_URL_EXAMPLES_CONF)) {
      return config.getStringList(CONNECTION_URL_EXAMPLES_CONF);
    }
    return super.getConnectionUrlExamples();
  }

  @Override
  public Collection<CommandHandler> getCommandHandlers(SqlLine sqlLine) {
    List<String> commandsToExclude = new ArrayList<>();

    // exclude connect command and then add it back to ensure connection url examples are updated
    boolean reloadConnect = config.hasPath(CONNECTION_URL_EXAMPLES_CONF);
    if (reloadConnect) {
      commandsToExclude.add("connect");
    }

    if (config.hasPath(COMMANDS_TO_EXCLUDE_CONF)) {
      commandsToExclude.addAll(config.getStringList(COMMANDS_TO_EXCLUDE_CONF));
    }

    if (commandsToExclude.isEmpty()) {
      return sqlLine.getCommandHandlers();
    }

    List<CommandHandler> commandHandlers = sqlLine.getCommandHandlers().stream()
        .filter(c -> c.getNames().stream()
            .noneMatch(commandsToExclude::contains))
        .collect(Collectors.toList());

    if (reloadConnect) {
      commandHandlers.add(new ReflectiveCommandHandler(sqlLine,
          new StringsCompleter(getConnectionUrlExamples()), "connect", "open"));
    }

    return commandHandlers;
  }

  @Override
  public SqlLineOpts getOpts(SqlLine sqlLine) {
    SqlLineOpts opts = sqlLine.getOpts();
    if (config.hasPath(OPTS_CONF)) {
      Config optsConfig = config.getConfig(OPTS_CONF);
      optsConfig.entrySet().forEach(
          e -> {
            String key = e.getKey();
            String value = String.valueOf(e.getValue().unwrapped());
            if (!opts.set(key, value, true)) {
              logger.warn("Unable to set SqlLine property [{}] to [{}].", key, value);
            }
          }
      );
    }
    return opts;
  }

  @Override
  public PromptHandler getPromptHandler(SqlLine sqlLine) {
    if (config.hasPath(PROMPT_WITH_SCHEMA) && config.getBoolean(PROMPT_WITH_SCHEMA)) {
      return new PromptHandler(sqlLine) {
        @Override
        protected String getDefaultPrompt(int connectionIndex, String url, String defaultPrompt) {
          StringBuilder builder = new StringBuilder();
          builder.append("apache drill");

          ConnectionMetadata meta = sqlLine.getConnectionMetadata();

          String currentSchema = meta.getCurrentSchema();
          if (currentSchema != null) {
            builder.append(" (").append(currentSchema).append(")");
          }
          return builder.append("> ").toString();
        }
      };
    }
    return super.getPromptHandler(sqlLine);
  }

  private Config loadConfig(String configName) {
    Set<URL> urls = ClassPathScanner.forResource(configName, false);
    if (urls.size() != 1) {
      if (logger.isDebugEnabled()) {
        urls.forEach(
            u -> logger.debug("Found duplicating [{}]: [{}].", configName, u.getPath())
        );
      }
      return ConfigFactory.empty();
    }

    URL url = urls.iterator().next();
    try {
      logger.debug("Parsing [{}] for the url: [{}].", configName, url.getPath());
      return ConfigFactory.parseURL(url);
    } catch (Exception e) {
      logger.warn("Was unable to parse [{}].", url.getPath(), e);
      return ConfigFactory.empty();
    }
  }

  private Config overrideConfig(String configName, Config config) {
    Config overrideConfig = loadConfig(configName);
    return overrideConfig.withFallback(config).resolve();
  }

}