/**
* Copyright (c) 2010 Yahoo! Inc. All rights reserved.
*
* 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. See accompanying
* LICENSE file.
*/
package com.impetus.kundera.ycsb.benchmark;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Random;
import java.util.Set;
import java.util.Vector;
import org.apache.cassandra.thrift.AuthenticationRequest;
import org.apache.cassandra.thrift.Cassandra;
import org.apache.cassandra.thrift.Column;
import org.apache.cassandra.thrift.ColumnOrSuperColumn;
import org.apache.cassandra.thrift.ColumnParent;
import org.apache.cassandra.thrift.ColumnPath;
import org.apache.cassandra.thrift.ConsistencyLevel;
import org.apache.cassandra.thrift.KeyRange;
import org.apache.cassandra.thrift.KeySlice;
import org.apache.cassandra.thrift.Mutation;
import org.apache.cassandra.thrift.SlicePredicate;
import org.apache.cassandra.thrift.SliceRange;
import org.apache.thrift.protocol.TBinaryProtocol;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.transport.TFramedTransport;
import org.apache.thrift.transport.TSocket;
import org.apache.thrift.transport.TTransport;
import org.scale7.cassandra.pelops.Bytes;
import com.yahoo.ycsb.ByteArrayByteIterator;
import com.yahoo.ycsb.ByteIterator;
import com.yahoo.ycsb.DB;
import com.yahoo.ycsb.DBException;
import com.yahoo.ycsb.StringByteIterator;
import common.Logger;
//XXXX if we do replication, fix the consistency levels
/**
* Cassandra 1.0.6 client for YCSB framework
*/
public class ThriftClient extends DB
{
private static Logger logger = Logger.getLogger(ThriftClient.class);
static Random random = new Random();
public static final int Ok = 0;
public static final int Error = -1;
public static final ByteBuffer emptyByteBuffer = ByteBuffer.wrap(new byte[0]);
public int ConnectionRetries;
public String column_family;
public static final String CONNECTION_RETRY_PROPERTY = "cassandra.connectionretries";
public static final String CONNECTION_RETRY_PROPERTY_DEFAULT = "300";
public static final String USERNAME_PROPERTY = "cassandra.username";
public static final String PASSWORD_PROPERTY = "cassandra.password";
public static final String COLUMN_FAMILY_PROPERTY = "columnfamilyOrTable";
// public static final String COLUMN_FAMILY_PROPERTY_DEFAULT = "data";
public static final String READ_CONSISTENCY_LEVEL_PROPERTY = "cassandra.readconsistencylevel";
public static final String READ_CONSISTENCY_LEVEL_PROPERTY_DEFAULT = "ONE";
public static final String WRITE_CONSISTENCY_LEVEL_PROPERTY = "cassandra.writeconsistencylevel";
public static final String WRITE_CONSISTENCY_LEVEL_PROPERTY_DEFAULT = "ONE";
public static final String SCAN_CONSISTENCY_LEVEL_PROPERTY = "cassandra.scanconsistencylevel";
public static final String SCAN_CONSISTENCY_LEVEL_PROPERTY_DEFAULT = "ONE";
public static final String DELETE_CONSISTENCY_LEVEL_PROPERTY = "cassandra.deleteconsistencylevel";
public static final String DELETE_CONSISTENCY_LEVEL_PROPERTY_DEFAULT = "ONE";
TTransport tr;
Cassandra.Client client;
boolean _debug = false;
String _table = "";
Exception errorexception = null;
List<Mutation> mutations = new ArrayList<Mutation>();
Map<String, List<Mutation>> mutationMap = new HashMap<String, List<Mutation>>();
Map<ByteBuffer, Map<String, List<Mutation>>> record = new HashMap<ByteBuffer, Map<String, List<Mutation>>>();
ColumnParent parent;
ConsistencyLevel readConsistencyLevel = ConsistencyLevel.ONE;
ConsistencyLevel writeConsistencyLevel = ConsistencyLevel.ONE;
ConsistencyLevel scanConsistencyLevel = ConsistencyLevel.ONE;
ConsistencyLevel deleteConsistencyLevel = ConsistencyLevel.ONE;
/**
* Initialize any state for this DB. Called once per DB instance; there is
* one DB instance per client thread.
*/
public void init() throws DBException
{
// System.out.println("**************************" + column_family);
try
{
String hosts = getProperties().getProperty("hosts");
if (hosts == null)
{
throw new DBException("Required property \"hosts\" missing for CassandraClient");
}
column_family = getProperties().getProperty(COLUMN_FAMILY_PROPERTY);
parent = new ColumnParent(column_family);
_table = getProperties().getProperty("schema");
ConnectionRetries = Integer.parseInt(getProperties().getProperty(CONNECTION_RETRY_PROPERTY,
CONNECTION_RETRY_PROPERTY_DEFAULT));
String username = getProperties().getProperty(USERNAME_PROPERTY);
String password = getProperties().getProperty(PASSWORD_PROPERTY);
readConsistencyLevel = ConsistencyLevel.valueOf(getProperties().getProperty(
READ_CONSISTENCY_LEVEL_PROPERTY, READ_CONSISTENCY_LEVEL_PROPERTY_DEFAULT));
writeConsistencyLevel = ConsistencyLevel.valueOf(getProperties().getProperty(
WRITE_CONSISTENCY_LEVEL_PROPERTY, WRITE_CONSISTENCY_LEVEL_PROPERTY_DEFAULT));
scanConsistencyLevel = ConsistencyLevel.valueOf(getProperties().getProperty(
SCAN_CONSISTENCY_LEVEL_PROPERTY, SCAN_CONSISTENCY_LEVEL_PROPERTY_DEFAULT));
deleteConsistencyLevel = ConsistencyLevel.valueOf(getProperties().getProperty(
DELETE_CONSISTENCY_LEVEL_PROPERTY, DELETE_CONSISTENCY_LEVEL_PROPERTY_DEFAULT));
_debug = Boolean.parseBoolean(getProperties().getProperty("debug", "false"));
String[] allhosts = hosts.split(",");
String myhost = allhosts[random.nextInt(allhosts.length)];
Exception connectexception = null;
// System.out.println("Initiating connection");
if (tr == null || client == null)
{
for (int retry = 0; retry < ConnectionRetries; retry++)
{
tr = new TFramedTransport(new TSocket(myhost, 9160));
TProtocol proto = new TBinaryProtocol(tr);
// System.out.println("creating connection::");
client = new Cassandra.Client(proto);
try
{
tr.open();
connectexception = null;
break;
}
catch (Exception e)
{
connectexception = e;
logger.error(e);
}
try
{
Thread.sleep(1000);
}
catch (InterruptedException e)
{
logger.error(e);
}
}
}
if (connectexception != null)
{
logger.info("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)
{
e.printStackTrace();
logger.info(e);
throw new DBException(e);
}
}
client.set_keyspace(_table);
}
catch (Exception e)
{
logger.info(e);
}
}
/**
* Cleanup any state for this DB. Called once per DB instance; there is one
* DB instance per client thread.
*/
public void cleanup() throws DBException
{
// tr.close();
}
/**
* 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)
{
// System.out.println(key);
// System.out.println("key:" + key);
// System.out.println("_table" + _table);
// System.out.println("table" + table);
/*
* if (!_table.equals(table)) {
*/
try
{
/*
* _table = table;
*/}
catch (Exception e)
{
logger.error(e);
return Error;
}
// }
try
{
SlicePredicate predicate = new SlicePredicate();
if (fields == null)
{
// SlicePredicate predicate = new SlicePredicate();
predicate = new SlicePredicate();
predicate.setSlice_range(new SliceRange(Bytes.EMPTY.getBytes(), Bytes.EMPTY.getBytes(), true, 10000));
// 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, readConsistencyLevel);
if (_debug)
{
System.out.print("Reading key: " + key);
}
Column column;
String name;
ByteIterator value;
assert results != null && !results.isEmpty();
/* 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)
{
logger.info("(" + name + "=" + value + ")");
}
}
*/
if (_debug)
{
logger.info("ConsistencyLevel=" + readConsistencyLevel.toString());
}
return Ok;
}
catch (Exception e)
{
logger.error(e);
return Error;
}
}
/**
* 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)
{
logger.error(e);
return Error;
}
}
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, scanConsistencyLevel);
if (_debug)
{
logger.info("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)
{
logger.info("(" + name + "=" + value + ")");
}
}
result.add(tuple);
if (_debug)
{
logger.info("ConsistencyLevel=" + scanConsistencyLevel.toString());
}
}
return Ok;
}
catch (Exception e)
{
logger.error(e);
return Error;
}
}
/**
* Update 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, overwriting any existing values with the same field name.
*
* @param table
* The name of the table
* @param key
* The record key of the record to write.
* @param values
* A HashMap of field/value pairs to update in the record
* @return Zero on success, a non-zero error code on error
*/
public int update(String table, String key, HashMap<String, ByteIterator> values)
{
return insert(table, key, values);
// return Error;
}
/**
* 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)
{
logger.error(e);
return Error;
}
/*
* }
*/if (_debug)
{
// System.out.println("Inserting key: " + key);
}
try
{
ByteBuffer wrappedKey = ByteBuffer.wrap(key.getBytes("UTF-8"));
Column col;
ColumnOrSuperColumn column;
for (Map.Entry<String, ByteIterator> entry : values.entrySet())
{
col = new Column();
col.setName(ByteBuffer.wrap(entry.getKey().getBytes("UTF-8")));
col.setValue(ByteBuffer.wrap(entry.getValue().toArray()));
col.setTimestamp(System.currentTimeMillis());
column = new ColumnOrSuperColumn();
column.setColumn(col);
mutations.add(new Mutation().setColumn_or_supercolumn(column));
}
mutationMap.put(column_family, mutations);
record.put(wrappedKey, mutationMap);
client.batch_mutate(record, writeConsistencyLevel);
mutations.clear();
mutationMap.clear();
record.clear();
if (_debug)
{
logger.info("ConsistencyLevel=" + writeConsistencyLevel.toString());
}
return Ok;
}
catch (Exception e)
{
logger.error(e);
return Error;
}
}
/**
* 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)
{
logger.error(e);
return Error;
}
}
try
{
client.remove(ByteBuffer.wrap(key.getBytes("UTF-8")), new ColumnPath(column_family),
System.currentTimeMillis(), deleteConsistencyLevel);
if (_debug)
{
logger.info("Delete key: " + key);
logger.info("ConsistencyLevel=" + deleteConsistencyLevel.toString());
}
return Ok;
}
catch (Exception e)
{
logger.error(e);
return Error;
}
}
public static void main(String[] args)
{
ThriftClient cli = new ThriftClient();
Properties props = new Properties();
props.setProperty("hosts", "localhost");
cli.setProperties(props);
try
{
cli.init();
}
catch (Exception e)
{
logger.error(e);
System.exit(0);
}
HashMap<String, ByteIterator> vals = new HashMap<String, ByteIterator>();
vals.put("age", new StringByteIterator("57"));
vals.put("middlename", new StringByteIterator("bradley"));
vals.put("favoritecolor", new StringByteIterator("blue"));
int res = cli.insert("usertable", "BrianFrankCooper", vals);
logger.info("Result of insert: " + res);
HashMap<String, ByteIterator> result = new HashMap<String, ByteIterator>();
HashSet<String> fields = new HashSet<String>();
fields.add("middlename");
fields.add("age");
fields.add("favoritecolor");
res = cli.read("usertable", "BrianFrankCooper", null, result);
logger.info("Result of read: " + res);
for (String s : result.keySet())
{
logger.info("[" + s + "]=[" + result.get(s) + "]");
}
res = cli.delete("usertable", "BrianFrankCooper");
logger.info("Result of delete: " + res);
}
}