CqlKeyspaceImpl.java

package com.netflix.astyanax.cql;

import static com.datastax.driver.core.querybuilder.QueryBuilder.eq;

import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.Callable;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.codahale.metrics.MetricRegistryListener;
import com.datastax.driver.core.Cluster;
import com.datastax.driver.core.Configuration;
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.Function;
import com.google.common.collect.Lists;
import com.netflix.astyanax.AstyanaxConfiguration;
import com.netflix.astyanax.Clock;
import com.netflix.astyanax.ColumnMutation;
import com.netflix.astyanax.Keyspace;
import com.netflix.astyanax.KeyspaceTracerFactory;
import com.netflix.astyanax.MutationBatch;
import com.netflix.astyanax.SerializerPackage;
import com.netflix.astyanax.clock.MicrosecondsAsyncClock;
import com.netflix.astyanax.connectionpool.ConnectionPool;
import com.netflix.astyanax.connectionpool.ConnectionPoolConfiguration;
import com.netflix.astyanax.connectionpool.ConnectionPoolMonitor;
import com.netflix.astyanax.connectionpool.ConnectionPoolProxy.SeedHostListener;
import com.netflix.astyanax.connectionpool.Host;
import com.netflix.astyanax.connectionpool.Operation;
import com.netflix.astyanax.connectionpool.OperationResult;
import com.netflix.astyanax.connectionpool.TokenRange;
import com.netflix.astyanax.connectionpool.exceptions.ConnectionException;
import com.netflix.astyanax.connectionpool.exceptions.NotFoundException;
import com.netflix.astyanax.connectionpool.exceptions.OperationException;
import com.netflix.astyanax.connectionpool.impl.OperationResultImpl;
import com.netflix.astyanax.cql.direct.DirectCqlStatement;
import com.netflix.astyanax.cql.reads.CqlColumnFamilyQueryImpl;
import com.netflix.astyanax.cql.schema.CqlColumnFamilyDefinitionImpl;
import com.netflix.astyanax.cql.schema.CqlKeyspaceDefinitionImpl;
import com.netflix.astyanax.cql.util.CFQueryContext;
import com.netflix.astyanax.cql.writes.CqlColumnMutationImpl;
import com.netflix.astyanax.cql.writes.CqlMutationBatchImpl;
import com.netflix.astyanax.ddl.ColumnFamilyDefinition;
import com.netflix.astyanax.ddl.KeyspaceDefinition;
import com.netflix.astyanax.ddl.SchemaChangeResult;
import com.netflix.astyanax.ddl.impl.SchemaChangeResponseImpl;
import com.netflix.astyanax.model.ColumnFamily;
import com.netflix.astyanax.partitioner.BigInteger127Partitioner;
import com.netflix.astyanax.partitioner.Murmur3Partitioner;
import com.netflix.astyanax.partitioner.Partitioner;
import com.netflix.astyanax.query.ColumnFamilyQuery;
import com.netflix.astyanax.retry.RetryPolicy;
import com.netflix.astyanax.serializers.SerializerPackageImpl;
import com.netflix.astyanax.serializers.UnknownComparatorException;

/**
 * Java Driver based impl of {@link Keyspace} that implements ddl operations as well as row queries and mutation batches.
 * The class encapsulates a java driver cluster and session object to provide all the functionality. 
 *  
 * Note that due to the way the object is setup via AstyanaxContext and CqlFamilyFactory, it needs to implements 
 * a {@link SeedHostListener} so that it can construct the cluster and session object appropriately once the seed hosts
 * have been provided by the {@link HostSupplier} object.
 *  
 * @author poberai
 */
public class CqlKeyspaceImpl implements Keyspace, SeedHostListener {

	private static final Logger Logger = LoggerFactory.getLogger(CqlKeyspaceImpl.class);
	
	private final Clock clock;
	
	public volatile Cluster cluster;
	public volatile Session session;
	
	private final KeyspaceContext ksContext;
	private final String keyspaceName;
	private final AstyanaxConfiguration astyanaxConfig;
	private final KeyspaceTracerFactory tracerFactory; 
	private final Configuration javaDriverConfig;
	private final ConnectionPoolMonitor cpMonitor;
	private final MetricRegistryListener metricsRegListener;

