From af26bb6b316c399f7e9b2ec8bf936f8841c8717f Mon Sep 17 00:00:00 2001
From: Mark Paluch <mpaluch@vmware.com>
Date: Wed, 5 Jul 2023 10:24:51 +0200
Subject: [PATCH] Polishing.

Introduce limit(Limit) method to limit query results applying the Limit domain type.

See #4397
Original pull request: #4398
---
 .../data/mongodb/core/query/Query.java        | 36 +++++++++++++++----
 .../data/mongodb/core/query/QueryTests.java   | 13 +++++++
 2 files changed, 42 insertions(+), 7 deletions(-)

diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/query/Query.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/query/Query.java
index e631852a42..ae4bfee10a 100644
--- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/query/Query.java
+++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/query/Query.java
@@ -31,6 +31,7 @@
 
 import org.bson.Document;
 import org.springframework.data.domain.KeysetScrollPosition;
+import org.springframework.data.domain.Limit;
 import org.springframework.data.domain.OffsetScrollPosition;
 import org.springframework.data.domain.Pageable;
 import org.springframework.data.domain.ScrollPosition;
@@ -66,7 +67,7 @@ public class Query implements ReadConcernAware, ReadPreferenceAware {
 	private @Nullable Field fieldSpec = null;
 	private Sort sort = Sort.unsorted();
 	private long skip;
-	private int limit;
+	private Limit limit = Limit.unlimited();
 
 	private KeysetScrollPosition keysetScrollPosition;
 	private @Nullable ReadConcern readConcern;
@@ -155,10 +156,30 @@ public Query skip(long skip) {
 	 * @return this.
 	 */
 	public Query limit(int limit) {
-		this.limit = limit;
+		this.limit = limit > 0 ? Limit.of(limit) : Limit.unlimited();
 		return this;
 	}
 
+	/**
+	 * Limit the number of returned documents to {@link Limit}.
+	 *
+	 * @param limit number of documents to return.
+	 * @return this.
+	 * @since 4.2
+	 */
+	public Query limit(Limit limit) {
+
+		Assert.notNull(limit, "Limit must not be null");
+
+		if (limit.isUnlimited()) {
+			this.limit = limit;
+			return this;
+		}
+
+		// retain zero/negative semantics for unlimited.
+		return limit(limit.max());
+	}
+
 	/**
 	 * Configures the query to use the given hint when being executed. The {@code hint} can either be an index name or a
 	 * json {@link Document} representation.
@@ -254,7 +275,7 @@ public Query with(Pageable pageable) {
 			return this;
 		}
 
-		this.limit = pageable.getPageSize();
+		this.limit = pageable.toLimit();
 		this.skip = pageable.getOffset();
 
 		return with(pageable.getSort());
@@ -457,7 +478,7 @@ public long getSkip() {
 	 * @since 4.1
 	 */
 	public boolean isLimited() {
-		return this.limit > 0;
+		return this.limit.isLimited();
 	}
 
 	/**
@@ -468,7 +489,7 @@ public boolean isLimited() {
 	 * @see #isLimited()
 	 */
 	public int getLimit() {
-		return this.limit;
+		return limit.isUnlimited() ? 0 : this.limit.max();
 	}
 
 	/**
@@ -683,7 +704,8 @@ public boolean isSorted() {
 		};
 
 		target.skip = source.getSkip();
-		target.limit = source.getLimit();
+
+		target.limit = source.isLimited() ? Limit.of(source.getLimit()) : Limit.unlimited();
 		target.hint = source.getHint();
 		target.collation = source.getCollation();
 		target.restrictedTypes = new HashSet<>(source.getRestrictedTypes());
@@ -746,7 +768,7 @@ public int hashCode() {
 		result += 31 * nullSafeHashCode(sort);
 		result += 31 * nullSafeHashCode(hint);
 		result += 31 * skip;
-		result += 31 * limit;
+		result += 31 * limit.hashCode();
 		result += 31 * nullSafeHashCode(meta);
 		result += 31 * nullSafeHashCode(collation.orElse(null));
 
diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/query/QueryTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/query/QueryTests.java
index c4f7deed12..452fb6178a 100644
--- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/query/QueryTests.java
+++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/query/QueryTests.java
@@ -22,6 +22,7 @@
 import org.bson.Document;
 import org.junit.jupiter.api.Test;
 import org.springframework.aop.framework.ProxyFactory;
+import org.springframework.data.domain.Limit;
 import org.springframework.data.domain.Sort;
 import org.springframework.data.domain.Sort.Direction;
 import org.springframework.data.domain.Sort.Order;
@@ -97,6 +98,18 @@ void testQueryWithLimit() {
 		assertThat(q.getQueryObject()).isEqualTo(Document
 				.parse("{ \"name\" : { \"$gte\" : \"M\" , \"$lte\" : \"T\"} , \"age\" : { \"$not\" : { \"$gt\" : 22}}}"));
 		assertThat(q.getLimit()).isEqualTo(50);
+
+		q.limit(Limit.unlimited());
+		assertThat(q.getLimit()).isZero();
+		assertThat(q.isLimited()).isFalse();
+
+		q.limit(Limit.of(10));
+		assertThat(q.getLimit()).isEqualTo(10);
+		assertThat(q.isLimited()).isTrue();
+
+		q.limit(Limit.of(-1));
+		assertThat(q.getLimit()).isZero();
+		assertThat(q.isLimited()).isFalse();
 	}
 
 	@Test