/**
* 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.
*/
/*
* This source file is based on code taken from SQLLine 1.0.2
* See SQLLine notice in LICENSE
*/
package org.apache.hive.beeline;
import java.io.IOException;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.Driver;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.Collection;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Set;
import java.util.TreeSet;
import org.apache.hive.jdbc.HiveConnection;
import jline.console.completer.ArgumentCompleter;
import jline.console.completer.Completer;
class DatabaseConnection {
private static final String HIVE_VAR_PREFIX = "hivevar:";
private static final String HIVE_CONF_PREFIX = "hiveconf:";
private final BeeLine beeLine;
private Connection connection;
private DatabaseMetaData meta;
private final String driver;
private final String url;
private final Properties info;
private Schema schema = null;
private Completer sqlCompleter = null;
public boolean isClosed() {
return (null == connection);
}
public DatabaseConnection(BeeLine beeLine, String driver, String url,
Properties info) throws SQLException {
this.beeLine = beeLine;
this.driver = driver;
this.url = url;
this.info = info;
}
@Override
public String toString() {
return getUrl() + "";
}
void setCompletions(boolean skipmeta) throws SQLException, IOException {
final String extraNameCharacters =
getDatabaseMetaData() == null || getDatabaseMetaData().getExtraNameCharacters() == null ? ""
: getDatabaseMetaData().getExtraNameCharacters();
// setup the completer for the database
sqlCompleter = new ArgumentCompleter(
new ArgumentCompleter.AbstractArgumentDelimiter() {
// delimiters for SQL statements are any
// non-letter-or-number characters, except
// underscore and characters that are specified
// by the database to be valid name identifiers.
@Override
public boolean isDelimiterChar(CharSequence buffer, int pos) {
char c = buffer.charAt(pos);
if (Character.isWhitespace(c)) {
return true;
}
return !(Character.isLetterOrDigit(c))
&& c != '_'
&& extraNameCharacters.indexOf(c) == -1;
}
},
new SQLCompleter(SQLCompleter.getSQLCompleters(beeLine, skipmeta)));
// not all argument elements need to hold true
((ArgumentCompleter) sqlCompleter).setStrict(false);
}
/**
* Connection to the specified data source.
*/
boolean connect() throws SQLException {
try {
if (driver != null && driver.length() != 0) {
Class.forName(driver);
}
} catch (ClassNotFoundException cnfe) {
return beeLine.error(cnfe);
}
boolean isDriverRegistered = false;
try {
isDriverRegistered = DriverManager.getDriver(getUrl()) != null;
} catch (Exception e) {
}
try {
close();
} catch (Exception e) {
return beeLine.error(e);
}
Map<String, String> hiveVars = beeLine.getOpts().getHiveVariables();
if (hiveVars != null){
for (Map.Entry<String, String> var : hiveVars.entrySet()) {
info.put(HIVE_VAR_PREFIX + var.getKey(), var.getValue());
}
}
Map<String, String> hiveConfVars = beeLine.getOpts().getHiveConfVariables();
if (hiveConfVars != null){
for (Map.Entry<String, String> var : hiveConfVars.entrySet()) {
info.put(HIVE_CONF_PREFIX + var.getKey(), var.getValue());
}
}
if (isDriverRegistered) {
// if the driver registered in the driver manager, get the connection via the driver manager
setConnection(DriverManager.getConnection(getUrl(), info));
} else {
beeLine.debug("Use the driver from local added jar file.");
setConnection(getConnectionFromLocalDriver(getUrl(), info));
}
setDatabaseMetaData(getConnection().getMetaData());
try {
beeLine.info(beeLine.loc("connected", new Object[] {
getDatabaseMetaData().getDatabaseProductName(),
getDatabaseMetaData().getDatabaseProductVersion()}));
} catch (Exception e) {
beeLine.handleException(e);
}
try {
beeLine.info(beeLine.loc("driver", new Object[] {
getDatabaseMetaData().getDriverName(),
getDatabaseMetaData().getDriverVersion()}));
} catch (Exception e) {
beeLine.handleException(e);
}
try {
getConnection().setAutoCommit(beeLine.getOpts().getAutoCommit());
// TODO: Setting autocommit should not generate an exception as long as it is set to false
// beeLine.autocommitStatus(getConnection());
} catch (Exception e) {
beeLine.handleException(e);
}
try {
beeLine.getCommands().isolation("isolation: " + beeLine.getOpts().getIsolation());
} catch (Exception e) {
beeLine.handleException(e);
}
return true;
}
public Connection getConnectionFromLocalDriver(String url, Properties properties) {
Collection<Driver> drivers = beeLine.getDrivers();
for (Driver d : drivers) {
try {
if (d.acceptsURL(url) && beeLine.isSupportedLocalDriver(d)) {
String clazzName = d.getClass().getName();
beeLine.debug("Driver name is " + clazzName);
Driver driver =
(Driver) Class.forName(clazzName, true, Thread.currentThread().getContextClassLoader())
.newInstance();
return driver.connect(url, properties);
}
} catch (Exception e) {
beeLine.error("Fail to connect with a local driver due to the exception:" + e);
beeLine.error(e);
}
}
return null;
}
public Connection getConnection() throws SQLException {
if (connection != null) {
return connection;
}
connect();
return connection;
}
public Connection getCurrentConnection() {
return connection;
}
public void reconnect() throws Exception {
close();
getConnection();
}
public void close() {
try {
try {
if (connection != null && !connection.isClosed()) {
beeLine.output(beeLine.loc("closing", connection));
connection.close();
}
} catch (Exception e) {
beeLine.handleException(e);
}
} finally {
setConnection(null);
setDatabaseMetaData(null);
}
}
public String[] getTableNames(boolean force) {
Schema.Table[] t = getSchema().getTables();
Set<String> names = new TreeSet<String>();
for (int i = 0; t != null && i < t.length; i++) {
names.add(t[i].getName());
}
return names.toArray(new String[names.size()]);
}
Schema getSchema() {
if (schema == null) {
schema = new Schema();
}
return schema;
}
void setConnection(Connection connection) {
this.connection = connection;
}
DatabaseMetaData getDatabaseMetaData() {
return meta;
}
void setDatabaseMetaData(DatabaseMetaData meta) {
this.meta = meta;
}
String getUrl() {
return url;
}
public String getConnectedUrl() {
if (connection instanceof HiveConnection) {
return ((HiveConnection) connection).getConnectedUrl();
}
return getUrl();
}
Completer getSQLCompleter() {
return sqlCompleter;
}
class Schema {
private Table[] tables = null;
Table[] getTables() {
if (tables != null) {
return tables;
}
List<Table> tnames = new LinkedList<Table>();
try {
ResultSet rs = getDatabaseMetaData().getTables(getConnection().getCatalog(),
null, "%", new String[] {"TABLE"});
try {
while (rs.next()) {
tnames.add(new Table(rs.getString("TABLE_NAME")));
}
} finally {
try {
rs.close();
} catch (Exception e) {
}
}
} catch (Throwable t) {
}
return tables = tnames.toArray(new Table[0]);
}
Table getTable(String name) {
Table[] t = getTables();
for (int i = 0; t != null && i < t.length; i++) {
if (name.equalsIgnoreCase(t[i].getName())) {
return t[i];
}
}
return null;
}
class Table {
final String name;
Column[] columns;
public Table(String name) {
this.name = name;
}
public String getName() {
return name;
}
class Column {
final String name;
boolean isPrimaryKey;
public Column(String name) {
this.name = name;
}
}
}
}
}