package com.netflix.astyanax.cql.schema;
import static com.datastax.driver.core.querybuilder.QueryBuilder.eq;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Properties;
import org.apache.commons.lang.StringUtils;
import org.codehaus.jettison.json.JSONException;
import org.codehaus.jettison.json.JSONObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import com.datastax.driver.core.ResultSet;
import com.datastax.driver.core.Row;
import com.datastax.driver.core.Session;
import com.datastax.driver.core.Statement;
import com.datastax.driver.core.querybuilder.QueryBuilder;
import com.google.common.base.Preconditions;
import com.google.common.collect.Maps;
import com.netflix.astyanax.connectionpool.OperationResult;
import com.netflix.astyanax.cql.CqlOperationResultImpl;
import com.netflix.astyanax.ddl.ColumnFamilyDefinition;
import com.netflix.astyanax.ddl.FieldMetadata;
import com.netflix.astyanax.ddl.KeyspaceDefinition;
import com.netflix.astyanax.ddl.SchemaChangeResult;
/**
* Impl for {@link KeyspaceDefinition} using the java driver.
*
* @author poberai
*
*/
public class CqlKeyspaceDefinitionImpl implements KeyspaceDefinition {
private static final Logger Log = LoggerFactory.getLogger(CqlKeyspaceDefinitionImpl.class);
private final Session session;
private boolean alterKeyspace;
private final Map<String, Object> options = new HashMap<String, Object>();
private final List<CqlColumnFamilyDefinitionImpl> cfDefList = new ArrayList<CqlColumnFamilyDefinitionImpl>();
public CqlKeyspaceDefinitionImpl(Session session) {
this.session = session;
}
public CqlKeyspaceDefinitionImpl(Session session, Map<String, Object> input) {
this.session = session;
checkOptionsMap(input);
}
public CqlKeyspaceDefinitionImpl(Session session, Properties props) {
this.session = session;
checkOptionsMap(propertiesToMap(props));
}
public CqlKeyspaceDefinitionImpl(Session session, Row row) {
this.session = session;
this.setName(row.getString("keyspace_name"));
this.setStrategyClass(row.getString("strategy_class"));
this.setStrategyOptionsMap(parseStrategyOptions(row.getString("strategy_options")));
this.options.put("durable_writes", row.getBool("durable_writes"));
}
public CqlKeyspaceDefinitionImpl alterKeyspace() {
alterKeyspace = true;
return this;
}
@Override
public CqlKeyspaceDefinitionImpl setName(String name) {
this.options.put("name", name.toLowerCase());
return this;
}
@Override
public String getName() {
return (String) options.get("name");
}
@Override
public CqlKeyspaceDefinitionImpl setStrategyClass(String strategyClass) {
getOrCreateReplicationMap().put("class", strategyClass);
return this;
}
@Override
public String getStrategyClass() {
return (String) getOrCreateReplicationMap().get("class");
}
@Override
public CqlKeyspaceDefinitionImpl setStrategyOptions(Map<String, String> strategyOptions) {
getOrCreateReplicationMap().putAll(strategyOptions);
return this;
}
public CqlKeyspaceDefinitionImpl setStrategyOptionsMap(Map<String, Object> strategyOptions) {
getOrCreateReplicationMap().putAll(strategyOptions);
return this;
}
@Override
public CqlKeyspaceDefinitionImpl addStrategyOption(String name, String value) {
this.getOrCreateReplicationMap().put(name, value);
return this;
}
@Override
public Map<String, String> getStrategyOptions() {
Map<String, String> map = new HashMap<String, String>();
Map<String, Object> repMap = getOrCreateReplicationMap();
for (String key : repMap.keySet()) {
map.put(key, (String) repMap.get(key));
}
return map;
}
@Override
public List<ColumnFamilyDefinition> getColumnFamilyList() {
Statement query = QueryBuilder.select().all()
.from("system", "schema_columnfamilies")
.where(eq("keyspace_name", getName()));
ResultSet rs = session.execute(query);
List<ColumnFamilyDefinition> cfDefs = new ArrayList<ColumnFamilyDefinition>();
List<Row> rows = rs.all();
if (rows != null) {
for (Row row : rows) {
cfDefs.add(new CqlColumnFamilyDefinitionImpl(session, row));
}
}
return cfDefs;
}
@Override
public ColumnFamilyDefinition getColumnFamily(String columnFamilyName) {
Statement query = QueryBuilder.select().all()
.from("system", "schema_columnfamilies")
.where(eq("keyspace_name", getName()))
.and(eq("columnfamily_name", columnFamilyName.toLowerCase()));
Row row = session.execute(query).one();
if (row == null) {
throw new RuntimeException("CF not found: " + columnFamilyName);
}
return new CqlColumnFamilyDefinitionImpl(session, row);
}
@Override
public KeyspaceDefinition addColumnFamily(ColumnFamilyDefinition cfDef) {
CqlColumnFamilyDefinitionImpl cqlCfDef = (CqlColumnFamilyDefinitionImpl) cfDef;
cqlCfDef.execute();
return this;
}
@Override
public Collection<String> getFieldNames() {
return options.keySet();
}
@Override
public Object getFieldValue(String name) {
return options.get(name);
}
@Override
public KeyspaceDefinition setFieldValue(String name, Object value) {
this.options.put(name, value);
return this;
}
@Override
public Collection<FieldMetadata> getFieldsMetadata() {
List<FieldMetadata> list = new ArrayList<FieldMetadata>();
for (String key : options.keySet()) {
Object value = options.get(key);
Class<?> clazz = value.getClass();
String name = key.toUpperCase();
String type = clazz.getSimpleName().toUpperCase();
boolean isContainer = Collection.class.isAssignableFrom(clazz) || Map.class.isAssignableFrom(clazz);
list.add(new FieldMetadata(name, type, isContainer));
}
return list;
}
@Override
public void setFields(Map<String, Object> optionsMap) {
checkOptionsMap(optionsMap);
}
@Override
public Properties getProperties() throws Exception {
return mapToProperties(options);
}
@Override
public void setProperties(Properties props) throws Exception {
options.clear();
options.putAll(propertiesToMap(props));
}
public OperationResult<SchemaChangeResult> execute() {
String query = getQuery();
if (Log.isDebugEnabled()) {
Log.debug("Query : " + query);
}
CqlOperationResultImpl<SchemaChangeResult> result = new CqlOperationResultImpl<SchemaChangeResult>(session.execute(query), null);
for (CqlColumnFamilyDefinitionImpl cfDef : cfDefList) {
cfDef.execute();
}
return result;
}
private String getQuery() {
String cmd = (alterKeyspace) ? "ALTER" : "CREATE";
StringBuilder sb = new StringBuilder(cmd);
sb.append(" KEYSPACE ");
sb.append(getName());
Map<String, Object> replicationOptions = (Map<String, Object>) options.get("replication");
appendReplicationOptions(sb, replicationOptions);
Object durableWrites = options.get("durable_writes");
if (durableWrites != null) {
sb.append(" AND durable_writes = ").append(durableWrites);
}
return sb.toString();
}
private void appendReplicationOptions(StringBuilder sb, Map<String, Object> replicationOptions) {
if (replicationOptions == null || replicationOptions.size() == 0) {
throw new RuntimeException("Missing properties for 'replication'");
}
sb.append(" WITH replication = {" );
Iterator<Entry<String, Object>> iter = replicationOptions.entrySet().iterator();
while (iter.hasNext()) {
Entry<String, Object> entry = iter.next();
sb.append("'").append(entry.getKey()).append("' : '").append(entry.getValue()).append("'");
if (iter.hasNext()) {
sb.append(", ");
}
}
sb.append("}");
}
private void checkOptionsMap(Map<String, Object> input) {
Object strategyOptions = input.get("strategy_options");
if (strategyOptions == null) {
Preconditions.checkArgument(input.get("replication") != null, "Invalid CREATE KEYSPACE properties");
options.clear();
options.putAll(input);
} else {
// this is an old style map. Convert to the new spec of CREATE KEYSPACE
options.clear();
Map<String, Object> replicationOptions = new HashMap<String, Object>();
options.put("replication", replicationOptions);
Map<String, Object> oldStrategyOptions = (Map<String, Object>) input.get("strategy_options");
replicationOptions.putAll(oldStrategyOptions);
String strategyClass = (String) input.get("strategy_class");
replicationOptions.put("class", strategyClass);
}
}
private Map<String, Object> getOrCreateReplicationMap() {
Map<String, Object> replicationMap = (Map<String, Object>) options.get("replication");
if (replicationMap == null) {
replicationMap = new HashMap<String, Object>();
options.put("replication", replicationMap);
}
return replicationMap;
}
private static Map<String, Object> propertiesToMap(Properties props) {
Map<String, Object> root = Maps.newTreeMap();
for (Entry<Object, Object> prop : props.entrySet()) {
String[] parts = StringUtils.split((String)prop.getKey(), ".");
Map<String, Object> node = root;
for (int i = 0; i < parts.length - 1; i++) {
if (!node.containsKey(parts[i])) {
node.put(parts[i], new LinkedHashMap<String, Object>());
}
node = (Map<String, Object>)node.get(parts[i]);
}
node.put(parts[parts.length-1], (String)prop.getValue());
}
return root;
}
private static Properties mapToProperties(Map<String, Object> map) {
Properties props = new Properties();
addProperties(props, null, map);
return props;
}
private static void addProperties(Properties props, String prefix, Map<String, Object> subMap) {
for (Entry<String, Object> entry : subMap.entrySet()) {
String key = (prefix != null) ? prefix + "." + entry.getKey() : entry.getKey();
if (entry.getValue() instanceof Map) {
addProperties(props, key, (Map<String, Object>) entry.getValue());
} else {
props.put(key, entry.getValue().toString());
}
}
}
private Map<String, Object> parseStrategyOptions(String jsonString) {
if (jsonString == null || jsonString.isEmpty()) {
return null;
}
Map<String, Object> map = new HashMap<String, Object>();
try {
JSONObject json = new JSONObject(jsonString);
Iterator<String> iter = json.keys();
while (iter.hasNext()) {
String key = iter.next();
Object obj = json.get(key);
map.put(key, obj);
}
return map;
} catch (JSONException e) {
throw new RuntimeException(e);
}
}
public String toString() {
return "CqlKeyspaceDefinition=[ " + options.toString() + " ]";
}
}