Support for Keyset scrolling.

Closes #2878
Original pull request #2885
See spring-projects/spring-data-commons#2151

Author: mp911de

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

@@ -0,0 +1,61 @@
+/*
+ * 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 java.util.List;
+
+/**
+ * Utility methods to obtain sublists.
+ *
+ * @author Mark Paluch
+ */
+class CollectionUtils {
+
+	/**
+	 * Return the first {@code count} items from the list.
+	 *
+	 * @param count the number of first elements to be included in the returned list.
+	 * @param list must not be {@literal null}
+	 * @return the returned sublist if the {@code list} is greater {@code count}.
+	 * @param <T>
+	 */
+	public static <T> List<T> getFirst(int count, List<T> list) {
+
+		if (count > 0 && list.size() > count) {
+			return list.subList(0, count);
+		}
+
+		return list;
+	}
+
+	/**
+	 * Return the last {@code count} items from the list.
+	 *
+	 * @param count the number of last elements to be included in the returned list.
+	 * @param list must not be {@literal null}
+	 * @return the returned sublist if the {@code list} is greater {@code count}.
+	 * @param <T>
+	 */
+	public static <T> List<T> getLast(int count, List<T> list) {
+
+		if (count > 0 && list.size() > count) {
+			return list.subList(list.size() - (count), list.size());
+		}
+
+		return list;
+	}
+
+}

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

@@ -0,0 +1,83 @@
+/*
+ * 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 jakarta.persistence.criteria.CriteriaBuilder;
+import jakarta.persistence.criteria.CriteriaQuery;
+import jakarta.persistence.criteria.Predicate;
+import jakarta.persistence.criteria.Root;
+
+import java.util.Collection;
+import java.util.LinkedHashSet;
+import java.util.Set;
+
+import org.springframework.data.domain.KeysetScrollPosition;
+import org.springframework.data.domain.Sort;
+import org.springframework.data.jpa.repository.support.JpaEntityInformation;
+import org.springframework.data.repository.query.ReturnedType;
+import org.springframework.data.repository.query.parser.PartTree;
+import org.springframework.lang.Nullable;
+
+/**
+ * Extension to {@link JpaQueryCreator} to create queries considering {@link KeysetScrollPosition keyset scrolling}.
+ *
+ * @author Mark Paluch
+ * @since 3.1
+ */
+class JpaKeysetScrollQueryCreator extends JpaQueryCreator {
+
+	private final JpaEntityInformation<?, ?> entityInformation;
+	private final KeysetScrollPosition scrollPosition;
+
+	public JpaKeysetScrollQueryCreator(PartTree tree, ReturnedType type, CriteriaBuilder builder,
+			ParameterMetadataProvider provider, JpaEntityInformation<?, ?> entityInformation,
+			KeysetScrollPosition scrollPosition) {
+		super(tree, type, builder, provider);
+		this.entityInformation = entityInformation;
+		this.scrollPosition = scrollPosition;
+	}
+
+	@Override
+	protected CriteriaQuery<?> complete(@Nullable Predicate predicate, Sort sort, CriteriaQuery<?> query,
+			CriteriaBuilder builder, Root<?> root) {
+
+		KeysetScrollSpecification<Object> keysetSpec = new KeysetScrollSpecification<>(scrollPosition, sort,
+				entityInformation);
+		Predicate keysetPredicate = keysetSpec.createPredicate(root, builder);
+
+		CriteriaQuery<?> queryToUse = super.complete(predicate, keysetSpec.sort(), query, builder, root);
+
+		if (keysetPredicate != null) {
+			if (queryToUse.getRestriction() != null) {
+				return queryToUse.where(builder.and(queryToUse.getRestriction(), keysetPredicate));
+			}
+			return queryToUse.where(keysetPredicate);
+		}
+
+		return queryToUse;
+	}
+
+	@Override
+	Collection<String> getRequiredSelection(Sort sort, ReturnedType returnedType) {
+
+		Sort sortToUse = KeysetScrollSpecification.createSort(scrollPosition, sort, entityInformation);
+
+		Set<String> selection = new LinkedHashSet<>(returnedType.getInputProperties());
+		sortToUse.forEach(it -> selection.add(it.getProperty()));
+
+		return selection;
+	}
+}

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

