/**
* Insert a record in the database. Any field/value pairs in the specified values HashMap will be
* written into the record with the specified record key.
*
* @param table The name of the table
* @param key The record key of the record to insert.
* @param values A HashMap of field/value pairs to insert in the record
* @return Zero on success, a non-zero error code on error
*/
public int insert(String table, String key, HashMap<String, ByteIterator> values) {
if (!_table.equals(table)) {
try {
client.set_keyspace(table);
_table = table;
} catch (Exception e) {
e.printStackTrace();
e.printStackTrace(System.out);
return Error;
}
}
for (int i = 0; i < OperationRetries; i++) {
if (_debug) {
System.out.println("Inserting key: " + key);
}
try {
ByteBuffer wrappedKey = ByteBuffer.wrap(key.getBytes("UTF-8"));
ColumnOrSuperColumn column;
for (Map.Entry<String, ByteIterator> entry : values.entrySet()) {
column = new ColumnOrSuperColumn();
Column subColumn = new Column(ByteBuffer.wrap(entry.getKey().getBytes("UTF-8")));
subColumn.setValue(ByteBuffer.wrap(entry.getValue().toArray()));
subColumn.setTimestamp(System.currentTimeMillis());
column.setColumn(subColumn);
mutations.add(new Mutation().setColumn_or_supercolumn(column));
}
mutationMap.put(column_family, mutations);
record.put(wrappedKey, mutationMap);
client.batch_mutate(record, ConsistencyLevel.ONE);
mutations.clear();
mutationMap.clear();
record.clear();
return Ok;
} catch (Exception e) {
errorexception = e;
}
try {
Thread.sleep(500);
} catch (InterruptedException e) {
}
}
errorexception.printStackTrace();
errorexception.printStackTrace(System.out);
return Error;
}
/**
* Return a client based on the given socket that points to the configured keyspace, and is logged
* in with the configured credentials.
*
* @param socket a socket pointing to a particular node, seed or otherwise
* @param conf a job configuration
* @return a cassandra client
* @throws InvalidRequestException
* @throws TException
* @throws AuthenticationException
* @throws AuthorizationException
*/
public static Cassandra.Client createAuthenticatedClient(TSocket socket, Configuration conf)
throws InvalidRequestException, TException, AuthenticationException, AuthorizationException {
TBinaryProtocol binaryProtocol = new TBinaryProtocol(new TFramedTransport(socket));
Cassandra.Client client = new Cassandra.Client(binaryProtocol);
socket.open();
client.set_keyspace(ConfigHelper.getOutputKeyspace(conf));
if (ConfigHelper.getOutputKeyspaceUserName(conf) != null) {
Map<String, String> creds = new HashMap<String, String>();
creds.put(SimpleAuthenticator.USERNAME_KEY, ConfigHelper.getOutputKeyspaceUserName(conf));
creds.put(SimpleAuthenticator.PASSWORD_KEY, ConfigHelper.getOutputKeyspacePassword(conf));
AuthenticationRequest authRequest = new AuthenticationRequest(creds);
client.login(authRequest);
}
return client;
}
private static Cassandra.Client createThriftClient(
String host, int port, String user, String passwd, ITransportFactory transportFactory)
throws Exception {
TTransport trans = transportFactory.openTransport(host, port);
TProtocol protocol = new TBinaryProtocol(trans);
Cassandra.Client client = new Cassandra.Client(protocol);
if (user != null && passwd != null) {
Map<String, String> credentials = new HashMap<>();
credentials.put(IAuthenticator.USERNAME_KEY, user);
credentials.put(IAuthenticator.PASSWORD_KEY, passwd);
AuthenticationRequest authenticationRequest = new AuthenticationRequest(credentials);
client.login(authenticationRequest);
}
return client;
}
/**
* Connects to the given server:port and returns a client based on the given socket that points to
* the configured keyspace, and is logged in with the configured credentials.
*
* @param host fully qualified host name to connect to
* @param port RPC port of the server
* @param conf a job configuration
* @return a cassandra client
* @throws Exception set of thrown exceptions may be implementation defined, depending on the used
* transport factory
*/
public static Cassandra.Client createAuthenticatedClient(
String host, int port, Configuration conf) throws Exception {
logger.debug("Creating authenticated client for CF output format");
TTransport transport = ConfigHelper.getClientTransportFactory(conf).openTransport(host, port);
TProtocol binaryProtocol = new TBinaryProtocol(transport, true, true);
Cassandra.Client client = new Cassandra.Client(binaryProtocol);
client.set_keyspace(ConfigHelper.getOutputKeyspace(conf));
if (ConfigHelper.getOutputKeyspaceUserName(conf) != null) {
Map<String, String> creds = new HashMap<String, String>();
creds.put(IAuthenticator.USERNAME_KEY, ConfigHelper.getOutputKeyspaceUserName(conf));
creds.put(IAuthenticator.PASSWORD_KEY, ConfigHelper.getOutputKeyspacePassword(conf));
AuthenticationRequest authRequest = new AuthenticationRequest(creds);
client.login(authRequest);
}
logger.debug("Authenticated client for CF output format created successfully");
return client;
}
public void init(String keyspace) {
outputHandler.output(
String.format(
"Starting client (and waiting %d seconds for gossip) ...",
StorageService.RING_DELAY / 1000));
try {
// Init gossip
StorageService.instance.initClient();
Set<InetAddress> hosts = Gossiper.instance.getLiveMembers();
hosts.remove(FBUtilities.getLocalAddress());
if (hosts.isEmpty())
throw new IllegalStateException(
"Cannot load any sstable, no live member found in the cluster");
// Query endpoint to ranges map and schemas from thrift
String host = hosts.iterator().next().toString().substring(1);
int port = DatabaseDescriptor.getRpcPort();
Cassandra.Client client = createThriftClient(host, port);
List<TokenRange> tokenRanges = client.describe_ring(keyspace);
List<KsDef> ksDefs = client.describe_keyspaces();
Token.TokenFactory tkFactory = StorageService.getPartitioner().getTokenFactory();
try {
for (TokenRange tr : tokenRanges) {
Range range =
new Range(tkFactory.fromString(tr.start_token), tkFactory.fromString(tr.end_token));
for (String ep : tr.endpoints) {
addRangeForEndpoint(range, InetAddress.getByName(ep));
}
}
} catch (UnknownHostException e) {
throw new RuntimeException("Got an unknow host from describe_ring()", e);
}
for (KsDef ksDef : ksDefs) {
Set<String> cfs = new HashSet<String>();
for (CfDef cfDef : ksDef.cf_defs) cfs.add(cfDef.name);
knownCfs.put(ksDef.name, cfs);
}
} catch (Exception e) {
throw new RuntimeException(e);
}
}
@Override
public void init(String keyspace) {
Iterator<InetAddress> hostiter = hosts.iterator();
while (hostiter.hasNext()) {
try {
// Query endpoint to ranges map and schemas from thrift
InetAddress host = hostiter.next();
Cassandra.Client client =
createThriftClient(
host.getHostAddress(), rpcPort, this.user, this.passwd, this.transportFactory);
setPartitioner(client.describe_partitioner());
Token.TokenFactory tkFactory = getPartitioner().getTokenFactory();
for (TokenRange tr : client.describe_ring(keyspace)) {
Range<Token> range =
new Range<>(
tkFactory.fromString(tr.start_token),
tkFactory.fromString(tr.end_token),
getPartitioner());
for (String ep : tr.endpoints) {
addRangeForEndpoint(range, InetAddress.getByName(ep));
}
}
String query =
String.format(
"SELECT * FROM %s.%s WHERE keyspace_name = '%s'",
Keyspace.SYSTEM_KS, SystemKeyspace.SCHEMA_COLUMNFAMILIES_CF, keyspace);
CqlResult result =
client.execute_cql3_query(
ByteBufferUtil.bytes(query), Compression.NONE, ConsistencyLevel.ONE);
for (CqlRow row : result.rows) {
CFMetaData metadata = CFMetaData.fromThriftCqlRow(row);
knownCfs.put(metadata.cfName, metadata);
}
break;
} catch (Exception e) {
if (!hostiter.hasNext())
throw new RuntimeException("Could not retrieve endpoint ranges: ", e);
}
}
}
/**
* Delete a record from the database.
*
* @param table The name of the table
* @param key The record key of the record to delete.
* @return Zero on success, a non-zero error code on error
*/
public int delete(String table, String key) {
if (!_table.equals(table)) {
try {
client.set_keyspace(table);
_table = table;
} catch (Exception e) {
e.printStackTrace();
e.printStackTrace(System.out);
return Error;
}
}
for (int i = 0; i < OperationRetries; i++) {
try {
client.remove(
ByteBuffer.wrap(key.getBytes("UTF-8")),
new ColumnPath(column_family),
System.currentTimeMillis(),
ConsistencyLevel.ONE);
if (_debug) {
System.out.println("Delete key: " + key);
}
return Ok;
} catch (Exception e) {
errorexception = e;
}
try {
Thread.sleep(500);
} catch (InterruptedException e) {
}
}
errorexception.printStackTrace();
errorexception.printStackTrace(System.out);
return Error;
}
protected void closeInternal() {
if (client != null) {
TTransport transport = client.getOutputProtocol().getTransport();
if (transport.isOpen()) transport.close();
}
}
/**
* Initialize any state for this DB. Called once per DB instance; there is one DB instance per
* client thread.
*/
public void init() throws DBException {
String hosts = getProperties().getProperty("hosts");
if (hosts == null) {
throw new DBException("Required property \"hosts\" missing for CassandraClient");
}
column_family =
getProperties().getProperty(COLUMN_FAMILY_PROPERTY, COLUMN_FAMILY_PROPERTY_DEFAULT);
parent = new ColumnParent(column_family);
ConnectionRetries =
Integer.parseInt(
getProperties()
.getProperty(CONNECTION_RETRY_PROPERTY, CONNECTION_RETRY_PROPERTY_DEFAULT));
OperationRetries =
Integer.parseInt(
getProperties()
.getProperty(OPERATION_RETRY_PROPERTY, OPERATION_RETRY_PROPERTY_DEFAULT));
String username = getProperties().getProperty(USERNAME_PROPERTY);
String password = getProperties().getProperty(PASSWORD_PROPERTY);
_debug = Boolean.parseBoolean(getProperties().getProperty("debug", "false"));
String[] allhosts = hosts.split(",");
String myhost = allhosts[random.nextInt(allhosts.length)];
Exception connectexception = null;
for (int retry = 0; retry < ConnectionRetries; retry++) {
tr = new TFramedTransport(new TSocket(myhost, 9160));
TProtocol proto = new TBinaryProtocol(tr);
client = new Cassandra.Client(proto);
try {
tr.open();
connectexception = null;
break;
} catch (Exception e) {
connectexception = e;
}
try {
Thread.sleep(1000);
} catch (InterruptedException e) {
}
}
if (connectexception != null) {
System.err.println(
"Unable to connect to " + myhost + " after " + ConnectionRetries + " tries");
System.out.println(
"Unable to connect to " + myhost + " after " + ConnectionRetries + " tries");
throw new DBException(connectexception);
}
if (username != null && password != null) {
Map<String, String> cred = new HashMap<String, String>();
cred.put("username", username);
cred.put("password", password);
AuthenticationRequest req = new AuthenticationRequest(cred);
try {
client.login(req);
} catch (Exception e) {
throw new DBException(e);
}
}
}
/**
* Perform a range scan for a set of records in the database. Each field/value pair from the
* result will be stored in a HashMap.
*
* @param table The name of the table
* @param startkey The record key of the first record to read.
* @param recordcount The number of records to read
* @param fields The list of fields to read, or null for all of them
* @param result A Vector of HashMaps, where each HashMap is a set field/value pairs for one
* record
* @return Zero on success, a non-zero error code on error
*/
public int scan(
String table,
String startkey,
int recordcount,
Set<String> fields,
Vector<HashMap<String, ByteIterator>> result) {
if (!_table.equals(table)) {
try {
client.set_keyspace(table);
_table = table;
} catch (Exception e) {
e.printStackTrace();
e.printStackTrace(System.out);
return Error;
}
}
for (int i = 0; i < OperationRetries; i++) {
try {
SlicePredicate predicate;
if (fields == null) {
predicate =
new SlicePredicate()
.setSlice_range(new SliceRange(emptyByteBuffer, emptyByteBuffer, false, 1000000));
} else {
ArrayList<ByteBuffer> fieldlist = new ArrayList<ByteBuffer>(fields.size());
for (String s : fields) {
fieldlist.add(ByteBuffer.wrap(s.getBytes("UTF-8")));
}
predicate = new SlicePredicate().setColumn_names(fieldlist);
}
KeyRange kr =
new KeyRange()
.setStart_key(startkey.getBytes("UTF-8"))
.setEnd_key(new byte[] {})
.setCount(recordcount);
List<KeySlice> results =
client.get_range_slices(parent, predicate, kr, ConsistencyLevel.ONE);
if (_debug) {
System.out.println("Scanning startkey: " + startkey);
}
HashMap<String, ByteIterator> tuple;
for (KeySlice oneresult : results) {
tuple = new HashMap<String, ByteIterator>();
Column column;
String name;
ByteIterator value;
for (ColumnOrSuperColumn onecol : oneresult.columns) {
column = onecol.column;
name =
new String(
column.name.array(),
column.name.position() + column.name.arrayOffset(),
column.name.remaining());
value =
new ByteArrayByteIterator(
column.value.array(),
column.value.position() + column.value.arrayOffset(),
column.value.remaining());
tuple.put(name, value);
if (_debug) {
System.out.print("(" + name + "=" + value + ")");
}
}
result.add(tuple);
if (_debug) {
System.out.println();
}
}
return Ok;
} catch (Exception e) {
errorexception = e;
}
try {
Thread.sleep(500);
} catch (InterruptedException e) {
}
}
errorexception.printStackTrace();
errorexception.printStackTrace(System.out);
return Error;
}
/**
* Read a record from the database. Each field/value pair from the result will be stored in a
* HashMap.
*
* @param table The name of the table
* @param key The record key of the record to read.
* @param fields The list of fields to read, or null for all of them
* @param result A HashMap of field/value pairs for the result
* @return Zero on success, a non-zero error code on error
*/
public int read(
String table, String key, Set<String> fields, HashMap<String, ByteIterator> result) {
if (!_table.equals(table)) {
try {
client.set_keyspace(table);
_table = table;
} catch (Exception e) {
e.printStackTrace();
e.printStackTrace(System.out);
return Error;
}
}
for (int i = 0; i < OperationRetries; i++) {
try {
SlicePredicate predicate;
if (fields == null) {
predicate =
new SlicePredicate()
.setSlice_range(new SliceRange(emptyByteBuffer, emptyByteBuffer, false, 1000000));
} else {
ArrayList<ByteBuffer> fieldlist = new ArrayList<ByteBuffer>(fields.size());
for (String s : fields) {
fieldlist.add(ByteBuffer.wrap(s.getBytes("UTF-8")));
}
predicate = new SlicePredicate().setColumn_names(fieldlist);
}
List<ColumnOrSuperColumn> results =
client.get_slice(
ByteBuffer.wrap(key.getBytes("UTF-8")), parent, predicate, ConsistencyLevel.ONE);
if (_debug) {
System.out.print("Reading key: " + key);
}
Column column;
String name;
ByteIterator value;
for (ColumnOrSuperColumn oneresult : results) {
column = oneresult.column;
name =
new String(
column.name.array(),
column.name.position() + column.name.arrayOffset(),
column.name.remaining());
value =
new ByteArrayByteIterator(
column.value.array(),
column.value.position() + column.value.arrayOffset(),
column.value.remaining());
result.put(name, value);
if (_debug) {
System.out.print("(" + name + "=" + value + ")");
}
}
if (_debug) {
System.out.println();
}
return Ok;
} catch (Exception e) {
errorexception = e;
}
try {
Thread.sleep(500);
} catch (InterruptedException e) {
}
}
errorexception.printStackTrace();
errorexception.printStackTrace(System.out);
return Error;
}
