Spring/spring-data-relational
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add fetchSize to ReactiveSelectOperationSupport.

Browse Source 
Closes #1652
Original pull request: #1898
This commit is contained in:
Mikhail2048
2024-09-26 12:08:41 +03:00
committed by Mark Paluch
parent 7cf81aed35
commit 96a4121058
5 changed files with 82 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -94,6 +94,7 @@ import org.springframework.util.Assert;
* @author Jose Luis Leon
* @author Robert Heim
* @author Sebastian Wieland
* @author Mikhail Polivakha
* @since 1.1
*/
public class R2dbcEntityTemplate implements R2dbcEntityOperations, BeanFactoryAware, ApplicationContextAware {
@@ -312,14 +313,14 @@ public class R2dbcEntityTemplate implements R2dbcEntityOperations, BeanFactoryAw
Assert.notNull(entityClass, "Entity class must not be null");
SqlIdentifier tableName = getTableName(entityClass);
return doSelect(query, entityClass, tableName, entityClass, RowsFetchSpec::all);
return doSelect(query, entityClass, tableName, entityClass, RowsFetchSpec::all, null);
}
@SuppressWarnings("unchecked")
<T, P extends Publisher<T>> P doSelect(Query query, Class<?> entityClass, SqlIdentifier tableName,
Class<T> returnType, Function<RowsFetchSpec<T>, P> resultHandler) {
Class<T> returnType, Function<RowsFetchSpec<T>, P> resultHandler, @Nullable Integer fetchSize) {
RowsFetchSpec<T> fetchSpec = doSelect(query, entityClass, tableName, returnType);
RowsFetchSpec<T> fetchSpec = doSelect(query, entityClass, tableName, returnType, fetchSize);
P result = resultHandler.apply(fetchSpec);
@@ -331,7 +332,7 @@ public class R2dbcEntityTemplate implements R2dbcEntityOperations, BeanFactoryAw
}
private <T> RowsFetchSpec<T> doSelect(Query query, Class<?> entityType, SqlIdentifier tableName,
Class<T> returnType) {
Class<T> returnType, @Nullable Integer fetchSize) {
StatementMapper statementMapper = dataAccessStrategy.getStatementMapper().forType(entityType);
@@ -358,13 +359,17 @@ public class R2dbcEntityTemplate implements R2dbcEntityOperations, BeanFactoryAw
PreparedOperation<?> operation = statementMapper.getMappedObject(selectSpec);
return getRowsFetchSpec(databaseClient.sql(operation), entityType, returnType);
return getRowsFetchSpec(
databaseClient.sql(operation).filter((statement) -> statement.fetchSize(Optional.ofNullable(fetchSize).orElse(0))),
entityType,
returnType
);
}
@Override
public <T> Mono<T> selectOne(Query query, Class<T> entityClass) throws DataAccessException {
return doSelect(query.isLimited() ? query : query.limit(2), entityClass, getTableName(entityClass), entityClass,
RowsFetchSpec::one);
RowsFetchSpec::one, null);
}
@Override

10
spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/ReactiveSelectOperation.java
View File

