tokens.cpp

#include <cassandra.h>

#include <iostream>
#include <algorithm>
#include <vector>
#include <memory>
#include <sstream>

// Requires this PR to reliably query the same coordinator: https://github.com/datastax/cpp-driver/pull/489

typedef std::unique_ptr<CassFuture, decltype(cass_future_free)*> FuturePtr;

void add_tokens(const FuturePtr& fut, std::vector<int64_t>& tokens) {
  std::unique_ptr<const CassResult, decltype(cass_result_free)*> result {
    cass_future_get_result(fut.get()),
    cass_result_free
  };

  std::unique_ptr<CassIterator, decltype(cass_iterator_free)*> rows_it {
    cass_iterator_from_result(result.get()),
    cass_iterator_free
  };

  while (cass_iterator_next(rows_it.get())) {
    const CassRow* row = cass_iterator_get_row(rows_it.get());
    const CassValue* tokens_col = cass_row_get_column_by_name(row, "tokens");

    std::unique_ptr<CassIterator, decltype(cass_iterator_free)*> tokens_it {
      cass_iterator_from_collection(tokens_col),
      cass_iterator_free
    };

    while (cass_iterator_next(tokens_it.get())) {
      const CassValue* token_val = cass_iterator_get_value(tokens_it.get());

      const char* token_str;
      size_t token_str_size;
      cass_value_get_string(token_val, &token_str, &token_str_size);

      int64_t token;
      std::istringstream ss(std::string(token_str, token_str_size));
      ss >> token;
      tokens.push_back(token);
    }
  }
}

FuturePtr query(CassSession* session, const char* query, const CassNode* node) {
  std::unique_ptr<CassStatement, decltype(cass_statement_free)*> statement {
    cass_statement_new(query, 0),
    cass_statement_free 
  };

  if (node != NULL) {
    cass_statement_set_node(statement.get(), node);
  }

  return FuturePtr {
    cass_session_execute(session, statement.get()),
    cass_future_free
  };
}

// Query `system.local` and `system.peers` for a sorted list of tokens
std::vector<int64_t> query_tokens(CassSession* session) {
  std::vector<int64_t> tokens;

  FuturePtr local {
    query(session, "SELECT tokens FROM system.local", NULL)
  };

  add_tokens(local, tokens);

  FuturePtr peers {
    query(session, "SELECT tokens FROM system.peers", cass_future_coordinator(local.get()))
  };

  add_tokens(peers, tokens);

  std::sort(tokens.begin(), tokens.end());

  return tokens;
}