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
This commit is contained in:
Mark Paluch
2022-03-04 13:59:19 +01:00
parent f386f82e5a
commit adb4c8c8d8
3 changed files with 115 additions and 38 deletions

View File

@@ -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;
@@ -51,6 +52,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;
@@ -824,6 +826,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) {
@@ -876,7 +891,7 @@ public class AsyncCassandraTemplate
if (PreparedStatementDelegate.canPrepare(isUsePreparedStatements(), statement, logger)) {
PreparedStatementHandler statementHandler = new PreparedStatementHandler(statement);
AsyncPreparedStatementHandler statementHandler = createPreparedStatementHandler(statement);
return getAsyncCqlOperations().query(statementHandler, statementHandler, rowMapper);
}
@@ -887,7 +902,7 @@ public class AsyncCassandraTemplate
if (PreparedStatementDelegate.canPrepare(isUsePreparedStatements(), statement, logger)) {
PreparedStatementHandler statementHandler = new PreparedStatementHandler(statement);
AsyncPreparedStatementHandler statementHandler = createPreparedStatementHandler(statement);
return getAsyncCqlOperations().query(statementHandler, statementHandler, callbackHandler);
}
@@ -902,7 +917,7 @@ public class AsyncCassandraTemplate
if (PreparedStatementDelegate.canPrepare(isUsePreparedStatements(), statement, logger)) {
PreparedStatementHandler statementHandler = new PreparedStatementHandler(statement);
AsyncPreparedStatementHandler statementHandler = createPreparedStatementHandler(statement);
return getAsyncCqlOperations().query(statementHandler, statementHandler,
(AsyncResultSetExtractor<T>) resultSet -> new AsyncResult<>(mappingFunction.apply(resultSet)));
}
@@ -1029,19 +1044,33 @@ 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
implements AsyncPreparedStatementCreator, PreparedStatementBinder, CqlProvider {
public static class PreparedStatementHandler implements AsyncPreparedStatementHandler {
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;
}
/*
@@ -1050,7 +1079,17 @@ public class AsyncCassandraTemplate
*/
@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);
}
/*

View File

@@ -885,6 +885,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 -> {});
}
@@ -921,7 +934,7 @@ public class CassandraTemplate implements CassandraOperations, ApplicationEventP
if (PreparedStatementDelegate.canPrepare(isUsePreparedStatements(), statement, logger)) {
PreparedStatementHandler statementHandler = new PreparedStatementHandler(statement);
PreparedStatementHandler statementHandler = createPreparedStatementHandler(statement);
return getCqlOperations().query(statementHandler, statementHandler, rowMapper);
}
@@ -936,7 +949,7 @@ public class CassandraTemplate implements CassandraOperations, ApplicationEventP
if (PreparedStatementDelegate.canPrepare(isUsePreparedStatements(), statement, logger)) {
PreparedStatementHandler statementHandler = new PreparedStatementHandler(statement);
PreparedStatementHandler statementHandler = createPreparedStatementHandler(statement);
return getCqlOperations().queryForStream(statementHandler, statementHandler, rowMapper);
}
@@ -955,7 +968,7 @@ public class CassandraTemplate implements CassandraOperations, ApplicationEventP
if (PreparedStatementDelegate.canPrepare(isUsePreparedStatements(), statement, logger)) {
PreparedStatementHandler statementHandler = new PreparedStatementHandler(statement);
PreparedStatementHandler statementHandler = createPreparedStatementHandler(statement);
return getCqlOperations().query(statementHandler, statementHandler, mappingFunction::apply);
}

View File

@@ -15,33 +15,17 @@
*/
package org.springframework.data.cassandra.core;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.SynchronousSink;
import java.util.Collections;
import java.util.function.BiConsumer;
import java.util.function.Function;
import com.datastax.oss.driver.api.core.CqlIdentifier;
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;
import com.datastax.oss.driver.api.core.cql.SimpleStatement;
import com.datastax.oss.driver.api.core.cql.Statement;
import com.datastax.oss.driver.api.querybuilder.QueryBuilder;
import com.datastax.oss.driver.api.querybuilder.delete.Delete;
import com.datastax.oss.driver.api.querybuilder.insert.Insert;
import com.datastax.oss.driver.api.querybuilder.insert.RegularInsert;
import com.datastax.oss.driver.api.querybuilder.select.Select;
import com.datastax.oss.driver.api.querybuilder.truncate.Truncate;
import com.datastax.oss.driver.api.querybuilder.update.Update;
import org.reactivestreams.Publisher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.SynchronousSink;
import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
@@ -82,6 +66,24 @@ import org.springframework.data.projection.SpelAwareProxyProjectionFactory;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
import com.datastax.oss.driver.api.core.CqlIdentifier;
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;
import com.datastax.oss.driver.api.core.cql.SimpleStatement;
import com.datastax.oss.driver.api.core.cql.Statement;
import com.datastax.oss.driver.api.querybuilder.QueryBuilder;
import com.datastax.oss.driver.api.querybuilder.delete.Delete;
import com.datastax.oss.driver.api.querybuilder.insert.Insert;
import com.datastax.oss.driver.api.querybuilder.insert.RegularInsert;
import com.datastax.oss.driver.api.querybuilder.select.Select;
import com.datastax.oss.driver.api.querybuilder.truncate.Truncate;
import com.datastax.oss.driver.api.querybuilder.update.Update;
/**
* Primary implementation of {@link ReactiveCassandraOperations}. It simplifies the use of Reactive Cassandra usage and
* helps to avoid common errors. It executes core Cassandra workflow. This class executes CQL queries or updates,
@@ -852,6 +854,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));
}
@@ -888,7 +903,7 @@ public class ReactiveCassandraTemplate
if (PreparedStatementDelegate.canPrepare(isUsePreparedStatements(), statement, logger)) {
PreparedStatementHandler statementHandler = new PreparedStatementHandler(statement);
ReactivePreparedStatementHandler statementHandler = createPreparedStatementHandler(statement);
return getReactiveCqlOperations().query(statementHandler, statementHandler, rowMapper);
}
@@ -899,7 +914,7 @@ public class ReactiveCassandraTemplate
if (PreparedStatementDelegate.canPrepare(isUsePreparedStatements(), statement, logger)) {
PreparedStatementHandler statementHandler = new PreparedStatementHandler(statement);
ReactivePreparedStatementHandler statementHandler = createPreparedStatementHandler(statement);
return getReactiveCqlOperations()
.query(statementHandler, statementHandler, rs -> Mono.just(mappingFunction.apply(rs))).next();
}
@@ -912,7 +927,7 @@ public class ReactiveCassandraTemplate
if (PreparedStatementDelegate.canPrepare(isUsePreparedStatements(), statement, logger)) {
PreparedStatementHandler statementHandler = new PreparedStatementHandler(statement);
ReactivePreparedStatementHandler statementHandler = createPreparedStatementHandler(statement);
return getReactiveCqlOperations().query(statementHandler, statementHandler, mappingFunction::apply).next();
}
@@ -948,9 +963,7 @@ public class ReactiveCassandraTemplate
}
}
return getReactiveCqlOperations()
.execute(new GetConfiguredPageSize())
.single();
return getReactiveCqlOperations().execute(new GetConfiguredPageSize()).single();
}
@SuppressWarnings("unchecked")
@@ -1020,14 +1033,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
implements ReactivePreparedStatementCreator, PreparedStatementBinder, CqlProvider {
public static class PreparedStatementHandler implements ReactivePreparedStatementHandler {
private final SimpleStatement statement;