@@ -36,6 +36,7 @@ import org.springframework.data.relational.core.sql.SqlIdentifier;
* <pre>
* <code>
* select(Human.class)
* .withFetchSize(10)
* .from("star_wars")
* .as(Jedi.class)
* .matching(query(where("firstname").is("luke")))
@@ -44,6 +45,7 @@ import org.springframework.data.relational.core.sql.SqlIdentifier;
* </pre>
*
* @author Mark Paluch
* @author Mikhail Polivakha
* @since 1.1
*/
public interface ReactiveSelectOperation {
@@ -115,6 +117,14 @@ public interface ReactiveSelectOperation {
*/
interface SelectWithQuery<T> extends TerminatingSelect<T> {
/**
* Specifies the fetch size for this query
*
* @param fetchSize
* @return
*/
SelectWithQuery<T> withFetchSize(int fetchSize);
/**
* Set the {@link Query} used as a filter in the {@code SELECT} statement.
*

28
spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/ReactiveSelectOperationSupport.java
View File

@@ -28,6 +28,7 @@ import org.springframework.util.Assert;
* Implementation of {@link ReactiveSelectOperation}.
*
* @author Mark Paluch
* @author Mikhail Polivakha
* @since 1.1
*/
class ReactiveSelectOperationSupport implements ReactiveSelectOperation {
@@ -43,7 +44,7 @@ class ReactiveSelectOperationSupport implements ReactiveSelectOperation {
Assert.notNull(domainType, "DomainType must not be null");
return new ReactiveSelectSupport<>(this.template, domainType, domainType, Query.empty(), null);
return new ReactiveSelectSupport<>(this.template, domainType, domainType, Query.empty(), null, null);
}
static class ReactiveSelectSupport<T> implements ReactiveSelect<T> {
@@ -54,14 +55,17 @@ class ReactiveSelectOperationSupport implements ReactiveSelectOperation {
private final Query query;
private final @Nullable SqlIdentifier tableName;
private final @Nullable Integer fetchSize;
ReactiveSelectSupport(R2dbcEntityTemplate template, Class<?> domainType, Class<T> returnType, Query query,
@Nullable SqlIdentifier tableName) {
@Nullable SqlIdentifier tableName, @Nullable Integer fetchSize) {
this.template = template;
this.domainType = domainType;
this.returnType = returnType;
this.query = query;
this.tableName = tableName;
this.fetchSize = fetchSize;
}
@Override
@@ -69,7 +73,7 @@ class ReactiveSelectOperationSupport implements ReactiveSelectOperation {
Assert.notNull(tableName, "Table name must not be null");
return new ReactiveSelectSupport<>(template, domainType, returnType, query, tableName);
return new ReactiveSelectSupport<>(template, domainType, returnType, query, tableName, fetchSize);
}
@Override
@@ -77,7 +81,15 @@ class ReactiveSelectOperationSupport implements ReactiveSelectOperation {
Assert.notNull(returnType, "ReturnType must not be null");
return new ReactiveSelectSupport<>(template, domainType, returnType, query, tableName);
return new ReactiveSelectSupport<>(template, domainType, returnType, query, tableName, fetchSize);
}
@Override
public SelectWithQuery<T> withFetchSize(int fetchSize) {
Assert.notNull(returnType, "FetchSize must not be null");
return new ReactiveSelectSupport<>(template, domainType, returnType, query, tableName, fetchSize);
}
@Override
@@ -85,7 +97,7 @@ class ReactiveSelectOperationSupport implements ReactiveSelectOperation {
Assert.notNull(query, "Query must not be null");
return new ReactiveSelectSupport<>(template, domainType, returnType, query, tableName);
return new ReactiveSelectSupport<>(template, domainType, returnType, query, tableName, fetchSize);
}
@Override
@@ -100,17 +112,17 @@ class ReactiveSelectOperationSupport implements ReactiveSelectOperation {
@Override
public Mono<T> first() {
return template.doSelect(query.limit(1), domainType, getTableName(), returnType, RowsFetchSpec::first);
return template.doSelect(query.limit(1), domainType, getTableName(), returnType, RowsFetchSpec::first, fetchSize);
}
@Override
public Mono<T> one() {
return template.doSelect(query.limit(2), domainType, getTableName(), returnType, RowsFetchSpec::one);
return template.doSelect(query.limit(2), domainType, getTableName(), returnType, RowsFetchSpec::one, fetchSize);
}
@Override
public Flux<T> all() {
return template.doSelect(query, domainType, getTableName(), returnType, RowsFetchSpec::all);
return template.doSelect(query, domainType, getTableName(), returnType, RowsFetchSpec::all, fetchSize);
}
private SqlIdentifier getTableName() {

28
spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/core/ReactiveSelectOperationUnitTests.java
View File

@@ -24,6 +24,7 @@ import io.r2dbc.spi.test.MockColumnMetadata;
import io.r2dbc.spi.test.MockResult;
import io.r2dbc.spi.test.MockRow;
import io.r2dbc.spi.test.MockRowMetadata;
import reactor.core.publisher.Flux;
import reactor.test.StepVerifier;
import org.junit.jupiter.api.BeforeEach;
@@ -38,6 +39,7 @@ import org.springframework.r2dbc.core.DatabaseClient;
* Unit test for {@link ReactiveSelectOperation}.
*
* @author Mark Paluch
* @author Mikhail Polivakha
*/
public class ReactiveSelectOperationUnitTests {
@@ -242,6 +244,32 @@ public class ReactiveSelectOperationUnitTests {
assertThat(statement.getSql()).isEqualTo("SELECT COUNT(*) FROM person WHERE person.THE_NAME = $1");
}
@Test // gh-1652
void shouldBeAbleToProvideFetchSize() {
MockRowMetadata metadata = MockRowMetadata.builder()
.columnMetadata(MockColumnMetadata.builder().name("id").type(R2dbcType.INTEGER).build())
.build();
MockResult result = MockResult.builder()
.row(MockRow.builder().identified("id", Object.class, "Walter").metadata(metadata).build())
.build();
recorder.addStubbing(s -> s.startsWith("SELECT"), result);
entityTemplate.select(Person.class) //
.withFetchSize(10)
.matching(query(where("name").is("Walter")).limit(10).offset(20)) //
.all() //
.as(StepVerifier::create) //
.expectNextCount(1) //
.verifyComplete();
StatementRecorder.RecordedStatement statement = recorder.getCreatedStatement(s -> s.startsWith("SELECT"));
assertThat(statement.getSql())
.isEqualTo("SELECT person.* FROM person WHERE person.THE_NAME = $1 LIMIT 10 OFFSET 20");
assertThat(statement.getFetchSize()).isEqualTo(10);
}
static class Person {
@Id String id;

13
spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/testing/StatementRecorder.java
View File

@@ -48,6 +48,7 @@ import org.springframework.r2dbc.core.Parameter;
* Recorder utility for R2DBC {@link Statement}s. Allows stubbing and introspection.
*
* @author Mark Paluch
* @author Mikhail Polivakha
*/
public class StatementRecorder implements ConnectionFactory {
@@ -273,6 +274,8 @@ public class StatementRecorder implements ConnectionFactory {
private final List<Result> results;
private int fetchSize;
private final Map<Object, Parameter> bindings = new LinkedHashMap<>();
public RecordedStatement(String sql, Result result) {
@@ -292,6 +295,10 @@ public class StatementRecorder implements ConnectionFactory {
return sql;
}
public int getFetchSize() {
return fetchSize;
}
@Override
public Statement add() {
return this;
@@ -321,6 +328,12 @@ public class StatementRecorder implements ConnectionFactory {
return this;
}
@Override
public Statement fetchSize(int rows) {
fetchSize = rows;
return this;
}
@Override
public Flux<Result> execute() {
return Flux.fromIterable(results).doOnSubscribe(subscription -> executedStatements.add(this));