Add consistencyLevel setter for batchOps

wtetsu · wtetsu · commit d806d32e38da · 2021-11-15T20:37:41.000+09:00
diff --git a/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/CassandraBatchTemplate.java b/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/CassandraBatchTemplate.java
@@ -15,6 +15,7 @@
  */
 package org.springframework.data.cassandra.core;
 
+import com.datastax.oss.driver.api.core.ConsistencyLevel;
 import java.util.Arrays;
 import java.util.concurrent.atomic.AtomicBoolean;
 
@@ -42,6 +43,7 @@
  * @since 1.5
  */
 class CassandraBatchTemplate implements CassandraBatchOperations {
+//public class CassandraBatchTemplate implements CassandraBatchOperations {
 
 	private final AtomicBoolean executed = new AtomicBoolean();
 
@@ -60,8 +62,7 @@ class CassandraBatchTemplate implements CassandraBatchOperations {
 	 *
 	 * @param operations must not be {@literal null}.
 	 */
-	CassandraBatchTemplate(CassandraOperations operations) {
-
+	public CassandraBatchTemplate(CassandraOperations operations) {
 		Assert.notNull(operations, "CassandraOperations must not be null");
 
 		this.operations = operations;
@@ -70,6 +71,16 @@ class CassandraBatchTemplate implements CassandraBatchOperations {
 		this.statementFactory = new StatementFactory(new UpdateMapper(converter));
 	}
 
+	/**
+	 * Create a new {@link CassandraBatchTemplate} given {@link CassandraOperations} with {@link ConsistencyLevel}.
+	 *
+	 * @param operations must not be {@literal null}.
+	 */
+	public CassandraBatchTemplate(CassandraOperations operations, ConsistencyLevel consistencyLevel) {
+		this(operations);
+		batch.setConsistencyLevel(consistencyLevel);
+	}
+
 	/**
 	 * Return a reference to the configured {@link CassandraConverter} used to map {@link Object Objects} to
 	 * {@link com.datastax.driver.core.Row Rows}.
diff --git a/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/CassandraOperations.java b/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/CassandraOperations.java
@@ -15,6 +15,7 @@
  */
 package org.springframework.data.cassandra.core;
 
+import com.datastax.oss.driver.api.core.ConsistencyLevel;
 import java.util.Iterator;
 import java.util.List;
 import java.util.stream.Stream;
@@ -58,6 +59,14 @@ public interface CassandraOperations extends FluentCassandraOperations {
 	 */
 	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.
+	 *
+	 * @return a new {@link CassandraBatchOperations} associated with the given entity class.
+	 */
+	CassandraBatchOperations batchOps(ConsistencyLevel consistencyLevel);
+
 	/**
 	 * Expose the underlying {@link CqlOperations} to allow CQL operations.
 	 *
diff --git a/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/CassandraTemplate.java b/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/CassandraTemplate.java
@@ -15,6 +15,7 @@
  */
 package org.springframework.data.cassandra.core;
 
+import com.datastax.oss.driver.api.core.ConsistencyLevel;
 import java.util.List;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -202,6 +203,12 @@ public CassandraBatchOperations batchOps() {
 		return new CassandraBatchTemplate(this);
 	}
 
+	@Override
+	public CassandraBatchOperations batchOps(ConsistencyLevel consistencyLevel) {
+		return new CassandraBatchTemplate(this, consistencyLevel);
+	}
+
+
 	/* (non-Javadoc)
 	 * @see org.springframework.context.ApplicationEventPublisherAware#setApplicationEventPublisher(org.springframework.context.ApplicationEventPublisher)
 	 */
