Skip to content

Commit 945afaf

Browse files
quaffquaff
committed
Supports scrolling base on keyset without id
id shouldn't be added to sort if sort property already provided. See GH-2996
1 parent 80916d4 commit 945afaf

File tree

8 files changed

+261
-28
lines changed

8 files changed

+261
-28
lines changed

Diff for: ‎spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/KeysetScrollSpecification.java

+8-14
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,6 @@
2222
import jakarta.persistence.criteria.Predicate;
2323
import jakarta.persistence.criteria.Root;
2424

25-
import java.util.ArrayList;
26-
import java.util.Collection;
2725
import java.util.List;
2826

2927
import org.springframework.data.domain.KeysetScrollPosition;
@@ -40,6 +38,7 @@
4038
*
4139
* @author Mark Paluch
4240
* @author Christoph Strobl
41+
* @author Yanming Zhou
4342
* @since 3.1
4443
*/
4544
public record KeysetScrollSpecification<T> (KeysetScrollPosition position, Sort sort,
@@ -63,21 +62,16 @@ public static Sort createSort(KeysetScrollPosition position, Sort sort, JpaEntit
6362

6463
KeysetScrollDelegate delegate = KeysetScrollDelegate.of(position.getDirection());
6564

66-
Collection<String> sortById;
67-
Sort sortToUse;
68-
if (entity.hasCompositeId()) {
69-
sortById = new ArrayList<>(entity.getIdAttributeNames());
70-
} else {
71-
sortById = new ArrayList<>(1);
72-
sortById.add(entity.getRequiredIdAttribute().getName());
65+
if (sort.isSorted()) {
66+
// assume sort applied on unique property
67+
return delegate.getSortOrders(sort);
7368
}
7469

75-
sort.forEach(it -> sortById.remove(it.getProperty()));
76-
77-
if (sortById.isEmpty()) {
78-
sortToUse = sort;
70+
Sort sortToUse;
71+
if (entity.hasCompositeId()) {
72+
sortToUse = sort.and(Sort.by(entity.getIdAttributeNames().toArray(new String[0])));
7973
} else {
80-
sortToUse = sort.and(Sort.by(sortById.toArray(new String[0])));
74+
sortToUse = sort.and(Sort.by(entity.getRequiredIdAttribute().getName()));
8175
}
8276

8377
return delegate.getSortOrders(sortToUse);

Diff for: ‎spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/ScrollDelegate.java

+10-7
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
* Delegate to run {@link ScrollPosition scroll queries} and create result {@link Window}.
3636
*
3737
* @author Mark Paluch
38+
* @author Yanming Zhou
3839
* @since 3.1
3940
*/
4041
public class ScrollDelegate<T> {
@@ -66,7 +67,9 @@ public Window<T> scroll(Query query, Sort sort, ScrollPosition scrollPosition) {
6667
List<T> result = query.getResultList();
6768

6869
if (scrollPosition instanceof KeysetScrollPosition keyset) {
69-
return createWindow(sort, limit, keyset.getDirection(), entity, result);
70+
// if unsorted then "id:ASC" will be used
71+
Sort sortToUse = KeysetScrollSpecification.createSort(keyset, sort, this.entity);
72+
return createWindow(sortToUse, limit, keyset.getDirection(), this.entity, result);
7073
}
7174

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

8285
KeysetScrollDelegate delegate = KeysetScrollDelegate.of(direction);
83-
List<T> resultsToUse = delegate.postProcessResults(result);
86+
List<T> resultsToUse = delegate.getResultWindow(delegate.postProcessResults(result), limit);
8487

8588
IntFunction<ScrollPosition> positionFunction = value -> {
86-
87-
T object = result.get(value);
89+
// if result other than resultsToUse used here, the index will be unexpected 1-based if direction is BACKWARD
90+
T object = resultsToUse.get(value);
8891
Map<String, Object> keys = entity.getKeyset(sort.stream().map(Order::getProperty).toList(), object);
89-
90-
return ScrollPosition.forward(keys);
92+
// the original direction should retain
93+
return ScrollPosition.of(keys, direction);
9194
};
9295

93-
return Window.from(delegate.getResultWindow(resultsToUse, limit), positionFunction, hasMoreElements(result, limit));
96+
return Window.from(resultsToUse, positionFunction, hasMoreElements(result, limit));
9497
}
9598

9699
private static <T> Window<T> createWindow(List<T> result, int limit,

Diff for: ‎spring-data-jpa/src/test/java/org/springframework/data/jpa/domain/sample/ScrollableEntity.java

+32
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
/*
2+
* Copyright 2019-2023 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package org.springframework.data.jpa.domain.sample;
17+
18+
import jakarta.persistence.*;
19+
import lombok.*;
20+
import org.springframework.data.jpa.domain.AbstractPersistable;
21+
22+
/**
23+
* @author Yanming Zhou
24+
*/
25+
@Entity
26+
@Setter
27+
@Getter
28+
public class ScrollableEntity extends AbstractPersistable<Integer> {
29+
30+
private @Column(unique = true) int seqNo;
31+
32+
}

Diff for: ‎spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/ScrollingIntegrationTests.java

+154
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
/*
2+
* Copyright 2019-2023 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package org.springframework.data.jpa.repository;
17+
18+
import org.junit.jupiter.api.extension.ExtendWith;
19+
import org.junit.jupiter.params.ParameterizedTest;
20+
import org.junit.jupiter.params.provider.Arguments;
21+
import org.junit.jupiter.params.provider.MethodSource;
22+
import org.springframework.beans.factory.annotation.Autowired;
23+
import org.springframework.data.domain.KeysetScrollPosition;
24+
import org.springframework.data.domain.ScrollPosition;
25+
import org.springframework.data.domain.Sort;
26+
import org.springframework.data.domain.Window;
27+
import org.springframework.data.jpa.domain.sample.ScrollableEntity;
28+
import org.springframework.data.jpa.repository.sample.ScrollableEntityRepository;
29+
import org.springframework.lang.Nullable;
30+
import org.springframework.test.context.ContextConfiguration;
31+
import org.springframework.test.context.junit.jupiter.SpringExtension;
32+
import org.springframework.transaction.annotation.Transactional;
33+
34+
import java.util.ArrayList;
35+
import java.util.List;
36+
import java.util.Map;
37+
import java.util.stream.Stream;
38+
39+
import static org.assertj.core.api.Assertions.assertThat;
40+
import static org.junit.jupiter.api.Assumptions.assumeTrue;
41+
42+
/**
43+
* @author Yanming Zhou
44+
*/
45+
@ExtendWith(SpringExtension.class)
46+
@ContextConfiguration("classpath:config/namespace-application-context.xml")
47+
@Transactional
48+
class ScrollingIntegrationTests {
49+
50+
private final static int pageSize = 10;
51+
52+
private final static String[][] sortKeys = new String[][] { null, { "id" }, { "seqNo" }, { "id", "seqNo" } };
53+
54+
private final static Integer[] totals = new Integer[] { 0, 5, 10, 15, 20, 25 };
55+
56+
@Autowired
57+
ScrollableEntityRepository repository;
58+
59+
void prepare(int total) {
60+
for (int i = 0; i < total; i++) {
61+
ScrollableEntity entity = new ScrollableEntity();
62+
entity.setSeqNo(i);
63+
this.repository.save(entity);
64+
}
65+
}
66+
67+
@ParameterizedTest
68+
@MethodSource("cartesian")
69+
void scroll(@Nullable String[] keys, Sort.Direction sortDirection, ScrollPosition.Direction scrollDirection, int total) {
70+
71+
prepare(total);
72+
73+
List<List<ScrollableEntity>> contents = new ArrayList<>();
74+
75+
Sort sort;
76+
if (keys != null) {
77+
sort = Sort.by(sortDirection, keys);
78+
}
79+
else {
80+
sort = Sort.unsorted();
81+
// implicit "id:ASC" will be used
82+
assumeTrue(sortDirection == Sort.Direction.ASC);
83+
}
84+
KeysetScrollPosition position = ScrollPosition.of(Map.of(), scrollDirection);
85+
if (scrollDirection == ScrollPosition.Direction.BACKWARD && position.getDirection() == ScrollPosition.Direction.FORWARD) {
86+
// remove this workaround if https://github.com/spring-projects/spring-data-commons/pull/2841 merged
87+
position = position.backward();
88+
}
89+
while (true) {
90+
ScrollPosition positionToUse = position;
91+
Window<ScrollableEntity> window = this.repository.findBy((root, query, cb) -> null,
92+
q -> q.limit(pageSize).sortBy(sort).scroll(positionToUse));
93+
if (!window.isEmpty()) {
94+
contents.add(window.getContent());
95+
}
96+
if (!window.hasNext()) {
97+
break;
98+
}
99+
int indexForNext = position.scrollsForward() ? window.size() - 1 : 0;
100+
position = (KeysetScrollPosition) window.positionAt(indexForNext);
101+
// position = window.positionForNext(); https://github.com/spring-projects/spring-data-commons/pull/2843
102+
}
103+
104+
if (total == 0) {
105+
assertThat(contents).hasSize(0);
106+
return;
107+
}
108+
109+
boolean divisible = total % pageSize == 0;
110+
111+
assertThat(contents).hasSize(divisible ? total / pageSize : total / pageSize + 1);
112+
for (int i = 0; i < contents.size() - 1; i++) {
113+
assertThat(contents.get(i)).hasSize(pageSize);
114+
}
115+
if (divisible) {
116+
assertThat(contents.get(contents.size() - 1)).hasSize(pageSize);
117+
}
118+
else {
119+
assertThat(contents.get(contents.size() - 1)).hasSize(total % pageSize);
120+
}
121+
122+
List<ScrollableEntity> first = contents.get(0);
123+
List<ScrollableEntity> last = contents.get(contents.size() - 1);
124+
125+
if (sortDirection == Sort.Direction.ASC) {
126+
if (scrollDirection == ScrollPosition.Direction.FORWARD) {
127+
assertThat(first.get(0).getSeqNo()).isEqualTo(0);
128+
assertThat(last.get(last.size() - 1).getSeqNo()).isEqualTo(total - 1);
129+
}
130+
else {
131+
assertThat(first.get(first.size() - 1).getSeqNo()).isEqualTo(total - 1);
132+
assertThat(last.get(0).getSeqNo()).isEqualTo(0);
133+
}
134+
}
135+
else {
136+
if (scrollDirection == ScrollPosition.Direction.FORWARD) {
137+
assertThat(first.get(0).getSeqNo()).isEqualTo(total - 1);
138+
assertThat(last.get(last.size() - 1).getSeqNo()).isEqualTo(0);
139+
}
140+
else {
141+
assertThat(first.get(first.size() - 1).getSeqNo()).isEqualTo(0);
142+
assertThat(last.get(0).getSeqNo()).isEqualTo(total - 1);
143+
}
144+
}
145+
}
146+
147+
private static Stream<Arguments> cartesian() {
148+
return Stream.of(sortKeys)
149+
.flatMap(keys -> Stream.of(Sort.Direction.class.getEnumConstants())
150+
.flatMap(sortDirection -> Stream.of(ScrollPosition.Direction.class.getEnumConstants())
151+
.flatMap(scrollDirection -> Stream.of(totals)
152+
.map(total -> Arguments.of(keys, sortDirection, scrollDirection, total)))));
153+
}
154+
}

Diff for: ‎spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/KeysetScrollSpecificationUnitTests.java

+26-5
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
* Unit tests for {@link KeysetScrollSpecification}.
3737
*
3838
* @author Mark Paluch
39+
* @author Yanming Zhou
3940
*/
4041
@ExtendWith(SpringExtension.class)
4142
@ContextConfiguration({ "classpath:infrastructure.xml" })
@@ -44,24 +45,44 @@ class KeysetScrollSpecificationUnitTests {
4445

4546
@PersistenceContext EntityManager em;
4647

48+
@Test
49+
void shouldUseIdentifierAsFallback() {
50+
51+
Sort sort = KeysetScrollSpecification.createSort(ScrollPosition.keyset(), Sort.unsorted(),
52+
new JpaMetamodelEntityInformation<>(User.class, em.getMetamodel(),
53+
em.getEntityManagerFactory().getPersistenceUnitUtil()));
54+
55+
assertThat(sort).isEqualTo(Sort.by("id"));
56+
}
57+
58+
@Test
59+
void shouldUseCompositeIdentifierAsFallback() {
60+
61+
Sort sort = KeysetScrollSpecification.createSort(ScrollPosition.keyset(), Sort.unsorted(),
62+
new JpaMetamodelEntityInformation<>(SampleWithIdClass.class, em.getMetamodel(),
63+
em.getEntityManagerFactory().getPersistenceUnitUtil()));
64+
65+
assertThat(sort).isEqualTo(Sort.by("first", "second"));
66+
}
67+
4768
@Test // GH-2996
48-
void shouldAddIdentifierToSort() {
69+
void shouldNotAddIdentifierToSort() {
4970

5071
Sort sort = KeysetScrollSpecification.createSort(ScrollPosition.keyset(), Sort.by("firstname"),
5172
new JpaMetamodelEntityInformation<>(User.class, em.getMetamodel(),
5273
em.getEntityManagerFactory().getPersistenceUnitUtil()));
5374

54-
assertThat(sort).extracting(Order::getProperty).containsExactly("firstname", "id");
75+
assertThat(sort).extracting(Order::getProperty).containsExactly("firstname");
5576
}
5677

5778
@Test // GH-2996
58-
void shouldAddCompositeIdentifierToSort() {
79+
void shouldNotAddCompositeIdentifierToSort() {
5980

6081
Sort sort = KeysetScrollSpecification.createSort(ScrollPosition.keyset(), Sort.by("first", "firstname"),
6182
new JpaMetamodelEntityInformation<>(SampleWithIdClass.class, em.getMetamodel(),
6283
em.getEntityManagerFactory().getPersistenceUnitUtil()));
6384

64-
assertThat(sort).extracting(Order::getProperty).containsExactly("first", "firstname", "second");
85+
assertThat(sort).extracting(Order::getProperty).containsExactly("first", "firstname");
6586
}
6687

6788
@Test // GH-2996
@@ -74,4 +95,4 @@ void shouldSkipExistingIdentifiersInSort() {
7495
assertThat(sort).extracting(Order::getProperty).containsExactly("id", "firstname");
7596
}
7697

77-
}
98+
}

Diff for: ‎spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/sample/ScrollableEntityRepository.java

+27
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
/*
2+
* Copyright 2019-2023 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package org.springframework.data.jpa.repository.sample;
17+
18+
import org.springframework.data.jpa.domain.sample.ScrollableEntity;
19+
import org.springframework.data.jpa.repository.JpaSpecificationExecutor;
20+
import org.springframework.data.repository.CrudRepository;
21+
22+
/**
23+
* @author Yanming Zhou
24+
*/
25+
public interface ScrollableEntityRepository extends CrudRepository<ScrollableEntity, Integer>, JpaSpecificationExecutor<ScrollableEntity> {
26+
27+
}

Diff for: ‎spring-data-jpa/src/test/resources/META-INF/persistence.xml

+2-2
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
<class>org.springframework.data.jpa.domain.sample.ConcreteType2</class>
1919
<class>org.springframework.data.jpa.domain.sample.CustomAbstractPersistable</class>
2020
<class>org.springframework.data.jpa.domain.sample.Customer</class>
21+
<class>org.springframework.data.jpa.domain.sample.ScrollableEntity</class>
2122
<class>org.springframework.data.jpa.domain.sample.EntityWithAssignedId</class>
2223
<class>org.springframework.data.jpa.domain.sample.EmbeddedIdExampleEmployeePK</class>
2324
<class>org.springframework.data.jpa.domain.sample.EmbeddedIdExampleEmployee</class>
@@ -35,8 +36,7 @@
3536
<class>org.springframework.data.jpa.domain.sample.Order</class>
3637
<class>org.springframework.data.jpa.domain.sample.Parent</class>
3738
<class>org.springframework.data.jpa.domain.sample.PersistableWithIdClass</class>
38-
<class>org.springframework.data.jpa.domain.sample.PersistableWithSingleIdClass
39-
</class>
39+
<class>org.springframework.data.jpa.domain.sample.PersistableWithSingleIdClass</class>
4040
<class>org.springframework.data.jpa.domain.sample.PrimitiveVersionProperty</class>
4141
<class>org.springframework.data.jpa.domain.sample.Product</class>
4242
<class>org.springframework.data.jpa.domain.sample.Role</class>

0 commit comments

Comments
 (0)