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.zeppelin.postgresql; import static org.apache.commons.lang.StringUtils.containsIgnoreCase; import java.io.IOException; import java.sql.Connection; import java.sql.DriverManager; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.Statement; import java.util.ArrayList; import java.util.List; import java.util.Properties; import java.util.Set; import org.apache.zeppelin.interpreter.Interpreter; import org.apache.zeppelin.interpreter.InterpreterContext; import org.apache.zeppelin.interpreter.InterpreterPropertyBuilder; import org.apache.zeppelin.interpreter.InterpreterResult; import org.apache.zeppelin.interpreter.InterpreterResult.Code; import org.apache.zeppelin.scheduler.Scheduler; import org.apache.zeppelin.scheduler.SchedulerFactory; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import com.google.common.base.Function; import com.google.common.collect.Lists; import com.google.common.collect.Sets; import com.google.common.collect.Sets.SetView; /** * PostgreSQL interpreter for Zeppelin. This interpreter can also be used for accessing HAWQ and * GreenplumDB. * * <ul> * <li>{@code postgresql.url} - JDBC URL to connect to.</li> * <li>{@code postgresql.user} - JDBC user name..</li> * <li>{@code postgresql.password} - JDBC password..</li> * <li>{@code postgresql.driver.name} - JDBC driver name.</li> * <li>{@code postgresql.max.result} - Max number of SQL result to display.</li> * </ul> * * <p> * How to use: <br/> * {@code %psql.sql} <br/> * {@code * SELECT store_id, count(*) * FROM retail_demo.order_lineitems_pxf * GROUP BY store_id; * } * </p> * * For SQL auto-completion use the (Ctrl+.) shortcut. */ public class PostgreSqlInterpreter extends Interpreter { private Logger logger = LoggerFactory.getLogger(PostgreSqlInterpreter.class); private static final char WhITESPACE = ' '; private static final char NEWLINE = '\n'; private static final char TAB = '\t'; private static final String TABLE_MAGIC_TAG = "%table "; private static final String EXPLAIN_PREDICATE = "EXPLAIN "; private static final String UPDATE_COUNT_HEADER = "Update Count"; static final String DEFAULT_JDBC_URL = "jdbc:postgresql://localhost:5432/"; static final String DEFAULT_JDBC_USER_PASSWORD = ""; static final String DEFAULT_JDBC_USER_NAME = "gpadmin"; static final String DEFAULT_JDBC_DRIVER_NAME = "org.postgresql.Driver"; static final String DEFAULT_MAX_RESULT = "1000"; static final String POSTGRESQL_SERVER_URL = "postgresql.url"; static final String POSTGRESQL_SERVER_USER = "postgresql.user"; static final String POSTGRESQL_SERVER_PASSWORD = "postgresql.password"; static final String POSTGRESQL_SERVER_DRIVER_NAME = "postgresql.driver.name"; static final String POSTGRESQL_SERVER_MAX_RESULT = "postgresql.max.result"; static final String EMPTY_COLUMN_VALUE = ""; static { Interpreter.register("sql", "psql", PostgreSqlInterpreter.class.getName(), new InterpreterPropertyBuilder() .add(POSTGRESQL_SERVER_URL, DEFAULT_JDBC_URL, "The URL for PostgreSQL.") .add(POSTGRESQL_SERVER_USER, DEFAULT_JDBC_USER_NAME, "The PostgreSQL user name") .add(POSTGRESQL_SERVER_PASSWORD, DEFAULT_JDBC_USER_PASSWORD, "The PostgreSQL user password") .add(POSTGRESQL_SERVER_DRIVER_NAME, DEFAULT_JDBC_DRIVER_NAME, "JDBC Driver Name") .add(POSTGRESQL_SERVER_MAX_RESULT, DEFAULT_MAX_RESULT, "Max number of SQL result to display.") .build()); } private Connection jdbcConnection; private Statement currentStatement; private Exception exceptionOnConnect; private int maxResult; private SqlCompleter sqlCompleter; 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 PostgreSqlInterpreter(Properties property) { super(property); } @Override public void open() { logger.info("Open psql connection!"); // Ensure that no previous connections are left open. close(); try { String driverName = getProperty(POSTGRESQL_SERVER_DRIVER_NAME); String url = getProperty(POSTGRESQL_SERVER_URL); String user = getProperty(POSTGRESQL_SERVER_USER); String password = getProperty(POSTGRESQL_SERVER_PASSWORD); maxResult = Integer.valueOf(getProperty(POSTGRESQL_SERVER_MAX_RESULT)); Class.forName(driverName); jdbcConnection = DriverManager.getConnection(url, user, password); sqlCompleter = createSqlCompleter(jdbcConnection); exceptionOnConnect = null; logger.info("Successfully created psql connection"); } catch (ClassNotFoundException | SQLException e) { logger.error("Cannot open connection", e); exceptionOnConnect = e; close(); } } 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; } @Override public void close() { logger.info("Close psql connection!"); try { if (getJdbcConnection() != null) { getJdbcConnection().close(); } } catch (SQLException e) { logger.error("Cannot close connection", e); } finally { exceptionOnConnect = null; } } private InterpreterResult executeSql(String sql) { try { if (exceptionOnConnect != null) { return new InterpreterResult(Code.ERROR, exceptionOnConnect.getMessage()); } currentStatement = getJdbcConnection().createStatement(); currentStatement.setMaxRows(maxResult); StringBuilder msg = null; boolean isTableType = false; if (containsIgnoreCase(sql, EXPLAIN_PREDICATE)) { msg = new StringBuilder(); } else { msg = new StringBuilder(TABLE_MAGIC_TAG); isTableType = true; } ResultSet resultSet = null; try { boolean isResultSetAvailable = currentStatement.execute(sql); if (isResultSetAvailable) { resultSet = currentStatement.getResultSet(); ResultSetMetaData md = resultSet.getMetaData(); for (int i = 1; i < md.getColumnCount() + 1; i++) { if (i > 1) { msg.append(TAB); } msg.append(replaceReservedChars(isTableType, md.getColumnName(i))); } msg.append(NEWLINE); int displayRowCount = 0; while (resultSet.next() && displayRowCount < getMaxResult()) { for (int i = 1; i < md.getColumnCount() + 1; i++) { msg.append(replaceReservedChars(isTableType, resultSet.getString(i))); if (i != md.getColumnCount()) { msg.append(TAB); } } msg.append(NEWLINE); displayRowCount++; } } else { // Response contains either an update count or there are no results. int updateCount = currentStatement.getUpdateCount(); msg.append(UPDATE_COUNT_HEADER).append(NEWLINE); msg.append(updateCount).append(NEWLINE); // In case of update event (e.g. isResultSetAvailable = false) update the completion // meta-data. if (sqlCompleter != null) { sqlCompleter.updateDataModelMetaData(getJdbcConnection()); } } } finally { try { if (resultSet != null) { resultSet.close(); } currentStatement.close(); } finally { currentStatement = null; } } return new InterpreterResult(Code.SUCCESS, msg.toString()); } catch (SQLException ex) { logger.error("Cannot run " + sql, ex); return new InterpreterResult(Code.ERROR, ex.getMessage()); } } /** * For %table response replace Tab and Newline characters from the content. */ private String replaceReservedChars(boolean isTableResponseType, String str) { if (str == null) { return EMPTY_COLUMN_VALUE; } return (!isTableResponseType) ? str : str.replace(TAB, WhITESPACE).replace(NEWLINE, WhITESPACE); } @Override public InterpreterResult interpret(String cmd, InterpreterContext contextInterpreter) { logger.info("Run SQL command '{}'", cmd); return executeSql(cmd); } @Override public void cancel(InterpreterContext context) { logger.info("Cancel current query statement."); if (currentStatement != null) { try { currentStatement.cancel(); } catch (SQLException ex) { } finally { currentStatement = null; } } } @Override public FormType getFormType() { return FormType.SIMPLE; } @Override public int getProgress(InterpreterContext context) { return 0; } @Override public Scheduler getScheduler() { return SchedulerFactory.singleton() .createOrGetFIFOScheduler(PostgreSqlInterpreter.class.getName() + this.hashCode()); } @Override public List<String> completion(String buf, int cursor) { List<CharSequence> candidates = new ArrayList<CharSequence>(); if (sqlCompleter != null && sqlCompleter.complete(buf, cursor, candidates) >= 0) { return Lists.transform(candidates, sequenceToStringTransformer); } else { return NO_COMPLETION; } } public int getMaxResult() { return maxResult; } // Test only method protected Connection getJdbcConnection() { return jdbcConnection; } }