	public CqlKeyspaceImpl(String ksName, AstyanaxConfiguration asConfig, KeyspaceTracerFactory tracerFactory, ConnectionPoolConfiguration cpConfig, ConnectionPoolMonitor cpMonitor) {
		this(null, ksName, asConfig, tracerFactory, cpConfig,cpMonitor);
	}

	public CqlKeyspaceImpl(KeyspaceContext ksContext) {
		this(ksContext.getSession(), ksContext.getKeyspace(), ksContext.getConfig(), ksContext.getTracerFactory(), null, ksContext.getConnectionPoolMonitor());
	}

	CqlKeyspaceImpl(Session session, String ksName, AstyanaxConfiguration asConfig, KeyspaceTracerFactory tracerFactory, ConnectionPoolMonitor cpMonitor) {
		this(session, ksName, asConfig, tracerFactory, null,cpMonitor);
	}

	private CqlKeyspaceImpl(Session session, String ksName, AstyanaxConfiguration asConfig, KeyspaceTracerFactory tracerFactory, ConnectionPoolConfiguration cpConfig, ConnectionPoolMonitor cpMonitor) {
		this.session = session;
		this.keyspaceName = ksName.toLowerCase();
		this.astyanaxConfig = asConfig;
		this.tracerFactory = tracerFactory;
		this.cpMonitor = cpMonitor;
		this.metricsRegListener = ((JavaDriverConnectionPoolMonitorImpl)cpMonitor).getMetricsRegistryListener();
		this.ksContext = new KeyspaceContext(this);
		
		if (asConfig.getClock() != null) {
			clock = asConfig.getClock();
		} else {
			clock = new MicrosecondsAsyncClock();
		}
		
		if (cpConfig != null) {
			javaDriverConfig = ((JavaDriverConnectionPoolConfigurationImpl)cpConfig).getJavaDriverConfig();
		} else {
			javaDriverConfig = null;
		}
	}
	
	
	@Override
	public AstyanaxConfiguration getConfig() {
		return astyanaxConfig;
	}

	@Override
	public String getKeyspaceName() {
		return keyspaceName;
	}

	@Override
	public Partitioner getPartitioner() throws ConnectionException {
		String pName = describePartitioner();
		if (pName.contains("Murmur3Partitioner")) {
			return Murmur3Partitioner.get();
		} else if (pName.contains("RandomPartitioner")) {
			return BigInteger127Partitioner.get();
		} else {
			throw new RuntimeException("Unrecognized partitioner: " + pName);
		}
	}

	@Override
	public String describePartitioner() throws ConnectionException {
		Statement q = QueryBuilder.select("partitioner").from("system", "local");
		ResultSet result = session.execute(q);
		com.datastax.driver.core.Row row = result.one();
		if (row == null) {
			throw new RuntimeException("Missing paritioner");
		}
		String pName = row.getString(0);
		return pName;
	}

	@Override
	public List<TokenRange> describeRing() throws ConnectionException {
		return CqlRingDescriber.getInstance().getTokenRanges(session, false);
	}

	@Override
	public List<TokenRange> describeRing(String dc) throws ConnectionException {
		return CqlRingDescriber.getInstance().getTokenRanges(session, dc, null);
	}

	@Override
	public List<TokenRange> describeRing(String dc, String rack) throws ConnectionException {
		return CqlRingDescriber.getInstance().getTokenRanges(session, dc, rack);
	}

	@Override
	public List<TokenRange> describeRing(boolean cached) throws ConnectionException {
		return CqlRingDescriber.getInstance().getTokenRanges(session, cached);
	}

	@Override
	public KeyspaceDefinition describeKeyspace() throws ConnectionException {
		
		Statement query = QueryBuilder.select().from("system", "schema_keyspaces").where(eq("keyspace_name", keyspaceName));
		Row row = session.execute(query).one();
		if (row == null) {
			throw new RuntimeException("Keyspace not found: " + keyspaceName);
		}
		return (new CqlKeyspaceDefinitionImpl(session, row));
	}

