Skip to content

Commit 50480ca

Browse files
trande4884trande4884
authored
Modifying Page Retrieval To Fix Total Page Count and Improve Overall Functionality of Pagination in Spring (Azure#30694)
* Modifying totalpage calculation for pageable objects. * Adding logic to grab the next set (page) of data when the end of a partition is reached. This will allow us to accurately report the total number of pages being returned and improve functionality to require fetching far less pages in cases where data is split across many partitions. * Fixing star imports. * Adding another unit test. * Fixing issues brought up in PR. Also removing unit test that is no longer applicable. * Updating the change log. * Simplifying logic for while loop. * More logic simplification and adding another unit test. * Removing function as we now only call the functionality once.
1 parent 8211ff5 commit 50480ca

File tree

9 files changed

+251
-82
lines changed

9 files changed

+251
-82
lines changed

‎sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/common/TestConstants.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@
1313
public final class TestConstants {
1414

1515
public static final String AUTOSCALE_MAX_THROUGHPUT = "4000";
16+
17+
public static final String MULTI_PARTITION_THROUGHPUT = "12000";
1618
private static final Address ADDRESS_1 = new Address("201107", "Zixing Road", "Shanghai");
1719
private static final Address ADDRESS_2 = new Address("200000", "Xuhui", "Shanghai");
1820
public static final String HOBBY1 = "photography";

‎sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/CosmosTemplateIT.java

Lines changed: 195 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@
5555

5656
import java.lang.reflect.InvocationTargetException;
5757
import java.lang.reflect.Method;
58+
import java.util.ArrayList;
5859
import java.util.Arrays;
5960
import java.util.Collections;
6061
import java.util.List;
@@ -81,6 +82,7 @@
8182
import static org.junit.Assert.assertEquals;
8283
import static org.junit.Assert.assertNotNull;
8384
import static org.junit.Assert.fail;
85+
import static org.springframework.data.domain.Sort.Direction.ASC;
8486

8587
@RunWith(SpringJUnit4ClassRunner.class)
8688
@ContextConfiguration(classes = TestRepositoryConfig.class)
@@ -440,13 +442,202 @@ public void testFindAllPageableMultiPagesPageSizeTwo() {
440442
final List<Person> expected2 = Lists.newArrayList(TEST_PERSON_3);
441443
assertThat(resultPage2.size()).isEqualTo(expected2.size());
442444
assertThat(resultPage2).containsAll(expected2);
443-
PageTestUtils.validateLastPage(page2, PAGE_SIZE_1);
445+
PageTestUtils.validateLastPage(page2, PAGE_SIZE_2);
446+
447+
assertThat(responseDiagnosticsTestUtils.getCosmosDiagnostics()).isNotNull();
448+
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics()).isNotNull();
449+
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics().getRequestCharge()).isGreaterThan(0);
450+
}
451+
452+
@Test
453+
public void testFindAllPageableMultiPagesMultiPartition() {
454+
cosmosTemplate.insert(TEST_PERSON_2,
455+
new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_2)));
456+
cosmosTemplate.insert(TEST_PERSON_3,
457+
new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_3)));
458+
final List<Person> expected = Lists.newArrayList(TEST_PERSON, TEST_PERSON_2, TEST_PERSON_3);
459+
460+
for (int i=4; i<=10; i++) {
461+
Person temp = new Person("id_" + i, "fred", LAST_NAME + "_" + i, HOBBIES,
462+
ADDRESSES, AGE, PASSPORT_IDS_BY_COUNTRY);
463+
insertPerson(temp);
464+
expected.add(temp);
465+
}
466+
467+
assertThat(responseDiagnosticsTestUtils.getCosmosDiagnostics()).isNotNull();
468+
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics()).isNull();
469+
470+
final CosmosPageRequest pageRequest = new CosmosPageRequest(0, 100, null);
471+
final Page<Person> page1 = cosmosTemplate.findAll(pageRequest, Person.class, containerName);
472+
473+
final List<Person> resultPage1 = TestUtils.toList(page1);
474+
assertThat(resultPage1.size()).isEqualTo(expected.size());
475+
assertThat(resultPage1).containsAll(expected);
476+
PageTestUtils.validateLastPage(page1, 100);
444477

