From 78da030b339cff4c481840591101166693d7c027 Mon Sep 17 00:00:00 2001
From: Mark Paluch <mpaluch@vmware.com>
Date: Wed, 5 Jul 2023 11:11:16 +0200
Subject: [PATCH] Add support for `Limit`.

Closes #1407
---
 .../data/cassandra/core/StatementFactory.java |  1 +
 .../data/cassandra/core/query/Query.java      | 36 ++++++++++++++-----
 .../CassandraParametersParameterAccessor.java | 11 ++++++
 .../query/ConvertingParameterAccessor.java    |  6 ++++
 .../query/QueryStatementCreator.java          |  6 ++++
 .../QueryDerivationIntegrationTests.java      | 11 ++++++
 ...veCassandraRepositoryIntegrationTests.java | 15 ++++++++
 .../reference/cassandra-repositories.adoc     | 21 ++++++-----
 src/main/asciidoc/reference/cassandra.adoc    |  3 +-
 9 files changed, 93 insertions(+), 17 deletions(-)

diff --git a/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/StatementFactory.java b/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/StatementFactory.java
index 77df286e6..ae1ac73fa 100644
--- a/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/StatementFactory.java
+++ b/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/StatementFactory.java
@@ -632,6 +632,7 @@ public class StatementFactory {
 
 		StatementBuilder<Select> select = createSelectAndOrder(selectors, tableName, filter, sort);
 
+		// TODO: Bind marker
 		if (query.getLimit() > 0) {
 			select.apply(it -> it.limit(Math.toIntExact(query.getLimit())));
 		}
diff --git a/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/query/Query.java b/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/query/Query.java
index 7249a341d..3b6a65ceb 100644
--- a/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/query/Query.java
+++ b/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/query/Query.java
@@ -27,6 +27,7 @@ import java.util.stream.Collectors;
 import java.util.stream.StreamSupport;
 
 import org.springframework.data.cassandra.core.cql.QueryOptions;
+import org.springframework.data.domain.Limit;
 import org.springframework.data.domain.PageRequest;
 import org.springframework.data.domain.Pageable;
 import org.springframework.data.domain.Sort;
@@ -46,7 +47,7 @@ import org.springframework.util.Assert;
 public class Query implements Filter {
 
 	private static final Query EMPTY = new Query(Collections.emptyList(), Columns.empty(), Sort.unsorted(),
-			Optional.empty(), Optional.empty(), Optional.empty(), false);
+			Optional.empty(), Optional.empty(), Limit.unlimited(), false);
 
 	private final boolean allowFiltering;
 
@@ -54,7 +55,7 @@ public class Query implements Filter {
 
 	private final List<CriteriaDefinition> criteriaDefinitions;
 
-	private final Optional<Long> limit;
+	private final Limit limit;
 
 	private final Optional<ByteBuffer> pagingState;
 
@@ -63,8 +64,9 @@ public class Query implements Filter {
 	private final Sort sort;
 
 	private Query(List<CriteriaDefinition> criteriaDefinitions, Columns columns, Sort sort,
-			Optional<ByteBuffer> pagingState, Optional<QueryOptions> queryOptions, Optional<Long> limit,
-			boolean allowFiltering) {
+			Optional<ByteBuffer> pagingState, Optional<QueryOptions> queryOptions, Limit limit, boolean allowFiltering) {
+
+		Assert.notNull(limit, "Limit must not be null");
 
 		this.criteriaDefinitions = criteriaDefinitions;
 		this.columns = columns;
@@ -110,7 +112,7 @@ public class Query implements Filter {
 		List<CriteriaDefinition> collect = StreamSupport.stream(criteriaDefinitions.spliterator(), false)
 				.collect(Collectors.toList());
 
-		return new Query(collect, Columns.empty(), Sort.unsorted(), Optional.empty(), Optional.empty(), Optional.empty(),
+		return new Query(collect, Columns.empty(), Sort.unsorted(), Optional.empty(), Optional.empty(), Limit.unlimited(),
 				false);
 	}
 
@@ -194,7 +196,7 @@ public class Query implements Filter {
 	/**
 	 * Create a {@link Query} initialized with a {@link PageRequest} to fetch the first page of results or advance in
 	 * paging along with sorting. Reads (and overrides, if set) {@link Pageable#getPageSize() page size} into
-	 * {@link QueryOptions#getPageSize()} and sets {@code pagingState} and {@link Sort}.
+	 * {@code QueryOptions#getPageSize()} and sets {@code pagingState} and {@link Sort}.
 	 *
 	 * @param pageable must not be {@literal null}.
 	 * @return a new {@link Query} object containing the former settings with {@link PageRequest} applied.
@@ -269,14 +271,32 @@ public class Query implements Filter {
 	 */
 	public Query limit(long limit) {
 		return new Query(this.criteriaDefinitions, this.columns, this.sort, this.pagingState, this.queryOptions,
-				Optional.of(limit), this.allowFiltering);
+				Limit.of(Math.toIntExact(limit)), this.allowFiltering);
+	}
+
+	/**
+	 * Limit the number of returned rows to {@link Limit}.
+	 *
+	 * @param limit
+	 * @return a new {@link Query} object containing the former settings with {@code limit} applied.
+	 */
+	public Query limit(Limit limit) {
+		return new Query(this.criteriaDefinitions, this.columns, this.sort, this.pagingState, this.queryOptions, limit,
+				this.allowFiltering);
 	}
 
 	/**
 	 * @return the maximum number of rows to be returned.
 	 */
 	public long getLimit() {
-		return this.limit.orElse(0L);
+		return this.limit.isLimited() ? this.limit.max() : 0;
+	}
+
+	/**
+	 * @return {@code true} if the query is limited.
+	 */
+	public boolean isLimited() {
+		return this.limit.isLimited();
 	}
 
 	/**
diff --git a/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/repository/query/CassandraParametersParameterAccessor.java b/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/repository/query/CassandraParametersParameterAccessor.java
index f506d0aa1..6e4b12a0e 100644
--- a/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/repository/query/CassandraParametersParameterAccessor.java
+++ b/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/repository/query/CassandraParametersParameterAccessor.java
@@ -18,6 +18,7 @@ package org.springframework.data.cassandra.repository.query;
 import org.springframework.data.cassandra.core.cql.QueryOptions;
 import org.springframework.data.cassandra.core.mapping.CassandraSimpleTypeHolder;
 import org.springframework.data.cassandra.core.mapping.CassandraType;
+import org.springframework.data.domain.Limit;
 import org.springframework.data.repository.query.ParameterAccessor;
 import org.springframework.data.repository.query.ParametersParameterAccessor;
 import org.springframework.lang.Nullable;
@@ -75,6 +76,16 @@ public class CassandraParametersParameterAccessor extends ParametersParameterAcc
 		return super.getValues();
 	}
 
+	@Override
+	public Limit getLimit() {
+
+		if (!getParameters().hasLimitParameter()) {
+			return Limit.unlimited();
+		}
+
+		return super.getLimit();
+	}
+
 	@Nullable
 	@Override
 	public QueryOptions getQueryOptions() {
diff --git a/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/repository/query/ConvertingParameterAccessor.java b/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/repository/query/ConvertingParameterAccessor.java
index b5fee87a6..cbc1cd777 100644
--- a/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/repository/query/ConvertingParameterAccessor.java
+++ b/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/repository/query/ConvertingParameterAccessor.java
@@ -20,6 +20,7 @@ import java.util.Iterator;
 import org.springframework.data.cassandra.core.convert.CassandraConverter;
 import org.springframework.data.cassandra.core.cql.QueryOptions;
 import org.springframework.data.cassandra.core.mapping.CassandraType;
+import org.springframework.data.domain.Limit;
 import org.springframework.data.domain.Pageable;
 import org.springframework.data.domain.Range;
 import org.springframework.data.domain.ScrollPosition;
@@ -64,6 +65,11 @@ class ConvertingParameterAccessor implements CassandraParameterAccessor {
 		return this.delegate.getSort();
 	}
 
+	@Override
+	public Limit getLimit() {
+		return this.delegate.getLimit();
+	}
+
 	@Nullable
 	@Override
 	public Class<?> findDynamicProjection() {
diff --git a/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/repository/query/QueryStatementCreator.java b/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/repository/query/QueryStatementCreator.java
index 123c753a8..a3b2f9369 100644
--- a/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/repository/query/QueryStatementCreator.java
+++ b/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/repository/query/QueryStatementCreator.java
@@ -30,6 +30,7 @@ import org.springframework.data.cassandra.core.mapping.CassandraPersistentProper
 import org.springframework.data.cassandra.core.query.Columns;
 import org.springframework.data.cassandra.core.query.Query;
 import org.springframework.data.cassandra.repository.Query.Idempotency;
+import org.springframework.data.domain.Limit;
 import org.springframework.data.mapping.context.MappingContext;
 import org.springframework.data.mapping.model.SpELExpressionEvaluator;
 import org.springframework.data.repository.query.QueryCreationException;
@@ -199,6 +200,11 @@ class QueryStatementCreator {
 				query = query.limit(tree.getMaxResults());
 			}
 
+			Limit limit = parameterAccessor.getLimit();
+			if (limit.isLimited()) {
+				query = query.limit(limit);
+			}
+
 			if (allowsFiltering()) {
 				query = query.withAllowFiltering();
 			}
diff --git a/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/repository/QueryDerivationIntegrationTests.java b/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/repository/QueryDerivationIntegrationTests.java
index a7c0c9250..437216178 100644
--- a/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/repository/QueryDerivationIntegrationTests.java
+++ b/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/repository/QueryDerivationIntegrationTests.java
@@ -52,6 +52,7 @@ import org.springframework.data.cassandra.repository.config.EnableCassandraRepos
 import org.springframework.data.cassandra.repository.support.AbstractSpringDataEmbeddedCassandraIntegrationTest;
 import org.springframework.data.cassandra.repository.support.IntegrationTestConfig;
 import org.springframework.data.cassandra.support.CassandraVersion;
+import org.springframework.data.domain.Limit;
 import org.springframework.data.domain.Pageable;
 import org.springframework.data.domain.Slice;
 import org.springframework.data.domain.Sort;
@@ -340,6 +341,14 @@ class QueryDerivationIntegrationTests extends AbstractSpringDataEmbeddedCassandr
 		assertThat(result).contains(walter, skyler, flynn);
 	}
 
+	@Test // GH-1407
+	public void shouldSelectWithLimit() {
+
+		List<Person> result = personRepository.findAllLimitedByLastname("White", Limit.of(2));
+
+		assertThat(result).hasSize(2);
+	}
+
 	@Test // DATACASS-512
 	public void shouldCountRecords() {
 
@@ -427,6 +436,8 @@ class QueryDerivationIntegrationTests extends AbstractSpringDataEmbeddedCassandr
 
 		Slice<Person> findAllSlicedByLastname(String lastname, Pageable pageable);
 
+		List<Person> findAllLimitedByLastname(String lastname, Limit limit);
+
 		Collection<PersonProjection> findPersonProjectedBy();
 
 		Collection<PersonDto> findPersonDtoBy();
diff --git a/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/repository/ReactiveCassandraRepositoryIntegrationTests.java b/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/repository/ReactiveCassandraRepositoryIntegrationTests.java
index 24842f0a1..6c7110534 100644
--- a/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/repository/ReactiveCassandraRepositoryIntegrationTests.java
+++ b/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/repository/ReactiveCassandraRepositoryIntegrationTests.java
@@ -42,6 +42,7 @@ import org.springframework.data.cassandra.repository.support.AbstractSpringDataE
 import org.springframework.data.cassandra.repository.support.IntegrationTestConfig;
 import org.springframework.data.cassandra.repository.support.ReactiveCassandraRepositoryFactory;
 import org.springframework.data.cassandra.repository.support.SimpleReactiveCassandraRepository;
+import org.springframework.data.domain.Limit;
 import org.springframework.data.domain.Pageable;
 import org.springframework.data.domain.Slice;
 import org.springframework.data.domain.Sort;
@@ -135,6 +136,18 @@ class ReactiveCassandraRepositoryIntegrationTests extends AbstractSpringDataEmbe
 				.expectNextMatches(users -> users.getSize() == 1 && users.hasNext()).verifyComplete();
 	}
 
+	@Test // GH-1407
+	void shouldFindWithLimitByLastName() {
+		repository.findByLastname(dave.getLastname(), Limit.of(1)).as(StepVerifier::create).expectNextCount(1)
+				.verifyComplete();
+
+		repository.findByLastname(dave.getLastname(), Limit.of(2)).as(StepVerifier::create).expectNextCount(2)
+				.verifyComplete();
+
+		repository.findByLastname(dave.getLastname(), Limit.unlimited()).as(StepVerifier::create).expectNextCount(2)
+				.verifyComplete();
+	}
+
 	@Test // DATACASS-529
 	void shouldFindEmptySliceByLastName() {
 		repository.findByLastname("foo", CassandraPageRequest.first(1)).as(StepVerifier::create)
@@ -235,6 +248,8 @@ class ReactiveCassandraRepositoryIntegrationTests extends AbstractSpringDataEmbe
 
 		Mono<Slice<User>> findByLastname(String firstname, Pageable pageable);
 
+		Flux<User> findByLastname(String lastname, Limit limit);
+
 		Mono<User> findFirstByLastname(String lastname);
 
 		Mono<User> findOneByLastname(String lastname);
diff --git a/src/main/asciidoc/reference/cassandra-repositories.adoc b/src/main/asciidoc/reference/cassandra-repositories.adoc
index 823fc2cb1..01e7d53d7 100644
--- a/src/main/asciidoc/reference/cassandra-repositories.adoc
+++ b/src/main/asciidoc/reference/cassandra-repositories.adoc
@@ -186,16 +186,19 @@ public interface PersonRepository extends CrudRepository<Person, String> {
 
     List<Person> findByFirstname(String firstname, Sort sort);              <4>
 
-    Person findByShippingAddress(Address address);                          <5>
+    List<Person> findByFirstname(String firstname, Limit limit);            <5>
 
-    Person findFirstByShippingAddress(Address address);                     <6>
+    Person findByShippingAddress(Address address);                          <6>
 
-    Stream<Person> findAllBy();                                             <7>
+    Person findFirstByShippingAddress(Address address);                     <7>
+
+    Stream<Person> findAllBy();                                             <8>
 
     @AllowFiltering
-    List<Person> findAllByAge(int age);                                     <8>
+    List<Person> findAllByAge(int age);                                     <9>
 }
 ----
+
 <1> The method shows a query for all people with the given `lastname`.
 The query is derived from parsing the method name for constraints, which can be concatenated with `And`.
 Thus, the method name results in a query expression of `SELECT * FROM person WHERE lastname = 'lastname'`.
@@ -204,12 +207,14 @@ You can equip your method signature with a `Pageable` parameter and let the meth
 <3> Passing a `QueryOptions` object applies the query options to the resulting query before its execution.
 <4> Applies dynamic sorting to a query.
 You can add a `Sort` parameter to your method signature, and Spring Data automatically applies ordering to the query.
-<5> Shows that you can query based on properties that are not a primitive type by using `Converter` instances registered in `CustomConversions`.
+<5> Applies dynamic result limiting to a query.
+Query results can be limited using `SELECT … LIMIT`.
+<6> Shows that you can query based on properties that are not a primitive type by using `Converter` instances registered in `CustomConversions`.
 Throws `IncorrectResultSizeDataAccessException` if more than one match is found.
-<6> Uses the `First` keyword to restrict the query to only the first result.
+<7> Uses the `First` keyword to restrict the query to only the first result.
 Unlike the preceding method, this method does not throw an exception if more than one match is found.
-<7> Uses a Java 8 `Stream` to read and convert individual elements while iterating the stream.
-<8> Shows a query method annotated with `@AllowFiltering`, to allow server-side filtering.
+<8> Uses a Java 8 `Stream` to read and convert individual elements while iterating the stream.
+<9> Shows a query method annotated with `@AllowFiltering`, to allow server-side filtering.
 ====
 
 NOTE: Querying non-primary key properties requires secondary indexes.
diff --git a/src/main/asciidoc/reference/cassandra.adoc b/src/main/asciidoc/reference/cassandra.adoc
index 26eb1da7c..e08f08444 100644
--- a/src/main/asciidoc/reference/cassandra.adoc
+++ b/src/main/asciidoc/reference/cassandra.adoc
@@ -1063,7 +1063,8 @@ The `Query` class has some additional methods that you can use to provide option
 * `Query` *by* `(CriteriaDefinition... criteria)`: Used to create a `Query` object.
 * `Query` *and* `(CriteriaDefinition criteria)`: Used to add additional criteria to the query.
 * `Query` *columns* `(Columns columns)`: Used to define columns to be included in the query results.
-* `Query` *limit* `(long limit)`: Used to limit the size of the returned results to the provided limit (used for paging).
+* `Query` *limit* `(Limit limit)`: Used to limit the size of the returned results to the provided limit (used `SELECT` limiting).
+* `Query` *limit* `(long limit)`: Used to limit the size of the returned results to the provided limit (used `SELECT` limiting).
 * `Query` *pageRequest* `(Pageable pageRequest)`: Used to associate `Sort`, `PagingState`, and `fetchSize` with the query (used for paging).
 * `Query` *pagingState* `(ByteBuffer pagingState)`: Used to associate a `ByteBuffer` with the query (used for paging).
 * `Query` *queryOptions* `(QueryOptions queryOptions)`: Used to associate `QueryOptions` with the query.