<commit_before>#include <stdlib.h>
#include <exception>
#include <sstream>

#include <boost/format.hpp>
#include "Rinternals.h"

#include "Bee/Client.h"

using namespace apache::thrift;
using namespace apache::hive::service::cli::thrift;

typedef struct {
  unsigned char bytes[4];
} Quad;

static Quad big_endian(int i) {
  return Quad {
    .bytes = {
      static_cast<unsigned char>((i & 0xff000000) >> 24),
      static_cast<unsigned char>((i & 0x00ff0000) >> 16),
      static_cast<unsigned char>((i & 0x0000ff00) >>  8),
      static_cast<unsigned char>((i & 0x000000ff) >>  0)
    }
  };
}

Bee::Client::Client(const std::string &hostname, int port, const std::string &user, const std::string &pass) :
  user_(user),
  pass_(pass),
  socket_(new transport::TSocket(hostname, port)),
  transport_(new transport::TFramedTransport(socket_)),
  protocol_(new protocol::TBinaryProtocol(transport_)),
  client_(new apache::hive::service::cli::thrift::TCLIServiceClient(protocol_)),
  session_handle_(),
  operation_handle_() {
    this->transport_->open();
    this->perform_sasl_handshake();
    this->open_session();
}

void Bee::Client::disconnect() {
  if (this->transport_->isOpen()) {
    this->transport_->close();
  }
}

void Bee::Client::execute(const std::string &sql) {
  this->ensure_connected();

  TExecuteStatementReq request;
  TExecuteStatementResp response;
  request.sessionHandle = this->session_handle_;
  request.statement = sql;

  this->client_->ExecuteStatement(response, request);
  this->operation_handle_ = response.operationHandle;
}

Rcpp::List Bee::Client::fetch(int num_rows) {
  this->ensure_connected();

  if (num_rows < 1) {
    throw std::runtime_error("num_rows must be > 1");
  }

  TFetchResultsReq result_request;
  TFetchResultsResp result_response;
  TGetResultSetMetadataReq metadata_request;
  TGetResultSetMetadataResp metadata_response;

  result_request.operationHandle = this->operation_handle_;
  result_request.orientation = TFetchOrientation::FETCH_NEXT;
  result_request.maxRows = num_rows;
  this->client_->FetchResults(result_response, result_request);

  // TFetchResultsResp.hasMoreRows is currently broken, so just keep going until
  // we fall short.
  this->has_more_rows_ = result_response.results.rows.size() >= num_rows;

  metadata_request.operationHandle = this->operation_handle_;
  this->client_->GetResultSetMetadata(metadata_response, metadata_request);

  return build_data_frame(metadata_response.schema, result_response.results);
}

bool Bee::Client::has_more_rows() const {
  return this->has_more_rows_;
}

std::string Bee::Client::inspect() {
  std::ostringstream out;
  out << "Bee::Client[" <<
    "hostname=" << this->socket_->getHost() <<
    ", port=" << this->socket_->getPort() <<
    ", user=" << this->user_ <<
    ", pass=" << this->pass_ <<
    (this->transport_->isOpen() ? ", connected" : "") <<
    (this->has_session_handle() ? ", in session" : "") <<
    (this->has_operation_handle() ? ", in operation" : "") <<
    "]";
  return out.str();
}

void Bee::Client::write_frame(const std::string &bytes) {
  Quad length = big_endian(bytes.size());
  this->socket_->write(length.bytes, 4);
  this->socket_->write(reinterpret_cast<const unsigned char *>(bytes.c_str()), bytes.size());
}

boost::shared_ptr<std::string> Bee::Client::read_frame() {
  unsigned int length = 0;
  unsigned char length_buffer[5];
  this->socket_->read(length_buffer, 4);
  length_buffer[4] = '\0';
  for (int i = 0; i < 4; ++i) {
    length += static_cast<int>(length_buffer[1 + i]) << i;
  }

  if (length == 0) {
    return boost::shared_ptr<std::string>(new std::string(""));
  } else {
    char message[length];
    this->socket_->read(reinterpret_cast<unsigned char *>(message), length);
    return boost::shared_ptr<std::string>(new std::string(message, length));
  }
}