	@Override
	public Properties getKeyspaceProperties() throws ConnectionException {
		try {
			return describeKeyspace().getProperties();
		} catch (Exception e) {
			throw new RuntimeException(e);
		}
	}

	@Override
	public Properties getColumnFamilyProperties(String columnFamily) throws ConnectionException {
        KeyspaceDefinition ksDef = this.describeKeyspace();
        ColumnFamilyDefinition cfDef = ksDef.getColumnFamily(columnFamily);
        if (cfDef == null)
            throw new NotFoundException(String.format("Column family '%s' in keyspace '%s' not found", columnFamily, getKeyspaceName()));
        try {
			return cfDef.getProperties();
		} catch (Exception e) {
			throw new RuntimeException();
		}
	}

	@Override
	public SerializerPackage getSerializerPackage(String cfName, boolean ignoreErrors) throws ConnectionException, UnknownComparatorException {
		
		ColumnFamilyDefinition cfDef = describeKeyspace().getColumnFamily(cfName);
		return new SerializerPackageImpl(cfDef, ignoreErrors);
	}

	@Override
	public MutationBatch prepareMutationBatch() {
		return new CqlMutationBatchImpl(ksContext, clock, astyanaxConfig.getDefaultWriteConsistencyLevel(), astyanaxConfig.getRetryPolicy());
	}

	@Override
	public <K, C> ColumnMutation prepareColumnMutation(ColumnFamily<K, C> columnFamily, K rowKey, C column) {
		return new CqlColumnMutationImpl<K,C>(ksContext, new CFQueryContext<K, C>(columnFamily, rowKey), column);
	}

	@Override
	public <K, C> ColumnFamilyQuery<K, C> prepareQuery(ColumnFamily<K, C> cf) {
		return new CqlColumnFamilyQueryImpl<K,C>(ksContext, cf);
	}

	@Override
	public OperationResult<SchemaChangeResult> createKeyspace(Map<String, Object> options) throws ConnectionException {
		return new CqlKeyspaceDefinitionImpl(session, options).setName(keyspaceName).execute();
	}

	@Override
	public OperationResult<SchemaChangeResult> createKeyspace(Properties properties) throws ConnectionException {
		return new CqlKeyspaceDefinitionImpl(session, properties).setName(keyspaceName).execute();
	}

	@SuppressWarnings("rawtypes")
	@Override
	public OperationResult<SchemaChangeResult> createKeyspace(Map<String, Object> options, Map<ColumnFamily, Map<String, Object>> cfs) throws ConnectionException {
		
		CqlKeyspaceDefinitionImpl ksDef = new CqlKeyspaceDefinitionImpl(session, options);
		if (ksDef.getName() == null) {
			ksDef.setName(keyspaceName);
		}
		
		OperationResult<SchemaChangeResult> result = ksDef.execute();
		
		for (ColumnFamily cf : cfs.keySet()) {
			CqlColumnFamilyDefinitionImpl cfDef = new CqlColumnFamilyDefinitionImpl(session, ksDef.getName(), cf, cfs.get(cf));
			ksDef.addColumnFamily(cfDef);
		}
		
		return result;
	}

	@Override
	public OperationResult<SchemaChangeResult> updateKeyspace(Map<String, Object> options) throws ConnectionException {
		return new CqlKeyspaceDefinitionImpl(session, options).setName(keyspaceName).alterKeyspace().execute();
	}

	@Override
	public OperationResult<SchemaChangeResult> updateKeyspace(Properties props) throws ConnectionException {
		return new CqlKeyspaceDefinitionImpl(session, props).setName(keyspaceName).alterKeyspace().execute();
	}

	@Override
	public OperationResult<SchemaChangeResult> dropKeyspace() throws ConnectionException {
		return new CqlOperationResultImpl<SchemaChangeResult>(session.execute("DROP KEYSPACE " + keyspaceName), null);
	}
	
