com.netflix.astyanax.entitystore.DefaultEntityManager Maven / Gradle / Ivy

Go to download
package com.netflix.astyanax.entitystore;

import java.util.Collection;
import java.util.List;
import java.util.Map;

import javax.persistence.PersistenceException;

import org.apache.commons.lang.StringUtils;

import com.google.common.base.Function;
import com.google.common.base.Preconditions;
import com.google.common.collect.Lists;
import com.netflix.astyanax.Keyspace;
import com.netflix.astyanax.MutationBatch;
import com.netflix.astyanax.connectionpool.exceptions.ConnectionException;
import com.netflix.astyanax.model.ColumnFamily;
import com.netflix.astyanax.model.ColumnList;
import com.netflix.astyanax.model.ConsistencyLevel;
import com.netflix.astyanax.model.CqlResult;
import com.netflix.astyanax.model.Row;
import com.netflix.astyanax.model.Rows;
import com.netflix.astyanax.partitioner.BigInteger127Partitioner;
import com.netflix.astyanax.partitioner.Partitioner;
import com.netflix.astyanax.query.ColumnFamilyQuery;
import com.netflix.astyanax.recipes.reader.AllRowsReader;
import com.netflix.astyanax.retry.RetryPolicy;
import com.netflix.astyanax.serializers.StringSerializer;

/**
 * Manager entities in a column famliy with any key type but columns that are
 * encoded as strings.
 */
public class DefaultEntityManager implements EntityManager {

	//////////////////////////////////////////////////////////////////
	// Builder pattern

	public static class Builder {

		private Class clazz = null;
		private EntityMapper entityMapper = null;
		private Keyspace keyspace = null;
		private ColumnFamily columnFamily = null;
		private ConsistencyLevel readConsitency = null;
		private ConsistencyLevel writeConsistency = null;
		private Integer ttl = null;
		private RetryPolicy retryPolicy = null;
		private LifecycleEvents lifecycleHandler = null;
		private String columnFamilyName = null;
		private boolean autoCommit = true;
		private Partitioner partitioner = DEFAULT_ENTITY_MANAGER_PARTITIONER;

		public Builder() {

		}

		/**
		 * mandatory
		 * @param clazz entity class type
		 */
		public Builder withEntityType(Class clazz) {
			Preconditions.checkNotNull(clazz);
			this.clazz = clazz;
			return this;
		}

		/**
		 * mandatory
		 * @param keyspace
		 */
		public Builder withKeyspace(Keyspace keyspace) {
			Preconditions.checkNotNull(keyspace);
			this.keyspace = keyspace;
			return this;
		}

		/**
		 * optional
		 * @param columnFamily column name type is fixed to String/UTF8
		 */
		public Builder withColumnFamily(ColumnFamily columnFamily) {
		    Preconditions.checkState(this.columnFamilyName == null && this.columnFamily == null , "withColumnFamily called multiple times");
			Preconditions.checkNotNull(columnFamily);
			this.columnFamily = columnFamily;
			return this;
		}
		
		/**
		 * optional
		 * @param columnFamilyName Name of column family to use.  
		 */
		public Builder withColumnFamily(String columnFamilyName) {
            Preconditions.checkState(this.columnFamilyName == null && columnFamily == null , "withColumnFamily called multiple times");
            Preconditions.checkNotNull(columnFamilyName);
		    this.columnFamilyName = columnFamilyName;
		    return this;
		}

		/**
		 * optional
		 * @param level
		 */
		public Builder withReadConsistency(ConsistencyLevel level) {
			Preconditions.checkNotNull(level);
			this.readConsitency = level;
			return this;
		}

		/**
		 * optional
		 * @param level
		 */
		public Builder withWriteConsistency(ConsistencyLevel level) {
			Preconditions.checkNotNull(level);
			this.writeConsistency = level;
			return this;
		}