void Bee::Client::perform_sasl_handshake() {
  // START PLAIN
  {
    unsigned char code = 0x01;
    this->socket_->write(&code, 1);
    this->write_frame("PLAIN");
  }

  // OK PLAIN user pass
  {
    unsigned char code = 0x02;
    this->socket_->write(&code, 1);
    std::ostringstream frame("[PLAIN]\x00");
    frame << user_ << '\0' << pass_;
    this->write_frame(frame.str());
  }

  this->socket_->flush();

  unsigned char status;
  this->socket_->read(&status, 1);

  switch (status) {
  case 3: // bad
  case 4: // error
      throw std::runtime_error(*this->read_frame());

  case 5: // complete
      this->read_frame();  // challenge -- ignored
      break;

  case 2: // ok
    {
      boost::shared_ptr<std::string> message(this->read_frame());
      throw std::runtime_error(std::string("ERROR: auth failed: ") + *message);
      break;
    }

  default:
    throw std::runtime_error("unexpected response from sasl handshake");
  }
}

void Bee::Client::open_session() {
  TOpenSessionReq request;
  TOpenSessionResp response;
  request.client_protocol = TProtocolVersion::HIVE_CLI_SERVICE_PROTOCOL_V3;

  client_->OpenSession(response, request);
  this->session_handle_ = response.sessionHandle;
}

void Bee::Client::ensure_connected() const {
  if (!this->transport_->isOpen())
    throw std::runtime_error("not connected");
}

bool Bee::Client::has_session_handle() const {
  return !this->session_handle_.sessionId.guid.empty();
}

