Introduce customization hook for PreparedStatementHandler.

CassandraTemplate and its asynchronous and reactive variants now define a createPreparedStatementHandler() template method to customize CQL preparation and binding.

Closes #1237

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

@@ -17,6 +17,7 @@ package org.springframework.data.cassandra.core;
 
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.CompletionStage;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -33,6 +34,7 @@ import org.springframework.context.ApplicationEventPublisherAware;
 import org.springframework.dao.DataAccessException;
 import org.springframework.dao.OptimisticLockingFailureException;
 import org.springframework.dao.support.DataAccessUtils;
+import org.springframework.dao.support.PersistenceExceptionTranslator;
 import org.springframework.data.cassandra.SessionFactory;
 import org.springframework.data.cassandra.core.EntityOperations.AdaptibleEntity;
 import org.springframework.data.cassandra.core.convert.CassandraConverter;
@@ -729,6 +731,19 @@ public class AsyncCassandraTemplate
 	// Implementation hooks and utility methods
 	// -------------------------------------------------------------------------
 
+	/**
+	 * Create a new statement-based {@link AsyncPreparedStatementHandler} using the statement passed in.
+	 * <p>
+	 * This method allows for the creation to be overridden by subclasses.
+	 *
+	 * @param statement the statement to be prepared.
+	 * @return the new {@link PreparedStatementHandler} to use.
+	 * @since 3.3.3
+	 */
+	protected AsyncPreparedStatementHandler createPreparedStatementHandler(Statement<?> statement) {
+		return new PreparedStatementHandler(statement, exceptionTranslator);
+	}
+
 	private <T> ListenableFuture<EntityWriteResult<T>> executeSave(T entity, CqlIdentifier tableName,
 			SimpleStatement statement) {
 
@@ -781,7 +796,7 @@ public class AsyncCassandraTemplate
 
 		if (PreparedStatementDelegate.canPrepare(isUsePreparedStatements(), statement, log)) {
 
-			PreparedStatementHandler statementHandler = new PreparedStatementHandler(statement);
+			AsyncPreparedStatementHandler statementHandler = createPreparedStatementHandler(statement);
 			return getAsyncCqlOperations().query(statementHandler, statementHandler, rowMapper);
 		}
 
@@ -792,7 +807,7 @@ public class AsyncCassandraTemplate
 
 		if (PreparedStatementDelegate.canPrepare(isUsePreparedStatements(), statement, log)) {
 
-			PreparedStatementHandler statementHandler = new PreparedStatementHandler(statement);
+			AsyncPreparedStatementHandler statementHandler = createPreparedStatementHandler(statement);
 			return getAsyncCqlOperations().query(statementHandler, statementHandler, callbackHandler);
 		}
 
@@ -807,7 +822,7 @@ public class AsyncCassandraTemplate
 
 		if (PreparedStatementDelegate.canPrepare(isUsePreparedStatements(), statement, log)) {
 
-			PreparedStatementHandler statementHandler = new PreparedStatementHandler(statement);
+			AsyncPreparedStatementHandler statementHandler = createPreparedStatementHandler(statement);
 			return getAsyncCqlOperations().query(statementHandler, statementHandler,
 					(AsyncResultSetExtractor<T>) resultSet -> new AsyncResult<>(mappingFunction.apply(resultSet)));
 		}
@@ -925,24 +940,48 @@ public class AsyncCassandraTemplate
 		}
 	}
 
+	/**
+	 * General callback interface used to create and bind prepared CQL statements.
+	 * <p>
+	 * This interface prepares the CQL statement and sets values on a {@link PreparedStatement} as union-type comprised
+	 * from {@link AsyncPreparedStatementCreator}, {@link PreparedStatementBinder}, and {@link CqlProvider}.
+	 *
+	 * @since 3.3.3
+	 */
+	public interface AsyncPreparedStatementHandler
+			extends AsyncPreparedStatementCreator, PreparedStatementBinder, CqlProvider {
+
+	}
+
 	/**
 	 * Utility class to prepare a {@link SimpleStatement} and bind values associated with the statement to a
 	 * {@link BoundStatement}.
 	 *
 	 * @since 3.2
 	 */