	@Override
	public <K, C> OperationResult<Void> truncateColumnFamily(ColumnFamily<K, C> columnFamily) throws OperationException, ConnectionException {
		ResultSet result = session.execute("TRUNCATE " + keyspaceName + "." + columnFamily.getName());
		return new CqlOperationResultImpl<Void>(result, null);
	}

	@Override
	public OperationResult<Void> truncateColumnFamily(String columnFamily) throws ConnectionException {
		ResultSet result = session.execute("TRUNCATE " + keyspaceName + "." + columnFamily);
		return new CqlOperationResultImpl<Void>(result, null);
	}

	@Override
	public <K, C> OperationResult<SchemaChangeResult> createColumnFamily(ColumnFamily<K, C> columnFamily, Map<String, Object> options) throws ConnectionException {
		return new CqlColumnFamilyDefinitionImpl(session, keyspaceName, columnFamily, options).execute();
	}

	@Override
	public OperationResult<SchemaChangeResult> createColumnFamily(Properties props) throws ConnectionException {
		return new CqlColumnFamilyDefinitionImpl(session, keyspaceName, props).execute();
	}

	@Override
	public OperationResult<SchemaChangeResult> createColumnFamily(Map<String, Object> options) throws ConnectionException {
		return new CqlColumnFamilyDefinitionImpl(session, keyspaceName, options).execute();
	}

	@Override
	public <K, C> OperationResult<SchemaChangeResult> updateColumnFamily(ColumnFamily<K, C> columnFamily, Map<String, Object> options) throws ConnectionException {
		return new CqlColumnFamilyDefinitionImpl(session, keyspaceName, columnFamily, options).alterTable().execute();
	}

	@Override
	public OperationResult<SchemaChangeResult> updateColumnFamily(Properties props) throws ConnectionException {
		return new CqlColumnFamilyDefinitionImpl(session, keyspaceName, props).alterTable().execute();
	}

	@Override
	public OperationResult<SchemaChangeResult> updateColumnFamily(Map<String, Object> options) throws ConnectionException {
		return new CqlColumnFamilyDefinitionImpl(session, keyspaceName, options).alterTable().execute();
	}

	@Override
	public OperationResult<SchemaChangeResult> dropColumnFamily(String columnFamilyName) throws ConnectionException {
		return new CqlOperationResultImpl<SchemaChangeResult>(session.execute("DROP TABLE " + keyspaceName + "." + columnFamilyName), null);
	}

	@Override
	public <K, C> OperationResult<SchemaChangeResult> dropColumnFamily(ColumnFamily<K, C> columnFamily) throws ConnectionException {
		return dropColumnFamily(columnFamily.getName());
	}

	@Override
	public Map<String, List<String>> describeSchemaVersions() throws ConnectionException {
		return new CqlSchemaVersionReader(session).exec();
	}

	@Override
	public CqlStatement prepareCqlStatement() {
		return new DirectCqlStatement(session);
	}

	@Override
	public ConnectionPool<?> getConnectionPool() throws ConnectionException {
		throw new UnsupportedOperationException("Operation not supported");
	}


	@Override
	public OperationResult<Void> testOperation(Operation<?, ?> operation) throws ConnectionException {
		throw new UnsupportedOperationException("Operation not supported");
	}

	@Override
	public OperationResult<Void> testOperation(Operation<?, ?> operation, RetryPolicy retry) throws ConnectionException {
		throw new UnsupportedOperationException("Operation not supported");
	}

