org.cassandraunit.DataLoader Maven / Gradle / Ivy

Go to download
package org.cassandraunit;

import java.util.ArrayList;
import java.util.List;

import me.prettyprint.cassandra.serializers.LongSerializer;
import me.prettyprint.hector.api.Cluster;
import me.prettyprint.hector.api.Keyspace;
import me.prettyprint.hector.api.beans.HColumn;
import me.prettyprint.hector.api.beans.HCounterColumn;
import me.prettyprint.hector.api.beans.HCounterSuperColumn;
import me.prettyprint.hector.api.beans.HSuperColumn;
import me.prettyprint.hector.api.ddl.ColumnFamilyDefinition;
import me.prettyprint.hector.api.ddl.ColumnType;
import me.prettyprint.hector.api.ddl.ComparatorType;
import me.prettyprint.hector.api.ddl.KeyspaceDefinition;
import me.prettyprint.hector.api.factory.HFactory;
import me.prettyprint.hector.api.mutation.Mutator;

import org.cassandraunit.dataset.DataSet;
import org.cassandraunit.model.ColumnFamilyModel;
import org.cassandraunit.model.ColumnModel;
import org.cassandraunit.model.KeyspaceModel;
import org.cassandraunit.model.RowModel;
import org.cassandraunit.model.SuperColumnModel;
import org.cassandraunit.serializer.GenericTypeSerializer;
import org.cassandraunit.type.GenericType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * 
 * @author Jeremy Sevellec
 * 
 */
public class DataLoader {
	Cluster cluster = null;

	private Logger log = LoggerFactory.getLogger(DataLoader.class);

	public DataLoader(String clusterName, String host) {
		super();
		cluster = HFactory.getOrCreateCluster(clusterName, host);
	}

	protected Cluster getCluster() {
		return cluster;
	}

	public void load(DataSet dataSet) {
		KeyspaceModel dataSetKeyspace = dataSet.getKeyspace();

		dropKeyspaceIfExist(dataSetKeyspace.getName());

		KeyspaceDefinition keyspaceDefinition = createKeyspaceDefinition(dataSet);

		cluster.addKeyspace(keyspaceDefinition, true);

		log.info("creating keyspace : {}", keyspaceDefinition.getName());
		Keyspace keyspace = HFactory.createKeyspace(dataSet.getKeyspace().getName(), cluster);
		log.info("loading data into keyspace : {}", keyspaceDefinition.getName());
		loadData(dataSet, keyspace);
	}

	private KeyspaceDefinition createKeyspaceDefinition(DataSet dataSet) {
		List columnFamilyDefinitions = createColumnFamilyDefinitions(dataSet);

		KeyspaceModel dataSetKeyspace = dataSet.getKeyspace();

		KeyspaceDefinition keyspaceDefinition = HFactory.createKeyspaceDefinition(dataSetKeyspace.getName(),
				dataSetKeyspace.getStrategy().value(), dataSetKeyspace.getReplicationFactor(), columnFamilyDefinitions);
		return keyspaceDefinition;
	}

	private void dropKeyspaceIfExist(String keyspaceName) {
		KeyspaceDefinition existedKeyspace = cluster.describeKeyspace(keyspaceName);
		if (existedKeyspace != null) {
			log.info("dropping existing keyspace : {}", existedKeyspace.getName());
			cluster.dropKeyspace(keyspaceName, true);
		}
	}

	private void loadData(DataSet dataSet, Keyspace keyspace) {
		for (ColumnFamilyModel columnFamily : dataSet.getColumnFamilies()) {
			loadColumnFamilyData(columnFamily, keyspace);
		}

	}

	private void loadColumnFamilyData(ColumnFamilyModel columnFamily, Keyspace keyspace) {
		Mutator mutator = HFactory.createMutator(keyspace, GenericTypeSerializer.get());
		for (RowModel row : columnFamily.getRows()) {
			switch (columnFamily.getType()) {
			case STANDARD:
				loadStandardColumnFamilyData(columnFamily, mutator, row);
				break;
			case SUPER:
				loadSuperColumnFamilyData(columnFamily, mutator, row);
				break;
			default:
				break;
			}

		}
		mutator.execute();

	}