-	private class PreparedStatementHandler
+	public static class PreparedStatementHandler implements AsyncPreparedStatementHandler {
-			implements AsyncPreparedStatementCreator, PreparedStatementBinder, CqlProvider {
 
 		private final SimpleStatement statement;
+		private final PersistenceExceptionTranslator exceptionTranslator;
 
-		public PreparedStatementHandler(Statement<?> statement) {
+		public PreparedStatementHandler(Statement<?> statement, PersistenceExceptionTranslator exceptionTranslator) {
 			this.statement = PreparedStatementDelegate.getStatementForPrepare(statement);
+			this.exceptionTranslator = exceptionTranslator;
 		}
 
 		@Override
 		public ListenableFuture<PreparedStatement> createPreparedStatement(CqlSession session) throws DriverException {
-			return new CassandraFutureAdapter<>(session.prepareAsync(statement), exceptionTranslator);
+			return new CassandraFutureAdapter<>(doPrepare(session), exceptionTranslator);
+		}
+
+		/**
+		 * Invokes the statement preparation.
+		 *
+		 * @param session
+		 * @return
+		 */
+		protected CompletionStage<PreparedStatement> doPrepare(CqlSession session) {
+			return session.prepareAsync(statement);
 		}
 
 		@Override

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

@@ -775,6 +775,19 @@ public class CassandraTemplate implements CassandraOperations, ApplicationEventP
 	// Implementation hooks and utility methods
 	// -------------------------------------------------------------------------
 
+	/**
+	 * Create a new statement-based {@link PreparedStatementHandler} using the statement passed in.
+	 * <p>
+	 * This method allows for the creation to be overridden by subclasses.
+	 *
+	 * @param statement the statement to be prepared.
+	 * @return the new {@link PreparedStatementHandler} to use.
+	 * @since 3.3.3
+	 */
+	protected PreparedStatementHandler createPreparedStatementHandler(Statement<?> statement) {
+		return new PreparedStatementHandler(statement);
+	}
+
 	private <T> EntityWriteResult<T> executeSave(T entity, CqlIdentifier tableName, SimpleStatement statement) {
 		return executeSave(entity, tableName, statement, ignore -> {});
 	}
@@ -811,7 +824,7 @@ public class CassandraTemplate implements CassandraOperations, ApplicationEventP
 
 		if (PreparedStatementDelegate.canPrepare(isUsePreparedStatements(), statement, log)) {
 
-			PreparedStatementHandler statementHandler = new PreparedStatementHandler(statement);
+			PreparedStatementHandler statementHandler = createPreparedStatementHandler(statement);
 			return getCqlOperations().query(statementHandler, statementHandler, rowMapper);
 		}
 
@@ -826,7 +839,7 @@ public class CassandraTemplate implements CassandraOperations, ApplicationEventP
 
 		if (PreparedStatementDelegate.canPrepare(isUsePreparedStatements(), statement, log)) {
 
-			PreparedStatementHandler statementHandler = new PreparedStatementHandler(statement);
+			PreparedStatementHandler statementHandler = createPreparedStatementHandler(statement);
 			return getCqlOperations().queryForStream(statementHandler, statementHandler, rowMapper);
 		}
 
@@ -845,7 +858,7 @@ public class CassandraTemplate implements CassandraOperations, ApplicationEventP
 
 		if (PreparedStatementDelegate.canPrepare(isUsePreparedStatements(), statement, log)) {
 
-			PreparedStatementHandler statementHandler = new PreparedStatementHandler(statement);
+			PreparedStatementHandler statementHandler = createPreparedStatementHandler(statement);
 			return getCqlOperations().query(statementHandler, statementHandler, mappingFunction::apply);
 		}
 

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

@@ -15,7 +15,6 @@
  */
 package org.springframework.data.cassandra.core;
 
-import org.springframework.data.projection.EntityProjection;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.SynchronousSink;
@@ -62,6 +61,7 @@ import org.springframework.data.domain.Slice;
 import org.springframework.data.domain.SliceImpl;
 import org.springframework.data.mapping.callback.EntityCallbacks;
 import org.springframework.data.mapping.callback.ReactiveEntityCallbacks;
+import org.springframework.data.projection.EntityProjection;
 import org.springframework.data.projection.ProjectionFactory;
 import org.springframework.data.projection.SpelAwareProxyProjectionFactory;
 import org.springframework.lang.Nullable;
@@ -753,6 +753,19 @@ public class ReactiveCassandraTemplate
 	// Implementation hooks and utility methods
 	// -------------------------------------------------------------------------
 
+	/**
+	 * Create a new statement-based {@link ReactivePreparedStatementHandler} using the statement passed in.
+	 * <p>
+	 * This method allows for the creation to be overridden by subclasses.
+	 *
+	 * @param statement the statement to be prepared.
+	 * @return the new {@link PreparedStatementHandler} to use.
+	 * @since 3.3.3
+	 */
+	protected ReactivePreparedStatementHandler createPreparedStatementHandler(Statement<?> statement) {
+		return new PreparedStatementHandler(statement);
+	}
+
 	private <T> Mono<EntityWriteResult<T>> executeSave(T entity, CqlIdentifier tableName, SimpleStatement statement) {
 		return executeSave(entity, tableName, statement, (writeResult, sink) -> sink.next(writeResult));
 	}
@@ -789,7 +802,7 @@ public class ReactiveCassandraTemplate
 
 		if (PreparedStatementDelegate.canPrepare(isUsePreparedStatements(), statement, log)) {
 
-			PreparedStatementHandler statementHandler = new PreparedStatementHandler(statement);
+			ReactivePreparedStatementHandler statementHandler = createPreparedStatementHandler(statement);
 			return getReactiveCqlOperations().query(statementHandler, statementHandler, rowMapper);
 		}
 
@@ -800,7 +813,7 @@ public class ReactiveCassandraTemplate
 
 		if (PreparedStatementDelegate.canPrepare(isUsePreparedStatements(), statement, log)) {
 
-			PreparedStatementHandler statementHandler = new PreparedStatementHandler(statement);
+			ReactivePreparedStatementHandler statementHandler = createPreparedStatementHandler(statement);
 			return getReactiveCqlOperations()
 					.query(statementHandler, statementHandler, rs -> Mono.just(mappingFunction.apply(rs))).next();
 		}
@@ -813,7 +826,7 @@ public class ReactiveCassandraTemplate
 
 		if (PreparedStatementDelegate.canPrepare(isUsePreparedStatements(), statement, log)) {
 
-			PreparedStatementHandler statementHandler = new PreparedStatementHandler(statement);
+			ReactivePreparedStatementHandler statementHandler = createPreparedStatementHandler(statement);
 			return getReactiveCqlOperations().query(statementHandler, statementHandler, mappingFunction::apply).next();
 		}
 
@@ -915,14 +928,26 @@ public class ReactiveCassandraTemplate
 		return Mono.just(object);
 	}
 
+	/**
+	 * General callback interface used to create and bind prepared CQL statements.
+	 * <p>
+	 * This interface prepares the CQL statement and sets values on a {@link PreparedStatement} as union-type comprised
+	 * from {@link ReactivePreparedStatementCreator}, {@link PreparedStatementBinder}, and {@link CqlProvider}.
+	 *
+	 * @since 3.3.3
+	 */
+	public interface ReactivePreparedStatementHandler
+			extends ReactivePreparedStatementCreator, PreparedStatementBinder, CqlProvider {
+
+	}
+
 	/**
 	 * Utility class to prepare a {@link SimpleStatement} and bind values associated with the statement to a
 	 * {@link BoundStatement}.
 	 *
 	 * @since 3.2
 	 */
-	private static class PreparedStatementHandler
+	public static class PreparedStatementHandler implements ReactivePreparedStatementHandler {
-			implements ReactivePreparedStatementCreator, PreparedStatementBinder, CqlProvider {
 
 		private final SimpleStatement statement;