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.ambari.view.hive.client; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.hive.shims.ShimLoader; import org.apache.hadoop.security.UserGroupInformation; import org.apache.hive.service.auth.HiveAuthFactory; import org.apache.hive.service.auth.KerberosSaslHelper; import org.apache.hive.service.auth.PlainSaslHelper; import org.apache.hive.service.auth.SaslQOP; import org.apache.hive.service.cli.thrift.*; import org.apache.thrift.TException; import org.apache.thrift.protocol.TBinaryProtocol; import org.apache.thrift.transport.TTransport; import org.apache.thrift.transport.TTransportException; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import javax.security.sasl.Sasl; import javax.security.sasl.SaslException; import java.io.IOException; import java.util.HashMap; import java.util.List; import java.util.Map; /** * Holds session */ public class Connection { private final static Logger LOG = LoggerFactory.getLogger(Connection.class); private String host; private int port; private Map<String, String> authParams; private TCLIService.Client client = null; private TSessionHandle sessHandle = null; private TProtocolVersion protocol = null; private TTransport transport; private DDLDelegator ddl; private String username; public Connection(String host, int port, Map<String, String> authParams, String username) throws HiveClientException { this.host = host; this.port = port; this.authParams = authParams; this.username = username; openConnection(); ddl = new DDLDelegator(this); } public DDLDelegator ddl() { return ddl; } public synchronized void openConnection() throws HiveClientException { try { transport = getTransport(); transport.open(); client = new TCLIService.Client(new TBinaryProtocol(transport)); } catch (TTransportException e) { throw new HiveClientException( "Could not establish connecton to " + host + ":" + port + ": " + e.toString(), e); } LOG.info("Hive connection opened"); } /** * Based on JDBC implementation of HiveConnection.createBinaryTransport * * @return transport * @throws HiveClientException */ protected TTransport getTransport() throws HiveClientException, TTransportException { TTransport transport; boolean assumeSubject = Utils.HiveAuthenticationParams.AUTH_KERBEROS_AUTH_TYPE_FROM_SUBJECT .equals(authParams.get(Utils.HiveAuthenticationParams.AUTH_KERBEROS_AUTH_TYPE)); try { if (!Utils.HiveAuthenticationParams.AUTH_SIMPLE .equalsIgnoreCase(authParams.get(Utils.HiveAuthenticationParams.AUTH_TYPE))) { // If Kerberos Map<String, String> saslProps = new HashMap<String, String>(); SaslQOP saslQOP = SaslQOP.AUTH; if (authParams.containsKey(Utils.HiveAuthenticationParams.AUTH_PRINCIPAL)) { if (authParams.containsKey(Utils.HiveAuthenticationParams.AUTH_QOP)) { try { saslQOP = SaslQOP.fromString(authParams.get(Utils.HiveAuthenticationParams.AUTH_QOP)); } catch (IllegalArgumentException e) { throw new HiveClientException("Invalid " + Utils.HiveAuthenticationParams.AUTH_QOP + " parameter. " + e.getMessage(), e); } } saslProps.put(Sasl.QOP, saslQOP.toString()); saslProps.put(Sasl.SERVER_AUTH, "true"); Configuration conf = new Configuration(); conf.set("hadoop.security.authentication", "kerberos"); UserGroupInformation.setConfiguration(conf); transport = KerberosSaslHelper.getKerberosTransport( authParams.get(Utils.HiveAuthenticationParams.AUTH_PRINCIPAL), host, HiveAuthFactory.getSocketTransport(host, port, 10000), saslProps, assumeSubject); } else { // If there's a delegation token available then use token based connection String tokenStr = getClientDelegationToken(authParams); if (tokenStr != null) { transport = KerberosSaslHelper.getTokenTransport(tokenStr, host, HiveAuthFactory.getSocketTransport(host, port, 10000), saslProps); } else { // we are using PLAIN Sasl connection with user/password String userName = getAuthParamDefault(Utils.HiveAuthenticationParams.AUTH_USER, getUsername()); String passwd = getAuthParamDefault(Utils.HiveAuthenticationParams.AUTH_PASSWD, Utils.HiveAuthenticationParams.ANONYMOUS_USER); // Note: Thrift returns an SSL socket that is already bound to the specified host:port // Therefore an open called on this would be a no-op later // Hence, any TTransportException related to connecting with the peer are thrown here. // Bubbling them up the call hierarchy so that a retry can happen in openTransport, // if dynamic service discovery is configured. if (isSslConnection()) { // get SSL socket String sslTrustStore = authParams.get(Utils.HiveAuthenticationParams.SSL_TRUST_STORE); String sslTrustStorePassword = authParams .get(Utils.HiveAuthenticationParams.SSL_TRUST_STORE_PASSWORD); if (sslTrustStore == null || sslTrustStore.isEmpty()) { transport = HiveAuthFactory.getSSLSocket(host, port, 10000); } else { transport = HiveAuthFactory.getSSLSocket(host, port, 10000, sslTrustStore, sslTrustStorePassword); } } else { // get non-SSL socket transport transport = HiveAuthFactory.getSocketTransport(host, port, 10000); } // Overlay the SASL transport on top of the base socket transport (SSL or non-SSL) transport = PlainSaslHelper.getPlainTransport(userName, passwd, transport); } } } else { //NOSASL return HiveAuthFactory.getSocketTransport(host, port, 10000); } } catch (SaslException e) { throw new HiveClientException("Could not create secure connection to " + host + ": " + e.getMessage(), e); } return transport; } private boolean isSslConnection() { return "true".equalsIgnoreCase(authParams.get(Utils.HiveAuthenticationParams.USE_SSL)); } // Lookup the delegation token. First in the connection URL, then Configuration private String getClientDelegationToken(Map<String, String> jdbcConnConf) throws HiveClientException { String tokenStr = null; if (Utils.HiveAuthenticationParams.AUTH_TOKEN .equalsIgnoreCase(jdbcConnConf.get(Utils.HiveAuthenticationParams.AUTH_TYPE))) { // check delegation token in job conf if any try { tokenStr = ShimLoader.getHadoopShims().getTokenStrForm(HiveAuthFactory.HS2_CLIENT_TOKEN); } catch (IOException e) { throw new HiveClientException("Error reading token ", e); } } return tokenStr; } private String getAuthParamDefault(String key, String defaultValue) { if (authParams.containsKey(key)) return authParams.get(key); return defaultValue; } private synchronized void openSession() throws HiveClientException { //It's possible to set proxy user configuration here TOpenSessionResp openResp = new HiveCall<TOpenSessionResp>(this) { @Override public TOpenSessionResp body() throws HiveClientException { TOpenSessionReq openReq = new TOpenSessionReq(); try { return client.OpenSession(openReq); } catch (TException e) { throw new HiveClientException("Unable to open Hive session", e); } } }.call(); Utils.verifySuccess(openResp.getStatus(), "Unable to open Hive session"); protocol = openResp.getServerProtocolVersion(); sessHandle = openResp.getSessionHandle(); LOG.info("Hive session opened"); } private synchronized void closeSession() throws HiveClientException { if (sessHandle == null) return; TCloseSessionReq closeReq = new TCloseSessionReq(sessHandle); //It's possible to set proxy user configuration here TCloseSessionResp closeResp = null; try { closeResp = client.CloseSession(closeReq); Utils.verifySuccess(closeResp.getStatus(), "Unable to close Hive session"); } catch (TException e) { throw new HiveClientException("Unable to close Hive session", e); } sessHandle = null; protocol = null; LOG.info("Hive session closed"); } public synchronized void closeConnection() throws HiveClientException { if (client == null) return; try { closeSession(); } catch (HiveClientException e) { LOG.error("Unable to close Hive session: " + e.getMessage()); } finally { transport.close(); transport = null; client = null; sessHandle = null; protocol = null; } LOG.info("Connection to Hive closed"); } /** * Execute query * @param cmd query * @param async wait till query finish? * @return handle of operation * @throws HiveClientException */ public TOperationHandle execute(final String cmd, final boolean async) throws HiveClientException { TOperationHandle handle = null; for (final String oneCmd : cmd.split(";")) { TExecuteStatementResp execResp = new HiveCall<TExecuteStatementResp>(this) { @Override public TExecuteStatementResp body() throws HiveClientException { TExecuteStatementReq execReq = null; openSession(); execReq = new TExecuteStatementReq(getSessHandle(), oneCmd); execReq.setRunAsync(async); execReq.setConfOverlay(new HashMap<String, String>()); //maybe it's hive configuration? use it, Luke! try { return client.ExecuteStatement(execReq); } catch (TException e) { throw new HiveClientException("Unable to submit statement " + cmd, e); } } }.call(); Utils.verifySuccess(execResp.getStatus(), "Unable to submit statement " + cmd); //TODO: check if status have results handle = execResp.getOperationHandle(); } if (handle == null) { throw new HiveClientException("Empty command given", null); } return handle; } public TOperationHandle executeAsync(String cmd) throws HiveClientException { return execute(cmd, true); } public TOperationHandle executeSync(String cmd) throws HiveClientException { return execute(cmd, false); } public String getLogs(TOperationHandle handle) { LogsCursor results = new LogsCursor(this, handle); results.reset(); // we have to read from FIRST line, to get // logs from beginning on every call this function List<String> logLineList = results.getValuesInColumn(0); StringBuilder log = new StringBuilder(); for (String line : logLineList) { log.append(line); log.append('\n'); } return log.toString(); } public Cursor getResults(TOperationHandle handle) { Cursor cursor = new Cursor(this, handle); cursor.reset(); // we have to read from FIRST line, to get // logs from beginning on every call this function return cursor; } /** * Retrieve status of operation * @param operationHandle handle * @return thrift status response object * @throws HiveClientException */ public TGetOperationStatusResp getOperationStatus(final TOperationHandle operationHandle) throws HiveClientException { return new HiveCall<TGetOperationStatusResp>(this) { @Override public TGetOperationStatusResp body() throws HiveClientException { TGetOperationStatusReq statusReq = new TGetOperationStatusReq(operationHandle); try { return client.GetOperationStatus(statusReq); } catch (TException e) { throw new HiveClientException("Unable to fetch operation status", e); } } }.call(); // transportLock.lock(); // try { // return client.GetOperationStatus(statusReq); // } catch (TException e) { // throw new HiveClientException("Unable to fetch operation status", e); // } finally { // transportLock.unlock(); // } } /** * Cancel operation * @param operationHandle operation handle */ public void cancelOperation(final TOperationHandle operationHandle) throws HiveClientException { TCancelOperationResp cancelResp = new HiveCall<TCancelOperationResp>(this) { @Override public TCancelOperationResp body() throws HiveClientException { TCancelOperationReq cancelReq = new TCancelOperationReq(operationHandle); try { return client.CancelOperation(cancelReq); } catch (TException e) { throw new HiveClientException("Unable to cancel operation", null); } } }.call(); Utils.verifySuccess(cancelResp.getStatus(), "Unable to cancel operation"); } public int getPort() { return port; } public void setPort(int port) { this.port = port; } public String getHost() { return host; } public void setHost(String host) { this.host = host; } public TSessionHandle getSessHandle() throws HiveClientException { if (sessHandle == null) openSession(); return sessHandle; } public void setSessHandle(TSessionHandle sessHandle) { this.sessHandle = sessHandle; } public TCLIService.Client getClient() { return client; } public void setClient(TCLIService.Client client) { this.client = client; } public TProtocolVersion getProtocol() { return protocol; } public void setProtocol(TProtocolVersion protocol) { this.protocol = protocol; } public Map<String, String> getAuthParams() { return authParams; } public void setAuthParams(Map<String, String> authParams) { this.authParams = authParams; } public String getUsername() { return username; } public void setUsername(String username) { this.username = username; } }