	@Override
	public void setHosts(Collection<Host> hosts, int port) {

		try {
			if (session != null) {
				Logger.info("Session has already been set, SKIPPING SET HOSTS");
				return;
			}
			List<Host> hostList = Lists.newArrayList(hosts);

			List<String> contactPoints = Lists.transform(hostList, new Function<Host, String>() {
				@Override
				public String apply(Host input) {
					if (input != null) {
						return input.getHostName(); 
					}
					return null;
				}
			});

			Configuration config = javaDriverConfig;
			
			// We really need a mechanism to easily override Configuration on the builder
			Logger.info("Using port: " + port);
			
			Cluster.Builder builder = Cluster.builder()
					.addContactPoints(contactPoints.toArray(new String[0]))
					.withPort(port)
					.withLoadBalancingPolicy(config.getPolicies().getLoadBalancingPolicy())
					.withReconnectionPolicy(config.getPolicies().getReconnectionPolicy())
					.withRetryPolicy(config.getPolicies().getRetryPolicy())
					.withCompression(config.getProtocolOptions().getCompression())
					.withPoolingOptions(config.getPoolingOptions())
					.withSocketOptions(config.getSocketOptions())
					.withQueryOptions(config.getQueryOptions());
			
			if (config.getMetricsOptions() == null) {
				builder.withoutMetrics();
			} else if (!config.getMetricsOptions().isJMXReportingEnabled()) {
				builder.withoutJMXReporting();
			}
					
			cluster = builder.build();
			if (!(this.cpMonitor instanceof JavaDriverConnectionPoolMonitorImpl))
				this.cluster.getMetrics().getRegistry().addListener((MetricRegistryListener) this.metricsRegListener);
			
			Logger.info("Connecting to cluster");
			session = cluster.connect();
			Logger.info("Done connecting to cluster, session object created");

		} catch (RuntimeException e) {
			Logger.error("Failed to set hosts for keyspace impl", e);
			
		} catch (Exception e) {
			Logger.error("Failed to set hosts for keyspace impl", e);
		}
	}
	
	@Override
	public void shutdown() {
		cluster.close();
	}


	public class KeyspaceContext {
		
		private final Keyspace ks; 
		
		public KeyspaceContext(Keyspace keyspaceCtx) {
			this.ks = keyspaceCtx;
		}
		public Session getSession() {
			return session;
		}
		public String getKeyspace() {
			return keyspaceName;
		}
		public AstyanaxConfiguration getConfig() {
			return astyanaxConfig;
		}
		public KeyspaceTracerFactory getTracerFactory() {
			return tracerFactory;
		}
		
		public Keyspace getKeyspaceContext() {
			return ks;
		}
		public ConnectionPoolMonitor getConnectionPoolMonitor(){
			return cpMonitor;
		}
	}


	@Override
	public OperationResult<SchemaChangeResult> createKeyspaceIfNotExists(final Map<String, Object> options) throws ConnectionException {

		return createKeyspaceIfNotExists(new Callable<OperationResult<SchemaChangeResult>>() {
			@Override
			public OperationResult<SchemaChangeResult> call() throws Exception {
				return createKeyspace(options);
			}
		});
	}

	@Override
	public OperationResult<SchemaChangeResult> createKeyspaceIfNotExists(final Properties properties) throws ConnectionException {
		
		return createKeyspaceIfNotExists(new Callable<OperationResult<SchemaChangeResult>>() {
			@Override
			public OperationResult<SchemaChangeResult> call() throws Exception {
				return createKeyspace(properties);
			}
		});
	}

	@Override
	public OperationResult<SchemaChangeResult> createKeyspaceIfNotExists(final Map<String, Object> options, final Map<ColumnFamily, Map<String, Object>> cfs) throws ConnectionException {
		
		return createKeyspaceIfNotExists(new Callable<OperationResult<SchemaChangeResult>>() {
			@Override
			public OperationResult<SchemaChangeResult> call() throws Exception {
				return createKeyspace(options, cfs);
			}
		});
	}
	
    
    private OperationResult<SchemaChangeResult> createKeyspaceIfNotExists(Callable<OperationResult<SchemaChangeResult>> createKeyspace) throws ConnectionException {
        
    	// Check if keyspace exists
    	ResultSet result = session.execute("select * from system.local where keyspace_name = '" + keyspaceName + "'");
    	List<Row> rows = result.all();
    	if (rows != null && rows.isEmpty()) {
    		return new OperationResultImpl<SchemaChangeResult>(Host.NO_HOST, new SchemaChangeResponseImpl().setSchemaId("no-op"), 0);
    	}

    	try {
    		return createKeyspace.call();
    	} catch (ConnectionException e) {
    		throw e;
    	} catch (Exception e) {
    		throw new RuntimeException(e);
    	}
    }
}