		/**
		 * set both read and write consistency
		 * optional
		 * @param level
		 */
		public Builder withConsistency(ConsistencyLevel level) {
			Preconditions.checkNotNull(level);
			this.readConsitency = level;
			this.writeConsistency = level;
			return this;
		}

		/**
		 * default TTL for all columns written to cassandra
		 * optional
		 * @return
		 */
		public Builder withTTL(Integer ttl) {
			this.ttl = ttl;
			return this;
		}

		/**
		 * optional
		 * @param level
		 */
		public Builder withRetryPolicy(RetryPolicy policy) {
			Preconditions.checkNotNull(policy);
			this.retryPolicy = policy;
			return this;
		}
		
		public Builder withAutoCommit(boolean autoCommit) {
			this.autoCommit = autoCommit;
			return this;
		}

		/**
		 * Partitioner used to determine token ranges and how to break token
		 * ranges into sub parts. The default is BigInteger127Partitioner in
		 * pre-cassandra 1.2.
		 * 
		 * @param partitioner
		 * @return
		 */
		public Builder withPartitioner(Partitioner partitioner) {
			this.partitioner = partitioner;
			return this;
		}

		@SuppressWarnings("unchecked")
		public DefaultEntityManager build() {
			// check mandatory fields
			Preconditions.checkNotNull(clazz, "withEntityType(...) is not set");
			Preconditions
					.checkNotNull(keyspace, "withKeyspace(...) is not set");

			// TODO: check @Id type compatibility
			// TODO: do we need to require @Entity annotation
			this.entityMapper = new EntityMapper(clazz, ttl);
			this.lifecycleHandler = new LifecycleEvents(clazz);

			if (columnFamily == null) {
				if (columnFamilyName == null)
					columnFamilyName = entityMapper.getEntityName();
				columnFamily = new ColumnFamily(columnFamilyName,
						(com.netflix.astyanax.Serializer) MappingUtils
								.getSerializerForField(this.entityMapper
										.getId()), StringSerializer.get());
			}
			// build object
			return new DefaultEntityManager(this);
		}
	}
	
	public static  Builder builder() {
	    return new Builder();
	}

	//////////////////////////////////////////////////////////////////
	// private members

	private final EntityMapper entityMapper;
	private final Keyspace keyspace;
	private final ColumnFamily columnFamily;
	private final ConsistencyLevel readConsitency;
	private final ConsistencyLevel writeConsistency;
	private final RetryPolicy retryPolicy;
	private final LifecycleEvents lifecycleHandler;
	private final boolean autoCommit;
	private final ThreadLocal tlMutation = new ThreadLocal();
	private static final Partitioner DEFAULT_ENTITY_MANAGER_PARTITIONER = BigInteger127Partitioner
			.get();
	private final Partitioner partitioner;

	private DefaultEntityManager(Builder builder) {
		entityMapper      = builder.entityMapper;
		keyspace          = builder.keyspace;
		columnFamily      = builder.columnFamily;
		readConsitency    = builder.readConsitency;
		writeConsistency  = builder.writeConsistency;
		retryPolicy       = builder.retryPolicy;
		lifecycleHandler  = builder.lifecycleHandler;
		autoCommit        = builder.autoCommit;
		partitioner       = builder.partitioner;
	}

	//////////////////////////////////////////////////////////////////
	// public APIs

	/**
	 * @inheritDoc
	 */
	public void put(T entity) throws PersistenceException {
		try {
		    lifecycleHandler.onPrePersist(entity);
            MutationBatch mb = newMutationBatch();
			entityMapper.fillMutationBatch(mb, columnFamily, entity);			
            if (autoCommit)
                mb.execute();
            lifecycleHandler.onPostPersist(entity);
		} catch(Exception e) {
			throw new PersistenceException("failed to put entity ", e);
		}
	}

