Skip to content

Commit

Permalink
Supports scrolling base on keyset without id
Browse files Browse the repository at this point in the history
id shouldn't be added to sort if sort property already provided.

See spring-projectsGH-2996
  • Loading branch information
quaff committed Jun 13, 2023
1 parent 80916d4 commit fa12cf0
Show file tree
Hide file tree
Showing 8 changed files with 259 additions and 28 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -40,6 +38,7 @@
*
* @author Mark Paluch
* @author Christoph Strobl
* @author Yanming Zhou
* @since 3.1
*/
public record KeysetScrollSpecification<T> (KeysetScrollPosition position, Sort sort,
Expand All @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
* Delegate to run {@link ScrollPosition scroll queries} and create result {@link Window}.
*
* @author Mark Paluch
* @author Yanming Zhou
* @since 3.1
*/
public class ScrollDelegate<T> {
Expand Down Expand Up @@ -66,7 +67,9 @@ public Window<T> scroll(Query query, Sort sort, ScrollPosition scrollPosition) {
List<T> result = query.getResultList();

if (scrollPosition instanceof KeysetScrollPosition keyset) {
return createWindow(sort, limit, keyset.getDirection(), entity, result);
// if unsorted then "id:ASC" will be used
Sort sortToUse = KeysetScrollSpecification.createSort(keyset, sort, this.entity);
return createWindow(sortToUse, limit, keyset.getDirection(), this.entity, result);
}

if (scrollPosition instanceof OffsetScrollPosition offset) {
Expand All @@ -80,17 +83,17 @@ private static <T> Window<T> createWindow(Sort sort, int limit, Direction direct
JpaEntityInformation<T, ?> entity, List<T> result) {

KeysetScrollDelegate delegate = KeysetScrollDelegate.of(direction);
List<T> resultsToUse = delegate.postProcessResults(result);
List<T> resultsToUse = delegate.getResultWindow(delegate.postProcessResults(result), limit);

IntFunction<ScrollPosition> positionFunction = value -> {

T object = result.get(value);
// if result other than resultsToUse used here, the index will be unexpected 1-based if direction is BACKWARD
T object = resultsToUse.get(value);
Map<String, Object> keys = entity.getKeyset(sort.stream().map(Order::getProperty).toList(), object);

return ScrollPosition.forward(keys);
// the original direction should retain
return ScrollPosition.of(keys, direction);
};

return Window.from(delegate.getResultWindow(resultsToUse, limit), positionFunction, hasMoreElements(result, limit));
return Window.from(resultsToUse, positionFunction, hasMoreElements(result, limit));
}

private static <T> Window<T> createWindow(List<T> result, int limit,
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
/*
* Copyright 2019-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.domain.sample;

import jakarta.persistence.*;
import lombok.*;
import org.springframework.data.jpa.domain.AbstractPersistable;

/**
* @author Yanming Zhou
*/
@Entity
@Setter
@Getter
public class ScrollableEntity extends AbstractPersistable<Integer> {

private @Column(unique = true) int seqNo;

}
Original file line number Diff line number Diff line change
@@ -0,0 +1,152 @@
/*
* Copyright 2019-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;

import org.junit.jupiter.api.extension.ExtendWith;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.domain.KeysetScrollPosition;
import org.springframework.data.domain.ScrollPosition;
import org.springframework.data.domain.Sort;
import org.springframework.data.domain.Window;
import org.springframework.data.jpa.domain.sample.ScrollableEntity;
import org.springframework.data.jpa.repository.sample.ScrollableEntityRepository;
import org.springframework.lang.Nullable;
import org.springframework.test.context.ContextConfiguration;
import org.springframework.test.context.junit.jupiter.SpringExtension;
import org.springframework.transaction.annotation.Transactional;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.stream.Stream;

import static org.assertj.core.api.Assertions.assertThat;

/**
* @author Yanming Zhou
*/
@ExtendWith(SpringExtension.class)
@ContextConfiguration("classpath:config/namespace-application-context.xml")
@Transactional
class ScrollingIntegrationTests {

private final static int pageSize = 10;

private final static String[][] sortKeys = new String[][] { null, { "id" }, { "seqNo" }, { "id", "seqNo" } };

private final static Integer[] totals = new Integer[] { 0, 5, 10, 15, 20, 25 };

@Autowired
ScrollableEntityRepository repository;

void prepare(int total) {
for (int i = 0; i < total; i++) {
ScrollableEntity entity = new ScrollableEntity();
entity.setSeqNo(i);
this.repository.save(entity);
}
}

@ParameterizedTest
@MethodSource("cartesian")
void scroll(@Nullable String[] keys, Sort.Direction sortDirection, ScrollPosition.Direction scrollDirection, int total) {

prepare(total);

List<List<ScrollableEntity>> contents = new ArrayList<>();

Sort sort;
if (keys != null) {
sort = Sort.by(sortDirection, keys);
}
else {
sort = Sort.unsorted();
sortDirection = Sort.Direction.ASC; // use actual sort "id:ASC" for asserts
}
KeysetScrollPosition position = ScrollPosition.of(Map.of(), scrollDirection);
if (scrollDirection == ScrollPosition.Direction.BACKWARD && position.getDirection() == ScrollPosition.Direction.FORWARD) {
// remove this workaround if https://github.com/spring-projects/spring-data-commons/pull/2841 merged
position = position.backward();
}
while (true) {
ScrollPosition positionToUse = position;
Window<ScrollableEntity> window = this.repository.findBy((root, query, cb) -> null,
q -> q.limit(pageSize).sortBy(sort).scroll(positionToUse));
if (!window.isEmpty()) {
contents.add(window.getContent());
}
if (!window.hasNext()) {
break;
}
int indexForNext = position.scrollsForward() ? window.size() - 1 : 0;
position = (KeysetScrollPosition) window.positionAt(indexForNext);
// position = window.positionForNext(); https://github.com/spring-projects/spring-data-commons/pull/2843
}

if (total == 0) {
assertThat(contents).hasSize(0);
return;
}

boolean divisible = total % pageSize == 0;

assertThat(contents).hasSize(divisible ? total / pageSize : total / pageSize + 1);
for (int i = 0; i < contents.size() - 1; i++) {
assertThat(contents.get(i)).hasSize(pageSize);
}
if (divisible) {
assertThat(contents.get(contents.size() - 1)).hasSize(pageSize);
}
else {
assertThat(contents.get(contents.size() - 1)).hasSize(total % pageSize);
}

List<ScrollableEntity> first = contents.get(0);
List<ScrollableEntity> last = contents.get(contents.size() - 1);

if (sortDirection == Sort.Direction.ASC) {
if (scrollDirection == ScrollPosition.Direction.FORWARD) {
assertThat(first.get(0).getSeqNo()).isEqualTo(0);
assertThat(last.get(last.size() - 1).getSeqNo()).isEqualTo(total - 1);
}
else {
assertThat(first.get(first.size() - 1).getSeqNo()).isEqualTo(total - 1);
assertThat(last.get(0).getSeqNo()).isEqualTo(0);
}
}
else {
if (scrollDirection == ScrollPosition.Direction.FORWARD) {
assertThat(first.get(0).getSeqNo()).isEqualTo(total - 1);
assertThat(last.get(last.size() - 1).getSeqNo()).isEqualTo(0);
}
else {
assertThat(first.get(first.size() - 1).getSeqNo()).isEqualTo(0);
assertThat(last.get(0).getSeqNo()).isEqualTo(total - 1);
}
}
}

private static Stream<Arguments> cartesian() {
return Stream.of(sortKeys)
.flatMap(keys -> Stream.of(Sort.Direction.class.getEnumConstants())
.flatMap(sortDirection -> Stream.of(ScrollPosition.Direction.class.getEnumConstants())
.flatMap(scrollDirection -> Stream.of(totals)
.map(total -> Arguments.of(keys, sortDirection, scrollDirection, total)))));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
* Unit tests for {@link KeysetScrollSpecification}.
*
* @author Mark Paluch
* @author Yanming Zhou
*/
@ExtendWith(SpringExtension.class)
@ContextConfiguration({ "classpath:infrastructure.xml" })
Expand All @@ -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
Expand All @@ -74,4 +95,4 @@ void shouldSkipExistingIdentifiersInSort() {
assertThat(sort).extracting(Order::getProperty).containsExactly("id", "firstname");
}

}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
/*
* Copyright 2019-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.sample;

import org.springframework.data.jpa.domain.sample.ScrollableEntity;
import org.springframework.data.jpa.repository.JpaSpecificationExecutor;
import org.springframework.data.repository.CrudRepository;

/**
* @author Yanming Zhou
*/
public interface ScrollableEntityRepository extends CrudRepository<ScrollableEntity, Integer>, JpaSpecificationExecutor<ScrollableEntity> {

}
4 changes: 2 additions & 2 deletions spring-data-jpa/src/test/resources/META-INF/persistence.xml
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
<class>org.springframework.data.jpa.domain.sample.ConcreteType2</class>
<class>org.springframework.data.jpa.domain.sample.CustomAbstractPersistable</class>
<class>org.springframework.data.jpa.domain.sample.Customer</class>
<class>org.springframework.data.jpa.domain.sample.ScrollableEntity</class>
<class>org.springframework.data.jpa.domain.sample.EntityWithAssignedId</class>
<class>org.springframework.data.jpa.domain.sample.EmbeddedIdExampleEmployeePK</class>
<class>org.springframework.data.jpa.domain.sample.EmbeddedIdExampleEmployee</class>
Expand All @@ -35,8 +36,7 @@
<class>org.springframework.data.jpa.domain.sample.Order</class>
<class>org.springframework.data.jpa.domain.sample.Parent</class>
<class>org.springframework.data.jpa.domain.sample.PersistableWithIdClass</class>
<class>org.springframework.data.jpa.domain.sample.PersistableWithSingleIdClass
</class>
<class>org.springframework.data.jpa.domain.sample.PersistableWithSingleIdClass</class>
<class>org.springframework.data.jpa.domain.sample.PrimitiveVersionProperty</class>
<class>org.springframework.data.jpa.domain.sample.Product</class>
<class>org.springframework.data.jpa.domain.sample.Role</class>
Expand Down
Loading

0 comments on commit fa12cf0

Please sign in to comment.