/**
* 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;
}
}