/**
* 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.hive.service.cli.thrift;
import java.lang.reflect.InvocationHandler;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.lang.reflect.UndeclaredThrowableException;
import java.net.SocketException;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import javax.security.sasl.SaslException;
import org.apache.hadoop.hive.conf.HiveConf;
import org.apache.hive.service.auth.HiveAuthFactory;
import org.apache.hive.service.auth.PlainSaslHelper;
import org.apache.hive.service.cli.CLIServiceClient;
import org.apache.hive.service.cli.FetchOrientation;
import org.apache.hive.service.cli.FetchType;
import org.apache.hive.service.cli.GetInfoType;
import org.apache.hive.service.cli.GetInfoValue;
import org.apache.hive.service.cli.HiveSQLException;
import org.apache.hive.service.cli.ICLIService;
import org.apache.hive.service.cli.OperationHandle;
import org.apache.hive.service.cli.OperationStatus;
import org.apache.hive.service.cli.RowSet;
import org.apache.hive.service.cli.SessionHandle;
import org.apache.hive.service.cli.TableSchema;
import org.apache.hive.service.rpc.thrift.TCLIService;
import org.apache.thrift.TApplicationException;
import org.apache.thrift.protocol.TBinaryProtocol;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.protocol.TProtocolException;
import org.apache.thrift.transport.TSocket;
import org.apache.thrift.transport.TTransport;
import org.apache.thrift.transport.TTransportException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/**
* RetryingThriftCLIServiceClient. Creates a proxy for a CLIServiceClient
* implementation and retries calls to it on failure.
*/
public class RetryingThriftCLIServiceClient implements InvocationHandler {
public static final Logger LOG = LoggerFactory.getLogger(RetryingThriftCLIServiceClient.class);
private ThriftCLIServiceClient base;
private final int retryLimit;
private final int retryDelaySeconds;
private HiveConf conf;
private TTransport transport;
public static class CLIServiceClientWrapper extends CLIServiceClient {
private final ICLIService cliService;
private TTransport tTransport;
public CLIServiceClientWrapper(ICLIService icliService, TTransport tTransport, HiveConf conf) {
super(conf);
cliService = icliService;
this.tTransport = tTransport;
}
@Override
public SessionHandle openSession(String username, String password) throws HiveSQLException {
return cliService.openSession(username, password, Collections.<String, String>emptyMap());
}
@Override
public String getDelegationToken(SessionHandle sessionHandle, HiveAuthFactory authFactory, String owner,
String renewer) throws HiveSQLException {
return cliService.getDelegationToken(sessionHandle, authFactory, owner, renewer);
}
@Override
public void cancelDelegationToken(SessionHandle sessionHandle, HiveAuthFactory authFactory,
String tokenStr) throws HiveSQLException {
cliService.cancelDelegationToken(sessionHandle, authFactory, tokenStr);
}
@Override
public void renewDelegationToken(SessionHandle sessionHandle, HiveAuthFactory authFactory,
String tokenStr) throws HiveSQLException {
cliService.renewDelegationToken(sessionHandle, authFactory, tokenStr);
}
@Override
public SessionHandle openSession(String username, String password, Map<String, String> configuration)
throws HiveSQLException {
return cliService.openSession(username, password, configuration);
}
@Override
public SessionHandle openSessionWithImpersonation(String username,
String password,
Map<String, String> configuration,
String delegationToken) throws HiveSQLException {
return cliService.openSessionWithImpersonation(username, password, configuration, delegationToken);
}
@Override
public void closeSession(SessionHandle sessionHandle) throws HiveSQLException {
cliService.closeSession(sessionHandle);
}
@Override
public GetInfoValue getInfo(SessionHandle sessionHandle, GetInfoType getInfoType) throws HiveSQLException {
return cliService.getInfo(sessionHandle, getInfoType);
}
@Override
public OperationHandle executeStatement(SessionHandle sessionHandle, String statement,
Map<String, String> confOverlay) throws HiveSQLException {
return cliService.executeStatement(sessionHandle, statement, confOverlay);
}
@Override
public OperationHandle executeStatement(SessionHandle sessionHandle, String statement,
Map<String, String> confOverlay, long queryTimeout) throws HiveSQLException {
return cliService.executeStatement(sessionHandle, statement, confOverlay, queryTimeout);
}
@Override
public OperationHandle executeStatementAsync(SessionHandle sessionHandle, String statement,
Map<String, String> confOverlay) throws HiveSQLException {
return cliService.executeStatementAsync(sessionHandle, statement, confOverlay);
}
@Override
public OperationHandle executeStatementAsync(SessionHandle sessionHandle, String statement,
Map<String, String> confOverlay, long queryTimeout) throws HiveSQLException {
return cliService.executeStatementAsync(sessionHandle, statement, confOverlay, queryTimeout);
}
@Override
public OperationHandle getTypeInfo(SessionHandle sessionHandle) throws HiveSQLException {
return cliService.getTypeInfo(sessionHandle);
}
@Override
public OperationHandle getCatalogs(SessionHandle sessionHandle) throws HiveSQLException {
return cliService.getCatalogs(sessionHandle);
}
@Override
public OperationHandle getSchemas(SessionHandle sessionHandle, String catalogName, String schemaName)
throws HiveSQLException {
return cliService.getSchemas(sessionHandle, catalogName, schemaName);
}
@Override
public OperationHandle getTables(SessionHandle sessionHandle, String catalogName, String schemaName,
String tableName, List<String> tableTypes) throws HiveSQLException {
return cliService.getTables(sessionHandle, catalogName, schemaName, tableName, tableTypes);
}
@Override
public OperationHandle getTableTypes(SessionHandle sessionHandle) throws HiveSQLException {
return null;
}
@Override
public OperationHandle getColumns(SessionHandle sessionHandle, String catalogName, String schemaName,
String tableName, String columnName) throws HiveSQLException {
return cliService.getColumns(sessionHandle, catalogName, schemaName, tableName, columnName);
}
@Override
public OperationHandle getFunctions(SessionHandle sessionHandle, String catalogName, String schemaName,
String functionName) throws HiveSQLException {
return cliService.getFunctions(sessionHandle, catalogName, schemaName, functionName);
}
@Override
public OperationStatus getOperationStatus(OperationHandle opHandle, boolean getProgressUpdate) throws HiveSQLException {
return cliService.getOperationStatus(opHandle, getProgressUpdate);
}
@Override
public void cancelOperation(OperationHandle opHandle) throws HiveSQLException {
cliService.cancelOperation(opHandle);
}
@Override
public void closeOperation(OperationHandle opHandle) throws HiveSQLException {
cliService.closeOperation(opHandle);
}
@Override
public TableSchema getResultSetMetadata(OperationHandle opHandle) throws HiveSQLException {
return cliService.getResultSetMetadata(opHandle);
}
@Override
public RowSet fetchResults(OperationHandle opHandle, FetchOrientation orientation, long maxRows,
FetchType fetchType) throws HiveSQLException {
return cliService.fetchResults(opHandle, orientation, maxRows, fetchType);
}
public void closeTransport() {
tTransport.close();
}
@Override
public OperationHandle getPrimaryKeys(SessionHandle sessionHandle,
String catalog, String schema, String table)
throws HiveSQLException {
return cliService.getPrimaryKeys(sessionHandle, catalog, schema, table);
}
@Override
public OperationHandle getCrossReference(SessionHandle sessionHandle,
String primaryCatalog, String primarySchema, String primaryTable,
String foreignCatalog, String foreignSchema, String foreignTable)
throws HiveSQLException {
return cliService.getCrossReference(sessionHandle, primaryCatalog, primarySchema,
primaryTable, foreignCatalog, foreignSchema, foreignTable);
}
}
protected RetryingThriftCLIServiceClient(HiveConf conf) {
this.conf = conf;
retryLimit = conf.getIntVar(HiveConf.ConfVars.HIVE_SERVER2_THRIFT_CLIENT_RETRY_LIMIT);
retryDelaySeconds = (int) conf.getTimeVar(HiveConf.ConfVars.HIVE_SERVER2_THRIFT_CLIENT_RETRY_DELAY_SECONDS,
TimeUnit.SECONDS);
}
public static CLIServiceClientWrapper newRetryingCLIServiceClient(HiveConf conf) throws HiveSQLException {
RetryingThriftCLIServiceClient retryClient = new RetryingThriftCLIServiceClient(conf);
TTransport tTransport = retryClient
.connectWithRetry(conf.getIntVar(HiveConf.ConfVars.HIVE_SERVER2_THRIFT_CLIENT_CONNECTION_RETRY_LIMIT));
ICLIService cliService =
(ICLIService) Proxy.newProxyInstance(RetryingThriftCLIServiceClient.class.getClassLoader(),
CLIServiceClient.class.getInterfaces(), retryClient);
return new CLIServiceClientWrapper(cliService, tTransport, conf);
}
protected TTransport connectWithRetry(int retries) throws HiveSQLException {
TTransportException exception = null;
for (int i = 0 ; i < retries; i++) {
try {
return connect(conf);
} catch (TTransportException e) {
exception = e;
LOG.warn("Connection attempt " + i, e);
}
try {
Thread.sleep(retryDelaySeconds * 1000);
} catch (InterruptedException e) {
LOG.warn("Interrupted", e);
}
}
throw new HiveSQLException("Unable to connect after " + retries + " retries", exception);
}
protected synchronized TTransport connect(HiveConf conf) throws HiveSQLException, TTransportException {
if (transport != null && transport.isOpen()) {
transport.close();
}
String host = conf.getVar(HiveConf.ConfVars.HIVE_SERVER2_THRIFT_BIND_HOST);
int port = conf.getIntVar(HiveConf.ConfVars.HIVE_SERVER2_THRIFT_PORT);
LOG.info("Connecting to " + host + ":" + port);
transport = new TSocket(host, port);
((TSocket) transport).setTimeout((int) conf.getTimeVar(HiveConf.ConfVars.SERVER_READ_SOCKET_TIMEOUT,
TimeUnit.SECONDS) * 1000);
try {
((TSocket) transport).getSocket().setKeepAlive(conf.getBoolVar(HiveConf.ConfVars.SERVER_TCP_KEEP_ALIVE));
} catch (SocketException e) {
LOG.error("Error setting keep alive to " + conf.getBoolVar(HiveConf.ConfVars.SERVER_TCP_KEEP_ALIVE), e);
}
String userName = conf.getVar(HiveConf.ConfVars.HIVE_SERVER2_THRIFT_CLIENT_USER);
String passwd = conf.getVar(HiveConf.ConfVars.HIVE_SERVER2_THRIFT_CLIENT_PASSWORD);
try {
transport = PlainSaslHelper.getPlainTransport(userName, passwd, transport);
} catch (SaslException e) {
LOG.error("Error creating plain SASL transport", e);
}
TProtocol protocol = new TBinaryProtocol(transport);
transport.open();
base = new ThriftCLIServiceClient(new TCLIService.Client(protocol), conf);
LOG.info("Connected!");
return transport;
}
protected class InvocationResult {
final boolean success;
final Object result;
final Throwable exception;
InvocationResult(boolean success, Object result, Throwable exception) {
this.success = success;
this.result = result;
this.exception = exception;
}
}
protected InvocationResult invokeInternal(Method method, Object[] args) throws Throwable {
InvocationResult result;
try {
Object methodResult = method.invoke(base, args);
result = new InvocationResult(true, methodResult, null);
} catch (UndeclaredThrowableException e) {
throw e.getCause();
} catch (InvocationTargetException e) {
if (e.getCause() instanceof HiveSQLException) {
HiveSQLException hiveExc = (HiveSQLException) e.getCause();
Throwable cause = hiveExc.getCause();
if ((cause instanceof TApplicationException) ||
(cause instanceof TProtocolException) ||
(cause instanceof TTransportException)) {
result = new InvocationResult(false, null, hiveExc);
} else {
throw hiveExc;
}
} else {
throw e.getCause();
}
}
return result;
}
@Override
public Object invoke(Object o, Method method, Object[] args) throws Throwable {
int attempts = 0;
while (true) {
attempts++;
InvocationResult invokeResult = invokeInternal(method, args);
if (invokeResult.success) {
return invokeResult.result;
}
// Error because of thrift client, we have to recreate base object
connectWithRetry(conf.getIntVar(HiveConf.ConfVars.HIVE_SERVER2_THRIFT_CLIENT_CONNECTION_RETRY_LIMIT));
if (attempts >= retryLimit) {
LOG.error(method.getName() + " failed after " + attempts + " retries.", invokeResult.exception);
throw invokeResult.exception;
}
LOG.warn("Last call ThriftCLIServiceClient." + method.getName() + " failed, attempts = " + attempts,
invokeResult.exception);
Thread.sleep(retryDelaySeconds * 1000);
}
}
public int getRetryLimit() {
return retryLimit;
}
public int getRetryDelaySeconds() {
return retryDelaySeconds;
}
}