private void refreshEndpointMap() {
String keyspace = ConfigHelper.getOutputKeyspace(conf);
try (Session session =
CqlConfigHelper.getOutputCluster(ConfigHelper.getOutputInitialAddress(conf), conf)
.connect(keyspace)) {
rangeMap = new HashMap<>();
metadata = session.getCluster().getMetadata();
Set<TokenRange> ranges = metadata.getTokenRanges();
for (TokenRange range : ranges) {
rangeMap.put(range, metadata.getReplicas(keyspace, range));
}
}
}
/** get prepared statement id from cache, otherwise prepare it from Cassandra server */
private PreparedStatement preparedStatement(Session client) {
PreparedStatement statement = preparedStatements.get(client);
if (statement == null) {
PreparedStatement result;
try {
result = client.prepare(cql);
} catch (NoHostAvailableException e) {
throw new RuntimeException("failed to prepare cql query " + cql, e);
}
PreparedStatement previousId = preparedStatements.putIfAbsent(client, result);
statement = previousId == null ? result : previousId;
}
return statement;
}
CqlRecordWriter(Configuration conf) {
this.conf = conf;
this.queueSize =
conf.getInt(ColumnFamilyOutputFormat.QUEUE_SIZE, 32 * FBUtilities.getAvailableProcessors());
batchThreshold = conf.getLong(ColumnFamilyOutputFormat.BATCH_THRESHOLD, 32);
this.clients = new HashMap<>();
try {
String keyspace = ConfigHelper.getOutputKeyspace(conf);
try (Session client =
CqlConfigHelper.getOutputCluster(ConfigHelper.getOutputInitialAddress(conf), conf)
.connect(keyspace)) {
ringCache = new NativeRingCache(conf);
if (client != null) {
TableMetadata tableMetadata =
client
.getCluster()
.getMetadata()
.getKeyspace(client.getLoggedKeyspace())
.getTable(ConfigHelper.getOutputColumnFamily(conf));
clusterColumns = tableMetadata.getClusteringColumns();
partitionKeyColumns = tableMetadata.getPartitionKey();
String cqlQuery = CqlConfigHelper.getOutputCql(conf).trim();
if (cqlQuery.toLowerCase().startsWith("insert"))
throw new UnsupportedOperationException(
"INSERT with CqlRecordWriter is not supported, please use UPDATE/DELETE statement");
cql = appendKeyWhereClauses(cqlQuery);
} else {
throw new IllegalArgumentException("Invalid configuration specified " + conf);
}
}
} catch (Exception e) {
throw new RuntimeException(e);
}
}
@Test
public void compressionTest() throws Exception {
// Same as executeTest, but with compression enabled
cluster
.getConfiguration()
.getProtocolOptions()
.setCompression(ProtocolOptions.Compression.SNAPPY);
try {
Session compressedSession = cluster.connect(TestUtils.SIMPLE_KEYSPACE);
// Simple calls to all versions of the execute/executeAsync methods
String key = "execute_compressed_test";
ResultSet rs =
compressedSession.execute(String.format(INSERT_FORMAT, TABLE, key, "foo", 42, 24.03f));
assertTrue(rs.isExhausted());
String SELECT_ALL = String.format(SELECT_ALL_FORMAT + " WHERE k = '%s'", TABLE, key);
// execute
checkExecuteResultSet(compressedSession.execute(SELECT_ALL), key);
checkExecuteResultSet(
compressedSession.execute(
new SimpleStatement(SELECT_ALL).setConsistencyLevel(ConsistencyLevel.ONE)),
key);
// executeAsync
checkExecuteResultSet(compressedSession.executeAsync(SELECT_ALL).getUninterruptibly(), key);
checkExecuteResultSet(
compressedSession
.executeAsync(
new SimpleStatement(SELECT_ALL).setConsistencyLevel(ConsistencyLevel.ONE))
.getUninterruptibly(),
key);
} finally {
cluster
.getConfiguration()
.getProtocolOptions()
.setCompression(ProtocolOptions.Compression.NONE);
}
}
protected void closeInternal() {
if (client != null) {
client.close();
}
}
/** Loops collecting cql binded variable values from the queue and sending to Cassandra */
public void run() {
outer:
while (run || !queue.isEmpty()) {
List<ByteBuffer> bindVariables;
try {
bindVariables = queue.take();
} catch (InterruptedException e) {
// re-check loop condition after interrupt
continue;
}
ListIterator<InetAddress> iter = endpoints.listIterator();
while (true) {
// send the mutation to the last-used endpoint. first time through, this will NPE
// harmlessly.
try {
int i = 0;
PreparedStatement statement = preparedStatement(client);
while (bindVariables != null) {
BoundStatement boundStatement = new BoundStatement(statement);
for (int columnPosition = 0;
columnPosition < bindVariables.size();
columnPosition++) {
boundStatement.setBytesUnsafe(columnPosition, bindVariables.get(columnPosition));
}
client.execute(boundStatement);
i++;
if (i >= batchThreshold) break;
bindVariables = queue.poll();
}
break;
} catch (Exception e) {
closeInternal();
if (!iter.hasNext()) {
lastException = new IOException(e);
break outer;
}
}
// attempt to connect to a different endpoint
try {
InetAddress address = iter.next();
String host = address.getHostName();
client = CqlConfigHelper.getOutputCluster(host, conf).connect();
} catch (Exception e) {
// If connection died due to Interrupt, just try connecting to the endpoint again.
// There are too many ways for the Thread.interrupted() state to be cleared, so
// we can't rely on that here. Until the java driver gives us a better way of knowing
// that this exception came from an InterruptedException, this is the best solution.
if (canRetryDriverConnection(e)) {
iter.previous();
}
closeInternal();
// Most exceptions mean something unexpected went wrong to that endpoint, so
// we should try again to another. Other exceptions (auth or invalid request) are
// fatal.
if ((e instanceof AuthenticationException || e instanceof InvalidQueryException)
|| !iter.hasNext()) {
lastException = new IOException(e);
break outer;
}
}
}
}
// close all our connections once we are done.
closeInternal();
}