Add slice query support in string based cosmos query (#25011)

* addSliceQuerysupportForStringQuery * add tests Co-authored-by: annie-mac <[email protected]>
Azure · Oct 26, 2021 · 1040fff · 1040fff
1 parent 596d3e5
commit 1040fff
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 0 deletions.
diff --git a/...t/src/test/java/com/azure/spring/data/cosmos/repository/integration/AnnotatedQueryIT.java b/...t/src/test/java/com/azure/spring/data/cosmos/repository/integration/AnnotatedQueryIT.java
@@ -19,6 +19,7 @@
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.data.domain.Page;
 import org.springframework.data.domain.PageRequest;
+import org.springframework.data.domain.Slice;
 import org.springframework.data.domain.Sort;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
@@ -150,6 +151,21 @@ public void testAnnotatedQueryWithJsonNodeAsPage() {
         assertAddressPostalCodes(actualPostalCodes, addresses);
     }
 
+    @Test
+    public void testAnnotatedQueryWithJsonNodeAsSlice() {
+        final List<Address> addresses = Arrays.asList(Address.TEST_ADDRESS1_PARTITION1, Address.TEST_ADDRESS2_PARTITION1);
+        addressRepository.saveAll(addresses);
+
+        final PageRequest cosmosPageRequest = CosmosPageRequest.of(0, 10);
+        final Slice<JsonNode> postalCodes = addressRepository.annotatedFindPostalCodesByCityAsSlice(TestConstants.CITY,
+            cosmosPageRequest);
+        final List<String> actualPostalCodes = postalCodes.getContent()
+            .stream()
+            .map(jsonNode -> jsonNode.get("postalCode").asText())
+            .collect(Collectors.toList());
+        assertAddressPostalCodes(actualPostalCodes, addresses);
+    }
+
     private void assertAddressPostalCodes(List<String> postalCodes, List<Address> expectedResults) {
         List<String> expectedPostalCodes = expectedResults.stream()
                                                           .map(Address::getPostalCode)

diff --git a/...t/src/test/java/com/azure/spring/data/cosmos/repository/repository/AddressRepository.java b/...t/src/test/java/com/azure/spring/data/cosmos/repository/repository/AddressRepository.java
@@ -8,6 +8,7 @@
 import com.fasterxml.jackson.databind.JsonNode;
 import org.springframework.data.domain.Page;
 import org.springframework.data.domain.Pageable;
+import org.springframework.data.domain.Slice;
 import org.springframework.data.domain.Sort;
 import org.springframework.data.repository.query.Param;
 import org.springframework.stereotype.Repository;
@@ -47,6 +48,9 @@ public interface AddressRepository extends CosmosRepository<Address, String> {
     @Query("select DISTINCT a.postalCode from a where a.city = @city")
     Page<JsonNode> annotatedFindPostalCodesByCity(@Param("city") String city, Pageable pageable);
 
+    @Query("select DISTINCT a.postalCode from a where a.city = @city")
+    Slice<JsonNode> annotatedFindPostalCodesByCityAsSlice(@Param("city") String city, Pageable pageable);
+
     @Query("select DISTINCT value a.postalCode from a where a.city = @city")
     List<String> annotatedFindPostalCodeValuesByCity(@Param("city") String city);
 }
diff --git a/...src/main/java/com/azure/spring/data/cosmos/repository/support/StringBasedCosmosQuery.java b/...src/main/java/com/azure/spring/data/cosmos/repository/support/StringBasedCosmosQuery.java
@@ -59,6 +59,12 @@ public Object execute(final Object[] parameters) {
         if (isPageQuery()) {
             return this.operations.runPaginationQuery(querySpec, accessor.getPageable(), processor.getReturnedType().getDomainType(),
                                                       processor.getReturnedType().getReturnedType());
+        } else if (isSliceQuery()) {
+            return this.operations.runSliceQuery(
+                querySpec,
+                accessor.getPageable(),
+                processor.getReturnedType().getDomainType(),
+                processor.getReturnedType().getReturnedType());
         } else if (isCountQuery()) {
             final String container = ((CosmosEntityMetadata<?>) getQueryMethod().getEntityInformation()).getContainerName();
             return this.operations.count(querySpec, container);