Added auto completion from PostgresqlInterpreter

jdbc/src/main/java/org/apache/zeppelin/jdbc/JDBCInterpreter.java

@@ -4,9 +4,9 @@
  * 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
- *
+ * <p>
  * http://www.apache.org/licenses/LICENSE-2.0
- *
+ * <p>
  * 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
@@ -50,7 +50,7 @@ import com.google.common.collect.Sets.SetView;
 /**
  * JDBC interpreter for Zeppelin. This interpreter can also be used for accessing HAWQ,
  * GreenplumDB, MariaDB, MySQL, Postgres and Redshit.
- * 
+ * <p>
  * <ul>
  * <li>{@code default.url} - JDBC URL to connect to.</li>
  * <li>{@code default.user} - JDBC user name..</li>
@@ -58,22 +58,21 @@ import com.google.common.collect.Sets.SetView;
  * <li>{@code default.driver.name} - JDBC driver name.</li>
  * <li>{@code common.max.result} - Max number of SQL result to display.</li>
  * </ul>
- * 
+ * <p>
  * <p>
  * How to use: <br/>
  * {@code %jdbc.sql} <br/>
- * {@code 
- *  SELECT store_id, count(*) 
- *  FROM retail_demo.order_lineitems_pxf 
- *  GROUP BY store_id;
+ * {@code
+ * SELECT store_id, count(*)
+ * FROM retail_demo.order_lineitems_pxf
+ * GROUP BY store_id;
  * }
  * </p>
- * 
  */
 public class JDBCInterpreter extends Interpreter {
 
   private Logger logger = LoggerFactory.getLogger(JDBCInterpreter.class);
- 
+
   static final String COMMON_KEY = "common";
   static final String MAX_LINE_KEY = "max_count";
   static final String MAX_LINE_DEFAULT = "1000";
@@ -84,7 +83,7 @@ public class JDBCInterpreter extends Interpreter {
   static final String USER_KEY = "user";
   static final String PASSWORD_KEY = "password";
   static final String DOT = ".";
-  
+
   private static final char WHITESPACE = ' ';
   private static final char NEWLINE = '\n';
   private static final char TAB = '\t';
@@ -100,21 +99,33 @@ public class JDBCInterpreter extends Interpreter {
   static final String DEFAULT_PASSWORD = DEFAULT_KEY + DOT + PASSWORD_KEY;
 
   static final String EMPTY_COLUMN_VALUE = "";
-  
+
   private final HashMap<String, Properties> propertiesMap;
   private final Map<String, Statement> paragraphIdStatementMap;
 
   private final Map<String, ArrayList<Connection>> propertyKeyUnusedConnectionListMap;
   private final Map<String, Connection> paragraphIdConnectionMap;
 
+  private final Map<String, SqlCompleter> propertyKeySqlCompleterMap;
+
+  private static final Function<CharSequence, String> sequenceToStringTransformer =
+      new Function<CharSequence, String>() {
+        public String apply(CharSequence seq) {
+          return seq.toString();
+        }
+      };
+
+  private static final List<String> NO_COMPLETION = new ArrayList<String>();
+
   public JDBCInterpreter(Properties property) {
     super(property);
     propertiesMap = new HashMap<>();
     propertyKeyUnusedConnectionListMap = new HashMap<>();
     paragraphIdStatementMap = new HashMap<>();
     paragraphIdConnectionMap = new HashMap<>();
+    propertyKeySqlCompleterMap = new HashMap<>();
   }
-  
+
   public HashMap<String, Properties> getPropertiesMap() {
     return propertiesMap;
   }
@@ -154,9 +165,38 @@ public class JDBCInterpreter extends Interpreter {
     }
 
     logger.debug("propertiesMap: {}", propertiesMap);
+
+    Connection connection = null;
+    SqlCompleter sqlCompleter = null;
+    for (String propertyKey : propertiesMap.keySet()) {
+      try {
+        connection = getConnection(propertyKey);
+        sqlCompleter = createSqlCompleter(connection);
+      } catch (Exception e) {
+        sqlCompleter = createSqlCompleter(null);
+      }
+      propertyKeySqlCompleterMap.put(propertyKey, sqlCompleter);
+    }
   }
-  
-  public Connection getConnection(String propertyKey)  throws ClassNotFoundException, SQLException {
+
+  private SqlCompleter createSqlCompleter(Connection jdbcConnection) {
+
+    SqlCompleter completer = null;
+    try {
+      Set<String> keywordsCompletions = SqlCompleter.getSqlKeywordsCompletions(jdbcConnection);
+      Set<String> dataModelCompletions =
+          SqlCompleter.getDataModelMetadataCompletions(jdbcConnection);
+      SetView<String> allCompletions = Sets.union(keywordsCompletions, dataModelCompletions);
+      completer = new SqlCompleter(allCompletions, dataModelCompletions);
+
+    } catch (IOException | SQLException e) {
+      logger.error("Cannot create SQL completer", e);
+    }
+
+    return completer;
+  }
+
+  public Connection getConnection(String propertyKey) throws ClassNotFoundException, SQLException {
     Connection connection = null;
     if (propertyKey == null || propertiesMap.get(propertyKey) == null) {
       return null;
@@ -180,7 +220,7 @@ public class JDBCInterpreter extends Interpreter {
     }
     return connection;
   }
-  
+
   public Statement getStatement(String propertyKey, String paragraphId)
       throws SQLException, ClassNotFoundException {
     Connection connection;
@@ -189,7 +229,7 @@ public class JDBCInterpreter extends Interpreter {
     } else {
       connection = getConnection(propertyKey);
     }
-    
+
     if (connection == null) {
       return null;
     }
@@ -213,7 +253,7 @@ public class JDBCInterpreter extends Interpreter {
       return false;
     }
   }
-  
+
   @Override
   public void close() {
 
@@ -241,13 +281,13 @@ public class JDBCInterpreter extends Interpreter {
 
   private InterpreterResult executeSql(String propertyKey, String sql,
       InterpreterContext interpreterContext) {
-    
+
     String paragraphId = interpreterContext.getParagraphId();
 
     try {
 
       Statement statement = getStatement(propertyKey, paragraphId);
-      
+
       if (statement == null) {
         return new InterpreterResult(Code.ERROR, "Prefix not found.");
       }
@@ -346,7 +386,7 @@ public class JDBCInterpreter extends Interpreter {
     if (null != propertyKey && !propertyKey.equals(DEFAULT_KEY)) {
       cmd = cmd.substring(propertyKey.length() + 2);
     }
-    
+
     cmd = cmd.trim();
 
     logger.info("PropertyKey: {}, SQL command: '{}'", propertyKey, cmd);
@@ -382,7 +422,7 @@ public class JDBCInterpreter extends Interpreter {
       return DEFAULT_KEY;
     }
   }
-  
+
   @Override
   public FormType getFormType() {
     return FormType.SIMPLE;
@@ -401,7 +441,13 @@ public class JDBCInterpreter extends Interpreter {
 
   @Override
   public List<String> completion(String buf, int cursor) {
-    return null;
+    List<CharSequence> candidates = new ArrayList<CharSequence>();
+    SqlCompleter sqlCompleter = propertyKeySqlCompleterMap.get(getPropertyKey(buf));
+    if (sqlCompleter != null && sqlCompleter.complete(buf, cursor, candidates) >= 0) {
+      return Lists.transform(candidates, sequenceToStringTransformer);
+    } else {
+      return NO_COMPLETION;
+    }
   }
 
   public int getMaxResult() {

jdbc/src/main/java/org/apache/zeppelin/jdbc/SqlCompleter.java

@@ -0,0 +1,250 @@
+/**
+ * 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
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * 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.zeppelin.jdbc;
+
+/*
+ * This source file is based on code taken from SQLLine 1.0.2 See SQLLine notice in LICENSE
+ */
+
+import com.google.common.base.Joiner;
+import com.google.common.collect.Sets;
+import com.google.common.collect.Sets.SetView;
+import jline.console.completer.ArgumentCompleter.ArgumentList;
+import jline.console.completer.ArgumentCompleter.WhitespaceArgumentDelimiter;
+import jline.console.completer.StringsCompleter;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.io.InputStreamReader;
+import java.sql.Connection;
+import java.sql.DatabaseMetaData;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.util.*;
+import java.util.regex.Pattern;
+
+import static org.apache.commons.lang.StringUtils.isBlank;
+
+/**
+ * SQL auto complete functionality for the PostgreSqlInterpreter.
+ */
+public class SqlCompleter extends StringsCompleter {
+
+  private static Logger logger = LoggerFactory.getLogger(SqlCompleter.class);
+
+  /**
+   * Delimiter that can split SQL statement in keyword list
+   */
+  private WhitespaceArgumentDelimiter sqlDelimiter = new WhitespaceArgumentDelimiter() {
+
+    private Pattern pattern = Pattern.compile("[\\.:;,]");
+
+    @Override
+    public boolean isDelimiterChar(CharSequence buffer, int pos) {
+      return pattern.matcher("" + buffer.charAt(pos)).matches()
+          || super.isDelimiterChar(buffer, pos);
+    }
+  };
+
+  private Set<String> modelCompletions = new HashSet<String>();
+
+  public SqlCompleter(Set<String> allCompletions, Set<String> dataModelCompletions) {
+    super(allCompletions);
+    this.modelCompletions = dataModelCompletions;
+  }
+
+  @Override
+  public int complete(String buffer, int cursor, List<CharSequence> candidates) {
+
+    if (isBlank(buffer) || (cursor > buffer.length() + 1)) {
+      return -1;
+    }
+
+    // The delimiter breaks the buffer into separate words (arguments), separated by the
+    // white spaces.
+    ArgumentList argumentList = sqlDelimiter.delimit(buffer, cursor);
+    String argument = argumentList.getCursorArgument();
+    // cursor in the selected argument
+    int argumentPosition = argumentList.getArgumentPosition();
+
+    if (isBlank(argument)) {
+      int argumentsCount = argumentList.getArguments().length;
+      if (argumentsCount <= 0 || ((buffer.length() + 2) < cursor)
+          || sqlDelimiter.isDelimiterChar(buffer, cursor - 2)) {
+        return -1;
+      }
+      argument = argumentList.getArguments()[argumentsCount - 1];
+      argumentPosition = argument.length();
+    }
+
+    int complete = super.complete(argument, argumentPosition, candidates);
+
+    logger.debug("complete:" + complete + ", size:" + candidates.size());
+
+    return complete;
+  }
+
+  public void updateDataModelMetaData(Connection connection) {
+
+    try {
+      Set<String> newModelCompletions = getDataModelMetadataCompletions(connection);
+      logger.debug("New model metadata is:" + Joiner.on(',').join(newModelCompletions));
+
+      // Sets.difference(set1, set2) - returned set contains all elements that are contained by set1
+      // and not contained by set2. set2 may also contain elements not present in set1; these are
+      // simply ignored.
+      SetView<String> removedCompletions = Sets.difference(modelCompletions, newModelCompletions);
+      logger.debug("Removed Model Completions: " + Joiner.on(',').join(removedCompletions));
+      this.getStrings().removeAll(removedCompletions);
+
+      SetView<String> newCompletions = Sets.difference(newModelCompletions, modelCompletions);
+      logger.debug("New Completions: " + Joiner.on(',').join(newCompletions));
+      this.getStrings().addAll(newCompletions);
+
+      modelCompletions = newModelCompletions;
+
+    } catch (SQLException e) {
+      logger.error("Failed to update the metadata conmpletions", e);
+    }
+  }
+
+  public static Set<String> getSqlKeywordsCompletions(Connection connection) throws IOException,
+      SQLException {
+
+    // Add the default SQL completions
+    String keywords =
+        new BufferedReader(new InputStreamReader(
+            SqlCompleter.class.getResourceAsStream("/ansi.sql.keywords"))).readLine();
+
+    Set<String> completions = new TreeSet<>();
+
+    if (null != connection) {
+      DatabaseMetaData metaData = connection.getMetaData();
+
+      // Add the driver specific SQL completions
+      String driverSpecificKeywords =
+          "/" + metaData.getDriverName().replace(" ", "-").toLowerCase() + "-sql.keywords";
+
+      logger.info("JDBC DriverName:" + driverSpecificKeywords);
+
+      if (SqlCompleter.class.getResource(driverSpecificKeywords) != null) {
+        String driverKeywords =
+            new BufferedReader(new InputStreamReader(
+                SqlCompleter.class.getResourceAsStream(driverSpecificKeywords))).readLine();
+        keywords += "," + driverKeywords.toUpperCase();
+      }
+
+
+      // Add the keywords from the current JDBC connection
+      try {
+        keywords += "," + metaData.getSQLKeywords();
+      } catch (Exception e) {
+        logger.debug("fail to get SQL key words from database metadata: " + e, e);
+      }
+      try {
+        keywords += "," + metaData.getStringFunctions();
+      } catch (Exception e) {
+        logger.debug("fail to get string function names from database metadata: " + e, e);
+      }
+      try {
+        keywords += "," + metaData.getNumericFunctions();
+      } catch (Exception e) {
+        logger.debug("fail to get numeric function names from database metadata: " + e, e);
+      }
+      try {
+        keywords += "," + metaData.getSystemFunctions();
+      } catch (Exception e) {
+        logger.debug("fail to get system function names from database metadata: " + e, e);
+      }
+      try {
+        keywords += "," + metaData.getTimeDateFunctions();
+      } catch (Exception e) {
+        logger.debug("fail to get time date function names from database metadata: " + e, e);
+      }
+
+      // Also allow lower-case versions of all the keywords
+      keywords += "," + keywords.toLowerCase();
+
+    }
+
+    StringTokenizer tok = new StringTokenizer(keywords, ", ");
+    while (tok.hasMoreTokens()) {
+      completions.add(tok.nextToken());
+    }
+
+    return completions;
+  }
+
+  public static Set<String> getDataModelMetadataCompletions(Connection connection)
+      throws SQLException {
+    Set<String> completions = new TreeSet<>();
+    if (null != connection) {
+      getColumnNames(connection.getMetaData(), completions);
+      getSchemaNames(connection.getMetaData(), completions);
+    }
+    return completions;
+  }
+
+  private static void getColumnNames(DatabaseMetaData meta, Set<String> names) throws SQLException {
+
+    try {
+      ResultSet columns = meta.getColumns(meta.getConnection().getCatalog(), null, "%", "%");
+      try {
+
+        while (columns.next()) {
+          // Add the following strings: (1) column name, (2) table name
+          String name = columns.getString("TABLE_NAME");
+          if (!isBlank(name)) {
+            names.add(name);
+            names.add(columns.getString("COLUMN_NAME"));
+            // names.add(columns.getString("TABLE_NAME") + "." + columns.getString("COLUMN_NAME"));
+          }
+        }
+      } finally {
+        columns.close();
+      }
+
+      logger.debug(Joiner.on(',').join(names));
+    } catch (Throwable t) {
+      logger.error("Failed to retrieve the column name", t);
+    }
+  }
+
+  private static void getSchemaNames(DatabaseMetaData meta, Set<String> names) throws SQLException {
+
+    try {
+      ResultSet schemas = meta.getSchemas();
+      try {
+        while (schemas.next()) {
+          String schemaName = schemas.getString("TABLE_SCHEM");
+          if (!isBlank(schemaName)) {
+            names.add(schemaName + ".");
+          }
+        }
+      } finally {
+        schemas.close();
+      }
+    } catch (Throwable t) {
+      logger.error("Failed to retrieve the column name", t);
+    }
+  }
+
+  // test purpose only
+  WhitespaceArgumentDelimiter getSqlDelimiter() {
+    return this.sqlDelimiter;
+  }
+}

jdbc/src/main/resources/ansi.sql.keywords

@@ -0,0 +1 @@
+ABSOLUTE,ACTION,ADD,ALL,ALLOCATE,ALTER,AND,ANY,ARE,AS,ASC,ASSERTION,AT,AUTHORIZATION,AVG,BEGIN,BETWEEN,BIT,BIT_LENGTH,BOTH,BY,CASCADE,CASCADED,CASE,CAST,CATALOG,CHAR,CHARACTER,CHAR_LENGTH,CHARACTER_LENGTH,CHECK,CLOSE,CLUSTER,COALESCE,COLLATE,COLLATION,COLUMN,COMMIT,CONNECT,CONNECTION,CONSTRAINT,CONSTRAINTS,CONTINUE,CONVERT,CORRESPONDING,COUNT,CREATE,CROSS,CURRENT,CURRENT_DATE,CURRENT_TIME,CURRENT_TIMESTAMP,CURRENT_USER,CURSOR,DATE,DAY,DEALLOCATE,DEC,DECIMAL,DECLARE,DEFAULT,DEFERRABLE,DEFERRED,DELETE,DESC,DESCRIBE,DESCRIPTOR,DIAGNOSTICS,DISCONNECT,DISTINCT,DOMAIN,DOUBLE,DROP,ELSE,END,END-EXEC,ESCAPE,EXCEPT,EXCEPTION,EXEC,EXECUTE,EXISTS,EXTERNAL,EXTRACT,FALSE,FETCH,FIRST,FLOAT,FOR,FOREIGN,FOUND,FROM,FULL,GET,GLOBAL,GO,GOTO,GRANT,GROUP,HAVING,HOUR,IDENTITY,IMMEDIATE,IN,INDICATOR,INITIALLY,INNER,INPUT,INSENSITIVE,INSERT,INT,INTEGER,INTERSECT,INTERVAL,INTO,IS,ISOLATION,JOIN,KEY,LANGUAGE,LAST,LEADING,LEFT,LEVEL,LIKE,LOCAL,LOWER,MATCH,MAX,MIN,MINUTE,MODULE,MONTH,NAMES,NATIONAL,NATURAL,NCHAR,NEXT,NO,NOT,NULL,NULLIF,NUMERIC,OCTET_LENGTH,OF,ON,ONLY,OPEN,OPTION,OR,ORDER,OUTER,OUTPUT,OVERLAPS,OVERWRITE,PAD,PARTIAL,PARTITION,POSITION,PRECISION,PREPARE,PRESERVE,PRIMARY,PRIOR,PRIVILEGES,PROCEDURE,PUBLIC,READ,REAL,REFERENCES,RELATIVE,RESTRICT,REVOKE,RIGHT,ROLLBACK,ROWS,SCHEMA,SCROLL,SECOND,SECTION,SELECT,SESSION,SESSION_USER,SET,SIZE,SMALLINT,SOME,SPACE,SQL,SQLCODE,SQLERROR,SQLSTATE,SUBSTRING,SUM,SYSTEM_USER,TABLE,TEMPORARY,THEN,TIME,TIMESTAMP,TIMEZONE_HOUR,TIMEZONE_MINUTE,TO,TRAILING,TRANSACTION,TRANSLATE,TRANSLATION,TRIM,TRUE,UNION,UNIQUE,UNKNOWN,UPDATE,UPPER,USAGE,USER,USING,VALUE,VALUES,VARCHAR,VARYING,VIEW,WHEN,WHENEVER,WHERE,WITH,WORK,WRITE,YEAR,ZONE,ADA,C,CATALOG_NAME,CHARACTER_SET_CATALOG,CHARACTER_SET_NAME,CHARACTER_SET_SCHEMA,CLASS_ORIGIN,COBOL,COLLATION_CATALOG,COLLATION_NAME,COLLATION_SCHEMA,COLUMN_NAME,COMMAND_FUNCTION,COMMITTED,CONDITION_NUMBER,CONNECTION_NAME,CONSTRAINT_CATALOG,CONSTRAINT_NAME,CONSTRAINT_SCHEMA,CURSOR_NAME,DATA,DATETIME_INTERVAL_CODE,DATETIME_INTERVAL_PRECISION,DYNAMIC_FUNCTION,FORTRAN,LENGTH,MESSAGE_LENGTH,MESSAGE_OCTET_LENGTH,MESSAGE_TEXT,MORE,MUMPS,NAME,NULLABLE,NUMBER,PASCAL,PLI,REPEATABLE,RETURNED_LENGTH,RETURNED_OCTET_LENGTH,RETURNED_SQLSTATE,ROW_COUNT,SCALE,SCHEMA_NAME,SERIALIZABLE,SERVER_NAME,SUBCLASS_ORIGIN,TABLE_NAME,TYPE,UNCOMMITTED,UNNAMED,LIMIT