@@ -118,7 +118,6 @@ public List<ParameterMetadata<?>> getParameterExpressions() {
 
 	@Override
 	protected Predicate create(Part part, Iterator<Object> iterator) {
-
 		return toPredicate(part, root);
 	}
 
@@ -158,9 +157,10 @@ protected CriteriaQuery<? extends Object> complete(@Nullable Predicate predicate
 
 		if (returnedType.needsCustomConstruction()) {
 
+			Collection<String> requiredSelection = getRequiredSelection(sort, returnedType);
 			List<Selection<?>> selections = new ArrayList<>();
 
-			for (String property : returnedType.getInputProperties()) {
+			for (String property : requiredSelection) {
 
 				PropertyPath path = PropertyPath.from(property, returnedType.getDomainType());
 				selections.add(toExpressionRecursively(root, path, true).alias(property));
@@ -195,6 +195,10 @@ protected CriteriaQuery<? extends Object> complete(@Nullable Predicate predicate
 		return predicate == null ? select : select.where(predicate);
 	}
 
+	Collection<String> getRequiredSelection(Sort sort, ReturnedType returnedType) {
+		return returnedType.getInputProperties();
+	}
+
 	/**
 	 * Creates a {@link Predicate} from the given {@link Part}.
 	 *

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

@@ -15,23 +15,25 @@
  */
 package org.springframework.data.jpa.repository.query;
 
-import java.lang.reflect.Method;
-import java.util.Collection;
-import java.util.List;
-import java.util.Optional;
-
 import jakarta.persistence.EntityManager;
 import jakarta.persistence.NoResultException;
 import jakarta.persistence.Query;
 import jakarta.persistence.StoredProcedureQuery;
 
+import java.lang.reflect.Method;
+import java.util.Collection;
+import java.util.List;
+import java.util.Optional;
+
 import org.springframework.core.convert.ConversionService;
 import org.springframework.core.convert.support.ConfigurableConversionService;
 import org.springframework.core.convert.support.DefaultConversionService;
 import org.springframework.dao.InvalidDataAccessApiUsageException;
 import org.springframework.data.domain.Pageable;
+import org.springframework.data.domain.ScrollPosition;
 import org.springframework.data.domain.Slice;
 import org.springframework.data.domain.SliceImpl;
+import org.springframework.data.domain.Sort;
 import org.springframework.data.jpa.provider.PersistenceProvider;
 import org.springframework.data.repository.core.support.SurroundingTransactionDetectorMethodInterceptor;
 import org.springframework.data.support.PageableExecutionUtils;
@@ -128,6 +130,33 @@ protected Object doExecute(AbstractJpaQuery query, JpaParametersParameterAccesso
 		}
 	}
 
+	/**
+	 * Executes the query to return a {@link org.springframework.data.domain.Window} of entities.
+	 *
+	 * @author Mark Paluch
+	 * @since 3.1
+	 */
+	static class ScrollExecution extends JpaQueryExecution {
+
+		private final Sort sort;
+		private final ScrollDelegate<?> delegate;
+
+		ScrollExecution(Sort sort, ScrollDelegate<?> delegate) {
+			this.sort = sort;
+			this.delegate = delegate;
+		}
+
+		@Override
+		@SuppressWarnings("unchecked")
+		protected Object doExecute(AbstractJpaQuery query, JpaParametersParameterAccessor accessor) {
+
+			ScrollPosition scrollPosition = accessor.getScrollPosition();
+			Query scrollQuery = query.createQuery(accessor);
+
+			return delegate.scroll(scrollQuery, sort.and(accessor.getSort()), scrollPosition);
+		}
+	}
+
 	/**
 	 * Executes the query to return a {@link Slice} of entities.
 	 *

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

@@ -18,6 +18,7 @@
 import jakarta.persistence.EntityManager;
 
 import org.springframework.data.jpa.repository.QueryRewriter;
+import org.springframework.data.repository.query.QueryCreationException;
 import org.springframework.data.repository.query.QueryMethodEvaluationContextProvider;
 import org.springframework.data.repository.query.RepositoryQuery;
 import org.springframework.expression.spel.standard.SpelExpressionParser;
@@ -49,6 +50,10 @@ AbstractJpaQuery fromMethodWithQueryString(JpaQueryMethod method, EntityManager
 			@Nullable String countQueryString, QueryRewriter queryRewriter,
 			QueryMethodEvaluationContextProvider evaluationContextProvider) {
 
+		if (method.isScrollQuery()) {
+			throw QueryCreationException.create(method, "Scroll queries are not supported using String-based queries");
+		}
+
 		return method.isNativeQuery()
 				? new NativeJpaQuery(method, em, queryString, countQueryString, queryRewriter, evaluationContextProvider,
 						PARSER)
@@ -64,6 +69,11 @@ AbstractJpaQuery fromMethodWithQueryString(JpaQueryMethod method, EntityManager
 	 * @return
 	 */
 	public StoredProcedureJpaQuery fromProcedureAnnotation(JpaQueryMethod method, EntityManager em) {
+
+		if (method.isScrollQuery()) {
+			throw QueryCreationException.create(method, "Scroll queries are not supported using stored procedures");
+		}
+
 		return new StoredProcedureJpaQuery(method, em);
 	}
 }