com.hmsonline.trident.cql.CassandraCqlMapStateFactory Maven / Gradle / Ivy

Go to download
package com.hmsonline.trident.cql;

import java.util.Map;

import com.datastax.driver.core.Session;
import storm.trident.state.State;
import storm.trident.state.StateFactory;
import storm.trident.state.StateType;
import storm.trident.state.map.CachedMap;
import storm.trident.state.map.MapState;
import storm.trident.state.map.NonTransactionalMap;
import storm.trident.state.map.OpaqueMap;
import storm.trident.state.map.SnapshottableMap;
import storm.trident.state.map.TransactionalMap;
import backtype.storm.task.IMetricsContext;
import backtype.storm.tuple.Values;

import com.hmsonline.trident.cql.CassandraCqlMapState.Options;
import com.hmsonline.trident.cql.mappers.CqlRowMapper;

/**
 * The class responsible for generating instances of
 * the {@link CassandraCqlMapState}.
 *
 * @author robertlee
 */
public class CassandraCqlMapStateFactory implements StateFactory {
    private static final long serialVersionUID = 1L;

	protected CqlClientFactory clientFactory;
	protected StateType stateType;
	protected Options options;

    @SuppressWarnings("rawtypes")
	protected CqlRowMapper mapper;

    @SuppressWarnings({"rawtypes"})
    public CassandraCqlMapStateFactory(CqlRowMapper mapper, StateType stateType, Options options) {
        this.stateType = stateType;
        this.options = options;
        this.mapper = mapper;
    }

    @SuppressWarnings({"rawtypes", "unchecked"})
    public State makeState(Map configuration, IMetricsContext metrics, int partitionIndex, int numPartitions) {

        if (clientFactory == null) {
            clientFactory = new MapConfiguredCqlClientFactory(configuration);
        }

        Session session;
        if(options.keyspace != null) {
            session = clientFactory.getSession(options.keyspace);
        } else {
            session = clientFactory.getSession();
        }

        CassandraCqlMapState state = new CassandraCqlMapState(session, mapper, options, configuration);
        state.registerMetrics(configuration, metrics, options.mapStateMetricName);

        CachedMap cachedMap = new CachedMap(state, options.localCacheSize);

        MapState mapState;
        if (stateType == StateType.NON_TRANSACTIONAL) {
            mapState = NonTransactionalMap.build(cachedMap);
        } else if (stateType == StateType.OPAQUE) {
            mapState = OpaqueMap.build(cachedMap);
        } else if (stateType == StateType.TRANSACTIONAL) {
            mapState = TransactionalMap.build(cachedMap);
        } else {
            throw new RuntimeException("Unknown state type: " + stateType);
        }

        return new SnapshottableMap(mapState, new Values(options.globalKey));
    }
}