From 2e489cb4fca1b2414e489936be4d3f2a55bdde79 Mon Sep 17 00:00:00 2001
From: Mark Paluch <mpaluch@vmware.com>
Date: Thu, 8 Jun 2023 09:47:40 +0200
Subject: [PATCH] Avoid duplicate sort orders through Keyset scrolling.

We now avoid adding sort properties if they are already specified by Sort.

Closes #2996
---
 .../query/KeysetScrollSpecification.java      | 16 +++-
 .../KeysetScrollSpecificationUnitTests.java   | 77 +++++++++++++++++++
 2 files changed, 91 insertions(+), 2 deletions(-)
 create mode 100644 spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/KeysetScrollSpecificationUnitTests.java

diff --git a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/KeysetScrollSpecification.java b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/KeysetScrollSpecification.java
index ee4979a637..42846ccc62 100644
--- a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/KeysetScrollSpecification.java
+++ b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/KeysetScrollSpecification.java
@@ -22,6 +22,8 @@
 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;
@@ -61,11 +63,21 @@ public static Sort createSort(KeysetScrollPosition position, Sort sort, JpaEntit
 
 		KeysetScrollDelegate delegate = KeysetScrollDelegate.of(position.getDirection());
 
+		Collection<String> sortById;
 		Sort sortToUse;
 		if (entity.hasCompositeId()) {
-			sortToUse = sort.and(Sort.by(entity.getIdAttributeNames().toArray(new String[0])));
+			sortById = new ArrayList<>(entity.getIdAttributeNames());
 		} else {
-			sortToUse = sort.and(Sort.by(entity.getRequiredIdAttribute().getName()));
+			sortById = new ArrayList<>(1);
+			sortById.add(entity.getRequiredIdAttribute().getName());
+		}
+
+		sort.forEach(it -> sortById.remove(it.getProperty()));
+
+		if (sortById.isEmpty()) {
+			sortToUse = sort;
+		} else {
+			sortToUse = sort.and(Sort.by(sortById.toArray(new String[0])));
 		}
 
 		return delegate.getSortOrders(sortToUse);
diff --git a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/KeysetScrollSpecificationUnitTests.java b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/KeysetScrollSpecificationUnitTests.java
new file mode 100644
index 0000000000..09362c5c5d
--- /dev/null
+++ b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/KeysetScrollSpecificationUnitTests.java
@@ -0,0 +1,77 @@
+/*
+ * Copyright 2023 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.jpa.repository.query;
+
+import static org.assertj.core.api.Assertions.*;
+
+import jakarta.persistence.EntityManager;
+import jakarta.persistence.PersistenceContext;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.springframework.data.domain.ScrollPosition;
+import org.springframework.data.domain.Sort;
+import org.springframework.data.domain.Sort.Order;
+import org.springframework.data.jpa.domain.sample.SampleWithIdClass;
+import org.springframework.data.jpa.domain.sample.User;
+import org.springframework.data.jpa.repository.support.JpaMetamodelEntityInformation;
+import org.springframework.test.context.ContextConfiguration;
+import org.springframework.test.context.junit.jupiter.SpringExtension;
+import org.springframework.transaction.annotation.Transactional;
+
+/**
+ * Unit tests for {@link KeysetScrollSpecification}.
+ *
+ * @author Mark Paluch
+ */
+@ExtendWith(SpringExtension.class)
+@ContextConfiguration({ "classpath:infrastructure.xml" })
+@Transactional
+class KeysetScrollSpecificationUnitTests {
+
+	@PersistenceContext EntityManager em;
+
+	@Test // GH-2996
+	void shouldAddIdentifierToSort() {
+
+		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");
+	}
+
+	@Test // GH-2996
+	void shouldAddCompositeIdentifierToSort() {
+
+		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");
+	}
+
+	@Test // GH-2996
+	void shouldSkipExistingIdentifiersInSort() {
+
+		Sort sort = KeysetScrollSpecification.createSort(ScrollPosition.keyset(), Sort.by("id", "firstname"),
+				new JpaMetamodelEntityInformation<>(User.class, em.getMetamodel(),
+						em.getEntityManagerFactory().getPersistenceUnitUtil()));
+
+		assertThat(sort).extracting(Order::getProperty).containsExactly("id", "firstname");
+	}
+
+}