From 5a2591f4df1cc4ee13b786e4769bbe9f44f34593 Mon Sep 17 00:00:00 2001 From: Mikhail2048 Date: Thu, 26 Sep 2024 12:08:41 +0300 Subject: [PATCH] GH-1652 Added fetchSize to ReactiveSelectOperationSupport --- .../data/r2dbc/core/R2dbcEntityTemplate.java | 17 +++++++---- .../r2dbc/core/ReactiveSelectOperation.java | 10 +++++++ .../core/ReactiveSelectOperationSupport.java | 28 +++++++++++++------ .../ReactiveSelectOperationUnitTests.java | 28 +++++++++++++++++++ .../data/r2dbc/testing/StatementRecorder.java | 13 +++++++++ 5 files changed, 82 insertions(+), 14 deletions(-) diff --git a/spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/R2dbcEntityTemplate.java b/spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/R2dbcEntityTemplate.java index d8c44ff779..467c7782f1 100644 --- a/spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/R2dbcEntityTemplate.java +++ b/spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/R2dbcEntityTemplate.java @@ -94,6 +94,7 @@ * @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 Flux select(Query query, Class entityClass) throws DataAccessEx 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") > P doSelect(Query query, Class entityClass, SqlIdentifier tableName, - Class returnType, Function, P> resultHandler) { + Class returnType, Function, P> resultHandler, @Nullable Integer fetchSize) { - RowsFetchSpec fetchSpec = doSelect(query, entityClass, tableName, returnType); + RowsFetchSpec fetchSpec = doSelect(query, entityClass, tableName, returnType, fetchSize); P result = resultHandler.apply(fetchSpec); @@ -331,7 +332,7 @@ > P doSelect(Query query, Class entityClass, SqlIde } private RowsFetchSpec doSelect(Query query, Class entityType, SqlIdentifier tableName, - Class returnType) { + Class returnType, @Nullable Integer fetchSize) { StatementMapper statementMapper = dataAccessStrategy.getStatementMapper().forType(entityType); @@ -358,13 +359,17 @@ private RowsFetchSpec doSelect(Query query, Class entityType, SqlIdent 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 Mono selectOne(Query query, Class entityClass) throws DataAccessException { return doSelect(query.isLimited() ? query : query.limit(2), entityClass, getTableName(entityClass), entityClass, - RowsFetchSpec::one); + RowsFetchSpec::one, null); } @Override diff --git a/spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/ReactiveSelectOperation.java b/spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/ReactiveSelectOperation.java index 66222fe726..ae1b9dbed7 100644 --- a/spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/ReactiveSelectOperation.java +++ b/spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/ReactiveSelectOperation.java @@ -36,6 +36,7 @@ * 
  *     
  *         select(Human.class)
+ *             .withFetchSize(10)
  *             .from("star_wars")
  *             .as(Jedi.class)
  *             .matching(query(where("firstname").is("luke")))
@@ -44,6 +45,7 @@
  *
* * @author Mark Paluch + * @author Mikhail Polivakha * @since 1.1 */ public interface ReactiveSelectOperation { @@ -115,6 +117,14 @@ interface SelectWithProjection extends SelectWithQuery { */ interface SelectWithQuery extends TerminatingSelect { + /** + * Specifies the fetch size for this query + * + * @param fetchSize + * @return + */ + SelectWithQuery withFetchSize(int fetchSize); + /** * Set the {@link Query} used as a filter in the {@code SELECT} statement. * diff --git a/spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/ReactiveSelectOperationSupport.java b/spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/ReactiveSelectOperationSupport.java index ad579a5cfc..d841a1c5ec 100644 --- a/spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/ReactiveSelectOperationSupport.java +++ b/spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/ReactiveSelectOperationSupport.java @@ -28,6 +28,7 @@ * Implementation of {@link ReactiveSelectOperation}. * * @author Mark Paluch + * @author Mikhail Polivakha * @since 1.1 */ class ReactiveSelectOperationSupport implements ReactiveSelectOperation { @@ -43,7 +44,7 @@ public ReactiveSelect select(Class domainType) { 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 implements ReactiveSelect { @@ -54,14 +55,17 @@ static class ReactiveSelectSupport implements ReactiveSelect { private final Query query; private final @Nullable SqlIdentifier tableName; + private final @Nullable Integer fetchSize; + ReactiveSelectSupport(R2dbcEntityTemplate template, Class domainType, Class 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 @@ public SelectWithProjection from(SqlIdentifier tableName) { 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 @@ public SelectWithQuery as(Class returnType) { 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 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 @@ public TerminatingSelect matching(Query query) { 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 @@ public Mono exists() { @Override public Mono 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 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 all() { - return template.doSelect(query, domainType, getTableName(), returnType, RowsFetchSpec::all); + return template.doSelect(query, domainType, getTableName(), returnType, RowsFetchSpec::all, fetchSize); } private SqlIdentifier getTableName() { diff --git a/spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/core/ReactiveSelectOperationUnitTests.java b/spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/core/ReactiveSelectOperationUnitTests.java index 014fb905c0..6524356467 100644 --- a/spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/core/ReactiveSelectOperationUnitTests.java +++ b/spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/core/ReactiveSelectOperationUnitTests.java @@ -24,6 +24,7 @@ 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 @@ * Unit test for {@link ReactiveSelectOperation}. * * @author Mark Paluch + * @author Mikhail Polivakha */ public class ReactiveSelectOperationUnitTests { @@ -242,6 +244,32 @@ void shouldSelectCount() { 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; diff --git a/spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/testing/StatementRecorder.java b/spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/testing/StatementRecorder.java index de993ff15c..395879d183 100644 --- a/spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/testing/StatementRecorder.java +++ b/spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/testing/StatementRecorder.java @@ -48,6 +48,7 @@ * 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 RecordedStatement implements Statement { private final List results; + private int fetchSize; + private final Map bindings = new LinkedHashMap<>(); public RecordedStatement(String sql, Result result) { @@ -292,6 +295,10 @@ public String getSql() { return sql; } + public int getFetchSize() { + return fetchSize; + } + @Override public Statement add() { return this; @@ -321,6 +328,12 @@ public Statement bindNull(String identifier, Class type) { return this; } + @Override + public Statement fetchSize(int rows) { + fetchSize = rows; + return this; + } + @Override public Flux execute() { return Flux.fromIterable(results).doOnSubscribe(subscription -> executedStatements.add(this));