bool Bee::Client::has_operation_handle() const {
  return !this->operation_handle_.operationId.guid.empty();<commit_msg>Clean up some noise.<commit_after>#include <stdlib.h>
#include <exception>
#include <sstream>

#include <boost/format.hpp>
#include "Rinternals.h"

#include "Bee/Client.h"

using namespace apache::thrift;
using namespace apache::hive::service::cli::thrift;

typedef struct {
  unsigned char bytes[4];
} Quad;

static Quad big_endian(int i) {
  return Quad {
    .bytes = {
      static_cast<unsigned char>((i & 0xff000000) >> 24),
      static_cast<unsigned char>((i & 0x00ff0000) >> 16),
      static_cast<unsigned char>((i & 0x0000ff00) >>  8),
      static_cast<unsigned char>((i & 0x000000ff) >>  0)
    }
  };
}

Bee::Client::Client(const std::string &hostname, int port, const std::string &user, const std::string &pass) :
  user_(user),
  pass_(<PASSWORD>),
  socket_(new transport::TSocket(hostname, port)),
  transport_(new transport::TFramedTransport(socket_)),
  protocol_(new protocol::TBinaryProtocol(transport_)),
  client_(new apache::hive::service::cli::thrift::TCLIServiceClient(protocol_)),
  session_handle_(),
  operation_handle_() {
    transport_->open();
    perform_sasl_handshake();
    open_session();
}

void Bee::Client::disconnect() {
  if (transport_->isOpen()) {
    transport_->close();
  }
}

void Bee::Client::execute(const std::string &sql) {
  ensure_connected();

  TExecuteStatementReq request;
  TExecuteStatementResp response;
  request.sessionHandle = session_handle_;
  request.statement = sql;

  client_->ExecuteStatement(response, request);
  operation_handle_ = response.operationHandle;
}

Rcpp::List Bee::Client::fetch(int num_rows) {
  ensure_connected();

  if (num_rows < 1) {
    throw std::runtime_error("num_rows must be > 1");
  }

  TFetchResultsReq result_request;
  TFetchResultsResp result_response;
  TGetResultSetMetadataReq metadata_request;
  TGetResultSetMetadataResp metadata_response;

  result_request.operationHandle = operation_handle_;
  result_request.orientation = TFetchOrientation::FETCH_NEXT;
  result_request.maxRows = num_rows;
  client_->FetchResults(result_response, result_request);

  // TFetchResultsResp.hasMoreRows is currently broken, so just keep going until
  // we fall short.
  has_more_rows_ = result_response.results.rows.size() >= num_rows;

  metadata_request.operationHandle = operation_handle_;
  client_->GetResultSetMetadata(metadata_response, metadata_request);

  return build_data_frame(metadata_response.schema, result_response.results);
}

bool Bee::Client::has_more_rows() const {
  return has_more_rows_;
}

std::string Bee::Client::inspect() {
  std::ostringstream out;
  out << "Bee::Client[" <<
    "hostname=" << socket_->getHost() <<
    ", port=" << socket_->getPort() <<
    ", user=" << user_ <<
    ", pass=" << pass_ <<
    (transport_->isOpen() ? ", connected" : "") <<
    (has_session_handle() ? ", in session" : "") <<
    (has_operation_handle() ? ", in operation" : "") <<
    "]";
  return out.str();
}

void Bee::Client::write_frame(const std::string &bytes) {
  Quad length = big_endian(bytes.size());
  socket_->write(length.bytes, 4);
  socket_->write(reinterpret_cast<const unsigned char *>(bytes.c_str()), bytes.size());
}

boost::shared_ptr<std::string> Bee::Client::read_frame() {
  unsigned int length = 0;
  unsigned char length_buffer[5];
  socket_->read(length_buffer, 4);
  length_buffer[4] = '\0';
  for (int i = 0; i < 4; ++i) {
    length += static_cast<int>(length_buffer[1 + i]) << i;
  }

  if (length == 0) {
    return boost::shared_ptr<std::string>(new std::string(""));
  } else {
    char message[length];
    socket_->read(reinterpret_cast<unsigned char *>(message), length);
    return boost::shared_ptr<std::string>(new std::string(message, length));
  }
}

void Bee::Client::perform_sasl_handshake() {
  // START PLAIN
  {
    unsigned char code = 0x01;
    socket_->write(&code, 1);
    write_frame("PLAIN");
  }

  // OK PLAIN user pass
  {
    unsigned char code = 0x02;
    socket_->write(&code, 1);
    std::ostringstream frame("[PLAIN]\x00");
    frame << user_ << '\0' << pass_;
    write_frame(frame.str());
  }

  socket_->flush();

  unsigned char status;
  socket_->read(&status, 1);

  switch (status) {
  case 3: // bad
  case 4: // error
      throw std::runtime_error(*read_frame());

  case 5: // complete
      read_frame();  // challenge -- ignored
      break;

  case 2: // ok
    {
      boost::shared_ptr<std::string> message(read_frame());
      throw std::runtime_error(std::string("ERROR: auth failed: ") + *message);
      break;
    }

  default:
    throw std::runtime_error("unexpected response from sasl handshake");
  }
}

void Bee::Client::open_session() {
  TOpenSessionReq request;
  TOpenSessionResp response;
  request.client_protocol = TProtocolVersion::HIVE_CLI_SERVICE_PROTOCOL_V3;

  client_->OpenSession(response, request);
  session_handle_ = response.sessionHandle;
}

void Bee::Client::ensure_connected() const {
  if (!transport_->isOpen())
    throw std::runtime_error("not connected");
}

bool Bee::Client::has_session_handle() const {
  return !session_handle_.sessionId.guid.empty();
}

bool Bee::Client::has_operation_handle() const {
  return !operation_handle_.operationId.guid.empty();
{'fraction_non_alphanumeric': 0.15743797754928782, 'fraction_numerical': 0.01207433261013112, 'mean_word_length': 4.6604378003203415, 'pattern_counts': {'":': 0, '<': 106, '<?xml version=': 0, '>': 132, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6299885', 'n_tokens_mistral': 3565, 'n_tokens_neox': 3390, 'n_words': 660}