	private void loadSuperColumnFamilyData(ColumnFamilyModel columnFamily, Mutator mutator, RowModel row) {
		if (columnFamily.isCounter()) {
			for (SuperColumnModel superColumnModel : row.getSuperColumns()) {
				HCounterSuperColumn superCounterColumn = HFactory.createCounterSuperColumn(
						superColumnModel.getName(), createHCounterColumnList(superColumnModel.getColumns()),
						GenericTypeSerializer.get(), GenericTypeSerializer.get());
				mutator.addCounter(row.getKey(), columnFamily.getName(), superCounterColumn);
			}
		} else {
			for (SuperColumnModel superColumnModel : row.getSuperColumns()) {
				HSuperColumn superColumn = HFactory.createSuperColumn(
						superColumnModel.getName(), createHColumnList(superColumnModel.getColumns()),
						GenericTypeSerializer.get(), GenericTypeSerializer.get(), GenericTypeSerializer.get());
				mutator.addInsertion(row.getKey(), columnFamily.getName(), superColumn);
			}
		}
	}

	private void loadStandardColumnFamilyData(ColumnFamilyModel columnFamily, Mutator mutator, RowModel row) {
		if (columnFamily.isCounter()) {
			for (HCounterColumn hCounterColumn : createHCounterColumnList(row.getColumns())) {
				mutator.addCounter(row.getKey(), columnFamily.getName(), hCounterColumn);
			}
		} else {
			for (HColumn hColumn : createHColumnList(row.getColumns())) {
				mutator.addInsertion(row.getKey(), columnFamily.getName(), hColumn);
			}
		}
	}

	private List> createHColumnList(List columnsModel) {
		List> hColumns = new ArrayList>();
		for (ColumnModel columnModel : columnsModel) {
			HColumn column = HFactory.createColumn(columnModel.getName(),
					columnModel.getValue(), GenericTypeSerializer.get(), GenericTypeSerializer.get());
			hColumns.add(column);
		}
		return hColumns;
	}

	private List> createHCounterColumnList(List columnsModel) {
		List> hColumns = new ArrayList>();
		for (ColumnModel columnModel : columnsModel) {
			HCounterColumn column = HFactory.createCounterColumn(columnModel.getName(), LongSerializer
					.get().fromByteBuffer(GenericTypeSerializer.get().toByteBuffer(columnModel.getValue())),
					GenericTypeSerializer.get());
			hColumns.add(column);
		}
		return hColumns;
	}

	private List createColumnFamilyDefinitions(DataSet dataSet) {
		KeyspaceModel dataSetKeyspace = dataSet.getKeyspace();
		List columnFamilyDefinitions = new ArrayList();
		for (ColumnFamilyModel columnFamily : dataSet.getColumnFamilies()) {
			ColumnFamilyDefinition cfDef = HFactory.createColumnFamilyDefinition(dataSetKeyspace.getName(),
					columnFamily.getName());
			cfDef.setColumnType(columnFamily.getType());

			cfDef.setKeyValidationClass(columnFamily.getKeyType().getClassName());
			cfDef.setComparatorType(ComparatorType.getByClassName(columnFamily.getComparatorType().getClassName()));
			if (columnFamily.isCounter()) {
				cfDef.setDefaultValidationClass(columnFamily.getDefaultColumnValueType().getClassName());
			}
			if (columnFamily.getType().equals(ColumnType.SUPER) && columnFamily.getSubComparatorType() != null) {
				cfDef.setSubComparatorType(columnFamily.getSubComparatorType());
			}

			columnFamilyDefinitions.add(cfDef);
		}
		return columnFamilyDefinitions;
	}
}