Supports scrolling base on keyset without id

id shouldn't be added to sort if sort property already provided. See spring-projectsGH-2996
quaff · Jun 16, 2023 · 1324b29 · 1324b29
1 parent e8d08dd
commit 1324b29
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 19 deletions.
diff --git a/...rc/main/java/org/springframework/data/jpa/repository/query/KeysetScrollSpecification.java b/...rc/main/java/org/springframework/data/jpa/repository/query/KeysetScrollSpecification.java
@@ -22,8 +22,6 @@
 import jakarta.persistence.criteria.Predicate;
 import jakarta.persistence.criteria.Root;
 
-import java.util.ArrayList;
-import java.util.Collection;
 import java.util.List;
 
 import org.springframework.data.domain.KeysetScrollPosition;
@@ -40,6 +38,7 @@
  *
  * @author Mark Paluch
  * @author Christoph Strobl
+ * @author Yanming Zhou
  * @since 3.1
  */
 public record KeysetScrollSpecification<T> (KeysetScrollPosition position, Sort sort,
@@ -63,21 +62,16 @@ public static Sort createSort(KeysetScrollPosition position, Sort sort, JpaEntit
 
 		KeysetScrollDelegate delegate = KeysetScrollDelegate.of(position.getDirection());
 
-		Collection<String> sortById;
-		Sort sortToUse;
-		if (entity.hasCompositeId()) {
-			sortById = new ArrayList<>(entity.getIdAttributeNames());
-		} else {
-			sortById = new ArrayList<>(1);
-			sortById.add(entity.getRequiredIdAttribute().getName());
+		if (sort.isSorted()) {
+			// assume sort applied on unique property
+			return delegate.getSortOrders(sort);
 		}
 
-		sort.forEach(it -> sortById.remove(it.getProperty()));
-
-		if (sortById.isEmpty()) {
-			sortToUse = sort;
+		Sort sortToUse;
+		if (entity.hasCompositeId()) {
+			sortToUse = sort.and(Sort.by(entity.getIdAttributeNames().toArray(new String[0])));
 		} else {
-			sortToUse = sort.and(Sort.by(sortById.toArray(new String[0])));
+			sortToUse = sort.and(Sort.by(entity.getRequiredIdAttribute().getName()));
 		}
 
 		return delegate.getSortOrders(sortToUse);

diff --git a/...ava/org/springframework/data/jpa/repository/query/KeysetScrollSpecificationUnitTests.java b/...ava/org/springframework/data/jpa/repository/query/KeysetScrollSpecificationUnitTests.java
@@ -36,6 +36,7 @@
  * Unit tests for {@link KeysetScrollSpecification}.
  *
  * @author Mark Paluch
+ * @author Yanming Zhou
  */
 @ExtendWith(SpringExtension.class)
 @ContextConfiguration({ "classpath:infrastructure.xml" })
@@ -44,24 +45,44 @@ class KeysetScrollSpecificationUnitTests {
 
 	@PersistenceContext EntityManager em;
 
+	@Test
+	void shouldUseIdentifierAsFallback() {
+
+		Sort sort = KeysetScrollSpecification.createSort(ScrollPosition.keyset(), Sort.unsorted(),
+				new JpaMetamodelEntityInformation<>(User.class, em.getMetamodel(),
+						em.getEntityManagerFactory().getPersistenceUnitUtil()));
+
+		assertThat(sort).isEqualTo(Sort.by("id"));
+	}
+
+	@Test
+	void shouldUseCompositeIdentifierAsFallback() {
+
+		Sort sort = KeysetScrollSpecification.createSort(ScrollPosition.keyset(), Sort.unsorted(),
+				new JpaMetamodelEntityInformation<>(SampleWithIdClass.class, em.getMetamodel(),
+						em.getEntityManagerFactory().getPersistenceUnitUtil()));
+
+		assertThat(sort).isEqualTo(Sort.by("first", "second"));
+	}
+
 	@Test // GH-2996
-	void shouldAddIdentifierToSort() {
+	void shouldNotAddIdentifierToSort() {
 
 		Sort sort = KeysetScrollSpecification.createSort(ScrollPosition.keyset(), Sort.by("firstname"),
 				new JpaMetamodelEntityInformation<>(User.class, em.getMetamodel(),
 						em.getEntityManagerFactory().getPersistenceUnitUtil()));
 
-		assertThat(sort).extracting(Order::getProperty).containsExactly("firstname", "id");
+		assertThat(sort).extracting(Order::getProperty).containsExactly("firstname");
 	}
 
 	@Test // GH-2996
-	void shouldAddCompositeIdentifierToSort() {
+	void shouldNotAddCompositeIdentifierToSort() {
 
 		Sort sort = KeysetScrollSpecification.createSort(ScrollPosition.keyset(), Sort.by("first", "firstname"),
 				new JpaMetamodelEntityInformation<>(SampleWithIdClass.class, em.getMetamodel(),
 						em.getEntityManagerFactory().getPersistenceUnitUtil()));
 
-		assertThat(sort).extracting(Order::getProperty).containsExactly("first", "firstname", "second");
+		assertThat(sort).extracting(Order::getProperty).containsExactly("first", "firstname");
 	}
 
 	@Test // GH-2996
@@ -74,4 +95,4 @@ void shouldSkipExistingIdentifiersInSort() {
 		assertThat(sort).extracting(Order::getProperty).containsExactly("id", "firstname");
 	}
 
-}
+}