	/**
	 * @inheritDoc
	 */
	public T get(K id) throws PersistenceException {
		try {
			ColumnFamilyQuery cfq = newQuery();            
			ColumnList cl = cfq.getKey(id).execute().getResult();
			// when a row is deleted in cassandra,
			// the row key remains (without any columns) until the next compaction.
			// simply return null (as non exist)
			if(cl.isEmpty())
				return null;
			T entity = entityMapper.constructEntity(id, cl);
			lifecycleHandler.onPostLoad(entity);
			return entity;
		} catch(Exception e) {
			throw new PersistenceException("failed to get entity " + id, e);
		}
	}

	/**
	 * @inheritDoc
	 */
	@Override
	public void delete(K id) throws PersistenceException {
		try {
			MutationBatch mb = getMutationBatch();
			mb.withRow(columnFamily, id).delete();
            if (autoCommit)
                mb.execute();
		} catch(Exception e) {
			throw new PersistenceException("failed to delete entity " + id, e);
		}
	}
	
    @Override
    public void remove(T entity) throws PersistenceException {
        K id = null;
        try {
            lifecycleHandler.onPreRemove(entity);
            id = entityMapper.getEntityId(entity);
            MutationBatch mb = newMutationBatch();
            mb.withRow(columnFamily, id).delete();
            if (autoCommit)
                mb.execute();
            lifecycleHandler.onPostRemove(entity);
        } catch(Exception e) {
            throw new PersistenceException("failed to delete entity " + id, e);
        }
    }

    /**
     * @inheritDoc
     */
    @Override
    public List getAll() throws PersistenceException {
        final List entities = Lists.newArrayList();
        visitAll(new Function() {
            @Override
            public synchronized Boolean apply(T entity) {
                entities.add(entity);
                return true;
            }
        });
        return entities;
    }

    /**
     * @inheritDoc
     */
    @Override
    public List get(Collection ids) throws PersistenceException {
        try {
            ColumnFamilyQuery cfq = newQuery();            
            Rows rows = cfq.getRowSlice(ids).execute().getResult();

            List entities = Lists.newArrayListWithExpectedSize(rows.size());
            for (Row row : rows) {
                if (!row.getColumns().isEmpty()) { 
                    T entity = entityMapper.constructEntity(row.getKey(), row.getColumns());
                    lifecycleHandler.onPostLoad(entity);
                    entities.add(entity);
                }
            }
            return entities;
        } catch(Exception e) {
            throw new PersistenceException("failed to get entities " + ids, e);
        }
    }

    /**
     * @inheritDoc
     */
    @Override
    public void delete(Collection ids) throws PersistenceException {
        MutationBatch mb = getMutationBatch();        
        try {
            for (K id : ids) {
                mb.withRow(columnFamily, id).delete();
            }
            if (autoCommit)
                mb.execute();
        } catch(Exception e) {
            throw new PersistenceException("failed to delete entities " + ids, e);
        }
    }

    @Override
    public void remove(Collection entities) throws PersistenceException {
        MutationBatch mb = getMutationBatch();        
        try {
            for (T entity : entities) {
                lifecycleHandler.onPreRemove(entity);
                K id = entityMapper.getEntityId(entity);
                mb.withRow(columnFamily, id).delete();
            }
            mb.execute();
            for (T entity : entities) {
                lifecycleHandler.onPostRemove(entity);
            }
        } catch(Exception e) {
            throw new PersistenceException("failed to delete entities ", e);
        }
    }

    /**
     * @inheritDoc
     */
    @Override
    public void put(Collection entities) throws PersistenceException {
        MutationBatch mb = getMutationBatch();        
        try {
            for (T entity : entities) {
                lifecycleHandler.onPrePersist(entity);
                entityMapper.fillMutationBatch(mb, columnFamily, entity);           
            }
            if (autoCommit)
                mb.execute();
            
            for (T entity : entities) {
                lifecycleHandler.onPostPersist(entity);
            }

        } catch(Exception e) {
            throw new PersistenceException("failed to put entities ", e);
        }
    }
    
