INTEXT-50: Cassandra Adapter Improvements

JIRA: https://jira.spring.io/browse/INTEXT-50

More improvements

`Void` returns handling

spring-integration-cassandra/build.gradle

@@ -34,9 +34,8 @@ if (project.hasProperty('platformVersion')) {
 sourceCompatibility = targetCompatibility = 1.7
 
 ext {
-	cassandraVersion = '2.1.2'
-	cassandraUnitVersion = '2.0.2.2'
-	gsCollectionsVersion = '6.1.0'
+	cassandraVersion = '2.1.5'
+	cassandraUnitVersion = '2.1.3.1'
 	jacocoVersion = '0.7.2.201409121644'
 	slf4jVersion = '1.7.11'
 	springDataCassandraVersion = '1.2.0.RELEASE'

spring-integration-cassandra/src/main/java/org/springframework/integration/cassandra/outbound/CassandraMessageHandler.java

@@ -16,181 +16,241 @@
 
 package org.springframework.integration.cassandra.outbound;
 
-import java.util.Collection;
+import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
-import org.apache.commons.logging.Log;
-import org.apache.commons.logging.LogFactory;
-
+import org.springframework.cassandra.core.CachedPreparedStatementCreator;
+import org.springframework.cassandra.core.PreparedStatementCreator;
 import org.springframework.cassandra.core.WriteOptions;
 import org.springframework.data.cassandra.core.CassandraOperations;
-import org.springframework.data.cassandra.core.WriteListener;
-import org.springframework.integration.cassandra.support.CassandraOutboundGatewayType;
+import org.springframework.expression.EvaluationContext;
+import org.springframework.expression.Expression;
+import org.springframework.expression.TypeLocator;
+import org.springframework.expression.spel.support.StandardEvaluationContext;
+import org.springframework.expression.spel.support.StandardTypeLocator;
+import org.springframework.integration.expression.IntegrationEvaluationContextAware;
 import org.springframework.integration.handler.AbstractReplyProducingMessageHandler;
+import org.springframework.integration.handler.ExpressionEvaluatingMessageProcessor;
+import org.springframework.integration.handler.MessageProcessor;
+import org.springframework.integration.util.AbstractExpressionEvaluator;
 import org.springframework.messaging.Message;
+import org.springframework.util.Assert;
+
+import com.datastax.driver.core.ColumnDefinitions;
+import com.datastax.driver.core.PreparedStatement;
+import com.datastax.driver.core.Statement;
 
 /**
  * @author Soby Chacko
+ * @author Artem Bilan
  */
-public class CassandraMessageHandler<T> extends AbstractReplyProducingMessageHandler {
+@SuppressWarnings("unchecked")
+public class CassandraMessageHandler<T> extends AbstractReplyProducingMessageHandler
+		implements IntegrationEvaluationContextAware {
 
-	private static final Log log = LogFactory.getLog(CassandraMessageHandler.class);
+	private final Map<String, Expression> parameterExpressions = new HashMap<>();
 
 	private final CassandraOperations cassandraTemplate;
 
-	private CassandraOutboundGatewayType gatewayType = CassandraOutboundGatewayType.INSERTING;
+	private Type queryType;
 
-	private WriteListener<T> writeListener;
-
-	private boolean producesReply = true;
+	private boolean producesReply;
 
 	/**
 	 * Prepared statement to use in association with high throughput ingestion.
 	 */
-	private String cqlIngest;
-
-	/**
-	 * Indicates whether the outbound operations need to be async.
-	 */
-	private boolean async;
+	private String ingestQuery;
 
 	/**
 	 * Various options that can be used for Cassandra writes.
 	 */
 	private WriteOptions writeOptions;
 
+	private MessageProcessor<Statement> statementProcessor;
+
+	private EvaluationContext evaluationContext;
+
 	public CassandraMessageHandler(CassandraOperations cassandraTemplate) {
+		this(cassandraTemplate, Type.INSERT);
+	}
+
+	public CassandraMessageHandler(CassandraOperations cassandraTemplate, CassandraMessageHandler.Type queryType) {
+		Assert.notNull(cassandraTemplate, "'cassandraTemplate' must not be null.");
+		Assert.notNull(queryType, "'queryType' must not be null.");
 		this.cassandraTemplate = cassandraTemplate;
+		this.queryType = queryType;
 	}
 
-	@Override
-	protected Object handleRequestMessage(Message<?> requestMessage) {
-		Object payload = requestMessage.getPayload();
-		Object result = null;
-
-		switch (gatewayType) {
-			case INSERTING:
-				if (cqlIngest != null) {
-					handleHighThroughputIngest(payload);
-				} else if (async) {
-					result = handleAsyncInsert(payload);
-				} else {
-					result = handleSynchronousInsert(requestMessage, payload);
-				}
-				break;
-			case UPDATING:
-				if (async) {
-					if (payload instanceof List) {
-						@SuppressWarnings("unchecked")
-						List<T> entities = (List<T>) payload;
-						result = cassandraTemplate.updateAsynchronously(entities, writeListener, writeOptions);
-					} else {
-						@SuppressWarnings("unchecked")
-						T typedPayload = (T) payload;
-						result = cassandraTemplate.updateAsynchronously(typedPayload, writeListener, writeOptions);
-					}
-				} else {
-					if (payload instanceof List) {
-						@SuppressWarnings("unchecked")
-						List<T> entities = (List<T>) payload;
-						result = cassandraTemplate.update(entities, writeOptions);
-					} else {
-						result = cassandraTemplate.update(payload, writeOptions);
-					}
-				}
-				break;
-			case DELETING:
-				result = null;
-				break;
-			default:
-				result = null;
-		}
-
-		if (result == null || !producesReply) {
-			return null;
-		}
-
-		return this.getMessageBuilderFactory().withPayload(result)
-				.copyHeaders(requestMessage.getHeaders()).build();
-	}
-
-	private Object handleSynchronousInsert(Message<?> message, Object payload) {
-		final Object result;
-		if (payload instanceof List) {
-			@SuppressWarnings("unchecked")
-			List<T> entities = (List<T>) payload;
-			result = cassandraTemplate.insert(entities, writeOptions);
-		} else {
-			result = cassandraTemplate.insert(message.getPayload(), writeOptions);
-		}
-		return result;
-	}
-
-	private Object handleAsyncInsert(Object payload) {
-		final Object result;
-		WriteListener<T> writeListener = getWriteListener();
-		if (payload instanceof List) {
-			@SuppressWarnings("unchecked")
-			List<T> entities = (List<T>) payload;
-			result = cassandraTemplate.insertAsynchronously(entities, writeListener, writeOptions);
-		} else {
-			@SuppressWarnings("unchecked")
-			T typedPayload = (T) payload;
-			result = cassandraTemplate.insertAsynchronously(typedPayload, writeListener, writeOptions);
-		}
-		return result;
-	}
-
-	private void handleHighThroughputIngest(Object payload) {
-		if (payload instanceof List) {
-			@SuppressWarnings("unchecked")
-			List<List<?>> data = (List<List<?>>) payload;
-			assert cqlIngest != null;
-			cassandraTemplate.ingest(cqlIngest, data, writeOptions);
-		}
-	}
-
-	private WriteListener<T> getWriteListener() {
-		return new WriteListener<T>() {
-
-			@Override
-			public void onWriteComplete(Collection<T> entities) {
-
-			}
-
-			@Override
-			public void onException(Exception x) {
-				log.debug("Exception thrown", x);
-			}
-		};
-	}
-
-	public void setCqlIngest(String cqlIngest) {
-		this.cqlIngest = cqlIngest;
-	}
-
-	public void setGatewayType(CassandraOutboundGatewayType gatewayType) {
-		this.gatewayType = gatewayType;
-	}
-
-	public void setWriteListener(WriteListener<T> writeListener) {
-		this.writeListener = writeListener;
+	public void setIngestQuery(String ingestQuery) {
+		Assert.hasText(ingestQuery, "'ingestQuery' must not be empty");
+		this.ingestQuery = ingestQuery;
 	}
 
 	public void setWriteOptions(WriteOptions writeOptions) {
 		this.writeOptions = writeOptions;
 	}
 
-	public void setAsync(boolean async) {
-		this.async = async;
-	}
-
 	public void setProducesReply(boolean producesReply) {
 		this.producesReply = producesReply;
 	}
 
+	public void setStatementExpression(Expression statementExpression) {
+		setStatementProcessor(new ExpressionEvaluatingMessageProcessor<Statement>(statementExpression,
+				Statement.class) {
+
+			@Override
+			protected StandardEvaluationContext getEvaluationContext() {
+				return (StandardEvaluationContext) CassandraMessageHandler.this.evaluationContext;
+			}
+
+		});
+	}
+
+	public void setQuery(String query) {
+		Assert.hasText(query, "'query' must not be empty");
+
+		final PreparedStatementCreator statementCreator = new CachedPreparedStatementCreator(query);
+
+		setStatementProcessor(new MessageProcessor<Statement>() {
+
+			@Override
+			public Statement processMessage(Message<?> message) {
+				PreparedStatement preparedStatement =
+						statementCreator.createPreparedStatement(cassandraTemplate.getSession());
+				ColumnDefinitions variables = preparedStatement.getVariables();
+				List<Object> values = new ArrayList<>(variables.size());
+				Map<String, Object> valueMap = new HashMap<>(variables.size());
+				for (ColumnDefinitions.Definition definition : variables) {
+					String name = definition.getName();
+					Object value = valueMap.get(name);
+					if (value == null) {
+						Expression expression = parameterExpressions.get(name);
+						Assert.state(expression != null, "No expression for parameter: " + name);
+						value = expression.getValue(evaluationContext, message);
+						valueMap.put(name, value);
+					}
+					values.add(value);
+				}
+				return preparedStatement.bind(values.toArray());
+			}
+
+		});
+	}
+
+	public void setParameterExpressions(Map<String, Expression> parameterExpressions) {
+		Assert.notEmpty(parameterExpressions, "'parameterExpressions' must not be empty.");
+		this.parameterExpressions.clear();
+		this.parameterExpressions.putAll(parameterExpressions);
+	}
+
+	public void setStatementProcessor(MessageProcessor<Statement> statementProcessor) {
+		Assert.notNull(statementProcessor, "'statementProcessor' must not be null.");
+		this.statementProcessor = statementProcessor;
+		this.queryType = Type.STATEMENT;
+	}
+
+	@Override
+	public void setIntegrationEvaluationContext(EvaluationContext evaluationContext) {
+		TypeLocator typeLocator = evaluationContext.getTypeLocator();
+		if (typeLocator instanceof StandardTypeLocator) {
+			/*
+			 * Register the Cassandra Query DSL package so they don't need a FQCN for QueryBuilder, for example.
+			 */
+			((StandardTypeLocator) typeLocator).registerImport("com.datastax.driver.core.querybuilder");
+		}
+		this.evaluationContext = evaluationContext;
+	}
+
+
 	@Override
 	public String getComponentType() {
-		return "cassandra:outbound-gateway";
+		return "cassandra:outbound-" + (this.producesReply ? "gateway" : "channel-adapter");
 	}
+
+	@Override
+	protected Object handleRequestMessage(Message<?> requestMessage) {
+		Object payload = requestMessage.getPayload();
+
+		Object result = payload;
+
+		Type queryType = this.queryType;
+
+		Statement statement = null;
+
+		if (payload instanceof Statement) {
+			statement = (Statement) payload;
+			queryType = Type.STATEMENT;
+		}
+
+		switch (queryType) {
+			case INSERT:
+				if (this.ingestQuery != null) {
+					Assert.isInstanceOf(List.class, payload,
+							"to perform 'ingest' the 'payload' must be of 'List<List<?>>' type.");
+					List<?> list = (List<?>) payload;
+					for (Object o : list) {
+						Assert.isInstanceOf(List.class, o,
+								"to perform 'ingest' the 'payload' must be of 'List<List<?>>' type.");
+					}
+					List<List<?>> rows = (List<List<?>>) payload;
+					this.cassandraTemplate.ingest(this.ingestQuery, rows, this.writeOptions);
+				}
+				else {
+					if (payload instanceof List) {
+						this.cassandraTemplate.insert((List<T>) payload, this.writeOptions);
+					}
+					else {
+						this.cassandraTemplate.insert(payload, this.writeOptions);
+					}
+				}
+				break;
+			case UPDATE:
+				if (payload instanceof List) {
+					this.cassandraTemplate.update((List<T>) payload, this.writeOptions);
+				}
+				else {
+					this.cassandraTemplate.update(payload, this.writeOptions);
+				}
+				break;
+			case DELETE:
+				if (payload instanceof List) {
+					this.cassandraTemplate.delete((List<T>) payload, this.writeOptions);
+				}
+				else {
+					this.cassandraTemplate.delete(payload, this.writeOptions);
+				}
+				break;
+			case STATEMENT:
+				if (statement == null) {
+					statement = this.statementProcessor.processMessage(requestMessage);
+				}
+
+				result = this.cassandraTemplate.executeAsynchronously(statement).getUninterruptibly();
+				break;
+		}
+
+		return this.producesReply ? result : null;
+	}
+
+	/**
+	 * Always return {@code false} to prevent a {@link com.datastax.driver.core.ResultSet}
+	 * draining on iteration.
+	 *
+	 * @param reply ignored.
+	 * @return {@code false}.
+	 */
+	@Override
+	protected boolean shouldSplitOutput(Iterable<?> reply) {
+		return false;
+	}
+
+
+	public enum Type {
+
+		INSERT, UPDATE, DELETE, STATEMENT;
+
+	}
+
 }

spring-integration-cassandra/src/main/java/org/springframework/integration/cassandra/support/CassandraOutboundGatewayType.java

@@ -1,24 +0,0 @@
-/*
- * Copyright 2015 the original author or authors.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.springframework.integration.cassandra.support;
-
-/**
- * @author Soby Chacko
- */
-public enum CassandraOutboundGatewayType {
-
-	INSERTING, UPDATING, DELETING;
-}

spring-integration-cassandra/src/main/java/org/springframework/integration/cassandra/support/package-info.java

@@ -1,4 +0,0 @@
-/**
- * Provides support classes for Cassandra Adapters.
- */
-package org.springframework.integration.cassandra.support;

spring-integration-cassandra/src/test/java/org/springframework/integration/cassandra/config/IntegrationTestConfig.java

@@ -36,8 +36,11 @@ import org.springframework.data.cassandra.config.java.AbstractCassandraConfigura
 @Configuration
 public class IntegrationTestConfig extends AbstractCassandraConfiguration {
 
+	public static final String HOST = "localhost";
+
 	//public static final SpringCassandraBuildProperties PROPS = new SpringCassandraBuildProperties();
 	public static final int PORT = 9043;//PROPS.getCassandraPort();
+
 	// public static final int RPC_PORT = PROPS.getCassandraRpcPort();
 
 	public String keyspaceName = randomKeyspaceName();

spring-integration-cassandra/src/test/java/org/springframework/integration/cassandra/outbound/CassandraMessageHandlerTests.java

@@ -16,44 +16,50 @@
 
 package org.springframework.integration.cassandra.outbound;
 
+import static org.hamcrest.Matchers.instanceOf;
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertThat;
 
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Date;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.UUID;
 
-import org.apache.cassandra.exceptions.ConfigurationException;
-import org.apache.thrift.transport.TTransportException;
-import org.cassandraunit.utils.EmbeddedCassandraServerHelper;
-import org.junit.AfterClass;
-import org.junit.Before;
-import org.junit.BeforeClass;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-
 import org.springframework.beans.factory.annotation.Autowired;
-import org.springframework.beans.factory.annotation.Qualifier;
 import org.springframework.cassandra.core.ConsistencyLevel;
 import org.springframework.cassandra.core.RetryPolicy;
 import org.springframework.cassandra.core.WriteOptions;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.data.cassandra.core.CassandraOperations;
+import org.springframework.expression.Expression;
+import org.springframework.expression.spel.standard.SpelExpressionParser;
 import org.springframework.integration.cassandra.config.IntegrationTestConfig;
 import org.springframework.integration.cassandra.test.domain.Book;
-import org.springframework.integration.channel.DirectChannel;
+import org.springframework.integration.channel.NullChannel;
+import org.springframework.integration.channel.QueueChannel;
+import org.springframework.integration.config.EnableIntegration;
 import org.springframework.integration.support.MessageBuilder;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageHandler;
-import org.springframework.messaging.MessagingException;
-import org.springframework.messaging.SubscribableChannel;
+import org.springframework.messaging.PollableChannel;
+import org.springframework.messaging.support.GenericMessage;
 import org.springframework.test.annotation.DirtiesContext;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
 
+import org.apache.cassandra.exceptions.ConfigurationException;
+import org.apache.thrift.transport.TTransportException;
+import org.cassandraunit.utils.EmbeddedCassandraServerHelper;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
 import com.datastax.driver.core.Cluster;
 import com.datastax.driver.core.ResultSet;
 import com.datastax.driver.core.Row;
@@ -70,7 +76,10 @@ import com.datastax.driver.core.querybuilder.Select;
 @DirtiesContext
 public class CassandraMessageHandlerTests {
 
+	private static final SpelExpressionParser PARSER = new SpelExpressionParser();
+
 	@Configuration
+	@EnableIntegration
 	public static class Config extends IntegrationTestConfig {
 
 		@Autowired
@@ -81,22 +90,21 @@ public class CassandraMessageHandlerTests {
 			return new String[]{Book.class.getPackage().getName()};
 		}
 
-		@Bean(name = "sync")
-		public MessageHandler cassandraOutboundGatewaySync() {
-			CassandraMessageHandler<Book> cassandraMessageHandler = new CassandraMessageHandler<Book>(template);
+		@Bean
+		public MessageHandler cassandraMessageHandler1() {
+			CassandraMessageHandler<Book> cassandraMessageHandler = new CassandraMessageHandler<>(this.template);
 			cassandraMessageHandler.setProducesReply(false);
 			return cassandraMessageHandler;
 		}
 
 		@Bean
-		SubscribableChannel messageChannel() {
-			return new DirectChannel();
+		public PollableChannel messageChannel() {
+			return new NullChannel();
 		}
 
-		@Bean(name = "async")
-		public MessageHandler cassandraOutboundGatewayAsync() {
-			CassandraMessageHandler<Book> cassandraMessageHandler = new CassandraMessageHandler<Book>(template);
-			cassandraMessageHandler.setAsync(true);
+		@Bean
+		public MessageHandler cassandraMessageHandler2() {
+			CassandraMessageHandler<Book> cassandraMessageHandler = new CassandraMessageHandler<>(this.template);
 
 			WriteOptions options = new WriteOptions();
 			options.setTtl(60);
@@ -110,36 +118,59 @@ public class CassandraMessageHandlerTests {
 			return cassandraMessageHandler;
 		}
 
-		@Bean(name = "ingest")
-		public MessageHandler cassandraOutboundGatewayIngest() {
-			CassandraMessageHandler<Book> cassandraMessageHandler = new CassandraMessageHandler<Book>(template);
+		@Bean
+		public MessageHandler cassandraMessageHandler3() {
+			CassandraMessageHandler<Book> cassandraMessageHandler =	new CassandraMessageHandler<>(this.template);
 			String cqlIngest = "insert into book (isbn, title, author, pages, saleDate, isInStock) values (?, ?, ?, ?, ?, ?)";
-			cassandraMessageHandler.setCqlIngest(cqlIngest);
+			cassandraMessageHandler.setIngestQuery(cqlIngest);
 			return cassandraMessageHandler;
 		}
+
+		@Bean
+		public PollableChannel resultChannel() {
+			return new QueueChannel();
+		}
+
+
+		@Bean
+		public MessageHandler cassandraMessageHandler4() {
+			CassandraMessageHandler<Book> cassandraMessageHandler = new CassandraMessageHandler<>(this.template);
+			//TODO https://jira.spring.io/browse/DATACASS-213
+			//cassandraMessageHandler.setQuery("SELECT * FROM book WHERE author = :author limit :size");
+			cassandraMessageHandler.setQuery("SELECT * FROM book limit :size");
+
+			Map<String, Expression> params = new HashMap<>();
+			params.put("author", PARSER.parseExpression("payload"));
+			params.put("size", PARSER.parseExpression("headers.limit"));
+
+			cassandraMessageHandler.setParameterExpressions(params);
+
+			cassandraMessageHandler.setOutputChannel(resultChannel());
+			cassandraMessageHandler.setProducesReply(true);
+			return cassandraMessageHandler;
+		}
+
 	}
 
 	@Autowired
-	@Qualifier("sync")
-	public MessageHandler messageHandlerSync;
+	public MessageHandler cassandraMessageHandler1;
 
 	@Autowired
-	@Qualifier("async")
-	public MessageHandler messageHandlerAsync;
+	public MessageHandler cassandraMessageHandler2;
 
 	@Autowired
-	@Qualifier("ingest")
-	public MessageHandler messageHandlerIngest;
+	public MessageHandler cassandraMessageHandler3;
+
+	@Autowired
+	public MessageHandler cassandraMessageHandler4;
 
 	@Autowired
 	public CassandraOperations template;
 
 	@Autowired
-	public SubscribableChannel channel;
+	public PollableChannel resultChannel;
 
-	protected static String CASSANDRA_CONFIG = "spring-cassandra.yaml";
-
-	protected static String CASSANDRA_HOST = "localhost";
+	protected static final String CASSANDRA_CONFIG = "spring-cassandra.yaml";
 
 	/**
 	 * The {@link Cluster} that's connected to Cassandra.
@@ -154,9 +185,12 @@ public class CassandraMessageHandlerTests {
 	@BeforeClass
 	public static void startCassandra() throws TTransportException, IOException, InterruptedException,
 			ConfigurationException {
-
-		EmbeddedCassandraServerHelper.startEmbeddedCassandra(CASSANDRA_CONFIG);
-		ensureClusterConnection();
+		EmbeddedCassandraServerHelper.startEmbeddedCassandra(CASSANDRA_CONFIG, "build/embeddedCassandra");
+		cluster = Cluster.builder()
+				.addContactPoint(IntegrationTestConfig.HOST)
+				.withPort(IntegrationTestConfig.PORT)
+				.build();
+		system = cluster.connect();
 	}
 
 	@AfterClass
@@ -165,31 +199,6 @@ public class CassandraMessageHandlerTests {
 		EmbeddedCassandraServerHelper.cleanEmbeddedCassandra();
 	}
 
-	@Before
-	public void setup() {
-		channel.subscribe(new MessageHandler() {
-			@Override
-			public void handleMessage(Message<?> message) throws MessagingException {
-			}
-		});
-	}
-
-	public static Cluster cluster() {
-		return Cluster.builder().addContactPoint(CASSANDRA_HOST).withPort(IntegrationTestConfig.PORT).build();
-	}
-
-
-	public static void ensureClusterConnection() {
-		// check cluster
-		if (cluster == null) {
-			cluster = cluster();
-		}
-
-		if (system == null) {
-			system = cluster.connect();
-		}
-	}
-
 	@Test
 	public void testBasicCassandraInsert() throws Exception {
 		Book b1 = new Book();
@@ -201,44 +210,44 @@ public class CassandraMessageHandlerTests {
 		b1.setInStock(true);
 
 		Message<Book> message = MessageBuilder.withPayload(b1).build();
-		messageHandlerSync.handleMessage(message);
+		this.cassandraMessageHandler1.handleMessage(message);
 
-		int n = 0;
 		Select select = QueryBuilder.select().all().from("book");
-		List<Book> books;
-		while ((books = template.select(select, Book.class)).isEmpty() && n++ < 10) {
-			Thread.sleep(100);
-		}
-		assertTrue(n < 10);
-		assertEquals(books.size(), 1);
+		List<Book> books = this.template.select(select, Book.class);
+		assertEquals(1, books.size());
 
-		template.delete(b1);
+		this.template.delete(b1);
 	}
 
 	@Test
-	public void testCassandraBatchInsert() throws Exception {
+	public void testCassandraBatchInsertAndSelectStatement() throws Exception {
 		List<Book> books = getBookList(5);
-		Message<List<Book>> message = MessageBuilder.withPayload(books).build();
-		messageHandlerAsync.handleMessage(message);
 
-		int n = 0;
-		Select select = QueryBuilder.select().all().from("book");
-		while ((books = template.select(select, Book.class)).isEmpty() && n++ < 10) {
-			Thread.sleep(100);
-		}
-		assertTrue(n < 10);
-		assertEquals(books.size(), 5);
+		this.cassandraMessageHandler2.handleMessage(new GenericMessage<>(books));
 
-		template.delete(books);
+		Message<?> message = MessageBuilder.withPayload("Cassandra Guru")
+				.setHeader("limit", 2)
+				.build();
+		this.cassandraMessageHandler4.handleMessage(message);
+
+		Message<?> receive = this.resultChannel.receive(10000);
+		assertNotNull(receive);
+		assertThat(receive.getPayload(), instanceOf(ResultSet.class));
+		ResultSet resultSet = (ResultSet) receive.getPayload();
+		assertNotNull(resultSet);
+		List<Row> rows = resultSet.all();
+		assertEquals(2, rows.size());
+
+		this.cassandraMessageHandler1.handleMessage(new GenericMessage<>(QueryBuilder.truncate("book")));
 	}
 
 	@Test
 	public void testCassandraBatchIngest() throws Exception {
 		List<Book> books = getBookList(5);
-		List<List<?>> ingestBooks = new ArrayList<List<?>>();
+		List<List<?>> ingestBooks = new ArrayList<>();
 		for (Book b : books) {
 
-			List<Object> l = new ArrayList<Object>();
+			List<Object> l = new ArrayList<>();
 			l.add(b.getIsbn());
 			l.add(b.getTitle());
 			l.add(b.getAuthor());
@@ -249,22 +258,18 @@ public class CassandraMessageHandlerTests {
 		}
 
 		Message<List<List<?>>> message = MessageBuilder.withPayload(ingestBooks).build();
-		messageHandlerIngest.handleMessage(message);
+		this.cassandraMessageHandler3.handleMessage(message);
 
-		int n = 0;
 		Select select = QueryBuilder.select().all().from("book");
-		while ((books = template.select(select, Book.class)).isEmpty() && n++ < 10) {
-			Thread.sleep(100);
-		}
-		assertTrue(n < 10);
-		assertEquals(books.size(), 5);
+		books = this.template.select(select, Book.class);
+		assertEquals(5, books.size());
 
-		template.delete(books);
+		this.template.delete(books);
 	}
 
 	private List<Book> getBookList(int numBooks) {
 
-		List<Book> books = new ArrayList<Book>();
+		List<Book> books = new ArrayList<>();
 
 		Book b;
 		for (int i = 0; i < numBooks; i++) {

spring-integration-cassandra/src/test/java/org/springframework/integration/cassandra/test/domain/Book.java

@@ -18,6 +18,7 @@ package org.springframework.integration.cassandra.test.domain;
 
 import java.util.Date;
 
+import org.springframework.data.cassandra.mapping.Indexed;
 import org.springframework.data.cassandra.mapping.PrimaryKey;
 import org.springframework.data.cassandra.mapping.Table;
 
@@ -33,6 +34,7 @@ public class Book {
 
 	private String title;
 
+	@Indexed
 	private String author;
 
 	private int pages;