thrift-client.h

Go to the documentation of this file.

// Copyright 2012 Cloudera Inc.
//
// Licensed 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.


#ifndef IMPALA_RPC_THRIFT_CLIENT_H
#define IMPALA_RPC_THRIFT_CLIENT_H

#include <ostream>
#include <sstream>
#include <boost/shared_ptr.hpp>
#include <common/status.h>
#include <thrift/Thrift.h>
#include <thrift/transport/TSocket.h>
#include <thrift/transport/TSSLSocket.h>
#include <thrift/transport/TBufferTransports.h>
#include <thrift/protocol/TBinaryProtocol.h>
#include <sstream>
#include <gflags/gflags.h>

#include "transport/TSaslClientTransport.h"
#include "transport/TSasl.h"
#include "rpc/authentication.h"
#include "rpc/thrift-server.h"
#include "gen-cpp/Types_types.h"

DECLARE_string(principal);
DECLARE_string(hostname);

namespace impala {
class ThriftClientImpl {
 public:
  ~ThriftClientImpl() {
    Close();
  }

  const TNetworkAddress& address() const { return address_; }

  Status Open();

  Status OpenWithRetry(uint32_t num_retries, uint64_t wait_ms);

  void Close();

  void setRecvTimeout(int32_t ms) { socket_->setRecvTimeout(ms); }

  void setSendTimeout(int32_t ms) { socket_->setSendTimeout(ms); }

 protected:
  ThriftClientImpl(const std::string& ipaddress, int port, bool ssl)
      : address_(MakeNetworkAddress(ipaddress, port)), ssl_(ssl) {
    socket_create_status_ = CreateSocket();
  }

  Status CreateSocket();

  TNetworkAddress address_;

  bool ssl_;

  Status socket_create_status_;

  boost::shared_ptr<sasl::TSasl> sasl_client_;

  boost::shared_ptr<apache::thrift::transport::TSocket> socket_;
  boost::shared_ptr<apache::thrift::transport::TTransport> transport_;
  boost::shared_ptr<apache::thrift::protocol::TBinaryProtocol> protocol_;
};


template <class InterfaceType>
class ThriftClient : public ThriftClientImpl {
 public:
  ThriftClient(const std::string& ipaddress, int port,
      const std::string& service_name = "", AuthProvider* auth_provider = NULL,
      bool ssl = false);

  InterfaceType* iface() { return iface_.get(); }

 private:
  boost::shared_ptr<InterfaceType> iface_;

  AuthProvider* auth_provider_;
};

template <class InterfaceType>
ThriftClient<InterfaceType>::ThriftClient(const std::string& ipaddress, int port,
    const std::string& service_name,
    AuthProvider* auth_provider, bool ssl)
    : ThriftClientImpl(ipaddress, port, ssl),
      iface_(new InterfaceType(protocol_)),
      auth_provider_(auth_provider) {

  transport_.reset(new apache::thrift::transport::TBufferedTransport(socket_));

  if (auth_provider_ == NULL) {
    auth_provider_ = AuthManager::GetInstance()->GetInternalAuthProvider();
  }

  auth_provider_->WrapClientTransport(address_.hostname, transport_, service_name,
      &transport_);

  protocol_.reset(new apache::thrift::protocol::TBinaryProtocol(transport_));
  iface_.reset(new InterfaceType(protocol_));
}

}
#endif