Adding sum support for annotated queries

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

@@ -4,6 +4,7 @@
 
 #### Features Added
 * Improved the Exception Handling of 'azure-spring-data-cosmos' to throw more detailed exceptions and not always the same exception - See [PR 42902](https://github.com/Azure/azure-sdk-for-java/pull/42902).
+* Implemented sum() support for annotated queries - See [PR 43239](https://github.com/Azure/azure-sdk-for-java/pull/43239).
 
 #### Breaking Changes
 

sdk/spring/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/CosmosOperations.java

@@ -367,6 +367,16 @@ public interface CosmosOperations {
      */
     <T> long count(SqlQuerySpec querySpec, String containerName);
 
+    /**
+     * Sum
+     *
+     * @param querySpec the document query spec
+     * @param containerName the container name
+     * @param <T> type class of domainType
+     * @return sum result
+     */
+    <T> long sum(SqlQuerySpec querySpec, String containerName);
+
     /**
      * To get converter
      *

sdk/spring/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/CosmosTemplate.java

@@ -1008,7 +1008,7 @@ private <T> Page<T> paginationQuery(SqlQuerySpec querySpec, SqlQuerySpec countQu
                                         Optional<Object> partitionKeyValue) {
         containerName = getContainerNameOverride(containerName);
         Slice<T> response = sliceQuery(querySpec, pageable, sort, returnType, containerName, partitionKeyValue);
-        final long total = getCountValue(countQuerySpec, containerName);
+        final long total = getNumericValue(countQuerySpec, containerName);
         return new CosmosPageImpl<>(response.getContent(), response.getPageable(), total);
     }
 
@@ -1132,12 +1132,21 @@ public <T> long count(CosmosQuery query, String containerName) {
 
     @Override
     public <T> long count(SqlQuerySpec querySpec, String containerName) {
+        return this.numeric(querySpec, containerName);
+    }
+
+    @Override
+    public <T> long sum(SqlQuerySpec querySpec, String containerName) {
+        return this.numeric(querySpec, containerName);
+    }
+
+    private <T> long numeric(SqlQuerySpec querySpec, String containerName) {
         containerName = getContainerNameOverride(containerName);
         Assert.hasText(containerName, "container name should not be empty");
 
-        final Long count = getCountValue(querySpec, containerName);
-        assert count != null;
-        return count;
+        final Long numericResult = getNumericValue(querySpec, containerName);
+        assert numericResult != null;
+        return numericResult;
     }
 
     @Override
@@ -1167,10 +1176,10 @@ private void markAuditedIfConfigured(Object object) {
 
     private Long getCountValue(CosmosQuery query, String containerName) {
         final SqlQuerySpec querySpec = new CountQueryGenerator().generateCosmos(query);
-        return getCountValue(querySpec, containerName);
+        return getNumericValue(querySpec, containerName);
     }
 
-    private Long getCountValue(SqlQuerySpec querySpec, String containerName) {
+    private Long getNumericValue(SqlQuerySpec querySpec, String containerName) {
         final CosmosQueryRequestOptions options = new CosmosQueryRequestOptions();
         options.setQueryMetricsEnabled(this.queryMetricsEnabled);
         options.setIndexMetricsEnabled(this.indexMetricsEnabled);

sdk/spring/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/ReactiveCosmosOperations.java

@@ -336,6 +336,15 @@ Mono<CosmosContainerProperties> replaceContainerProperties(String containerName,
      */
     Mono<Long> count(SqlQuerySpec querySpec, String containerName);
 
+    /**
+     * Sum
+     *
+     * @param querySpec the document query spec
+     * @param containerName the container name
+     * @return sum result
+     */
+    Mono<Long> sum(SqlQuerySpec querySpec, String containerName);
+
     /**
      * To get converter
      * @return MappingCosmosConverter

sdk/spring/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/ReactiveCosmosTemplate.java

@@ -920,7 +920,7 @@ public Mono<Long> count(String containerName) {
     @Override
     public Mono<Long> count(CosmosQuery query, String containerName) {
         final SqlQuerySpec querySpec = new CountQueryGenerator().generateCosmos(query);
-        return getCountValue(querySpec, containerName);
+        return getNumericValue(querySpec, containerName);
     }
 
     /**
@@ -932,7 +932,19 @@ public Mono<Long> count(CosmosQuery query, String containerName) {
      */
     @Override
     public Mono<Long> count(SqlQuerySpec querySpec, String containerName) {
-        return getCountValue(querySpec, containerName);
+        return getNumericValue(querySpec, containerName);
+    }
+
+    /**
+     * Sum
+     *
+     * @param querySpec the document query spec
+     * @param containerName the container name
+     * @return Mono with sum or error
+     */
+    @Override
+    public Mono<Long> sum(SqlQuerySpec querySpec, String containerName) {
+        return getNumericValue(querySpec, containerName);
     }
 
     @Override
@@ -975,7 +987,7 @@ private Flux<JsonNode> runQuery(SqlQuerySpec querySpec, Class<?> domainType) {
                                           this.responseDiagnosticsProcessor));
     }
 
-    private Mono<Long> getCountValue(SqlQuerySpec querySpec, String containerName) {
+    private Mono<Long> getNumericValue(SqlQuerySpec querySpec, String containerName) {
         final CosmosQueryRequestOptions options = new CosmosQueryRequestOptions();
         options.setQueryMetricsEnabled(this.queryMetricsEnabled);
         options.setIndexMetricsEnabled(this.indexMetricsEnabled);
@@ -987,7 +999,7 @@ private Mono<Long> getCountValue(SqlQuerySpec querySpec, String containerName) {
             .doOnNext(feedResponse -> CosmosUtils.fillAndProcessResponseDiagnostics(this.responseDiagnosticsProcessor,
                 feedResponse.getCosmosDiagnostics(), feedResponse))
             .onErrorResume(throwable ->
-                CosmosExceptionUtils.exceptionHandler("Failed to get count value", throwable,
+                CosmosExceptionUtils.exceptionHandler("Failed to get numeric value", throwable,
                     this.responseDiagnosticsProcessor))
             .next()
             .map(r -> r.getResults().get(0).asLong());

sdk/spring/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/repository/support/StringBasedCosmosQuery.java

@@ -31,6 +31,7 @@
  */
 public class StringBasedCosmosQuery extends AbstractCosmosQuery {
     private static final Pattern COUNT_QUERY_PATTERN = Pattern.compile("^\\s*select\\s+value\\s+count.*", Pattern.CASE_INSENSITIVE);
+    private static final Pattern SUM_QUERY_PATTERN = Pattern.compile("^\\s*select\\s+value\\s+sum.*", Pattern.CASE_INSENSITIVE);
 
     private final String query;
 
@@ -104,9 +105,12 @@ public Object execute(final Object[] parameters) {
         } else if (isCountQuery()) {
             final String container = ((CosmosEntityMetadata<?>) getQueryMethod().getEntityInformation()).getContainerName();
             return this.operations.count(querySpec, container);
+        } else if (isSumQuery()) {
+            final String container = ((CosmosEntityMetadata<?>) getQueryMethod().getEntityInformation()).getContainerName();
+            return this.operations.sum(querySpec, container);
         } else {
             return this.operations.runQuery(querySpec, accessor.getSort(), processor.getReturnedType().getDomainType(),
-                                            processor.getReturnedType().getReturnedType());
+                processor.getReturnedType().getReturnedType());
         }
     }
 
@@ -129,15 +133,27 @@ protected boolean isCountQuery() {
         return isCountQuery(query, getQueryMethod().getReturnedObjectType());
     }
 
+    protected boolean isSumQuery() {
+        return isSumQuery(query, getQueryMethod().getReturnedObjectType());
+    }
+
     static boolean isCountQuery(String query, Class<?> returnedType) {
-        if (isCountQueryReturnType(returnedType)) {
+        if (isNumericQueryReturnType(returnedType)) {
             return COUNT_QUERY_PATTERN.matcher(query).matches();
         } else {
             return false;
         }
     }
 
-    private static boolean isCountQueryReturnType(Class<?> returnedType) {
+    static boolean isSumQuery(String query, Class<?> returnedType) {
+        if (isNumericQueryReturnType(returnedType)) {
+            return SUM_QUERY_PATTERN.matcher(query).matches();
+        } else {
+            return false;
+        }
+    }
+
+    private static boolean isNumericQueryReturnType(Class<?> returnedType) {
         return returnedType == Long.class
             || returnedType == long.class
             || returnedType == Integer.class

sdk/spring/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/repository/support/StringBasedReactiveCosmosQuery.java

@@ -16,7 +16,6 @@
 import org.springframework.data.repository.query.Parameter;
 import org.springframework.data.repository.query.ResultProcessor;
 import reactor.core.publisher.Flux;
-import reactor.core.publisher.Mono;
 
 import java.util.ArrayList;
 import java.util.Collection;
@@ -96,8 +95,10 @@ public Object execute(final Object[] parameters) {
         SqlQuerySpec querySpec = new SqlQuerySpec(stripExtraWhitespaceFromString(expandedQuery), sqlParameters);
         if (isCountQuery()) {
             final String container = ((SimpleReactiveCosmosEntityMetadata<?>) getQueryMethod().getEntityInformation()).getContainerName();
-            final Mono<Long> mono = this.operations.count(querySpec, container);
-            return mono;
+            return this.operations.count(querySpec, container);
+        } else if (isSumQuery()) {
+            final String container = ((SimpleReactiveCosmosEntityMetadata<?>) getQueryMethod().getEntityInformation()).getContainerName();
+            return this.operations.sum(querySpec, container);
         } else {
             Flux<?> flux = this.operations.runQuery(querySpec, accessor.getSort(), processor.getReturnedType().getDomainType(),
                                                     processor.getReturnedType().getReturnedType());
@@ -123,4 +124,8 @@ protected boolean isCountQuery() {
         return StringBasedCosmosQuery.isCountQuery(query, getQueryMethod().getReturnedObjectType());
     }
 
+    protected boolean isSumQuery() {
+        return StringBasedCosmosQuery.isSumQuery(query, getQueryMethod().getReturnedObjectType());
+    }
+
 }

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

@@ -248,6 +248,18 @@ public void testAnnotatedQueryWithValueAsList() {
         assertAddressPostalCodesUnordered(postalCodes, addresses);
     }
 
+    @Test
+    public void testAnnotatedQueryWithSum() {
+        Address.TEST_ADDRESS1_PARTITION1.setLongId(1L);
+        Address.TEST_ADDRESS2_PARTITION1.setLongId(2L);
+        final List<Address> addresses = Arrays.asList(Address.TEST_ADDRESS1_PARTITION1, Address.TEST_ADDRESS2_PARTITION1);
+        addressRepository.saveAll(addresses);
+
+        final Long sumResult = addressRepository.annotatedSumLongIdValuesByCity(TestConstants.CITY);
+
+        assertThat(sumResult).isEqualTo(3);
+    }
+
     @Test
     public void testAnnotatedQueryWithJsonNodeAsPage() {
         final List<Address> addresses = Arrays.asList(Address.TEST_ADDRESS1_PARTITION1, Address.TEST_ADDRESS2_PARTITION1);

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

@@ -35,8 +35,11 @@ public class ReactiveLongIdDomainPartitionPartitionRepositoryIT {
     private static final Long ID_2 = 67890L;
     private static final String NAME_2 = "camille";
 
+    private static final Long ID_3 = 98765L;
+
     private static final LongIdDomainPartition DOMAIN_1 = new LongIdDomainPartition(ID_1, NAME_1);
     private static final LongIdDomainPartition DOMAIN_2 = new LongIdDomainPartition(ID_2, NAME_2);
+    private static final LongIdDomainPartition DOMAIN_3 = new LongIdDomainPartition(ID_3, NAME_1);
 
     @ClassRule
     public static final ReactiveIntegrationTestCollectionManager collectionManager = new ReactiveIntegrationTestCollectionManager();
@@ -192,6 +195,18 @@ public void testFindAllSort() {
         StepVerifier.create(descAllFlux).expectNext(DOMAIN_2, other, DOMAIN_1).verifyComplete();
     }
 
+    @Test
+    public void testSum() {
+        Mono<Long> sum1 = this.repository.annotatedSumLongIdValuesByName(NAME_1);
+        StepVerifier.create(sum1).expectNext(12345L).verifyComplete();
+
+        Mono<LongIdDomainPartition> saveMono = this.repository.save(DOMAIN_3);
+        StepVerifier.create(saveMono).expectNext(DOMAIN_3).expectComplete().verify();
+
+        Mono<Long> sum2 = this.repository.annotatedSumLongIdValuesByName(NAME_1);
+        StepVerifier.create(sum2).expectNext(111110L).verifyComplete();
+    }
+
     private static class InvalidDomain {
 
         private long count;

sdk/spring/azure-spring-data-cosmos/src/test/java/com/azure/spring/data/cosmos/repository/repository/AddressRepository.java

@@ -59,6 +59,9 @@ public interface AddressRepository extends CosmosRepository<Address, String> {
     @Query("select DISTINCT value a.postalCode from a where a.city = @city")
     List<String> annotatedFindPostalCodeValuesByCity(@Param("city") String city);
 
+    @Query("SELECT VALUE SUM(a.longId) from a where a.city = @city")
+    Long annotatedSumLongIdValuesByCity(@Param("city") String city);
+
     @Query(value = "select * from a where a.city IN (@cities)")
     List<Address> annotatedFindByCityIn(@Param("cities") List<String> cities, Sort sort);
 

sdk/spring/azure-spring-data-cosmos/src/test/java/com/azure/spring/data/cosmos/repository/repository/ReactiveLongIdDomainPartitionRepository.java

@@ -3,10 +3,16 @@
 package com.azure.spring.data.cosmos.repository.repository;
 
 import com.azure.spring.data.cosmos.domain.LongIdDomainPartition;
+import com.azure.spring.data.cosmos.repository.Query;
 import com.azure.spring.data.cosmos.repository.ReactiveCosmosRepository;
+import org.springframework.data.repository.query.Param;
 import org.springframework.stereotype.Repository;
+import reactor.core.publisher.Mono;
 
 @Repository
 public interface ReactiveLongIdDomainPartitionRepository extends ReactiveCosmosRepository<LongIdDomainPartition, Long> {
 
+    @Query("SELECT VALUE SUM(a.number) from a where a.name = @name")
+    Mono<Long> annotatedSumLongIdValuesByName(@Param("name") String name);
+
 }