    /**
     * @inheritDoc
     */
    @Override
    public void visitAll(final Function callback) throws PersistenceException {
        try {
            new AllRowsReader.Builder(keyspace, columnFamily)
                    .withIncludeEmptyRows(false)
                    .withPartitioner(partitioner)
                    .forEachRow(new Function, Boolean>() {
                        @Override
                        public Boolean apply(Row row) {
                            if (row.getColumns().isEmpty())
                                return true;
                            T entity = entityMapper.constructEntity(row.getKey(), row.getColumns());
                            try {
                                lifecycleHandler.onPostLoad(entity);
                            } catch (Exception e) {
                                // TODO:
                            }
                            return callback.apply(entity);
                        }
                    })
                    .build()
                    .call();
        } catch (Exception e) {
            throw new PersistenceException("Failed to fetch all entites", e);
        }
    }
    
    @Override
    public List find(String cql) throws PersistenceException {
        Preconditions.checkArgument(StringUtils.left(cql, 6).equalsIgnoreCase("SELECT"), "CQL must be SELECT statement");
        
        try {
            CqlResult results = newQuery().withCql(cql).execute().getResult();
            List entities = Lists.newArrayListWithExpectedSize(results.getRows().size());
            for (Row row : results.getRows()) {
                if (!row.getColumns().isEmpty()) { 
                    T entity = entityMapper.constructEntity(row.getKey(), row.getColumns());
                    lifecycleHandler.onPostLoad(entity);
                    entities.add(entity);
                }
            }
            return entities;
        } catch (Exception e) {
            throw new PersistenceException("Failed to execute cql query", e);
        }
    }
    
    private MutationBatch newMutationBatch() {
        MutationBatch mb = keyspace.prepareMutationBatch();
        if(writeConsistency != null)
            mb.withConsistencyLevel(writeConsistency);
        if(retryPolicy != null)
            mb.withRetryPolicy(retryPolicy);
        return mb;
    }
    
    private MutationBatch getMutationBatch() {
        if (autoCommit) {
            return newMutationBatch();
        }
        else {
            MutationBatch mb = tlMutation.get();
            if (mb == null) {
                mb = newMutationBatch();
                tlMutation.set(mb);
            }
            return mb;
        }
    }
    
    private ColumnFamilyQuery newQuery() {
        ColumnFamilyQuery cfq = keyspace.prepareQuery(columnFamily);
        if(readConsitency != null)
            cfq.setConsistencyLevel(readConsitency);
        if(retryPolicy != null)
            cfq.withRetryPolicy(retryPolicy);
        return cfq;
    }

    @Override
    public void createStorage(Map options) throws PersistenceException {
        try {
            keyspace.createColumnFamily(this.columnFamily, options);
        } catch (ConnectionException e) {
            if (e.getMessage().contains("already exist")) 
                return;
            throw new PersistenceException("Unable to create column family " + this.columnFamily.getName(), e);
        }
    }

    @Override
    public void deleteStorage() throws PersistenceException {
        try {
            keyspace.dropColumnFamily(this.columnFamily);
        } catch (ConnectionException e) {
            throw new PersistenceException("Unable to drop column family " + this.columnFamily.getName(), e);
        }
    }

    @Override
    public void truncate() throws PersistenceException {
        try {
            keyspace.truncateColumnFamily(this.columnFamily);
        } catch (ConnectionException e) {
            throw new PersistenceException("Unable to drop column family " + this.columnFamily.getName(), e);
        }
    }

    @Override
    public void commit() throws PersistenceException {
        MutationBatch mb = tlMutation.get();
        if (mb != null) {
            try {
                mb.execute();
            } catch (ConnectionException e) {
                throw new PersistenceException("Failed to commit mutation batch", e);
            }
        }
    }

    @Override
    public NativeQuery createNativeQuery() {
        throw new UnsupportedOperationException("Not implemented yet");
    }

}