add: basic keyset pagination support (without directions)

Signed-off-by: Artemiy Degtyarev <[email protected]>

Author: Artemiy Degtyarev

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/repository/query/PartTreeJdbcQuery.java

@@ -15,17 +15,6 @@
  */
 package org.springframework.data.jdbc.repository.query;
 
-import static org.springframework.data.jdbc.repository.query.JdbcQueryExecution.*;
-
-import java.sql.ResultSet;
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.List;
-import java.util.function.Function;
-import java.util.function.IntFunction;
-import java.util.function.LongSupplier;
-import java.util.function.Supplier;
-
 import org.jspecify.annotations.Nullable;
 import org.springframework.core.convert.converter.Converter;
 import org.springframework.data.domain.*;
@@ -50,6 +39,17 @@
 import org.springframework.jdbc.core.namedparam.SqlParameterSource;
 import org.springframework.util.Assert;
 
+import java.lang.reflect.Field;
+import java.sql.ResultSet;
+import java.util.*;
+import java.util.function.Function;
+import java.util.function.IntFunction;
+import java.util.function.LongSupplier;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
+
+import static org.springframework.data.jdbc.repository.query.JdbcQueryExecution.ResultProcessingConverter;
+
 /**
  * An {@link AbstractJdbcQuery} implementation based on a {@link PartTree}.
  *
@@ -73,61 +73,61 @@ public class PartTreeJdbcQuery extends AbstractJdbcQuery {
 	/**
 	 * Creates a new {@link PartTreeJdbcQuery}.
 	 *
-	 * @param queryMethod must not be {@literal null}.
-	 * @param operations must not be {@literal null}.
+	 * @param queryMethod      must not be {@literal null}.
+	 * @param operations       must not be {@literal null}.
 	 * @param rowMapperFactory must not be {@literal null}.
 	 * @since 4.0
 	 */
 	public PartTreeJdbcQuery(JdbcQueryMethod queryMethod, JdbcAggregateOperations operations,
-			org.springframework.data.jdbc.repository.query.RowMapperFactory rowMapperFactory) {
+							 org.springframework.data.jdbc.repository.query.RowMapperFactory rowMapperFactory) {
 		this(operations.getConverter().getMappingContext(), queryMethod, operations.getDataAccessStrategy().getDialect(),
-				operations.getConverter(), operations.getDataAccessStrategy().getJdbcOperations(), rowMapperFactory);
+			operations.getConverter(), operations.getDataAccessStrategy().getJdbcOperations(), rowMapperFactory);
 	}
 
 	/**
 	 * Creates a new {@link PartTreeJdbcQuery}.
 	 *
-	 * @param queryMethod must not be {@literal null}.
-	 * @param dialect must not be {@literal null}.
-	 * @param converter must not be {@literal null}.
-	 * @param operations must not be {@literal null}.
+	 * @param queryMethod      must not be {@literal null}.
+	 * @param dialect          must not be {@literal null}.
+	 * @param converter        must not be {@literal null}.
+	 * @param operations       must not be {@literal null}.
 	 * @param rowMapperFactory must not be {@literal null}.
 	 */
 	public PartTreeJdbcQuery(JdbcQueryMethod queryMethod, Dialect dialect, JdbcConverter converter,
-			NamedParameterJdbcOperations operations,
-			org.springframework.data.jdbc.repository.query.RowMapperFactory rowMapperFactory) {
+							 NamedParameterJdbcOperations operations,
+							 org.springframework.data.jdbc.repository.query.RowMapperFactory rowMapperFactory) {
 		this(converter.getMappingContext(), queryMethod, dialect, converter, operations, rowMapperFactory);
 	}
 
 	/**
 	 * Creates a new {@link PartTreeJdbcQuery}.
 	 *
-	 * @param context must not be {@literal null}.
+	 * @param context     must not be {@literal null}.
 	 * @param queryMethod must not be {@literal null}.
-	 * @param dialect must not be {@literal null}.
-	 * @param converter must not be {@literal null}.
-	 * @param operations must not be {@literal null}.
-	 * @param rowMapper must not be {@literal null}.
+	 * @param dialect     must not be {@literal null}.
+	 * @param converter   must not be {@literal null}.
+	 * @param operations  must not be {@literal null}.
+	 * @param rowMapper   must not be {@literal null}.
 	 */
 	public PartTreeJdbcQuery(RelationalMappingContext context, JdbcQueryMethod queryMethod, Dialect dialect,
-			JdbcConverter converter, NamedParameterJdbcOperations operations, RowMapper<Object> rowMapper) {
+							 JdbcConverter converter, NamedParameterJdbcOperations operations, RowMapper<Object> rowMapper) {
 		this(context, queryMethod, dialect, converter, operations, it -> rowMapper);
 	}
 
 	/**
 	 * Creates a new {@link PartTreeJdbcQuery}.
 	 *
-	 * @param context must not be {@literal null}.
-	 * @param queryMethod must not be {@literal null}.
-	 * @param dialect must not be {@literal null}.
-	 * @param converter must not be {@literal null}.
-	 * @param operations must not be {@literal null}.
+	 * @param context          must not be {@literal null}.
+	 * @param queryMethod      must not be {@literal null}.
+	 * @param dialect          must not be {@literal null}.
+	 * @param converter        must not be {@literal null}.
+	 * @param operations       must not be {@literal null}.
 	 * @param rowMapperFactory must not be {@literal null}.
 	 * @since 2.3
 	 */
 	public PartTreeJdbcQuery(RelationalMappingContext context, JdbcQueryMethod queryMethod, Dialect dialect,
-			JdbcConverter converter, NamedParameterJdbcOperations operations,
-			org.springframework.data.jdbc.repository.query.RowMapperFactory rowMapperFactory) {
+							 JdbcConverter converter, NamedParameterJdbcOperations operations,
+							 org.springframework.data.jdbc.repository.query.RowMapperFactory rowMapperFactory) {
 
 		super(queryMethod, operations);
 
@@ -146,7 +146,7 @@ public PartTreeJdbcQuery(RelationalMappingContext context, JdbcQueryMethod query
 		JdbcQueryCreator.validate(this.tree, this.parameters, this.converter.getMappingContext());
 
 		this.cachedRowMapperFactory = new CachedRowMapperFactory(tree, rowMapperFactory, converter,
-				queryMethod.getResultProcessor());
+			queryMethod.getResultProcessor());
 	}
 
 	private Sort getDynamicSort(RelationalParameterAccessor accessor) {
@@ -158,7 +158,7 @@ private Sort getDynamicSort(RelationalParameterAccessor accessor) {
 	public Object execute(Object[] values) {
 
 		RelationalParametersParameterAccessor accessor = new RelationalParametersParameterAccessor(getQueryMethod(),
-				values);
+			values);
 
 		if (tree.isDelete()) {
 			JdbcQueryExecution<?> execution = createModifyingQueryExecutor();
@@ -180,17 +180,18 @@ public Object execute(Object[] values) {
 	}
 
 	private JdbcQueryExecution<?> getQueryExecution(ResultProcessor processor,
-			RelationalParametersParameterAccessor accessor) {
+													RelationalParametersParameterAccessor accessor) {
 
 		ResultSetExtractor<Boolean> extractor = tree.isExistsProjection() ? (ResultSet::next) : null;
 		Supplier<RowMapper<?>> rowMapper = parameters.hasDynamicProjection()
-				? () -> cachedRowMapperFactory.getRowMapper(processor)
-				: cachedRowMapperFactory;
+			? () -> cachedRowMapperFactory.getRowMapper(processor)
+			: cachedRowMapperFactory;
 
 		JdbcQueryExecution<?> queryExecution = getJdbcQueryExecution(extractor, rowMapper);
 
 		if (getQueryMethod().isScrollQuery()) {
-			return new ScrollQueryExecution<>((JdbcQueryExecution<Collection<Object>>) queryExecution, accessor.getScrollPosition(), this.tree.getMaxResults());
+			//noinspection unchecked
+			return new ScrollQueryExecution<>((JdbcQueryExecution<Collection<Object>>) queryExecution, accessor.getScrollPosition(), this.tree.getMaxResults(), tree.getSort());
 		}
 
 		if (getQueryMethod().isSliceQuery()) {
@@ -202,23 +203,23 @@ private JdbcQueryExecution<?> getQueryExecution(ResultProcessor processor,
 
 			// noinspection unchecked
 			return new PageQueryExecution<>((JdbcQueryExecution<Collection<Object>>) queryExecution, accessor.getPageable(),
-					() -> {
+				() -> {
 
-						RelationalEntityMetadata<?> entityMetadata = getQueryMethod().getEntityInformation();
+					RelationalEntityMetadata<?> entityMetadata = getQueryMethod().getEntityInformation();
 
-						JdbcCountQueryCreator queryCreator = new JdbcCountQueryCreator(context, tree, converter, dialect,
-								entityMetadata, accessor, false, processor.getReturnedType(), getQueryMethod().lookupLockAnnotation(), false);
+					JdbcCountQueryCreator queryCreator = new JdbcCountQueryCreator(context, tree, converter, dialect,
+						entityMetadata, accessor, false, processor.getReturnedType(), getQueryMethod().lookupLockAnnotation(), false);
 
-						ParametrizedQuery countQuery = queryCreator.createQuery(Sort.unsorted());
-						Object count = singleObjectQuery(new SingleColumnRowMapper<>(Number.class)).execute(countQuery.getQuery(),
-								countQuery.getParameterSource(dialect.getLikeEscaper()));
+					ParametrizedQuery countQuery = queryCreator.createQuery(Sort.unsorted());
+					Object count = singleObjectQuery(new SingleColumnRowMapper<>(Number.class)).execute(countQuery.getQuery(),
+						countQuery.getParameterSource(dialect.getLikeEscaper()));
 
-						Long converted = converter.getConversionService().convert(count, Long.class);
+					Long converted = converter.getConversionService().convert(count, Long.class);
 
-						Assert.state(converted != null, "Count must not be null");
+					Assert.state(converted != null, "Count must not be null");
 
-						return converted;
-					});
+					return converted;
+				});
 		}
 
 		return queryExecution;
@@ -229,7 +230,7 @@ ParametrizedQuery createQuery(RelationalParametersParameterAccessor accessor, Re
 		RelationalEntityMetadata<?> entityMetadata = getQueryMethod().getEntityInformation();
 
 		JdbcQueryCreator queryCreator = new JdbcQueryCreator(context, tree, converter, dialect, entityMetadata, accessor,
-				getQueryMethod().isSliceQuery(), returnedType, this.getQueryMethod().lookupLockAnnotation(), getQueryMethod().isScrollQuery());
+			getQueryMethod().isSliceQuery(), returnedType, this.getQueryMethod().lookupLockAnnotation(), getQueryMethod().isScrollQuery());
 		return queryCreator.createQuery(getDynamicSort(accessor));
 	}
 
@@ -238,12 +239,12 @@ private List<ParametrizedQuery> createDeleteQueries(RelationalParametersParamete
 		RelationalEntityMetadata<?> entityMetadata = getQueryMethod().getEntityInformation();
 
 		JdbcDeleteQueryCreator queryCreator = new JdbcDeleteQueryCreator(context, tree, converter, dialect, entityMetadata,
-				accessor);
+			accessor);
 		return queryCreator.createQuery();
 	}
 
 	private JdbcQueryExecution<?> getJdbcQueryExecution(@Nullable ResultSetExtractor<Boolean> extractor,
-			Supplier<RowMapper<?>> rowMapper) {
+														Supplier<RowMapper<?>> rowMapper) {
 
 		if (getQueryMethod().isPageQuery() || getQueryMethod().isSliceQuery() || getQueryMethod().isScrollQuery()) {
 			return collectionQuery(rowMapper.get());
@@ -260,12 +261,14 @@ private JdbcQueryExecution<?> getJdbcQueryExecution(@Nullable ResultSetExtractor
 	static class ScrollQueryExecution<T> implements JdbcQueryExecution<Window<T>> {
 		private final JdbcQueryExecution<? extends Collection<T>> delegate;
 		private final ScrollPosition position;
-		@Nullable private final Integer maxResults;
+		private final @Nullable Integer maxResults;
+		private final Sort sort;
 
-		ScrollQueryExecution(JdbcQueryExecution<? extends Collection<T>> delegate, ScrollPosition position, @Nullable Integer maxResults) {
+		ScrollQueryExecution(JdbcQueryExecution<? extends Collection<T>> delegate, ScrollPosition position, @Nullable Integer maxResults, Sort sort) {
 			this.delegate = delegate;
 			this.position = position;
 			this.maxResults = maxResults;
+			this.sort = sort;
 		}
 
 		@Override
@@ -277,10 +280,46 @@ static class ScrollQueryExecution<T> implements JdbcQueryExecution<Window<T>> {
 			if (position instanceof OffsetScrollPosition)
 				positionFunction = ((OffsetScrollPosition) position).positionFunction();
 
+			if (position instanceof KeysetScrollPosition) {
+				Map<String, Object> keys = ((KeysetScrollPosition) position).getKeys();
+
+				if (keys.isEmpty()) {
+					List<String> orders = sort.get().map(Sort.Order::getProperty).toList();
+
+					keys = extractKeys(resultList, orders);
+				}
+
+				Map<String, Object> finalKeys = keys;
+				positionFunction = (ignoredI) -> ScrollPosition.of(finalKeys, ((KeysetScrollPosition) position).getDirection()) ;
+			}
+
 			boolean hasNext = resultList.size() >= maxResults;
 
 			return Window.from(resultList, positionFunction, hasNext);
 		}
+
+		private Map<String, Object> extractKeys(List<T> resultList, List<String> orders) {
+			T last = resultList.get(resultList.size() - 1);
+
+			Field[] fields = last.getClass().getDeclaredFields();
+
+			return Arrays
+				.stream(fields)
+				.filter(it -> orders.contains(it.getName()))
+				.peek(it -> it.setAccessible(true))
+				.collect(
+					Collectors.toMap(
+						Field::getName,
+						it -> {
+							try {
+								return it.get(last);
+							} catch (Exception e) {
+								throw new RuntimeException(e);
+							}
+						}
+					)
+				);
+		}
 	}
 
 	/**
@@ -329,7 +368,7 @@ static class PageQueryExecution<T> implements JdbcQueryExecution<Slice<T>> {
 		private final LongSupplier countSupplier;
 
 		PageQueryExecution(JdbcQueryExecution<? extends Collection<T>> delegate, Pageable pageable,
-				LongSupplier countSupplier) {
+						   LongSupplier countSupplier) {
 			this.delegate = delegate;
 			this.pageable = pageable;
 			this.countSupplier = countSupplier;
@@ -341,7 +380,7 @@ public Slice<T> execute(String query, SqlParameterSource parameter) {
 			Collection<T> result = delegate.execute(query, parameter);
 
 			return PageableExecutionUtils.getPage(result instanceof List ? (List<T>) result : new ArrayList<>(result),
-					pageable, countSupplier);
+				pageable, countSupplier);
 		}
 
 	}
@@ -356,18 +395,18 @@ class CachedRowMapperFactory implements Supplier<RowMapper<?>> {
 		private final Function<ResultProcessor, RowMapper<?>> rowMapperFunction;
 
 		public CachedRowMapperFactory(PartTree tree,
-				RowMapperFactory rowMapperFactory, RelationalConverter converter,
-				ResultProcessor defaultResultProcessor) {
+									  RowMapperFactory rowMapperFactory, RelationalConverter converter,
+									  ResultProcessor defaultResultProcessor) {
 
 			this.rowMapperFunction = processor -> {
 
 				if (tree.isCountProjection() || tree.isExistsProjection()) {
 					return rowMapperFactory.create(resolveTypeToRead(processor));
 				}
 				Converter<Object, Object> resultProcessingConverter = new ResultProcessingConverter(processor,
-						converter.getMappingContext(), converter.getEntityInstantiators());
+					converter.getMappingContext(), converter.getEntityInstantiators());
 				return new ConvertingRowMapper(
-						rowMapperFactory.create(processor.getReturnedType().getDomainType()), resultProcessingConverter);
+					rowMapperFactory.create(processor.getReturnedType().getDomainType()), resultProcessingConverter);
 			};
 
 			this.rowMapper = Lazy.of(() -> this.rowMapperFunction.apply(defaultResultProcessor));

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/repository/query/StatementFactory.java

@@ -15,10 +15,7 @@
  */
 package org.springframework.data.jdbc.repository.query;
 
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collection;
-import java.util.List;
+import java.util.*;
 import java.util.function.Predicate;
 
 import org.jspecify.annotations.Nullable;
@@ -208,6 +205,8 @@ public String build(MapSqlParameterSource parameterSource) {
 
 			SelectBuilder.SelectLimitOffset limitOffsetBuilder = createSelectClause(entity, table);
 			SelectBuilder.SelectWhere whereBuilder = applyLimitAndOffset(limitOffsetBuilder);
+			criteria = applyScrollCriteria(criteria, scrollPosition);
+
 			SelectBuilder.SelectOrdered selectOrderBuilder = applyCriteria(criteria, entity, table, parameterSource,
 					whereBuilder);
 			selectOrderBuilder = applyOrderBy(sort, entity, table, selectOrderBuilder);
@@ -222,6 +221,31 @@ public String build(MapSqlParameterSource parameterSource) {
 			return SqlRenderer.create(renderContextFactory.createRenderContext()).render(select);
 		}
 
+		@Nullable Criteria applyScrollCriteria(@Nullable Criteria criteria, @Nullable ScrollPosition position) {
+			if (!(position instanceof KeysetScrollPosition) || position.isInitial() || ((KeysetScrollPosition) position).getKeys().isEmpty())
+				return criteria;
+
+			criteria = criteria == null ? Criteria.empty() : criteria;
+
+			Map<String, Object> keys = ((KeysetScrollPosition) position).getKeys();
+			List<String> columns = new ArrayList<>(keys.keySet());
+			List<Object> values = new ArrayList<>(keys.values());
+
+			Criteria result = null;
+
+			for (int i = 0; i < keys.size(); i++) {
+				Criteria orCriteria = Criteria.where(columns.get(i)).greaterThan(values.get(i));
+
+				for (int j = 0; j < i; j++) {
+					orCriteria = Criteria.where(columns.get(j)).is(values.get(j)).and(orCriteria);
+				}
+
+				result = (result == null) ? orCriteria : result.or(orCriteria);
+			}
+
+			return criteria.and(result);
+		}
+
 		SelectBuilder.SelectOrdered applyOrderBy(Sort sort, RelationalPersistentEntity<?> entity, Table table,
 				SelectBuilder.SelectOrdered selectOrdered) {