Add QueryWrapper support for R2DBC query customization

  - Introduce QueryWrapper interface for query modification
  - Add DefaultQueryWrapper as no-op implementation
  - Integrate query wrapper into R2dbcEntityTemplate operations
  - Support query customization for count, exists, select, update, delete operations

Signed-off-by: finger <[email protected]>

Author: yuyu1025

spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/DefaultQueryWrapper.java

@@ -0,0 +1,10 @@
+package org.springframework.data.r2dbc.core;
+
+import org.springframework.data.relational.core.query.Query;
+
+public class DefaultQueryWrapper implements QueryWrapper{
+    @Override
+    public Query wrapper(Query query, Class<?> domainType) {
+        return query;
+    }
+}

spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/QueryWrapper.java

@@ -0,0 +1,8 @@
+package org.springframework.data.r2dbc.core;
+
+import org.springframework.data.relational.core.query.Query;
+
+public interface QueryWrapper {
+
+    Query wrapper(Query query, Class<?> domainType);
+}

spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/R2dbcEntityTemplate.java

@@ -114,6 +114,7 @@ public class R2dbcEntityTemplate implements R2dbcEntityOperations, BeanFactoryAw
 
 	private Function<Statement, Statement> statementFilterFunction = Function.identity();
 
+	private final QueryWrapper queryWrapper;
 	/**
 	 * Create a new {@link R2dbcEntityTemplate} given {@link ConnectionFactory}.
 	 *
@@ -132,6 +133,7 @@ public R2dbcEntityTemplate(ConnectionFactory connectionFactory) {
 		this.converter = dataAccessStrategy.getConverter();
 		this.mappingContext = converter.getMappingContext();
 		this.projectionFactory = new SpelAwareProxyProjectionFactory();
+		this.queryWrapper=new DefaultQueryWrapper();
 	}
 
 	/**
@@ -157,14 +159,22 @@ public R2dbcEntityTemplate(DatabaseClient databaseClient, R2dbcDialect dialect)
 	public R2dbcEntityTemplate(DatabaseClient databaseClient, R2dbcDialect dialect, R2dbcConverter converter) {
 		this(databaseClient, new DefaultReactiveDataAccessStrategy(dialect, converter));
 	}
-
 	/**
 	 * Create a new {@link R2dbcEntityTemplate} given {@link DatabaseClient} and {@link ReactiveDataAccessStrategy}.
 	 *
 	 * @param databaseClient must not be {@literal null}.
 	 * @since 1.2
 	 */
 	public R2dbcEntityTemplate(DatabaseClient databaseClient, ReactiveDataAccessStrategy strategy) {
+		this(databaseClient,strategy,new DefaultQueryWrapper());
+	}
+	/**
+	 * Create a new {@link R2dbcEntityTemplate} given {@link DatabaseClient} and {@link ReactiveDataAccessStrategy}.
+	 *
+	 * @param databaseClient must not be {@literal null}.
+	 * @since 1.2
+	 */
+	public R2dbcEntityTemplate(DatabaseClient databaseClient, ReactiveDataAccessStrategy strategy,QueryWrapper queryWrapper) {
 
 		Assert.notNull(databaseClient, "DatabaseClient must not be null");
 		Assert.notNull(strategy, "ReactiveDataAccessStrategy must not be null");
@@ -174,6 +184,7 @@ public R2dbcEntityTemplate(DatabaseClient databaseClient, ReactiveDataAccessStra
 		this.converter = dataAccessStrategy.getConverter();
 		this.mappingContext = strategy.getConverter().getMappingContext();
 		this.projectionFactory = new SpelAwareProxyProjectionFactory();
+		this.queryWrapper = queryWrapper;
 	}
 
 	/**
@@ -274,7 +285,7 @@ public Mono<Long> count(Query query, Class<?> entityClass) throws DataAccessExce
 	}
 
 	Mono<Long> doCount(Query query, Class<?> entityClass, SqlIdentifier tableName) {
-
+		query = queryWrapper.wrapper(query,entityClass);
 		StatementMapper statementMapper = dataAccessStrategy.getStatementMapper().forType(entityClass);
 
 		StatementMapper.SelectSpec selectSpec = statementMapper //
@@ -305,7 +316,7 @@ public Mono<Boolean> exists(Query query, Class<?> entityClass) throws DataAccess
 	}
 
 	Mono<Boolean> doExists(Query query, Class<?> entityClass, SqlIdentifier tableName) {
-
+		query = queryWrapper.wrapper(query,entityClass);
 		StatementMapper statementMapper = dataAccessStrategy.getStatementMapper().forType(entityClass);
 		StatementMapper.SelectSpec selectSpec = statementMapper.createSelect(tableName).limit(1)
 				.withProjection(Expressions.just("1"));
@@ -352,12 +363,12 @@ <T, P extends Publisher<T>> P doSelect(Query query, Class<?> entityClass, SqlIde
 
 	private <T> RowsFetchSpec<T> doSelect(Query query, Class<?> entityType, SqlIdentifier tableName,
 			Class<T> returnType, Function<? super Statement, ? extends Statement> filterFunction) {
-
+		query = queryWrapper.wrapper(query,entityType);
 		StatementMapper statementMapper = dataAccessStrategy.getStatementMapper().forType(entityType);
-
+		final Query finalQuery = query;
 		StatementMapper.SelectSpec selectSpec = statementMapper //
 				.createSelect(tableName) //
-				.doWithTable((table, spec) -> spec.withProjection(getSelectProjection(table, query, entityType, returnType)));
+				.doWithTable((table, spec) -> spec.withProjection(getSelectProjection(table, finalQuery, entityType, returnType)));
 
 		if (query.getLimit() > 0) {
 			selectSpec = selectSpec.limit(query.getLimit());
@@ -402,7 +413,7 @@ public Mono<Long> update(Query query, Update update, Class<?> entityClass) throw
 	}
 
 	Mono<Long> doUpdate(Query query, Update update, Class<?> entityClass, SqlIdentifier tableName) {
-
+		query = queryWrapper.wrapper(query, entityClass);
 		StatementMapper statementMapper = dataAccessStrategy.getStatementMapper().forType(entityClass);
 
 		StatementMapper.UpdateSpec selectSpec = statementMapper //
@@ -427,7 +438,7 @@ public Mono<Long> delete(Query query, Class<?> entityClass) throws DataAccessExc
 	}
 
 	Mono<Long> doDelete(Query query, Class<?> entityClass, SqlIdentifier tableName) {
-
+		query = queryWrapper.wrapper(query, entityClass);
 		StatementMapper statementMapper = dataAccessStrategy.getStatementMapper().forType(entityClass);
 
 		StatementMapper.DeleteSpec deleteSpec = statementMapper //