Exclude id if provided sort properties is keyset qualified

See GH-2996

Author: quaff

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/KeysetScrollSpecification.java

@@ -40,6 +40,7 @@
  *
  * @author Mark Paluch
  * @author Christoph Strobl
+ * @author Yanming Zhou
  * @since 3.1
  */
 public record KeysetScrollSpecification<T> (KeysetScrollPosition position, Sort sort,
@@ -63,21 +64,26 @@ 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());
-		}
-
-		sort.forEach(it -> sortById.remove(it.getProperty()));
 
-		if (sortById.isEmpty()) {
+		if (entity.isKeysetQualified(sort.stream().map(Order::getProperty).toList())) {
 			sortToUse = sort;
 		} else {
-			sortToUse = sort.and(Sort.by(sortById.toArray(new String[0])));
+			Collection<String> sortById;
+			if (entity.hasCompositeId()) {
+				sortById = new ArrayList<>(entity.getIdAttributeNames());
+			} else {
+				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);

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/support/JpaEntityInformation.java

@@ -30,6 +30,7 @@
  * @author Oliver Gierke
  * @author Thomas Darimont
  * @author Mark Paluch
+ * @author Yanming Zhou
  */
 public interface JpaEntityInformation<T, ID> extends EntityInformation<T, ID>, JpaEntityMetadata<T> {
 
@@ -79,12 +80,24 @@ public interface JpaEntityInformation<T, ID> extends EntityInformation<T, ID>, J
 	Object getCompositeIdAttributeValue(Object id, String idAttribute);
 
 	/**
-	 * Extract a keyset for {@code propertyPaths} and the primary key (including composite key components if applicable).
+	 * Extract a keyset for {@code propertyPaths}, and the primary key (including composite key components if applicable)
+	 * if {@code propertyPaths} is not qualified.
 	 *
 	 * @param propertyPaths the property paths that make up the keyset in combination with the composite key components.
 	 * @param entity the entity to extract values from
 	 * @return a map mapping String representations of the paths to values from the entity.
 	 * @since 3.1
 	 */
 	Map<String, Object> getKeyset(Iterable<String> propertyPaths, T entity);
+
+	/**
+	 * Determine whether propertyPaths is qualified for keyset.
+	 *
+	 * @param propertyPaths the property paths that make up the keyset in combination with the composite key components.
+	 * @return {@code propertyPaths} is qualified for keyset.
+	 * @since 3.2
+	 */
+	default boolean isKeysetQualified(Iterable<String> propertyPaths) {
+		return false;
+	}
 }

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/support/JpaMetamodelEntityInformation.java

@@ -15,6 +15,7 @@
  */
 package org.springframework.data.jpa.repository.support;
 
+import jakarta.persistence.Column;
 import jakarta.persistence.IdClass;
 import jakarta.persistence.PersistenceUnitUtil;
 import jakarta.persistence.Tuple;
@@ -44,6 +45,7 @@
 import org.springframework.data.util.DirectFieldAccessFallbackBeanWrapper;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
+import org.springframework.util.StringUtils;
 
 /**
  * Implementation of {@link org.springframework.data.repository.core.EntityInformation} that uses JPA {@link Metamodel}
@@ -55,6 +57,7 @@
  * @author Mark Paluch
  * @author Jens Schauder
  * @author Greg Turnquist
+ * @author Yanming Zhou
  */
 public class JpaMetamodelEntityInformation<T, ID> extends JpaEntityInformationSupport<T, ID> {
 
@@ -236,12 +239,14 @@ public Map<String, Object> getKeyset(Iterable<String> propertyPaths, T entity) {
 
 		Map<String, Object> keyset = new LinkedHashMap<>();
 
-		if (hasCompositeId()) {
-			for (String idAttributeName : getIdAttributeNames()) {
-				keyset.put(idAttributeName, getter.apply(idAttributeName));
+		if(!isKeysetQualified(propertyPaths)) {
+			if (hasCompositeId()) {
+				for (String idAttributeName : getIdAttributeNames()) {
+					keyset.put(idAttributeName, getter.apply(idAttributeName));
+				}
+			} else {
+				keyset.put(getIdAttribute().getName(), getId(entity));
 			}
-		} else {
-			keyset.put(getIdAttribute().getName(), getId(entity));
 		}
 
 		for (String propertyPath : propertyPaths) {
@@ -251,6 +256,52 @@ public Map<String, Object> getKeyset(Iterable<String> propertyPaths, T entity) {
 		return keyset;
 	}
 
+	@Override
+	public boolean isKeysetQualified(Iterable<String> propertyPaths) {
+
+		if (propertyPaths.iterator().hasNext()) {
+			for (String property : propertyPaths) {
+				if (isUnique(property)) {
+					return true;
+				}
+			}
+		}
+
+		return false;
+	}
+
+	@Nullable
+	private boolean isUnique(String property) {
+
+		Class<?> clazz = getJavaType();
+
+		while (clazz != Object.class) {
+
+			try {
+				Column column =  clazz.getDeclaredField(property).getAnnotation(Column.class);
+				if (column != null) {
+					return column.unique();
+				}
+			} catch (NoSuchFieldException ex) {
+				// ignore
+			}
+
+			try {
+				String getter = "get" + StringUtils.capitalize(property);
+				Column column = clazz.getDeclaredMethod(getter).getAnnotation(Column.class);
+				if (column != null) {
+					return column.unique();
+				}
+			} catch (NoSuchMethodException ex) {
+				// ignore
+			}
+
+			clazz = clazz.getSuperclass();
+		}
+
+		return false;
+	}
+
 	private Function<String, Object> getPropertyValueFunction(Object entity) {
 
 		if (entity instanceof Tuple t) {

spring-data-jpa/src/test/java/org/springframework/data/jpa/domain/sample/Product.java

@@ -1,5 +1,6 @@
 package org.springframework.data.jpa.domain.sample;
 
+import jakarta.persistence.Column;
 import jakarta.persistence.Entity;
 import jakarta.persistence.GeneratedValue;
 import jakarta.persistence.Id;
@@ -9,6 +10,9 @@ public class Product {
 
 	@Id @GeneratedValue private Long id;
 
+	@Column(unique = true)
+	private String code;
+
 	public Long getId() {
 		return id;
 	}

spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/KeysetScrollSpecificationUnitTests.java

@@ -25,17 +25,22 @@
 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.Product;
 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;
 
+import java.util.List;
+import java.util.Map;
+
 /**
  * Unit tests for {@link KeysetScrollSpecification}.
  *
  * @author Mark Paluch
+ * @author Yanming Zhou
  */
 @ExtendWith(SpringExtension.class)
 @ContextConfiguration({ "classpath:infrastructure.xml" })
@@ -74,4 +79,17 @@ void shouldSkipExistingIdentifiersInSort() {
 		assertThat(sort).extracting(Order::getProperty).containsExactly("id", "firstname");
 	}
 
+	@Test // GH-3013
+	void shouldSkipIdentifiersInSortIfUniquePropertyPresent() {
+
+		JpaMetamodelEntityInformation<Product, Long> info = new JpaMetamodelEntityInformation<>(Product.class, em.getMetamodel(),
+				em.getEntityManagerFactory().getPersistenceUnitUtil());
+		Map<String, Object> keyset = info.getKeyset(List.of("code"), new Product());
+
+		assertThat(keyset).containsOnlyKeys("code");
+
+		Sort sort = KeysetScrollSpecification.createSort(ScrollPosition.keyset(), Sort.by("code"), info);
+
+		assertThat(sort).extracting(Order::getProperty).containsExactly("code");
+	}
 }