Support configurable batch types for batch operations.

Closes #1174
Original pull request: #1175.

Author: samueldlightfoot

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/CassandraBatchTemplate.java

@@ -39,13 +39,14 @@
  * @author Mark Paluch
  * @author John Blum
  * @author Anup Sabbi
+ * @author Sam Lightfoot
  * @since 1.5
  */
 class CassandraBatchTemplate implements CassandraBatchOperations {
 
 	private final AtomicBoolean executed = new AtomicBoolean();
 
-	private final BatchStatementBuilder batch = BatchStatement.builder(BatchType.LOGGED);
+	private final BatchStatementBuilder batch;
 
 	private final CassandraConverter converter;
 
@@ -61,10 +62,23 @@ class CassandraBatchTemplate implements CassandraBatchOperations {
 	 * @param operations must not be {@literal null}.
 	 */
 	CassandraBatchTemplate(CassandraOperations operations) {
+		this(operations, BatchType.LOGGED);
+	}
+
+	/**
+	 * Create a new {@link CassandraBatchTemplate} given {@link CassandraOperations} and {@link BatchType}.
+	 *
+	 * @param operations must not be {@literal null}.
+	 * @param batchType must not be {@literal null}.
+	 * @since 3.3.0
+	 */
+	CassandraBatchTemplate(CassandraOperations operations, BatchType batchType) {
 
 		Assert.notNull(operations, "CassandraOperations must not be null");
+		Assert.notNull(batchType, "BatchType must not be null");
 
 		this.operations = operations;
+		this.batch = BatchStatement.builder(batchType);
 		this.converter = operations.getConverter();
 		this.mappingContext = this.converter.getMappingContext();
 		this.statementFactory = new StatementFactory(new UpdateMapper(converter));

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/CassandraOperations.java

@@ -31,6 +31,7 @@
 import org.springframework.lang.Nullable;
 
 import com.datastax.oss.driver.api.core.CqlIdentifier;
+import com.datastax.oss.driver.api.core.cql.BatchType;
 import com.datastax.oss.driver.api.core.cql.ResultSet;
 import com.datastax.oss.driver.api.core.cql.Statement;
 
@@ -42,6 +43,7 @@
  * @author David Webb
  * @author Matthew Adams
  * @author Mark Paluch
+ * @author Sam Lightfoot
  * @see CassandraTemplate
  * @see CqlOperations
  * @see Statement
@@ -52,12 +54,22 @@ public interface CassandraOperations extends FluentCassandraOperations {
 
 	/**
 	 * Returns a new {@link CassandraBatchOperations}. Each {@link CassandraBatchOperations} instance can be executed only
-	 * once so you might want to obtain new {@link CassandraBatchOperations} instances for each batch.
+	 * once, so you might want to obtain new {@link CassandraBatchOperations} instances for each batch.
 	 *
 	 * @return a new {@link CassandraBatchOperations} associated with the given entity class.
 	 */
 	CassandraBatchOperations batchOps();
 
+	/**
+	 * Returns a new {@link CassandraBatchOperations}. Each {@link CassandraBatchOperations} instance can be executed only
+	 * once, so you might want to obtain new {@link CassandraBatchOperations} instances for each batch.
+	 *
+	 * @param batchType must not be {@literal null}.
+	 * @return a new {@link ReactiveCassandraBatchOperations} associated with the given entity class.
+	 * @since 3.3.0
+	 */
+	CassandraBatchOperations batchOps(BatchType batchType);
+
 	/**
 	 * Expose the underlying {@link CqlOperations} to allow CQL operations.
 	 *

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/CassandraTemplate.java

@@ -73,6 +73,7 @@
 import com.datastax.oss.driver.api.core.CqlSession;
 import com.datastax.oss.driver.api.core.DriverException;
 import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
+import com.datastax.oss.driver.api.core.cql.BatchType;
 import com.datastax.oss.driver.api.core.cql.BoundStatement;
 import com.datastax.oss.driver.api.core.cql.PreparedStatement;
 import com.datastax.oss.driver.api.core.cql.ResultSet;
@@ -109,6 +110,7 @@
  * @author Mark Paluch
  * @author John Blum
  * @author Lukasz Antoniak
+ * @author Sam Lightfoot
  * @see org.springframework.data.cassandra.core.CassandraOperations
  * @since 2.0
  */
@@ -202,6 +204,14 @@ public CassandraBatchOperations batchOps() {
 		return new CassandraBatchTemplate(this);
 	}
 
+	/* (non-Javadoc)
+	 * @see org.springframework.data.cassandra.core.CassandraOperations#batchOps(com.datastax.oss.driver.api.core.cql.BatchType)
+	 */
+	@Override
+	public CassandraBatchOperations batchOps(BatchType batchType) {
+		return new CassandraBatchTemplate(this, batchType);
+	}
+
 	/* (non-Javadoc)
 	 * @see org.springframework.context.ApplicationEventPublisherAware#setApplicationEventPublisher(org.springframework.context.ApplicationEventPublisher)
 	 */

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/ReactiveCassandraBatchTemplate.java

@@ -46,13 +46,14 @@
  *
  * @author Oleh Dokuka
  * @author Mark Paluch
+ * @author Sam Lightfoot
  * @since 2.1
  */
 class ReactiveCassandraBatchTemplate implements ReactiveCassandraBatchOperations {
 
 	private final AtomicBoolean executed = new AtomicBoolean();
 
-	private final BatchStatementBuilder batch = BatchStatement.builder(BatchType.LOGGED);
+	private final BatchStatementBuilder batch;
 
 	private final CassandraConverter converter;
 
@@ -70,10 +71,23 @@ class ReactiveCassandraBatchTemplate implements ReactiveCassandraBatchOperations
 	 * @param operations must not be {@literal null}.
 	 */
 	ReactiveCassandraBatchTemplate(ReactiveCassandraOperations operations) {
+		this(operations, BatchType.LOGGED);
+	}
+
+	/**
+	 * Create a new {@link CassandraBatchTemplate} given {@link CassandraOperations} and {@link BatchType}.
+	 *
+	 * @param operations must not be {@literal null}.
+	 * @param batchType must not be {@literal null}.
+	 * @since 3.3.0
+	 */
+	ReactiveCassandraBatchTemplate(ReactiveCassandraOperations operations, BatchType batchType) {
 
 		Assert.notNull(operations, "CassandraOperations must not be null");
+		Assert.notNull(batchType, "BatchType must not be null");
 
 		this.operations = operations;
+		this.batch = BatchStatement.builder(batchType);
 		this.converter = operations.getConverter();
 		this.mappingContext = this.converter.getMappingContext();
 		this.statementFactory = new StatementFactory(new UpdateMapper(converter));
@@ -155,7 +169,6 @@ public Mono<WriteResult> execute() {
 	public ReactiveCassandraBatchOperations withTimestamp(long timestamp) {
 
 		assertNotExecuted();
-
 		this.batch.setQueryTimestamp(timestamp);
 
 		return this;

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/ReactiveCassandraOperations.java

@@ -29,6 +29,7 @@
 import org.springframework.data.cassandra.core.query.Update;
 import org.springframework.data.domain.Slice;
 
+import com.datastax.oss.driver.api.core.cql.BatchType;
 import com.datastax.oss.driver.api.core.cql.Statement;
 
 /**
@@ -38,6 +39,7 @@
  * @author Mark Paluch
  * @author Hleb Albau
  * @author Oleh Dokuka
+ * @author Sam Lightfoot
  * @since 2.0
  * @see ReactiveCassandraTemplate
  * @see ReactiveCqlOperations
@@ -51,14 +53,25 @@ public interface ReactiveCassandraOperations extends ReactiveFluentCassandraOper
 
 	/**
 	 * Returns a new {@link ReactiveCassandraBatchOperations}. Each {@link ReactiveCassandraBatchOperations} instance can
-	 * be executed only once so you might want to obtain new {@link ReactiveCassandraBatchOperations} instances for each
+	 * be executed only once, so you might want to obtain new {@link ReactiveCassandraBatchOperations} instances for each
 	 * batch.
 	 *
 	 * @return a new {@link ReactiveCassandraBatchOperations} associated with the given entity class.
 	 * @since 2.1
 	 */
 	ReactiveCassandraBatchOperations batchOps();
 
+	/**
+	 * Returns a new {@link ReactiveCassandraBatchOperations}. Each {@link ReactiveCassandraBatchOperations} instance can
+	 * be executed only once, so you might want to obtain new {@link ReactiveCassandraBatchOperations} instances for each
+	 * batch.
+	 *
+	 * @param batchType must not be {@literal null}.
+	 * @return a new {@link ReactiveCassandraBatchOperations} associated with the given entity class.
+	 * @since 3.3.0
+	 */
+	ReactiveCassandraBatchOperations batchOps(BatchType batchType);
+
 	/**
 	 * Expose the underlying {@link ReactiveCqlOperations} to allow CQL operations.
 	 *

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/ReactiveCassandraTemplate.java

@@ -69,6 +69,7 @@
 import com.datastax.oss.driver.api.core.DriverException;
 import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
 import com.datastax.oss.driver.api.core.context.DriverContext;
+import com.datastax.oss.driver.api.core.cql.BatchType;
 import com.datastax.oss.driver.api.core.cql.BoundStatement;
 import com.datastax.oss.driver.api.core.cql.PreparedStatement;
 import com.datastax.oss.driver.api.core.cql.Row;
@@ -105,6 +106,7 @@
  * @author John Blum
  * @author Lukasz Antoniak
  * @author Hleb Albau
+ * @author Sam Lightfoot
  * @since 2.0
  */
 public class ReactiveCassandraTemplate
@@ -199,6 +201,14 @@ public ReactiveCassandraBatchOperations batchOps() {
 		return new ReactiveCassandraBatchTemplate(this);
 	}
 
+	/* (non-Javadoc)
+	 * @see org.springframework.data.cassandra.core.ReactiveCassandraOperations#batchOps(com.datastax.oss.driver.api.core.cql.BatchType)
+	 */
+	@Override
+	public ReactiveCassandraBatchOperations batchOps(BatchType batchType) {
+		return new ReactiveCassandraBatchTemplate(this, batchType);
+	}
+
 	/* (non-Javadoc)
 	 * @see org.springframework.context.ApplicationEventPublisherAware#setApplicationEventPublisher(org.springframework.context.ApplicationEventPublisher)
 	 */