445478
assertThat(responseDiagnosticsTestUtils.getCosmosDiagnostics()).isNotNull();
446479
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics()).isNotNull();
447480
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics().getRequestCharge()).isGreaterThan(0);
448481
}
449482

483+
@Test
484+
public void testFindAllPageableMultiPagesMultiPartition2() {
485+
cosmosTemplate.insert(TEST_PERSON_2,
486+
new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_2)));
487+
cosmosTemplate.insert(TEST_PERSON_3,
488+
new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_3)));
489+
final List<Person> expected = Lists.newArrayList(TEST_PERSON, TEST_PERSON_2, TEST_PERSON_3);
490+
491+
for (int i=4; i<=10; i++) {
492+
Person temp = new Person("id_" + i, "fred", LAST_NAME + "_" + i, HOBBIES,
493+
ADDRESSES, AGE, PASSPORT_IDS_BY_COUNTRY);
494+
insertPerson(temp);
495+
expected.add(temp);
496+
}
497+
498+
assertThat(responseDiagnosticsTestUtils.getCosmosDiagnostics()).isNotNull();
499+
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics()).isNull();
500+
501+
final CosmosPageRequest pageRequest = new CosmosPageRequest(0, 7, null);
502+
final Page<Person> page1 = cosmosTemplate.findAll(pageRequest, Person.class, containerName);
503+
504+
final List<Person> resultPage1 = TestUtils.toList(page1);
505+
assertThat(resultPage1.size()).isEqualTo(7);
506+
PageTestUtils.validateNonLastPage(page1, 7);
507+
508+
assertThat(responseDiagnosticsTestUtils.getCosmosDiagnostics()).isNotNull();
509+
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics()).isNotNull();
510+
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics().getRequestCharge()).isGreaterThan(0);
511+
512+
final Page<Person> page2 = cosmosTemplate.findAll(page1.nextPageable(), Person.class, containerName);
513+
514+
final List<Person> resultPage2 = TestUtils.toList(page2);
515+
assertThat(resultPage2.size()).isEqualTo(3);
516+
PageTestUtils.validateLastPage(page2, 7);
517+
518+
assertThat(responseDiagnosticsTestUtils.getCosmosDiagnostics()).isNotNull();
519+
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics()).isNotNull();
520+
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics().getRequestCharge()).isGreaterThan(0);
521+
522+
final List<Person> allResults = new ArrayList<>();
523+
allResults.addAll(resultPage1);
524+
allResults.addAll(resultPage2);
525+
assertThat(allResults).containsAll(expected);
526+
}
527+
528+
@Test
529+
public void testFindAllPageableMultiPagesMultiPartition3() {
530+
cosmosTemplate.insert(TEST_PERSON_2,
531+
new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_2)));
532+
cosmosTemplate.insert(TEST_PERSON_3,
533+
new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_3)));
534+
final List<Person> expected = Lists.newArrayList(TEST_PERSON, TEST_PERSON_2, TEST_PERSON_3);
535+
536+
for (int i=4; i<=10; i++) {
537+
Person temp = new Person("id_" + i, "fred", LAST_NAME + "_" + i, HOBBIES,
538+
ADDRESSES, AGE, PASSPORT_IDS_BY_COUNTRY);
539+
insertPerson(temp);
540+
expected.add(temp);
541+
}
542+
543+
assertThat(responseDiagnosticsTestUtils.getCosmosDiagnostics()).isNotNull();
544+
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics()).isNull();
545+
546+
final CosmosPageRequest pageRequest = new CosmosPageRequest(0, 3, null);
547+
final Page<Person> page1 = cosmosTemplate.findAll(pageRequest, Person.class, containerName);
548+
549+
final List<Person> resultPage1 = TestUtils.toList(page1);
550+
assertThat(resultPage1.size()).isEqualTo(3);
551+
PageTestUtils.validateNonLastPage(page1, 3);
552+
553+
assertThat(responseDiagnosticsTestUtils.getCosmosDiagnostics()).isNotNull();
554+
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics()).isNotNull();
555+
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics().getRequestCharge()).isGreaterThan(0);
556+
557+
final Page<Person> page2 = cosmosTemplate.findAll(page1.nextPageable(), Person.class, containerName);
558+
559+
final List<Person> resultPage2 = TestUtils.toList(page2);
560+
assertThat(resultPage2.size()).isEqualTo(3);
561+
PageTestUtils.validateNonLastPage(page2, 3);
562+
563+
assertThat(responseDiagnosticsTestUtils.getCosmosDiagnostics()).isNotNull();
564+
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics()).isNotNull();
565+
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics().getRequestCharge()).isGreaterThan(0);
566+
567+
final Page<Person> page3 = cosmosTemplate.findAll(page2.nextPageable(), Person.class, containerName);
568+
569+
final List<Person> resultPage3 = TestUtils.toList(page3);
570+
assertThat(resultPage3.size()).isEqualTo(3);
571+
PageTestUtils.validateNonLastPage(page3, 3);
572+
573+
assertThat(responseDiagnosticsTestUtils.getCosmosDiagnostics()).isNotNull();
574+
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics()).isNotNull();
575+
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics().getRequestCharge()).isGreaterThan(0);
576+
577+
final Page<Person> page4 = cosmosTemplate.findAll(page3.nextPageable(), Person.class, containerName);
578+
579+
final List<Person> resultPage4 = TestUtils.toList(page4);
580+
assertThat(resultPage4.size()).isEqualTo(1);
581+
PageTestUtils.validateLastPage(page4, 3);
582+
583+
assertThat(responseDiagnosticsTestUtils.getCosmosDiagnostics()).isNotNull();
584+
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics()).isNotNull();
585+
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics().getRequestCharge()).isGreaterThan(0);
586+
587+
final List<Person> allResults = new ArrayList<>();
588+
allResults.addAll(resultPage1);
589+
allResults.addAll(resultPage2);
590+
allResults.addAll(resultPage3);
591+
allResults.addAll(resultPage4);
592+
assertThat(allResults).containsAll(expected);
593+
}
594+
595+
@Test
596+
public void testFindAllPageableMultiPagesMultiPartitionWithOffset() {
597+
cosmosTemplate.insert(TEST_PERSON_2,
598+
new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_2)));
599+
cosmosTemplate.insert(TEST_PERSON_3,
600+
new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_3)));
601+
final List<Person> expected = Lists.newArrayList(TEST_PERSON_2, TEST_PERSON_3);
602+
603+
for (int i=4; i<=10; i++) {
604+
Person temp = new Person("id_" + i, "fred", LAST_NAME + "_" + i, HOBBIES,
605+
ADDRESSES, AGE, PASSPORT_IDS_BY_COUNTRY);
606+
insertPerson(temp);
607+
expected.add(temp);
608+
}
609+
610+
assertThat(responseDiagnosticsTestUtils.getCosmosDiagnostics()).isNotNull();
611+
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics()).isNull();
612+
613+
final CosmosPageRequest pageRequest = CosmosPageRequest.of(1, 0, 7,
614+
null, Sort.by(ASC, "id"));
615+
final Page<Person> page1 = cosmosTemplate.findAll(pageRequest, Person.class, containerName);
616+
617+
final List<Person> resultPage1 = TestUtils.toList(page1);
618+
assertThat(resultPage1.size()).isEqualTo(7);
619+
PageTestUtils.validateNonLastPage(page1, 7);
620+
621+
assertThat(responseDiagnosticsTestUtils.getCosmosDiagnostics()).isNotNull();
622+
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics()).isNotNull();
623+
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics().getRequestCharge()).isGreaterThan(0);
624+
625+
final Page<Person> page2 = cosmosTemplate.findAll(page1.nextPageable(), Person.class, containerName);
626+
627+
final List<Person> resultPage2 = TestUtils.toList(page2);
628+
assertThat(resultPage2.size()).isEqualTo(2);
629+
PageTestUtils.validateLastPage(page2, 7);
630+
631+
assertThat(responseDiagnosticsTestUtils.getCosmosDiagnostics()).isNotNull();
632+
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics()).isNotNull();
633+
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics().getRequestCharge()).isGreaterThan(0);
634+
635+
final List<Person> allResults = new ArrayList<>();
636+
allResults.addAll(resultPage1);
637+
allResults.addAll(resultPage2);
638+
assertThat(allResults).containsAll(expected);
639+
}
640+
450641
@Test
451642
public void testPaginationQuery() {
452643
cosmosTemplate.insert(TEST_PERSON_2,
@@ -462,7 +653,7 @@ public void testPaginationQuery() {
462653

463654
final Page<Person> page = cosmosTemplate.paginationQuery(query, Person.class, containerName);
464655
assertThat(page.getContent().size()).isEqualTo(1);
465-
PageTestUtils.validateLastPage(page, page.getContent().size());
656+
PageTestUtils.validateLastPage(page, PAGE_SIZE_2);
466657

467658
// add ignore case testing
468659
final Criteria criteriaIgnoreCase = Criteria.getInstance(CriteriaType.IS_EQUAL, "firstName",
@@ -472,7 +663,7 @@ public void testPaginationQuery() {
472663
final Page<Person> pageIgnoreCase = cosmosTemplate.paginationQuery(queryIgnoreCase, Person.class,
473664
containerName);
474665
assertThat(pageIgnoreCase.getContent().size()).isEqualTo(1);
475-
PageTestUtils.validateLastPage(pageIgnoreCase, pageIgnoreCase.getContent().size());
666+
PageTestUtils.validateLastPage(pageIgnoreCase, PAGE_SIZE_2);
476667

477668
assertThat(responseDiagnosticsTestUtils.getCosmosDiagnostics()).isNotNull();
478669
assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics()).isNotNull();
@@ -601,7 +792,7 @@ public void testFindAllWithTwoPagesAndVerifySortOrder() {
601792
containerName);
602793

603794
assertThat(secondPage.getContent().size()).isEqualTo(2);
604-
PageTestUtils.validateLastPage(secondPage, secondPage.getContent().size());
795+
PageTestUtils.validateLastPage(secondPage, PAGE_SIZE_3);
605796

606797
final List<Person> secondPageResults = secondPage.getContent();
607798
assertThat(secondPageResults.get(0).getFirstName()).isEqualTo(NEW_FIRST_NAME);

‎sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/CosmosTemplatePartitionIT.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -340,7 +340,7 @@ public void testPartitionedPaginationQuery() {
340340

341341
final Page<PartitionPerson> page = cosmosTemplate.paginationQuery(query, PartitionPerson.class, containerName);
342342
assertThat(page.getContent().size()).isEqualTo(1);
343-
PageTestUtils.validateLastPage(page, page.getContent().size());
343+
PageTestUtils.validateLastPage(page, PAGE_SIZE_2);
344344
}
345345

346346
@Test
@@ -356,7 +356,7 @@ public void testPartitionedPaginationQueryWithOneOffset() {
356356
final Page<PartitionPerson> page = cosmosTemplate.paginationQuery(query, PartitionPerson.class, containerName);
357357
assertThat(page.getContent().size()).isEqualTo(PAGE_SIZE_1);
358358
assertThat(page.getContent().get(0).getId()).isEqualTo(ID_3);
359-
PageTestUtils.validateLastPage(page, page.getContent().size());
359+
PageTestUtils.validateLastPage(page, PAGE_SIZE_2);
360360
}
361361

362362
@Test
@@ -423,6 +423,6 @@ public void testPartitionedPaginationQueryIgnoreCase() {
423423
final Page<PartitionPerson> pageIgnoreCase = cosmosTemplate
424424
.paginationQuery(queryIgnoreCase, PartitionPerson.class, containerName);
425425
assertThat(pageIgnoreCase.getContent().size()).isEqualTo(1);
426-
PageTestUtils.validateLastPage(pageIgnoreCase, pageIgnoreCase.getContent().size());
426+
PageTestUtils.validateLastPage(pageIgnoreCase, PAGE_SIZE_2);
427427
}
428428
}

‎sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/domain/Person.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
package com.azure.spring.data.cosmos.domain;
55

6+
import com.azure.spring.data.cosmos.common.TestConstants;
67
import com.azure.spring.data.cosmos.core.mapping.Container;
78
import com.azure.spring.data.cosmos.core.mapping.CosmosIndexingPolicy;
89
import com.azure.spring.data.cosmos.core.mapping.PartitionKey;
@@ -12,7 +13,7 @@
1213
import java.util.Map;
1314
import java.util.Objects;
1415

15-
@Container()
16+
@Container(ru = TestConstants.MULTI_PARTITION_THROUGHPUT, autoScale = true)
1617
@CosmosIndexingPolicy()
1718
public class Person {
1819
private String id;

‎sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/repository/integration/PageableAddressRepositoryIT.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ public void testFindAllByPage() {
8585

8686
final Page<Address> nextPage = repository.findAll(page.nextPageable());
8787
assertThat(nextPage.getContent().size()).isLessThanOrEqualTo(PAGE_SIZE_3);
88-
validateLastPage(nextPage, nextPage.getContent().size());
88+
validateLastPage(nextPage, PAGE_SIZE_3);
8989
}
9090

9191
@Test
@@ -95,7 +95,7 @@ public void testFindWithPartitionKeySinglePage() {
9595

9696
assertThat(page.getContent().size()).isEqualTo(2);
9797
validateResultCityMatch(page, TestConstants.CITY);
98-
validateLastPage(page, page.getContent().size());
98+
validateLastPage(page, PAGE_SIZE_3);
9999
}
100100

101101
@Test
@@ -121,7 +121,7 @@ public void testFindWithoutPartitionKeySinglePage() {
121121

122122
assertThat(page.getContent().size()).isEqualTo(2);
123123
validateResultStreetMatch(page, TestConstants.STREET);
124-
validateLastPage(page, page.getContent().size());
124+
validateLastPage(page, PAGE_SIZE_3);
125125
}
126126

127127
@Test

‎sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/repository/integration/PageablePersonRepositoryIT.java

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -80,14 +80,6 @@ public void setUp() {
8080
// With documents more than 10 KB, and collection RUs 400,
8181
// it usually return documents less than total 12 documents.
8282

83-
// This test covers the case where page size is greater than returned documents limit
84-
@Test
85-
public void testFindAllWithPageSizeGreaterThanReturnedLimit() {
86-
final Set<PageablePerson> outputSet = findAllWithPageSize(15, true);
87-
boolean equals = outputSet.equals(personSet);
88-
assertThat(equals).isTrue();
89-
}
90-
9183
// This test covers the case where page size is less than returned documents limit
9284
@Test
9385
public void testFindAllWithPageSizeLessThanReturnedLimit() {

‎sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/repository/integration/ProjectRepositorySortIT.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -227,7 +227,7 @@ public void testFindWithPageableAndSort() {
227227

228228
Assert.assertEquals(references.size(), result.getContent().size());
229229
Assert.assertEquals(references, result.getContent());
230-
validateLastPage(result, result.getContent().size());
230+
validateLastPage(result, 5);
231231
}
232232
}
233233

‎sdk/cosmos/azure-spring-data-cosmos/CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
#### Breaking Changes
88

99
#### Bugs Fixed
10+
* Fixing pagination bug when performing a cross-partition query to fill every page and fix the total page count reporting. - See [PR 30694](https://github.com/Azure/azure-sdk-for-java/pull/30694)
1011

1112
#### Other Changes
1213

0 commit comments

Comments
 (0)